2. AI计算基础

在深入AI编译器与NPU指令调度之前,必须建立对AI计算本身的精确认知。AI编译器本质上是一个“翻译与优化器”,它将高层神经网络描述(如PyTorch、TensorFlow模型)映射到低层硬件指令(如NPU的标量、向量、矩阵指令)。如果对算子语义、张量布局、计算图拓扑或数据流方向理解有偏差,后续的调度优化将失去依据。本章将系统梳理这五个核心概念,并建立它们与NPU硬件执行之间的映射关系。

AI计算基础概念与NPU映射关系 常见AI模型结构 CNN / Transformer / LSTM 计算图 (Graph) 节点=算子,边=张量依赖 算子 (Operator) Conv / MatMul / Softmax 张量 (Tensor) Shape / Stride / Dtype 数据流 (Dataflow) 前向 / 反向 / 梯度传播 NPU指令调度优化 算子融合 / 内存分配 / 指令流水 / 数据搬运 硬件约束 - 内存带宽 - 算力峰值 - 缓存容量 - 指令流水 优化目标 - 延迟降低 - 吞吐提升 - 功耗下降 - 内存占用

上图展示了本章核心概念之间的层级关系:模型结构由计算图描述,计算图由算子和张量构成,数据流驱动算子的执行顺序,最终所有信息汇聚到NPU指令调度层。硬件约束(内存带宽、算力峰值)与优化目标(延迟、吞吐)共同决定了调度策略的取舍。下面逐一展开。

2.1 张量:AI计算的基本数据单元

核心原理

张量(Tensor)是多维数组的泛化,是AI模型中所有数据(输入、权重、中间激活、梯度)的载体。在NPU视角下,张量不仅仅是逻辑上的多维数组,更关键的是其在物理内存中的布局方式。一个张量由以下属性完整描述:

  • Shape(形状):各维度大小,如 [N, C, H, W] 表示批大小、通道数、高度、宽度。
  • Stride(步长):每个维度上相邻元素之间的内存偏移量,决定了张量的内存布局是连续还是非连续。
  • Dtype(数据类型):如 FP32、FP16、INT8、BF16,直接影响计算精度和内存带宽需求。
  • Layout(布局):如 NCHW、NHWC、NC1HWC0(华为NPU的5D格式),不同布局对应不同的内存访问模式。

NPU编译器的一个核心任务就是将张量从框架默认布局转换为硬件最优布局。例如,卷积运算在NPU上通常需要将通道维对齐到硬件向量单元宽度(如32通道对齐),否则会产生大量padding浪费。

判断方法:张量布局对性能的影响

在评估一个模型在NPU上的潜在性能时,需要快速判断张量布局是否合理。以下是一个三步检查法:

  1. 检查内存连续性:使用 tensor.is_contiguous() 或查看 stride 是否与 shape 匹配。非连续张量会导致NPU无法直接发起DMA搬运,需要额外的拷贝操作。
  2. 检查通道对齐:对于卷积和矩阵乘算子,确认通道维度是否满足NPU的硬件对齐要求(如16、32、64的倍数)。若不满足,编译器会插入填充(padding)操作,增加计算量。
  3. 检查数据类型:确认是否使用了NPU高效支持的类型。例如,INT8量化在多数NPU上可获得2-4倍于FP16的吞吐,但需要确认量化精度损失是否在可接受范围。

案例拆解:NCHW vs NHWC 布局选择

以经典的ResNet-50第一层卷积为例,输入张量为 [1, 3, 224, 224](NCHW)。在GPU上,NCHW是默认布局;但在多数NPU上,NHWC布局更优,原因在于:

  • NHWC将通道维放在最后,使得内存访问时同一像素的多个通道可以连续读取,便于向量化加载。
  • NPU的卷积单元通常按行处理输出像素,NHWC布局下输出行数据在内存中连续,减少写回延迟。

实际测试数据表明,在同等算力下,NHWC布局比NCHW布局在NPU上的卷积算子性能提升约15%-30%,主要收益来自内存访问效率的提升。但NHWC布局在BatchNorm等逐通道操作中需要额外的维度变换,编译器需要权衡整体收益。

经验提醒:在编写NPU算子时,不要盲目沿用GPU的布局习惯。建议先查阅目标NPU的指令集文档,了解其向量单元和矩阵单元的原生数据格式,再决定张量布局。

2.2 算子:计算的基本单元

核心原理

算子是神经网络中的基本计算单元,对应计算图中的一个节点。从NPU指令调度的角度看,算子可以按计算特征分为三类:

