Tensor Is the Might

Tensor is the might Every good abstraction solves a problem, and this post will cover everything I know so far about a brilliant math abstraction - tensors. Neural networks, from a simple 2-layer MLP to GPT-5, all boil down to the same thing: floating-point numbers flowing through a graph of operations. This post builds a complete, accelerated tensor library from scratch in C. It is heavily inspired by Bellard’s libnc, which unfortunately has not been open-sourced yet. A tensor is nothing but a flat array of numbers, plus some metadata telling you how to interpret those numbers as a multi-dimensional object. We all learned that 2D arrays can be better represented as 1D array plus a number of rows/columns - this is essentially what a tensor is. But going beyond two dimensions - we might need some other metadata, such as a generalised shape: float data[32 * 3 * 28 * 28]; // 32 images, 3 channels, 28x28 pixels int shape[4] = {32, 3, 28, 28}; // shape of the tensor int ndim = 4; // number of dimensions Having a shape we can figure out that an element at position [n,c,h,w] in a 4D tensor lives at offset data[n*(3*28*28)+c*(28*28)+h*28+w], if we keep our tensor in a row-major C-order format (often the default in most tensor libraries today). Now, calculating each time an index of an element like this is inefficient, so we can precompute the strides once the shape is known. Strides tell how many elements to skip if we want to advance for one element in the given dimension. We could also group all shape-related fields together: struct ut_shape { int ndim; // number of dimensions int nelem; // number of elements int shape[4]; // shape of the tensor int strides[4]; // strides for each dimension }; struct ut_tensor { struct ut_shape shape; // shape of the tensor float *data; // pointer to the data ... // more fields to come later }; We can add some helpers to create a shape and get flat index from a multi-dimensional index: ut_shape ut_shape_new(int ndims, int* dims) { ut_shape s = {.ndims = ndims, .nelems = 1}; for (int i = 0; i < ndim; i++) { s.shape[i] = dims[i]; s.nelem *= dims[i]; } return s; } int ut_index(ut_shape s, const int* idx) { int flat = 0, stride = 1; for (int i = s.ndim - 1; i >= 0; i--) { flat += idx[i] * stride; stride *= s.shape[i]; } return flat; } ut_shape s = ut_shape_new(3, (int[]){2, 3, 4}); assert(s.nelem == 24); assert(s.ndim == 3); // element [1][2][3] should be at offset 1*12 + 2*4 + 3 = 23 assert(ut_index(s, (int[]){1, 2, 3}) == 23); Tensors are usually dynamically allocated, so we should provide a way to create and destroy them, nothing but a wrapper on top of malloc/free. Sometimes we want to create a tensor that shares the same data with another tensor, for example to get a “view” of a part of a tensor, to transpose a tensor without copying data, or to modify its shape (flatten). In this case we need to track ownership of the data. It’d be also useful to “retain” tensors, so that they could outlive their original scope, like during iterative training to avoid the same allocation over and over. So, we add a reference counter field to the tensor struct: struct ut_tensor { struct ut_shape shape; // shape of the tensor float *data; // pointer to the data int refcount; // reference count for shared ownership, free when reaches 0 struct ut_tensor *owner; // pointer to the owner tensor if this tensor is a view }; We can optimise memory management further, adding arena allocator or memory pool to avoid frequent malloc/free calls, but what we already have is a good start. However, our tensors are hardly useful without operations on them. Elementwise The most basic operations on tensors are elementwise – a single loop over all elements, applying a function to each element (unary) or a pair of elements (binary). We can implement them just like that: static void ew_neg(float* out, const float* a, int n) { for (int i = 0; i < n; i++) out[i] = -a[i]; } // ...more unary ops... static void ew_relu(float* out, const float* a, int n) { for (int i = 0; i < n; i++) out[i] = fmaxf(0.f, a[i]); } static void ew_add(float* out, const float* a, const float* b, int n) { for (int i = 0; i < n; i++) out[i] = a[i] + b[i]; } // ...more binary ops... static ut_tensor* ew_unary(ut_tensor* a, void (*fn)(float*, const float*, int)) { ut_tensor* out = ut_alloc(a->shape.ndim, a->shape.shape); fn(out->data, a->data, a->shape.nelem); return out; } static ut_tensor* ew_binary(ut_tensor* a, ut_tensor* b, void (*fn)(float*, const float*, const float*, int)) { ut_tensor* out = ut_alloc(a->shape.ndim, a->shape.shape); fn(out->data, a->data, b->data, a->shape.nelem); return out; } ut_tensor* ut_neg(ut_tensor* a) { return ew_unary(a, ew_neg); } ut_tensor* ut_exp(ut_tensor* a) { return ew_unary(a, ew_exp); } ut_tensor* ut_sigmoid(ut_tensor* a) { return ew_unary(a, ew_sigmoid); } ut_tensor* ut_tanh(ut_tensor* a) { return ew_unary(a, ew_tanh); } ut_tensor* ut_relu(ut_tensor* a) { return ew_unary(a, ew_relu); } ut_tensor* ut_add(ut_tensor* a, ut_tensor* b) { return ew_binary(a, b, ew_add); } ut_tensor* ut_sub(ut_tensor* a, ut_tensor* b) { return ew_binary(a, b, ew_sub); } ut_tensor* ut_mul(ut_tensor* a, ut_tensor* b) { return ew_binary(a, b, ew_mul); } ut_tensor* ut_div(ut_tensor* a, ut_tensor* b) { return ew_binary(a, b, ew_div); } ut_tensor* ut_scale(ut_tensor* a, float s) { ut_tensor* out = ut_alloc(a->shape.ndim, a->shape.shape); for (int i = 0; i < a->shape.nelem; i++) out->data[i] = a->data[i] * s; return out; } It’d be nice to assert() that the shapes of two tensors are the same before doing a binary operation. Alternatively, at this point we might decide that we want to support “broadcasting” – extending the smaller tensor to match the shape of the larger tensor. For example a tensor of shape [2, 3, 4] and a tensor of shape [3, 4] can be added together by “stretching” the second tensor along the first dimension. While being a common feature in many tensor libraries, I decided to leave it out for now. Most models I’m aiming for have their underlying tensors perfectly aligned. If not - we can either duplicate data explicitly aligning tensor shapes, or adding broadcasting index calculation later. It might be tempting at this point to implement more operations, matrix multiplication, convolutions, etc. But this is the moment where we have to ask – do we want all operations to run on CPU? Golden Processing Unit (GPU) Looking at modern GPU prices, it is clear that they probably have a significant value when it comes to crunching arrays of numbers. Thus, instead of limiting our library to a CPU, we might consider offloading some operations to a GPU. This is where things get interesting, because we now need to manage memory across both CPU and GPU, transfer data efficiently, learn how to write GPU kernels, and much more. To make things worse, there is no single “GPU accelerator” – there are many vendors, each with their own APIs and quirks: CUDA, OpenCL, WebGPU, Metal, Vulkan, etc. I tried many times to get the most out of CPU-only tensors, but even with BLAS, LAPACK, OpenMP – I could not reach the same magnitude of performance as GPU-accelerated frameworks. To keep things manageable, I decided to start with Metal. On the one hand it limits us to Apple devices, but on the other hand it’s a fairly modern API and simplifies things due to “unified memory” of Apple Silicon. Metal uses its own language to define GPU kernels (MSL), which is similar to modern C++, and it compiles kernels at runtime, so we can define them as strings in our C code: static const char *shaderSource = "#include <metal_stdlib>\n" "using namespace metal;\n" "\n" "kernel void relu(device const float *in [[buffer(0)]],\n" " device float *out [[buffer(1)]],\n" " uint id [[thread_position_in_grid]]) {\n" " float val = in[id];\n" " out[id] = fmax(val, 0.0f);\n" "}\n"; To make this kernel run on GPU we need to create a device, build a command queue (GPUs are asynchronous), compile the kernel, create buffers for input and output, and