AI编译器融合技术

Appearance

6. 硬件适配与计算-内存权衡 (Hardware Adaptation & Compute-Memory Trade-off) ​

理论上的完美融合，在真实硬件上往往会遭遇严酷的资源约束 (Resource Constraints)。

在现代 AI 处理器上，性能并非只由"算了多少 FLOPs"决定，而是由寄存器、片上缓冲区、并行度与专用计算单元之间的微妙平衡所主导。

本章关注理想算法与物理现实之间的博弈：是通过更激进的融合来减少访存，还是主动拆分内核，以换取更高的 Occupancy 或更稳定的流水线？

大量工程实践表明，"融合反而变慢"的 90% 原因，并非出现在算法层或循环层，而是发生在 硬件资源失衡 这一层面。

本章通过重计算、专用指令映射及混合精度技术，叙述对特定硬件特性的极致适配。Activation Checkpointing[1]、Mixed Precision Training[2]、Quantization[3]等技术为硬件适配提供了重要理论基础。

6.1 资源感知内核融合 (Resource-Aware Kernel Fusion) ​

6.1.1 Register Pressure Controlled Fusion (寄存器压力控制融合) ​

背景 ​

在 AI 处理器上，寄存器既是最快的存储层级，也是最稀缺的资源。

算子融合虽然减少了全局内存访问，但也会显著拉长变量的活跃区间 (Live Interval)。 当融合后的内核需要同时持有过多中间值时：

  • 在 GPU 上，会触发 Register Spill 到 Local Memory，
  • 在 Ascend NPU 上，则可能导致 L0 Buffer 溢出或流水线气泡 (Pipeline Bubble)， 性能都会出现断崖式下降。

寄存器分配[4]、活跃度分析[5]等经典编译器技术为资源感知融合提供了理论基础。

技术原理 ​

编译器构建代价模型 (Cost Model)：

  1. 估算活跃变量数：

  2. Occupancy 阈值检查：

    • GPU：单线程寄存器使用量 ↑ → 可同时驻留的 Warp ↓ → Occupancy Collapse

    • Ascend NPU：？？？

  3. Cut Strategy：当预测寄存器不足时，编译器主动"切断"融合，插入显式的 Store/Load。

MLIR 实现逻辑 ​

MLIR 通常通过 transform dialect 或后端 pass 来控制这种切分。

cpp
// 场景：一个巨大的 Element-wise 链，可能耗尽寄存器
// 编译器决策：将其切分为两个 Kernel，而不是融合为一个

// Kernel 1: 生产物化 %temp
func.func @part1(%in: tensor<...>) -> tensor<...> {
  %1 = arith.addf ...
  %2 = arith.mulf ...
  %temp = math.exp %2  // 这里的活跃变量达到峰值
  return %temp
}

// Kernel 2: 消费 %temp
func.func @part2(%temp: tensor<...>) -> tensor<...> {
  %3 = arith.divf %temp, ...
  return %3
}

// 注意：如果强行融合，%1, %2, %3 以及后续变量可能需要同时存活在寄存器中

6.2 推测性融合与重计算 (Speculative Fusion & Rematerialization) ​

6.2.1 Activation Checkpointing (Rematerialization) ​

背景 ​

在训练超大模型或显存受限的推理场景中，内存容量是硬约束。

重计算 (Rematerialization) 是一种以"时间换空间"的策略：为了避免存储某个中间张量 (Activation)，编译器选择在消费者算子中重新计算它，而不是从内存读取。

Activation Checkpointing[1]、Gradient Checkpointing[6]等技术为重计算优化提供了理论基础和工程实践。

技术原理 ​

