一、硬件架构概览：张量加速核在GPU/TPU/NPU中的位置
各位同学，咱们今天聊聊硬件架构。说实话，很多人一上来就扎进算子优化、内存搬运这些细节，结果发现跑出来的性能跟理论峰值差一大截。为什么？因为你不了解你的"武器"长什么样。
我个人习惯，拿到一块加速芯片，第一件事不是看它的算力有多高，而是先搞清楚三件事：计算单元在哪、数据怎么流、存储怎么分。这三点搞明白了，后面写代码心里才有底。
1.1 张量加速核的"生态位"
先看一张图，这是我手绘的三种主流芯片的架构对比。别嫌丑，重点看结构。
这张图里，我刻意把三种芯片放在一起对比。你想想看，GPU的SM里藏着Tensor Core，TPU的核心就是那个脉动阵列，NPU则是一堆专用核。它们的位置不同，决定了编程模型完全不同。
核心观点：张量加速核不是凭空出现的，它是在通用计算单元（如GPU的CUDA Core）基础上，专门为矩阵乘法这类规则运算设计的"加速器中的加速器"。
1.2 数据流架构：脉动阵列（Systolic Array）
好，接下来咱们聊聊脉动阵列。这东西听起来玄乎，说白了就是让数据像心跳一样，有节奏地在计算单元之间流动。
我在做TPU模拟器的时候，第一次看到脉动阵列的波形图，差点没反应过来——数据居然可以这样"流"着算！
来看一个简单的4x4脉动阵列示例：
// 伪代码：脉动阵列的矩阵乘法 (C = A * B)
// A矩阵按行流入，B矩阵按列流入
// 初始化
for i in 0..3:
for j in 0..3:
C[i][j] = 0
// 脉动计算
for k in 0..3: // 第k次脉动
for i in 0..3: // 行方向
for j in 0..3: // 列方向
// 每个PE执行乘加
C[i][j] += A[i][k] * B[k][j]
// 数据向右传递
A[i][k+1] = A[i][k] // 实际是寄存器传递
// 数据向下传递
B[k+1][j] = B[k][j] // 实际是寄存器传递
你可能会问：这不就是三重循环吗？有什么特别的？
嗯，这里要注意：脉动阵列的关键在于数据复用。传统CPU/GPU做矩阵乘法，每个数据要从内存读很多次。但在脉动阵列里，数据一旦从内存读进来，就在阵列内部"流动"，被多个PE重复使用。
避坑指南：我曾经在部署一个BERT模型时，发现TPU利用率只有30%。后来一查，是因为矩阵尺寸不是128的倍数，导致脉动阵列有很多"空泡"。解决办法是补零对齐，利用率直接飙到85%。
脉动阵列的三种常见数据流模式：
| 数据流模式 | 描述 | 典型应用 |
|---|---|---|
| 权重固定（Weight Stationary） | 权重数据固定在PE中，输入数据流动 | 卷积层（权重复用度高） |
| 输入固定（Input Stationary） | 输入数据固定在PE中，权重流动 | 全连接层（输入复用度高） |
| 输出固定（Output Stationary） | 部分和在PE中累加，输入和权重都流动 | 矩阵乘法（减少写回次数） |
我个人习惯，在写算子时先判断哪个数据复用度最高，然后选择对应的数据流模式。比如卷积层，权重在整个计算过程中只加载一次，那就用权重固定模式。
1.3 存储层次结构
存储层次，说白了就是离计算单元越近，越快也越小。这个道理大家都懂，但真正写代码时，很多人还是会把数据放在Global Memory里反复读写。
来看一个典型的GPU存储层次：
这张金字塔图，我每次培训都要画一遍。你想想看，从Global Memory读一次数据要400个周期，而从Register File读只要1个周期。这差距，比北京到上海的高铁和走路还大。
1.3.1 Register File（寄存器文件）
寄存器文件是离计算单元最近的地方。在GPU里，每个线程都有自己的寄存器，但数量有限——NVIDIA的GPU每个线程最多255个寄存器。
注意：我曾经遇到一个坑，写了一个算子用了大量局部变量，结果寄存器溢出（spill），数据被挤到Local Memory里，性能直接腰斩。后来我手动限制了每个线程的寄存器使用量，才把性能拉回来。
1.3.2 Shared Memory（共享内存）
共享内存是同一个Block内的线程共享的。它的延迟比Global Memory低一个数量级，但容量有限——通常每个SM只有48KB到128KB。
我建议你在写算子时，优先把频繁访问的数据放到Shared Memory里。比如矩阵乘法中的分块（tiling），就是把大矩阵切成小块，先搬到Shared Memory里，再让线程去算。
// CUDA示例：使用Shared Memory做矩阵分块
__global__ void matmul_tiled(float* A, float* B, float* C, int N) {
__shared__ float As[TILE_SIZE][TILE_SIZE];
__shared__ float Bs[TILE_SIZE][TILE_SIZE];
int row = blockIdx.y * TILE_SIZE + threadIdx.y;
int col = blockIdx.x * TILE_SIZE + threadIdx.x;
float sum = 0.0f;
for (int t = 0; t < N / TILE_SIZE; t++) {
// 加载数据到Shared Memory
As[threadIdx.y][threadIdx.x] = A[row * N + t * TILE_SIZE + threadIdx.x];
Bs[threadIdx.y][threadIdx.x] = B[(t * TILE_SIZE + threadIdx.y) * N + col];
__syncthreads(); // 同步，确保所有线程都加载完毕
// 计算
for (int k = 0; k < TILE_SIZE; k++) {
sum += As[threadIdx.y][k] * Bs[k][threadIdx.x];
}
__syncthreads(); // 同步，确保所有线程都计算完毕
}
C[row * N + col] = sum;
}
这段代码里，__syncthreads()是个关键点。我刚开始写CUDA时经常忘记加，结果数据还没加载完就开始算，算出来的结果全是错的。嗯，这个坑我踩过好几次。
1.3.3 Global Memory（全局内存）
Global Memory是最大的存储空间，但也是最慢的。在GPU里，它通常是HBM（高带宽内存），带宽很高（可达2TB/s），但延迟也高（400-800周期）。
对于Global Memory，我个人的经验是：能合并访问就合并访问，能少读就少读。合并访问的意思是，让连续的线程访问连续的内存地址，这样硬件可以把多次访问合并成一次。
关键技巧：在TPU/NPU上，Global Memory的访问模式更加重要。因为脉动阵列对数据流的规律性要求极高，一旦出现非对齐访问，整个阵列都要"空转"等待。
1.4 三种芯片的存储层次对比
| 层级 | GPU (NVIDIA A100) | TPU (v4) | NPU (华为昇腾910) |
|---|---|---|---|
| 寄存器 | 256KB/SM (每个线程255个) | 无独立寄存器，PE内嵌 | 每个Core内嵌 |
| 片上存储 | Shared Memory 192KB/SM | Unified Buffer 144MB | Local Memory 1MB/Core |
| 全局内存 | HBM2e 80GB (2TB/s) | HBM2 32GB (1.2TB/s) | HBM2e 32GB (1.2TB/s) |
| 延迟差异 | 寄存器:1周期 → 全局:400周期 | UB:5周期 → HBM:200周期 | Local:3周期 → HBM:300周期 |
你看这张表，GPU的寄存器到全局内存延迟差了400倍，TPU和NPU虽然好一些，但也有几十倍的差距。所以，优化张量加速核的核心，就是让数据尽量待在"快"的地方。
最后说一句，存储层次的理解，直接决定了你能写出多高效的算子。我见过太多人，算法写得天花乱坠，结果一跑性能，全卡在内存搬运上。记住：计算是免费的，数据搬运才是成本。
个人建议：刚开始接触张量加速核时，先别急着写复杂算子。拿一个简单的矩阵乘法，手动做分块、做数据搬移，把每个存储层次都用上。跑通了，你就入门了。
想要更多资料，联系可可