算子类别 典型示例 计算特征 NPU硬件映射
矩阵/向量密集计算 Conv2D、MatMul、Gemm 高算术强度,数据可复用 矩阵单元(如脉动阵列)
逐元素/归约操作 ReLU、Add、Softmax、LayerNorm 低算术强度,内存带宽受限 向量单元(SIMD)
数据搬运/变形 Reshape、Transpose、Concat、Split 无计算,纯内存操作 DMA引擎、内存重映射

理解算子类别是编译器进行算子融合(Operator Fusion)的基础。例如,Conv + ReLU + Add 可以融合为单个NPU指令序列,避免中间张量的内存写回和读取。

操作流程:算子性能分析四步法

  1. 计算量估算:使用FLOPs公式计算算子的理论计算量。例如,Conv2D的FLOPs = 2 × Cout × Hout × Wout × Cin × K × K。
  2. 访存量估算:计算输入、输出、权重的总字节数。访存量 = (输入字节 + 输出字节 + 权重字节) × 重复读取次数。
  3. 算术强度计算:算术强度 = 计算量(FLOPs) / 访存量(Bytes)。若算术强度低于硬件平衡点(Roofline模型中的转折点),则为访存受限算子。
  4. 瓶颈判定:对比算子的算术强度与NPU的硬件平衡点。若低于平衡点,优化方向是减少访存(如算子融合、数据复用);若高于平衡点,优化方向是提高计算效率(如指令流水、向量化)。

案例拆解:Softmax算子的访存优化

Softmax在Transformer模型中频繁出现,其计算过程为:对每个向量求最大值(数值稳定性)、计算指数和、归一化。在NPU上,Softmax的典型实现需要三次遍历数据:第一次求最大值,第二次计算指数和,第三次归一化。这意味着输入数据需要从内存读取三次。

优化方案是使用在线Softmax算法:在单次遍历中同时维护最大值和指数和,将三次遍历压缩为一次。对于长度为1024的向量,该优化可将访存量减少67%。在NPU上,由于向量单元和内存之间带宽有限,这种优化带来的性能提升往往比计算优化更显著。

风险提示:在线Softmax虽然减少了访存,但引入了额外的分支判断和状态维护,在向量化程度较低的NPU上可能导致指令流水线停顿。实施前应通过profiling确认当前瓶颈确实是访存而非计算。

2.3 计算图:AI模型的中间表示

核心原理

计算图(Computational Graph)是深度学习框架的中间表示(IR),它将模型描述为有向无环图(DAG)。图中的节点代表算子,边代表张量依赖关系。计算图分为静态图和动态图两种:

  • 静态图(如TensorFlow 1.x、TorchScript):先定义完整图结构,再执行。优点是可进行全局优化(如内存规划、算子融合),缺点是调试困难。
  • 动态图(如PyTorch Eager模式):边执行边建图。优点是灵活易调试,缺点是难以进行跨算子全局优化。

AI编译器(如TVM、XLA、MindSpore Graph Kernel)的核心工作之一就是将动态图转换为静态图表示,以便进行全局调度优化。NPU编译器接收的输入通常是静态计算图(如ONNX、MindIR格式)。

判断方法:计算图优化的关键检查点

在将计算图映射到NPU之前,编译器需要执行以下检查:

  1. 图是否可微分:反向传播所需的梯度算子是否已包含在图中。
  2. 是否存在数据依赖瓶颈:图中是否存在长串的串行依赖链,导致NPU并行度无法充分利用。
  3. 是否有可融合的算子对:相邻的逐元素算子(如ReLU+Add)是否可融合为单个算子。
  4. 是否有可消除的中间张量:某些算子的输出仅被一个算子使用,且两者可融合,则中间张量可消除。

案例拆解:Transformer Block的计算图优化

以的一个Transformer Block为例,其计算图包含约30个算子。在NPU上直接按原始图执行时,存在以下问题:

  • LayerNorm被拆分为多个小算子(Mean、Sub、Square、Add、Rsqrt、Mul),每个算子都需要读写中间张量。
  • 多头注意力中的Reshape和Transpose操作产生大量非连续内存访问。

优化策略:将LayerNorm的6个小算子融合为单个LayerNorm算子,将Reshape+Transpose+Reshape序列融合为一次内存重映射操作。优化后,算子数量从30个减少到18个,中间张量内存占用减少约40%,端到端延迟降低约25%。

关键认知:计算图优化的本质是减少数据搬运和中间存储,而非减少计算量。NPU的算力通常不是瓶颈,内存带宽才是。因此,图优化应优先关注如何减少张量的内存读写次数。

2.4 数据流:前向与反向传播的执行逻辑

