张量即力量

# 张量即力量 来源：https://zserge.com/posts/tensor/ 每个好的抽象都能解决一个问题，这篇文章将涵盖我目前所知的关于一个杰出数学抽象——张量——的一切。神经网络，从简单的两层MLP (https://en.wikipedia.org/wiki/Multilayer_perceptron)到GPT-5，最终都归结为同一件事：浮点数通过操作图流动。这篇文章从头开始用C语言构建了一个完整且加速的张量库。它深受 Bellard 的 libnc (https://bellard.org/libnc/) 启发，遗憾的是该库尚未开源。张量不过是一个扁平的数组，加上一些元数据，告诉你如何将这些数字解释为多维对象。我们都知道，二维数组可以更好地表示为一位数组加上行数/列数——这本质上就是张量。张量但扩展到超过二维——我们可能需要其他元数据，比如通用的形状： `` float data[32 * 3 * 28 * 28]; // 32 张图像，3 个通道，28x28 像素 int shape[4] = {32, 3, 28, 28}; // 张量的形状 int ndim = 4; // 维度数量 `` 有了形状，我们可以得知，在四维张量中，位置`[n,c,h,w]`处的元素位于偏移量`data[n*(3*28*28)+c*(28*28)+h*28+w]`，如果我们以*行主序*C顺序格式存储张量（目前大多数张量库的默认方式）。现在，每次像这样计算元素索引效率不高，因此我们可以在已知形状后预计算*步幅*。步幅告诉我们要在给定维度上前进一个元素需要跳过多少个元素。我们还可以将所有与形状相关的字段分组： `` struct ut_shape { int ndim; // 维度数量 int nelem; // 元素数量 int shape[4]; // 张量的形状 int strides[4]; // 每个维度的步幅 }; struct ut_tensor { struct ut_shape shape; // 张量的形状 float *data; // 指向数据的指针 ... // 稍后添加更多字段 }; `` 我们可以添加一些辅助函数来创建形状，并从多维索引获取扁平索引： `` 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); // 元素 [1][2][3] 应位于偏移量 1*12 + 2*4 + 3 = 23 assert(ut_index(s, (int[]){1, 2, 3}) == 23); `` 张量通常是动态分配的，因此我们应该提供创建和销毁的方法，这不过是`malloc`/`free`的包装。有时我们希望创建一个与其他张量共享相同数据的张量，例如获取张量一部分的“视图”、在不复制数据的情况下转置张量，或修改其形状（展平）。在这种情况下，我们需要跟踪数据的所有权。同样有用的还有“保留”张量，这样它们可以比原始作用域活得更久，比如在迭代训练中避免重复分配。因此，我们在张量结构体中添加一个引用计数器字段： `` struct ut_tensor { struct ut_shape shape; // 张量的形状 float *data; // 指向数据的指针 int refcount; // 引用计数，用于共享所有权，达到 0 时释放 struct ut_tensor *owner; // 如果此张量是一个视图，指向所有者张量 }; `` 我们可以进一步优化内存管理，添加 arena 分配器或内存池以避免频繁的`malloc`/`free`调用，但目前已有的已经是一个好的开始。然而，没有操作的张量几乎没有用处。 ## 逐元素操作 最基本的张量操作是逐元素的——对所有元素进行单次循环，将一个函数应用于每个元素（一元）或一对元素（二元）。我们可以直接实现它们： `` static void ew_neg(float* out, const float* a, int n) { for (int i = 0; i < n; i++) out[i] = -a[i]; } // ...更多一元操作... 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]; } // ...更多二元操作... 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; } `` 在执行二元操作之前，最好`assert()`两个张量的形状相同。或者，此时我们可能决定支持“广播”——将较小的张量扩展以匹配较大张量的形状。例如，形状为`[2, 3, 4]`的张量与形状为`[3, 4]`的张量可以通过沿第一维度“拉伸”第二个张量来相加。虽然这是许多张量库的常见特性，我决定暂时不包含它。我目标中的大多数模型，其底层张量都是完美对齐的。如果不是——我们可以显式地复制数据来对齐张量形状，或者稍后添加广播索引计算。此时可能会很想实现更多操作，如矩阵乘法、卷积等。但这时我们必须问——我们是否希望所有操作都在CPU上运行？ ## 黄金处理单元 (GPU) 看看现代GPU的价格，很明显它们在处理数字数组方面具有显著价值。因此，我们不应将库限制在CPU上，而应考虑将某些操作卸载到GPU。事情变得有趣起来，因为我们现在需要在CPU和GPU之间管理内存，高效传输数据，学习如何编写GPU内核等等。更糟糕的是，并没有单一的“GPU加速器”——有许多供应商，各自拥有自己的API和怪癖：CUDA、OpenCL、WebGPU、Metal、Vulkan等。我多次尝试优化纯CPU张量，即使使用BLAS、LAPACK、OpenMP——我也无法达到与GPU加速框架相同的性能水平。为了保持可控，我决定从Metal开始。一方面，这限制了我们使用Apple设备，但另一方面，它是一个相当现代的API，并且由于Apple Silicon的“统一内存”而简化了事情。Metal使用自己的语言定义GPU内核（MSL），类似于现代C++，并且在运行时编译内核，因此我们可以将它们定义为C代码中的字符串： `` static const char *shaderSource = "#include \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"; `` 要使此内核在GPU上运行，我们需要创建设备、构建命令队列（GPU是异步的）、编译内核、为输入和输出创建缓冲区，最后调度内核在GPU上运行。以下是一个简单的GPU加速ReLU操作的样子： `` // 准备设备、队列和编译内核 id device = MTLCreateSystemDefaultDevice(); id queue = [device newCommandQueue]; NSError *err = nil; id library = [device newLibraryWithSource:[NSString stringWithUTF8String:shaderSource] options:nil error:&err]; id reluFn = [library newFunctionWithName:@"relu"]; id pipeline = [device newComputePipelineStateWithFunction:reluFn error:&err]; // 准备数据，调度内核 float input[8] = {-2.0, -1.0, 0.0, 1.0, 2.0, -0.5, 3.0, -3.0}; float output[8] = {0}; id bufIn = [device newBufferWithBytes:input length:sizeof(input) options:MTLResourceStorageModeShared]; id bufOut = [device newBufferWithLength:sizeof(output) options:MTLResourceStorageModeShared]; id cmdBuf = [queue commandBuffer]; id 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]; // 将结果复制回来 memcpy(output, bufOut.contents, sizeof(output)); `` 我们可以将这部分代码的第一部分包装成一个“Metal context”结构体，创建一次并在整个库的生命周期内保持存活。然后我们可以添加辅助函数来分配缓冲区、读取和写入它们，以及调度内核。虽然官方的Metal API是Objective-C或Swift，但我们可以使用底层的ObjC运行时函数 (https://developer.apple.com/documentation/objectivec/objective-c-runtime) 从C调用它，以保持库的纯C性（到处都是`objc_msgSend`）。现在，每个张量操作都必须实现两次——一次是朴素的CPU实现，一次是GPU内核。根据张量的来源，我们将调用其中一个实现。注意，这也意味着张量带有“device”字段，并且我们在分配张量时必须指定设备——CPU张量只有正常的数据缓冲区，但Metal张量还会有一个指向Metal缓冲对象的指针。跟踪GPU和CPU缓冲区何时不同也变得重要。一种简单的方法是保持CPU/GPU缓冲区的“脏”标志，并提供“同步”函数来在缓冲区之间复制数据。即使某些操作可以在GPU上完成或应该完成，有时为了简化，我们仍然先在CPU上创建它，然后进行处理，并将数据发送/复制回GPU。总之，我们需要实现以下内容： `` static const char *_mtl_src = "..."; // 所有内核的MSL源代码 void *ut_mtl_init(); // 创建单例Metal上下文 void ut_mtl_dispatch(void *ctx, char *kernel, void **bufs, int nbufs, void *bytes, int blen, int n); // 使用缓冲区调度内核 void *ut_mtl_buf_alloc(void *ctx, void *data, int len); // 分配Metal缓冲区 void ut_mtl_buf_free(void *ctx, void *buf); // 释放Metal缓冲区 void ut_mtl_buf_read(void *ctx, void *buf, void *data, int len); // 从Metal缓冲区读取到CPU void ut_mtl_buf_write(void *ctx, void *buf, void *data, int len); // 从CPU缓冲区写入Metal void ut_sync_cpu(ut_tensor *t); // 从GPU同步CPU缓冲区 void ut_sync_gpu(ut_tensor *t); // 从CPU同步GPU缓冲区 void ut_to_device(ut_tensor *t, ut_dev device); // 将张量移动到指定设备（CPU/GPU） `` 完成逐元素操作及其GPU内核等价实现后，我们可以转向所有神经网络的核心——矩阵乘法。 ## 矩阵乘法 一个临时编造的统计告诉我，在任何神经网络中，80%的FLOPs都是矩阵乘法。正确且快速地实现它们对于任何深度学习框架都至关重要。有许多方法可以实现矩阵乘法，最有效的方法取决于矩阵的大小、硬件和内存布局。矩阵乘法首先，考虑两个二维矩阵A和B相乘的最简单情况，其中A的形状为`[M, K]`，B的形状为`[K, N]`。结果C的形状为`[M, N]`。朴素的实现可能如下所示： `` void gemm(const float *A, const float *B, float *C, int m, int n, int k, bool ta, bool tb) { for (int i = 0; i < m; i++) for (int j = 0; j < n; j++) { float sum = 0.f; for (int l = 0; l < k; l++) { float a = ta ? A[l*m + i] : A[i*k + l]; float b = tb ? B[j*k + l] : B[l*n + j]; sum += a * b; } C[i*n + j] = sum; } } `` 不用说，这会非常慢——三个嵌套循环且没有缓存优化。在Apple Silicon上，我们可以直接使用Accelerate框架中的`cblas_sgemm`作为替代品。它使用AMX（矩阵协处理器），对于任何大于32x32的矩阵，速度是普通实现的50到100倍： `` cblas_sgemm(CblasRowMajor, ta ? CblasTrans : CblasNoTrans, tb ? CblasTrans : CblasNoTrans, m, n, k, 1.f, A, ta ? m : k, B, tb ? k : n, 0.f, C, n); `` 但我们甚至可以更进一步，使用Metal进行GPU矩阵乘法来处理大型矩阵。Metal带有`MPSMatrix`类和几个执行乘法的方法。但它仅支持二维矩阵，因此我们必须实现自己的内核来处理超过二维的张量。 `` kernel void bmatmul(device const float * a, device const float * b_, device float * c, constant int * p, uint idx[[thread_position_in_grid]]) { int M = p[1], N = p[2], K = p[3]; int tot = p[0] * M * N; if ((int) idx >= tot) return; int n = (int) idx % N, t = (int) idx / N, m = t % M, bat = t / M; float s = 0; int ao = bat * M * K + m * K, bo = bat * K * N + n; for (int k = 0; k < K; k++) s += a[ao + k] * b_[bo + k * N]; c[idx] = s; } `` 整体的`ut_matmul`函数现在看起来像这样： `` ut_tensor *ut_matmul(ut_tensor *a, ut_tensor *b) { // 特殊情况：2D x 2D if (a->shape.ndim == 2 && b->shape.ndim == 2) { int n = ..., m = ..., k = ...; // 维度：a: [m, k], b: [k, n] if (a->dev == UT_METAL || b->dev == UT_METAL) { // GPU 矩阵乘法 ut_tensor *out = ut_alloc(2, (int[]){m, n}, UT_METAL); ut_mtl_matmul(a, b, out); return out; } else { // CPU 矩阵乘法 ut_tensor *out = ut_alloc(2, (int[]){m, n}, UT_CPU); cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, m, n, k, 1.f, a->data, k, b->data, n, 0.f, out->data, n); return out; } } // 特殊情况：3D x 3D（批次） if (a->shape.ndim == 3 && b->shape.ndim == 3) { int B, n, m, k; // 维度：a: [B, m, k], b: [B, k, n] // 类似地，对于Metal：调用自定义内核，对于CPU：在每个批次上循环执行 sgemm } } `` 缺少的一点是转置矩阵的矩阵乘法。之前我们找到了一种在不复制数据的情况下“转置”矩阵的方法（仅更新步幅），但对于乘法，我们也应该以特殊方式处理它。幸运的是，所有实现似乎都支持转置，只需少量代码更改（`cblas_sgemm`中的`ta`和`tb`标志，以及Metal内核中的几行额外代码）。 ## 手动自动微分 至此，我们可以在张量上执行算术运算，并且添加新操作基本变得直接。这是构建神经网络的良好基础，但我们仍然缺少一个关键特性——正确的微分。当网络“学习”时，它会调整参数以最小化损失函数。为此，我们需要计算损失函数相对于参数的梯度。这就是自动微分通常发挥作用的地方。在PyTorch或TensorFlow等“大型”框架中，微分引擎在前向传播期间记录计算图，然后反向遍历以计算梯度。我将做一个更简单的事情，手动反向传播。每个