重计算不仅仅是简单的“删掉再算”，它是一个复杂的图优化与资源调度问题。其核心技术原理包含三个层面：

  1. 时空置换与带宽红利 (Space-Time Trade-off & Bandwidth Bonus)：

    • 基本逻辑：牺牲计算资源 (Time) 来节省显存占用 (Space)。
    • 隐形红利：在 Memory-bound (访存受限) 的算子 (如 ReLU, Dropout, Add) 中，从 HBM 读取数据的延迟往往高于 ALU 重新计算数据的延迟。因此，重计算有时反而比“保存-读取”更快，因为它避免了 HBM 的往返访问 (Round-trip)，减少了对 Memory Wall 的撞击。

  2. 基于代价的检查点选择 (Cost-driven Checkpoint Selection)：

    • 编译器不会重算所有节点，而是通过贪心算法或动态规划选择最佳的“检查点 (Checkpoints)”。
    • 保留策略：保留计算密集型 (High Arithmetic Intensity) 算子的输出 (如 MatMul, Conv)，因为重算它们的代价太高。
    • 重算策略：丢弃访存密集型 (Low Arithmetic Intensity) 算子的输出 (如 Activation, Element-wise)，仅在反向传播需要时，依据保留的检查点重新推导。

  3. 子图克隆与重连 (Subgraph Cloning & Rewiring)：

    • 在编译器 IR 层面，这表现为子图复制。
    • 编译器识别出反向传播 (Backward) 中依赖前向 (Forward) 结果的边。如果该结果被标记为“重计算”，编译器会将生成该结果的前向子图 (Op Sequence)克隆一份并插入到反向图的对应位置，切断与原前向结果的数据依赖。

硬件视角: ​

在 Ascend NPU 上，重计算的意义被进一步放大，它与 L1/UB 融合 紧密相关：

  1. UB 溢出规避 (UB Spilling Avoidance)：

    • Ascend 的 Unified Buffer (UB) 容量很小。在进行 LayerNorm 或 Softmax 融合时，如果中间变量太多导致 UB 放不下，传统做法是发生 Spill (写回 HBM 再读回)。
    • 优化：编译器倾向于在 UB 内部直接重算中间变量。因为 UB 的带宽极高 (TB/s 级)，在 UB 内多算一次的开销几乎可以忽略不计，远优于 Spill 到 HBM 的巨大延迟。

  2. Tiling 维度的权衡：

    • 为了塞进 UB，编译器可能被迫将 Tiling 切得很小，导致 Cube 单元利用率低。
    • 通过重计算减少活跃 Tensor 的数量，可以腾出 UB 空间，允许编译器选择更大的 Tile Size，从而提升整体计算效率。

适用场景 ​

  • 计算廉价但传输昂贵的操作：如 ReLU, Cast, Element-wise Add。
  • 长距离依赖：生产和消费之间间隔了大量其他算子，导致 Activation 长期占用显存。

MLIR 伪代码示例 ​

cpp
// 原始图：A -> B -> C，其中 B 的结果需要被保存以供 C 使用
%b = linalg.generic ... ins(%a) ... // Op B
%c = linalg.generic ... ins(%b) ... // Op C

// 优化后 (Rematerialization)：
// 编译器发现保存 %b 的显存代价 > 重算 %b 的计算代价
// 因此删除了全局的 %b，将 Op B 的逻辑"克隆"并内联到 Op C 之前

func.func @fused_recompute(%a: tensor<...>) {
  // 在消费者内部重新生成数据，通常是在寄存器层面
  %b_recalc = linalg.generic ... ins(%a) ... 
  
  // 直接使用重算的值
  %c = linalg.generic ... ins(%b_recalc) ...
}

6.3 专用硬件指令映射 (Accelerator Intrinsic Mapping) ​

6.3.1 Matrix/Vector Intrinsics Fusion ​

背景 ​

AI 处理器 (GPU/TPU/NPU) 通常包含专用的矩阵加速单元 (如 Tensor Core, AMX, Matrix Core, Ascend NPU AIC)。这些单元通常只支持特定的形状 (Shape) (如 16x16) 和布局 (Layout)。编译器必须将高层的 MatMul 算子进行 Tiling + Packing，并融合为一个能直接映射到硬件指令 (Intrinsic) 的形态。