核心原理

数据流描述了张量在计算图中的传播方向。在训练场景下,数据流分为前向传播(Forward)和反向传播(Backward)。前向传播从输入张量开始,逐层计算输出;反向传播从损失函数开始,逐层计算梯度。NPU编译器需要同时优化两个方向的数据流。

数据流的一个关键特性是数据依赖:一个算子必须等待其所有输入张量就绪后才能执行。这种依赖关系决定了NPU指令调度的并行度。例如,在Transformer中,多头注意力的多个头之间没有数据依赖,可以并行执行;但注意力输出与FFN(前馈网络)之间存在顺序依赖。

操作流程:数据流分析在NPU调度中的应用

  1. 构建依赖图:从计算图中提取每个算子的输入/输出张量,构建数据依赖关系。
  2. 关键路径分析:找出从输入到输出的最长依赖链,这条链决定了理论最小延迟。
  3. 并行度识别:标记出没有依赖关系的算子组,这些算子可以并行调度到NPU的不同计算单元。
  4. 内存生命周期管理:根据数据流方向,确定每个张量的创建和销毁时间点,优化内存复用策略。

案例拆解:梯度累积与数据流优化

在大模型训练中,由于NPU内存有限,无法一次性装入整个batch,通常采用梯度累积策略。数据流视角下,梯度累积将原本的“前向-反向-更新”串行流变为“前向-反向×N-更新”的流水线。这带来一个优化机会:

在标准实现中,每个micro-batch的反向传播需要等待其前向传播完成。但多个micro-batch的前向传播可以并行执行(如果内存允许)。通过将前向传播和反向传播解耦,可以实现类似CPU流水线的重叠执行。在NPU上,这意味着前向计算的矩阵单元和反向计算的向量单元可以同时工作,提高硬件利用率。

实践建议:在配置NPU训练任务时,不要简单地将batch size调大。应优先考虑梯度累积步数与NPU内存容量的匹配关系,通过数据流分析找到前向和反向计算重叠的最佳比例。

2.5 常见AI模型结构及其计算特征

核心原理

不同模型结构在NPU上表现出截然不同的计算特征。理解这些特征是进行指令调度优化的前提。下表对比了三种主流模型结构的计算特征:

模型结构 代表模型 主要算子 计算特征 NPU优化重点
CNN ResNet、MobileNet Conv2D、BatchNorm、Pooling 高计算密度,权重复用率高 卷积算子优化、通道对齐、Winograd变换
Transformer BERT、 MatMul、Softmax、LayerNorm 访存密集,注意力计算占比高 矩阵分块、Softmax融合、KV Cache管理
LSTM/RNN LSTM、GRU MatMul、Sigmoid、Tanh 串行依赖强,时间步间无法并行 时间步展开、门控融合、批处理

判断方法:模型结构识别与优化方向选择

当拿到一个模型时,通过以下特征快速判断其类型并确定优化方向:

  • 若模型包含大量Conv2D且输入为图像:优先优化卷积算子,检查是否可用Winograd或FFT加速。
  • 若模型包含Multi-Head Attention:优先优化MatMul和Softmax,检查KV Cache的访存效率。
  • 若模型包含循环结构:优先考虑时间步展开(Unrolling)和门控算子融合。

案例拆解:MobileNet在NPU上的深度优化

MobileNet使用深度可分离卷积(Depthwise Separable Convolution)替代标准卷积,大幅减少计算量。但在NPU上,深度卷积(Depthwise Conv)的算术强度极低,因为每个卷积核只处理一个输入通道,权重复用率低,导致访存受限。

优化方案:将深度卷积与后续的Pointwise卷积(1×1卷积)融合。在NPU上,深度卷积的输出可以直接留在片上缓存(SRAM)中,作为Pointwise卷积的输入,避免中间张量写回DRAM。实测表明,这种融合可将MobileNetV2在NPU上的推理延迟降低约35%。

常见误区与总结复盘

误区一:认为FLOPs越低模型越快。实际上,在NPU上,访存量往往比计算量更影响延迟。MobileNet的FLOPs远低于ResNet,但在某些NPU上推理速度反而更慢,原因就是深度卷积的访存效率太低。
误区二:忽略张量布局对性能的影响。同一个模型,仅将NCHW改为NHWC布局,在NPU上可能带来20%以上的性能差异。编译器无法自动修复所有布局问题,算子开发者需要主动适配。
误区三:盲目追求算子融合。算子融合减少了中间张量访存,但可能增加寄存器压力或导致指令缓存溢出。在融合前应评估融合后算子的指令长度和寄存器使用量。