第1章：数据流与内存层次优化：从寄存器到HBM

大家好，我是你们的AI芯片架构课讲师。今天咱们聊一个特别实在的话题——数据搬移的代价。

说实话，我做了这么多年AI芯片，见过太多团队在算力上堆得飞起，结果性能就是上不去。为什么？说白了，数据没喂饱计算单元。你想想看，一个MAC阵列再强，如果数据搬不过来，那就是在空转。

1.1 存储层次：从寄存器到HBM的距离

先看一张图，这是我个人习惯画的存储层次结构：

这张图我每次讲课都会放。为什么？因为数据搬移的代价是指数级增长的。从寄存器到HBM，延迟差了三个数量级。我在项目中遇到过，有人把权重放在HBM里，每次计算都去读，结果功耗直接爆表。

1.2 数据搬移的真实代价

咱们算笔账。一个典型的卷积层，假设输入是224x224x3，权重是3x3x3x64。如果每次计算都从HBM读数据：

• 输入特征图：224*224*3 ≈ 150KB
• 权重：3*3*3*64 ≈ 1.7KB
• 输出特征图：224*224*64 ≈ 3.1MB

你想想看，3.1MB的数据从HBM搬到计算单元，一次就要几千个cycle。一个ResNet-50有上百层，这得搬多少次？

1.3 数据复用策略：三种固定模式

怎么解决？说白了就是让数据在片上多待一会儿。我总结了三种经典的数据复用模式：

1.3.1 权重固定（Weight Stationary）

这是我最常用的策略。把权重加载到局部存储里，然后反复使用。比如一个3x3的卷积核，权重只有9个数，加载一次后可以跟整个输入特征图做卷积。

// 权重固定伪代码
for (int oc = 0; oc < output_channels; oc++) {
    // 只加载一次权重
    load_weight_to_local(weight[oc], local_weight_buf);
    
    for (int oh = 0; oh < output_height; oh++) {
        for (int ow = 0; ow < output_width; ow++) {
            // 从局部存储读权重，从HBM读输入
            int sum = 0;
            for (int kh = 0; kh < kernel_h; kh++) {
                for (int kw = 0; kw < kernel_w; kw++) {
                    sum += local_weight_buf[kh][kw] * 
                           input[oh+kh][ow+kw];
                }
            }
            output[oc][oh][ow] = sum;
        }
    }
}

这样做的好处是权重只搬一次，后面都在局部存储里。我在做第一代NPU时，就是用这个策略把带宽需求降了10倍。

1.3.2 输入固定（Input Stationary）

反过来，把输入特征图固定在片上。适合那种输入通道数少、输出通道数多的场景。比如depthwise卷积，输入就一个通道，但输出通道很多。

1.3.3 输出固定（Output Stationary）

把部分和（partial sum）留在片上累加。这个在矩阵乘法里特别常见。比如做矩阵乘法C = A * B，每次算出一部分结果，先存在局部存储里，等所有部分和算完再写回HBM。

// 输出固定伪代码
for (int i = 0; i < M; i += tile_m) {
    for (int j = 0; j < N; j += tile_n) {
        // 初始化局部输出缓冲区
        clear_local_output(local_out_buf);
        
        for (int k = 0; k < K; k += tile_k) {
            // 加载A和B的分块
            load_tile(A, i, k, local_a_buf);
            load_tile(B, k, j, local_b_buf);
            
            // 在局部累加
            matmul_tile(local_a_buf, local_b_buf, local_out_buf);
        }
        
        // 只写回一次最终结果
        store_output(local_out_buf, C, i, j);
    }
}

1.4 如何选择复用策略？

没有银弹。我一般根据这几个维度来选：

策略	适用场景	局部存储需求	带宽节省效果
权重固定	卷积层、全连接层	中等（存权重）	★★★★★
输入固定	Depthwise卷积、逐点卷积	大（存输入特征图）	★★★★
输出固定	矩阵乘法、大卷积核	中等（存部分和）	★★★★

1.5 实战中的避坑指南

我做了这么多年，踩过的坑不少。分享几个：

• 别盲目追求复用：我曾经在一个项目里，为了最大化数据复用，把局部存储塞得满满当当。结果调度器复杂度爆炸，控制逻辑的功耗比数据搬移还高。
• 注意数据生命周期：权重固定模式下，权重在局部存储里待多久？如果网络层数太多，权重可能被覆盖，那就白费功夫了。
• 混合策略才是王道：我现在的做法是，不同层用不同策略。比如前几层用权重固定，中间层用输出固定，最后几层用输入固定。灵活搭配，效果最好。

好了，这一章就到这里。记住，别光看理论，回去看看你的设计里，数据搬移占了多大比例。动手优化一下，效果立竿见影。