Module ml::tensor
Tensor: CPU tensors with reverse-mode Autograd::get.
PyTorch-style ergonomics (build a graph implicitly by calling ops; Tensor::backward() walks it) over a plain flat f32 buffer — no unified- memory story here (that's MPS's job; see gpu_mps.h), just row-major single-precision arrays with a recorded computation graph. Scope: 1D and 2D tensors (vectors and matrices — everything a small MLP/attention block needs), add/sub/mul/relu/sum/matmul, and reverse-mode autodiff over that op set. Not a general n-dimensional array library — no broadcasting beyond 2D, no int/complex dtypes, no in-place ops.
float32, not float64: PyTorch defaults to float32, and this codebase's training-loop shapes measured ~30% faster on Accelerate's cblas_sgemm than cblas_dgemm for identical work, on top of halving memory bandwidth for every elementwise pass -- see tensor_blas.h. Also the same precision Metal itself supports (Metal has no double at all), so this removes the double<->float conversion every GPU-backed op used to pay at its CPU/GPU boundary (see gpu_mps.h's history) instead of just moving where the precision loss happens.
Internal helpers (__tensor_alloc, __tensor_link1/2, every _X_backward function, etc.) are defined right here in the header as 'static' functions, not in tensor.sc — this is what lets tensor.sc, tensor_nn.sc, and the tensor<gpu-backend>.sc family compile as genuinely separate translation units (each #includes this header, gets its own private copy of these helpers, and is linked against the others' public symbols — tensor_add, tensor_matmul, etc. — at build time) instead of every file needing to textually #include tensor.sc's entire body just to reach a handful of internal-linkage helpers it happens to reuse.
Every op here is its own small function with one loop -- not just style. A hand-rolled training loop with everything (BLAS calls, elementwise passes) inlined into one large function was measured to make this compiler's -O2 optimizer hang (40+ minutes, not just slow) on exactly this kind of training-loop shape; splitting into small single-purpose functions avoided it entirely and let -O2's auto-vectorization actually run. Keeping every op as its own function isn't just for consistency -- it's load-bearing.
Structs
| TensorArena | Every op allocates its output. In a training loop that is a malloc (and, for fresh memory, page faults) per op per step: on a small MLP that cost more than the arithmetic. A TensorArena is a chunked bump allocator for a step's temporaries: between a.begin() and tensor_arena_end() every tensor created on this thread (data, shape, gradient) comes from it, and a.reset() drops them all at once while keeping the memory -- already faulted in -- for the next step. |
| DenseTensor | A DenseTensor<E> stores its values as E (element.h) and computes in float32; its gradient is always float32. FloatTensor (E = float) is the float32 tensor every op has always used; HalfTensor (F16) and BF16Tensor (BF16) keep half the memory per value and run every op and its backward pass through float32. Generic code is written against trait Tensor and works with all three. |
| GradGuard | Global, not per-TU: every tensor_<op> across every separately-compiled file linked into the same program shares one flag (that's why these are ordinary externally-linked functions, defined once in tensor.sc, not 'static' — unlike the internal helpers below, this state must be visible and consistent everywhere, not private per including file). |
| Autograd |
Traits
| Tensor | What generic ML code asks of a tensor: FloatTensor, HalfTensor and BF16Tensor all satisfy it ('generic<T: Tensor>'). |
| DeviceTensor | What a tensor that can live on a device adds (device.h): DenseTensor and SparseTensor (sparse.h) conform. device() where the values live (DEVICE_CPU, DEVICE_ROCM, DEVICE_CUDA, DEVICE_MPS) place(d) moves the tensor itself to d (parameters, leaves); 1 on success sync() makes the host copy current; 1 on success |
| RocmTensor | One trait per GPU backend, each inheriting DeviceTensor: code written for one backend's buffers (a hand-written kernel, a vendor library call) asks for its tensor. device_buffer() is the values' buffer on that backend -- a hipDeviceptr, a CUdeviceptr, an MTLBuffer's GPU address, a Vulkan buffer device address -- and 0 while the tensor is elsewhere: check device() first, or place() it. |
| CudaTensor | |
| MpsTensor | |
| VulkanTensor | |
| WebGpuTensor | WebGPU's device_buffer() is a byte offset into the backend's device heap (gpu_webgpu.h's WebGpuBuf.handle), the element type in its top bits. |
Type aliases
| FloatTensor | |
| HalfTensor | |
| BF16Tensor | |
| TensorBackwardFn | A node's backward function. Its argument is the node viewed as a float32 tensor: the graph fields (and the float32 gradient) have one layout for every element type, so one function serves a graph that mixes them, and each function reads only those fields of its node (and its parents). |
Functions
| malloc | The static helpers below call libc malloc/free directly (not std::alloc/dealloc, so no dependency on mem.sc's allocator machinery for this hot path) -- declared here rather than pulled in via std/mem.sc, matching every other .sc file in std/ that touches libc memory functions directly (see mem.sc, heap.sc, result.sc, etc., each of which declares its own copy of these same externs rather than including one another's .sc file for them). |
| tensor_arena_end | back to the heap |
| tensor_add | requires grad) ──────────────────────────────────────────────────────────── Generic over the element type (defined at the end of this header); a HalfTensor / BF16Tensor op computes in float32 and rounds its result. |
| tensor_sub | |
| tensor_mul | elementwise |
| tensor_matmul | 2D x 2D |
Macros
| TS_HOST | Coherence bits ('coh'): which copies are current. Both the values' bits may be set (the copies agree); a gradient that doesn't exist yet counts as current zeros on both sides. |
| TS_DEV | the values on the device |
| TS_GHOST | the gradient on the host |
| TS_GDEV | the gradient on the device |
| TS_DEVOP | made by a device op: its backward runs on the device |
| TS_GZERO | the device gradient is all zeros, not yet written (its buffer may hold anything) |
| TDEV_POOL | Spatial ops over [batch, CHW] planes (convnet.h's layout). |
| TDEV_GPOOL | global pooling: mode 0 max, 1 avg |
| TDEV_UPSAMPLE | nearest by the integer factor kH |
| TDEV_RESIZE | to oH x oW, mode RESIZE_*, align |
| EW_EXP | The EW codes past tensor.h's own, shared with device_kernels.cu. |
| EW_LOG | |
| EW_ABS | |
| EW_LEAKY | |
| EW_CLAMP | |
| EW_SIGMOID | |
| EW_TANH | |
| EW_GELU | |
| EW_SILU | |
| EW_ADDK | |
| EW_MIN | |
| EW_ELU | |
| EW_LERP | out += k * (a - out), in place |
| EW_MULS | out = a * b[0], b a device scalar |
| EW_ACC_PK | |
| EW_ACC_EXPB | |
| EW_ACC_LOGB | |
| EW_ACC_ABSB | |
| EW_ACC_LEAKYB | |
| EW_ACC_CLAMPB | |
| EW_ACC_GELUB | |
| EW_ACC_SILUB | |
| EW_ACC_MINA | |
| EW_ACC_MINB | |
| EW_ACC_ELUB | |
| EW_ACC_PSCALE | |
| TENSOR_EW_PAR | One loop per op over raw arrays, which the compiler vectorizes; arrays of TENSOR_EW_PAR or more elements are split across std/sync/parallel.h's thread pool. Used by the forward ops in tensor.sc / tensor_nn.sc and by the backward functions below. |
| EW_ADD | out = a + b |
| EW_SUB | out = a - b |
| EW_MUL | out = a * b |
| EW_SCALE | out = a * k |
| EW_RELU | out = max(a, 0) |
| EW_FILL | out = k |
| EW_COPY | out = a |
| EW_ACC | out += a |
| EW_DEC | out -= a |
| EW_ACC_MUL | out += a * b |
| EW_ACC_SCALE | out += a * k |
| EW_ACC_RELU | out += b > 0 ? a : 0 |
| EW_ACC_K | out += k |
| EW_ACC_SIGB | out += a * b * (1 - b) (sigmoid backward, b = output) |
| EW_ACC_TANHB | out += a * (1 - b * b) (tanh backward, b = output) |