硬件基础与计算原语：AI芯片的核心计算单元
各位同学，咱们今天聊点硬核的。做AI算子加速，说白了就是跟芯片硬件打交道。你写的那一行代码，最终都要映射到芯片上某个晶体管的状态变化。我刚开始做这行时，总觉得软件层抽象一下就行了，直到有一次在FPGA上调试一个卷积算子，死活跑不出预期性能……嗯，后来发现是对硬件计算单元的理解不够深。
这一节，我们就把AI芯片的「家底」翻出来看看。你想想看，一个算子要跑得快，靠的是什么？靠的是计算单元够猛、数据搬得够快、存储层次够合理。这三样缺一不可。
核心计算单元：MAC与PE
先说最基础的——MAC（Multiply-Accumulate，乘累加）。这玩意儿是AI计算的灵魂。一个MAC单元，一次完成一个乘法和一个加法。就这么简单，但架不住数量多啊。
我习惯把MAC比作「数字世界的砖头」。你盖一栋楼需要几万块砖，跑一个神经网络需要几亿甚至几十亿次MAC操作。芯片设计者干的事，就是在芯片上堆MAC单元，越多越好。
关键点：一个MAC = 一次乘法 + 一次加法。现代AI芯片里，MAC单元通常以PE（Processing Element，处理单元）的形式组织。一个PE里可能包含多个MAC。
举个例子，英伟达的Tensor Core，一个PE里能同时做多个低精度MAC。华为的达芬奇架构，一个Cube单元里塞了密密麻麻的MAC阵列。说白了，算力就是靠堆MAC堆出来的。
我在项目中遇到过一个问题：MAC单元利用率上不去。明明芯片理论算力很高，实际跑起来只有30%的利用率。后来发现是数据供给跟不上，MAC单元在「等米下锅」。这就是我们后面要讲的存储层次的问题。
数据流架构：行/列/脉动阵列
MAC单元有了，怎么把它们组织起来干活？这就涉及到数据流架构了。常见的三种：行式、列式、脉动阵列。
| 架构类型 | 特点 | 典型芯片 |
|---|---|---|
| 行式（Row-wise） | 数据按行流动，每个PE处理一行 | 早期DSP |
| 列式（Column-wise） | 数据按列流动，适合矩阵乘法 | 某些FPGA实现 |
| 脉动阵列（Systolic Array） | 数据像心跳一样在PE间传递 | Google TPU、华为昇腾 |
脉动阵列是我个人最喜欢的架构。为什么？因为它把数据复用做到了极致。你想想看，一个权重数据进来，可以在阵列里被多个PE重复使用，不用反复从内存里搬。这正好解决了前面说的「数据供给跟不上」的问题。
我的经验：脉动阵列的调度是个技术活。我曾经在TPU上优化一个矩阵乘算子，发现数据流方向搞反了，导致PE间通信拥堵。后来改成「权重固定、输入流动」的模式，性能直接翻倍。
这里我画了一张图，帮你理解脉动阵列的数据流动方式：
存储层次：L1/L2/Global Memory
好了，计算单元有了，数据流架构也定了。但数据从哪来？这就得聊存储层次了。
AI芯片的存储，我习惯把它看成「金字塔」结构：
- L1（寄存器/本地存储）：离PE最近，速度最快，但容量最小。通常几十KB到几百KB。
- L2（共享存储）：多个PE共享，容量几MB，速度比L1慢一些。
- Global Memory（全局存储）：通常是HBM或DDR，容量大（几GB到几十GB），但访问延迟高。
你想想看，一个数据从Global Memory搬到L1，可能要几百个时钟周期。但从L1读到PE里，只要1个时钟周期。这差距有多大？所以算子优化的核心，说白了就是「让数据尽量待在L1里，别老往Global Memory跑」。
避坑指南：我曾经在一个项目里，把所有的权重都放在Global Memory里，每次计算都去读。结果MAC利用率不到10%。后来改成把权重提前搬到L2，再分块搬到L1，利用率直接拉到80%以上。记住：存储层次用得好，算力才能发挥出来。
这里有个简单的代码示例，展示如何在算子中手动管理存储层次：
// 伪代码：矩阵乘法的存储层次优化
void matmul_optimized(float* A, float* B, float* C, int M, int N, int K) {
// 分块大小，根据L1容量决定
const int BLOCK_M = 64;
const int BLOCK_N = 64;
const int BLOCK_K = 16;
// 在L1中分配本地存储
__local__ float A_local[BLOCK_M][BLOCK_K];
__local__ float B_local[BLOCK_K][BLOCK_N];
for (int i = 0; i < M; i += BLOCK_M) {
for (int j = 0; j < N; j += BLOCK_N) {
// 从Global Memory搬数据到L1
load_block(A, A_local, i, j, M, N, K);
load_block(B, B_local, i, j, M, N, K);
// 在L1中完成计算
for (int k = 0; k < K; k += BLOCK_K) {
// PE阵列计算
systolic_compute(A_local, B_local, C, i, j, k);
}
}
}
}片上互联总线
最后说说互联总线。这东西就像芯片的「血管」，负责把各个部件连起来。常见的互联方式有：
- 总线型（Bus）：所有设备共享一条总线，简单但容易拥堵。
- 交叉开关（Crossbar）：任意两个设备可以直连，速度快但面积大。
- 网状网络（Mesh）：每个节点只和邻居相连，适合大规模芯片。
我个人比较喜欢Mesh结构。为什么？因为AI计算的数据流是有规律的，邻居之间的通信模式很固定。你想想看，脉动阵列里PE之间就是邻居通信，用Mesh天然匹配。
我的经验：在设计互联总线时，带宽和延迟要平衡。我曾经见过一个芯片，计算单元很强，但互联总线带宽不够，结果数据堵在总线上，计算单元闲着。这就好比高速公路修得很宽，但收费站只有两个口——全堵在那了。
好了，这一节的内容就这些。记住三个关键词：MAC单元是基础，脉动阵列是灵魂，存储层次是关键。下一节我们开始讲具体的算子实现，到时候你会感谢今天打下的硬件基础。
公众号:蓝海资料掘金营，微信deep3321