Tensor芯片展现强劲实力

背景

Tensor（张量）是深度学习的数学基石，从简单的两层MLP到GPT-5，所有神经网络本质上都是浮点数在操作图中流动。本文作者从零开始用C语言构建了一个完整的加速张量库，灵感主要来自Fabrice Bellard的libnc（遗憾的是该库尚未开源）。文章旨在揭示张量这一优雅数学抽象的本质：它解决了如何高效表示和操作多维数据的问题。

核心内容

张量的本质

张量本质上就是一个扁平的浮点数数组，加上一些元数据（metadata），告诉如何将这些数解释为多维对象。二维数组可以更好地表示为一维数组加上行/列数——张量本质上就是这一思想的推广。

对于超过二维的情况，需要额外的元数据，例如通用的形状（shape）：

float data[32 * 3 * 28 * 28]; // 32张图片，3个通道，28x28像素
int shape[4] = {32, 3, 28, 28}; // 张量的形状
int ndim = 4; // 维度数量

有了形状，在行优先C顺序（row-major C-order，大多数现代张量库的默认布局）下，4D张量中位置[n,c,h,w]的元素位于偏移量data[n*(3*28*28)+c*(28*28)+h*28+w]。

每次都这样手动计算索引效率低下，因此可以在已知形状后预计算步长（strides）。步长告诉我们在给定维度上每前进一个元素需要跳过多少个元素。

张量结构体定义

将形状相关字段组织到一起：

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_new(3, (int[]){2, 3, 4}) 得到nelem=24，元素[1][2][3]的扁平索引为1*12 + 2*4 + 3 = 23。

内存管理

张量通常动态分配，提供创建和销毁的封装（基于malloc/free）。有时需要创建共享同一数据的张量（如视图、转置、改变形状），这就要求跟踪数据所有权。为此增加引用计数（refcount）字段，以及指向所有者张量的指针：

struct ut_tensor {
    struct ut_shape shape;
    float *data;
    int refcount;          // 引用计数，减到0时释放
    struct ut_tensor *owner; // 如果此张量是视图，指向所有者
};

还可进一步使用arena分配器或内存池优化，但当前设计已是不错的起点。

逐元素操作

最基本的张量操作是逐元素（elementwise）操作——单循环遍历所有元素，对每个元素应用函数（一元）或每对元素应用函数（二元）。

一元操作示例：ew_neg, ew_relu 二元操作示例：ew_add

统一封装为：

• static ut_tensor* ew_unary(ut_tensor* a, void (*fn)(float*, const float*, int))
• static ut_tensor* ew_binary(ut_tensor* a, ut_tensor* b, void (*fn)(float*, const float*, const float*, int))

进而实现高层API：ut_neg, ut_exp, ut_sigmoid, ut_tanh, ut_relu, ut_add, ut_sub, ut_mul, ut_div, ut_scale。

在二元操作前应assert形状相同。虽然可以支持广播（broadcasting），但作者决定暂时不做——因为目标模型中的张量形状通常完美对齐，若不齐可显式复制数据或后续添加广播索引计算。

GPU加速

单纯在CPU上运行操作很快就会遇到性能瓶颈。即使使用BLAS、LAPACK、OpenMP，CPU-only的张量性能也无法与GPU加速框架相提并论。因此需要将部分操作卸载到GPU。

难点在于GPU厂商众多，各自有不同API和特性：CUDA、OpenCL、WebGPU、Metal、Vulkan等。作者选择从Apple的Metal开始，虽然限制在Apple设备上，但Apple Silicon的统一内存（unified memory）简化了问题。

Metal使用自己的语言MSL（Metal Shading Language）定义GPU内核，类似于现代C++。内核可在运行时编译，因此可以直接在C代码中以字符串形式定义：

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";

要让内核在GPU上运行，需要创建设备（device）、构建命令队列（command queue，GPU是异步的）、编译内核、创建缓冲区等。

关键要点

• 张量是扁平数组+形状元数据：多维数据只需一个连续内存块和描述各维度大小的shape数组即可表示，步长可预计算加速索引。
• 行优先C顺序是默认布局：大多数现代张量库采用row-major C-order，多维索引转换为扁平索引的公式为：从最后一个维度开始累加。
• 引用计数实现数据共享：通过refcount和owner指针支持视图、转置、形状修改而不复制数据，并允许张量在原始作用域之外存活（如迭代训练中避免重复分配）。
• 逐元素操作是基础：所有一元和二元操作可以通过统一循环实现，对外提供简洁的API。
• 广播支持暂不实现：当前目标模型张量形状已对齐，广播可以后续通过显式对齐或广播索引计算添加。
• GPU加速是性能关键：即使CPU优化（BLAS、OpenMP等）也无法达到GPU性能；选择从Metal开始，利用Apple Silicon的统一内存降低复杂度。
• Metal内核用MSL字符串定义：MSL类似现代C++，运行时编译，可直接嵌入C代码字符串。

意义与影响

本文以实战方式清晰阐明了张量抽象的核心设计原理——从结构定义、索引计算、内存管理到逐元素操作，再到GPU加速的选择。它不仅是一篇技术教程，更展示了如何从零构建一个最小可行张量库（类似微型PyTorch/TensorFlow），让读者理解现代深度学习框架底层的工作机制。对于想深入理解张量库实现（如CUDA内存管理、步长计算）的开发者，这是一份宝贵的入门指南。同时，作者对Metal的选择也反映了在Apple生态中构建高效张量运算的务实路径，为后续扩展至其他GPU API（CUDA、Vulkan等）奠定了基础。