AI编译器融合技术

Appearance

5. 并行性与分布式融合 (Parallelism & Distributed Fusion) ​

随着摩尔定律的演进，算力的增长主要源于 并行度 (Parallelism) 的提升。

本章关注如何将融合后的内核映射到硬件海量的执行单元上。优化的跨度极其宽广：从微观的指令级 (SIMD/Vector) 和线程级 (Thread/Warp)，扩展到宏观的任务级 (Streams) 以及跨芯片的分布式集群 (SPMD)。核心目标是解决"如何切分数据"与"如何分配任务"的问题，通过融合通信与计算，确保成千上万个核心能够协同工作且互不阻塞。

SIMD向量化[1,2]、GPU线程级并行[3,4]、Tensor Core[5,6]、SPMD分布式训练[7,8]等技术构成了现代AI编译器并行优化的基础。CUTLASS[9]、GSPMD[7]、Megatron-LM[8]等系统为这些技术的工程实现提供了重要参考。

5.1 指令与向量级并行 (Instruction/Vector Parallelism) ​

5.1.1 SIMD Vectorization Fusion ​

背景 ​

SIMD (Single Instruction Multiple Data) 允许单条指令处理多个数据元素。这是 CPU (AVX/SVE)、DSP (Hexagon) 以及部分 NPU 的基础并行方式。

编译器通过 Loop Vectorizer 将多个标量操作融合为向量指令，并利用 Predication/Masking 技术处理非对齐的边界条件。SLP Vectorizer[1]、Auto-Vectorization[2]等研究为SIMD优化奠定了理论基础。

技术原理 ​

  1. Loop Strip-mining (循环条带化)： 向量寄存器有固定宽度 (如 256-bit 或 2048-bit)。编译器必须将逻辑循环切分为步长等于向量宽度的"条带"。
  2. Masked Execution (掩码执行)： 处理循环尾部 (Loop Tail) 或条件分支时，为了避免标量回退，编译器生成 带掩码 (Predicate Mask) 的向量指令，仅对 Mask 为 1 的 Lane 进行计算和写回。

硬件视角: ​

Ascend 的 Vector Unit 非常宽 (通常 256 Bytes，即一次处理 128 个 FP16)。

  • 全掩码模式：Ascend 指令集强依赖 mask 参数 (如 vector_add (..., mask, ...))。编译器通过设置特殊的 Mask 寄存器 (64-bit)，精确控制哪部分数据参与计算。
  • 指令级 Repeat：如前所述，Ascend 编译器倾向于将 SIMD 进一步融合为 Repeat 指令，一条指令执行多达 255 个 Vector 宽度的计算，极大降低取指开销。

MLIR Vector Dialect ​

cpp
// 融合后的向量化操作
// vector<256xf32> 对应硬件的向量寄存器宽度 (如 AVX-512 或 SVE)
%v_a = vector.load %A[%i] : memref<?xf32>, vector<256xf32>
%v_b = vector.load %B[%i] : memref<?xf32>, vector<256xf32>

// 融合 FMA (Fused Multiply-Add) 指令
%v_sum = vector.fma %v_a, %v_b, %v_c : vector<256xf32>

// 融合 Masked Store (处理 Loop Tail 边界)
// %mask 是通过 vector.create_mask 生成的谓词掩码
vector.transfer_write %v_sum, %C[%i], %mask : vector<256xf32>, memref<?xf32>

5.2 线程级并行融合 (Thread-level Parallelism) ​

5.2.1 Workgroup & Subgroup Fusion ​

背景 ​

在 SIMT 架构 (如 GPU) 中，线程被组织为层级结构：Workgroup (对应 CUDA Block) 和 Subgroup (对应 CUDA Warp 或 AMD Wavefront)。

Subgroup-level Fusion 利用硬件提供的 Shuffle/Permute 指令，使同一 Subgroup 内的线程能直接交换寄存器数据，无需经过 Shared Memory，从而实现极低延迟的归约 (Reduction) 融合。CUDA Warp Shuffle[3]、AMD Wavefront优化[4]等技术为GPU线程级并行提供了基础。