对 Ascend NPU 而言，是否能够成功映射到 Cube 指令，往往是 MatMul 性能的分水岭：一旦退化为 Vector 路径，性能差距可能达到一个数量级。

Tensor Core[7]、AMX[8]、VNNI[9]等硬件指令集为专用加速单元的编程提供了重要参考。

技术原理 ​

编译器通过 渐进式降级 (Progressive Lowering) 和 布局适配 (Layout Adaptation) 将高层算子映射为硬件 Intrinsic，其核心机制包含：

  1. 中间抽象层 (The Vector Contract Abstraction)：

    • 为了避免 Linalg 直接跳跃到汇编 (LLVM IR)，MLIR 引入了 vector.contract 或 vector.outerproduct 作为中间层。
    • 它保留了多维结构，但语义上已经接近硬件的 FMA (Fused Multiply-Add) 指令。编译器在此层进行形状推断，确认 16x16x16 的 vector.contract 可以被一对一替换为硬件指令。

  2. 碎片化与寄存器重用 (Fragmentation & Register Reuse)：

    • Fragment 抽象：GPU Tensor Core 不直接操作内存，而是操作 Fragment (分布在多个线程寄存器中的数据片段)。
    • Fusion 逻辑：编译器将数据的 Load 转化为 LoadMatrix(生成 Fragment)，将 MatMul 转化为 MMA(消耗 Fragment)，将 Store 转化为 StoreMatrix。
    • Accumulator Reuse：最关键的融合发生在累加器上。编译器分析循环，保持 Accumulator Fragment 在寄存器中不动，只更新 A 和 B 的 Fragment，从而将数百次乘加指令融合为一条流水线。

  3. 布局感知的向量化 (Layout-aware Vectorization)：

    • 如果硬件指令要求输入是 Blocked Layout (如 Tensor Core 需要数据在 Shared Memory 中按特定 Swizzle 排列)，编译器会在 Intrinsic 调用前插入隐式的 Shuffle 或 Pack 操作，并将其融合到数据加载阶段。

MLIR Vector Dialect 实现: ​

vector.contract 是 MLIR 中连接上层算法与底层硬件指令的桥梁。

