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.ElementwiseThe 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 finally dispatch the kernel to run on GPU.

Here’s how a simple GPU-accelerated ReLU operation might look like:// prepare device, queue, and compile kernel id<MTLDevice> device = MTLCreateSystemDefaultDevice(); id<MTLCommandQueue> queue = [device newCommandQueue]; NSError *err = nil; id<MTLLibrary> library = [device newLibraryWithSource:[NSString stringWithUTF8String:shaderSource] options:nil error:&err]; id<MTLFunction> reluFn = [library newFunctionWithName:@"relu"]; id<MTLComputePipelineState> pipeline = [device newComputePipelineStateWithFunction:reluFn error:&err];

// prepare data, dispatch kernel float input[8] = {-2.0, -1.0, 0.0, 1.0, 2.0, -0.5, 3.0, -3.0}; float output[8] = {0}; id<MTLBuffer> bufIn = [device newBufferWithBytes:input length:sizeof(input) options:MTLResourceStorageModeShared]; id<MTLBuffer> bufOut = [device newBufferWithLength:sizeof(output) options:MTLResourceStorageModeShared]; id<MTLCommandBuffer> cmdBuf = [queue commandBuffer]; id<MTLComputeCommandEncoder> enc = [cmdBuf computeCommandEncoder]; [enc setComputePipelineState:pipeline]; [enc setBuffer:bufIn offset:0 atIndex:0]; [enc setBuffer:bufOut offset:0 atIndex:1]; MTLSize gridSize = MTLSizeMake(8, 1, 1); MTLSize tgSize = MTLSizeMake(pipeline.maxTotalThreadsPerThreadgroup, 1, 1); [enc dispatchThreads:gridSize threadsPerThreadgroup:tgSize]; [enc endEncoding]; [cmdBuf commit]; [cmdBuf waitUntilCompleted]; // copy results back memcpy(output, bufOut.contents, sizeof(output)); We can wrap the first part of this code into a “Metal context” struct, that is created once and is kept alive for the lifetime of the whole library. Then we can add helpers to allocate buffers, read and write them and dispatch kernels.

While the official Metal API is in Objective-C or Swift, we can use low-level ObjC runtime functions to call it from C, to keep the library pure (objC_msgSend all over the code).Now every operation on tensors has to be implemented twice – as a naïve CPU implementation, and as a GPU kernel. Depending on the origin of the tensor we would call one of the implementations. Note that this also means that tensors carry the “device” field, and we must specify the device when the tensor is allocated – CPU tensors would only have a normal data buffer, but Metal tensors would also have a pointer to a Metal buffer object.It’s also becoming important that we need to track when GPU and CPU buffers diverge. An easy way is to keep “dirty” flags for CPU/GPU buffers and provide “sync” functions to copy data from one buffer to another.Even though some operations can or should be done on a GPU, in some cases to simplify things we still create it on a CPU first, then do the processing and send/copy data back to GPU.All in all, we’d need to implement the following:static const char *_mtl_src = "..."; // MSL source code for all kernels void *ut_mtl_init(); // create Metal context as a singleton void ut_mtl_dispatch(void *ctx, char *kernel, void **bufs, int nbufs, void *bytes, int blen, int n); // dispatch kernel with buffers void *ut_mtl_buf_alloc(void *ctx, void *data, int len); // allocate Metal buffer void ut_mtl_buf_free(void *ctx, void *buf); // free Metal buffer void ut_mtl_buf_read(void *ctx, void *buf, void *data, int len); // read Metal buffer to CPU void ut_mtl_buf_write(void *ctx, void *buf, void *data, int len); // write CPU buffer to Metal

void ut_sync_cpu(ut_tensor *t); // sync CPU buffer from GPU void ut_sync_gpu(ut_tensor *t); // sync GPU buffer from CPU void ut_to_device(ut_tensor *t, ut_dev device); // move tensor to specified device (CPU/GPU) Done with element-wise ops and their GPU kernel equivalents we can move on to the core of all neural networks - matrix multiplication.