技术原理 ​

  1. 资源划分 (Resource Partitioning)： 编译器根据计算图的并行度，将任务划分为多个 Workgroup (Grid)。每个 Workgroup 内部共享 L1/Shared Memory。
  2. Barrier Optimization (屏障优化)： 在 Subgroup (如 Warp/Wavefront) 内部，编译器利用硬件隐式同步特性消除显式 Barrier；在 Workgroup 内部，编译器尝试移动 Barrier 位置以最大化指令流水线重叠。

硬件视角: ​

Ascend 没有 CUDA 那样显式的 "Warp" 概念。

  • Block = AI Core：MLIR 的 Workgroup 通常直接映射为一个 AI Core (Block) 任务。
  • Intra-Core Parallelism (核内并行)：Ascend 的"Subgroup 融合"实际上体现为 Cube、Vector、MTE 三条流水线的并行。
    • 融合策略：编译器通过指令重排 (Instruction Reordering)，让 MTE 搬运数据、Vector 处理数据、Cube 计算矩阵在同一个时间窗口内并发运行。同步通过 Set/Wait Event 指令 (而非简单的 __syncthreads) 实现。

MLIR GPU/Vector Dialect 实现 ​

cpp
// Subgroup-level Reduction Fusion
// 场景：融合 加法 和 线程间通信 (AllReduce within a subgroup)
%val = ... : f32

// 使用通用的 shuffle 指令在 lanes 之间交换数据
// 下层会映射为 NVVM shfl (NVIDIA) 或 ROCDL ds_bpermute (AMD)
%shuffled = gpu.shuffle xor %val, %offset, %width : f32

// 融合计算
%reduced = arith.addf %val, %shuffled : f32

5.2.2 Cooperative Matrix Fusion (MMA) ​

背景 ​

现代 AI 处理器通常配备专用的矩阵加速单元 (如 NVIDIA Tensor Core, Google TPU MXU, Intel AMX, ARM SME, Ascend NPU AIC)。

为了利用这些单元，编译器必须执行 MMA Fusion，将标准的矩阵乘法循环重构为 协作式 (Cooperative) 指令，让一组线程 (Subgroup) 协同完成一个小块矩阵的加载与计算。

NVIDIA Tensor Core[5]、CUTLASS[6]等系统为矩阵加速单元的编程模型提供了重要参考，成为GPU矩阵计算优化的标准实践。

技术原理 ​

  1. Intrinsic Mapping (内建指令映射)： 编译器将高层的 linalg.matmul 识别为特定硬件的协作矩阵指令 (如 mma.sync 或 wmma)。
  2. Layout Conformation (布局适配)： 协作指令通常要求输入数据满足特定的布局 (如 Fragment Layout)。编译器在 Lowering 阶段插入隐式的 Pack/Unpack 或利用寄存器重用策略。

硬件视角: ​

Ascend 的核心就是 Cube Unit (矩阵加速器)。

  • Fractal Computing (分形计算)：Cube 单元只能计算 的分形块。
  • L0 Buffer Fusion：融合的关键在于L0A/L0B/L0C 缓存的管理。
    • 编译器生成代码，指示 MTE 将数据直接搬运到 L0 缓存 (并自动完成分形转换)。
    • Cube 指令连续发射，直接复用 L0C 中的累加结果 (Accumulator)，避免写回 UB 或 GM。这就是 Ascend 上最高效的 MMA 融合。

MLIR Vector Dialect 实现 ​

