2. AI计算基础
深度学习模型从数学本质上看，是一张由张量（Tensor）流动构成的计算图。无论是ResNet、Transformer还是YOLO，其底层计算都可以拆解为有限种类的核心算子。AI芯片设计的首要任务，就是深刻理解这些算子的计算模式、访存特征与并行潜力，因为芯片的微架构、存储层次、数据流策略乃至指令集设计，都直接由算子特性驱动。本章将系统剖析卷积、矩阵乘、激活、归一化等核心算子的数学定义、计算特性，并建立从算法到硬件映射的桥梁。
2.1 卷积算子：AI芯片的“第一公民”
核心原理：卷积的数学本质与计算展开
卷积神经网络（CNN）中的二维卷积（Conv2D）本质上是滑窗内的乘累加（MAC）运算。对于输入特征图 \( X \in \mathbb{R}^{C_{in} \times H \times W} \) 和卷积核 \( W \in \mathbb{R}^{C_{out} \times C_{in} \times K_h \times K_w} \)，输出特征图 \( Y \in \mathbb{R}^{C_{out} \times H_{out} \times W_{out}} \) 的每个元素计算为：
Y[n][co][oh][ow] = bias[co] + Σ_{ci=0}^{Cin-1} Σ_{kh=0}^{Kh-1} Σ_{kw=0}^{Kw-1}
X[n][ci][oh*stride + kh - pad][ow*stride + kw - pad] * W[co][ci][kh][kw]从硬件视角看，卷积计算具有三个关键特性：高计算密度（每个输出点需要 \( C_{in} \times K_h \times K_w \) 次MAC）、权重复用（同一卷积核在空间维度滑动时被反复使用）、输入数据复用（相邻输出窗口共享大量输入像素）。这些特性决定了AI芯片中数据流（Dataflow）设计的核心权衡：是优先复用权重（Weight Stationary）、优先复用输入（Input Stationary）还是优先累积输出（Output Stationary）。
判断方法：卷积的计算强度与并行维度分析
在决定硬件资源分配前，工程师需要量化卷积算子的计算强度（Arithmetic Intensity, AI = FLOPs / Bytes）。以典型的3×3卷积、256通道、112×112特征图为例：
| 参数 | 数值 | 说明 |
|---|---|---|
| 单输出点MAC数 | 3×3×256 = 2,304 | 计算密度极高 |
| 输出特征图大小 | 256×112×112 = 3.2M | 总FLOPs ≈ 7.4 GFLOPs |
| 权重总量 | 256×256×3×3 = 590K | 约2.36 MB (FP16) |
| 输入特征图 | 256×112×112 = 3.2M | 约6.4 MB (FP16) |
| 计算强度 | 7.4G / (2.36M+6.4M) ≈ 845 FLOPs/Byte | 远高于roofline拐点 |
这意味着卷积是计算密集型算子，芯片设计应优先保证MAC阵列的利用率，而非过度追求带宽。并行维度上，卷积天然具备四重并行性：Batch维（N）、输出通道维（Cout）、空间维（H/W）、输入通道维（Cin），为硬件并行调度提供了丰富选择。
案例拆解：YOLOv5s骨干网络的卷积硬件映射
以YOLOv5s的Focus模块和CSP结构为例，其第一层卷积输入为640×640×3，输出32通道，卷积核3×3。在部署到某AI芯片时，我们面临以下约束：MAC阵列规模为128×128，片上SRAM为2MB，外部带宽为64GB/s。
初始现象：直接按输出通道顺序映射，MAC利用率仅62%，SRAM带宽瓶颈明显。分析过程：通过profiling发现，输入特征图在行方向上的滑动窗口复用未被充分利用，导致输入数据重复从DRAM读取。优化操作：采用行缓冲（Line Buffer）机制，在片上缓存输入特征图的连续3行，配合权重的全缓存（32×3×3×3=864个权重仅需1.7KB），实现输入数据的行内复用。结果对比：MAC利用率提升至91%，DRAM访问量减少约40%。复盘经验：小卷积核（3×3）场景下，行缓冲+权重驻留是最优数据流策略，但需注意行缓冲的SRAM开销与卷积核尺寸的匹配。
2.2 矩阵乘（GEMM）：Transformer与全连接层的基石
核心原理：GEMM的硬件友好性
矩阵乘 \( C = A \times B \) 是Transformer中Self-Attention（QK^T、PV）、MLP层以及CNN中全连接层的核心运算。从硬件角度看，GEMM具有极其规整的计算模式：输出矩阵的每个元素是输入矩阵一行与一列的内积。这种规整性使得GEMM可以高效映射到脉动阵列（Systolic Array）或SIMD阵列上。以Google TPU的脉动阵列为例，256×256的MAC阵列可以在一个时钟周期内完成65,536次MAC运算，数据以“波前”形式在阵列中流动，极大减少了寄存器访问开销。
判断方法：GEMM的切分（Tiling）策略
实际矩阵规模往往远超片上存储容量，因此必须进行分块计算。以BERT-Large的QK^T为例，序列长度512，头维度64，批量32，则矩阵规模为（32×512）×（512×64）。切分策略需同时考虑三个维度：
- M维度切分（Batch×Seq）：决定并行度，通常切分为与MAC阵列行数匹配的块。
- N维度切分（Head dim）：决定输出累加粒度，影响部分和（Partial Sum）的缓存需求。
- K维度切分（Seq dim）：决定数据复用距离，K块越大，A矩阵复用率越高，但需要更大SRAM。
一个实用的经验法则是：K维度切分大小应使得A矩阵的单个分块恰好放入片上SRAM，同时B矩阵分块流式加载，从而最大化计算强度。
案例拆解：推理时的GEMM性能优化
在部署（117M参数）时，我们观察到MLP层的GEMM（输入3072×768，权重768×3072）在GPU上表现良好，但在某AI芯片上仅达到理论峰值的45%。关键输入数据：MAC阵列64×64，SRAM 1MB，DRAM带宽32GB/s。分析判断：通过Roofline模型分析，发现该GEMM的计算强度为 2×3072×768×3072 / (3072×768×2 + 768×3072×2) ≈ 1536 FLOPs/Byte，远高于硬件拐点，理论上应接近峰值。问题出在权重矩阵的存储布局——原始权重按行主序存储，导致读取B矩阵分块时DRAM访问效率低下。实际操作：在编译器中增加权重重排（Weight Reshuffle）步骤，将权重转换为分块列主序（Blocked Column-Major），并调整DMA描述符实现双缓冲。结果对比：MAC利用率从45%提升至83%，端到端推理延迟降低38%。风险提示：权重重排增加了编译时间（约2.3秒），但推理阶段无额外开销。
2.3 激活函数与归一化：内存密集型算子的挑战
核心原理：非线性与统计特性的硬件实现
激活函数（ReLU、GELU、Swish）和归一化（BatchNorm、LayerNorm）虽然计算量远小于卷积和GEMM，但其内存访问模式对芯片设计提出了独特挑战。ReLU只需比较操作，但GELU涉及指数和误差函数运算；BatchNorm在推理时可融合进卷积权重，但训练时需要计算均值/方差并更新；LayerNorm需要对每个样本的特征维度做归约，涉及跨通道的均值/方差计算。
判断方法：算子分类与硬件资源分配
| 算子类型 | 计算特性 | 访存特性 | 硬件映射建议 |
|---|---|---|---|
| ReLU/Clip | 极低（1次比较） | 逐元素访存 | 融合进前序算子，零额外开销 |
| GELU/Swish | 中等（含exp/erf） | 逐元素访存 | 查表法（LUT）或多项式近似 |
| BatchNorm（推理） | 低（2次乘加） | 逐通道访存 | 折叠进卷积权重，消除独立算子 |
| LayerNorm | 中等（均值/方差） | 跨通道归约 | 专用归约单元，避免通用ALU |
| Softmax | 中等（exp+归约） | 跨行归约 | 分段计算，在线Softmax算法 |
关键判断原则：激活函数应尽量“融合”到计算算子中，避免产生额外的片上-片外数据搬移。例如，在卷积输出写回SRAM前，直接在MAC阵列输出端插入ReLU逻辑，可节省一次完整的特征图读写。
案例拆解：Transformer LayerNorm的在线计算优化
在训练一个12层Transformer模型时，LayerNorm算子占用了约15%的端到端时间。初始现象：LayerNorm需要先计算均值再计算方差，导致两次完整的数据读取。关键参数：序列长度512，隐藏维度768，批量32。分析判断：传统两遍算法需要读取2×512×768×32×2字节 ≈ 50MB数据，而片上SRAM仅2MB。实际操作：采用Welford在线算法，单次遍历同时更新均值和方差，将数据读取量减半；同时利用向量ALU的SIMD能力，一次处理8个维度。结果对比：LayerNorm耗时降低42%，整体训练速度提升6%。复盘经验：归一化算子的优化核心是减少数据遍历次数，而非提升MAC利用率。
2.4 算子融合与计算图优化：系统级设计视角
核心原理：融合的收益与约束
算子融合（Operator Fusion）是将多个连续算子合并为一个内核（Kernel）执行，以减少中间数据的片上-片外搬移。典型的融合模式包括：Conv+BN+ReLU融合、QKV投影GEMM融合、Attention中的Softmax+Scale融合。融合的收益直接体现在内存访问量减少上：以ResNet-50为例，若不融合，中间特征图需反复写回DRAM，融合后可将DRAM流量降低约30%。
判断方法：融合可行性分析框架
- 数据依赖检查：算子间是否存在元素级映射（Element-wise）或归约依赖？ReLU是元素级，可融合；Pooling是窗口级，需谨慎。
- 计算强度匹配：融合后算子的计算强度是否仍高于硬件拐点？若融合导致计算强度下降，可能得不偿失。
- 片上存储预算：融合需要将中间结果暂存于SRAM，需评估是否超出容量。
- 数值精度影响：BN折叠到卷积时，需确保FP16/INT8精度下数值稳定。
案例拆解：MobileNetV3的硬件感知融合
MobileNetV3使用深度可分离卷积（Depthwise Conv + Pointwise Conv），其计算密度远低于标准卷积。初始目标：在边缘AI芯片（算力4 TOPS，带宽12.8GB/s）上达到实时推理（30 FPS）。关键约束：深度卷积的计算强度仅约50 FLOPs/Byte，属于内存密集型。分析判断：若将深度卷积与后续的Pointwise卷积分离执行，DRAM流量将翻倍。实际操作：在芯片中设计专用的“深度卷积+激活+Pointwise”融合流水线，深度卷积的输出直接留在SRAM中，不写回DRAM；同时将BN参数折叠进Pointwise卷积权重。结果对比：DRAM流量降低52%，实际帧率从18 FPS提升至31 FPS。风险与代价：融合流水线增加了控制逻辑复杂度，且对动态形状（Dynamic Shape）支持较差。复盘经验：对于内存密集型算子，融合是比提升MAC利用率更有效的优化手段。