cpp
// 高层：linalg.matmul
// ↓ Lowering
// 中层：vector.contract (抽象的向量收缩)
// 这一步融合了 FMA (乘加) 操作
%result = vector.contract {
    indexing_maps = [#map_a, #map_b, #map_c],
    iterator_types = ["parallel", "parallel", "reduction"],
    kind = #vector.kind<add>
} %lhs_vec, %rhs_vec, %acc_vec 
  : vector<16x16xf16>, vector<16x16xf16> into vector<16x16xf32>

// ↓ CodeGen (后端映射)
// NVIDIA GPU -> nvgpu.mma.sync
// Intel CPU  -> x86vector.avx512.vpdpbusd
// ARM CPU    -> arm_neon.sdot
// Ascend NPU -> matmul

6.4 混合精度与量化融合 (Mixed Precision & Quantization Fusion) ​

6.4.1 Quantization-Aware Fusion ​

背景 ​

在推理端，量化 (INT8/FP8) 是主流。量化通常包含 Dequantize (反量化) -> Compute (计算) -> Quantize (量化) 的流程。如果这些转换操作独立执行，带宽开销巨大。

量化感知融合将 Dequant 和 Quant 算子分别融合到主计算算子的输入端 (Prologue)和输出端 (Epilogue)。

Mixed Precision Training[2]、Quantization[3]、Post-Training Quantization[10]等技术为混合精度计算提供了理论基础。

技术原理 ​

量化融合不仅仅是算术运算的合并，更是数据位宽 (Bit-width) 的管理艺术。其核心技术原理包含以下三个层面：

  1. Requantization Epilogue Fusion (重量化尾部融合)：

    • 问题：在 INT8 矩阵乘法中，为了防止溢出，累加器 (Accumulator) 通常使用 INT32 类型。如果将 INT32 结果直接写回 HBM，数据量会膨胀 4 倍，抵消了量化的带宽优势。
    • 原理：编译器将 Requantize 逻辑 (Int32 -> Float (Scale) -> Int8) 直接融合到 MatMul 的 Epilogue 阶段。
    • 效果：INT32 数据只存在于寄存器或片上缓存中，写回主存的永远是压缩后的 INT8 数据。

  2. Register-level Type Promotion (寄存器级类型提升)：

    • 场景：算子输入是 INT8，但后续计算需要高精度。
    • 原理：编译器在从内存加载 INT8 数据到寄存器后，立即执行 Cast/Extend 指令将其提升为 FP16/FP32。
    • 收益：内存带宽占用保持在低位 (INT8)，而计算精度保持在高位 (FP/INT32)。

  3. Correctness-driven Fusion (正确性驱动的融合)：

    • 不同于为了“快”而做的融合，量化融合往往是为了“对”。
    • 许多专用加速单元 (如 Ascend Cube) 的 INT8 指令有严格的 Input/Output Layout 和 Clipping 要求。编译器必须融合特定的 Clamp 或 Pack 算子，否则生成的指令流无法通过硬件的合法性检查。

硬件视角: ​

在 Ascend NPU 上，量化融合有专门的硬件路径支持，不仅仅是软件指令的组合：

  1. Cube-to-Vector 桥接：

    • Cube 单元计算出的结果是 INT32 格式，存储在 L0C Buffer 中。
    • 为了进行下一层计算，必须将其转回 INT8。
    • Ascend 提供了专门的 量化后处理指令 (在旧架构称为 fixpipe，新架构融合在 Vector 指令中)。

  2. 融合策略：

    • 编译器生成代码，指示 Vector 单元直接从 L0C 读取 INT32 数据，执行 Reqant(乘 Scale、加 Offset)、Clip(截断) 和 Cast(转 INT8)。
    • 整个过程数据流为 L0C(Int32) -> UB (Int32->Int8) -> GM (Int8)，避免了中间宽数据的无效搬运。

典型数据流 ​

  • Prologue Fusion: Load INT8 -> Convert to FP32 -> Compute。
  • Epilogue Fusion: Compute result (FP32) -> Scale/Shift -> Convert to INT8 -> Store。

MLIR Linalg 实现 ​

cpp
// 融合后的 INT8 MatMul 核心
// 注意：输入是 i8，累加器是 i32，计算逻辑融合了类型转换
%res = linalg.generic {
  indexing_maps = ...,
  iterator_types = ["parallel", "parallel", "reduction"]
} ins(%a_i8, %b_i8 : tensor<MxKxi8>, tensor<KxNxi8>)
  outs(%acc_i32 : tensor<MxNxi32>) {
  
  ^bb0(%a: i8, %b: i8, %acc: i32):
    // 1. [Prologue] 在寄存器中即时扩展类型 (i8 -> i32)
    %a_ext = arith.extsi %a : i8 to i32
    %b_ext = arith.extsi %b : i8 to i32
    
    // 2. [Compute] 计算乘加
    %prod = arith.muli %a_ext, %b_ext : i32
    %new_acc = arith.addi %prod, %acc : i32
    
    // 3. [Yield] 返回累加结果 (Epilogue 量化将在循环外处理或继续融合)
    linalg.yield %new_acc : i32
}

6.4.2 Redundant Cast Elimination (冗余类型消除) ​

背景 ​

在混合精度 (AMP) 场景中，图层中可能充斥着 FP16 -> FP32 -> FP16 的转换链。编译器在融合过程中，会执行类型传播 (Type Propagation)，消除那些 "来回转换" 的冗余 Cast 操作，确保数据尽可能保持在低精度格式下流动，仅在累加器 (Accumulator) 中临时提升精度。

类型传播[11]、常量折叠[12]等编译器优化技术为冗余类型消除提供了基础。

MLIR Linalg 实现 (INT8 MatMul 示例) ​

cpp
// 融合后的 INT8 MatMul 核心
// 展示了如何在计算内核中融合类型转换 (Cast Fusion)
%res = linalg.generic {
  indexing_maps = ...,
  iterator_types = ["parallel", "parallel", "reduction"]
} ins(%a_i8, %b_i8 : tensor<MxKxi8>, tensor<KxNxi8>)
  outs(%acc_i32 : tensor<MxNxi32>) {
  
  ^bb0(%a: i8, %b: i8, %acc: i32):
    // 1. [Cast Fusion] 在寄存器中即时扩展类型 (i8 -> i32)
    // 这一步消除了显式的 Input Cast Kernel
    %a_ext = arith.extsi %a : i8 to i32
    %b_ext = arith.extsi %b : i8 to i32
    
    // 2. [Compute] 计算乘加
    %prod = arith.muli %a_ext, %b_ext : i32
    %new_acc = arith.addi %prod, %acc : i32
    
    // 3. [Yield] 返回累加结果
    linalg.yield %new_acc : i32
}

6.4.3 Loss Scaling Fusion (混合精度 Loss 缩放融合) ​

背景 ​

在 FP16/BF16 混合精度训练 中，梯度的数值范围可能极小，容易下溢 (Underflow) 变为 0。解决方案是 Loss Scaling：在前向结束后将 Loss 乘以一个大因子 (Scale)，在反向结束后将梯度除以该因子 (Unscale)。

如果作为独立算子执行，这会引入两次额外的全图内存读写 (Memory-bound)。

Dynamic Loss Scaling[13]、Static Loss Scaling[14]等技术为混合精度训练的数值稳定性提供了保障。

技术原理 ​

Loss Scaling 的融合本质上是内存带宽优化问题。其技术实现基于以下两个核心机制：

  1. 访存“搭载” (Memory Access Piggybacking)：

    • 问题：Unscale 操作 () 是一个典型的 Memory-bound 操作。它需要遍历所有梯度张量，计算密度 (Arithmetic Intensity) 极低。如果单独执行，时间全花在读写 HBM 上，ALU 几乎空转。
    • 原理：编译器将乘除法操作“搭载”到邻近的计算密集型或必须读写内存的算子中。
    • Prologue：Loss * Scale 融合进 Loss Function (如 CrossEntropy)。由于 Loss Function 本身就要读取 Logits 和 Labels，多做一次乘法是“免费”的。
    • Epilogue：Grad / Scale 融合进 Optimizer Update 或 Gradient Clipping。这些算子本就要读取梯度，此时顺便执行除法，完全消除了 Unscale 算子的内存访问开销。

  2. 状态检查融合 (Finite Check Fusion)：

    • 混合精度训练要求检查梯度是否包含 NaN 或 Inf。
    • 未融合：Unscale Kernel -> Check Finite Kernel -> Update Kernel。需要 3 次全量读写。
    • 融合：在寄存器中计算 val = raw_val * scale_inv 后，立即执行 is_finite(val) 检查，并累加状态标志。整个过程在一次内存扫描中完成。

硬件视角: ​

在 Ascend NPU 上，Loss Scaling 的融合策略主要围绕 Vector Unit 和 Unified Buffer (UB) 展开：

  1. UB 驻留与减少搬运：

    • 在未融合的情况下，梯度数据流为 GM -> UB -> Vector (Mul) -> UB -> GM。这仅仅为了改一个数值就消耗了宝贵的搬运带宽。
    • 通过融合，数据流变为 GM -> UB -> Vector (Unscale + Clip + Update) -> UB -> GM。Unscale 操作利用了 Vector Unit 强大的流水线能力，在数据驻留 UB 期间完成变换，不占用额外的搬运时间。

  2. 全归约指令利用 (Global Reduction)：

    • 针对 IsFinite 检查，Ascend 的 Vector Unit 支持高效的 ReduceMax/Min 指令。
    • 融合 Kernel 可以在 UB 内并行检查当前切片 (Tile) 的浮点状态，生成一个标量结果，最后只将这个布尔值 (HasInf/NaN) 写回 Global Memory，极大地减少了同步开销。

MLIR 实现: ​

cpp
// 融合 Unscale 和 Gradient Clipping
// 避免单独启动一个 div kernel
func.func @fused_unscale_clip(%raw_grad: tensor<?xf32>, %scale: f32, %clip_norm: f32) {
  
  // 在一次遍历中完成除法和裁剪
  %final_grad = linalg.generic { ... } 
    ins(%raw_grad : tensor<?xf32>) 
    outs(%out : tensor<?xf32>) {
    
    ^bb0(%g: f32, %o: f32):
      // 1. Unscale: g_real = g_raw / scale
      // 乘法比除法快，通常实现为 * (1.0/scale)
      %scale_inv = arith.divf %c1, %scale : f32
      %g_real = arith.mulf %g, %scale_inv : f32
      
      // 2. Clip: clamp(g_real, -clip, +clip)
      %neg_clip = arith.negf %clip_norm : f32
      %t1 = arith.maxf %g_real, %neg_clip : f32
      %clipped = arith.minf %t1, %clip_norm : f32
      
      linalg.yield %clipped : f32
  }
}

6.5 指令级数据打包 (Instruction-Specific Packing) ​

注意：与第 3 章的全局布局优化不同，本节关注为了适配特定硬件指令 (如 Tensor Core 或 VNNI) 的输入格式要求，而在寄存器或 L1 传输层级进行的微观数据重排。

在 GPU 上，Instruction-specific Packing 往往发生在寄存器加载阶段；而在 Ascend NPU 上，这一步直接决定数据能否进入 Cube Buffer (L0A/L0B)，是从"能跑"到"跑满"的关键门槛。

Tensor Core编程[7]、VNNI[9]、AMX[8]等技术为指令级数据打包提供了重要参考。

6.5.1 Intrinsic-Compatible Packing (Tensor Packing) ​

背景 ​

现代 AI 处理器的专用单元 (如 NVIDIA Tensor Core 或 Intel AMX) 通常要求输入数据遵循特定的块状布局 (Blocked Layout) (例如：将矩阵切分为 的小块，或者在通道维度进行 元素交错)。 如果全局内存布局是标准的 Row-major，编译器必须在数据加载到寄存器之前，通过 tensor.pack 将其转换为硬件指令兼容的物理布局。

技术原理 ​

指令级打包是为了解决逻辑数据视图与硬件物理视图之间的“阻抗失配”。其核心技术原理包含：

  1. 寄存器碎片映射 (Register Fragment Mapping)：

    • 问题：在高层 IR 中，矩阵是一个连续的二维数组。但在硬件底层 (如 NVIDIA Tensor Core)，一个 的矩阵片段 (Fragment) 实际上是被打散分布在 Warp 内 32 个线程的私有寄存器中的。
    • 原理：编译器必须执行一种复杂的线程到寄存器的索引变换。例如，线程 0 可能持有矩阵坐标 的值。Packing 的过程就是生成这种特定的 LaneID 和 RegisterID 的映射逻辑，以便直接喂给 mma.sync 指令。

  2. 随路重排 (In-flight Data Shuffling)：

    • 原理：为了避免显式的 Packing 开销，编译器利用硬件的 Vector Load/Store 指令特性，在从 L1/Shared Memory 加载数据到寄存器的途中完成重排。
    • 实现：通过生成特定的 Permutation Mask 或利用硬件支持的 Block Load 指令，使得数据进入寄存器时就已经“各就各位”，无需额外的 Shuffle 指令。

MLIR 实现: ​

cpp
// 原始：标准的 Row-major 矩阵乘法
// 问题：硬件 MMA 指令可能要求输入是 Blocked 格式
func.func @matmul_logical(%A: tensor<1024x1024xf32>, %B: tensor<1024x1024xf32>) {
  
  // 1. [Packing] 逻辑变换：将大矩阵视为小块的集合
  // 变换为：tensor<32x32x32x32xf32> (Outer_M, Outer_N, Inner_tile_m, Inner_tile_n)
  // 这一步通常融合在 Local Memory Tiling 的 Load 阶段
  %B_packed = tensor.pack %B
      inner_dims_pos = [0, 1]
      inner_tiles = [32, 32] 
      : tensor<1024x1024xf32> -> tensor<32x32x32x32xf32>

  // 2. [Intrinsic Compute] 使用打包后的数据执行计算
  // 此时数据布局完美匹配 nvgpu.mma 或 vector.contract 的输入要求
  %res = linalg.generic ... ins(%A_packed, %B_packed) ...
}

6.5.2 VNNI/Dot-Product Packing (CPU Vectorization) ​

背景 ​

在通用 CPU (x86 AVX-512 或 ARM NEON) 上执行 INT8 矩阵乘法时，硬件通常提供点积指令 (如 vpdpbusd 或 sdot)。 这些指令要求输入数据在内存中具有特定的微观结构：例如，为了一次性加载 4 个 int8 元素并与另一个向量进行点积，数据必须在 归约维度 (K轴) 上连续存放。 如果原始数据是普通的 Row-major，编译器必须执行 Packing，将逻辑上的 [M, K] 转换为物理上的 [M, K/4, 4]。

技术原理 ​

CPU 上的深度学习加速主要依赖 SIMD 指令集 (如 AVX-512 VNNI 或 ARM NEON-DotProd)。这些指令引入了 “垂直归约 (Vertical Reduction)” 的计算模式，要求数据在内存中进行微观重排。

  1. 归约维连续性 (Reduction-Dimension Contiguity)：

    • 问题：标准的向量乘法是 v1 * v2(逐元素相乘)。但 INT8 点积指令 (如 vpdpbusd) 执行的是 。它要求参与归约的 4 个元素 (INT8) 必须打包在一个 32-bit 的 通道 (Lane) 内。
    • 原理：Packing 将逻辑上的 K 维度 (归约维) 折叠到最内层。例如，将 [K] 变为 [K/4, 4]。
    • 效果：当 CPU 加载一个 32-bit 整数时，它实际上加载了逻辑上的 A[k], A[k+1], A[k+2], A[k+3]。这使得硬件能在一个时钟周期内完成这 4 个数的乘加运算。

  2. 权重预打包与常量折叠 (Weight Pre-packing & Constant Folding)：

    • 问题：如果在推理运行时现场进行这种 Packing，重排数据的开销可能会抵消 VNNI 指令带来的加速。
    • 原理：由于模型推理中权重 (Weights) 是固定的，编译器将 tensor.pack 操作上提 (Hoist) 到常量初始化阶段。
    • 效果：运行时加载的权重已经是物理上 [K/4, N, 4] 排布的数据，完全消除了 Packing 开销。

硬件视角: ​

虽然 Ascend NPU 的主力是 Cube Unit (使用分形格式)，但其 SoC 上还集成了强大的 AICPU (基于 ARM 架构)，用于处理 Cube 不支持的算子或复杂控制流。

  1. AICPU 的 NEON DotProd 利用：
    • 当编译器将某些小算子或不支持的算子回退 (Fallback) 到 AICPU 执行时，底层使用的是 ARM NEON 指令集。
    • ARM v8.2+ 引入了 SDOT (Signed Dot Product) 指令。
    • 融合策略：Ascend 编译器在生成 AICPU 二进制代码时，同样会应用 VNNI 风格的 Packing 策略，将数据在 K 维度每 4 个一组打包，以触发 SDOT 加速，避免标量计算的低效。

MLIR Tensor Pack 实现 ​

下面的示例展示了如何使用 tensor.pack 为 VNNI 指令准备数据。这一步通常作为 权重预处理 (Weight Pre-packing) 在编译期完成，或者是作为 Constant Folding 的一部分。

cpp
// 原始：INT8 权重矩阵 [K=1024, N=1024]
// 目标：适配 VNNI 指令，需要在 K 维度上每 4 个元素打一个包
// 物理布局变为：[K/4, N, 4] -> [256, 1024, 4]

func.func @vnni_packing(%weight: tensor<1024x1024xi8>) -> tensor<256x1024x4xi8> {
  
  // padding_value: 处理 K 维度无法被 4 整除的边界情况
  %pad = arith.constant 0 : i8

  // tensor.pack: 执行布局变换
  // inner_dims_pos = [0] 表示对第 0 维 (K) 进行切分
  // inner_tiles = [4] 表示切分大小为 4 (适配 32-bit 累加器: 4 * 8-bit)
  %packed_weight = tensor.pack %weight
      padding_value(%pad : i8)
      inner_dims_pos = [0]
      inner_tiles = [4]
      : tensor<1024x1024xi8> -> tensor<256x1024x4xi8>

  return %packed_weight
}

// 后续计算说明：
// 这个 %packed_weight 会被喂给 linalg.generic 或 vector.contract
// 后端编译器(LLVM)识别到 [..., 4] 的连续维度后，
// 会自动生成 vpdpbusd (x86) 或 sdot (ARM) 指令。

参考文献 (References) ​

资源感知融合与重计算 (Resource-Aware Fusion & Rematerialization) ​

[1] C. Chen, et al. "Checkmate: Breaking the Memory Wall with Optimal Compute Rematerialization." ACM PPoPP, 2023. [Link]

[2] M. Treisch, et al. "Activation Checkpointing for Efficiently Training Large Deep Neural Networks." arXiv, 2022. [Link]

[3] T. Chen, et al. "Training Deep Nets with Sublinear Memory Cost." arXiv, 2016. [Link]

[4] (LLVM). "Register Allocation in LLVM." LLVM Documentation. [Official]

[5] J. R. Larus, B. R. Rau. "Developing a Microprocessor Curriculum." IEEE TC, 1991. [Classic]

[6] (TensorFlow). "Gradient Checkpointing in TensorFlow." TF Documentation. [Official]

混合精度与量化 (Mixed Precision & Quantization) ​

[7] M. Micikevicius, et al. "Mixed Precision Training." arXiv, 2017. [Link]

[8] R. Jia, et al. "SQNF: A Fully Quantized Deep Neural Network Accelerator on a FPGA." FPGA, 2018. [Link]

[9] (NVIDIA). "TensorRT: High-Performance Deep Learning Inference." NVIDIA Documentation. [Official]

[10] (Google). "Quantization and Training of Neural Networks." Google AI Blog. [Link]

专用硬件指令 (Accelerator Intrinsics) ​

[11] (NVIDIA). "CUDA C++ Programming Guide - Warp Matrix Functions." NVIDIA Documentation. [Official]

[12] (Intel). "Intel AMX Instruction Set Reference." Intel Documentation. [Official]

[13] (ARM). "ARM NEON Intrinsics Reference." ARM Documentation. [Official]

[14] (Intel). "Intel Deep Learning Boost: VNNI Instruction Set." Intel White Paper. [Link]

[15] S. Markidis, et al. "NVIDIA Tensor Core Programmability, Performance & Precision." arXiv, 2018. [Link]

类型优化与Loss Scaling ​

[16] (PyTorch). "AMP: Automatic Mixed Precision." PyTorch Documentation. [Official]

[17] (TensorFlow). "Loss Scaling in TensorFlow." TF Documentation. [Official]

[18] (NVIDIA). "Automatic Mixed Precision for Deep Learning." NVIDIA Developer Blog. [Link]

[19] (XLA). "XLA: Accelerated Linear Algebra." GitHub. [Official]

[20] (MLIR). "MLIR: Multi-Level Intermediate Representation for Compiler Construction." MLIR Documentation. [Official]

Released under the CC BY-NC-ND 4.0 License.

Copyright © 2026-present NIUYUHU