cpp
// 使用 vector.contract 或 nvgpu/amx dialect 描述协作计算
// 这里的语义是：一组线程共同持有一个 vector 片段 (Fragment)
%res = vector.contract {
    indexing_maps = [#map_a, #map_b, #map_c],
    iterator_types = ["parallel", "parallel", "reduction"],
    kind = #vector.kind<add>
} %a_vec, %b_vec, %acc_vec : vector<4x8xf16>, vector<8x4xf16> into vector<4x4xf32>

5.3 分布式与张量并行 (Distributed & Tensor Parallelism) ​

5.3.1 SPMD Sharding Propagation (SPMD 切分传播) ​

背景 ​

在分布式训练 (如 Megatron-LM/GSPMD) 中，大张量被 切分 (Sharded) 分布在设备集群上。

切分感知融合 (Sharding-aware Fusion) 的核心逻辑是：如果相邻算子的 分布式布局 (Mesh Layout) 一致，则它们可以在本地融合执行，完全消除中间的通信 (Resharding)。

GSPMD[7]、Megatron-LM[8]等系统为分布式训练的切分策略和通信优化提供了理论基础和工程实践。OneFlow[10]等框架进一步推动了分布式编译技术的发展。

技术原理 ​

  1. Slice Propagation (切片传播)： 编译器根据设备网格 (Mesh) 和用户指定的 Sharding 策略 (如 Shard ([0], [1]))，推导图中每个 Tensor 的分布式布局。如果生产者和消费者的切分方式一致 (Compatible)，则它们之间的边被标记为 Local，无需通信。
  2. Communication Insertion (通信插入)： 当布局不一致时 (如 Row-parallel 转 Col-parallel)，编译器插入 AllGather 或 AllToAll 算子。

硬件视角: ​

  • HCCS (Huawei Cache Coherent System)：Ascend 芯片间的高速互连。
  • Alignment Constraints (对齐约束)：在切分 Tensor 时，Ascend 编译器必须保证切分后的 Local Shape 依然满足 32-byte 对齐 或 C0=16 对齐。
    • 融合策略：如果切分导致数据不对齐 (例如切分 Channel=32 为两份 16)，编译器会接受；但如果切分 Channel=16 为 8，编译器会拒绝该 Sharding 策略，或者强制插入 Padding，因为不对齐会导致 Cube 单元无法计算 Local Slice。

MLIR Mesh Dialect 示例 ​

cpp
// 定义通用的设备拓扑：2x4 的处理器集群
%mesh = mesh.cluster @device_mesh(rank = 2, dim_sizes = [2, 4])

// 切分感知融合：
// MatMul 和 Elementwise 算子共享相同的 Sharding [[Link]]
// 编译器不仅融合了计算，还消除了它们之间的网络通信
func.func @spmd_fused(%A: tensor<1024x1024xf32, #mesh.shard<@device_mesh, [[Link]]>>,
                      %B: tensor<1024x1024xf32>) {
  // 1. 本地计算 MatMul (Local Compute)
  %local_matmul = linalg.matmul ... 

  // 2. 本地融合 Elementwise (Zero-cost fusion)
  // 因为布局一致，无需 AllGather，直接在切片上计算
  %local_result = arith.maxf %local_matmul, %c0 ...
  
  // 3. 仅在必要时插入集合通信
  mesh.all_reduce %local_result on @device_mesh ...
}

5.3.2 Collective Operation Fusion (通信算子融合) ​

背景 ​

在分布式数据并行训练中，模型往往包含数千个参数张量。如果对每个张量单独执行 AllReduce，通信握手 (Handshake) 和 内核启动 (Kernel Launch) 的延迟将远超实际的数据传输时间。

分桶 (Bucketing) 融合技术：编译器将多个逻辑上独立的小张量通信，通过 内存拷贝 (Memcpy) 打包到一个大的连续缓冲区 (Buffer) 中，执行单次大通信，然后再解包。这能显著提升带宽利用率。

Horovod[11]、NCCL[12]、HCCL等通信库为集合通信优化提供了高效实现。Gradient Bucketing[13]等技术成为分布式训练的标准实践。

技术原理 ​

  1. Buffer Coalescing (缓冲区合并)： 小包通信受限于 延迟 (Latency) 而非带宽。编译器或运行时将多个小的通信 Tensor (如不同层的梯度) 通过 Memcpy 拼接到一个大的连续 Bucket 中。
  2. Deterministic Scheduling (确定性调度)： 为了防止死锁，编译器必须保证所有设备上执行 AllReduce (Bucket) 的顺序严格一致。

硬件视角: ​

  • Task Launch Overhead：在 Ascend 上，启动一个 HCCL 任务 (通信任务) 需要 CPU 下发描述符，开销不可忽视。
  • HCCL Fusion：
    • Ascend 编译器计算最优的 Bucket Size (例如 32MB)。
    • 利用专门的 DMA 引擎 在后台进行 Bucket 的打包 (Packing)。
    • 由于 HCCS 带宽极高 (通常 30GB/s+)，打包 (Memory Copy) 往往成为瓶颈。因此融合策略会尝试使用 Vector Unit 协助打包，或者在计算结束时直接写入 Bucket 预留的地址 (Zero-copy Bucketing)。

逻辑示意 ​

cpp
// 原始：
AllReduce(Grad_Layer1) -> Wait -> AllReduce(Grad_Layer2) -> Wait

// 融合后 (Bucketing)：
Buffer = Concat(Grad_Layer1, Grad_Layer2)
AllReduce(Buffer)  <-- 融合为单次通信内核调用
Grad_Layer1, Grad_Layer2 = Split(Buffer)

MLIR 伪代码示例 ​

这段代码展示了编译器如何将两个独立梯度的 AllReduce 融合为一次通信：

cpp
func.func @bucket_allreduce(%grad_a: tensor<1024xf32>, %grad_b: tensor<2048xf32>) 
    -> (tensor<1024xf32>, tensor<2048xf32>) {
  
  // 1. [Packing] 将分散的小张量平铺并拼接 (Concat) 到一个大 Buffer
  // 这一步通常在本地内存通过 memref.copy 完成
  %packed_buffer = tensor.concat %grad_a, %grad_b dim(0) 
                 : tensor<1024xf32>, tensor<2048xf32> -> tensor<3072xf32>

  // 2. [Fused Communication] 执行单次集合通信
  // 大消息体能更好地利用互连网络带宽 (NVLink/Infiniband/HCCS)
  %reduced_buffer = mesh.all_reduce %packed_buffer on @device_mesh reduction("sum")
                  : tensor<3072xf32>

  // 3. [Unpacking] 将结果切片 (Slice) 回原始形状
  %res_a = tensor.extract_slice %reduced_buffer[0] [1024] [1] 
         : tensor<3072xf32> to tensor<1024xf32>
  %res_b = tensor.extract_slice %reduced_buffer[1024] [2048] [1] 
         : tensor<3072xf32> to tensor<2048xf32>

  return %res_a, %res_b
}

5.3.3 Communication-Computation Overlap (通信-计算重叠融合) ​

背景 ​

将通信算子与无依赖的计算算子进行流水线融合。通过将大算子切分为微操作 (Micro-ops)，在执行计算流 (Compute Stream) 的同时，并行触发通信流 (Comm Stream)。

通信-计算重叠是分布式系统优化的核心技术[14,15]，PipeDream[14]、1F1B调度[15]等研究为流水线并行提供了理论基础。

技术原理 ​

通信-计算重叠的核心在于打破"计算完成才能传输"的串行依赖，利用硬件的独立部件并行工作。其实现依赖两个关键机制：

  1. 多流流水线 (Multi-Stream Pipelining)：

    • 将硬件指令队列划分为 Compute Stream (负责矩阵乘、向量计算) 和 Communication Stream (负责 Send/Recv/AllReduce)。
    • 两个 Stream 只有在显式的 同步点 (Synchronization Point) 才会互相等待，其余时间并行执行。

  2. 微批次与双缓冲 (Micro-batching & Double Buffering)：

    • 为了实现重叠，数据必须切分。当计算单元正在处理第 个数据块 (Micro-batch) 时，通信单元正在传输第 个数据块的结果。
    • 这要求内存分配上采用 Ping-Pong Buffer，防止读写冲突 (Race Condition)。

硬件视角: ​

在 Ascend 达芬奇架构中，重叠的实现高度依赖于硬件调度器 Task Scheduler (TS) 和互连架构 HCCS：

  1. 物理独立的执行引擎：

    • AICore：执行 Cube 和 Vector 指令。
    • HCCL Engine：独立的通信控制器，直接通过 HCCS 链路搬运 HBM 中的数据，不占用 AICore 的算力。
    • 原理：编译器将计算算子下发到 Stream 0，将 HCCL 算子下发到 Stream 1。TS (Task Scheduler) 会同时从两个队列中提取任务并分发给对应的硬件引擎，实现真正的物理并行。

  2. 硬件级事件同步 (Hardware Event Synchronization)：

    • 为了保证数据一致性 (例如：必须等计算写完 HBM，通信才能读)，编译器插入 Event Record 和 Stream Wait 指令。
    • Ascend 特性：这种等待是硬件级的。TS 会挂起依赖流的执行，直到 Event 信号到达，而不会阻塞 Host CPU 的下发线程，也不会阻塞其他无依赖的 Stream。

伪代码逻辑 ​

cpp
// 融合优化(Pipelined Schedule)：
// 编译器自动插入 Event Record/Wait 实现流同步
Stream_Compute:  [Comp A0]  [Comp A1]  [Comp A2]
Stream_Comm:                [Send A0]  [Send A1]  [Send A2]
// 此时 Comp A1 与 Send A0 实现时间轴上的融合(重叠)

MLIR 实现: ​

在 MLIR 中，通过 async.execute 创建的 Token 依赖关系，在 Ascend 后端需要 Lowering 为 Stream 和 Event 操作。

cpp
// 场景：流水线并行中的 1F1B 模式
// 目标：在计算 MB[i] 的同时，发送 MB[i-1]

func.func @ascend_overlap(%input_i: tensor<...>, %result_i_minus_1: tensor<...>) {
  
  // Stream 0: Compute (映射为 AICore 任务)
  %compute_token = async.execute {
    // 这是一个计算密集型算子
    %res_i = linalg.matmul ins(%input_i, ...) ...
    
    // 记录 Event，表示计算完成，数据已准备好
    // (隐式: gpu.record_event %evt_compute)
    async.yield %res_i : tensor<...>
  }

  // Stream 1: Communication (映射为 HCCL 任务)
  // 必须等待上一轮的 buffer 可用，但不阻塞当前轮的计算
  %comm_token = async.execute {
    // 这是一个通信密集型算子
    // 它在 HCCS 链路上运行，与 AICore 并行
    mesh.send %result_i_minus_1 to @next_stage ...
    
    async.yield
  }

  // 屏障：确保本轮所有操作提交完成
  async.await %compute_token, %comm_token
}

5.3.4 Pipeline P2P Fusion (流水线点对点融合) ​

背景 ​

在超大模型训练 (如 GPT-4, DeepSeek) 中，通常采用 Pipeline Parallelism (PP) 将模型的不同层分布在不同设备上。PP 的核心挑战是 Bubble (流水线气泡)。

为了掩盖跨设备的 P2P 通信 (Send/Recv) 延迟，编译器需要实现 1F1B (One-Forward-One-Backward) 调度的自动化融合：即在执行当前 Micro-Batch 计算的同时，异步发送/接收上一个 Micro-Batch 的边界数据。

GPipe[16]、PipeDream[14]、PipeCleaner[17]等系统为流水线并行提供了重要参考。DeepSpeed[18]等框架将1F1B调度应用于大规模语言模型训练。

技术原理 ​

计算-通信交错 (Inter-op Overlap)： 不同于 Tensor Parallel 的层内通信，PP 的通信发生在层间。编译器将 Send/Recv 指令下沉到计算图中，并使用异步原语将其与无依赖的计算任务 (通常是不同 Micro-Batch 的计算) 并行化。

MLIR 示例: ​

cpp
// 模拟 Pipeline Stage N 的 1F1B 调度融合
// 同时处理：计算 MicroBatch[i] 和 发送 MicroBatch[i-1] 的结果

func.func @pipeline_stage_fused(%mb_curr: tensor<...>, %mb_prev_result: tensor<...>) {
  
  // Stream 1: 计算当前 Micro-Batch (计算密集)
  %compute_token = async.execute {
    // 前向传播计算
    %res = call @forward_layers(%mb_curr)
    async.yield %res : tensor<...>
  }

  // Stream 2: 发送上一个 Micro-Batch 的结果给 Stage N+1 (网络密集)
  // 这个操作与 Stream 1 并行执行，掩盖了 P2P 延迟
  %comm_token = async.execute {
    // P2P 发送指令
    mesh.send %mb_prev_result to @stage_next_mesh ...
    async.yield
  }

  // 同步点：等待两者完成，准备进入下一个 Step
  async.await %compute_token, %comm_token
  return
}

5.4 任务级与异构并行 (Task & Heterogeneous Parallelism) ​

5.4.1 Async & Multi-stream Fusion ​

背景 ​

在异构系统 (Host CPU + Device Accelerator) 中，利用 事件 (Event) 和 流 (Stream/Queue) 机制，将 CPU 逻辑、DMA 数据搬运和 Device 计算融合在同一个时间窗口内并发执行。

CUDA Streams[19]、异构计算[20]等为任务级并行提供了编程模型基础。现代AI加速器（如Ascend NPU）通过硬件任务调度器实现了多流并行[21]。

技术原理 ​

多流融合的核心是将无数据依赖的子图 (Subgraphs) 调度到独立的执行队列中，利用硬件的资源冗余来实现任务级并行。

  1. DAG 分区与流映射 (DAG Partitioning & Stream Mapping)：

    • 编译器对计算图进行依赖分析 (Reachability Analysis)。如果图在某点分叉为两个互不依赖的分支 (Branch A 和 Branch B)，编译器将它们标记为可并行。
    • 流分配：Branch A 分配给 Stream 0，Branch B 分配给 Stream 1。
    • 同步插入：在两个分支汇合 (Join) 的地方插入 StreamWaitEvent，确保结果正确性。

  2. 异构引擎并行 (Heterogeneous Engine Parallelism)：

    • 这是 AI 加速器区别于 CPU 的关键。硬件内部包含多种专用引擎。
    • DMA 引擎：负责搬运。
    • Compute 引擎：负责计算。
    • Scalar/Host 引擎：负责复杂逻辑控制。
    • 原理：编译器将不同类型的算子 (Copy vs MatMul vs Shape Inference) 下发到对应的流，使得 DMA 在搬运数据的同时，Compute 引擎在计算，Scalar 引擎在处理动态 Shape 逻辑。

硬件视角: ​

在 Ascend 达芬奇架构中，多流融合的实现高度依赖于 Task Scheduler (TS) 和 AICPU 的协同：

  1. TS 硬件分发 (Task Scheduler Dispatch)：

    • Ascend 芯片内置了一个硬件级的 TS (Task Scheduler) 模块。
    • Host CPU 只需将任务描述符 (Task Descriptor) 推送到内存中的 SQ (Submission Queue)。
    • 并行分发：TS 会自动从不同的 Queue 中抓取任务，如果任务依赖满足 (Event 信号到达)，TS 会将 AICore 任务发给 Cube/Vector，将系统任务发给 AICPU，将拷贝任务发给 DMA。这种分发完全由芯片硬件完成，无 Host CPU 负载。

  2. AICore 与 AICPU 的异步并行：

    • AICore：擅长张量计算。
    • AICPU：擅长标量逻辑、OS 交互、打印日志等。
    • 融合策略：对于无法被 AICore 融合的复杂算子 (如 Unique, Where, TopK 的某些变体)，编译器将其调度到 AICPU 上异步执行。与此同时，AICore 继续执行后续无依赖的矩阵计算。两者通过 Stream/Event 机制同步。

MLIR 实现: ​

MLIR 的 async dialect 提供了完美的抽象。在 Ascend 后端，async.execute 需要映射为不同的 Stream ID。

cpp
// 场景：Inception 模块的多分支并行 + 异构计算
func.func @ascend_multi_stream(%input: tensor<...>) {
  
  // Stream 0: 密集计算分支 (映射为 AICore Task)
  %token0 = async.execute {
    %conv = linalg.conv_2d ... ins(%input) ...
    async.yield %conv : tensor<...>
  }

  // Stream 1: 复杂逻辑/预处理分支 (映射为 AICPU Task)
  %token1 = async.execute {
    // 假设这是一个 AICore 不支持的复杂 CPU 算子
    %processed = "tf.Unique"(%input) ... 
    async.yield %processed : tensor<...>
  }

  // Stream 2: 独立的数据搬运 (映射为 DMA Task)
  %token2 = async.execute {
    %copy = memref.copy %input, %host_buffer ...
    async.yield
  }

  // 融合点：TS 硬件调度器会等待所有 Stream 完成
  async.await %token0, %token1, %token2
  
  // 后续处理...
}

5.4.2 Host-Device Prefetch Fusion (预取融合) ​

背景 ​

在推荐系统 (Embedding Lookup) 或图神经网络等场景中，Host (CPU) 到 Device (处理器) 的数据搬运往往是瓶颈。

如果采用简单的串行模式 (Copy -> Compute)，处理器会频繁处于 饥饿 (Starvation) 状态。

编译器通过 软件流水线 (Software Pipelining) 技术，将"当前批次的计算"与"下一批次的搬运"融合在同一个时间窗口内，实现 Host 与 Device 的全并行。

Double Buffering[22]、Prefetching[23]等技术为Host-Device重叠提供了优化基础。推荐系统[24]、图神经网络[25]等场景广泛应用了这些技术。

技术原理 ​

编译器进行软件流水线化 (Software Pipelining)：

  1. Prologue: 启动 Batch 0 的拷贝。
  2. Steady State (Kernel):
    • 并行执行：Compute (Batch i) && Copy (Batch i+1)
  3. Epilogue: 计算最后的 Batch。

MLIR 伪代码示例 ​

这里展示了一个典型的 Double Buffering (双缓冲) 预取流水线：

cpp
// 预先分配两个 Device Buffer 用于乒乓操作
%buf0 = memref.alloc() : memref<...>
%buf1 = memref.alloc() : memref<...>

// [Prologue] 预取第 0 个 Batch 的数据
%token_init = gpu.memcpy async %buf0, %host_data[0] : memref<...>, memref<...>

// 主循环：处理 Batch i，同时预取 Batch i+1
// iter_args 携带当前的 token 和 buffer 索引
scf.for %i = 0 to %num_batches step 1 
    iter_args(%prev_token = %token_init, %curr_buf = %buf0, %next_buf = %buf1) {
  
  // 1. [Wait] 等待当前 Buffer 的数据搬运完成
  // 这里的等待时间通常被上一轮的计算掩盖了
  async.await %prev_token

  // 2. [Prefetch] 立即启动下一 Batch 的搬运 (i+1) 到 spare buffer
  // 这是一个异步操作，CPU 立即返回，DMA 开始工作
  %next_token = gpu.memcpy async %next_buf, %host_data[%i + 1] 
  
  // 3. [Compute] 在当前 Buffer 上全速计算
  // 此时 Device 在计算 %curr_buf，DMA 在搬运 %next_buf
  // 实现了 Host-Device Overlap
  linalg.generic ... ins(%curr_buf) ...

  // 交换 Buffer 和 Token 进入下一次迭代
  scf.yield %next_token, %next_buf, %curr_buf
}

参考文献 (References) ​

SIMD向量化 (SIMD Vectorization) ​

[1] L. Larsen, E. Amarasinghe. "Superword Level Parallelism and Floating Point SLP Vectorization." ACM CGO, 2000. [Classic]

[2] A. Cohen, S. Girbal, et al. "A Graph-Based Framework for the Automatic Vectorization of Parallel Computations." IEEE TCAD, 2006. [Link]

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

[4] (AMD). "ROCm Optimization Guide - Wavefront Operations." AMD Documentation. [Official]

GPU线程级并行与Tensor Core (GPU Parallelism & Tensor Core) ​

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

[6] (NVIDIA). "CUTLASS: CUDA Templates for Linear Algebra Subroutines." GitHub. [Official]

[7] J. Yan, et al. "Demystifying Tensor Cores to Optimize Half-Precision Matrix Multiplication." IEEE IPDPS, 2020. [Link]

[8] (NVIDIA). "CUDA C++ Programming Guide - Warp Matrix Multiply Accumulate." NVIDIA Documentation. [Official]

SPMD与分布式训练 (SPMD & Distributed Training) ​

[9] Google. "GSPMD: General and Scalable Parallelization for ML Computation Graphs." MLSys, 2022. [Link]

[10] M. Shoeybi, et al. "Megatron-LM: Training Multi-Billion Parameter Language Models." arXiv, 2019. [Link]

[11] (OneFlow). "OneFlow: A Distributed Deep Learning Framework." GitHub. [Link]

通信融合与优化 (Communication Fusion & Optimization) ​

[12] A. Sergeev, M. Del Balso. "Horovod: Fast and Easy Distributed Deep Learning in TensorFlow." arXiv, 2018. [Link]

[13] (NVIDIA). "NCCL: NVIDIA Collective Communications Library." GitHub. [Official]

[14] (DeepSpeed). "Gradient Bucketing in DeepSpeed." Microsoft Research. [[Link]](https://www.microsoft.com/en-us/research/blog/deepspeed/

[15] Y. Peng, et al. "Optimizing Distributed Training with Communication-Computation Overlap." MLSys, 2023. [Link]

流水线并行与异步执行 (Pipeline Parallelism & Async Execution) ​

[16] Y. Huang, et al. "GPipe: Efficient Training of Giant Neural Networks." NeurIPS, 2019. [Link]

[17] D. Narayanan, et al. "PipeDream: Pipeline Parallelism for Large-Scale Model Training." MLSys, 2021. [Link]

[18] L. Zong, et al. "PipeCleaner: Optimizing Pipeline Parallelism by Fusing Large and Small Batches." OSDI, 2022. [Link]

[19] M. Rajbhandari, et al. "ZeRO: Memory Optimizations for Large-Scale Deep Learning." SC, 2020. [Link]

[20] (NVIDIA). "CUDA C++ Programming Guide - Streams and Events." NVIDIA Documentation. [Official]

异构计算与任务调度 (Heterogeneous Computing & Task Scheduling) ​

[21] (AMD). "HIP Programming Guide - Streams and Events." ROCm Documentation. [Official]

[22] J. Sanders, E. Kandrot. "CUDA by Example: An Introduction to General-Purpose GPU Programming." Addison-Wesley, 2010. [Book]

[23] (Huawei). "Ascend C Programming Guide - Task Scheduler." Huawei Developer. [Official]

[24] W. Wang, et al. "Deep Gradient Compression: Reducing the Communication Bandwidth for Distributed Training." ICLR, 2018. [Link]

[25] T. Ben-Nun, T. Hoefler. "Demystifying Parallel and Distributed Deep Learning: An In-Depth Concurrency Analysis." ACM Computing Surveys, 2019. [Link]

Double Buffering与Prefetching ​

[26] J. Demmel, et al. "Communication-Optimal Parallel Recursive Prefix Summation." SPAA, 2008. [Link]

[27] (NVIDIA). "CUDA C Best Practices Guide - Device Memory Overlap." NVIDIA Documentation. [Official]

[28] S. Chetlur, et al. "cuBLAS: A Library for GPU-Accelerated Deep Learning." arXiv, 2014. [Link]

推荐系统与图神经网络优化 ​

[29] H. Tang, et al. "Recommender Systems with Deep Learning: A Survey." IJCAI, 2019. [Link]

[30] J. You, et al. "GraphBERT: Graph Representation Learning with Attention." ICLR, 2020. [Link]

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

Copyright © 2026-present NIUYUHU