cheatah
Source BigBrain
Source

stdlib/linalg/backend.hpp

1// Copyright (c) 2026 BigBrain LLC. MIT-licensed (see LICENSE).
2// Original work; see ACKNOWLEDGMENTS.md for the open-source ideas we build upon.
3#pragma once
5/**
6 * @file backend.hpp
7 * @brief cheatah `linalg` — the two-layer (element T + container Array) generic fronts.
8 *
9 * Every routine takes TWO template layers: the element `T` and the container template
10 * `Array`, with a `requires` concept enforcing the container. Both operands are spelled
11 * `Array<T>`, so a host⊗device or f64⊗f32 mix cannot deduce a single `Array`/`T` and is a
12 * compile error — the location/element firewall is FREE, via deduction, with no runtime
13 * check and no `SameLocation` clause on the common binary ops.
14 *
15 * Each op is a pair of same-named overloads:
16 * - a 2-arg **allocating front** `Array<T> op(const Array<T>&, const Array<T>&)` that
17 * allocates the result with `Array<T>::uninitialized(...)` and calls the out-param form;
18 * - a 3-arg **out-parameter kernel** `void op(Array<T>& out, …)`, split by concept: the
19 * HOST overload (declared here, defined in routines.cpp) runs the raw-pointer SIMD kernel;
20 * a device extension supplies a `requires DeviceArray<Array<T>>` overload in ITS namespace,
21 * found by ADL. Mutually exclusive concepts → no ambiguity, and cheatah never names the
22 * extension. (No CPO objects, no tag_invoke — plain concept-constrained overloads.)
23 */
24#include <stdexcept>
25#include <vector>
27#include "concepts.hpp"
29namespace cheatah::linalg {
31/// @cond INTERNAL — the allocation-free out-parameter kernel (HOST overload; a device
32/// extension adds its own `requires DeviceArray<Array<T>>` overload). Declared here so the
33/// allocating front below can call it; defined + explicitly instantiated in routines.cpp.
34/**
35 * Matmul into the CALLER'S buffer @p out (out FIRST) — no result allocation (a hot loop hands
36 * the same scratch every call). ONE two-layer overload over `Array<T>` unifying the former real
37 * and complex, host `NDArray`/`CNDArray` out-param functions.
38 * @tparam T the element type; @tparam Array the (host) container template.
39 * @param out contiguous [a.rows, b.cols] destination, overwritten; must NOT alias @p a or @p b.
40 * @param a,b the operands.
41 * @complexity O(n³) (× B for a batch).
42 * @alloc none for contiguous operands (product written straight into @p out); a
43 * non-contiguous operand is packed once into scratch.
44 * @test LinalgRoutines.MatmulIntoReusesBuffer
45 * @test LinalgRoutines.ComplexMatmulIntoReusesBuffer
46 */
47template <ndarray::Field T, template <typename> class Array>
48 requires HostArray<Array<T>>
49void matmul(Array<T>& out, const Array<T>& a, const Array<T>& b);
50/// @endcond
52/**
53 * Matrix multiply — the allocating front. Both operands are `Array<T>` (so host⊗device / element
54 * mixes fail to deduce and are compile errors); requires both to be 2-D with matching inner
55 * dimensions — or both 3-D for the BATCHED product `[B,M,K] @ [B,K,N] → [B,M,N]` (equal batch
56 * counts, strict: no broadcast batching). Allocates the result via `Array<T>::uninitialized` (no
57 * throwaway zero-fill) and fills it through the out-parameter kernel — the host SIMD path, or a
58 * device shader when `Array` is a device container (selected by concept at compile time).
59 * @tparam T the element type (`double` / `std::complex<double>`), @tparam Array the container template.
60 * @param a m×k matrix, or a B×m×k batch of matrices.
61 * @param b k×p matrix, or a B×k×p batch.
62 * @return m×p product (or the B×m×p batch), an `Array<T>` of the same container and element.
63 * @complexity O(n³) (× B for a batch).
64 * @alloc allocates only the result; operands read in place (a strided host view packs once).
65 * @concurrency deliberately single-threaded (the fastest-per-core contract); parallelize
66 * across independent products in the caller.
67 * @test LinalgRoutines.ProductsAndTrace
68 * @test LinalgRoutines.BatchedMatmul
69 * @crtest LinalgCompileRun.Matmul
70 * @systest StdlibE2E.Linalg
71 */
72template <ndarray::Field T, template <typename> class Array>
73 requires NumericArray<Array<T>>
74[[nodiscard]] Array<T> matmul(const Array<T>& a, const Array<T>& b) {
75 if (a.ndim() == 3 || b.ndim() == 3) {
76 if (a.ndim() != 3 || b.ndim() != 3)
77 throw std::runtime_error("linalg: batched matmul expects two 3-D operands");
78 if (a.shape()[0] != b.shape()[0])
79 throw std::runtime_error("linalg: batched matmul batch-count mismatch");
80 if (a.shape()[2] != b.shape()[1])
81 throw std::runtime_error("linalg: matmul inner dimension mismatch");
82 Array<T> out = Array<T>::uninitialized({a.shape()[0], a.shape()[1], b.shape()[2]});
83 matmul(out, a, b);
84 return out;
85 }
86 if (a.ndim() != 2 || b.ndim() != 2)
87 throw std::runtime_error("linalg: matmul expects 2-D matrices");
88 if (a.shape()[1] != b.shape()[0])
89 throw std::runtime_error("linalg: matmul inner dimension mismatch");
90 Array<T> out = Array<T>::uninitialized({a.shape()[0], b.shape()[1]});
91 matmul(out, a, b);
92 return out;
95/**
96 * The flattened length of a vector-shaped operand — 1-D, or 2-D with a size-1 row/column
97 * (throws otherwise). Reads only host-resident shape metadata, so it is valid for ANY located
98 * container, device arrays included; the shared validation step of every vector front below.
99 * @tparam A the (located) container type.
100 * @param a the operand whose vector length is wanted.
101 * @return the element count of the flattened vector.
102 * @complexity O(1).
103 * @alloc none.
104 * @test LinalgRoutines.ProductsAndTrace
105 */
106template <NumericArray A>
107[[nodiscard]] inline std::size_t vector_len(const A& a) {
108 if (a.ndim() == 1) return a.shape()[0];
109 if (a.ndim() == 2 && (a.shape()[0] == 1 || a.shape()[1] == 1)) return a.size();
110 throw std::runtime_error("linalg: expected a 1-D vector");
113// ---- reductions (dot / vdot / inner / trace): the scalar-out kernel pattern ----
114// Same two-layer seam as matmul, with a SCALAR out-parameter: the front validates and calls the
115// unqualified `op(out, …)`, which resolves to the HOST kernel below (routines.cpp) or a device
116// extension's `requires DeviceArray<Array<T>>` overload via ADL.
118/// @cond INTERNAL — the scalar-out reduction kernels (HOST overloads; a device extension adds its
119/// own `requires DeviceArray<Array<T>>` overloads). Declared here so the allocating fronts below
120/// can call them; defined + explicitly instantiated in routines.cpp.
121/**
122 * Bilinear dot product Σ aᵢbᵢ into the caller's scalar @p out (out FIRST) — the reduction analogue
123 * of the out-param matmul kernel, shared by real and complex elements.
124 * @tparam T the element type; @tparam Array the (host) container template.
125 * @param out receives the scalar sum. @param a,b same-length vectors (validated by the front).
126 * @test LinalgRoutines.ProductsAndTrace
127 */
128template <ndarray::Field T, template <typename> class Array>
129 requires HostArray<Array<T>>
130void dot(T& out, const Array<T>& a, const Array<T>& b);
131/**
132 * Hermitian inner product Σ conj(aᵢ)·bᵢ into @p out (bilinear for a real element — the
133 * conjugation is an `if constexpr` branch in the host kernel).
134 * @tparam T the element type; @tparam Array the (host) container template.
135 * @param out receives the scalar sum. @param a,b same-length vectors (validated by the front).
136 * @test LinalgRoutines.VdotInnerOuterKron
137 */
138template <ndarray::Field T, template <typename> class Array>
139 requires HostArray<Array<T>>
140void vdot(T& out, const Array<T>& a, const Array<T>& b);
141/**
142 * Bilinear inner product Σ aᵢbᵢ into @p out (numpy's `inner`; identical to @ref dot for
143 * flattened vectors).
144 * @tparam T the element type; @tparam Array the (host) container template.
145 * @param out receives the scalar sum. @param a,b same-length vectors (validated by the front).
146 * @test LinalgRoutines.VdotInnerOuterKron
147 */
148template <ndarray::Field T, template <typename> class Array>
149 requires HostArray<Array<T>>
150void inner(T& out, const Array<T>& a, const Array<T>& b);
151/**
152 * Trace (diagonal sum) into @p out — strided diagonal read, no copy.
153 * @tparam T the element type; @tparam Array the (host) container template.
154 * @param out receives the diagonal sum. @param a a 2-D matrix (validated by the front).
155 * @test LinalgRoutines.ProductsAndTrace
156 */
157template <ndarray::Field T, template <typename> class Array>
158 requires HostArray<Array<T>>
159void trace(T& out, const Array<T>& a);
160/// @endcond
162/**
163 * Dot product: 1-D inner product (vectors flattened) — the bilinear Σ aᵢbᵢ. ONE two-layer
164 * template over the element `T` and container `Array` serving real, complex, host and — via a
165 * device extension — device operands. Flattens each operand to a vector (1-D, or 2-D with a
166 * size-1 row/column) and throws if either is not vector-shaped or the lengths differ. Both
167 * operands are `Array<T>` (the deduction firewall).
168 * @tparam T the element type; @tparam Array the container template.
169 * @param a,b same-length vectors.
170 * @return Σ aᵢbᵢ as the scalar `T`.
171 * @complexity O(n).
172 * @alloc none for contiguous operands (read in place); a non-contiguous view packs once O(n).
173 * @test LinalgRoutines.ProductsAndTrace
174 * @test LinalgRoutines.ComplexProducts
175 * @crtest LinalgCompileRun.Dot
176 * @crtest LinalgCompileRun.ComplexDot
177 * @systest StdlibE2E.Linalg
178 * @systest StdlibE2E.LinalgComplex
179 */
180template <ndarray::Field T, template <typename> class Array>
181 requires NumericArray<Array<T>>
182[[nodiscard]] T dot(const Array<T>& a, const Array<T>& b) {
183 if (vector_len(a) != vector_len(b))
184 throw std::runtime_error("linalg: dot dimension mismatch");
185 T out;
186 dot(out, a, b);
187 return out;
190/**
191 * Vector dot product. For a REAL element this is the bilinear Σ aᵢbᵢ (identical to @ref dot and
192 * @ref inner); for a **complex** element it is the conjugate-linear Hermitian inner product
193 * ⟨a, b⟩ = Σ conj(aᵢ)·bᵢ (numpy's `vdot`, conjugating the first argument) — one two-layer template,
194 * the conjugation chosen at compile time by `if constexpr`. `vdot(a, a)` is the real ‖a‖².
195 * @tparam T the element type; @tparam Array the container template.
196 * @param a,b same-length vectors.
197 * @return Σ aᵢbᵢ (real) or Σ conj(aᵢ)·bᵢ (complex), as the scalar `T`.
198 * @complexity O(n).
199 * @alloc none for contiguous operands; a non-contiguous view packs once O(n).
200 * @test LinalgRoutines.VdotInnerOuterKron
201 * @test LinalgRoutines.ComplexProducts
202 * @crtest LinalgCompileRun.Vdot
203 * @crtest LinalgCompileRun.ComplexVdot
204 * @systest StdlibE2E.Linalg
205 * @systest StdlibE2E.LinalgComplex
206 */
207template <ndarray::Field T, template <typename> class Array>
208 requires NumericArray<Array<T>>
209[[nodiscard]] T vdot(const Array<T>& a, const Array<T>& b) {
210 if (vector_len(a) != vector_len(b))
211 throw std::runtime_error("linalg: dot dimension mismatch");
212 T out;
213 vdot(out, a, b);
214 return out;
217/**
218 * Inner product of two vectors — the bilinear Σ aᵢbᵢ (numpy's `inner`; same as @ref dot for
219 * flattened vectors). One two-layer template over the element and container.
220 * @tparam T the element type; @tparam Array the container template.
221 * @param a,b same-length vectors.
222 * @return Σ aᵢbᵢ as the scalar `T`.
223 * @complexity O(n).
224 * @alloc none for contiguous operands; a non-contiguous view packs once O(n).
225 * @test LinalgRoutines.VdotInnerOuterKron
226 * @crtest LinalgCompileRun.Inner
227 * @systest StdlibE2E.Linalg
228 */
229template <ndarray::Field T, template <typename> class Array>
230 requires NumericArray<Array<T>>
231[[nodiscard]] T inner(const Array<T>& a, const Array<T>& b) {
232 if (vector_len(a) != vector_len(b))
233 throw std::runtime_error("linalg: dot dimension mismatch");
234 T out;
235 inner(out, a, b);
236 return out;
239/**
240 * Trace: the sum of the matrix diagonal, as the scalar `T`. Requires a 2-D matrix (throws
241 * otherwise); rectangular matrices sum min(r, c) diagonal entries.
242 * @tparam T the element type; @tparam Array the container template.
243 * @param a a 2-D matrix.
244 * @return Σ aᵢᵢ as the scalar `T`.
245 * @complexity O(min(r, c)).
246 * @alloc none (strided diagonal read straight from the buffer).
247 * @test LinalgRoutines.ProductsAndTrace
248 * @crtest LinalgCompileRun.Trace
249 * @systest StdlibE2E.Linalg
250 */
251template <ndarray::Field T, template <typename> class Array>
252 requires NumericArray<Array<T>>
253[[nodiscard]] T trace(const Array<T>& a) {
254 if (a.ndim() != 2) throw std::runtime_error("linalg: expected a 2-D matrix");
255 T out;
256 trace(out, a);
257 return out;
260// ---- products with array results (outer / conj_transpose / kron): the matmul pattern ----
262/// @cond INTERNAL — the allocation-free out-parameter kernels (HOST overloads; a device extension
263/// adds its own `requires DeviceArray<Array<T>>` overloads). Declared here so the allocating
264/// fronts below can call them; defined + explicitly instantiated in routines.cpp.
265/**
266 * Outer product into the caller's buffer @p out (out FIRST) — the buffer-reuse overload of
267 * @ref outer, writing the rank-1 result straight into @p out with no allocation.
268 * @param out destination; a contiguous n×m matrix, overwritten. Must NOT alias @p a or @p b.
269 * @param a length-n vector.
270 * @param b length-m vector.
271 * @complexity O(n·m).
272 * @alloc none for contiguous operands (result written straight into @p out); a
273 * non-contiguous operand is packed once into scratch.
274 * @test LinalgRoutines.OuterIntoReusesBuffer
275 */
276template <ndarray::Field T, template <typename> class Array>
277 requires HostArray<Array<T>>
278void outer(Array<T>& out, const Array<T>& a, const Array<T>& b);
279/**
280 * Conjugate transpose into the caller's buffer @p out (out FIRST) — the buffer-reuse overload of
281 * @ref conj_transpose, writing the c×r adjoint straight into @p out with no allocation.
282 * @param out destination; a contiguous c×r matrix (for an r×c input), overwritten. Must NOT
283 * alias @p a (it reads A while writing the transpose — not an in-place op).
284 * @param a a 2-D matrix.
285 * @complexity O(r·c).
286 * @alloc none for a contiguous operand (adjoint written straight into @p out); a
287 * non-contiguous operand is packed once into scratch.
288 * @test LinalgRoutines.ConjTransposeIntoReusesBuffer
289 */
290template <ndarray::Field T, template <typename> class Array>
291 requires HostArray<Array<T>>
292void conj_transpose(Array<T>& out, const Array<T>& a);
293/**
294 * Kronecker product into the caller's buffer @p out (out FIRST) — the buffer-reuse overload of
295 * @ref kron, writing the block product straight into @p out with no allocation.
296 * @param out destination; a contiguous (m·p)×(k·q) matrix, overwritten. Must NOT alias @p a or @p b.
297 * @param a m×k matrix.
298 * @param b p×q matrix.
299 * @complexity O(n⁴) in the output area.
300 * @alloc none for contiguous operands (block product written straight into @p out); a
301 * non-contiguous operand is packed once into scratch.
302 * @test LinalgRoutines.KronIntoReusesBuffer
303 */
304template <ndarray::Field T, template <typename> class Array>
305 requires HostArray<Array<T>>
306void kron(Array<T>& out, const Array<T>& a, const Array<T>& b);
307/// @endcond
309/**
310 * Outer product of two vectors.
311 *
312 * Flattens both operands to vectors and forms the full rank-1 matrix; any pair
313 * of vector lengths is accepted (no matching constraint). Allocates the result via
314 * `Array<T>::uninitialized` and fills it through the out-parameter kernel — the host SIMD
315 * path, or a device shader when `Array` is a device container (selected by concept).
316 * @param a length-n vector.
317 * @param b length-m vector.
318 * @return n×m matrix aᵢbⱼ.
319 * @complexity O(n·m).
320 * @alloc allocates only the n×m result; operands read in place when contiguous.
321 * @test LinalgRoutines.VdotInnerOuterKron
322 * @crtest LinalgCompileRun.Outer
323 * @systest StdlibE2E.Linalg
324 */
325template <ndarray::Field T, template <typename> class Array>
326 requires NumericArray<Array<T>>
327[[nodiscard]] Array<T> outer(const Array<T>& a, const Array<T>& b) {
328 Array<T> out = Array<T>::uninitialized({vector_len(a), vector_len(b)});
329 outer(out, a, b);
330 return out;
333/**
334 * Conjugate transpose (Hermitian adjoint) Aᴴ: transpose, then conjugate every entry (a plain
335 * transpose for a real element — the conjugation is compiled out). A matrix is Hermitian iff
336 * `conj_transpose(A) == A`.
337 * @param a a 2-D matrix.
338 * @return the c×r adjoint of an r×c input; throws on non-2-D input.
339 * @complexity O(r·c).
340 * @alloc allocates only the c×r result; a non-contiguous operand is packed once into scratch.
341 * @test LinalgRoutines.ComplexProducts
342 * @crtest LinalgCompileRun.ConjTranspose
343 * @systest StdlibE2E.LinalgComplex
344 */
345template <ndarray::Field T, template <typename> class Array>
346 requires NumericArray<Array<T>>
347[[nodiscard]] Array<T> conj_transpose(const Array<T>& a) {
348 if (a.ndim() != 2) throw std::runtime_error("linalg: expected a 2-D matrix");
349 Array<T> out = Array<T>::uninitialized({a.shape()[1], a.shape()[0]});
350 conj_transpose(out, a);
351 return out;
354/**
355 * Kronecker product.
356 *
357 * Requires both operands to be 2-D (throws otherwise) and replaces each entry of
358 * @p a with that scalar times the whole of @p b, giving the (m·p)×(k·q) block
359 * matrix; no dimension matching is needed.
360 * @param a m×k matrix.
361 * @param b p×q matrix.
362 * @return (m·p)×(k·q) block product.
363 * @complexity O(n⁴) in the output area.
364 * @alloc allocates only the (m·p)×(k·q) result; a non-contiguous operand is packed once
365 * into scratch.
366 * @test LinalgRoutines.VdotInnerOuterKron
367 * @crtest LinalgCompileRun.Kron
368 * @systest StdlibE2E.Linalg
369 */
370template <ndarray::Field T, template <typename> class Array>
371 requires NumericArray<Array<T>>
372[[nodiscard]] Array<T> kron(const Array<T>& a, const Array<T>& b) {
373 if (a.ndim() != 2 || b.ndim() != 2)
374 throw std::runtime_error("linalg: kron expects 2-D matrices");
375 // Each output dimension is a PRODUCT of two input dims, so it must be overflow-checked
376 // BEFORE it collapses into the shape vector: `product({m*p, k*q})` would only catch a wrap
377 // of the final area, not a wrap of m*p (or k*q) alone, which under-allocates and lets the
378 // kernel write out of bounds. `product({x, y})` is the shared checked multiply — it throws
379 // on wrap instead. (See ndarray::detail::product; same class as the ndarray shape-overflow fix.)
380 const std::size_t orows = ndarray::detail::product({a.shape()[0], b.shape()[0]});
381 const std::size_t ocols = ndarray::detail::product({a.shape()[1], b.shape()[1]});
382 Array<T> out = Array<T>::uninitialized({orows, ocols});
383 kron(out, a, b);
384 return out;
387} // namespace cheatah::linalg