面向异构硬件的代码生成

前面两篇讲了 Tensor/Linalg（计算抽象层）和 Affine/SCF（循环与调度层）。这一篇讲这些循环和计算最终怎么映射到 GPU——MLIR 的 GPU 方言、SPIR-V 出口和内存层次抽象。

一、GPU 方言的设计哲学

MLIR 的 GPU 方言不同于 CUDA C++ 的编程模型。它的设计原则是：用 MLIR 的方言体系表示 GPU 编程的每个概念，允许编译器在生成最终代码前对 GPU 层次的 IR 做分析、变换和优化。

与传统 GPU 编译的区别：

GPGPU 编程回顾：SIMT 的两个层次

GPU 编程的核心理念是将计算映射到两级并行组织：线程块（thread blocks）组成网格（grid），块内线程（threads）以 warp 为单位同步执行。每个线程通过内置变量（如 CUDA 的 blockIdx、threadIdx）获取其在网格中的位置索引——这是将数据并行分解到硬件执行单元的关键机制。

MLIR 的 GPU 方言将这些概念直接编码为 IR 中的操作和属性，使得编译器可以看到每一层的并行和内存语义。

二、GPU 方言的核心操作

2.1 gpu.launch：并行启动

gpu.launch
    blocks(%bx, %by, %bz) in (%grid_x = %c256, %grid_y = %c1, %grid_z = %c1)
    threads(%tx, %ty, %tz) in (%block_x = %c256, %block_y = %c1, %block_z = %c1)
{
  // 在此 Region 中，%bx/%by/%bz 是 block index
  // %tx/%ty/%tz 是 thread index
  gpu.printf "block (%d,%d) thread (%d,%d)\n" %bx, %by, %tx, %ty
  gpu.terminator
}

gpu.launch 的 Region 执行在 GPU 上——每个 block 中的每个 thread 都会执行一次 Region body。这与 CUDA 的 <<<grid, block>>> 启动语法等价，但显式地将索引作为 SSA 值传递。

2.2 内存操作

// 在 shared memory 中分配：
// memref<256xf32, 3> 的地址空间 3 = GPU workgroup memory
%shared = memref.alloc() : memref<256xf32, 3>

// 从 global 内存拷贝到 shared 内存
gpu.memcpy %shared, %global_tensor : memref<256xf32, 3>, memref<256xf32>

2.3 同步与原子操作

// block 内所有 thread 的同步屏障
gpu.barrier

// 原子操作
%old = gpu.atomic_rmw add %ptr, %value : memref<?xf32>

三、内存层次映射

MLIR 的内存地址空间建模来自 GPU 的实际硬件结构：

具体数字取决于 GPU 架构与编译器后端，上表仅作相对层次示意，不作 benchmark 结论。

在 MLIR 中，通过 MemRef 的 memory space 标注来区分：

// global memory: 默认地址空间 0
%global = memref.alloc() : memref<1024xf32>
// workgroup (shared) memory: 地址空间 3
%shared = memref.alloc() : memref<256xf32, 3>

// 从 global 加载到 private（寄存器）
%val = memref.load %global[%idx] : memref<1024xf32>

// 向 shared 写入
memref.store %val, %shared[%local_idx] : memref<256xf32, 3>

编译器可以根据地址空间做优化： - memref<*, 3>（shared memory）的访问可以被 stride 优化和 bank conflict 消除。 - memref<*, 5>（private/register）的访问可以减少到单周期操作。 - 不同地址空间间的 memref.load/memref.store 需要显式的 barrier（gpu.barrier）保证同步。

四、Tiling 策略与 GPU 映射

从 linalg.matmul 到 gpu.launch 的 tiling 策略是 GPU 代码生成的核心决策：

linalg.matmul(256x256)
    │
    ▼
tiling: [8, 8] per block, [4, 4] per thread
    │
    ▼
gpu.launch blocks(32x32) threads(8x8)
    └── for i_tile in 0..4:
         for j_tile in 0..4:
           linalg.matmul(64x64) tile
           (分到 shared memory 再进一步 tile)

典型的 GPU matmul tiling 分为三级：

1. Block-level tiling：将 256×256 的 matmul 分成 8×8 的 block，每个 block 计算一个 32×32 的 tile。
2. Shared memory tiling：将每个 block 的 32×32 tile 进一步分解为 4×4 的 warp-level tile。
3. Thread-level tiling：每个 warp 内的 32 个 thread 协作计算一个 4×4 tile 的乘加操作。

// 以下为伪代码级示意，省略 slice 下标推导、shared buffer 分配与类型细节
gpu.launch blocks(%bx, %by) in (%grid_x = 32, %grid_y = 32)
    threads(%tx, %ty) in (%block_x = 8, %block_y = 8) {
  %shared_A = memref.alloc() : memref<8x256xf32, 3>
  %shared_B = memref.alloc() : memref<256x8xf32, 3>
  // … 将 global tile 拷入 shared_A / shared_B …
  gpu.barrier
  %partial = arith.constant 0.0 : f32
  scf.for %k = %c0 to %c256 step %c1 {
    %a = memref.load %shared_A[%tx, %k] : memref<8x256xf32, 3>
    %b = memref.load %shared_B[%k, %ty] : memref<256x8xf32, 3>
    %mul = arith.mulf %a, %b : f32
    %partial = arith.addf %partial, %mul : f32
  }
  memref.store %partial, %C[%bx*8 + %tx, %by*8 + %ty] : memref<256x256xf32>
  gpu.terminator
}

实际 IREE 和 Triton 的 GPU matmul 编译都比这个简化版本复杂得多——包括 warp-level 矩阵乘（TensorCore/WMMA）的调度、双缓冲（double buffering）来隐藏延迟、寄存器压力和 occupancy 的平衡。但这些原理建立在相同的 tiling 层次上。

五、SPIR-V 出口路径

SPIR-V 是 Vulkan/OpenCL 生态的标准中间语言。MLIR 有完整的 gpu → spirv 降阶路径：

gpu.launch ──→ spirv.module
   │                └── spirv.func (device kernel)
   │                     └── spirv.ControlBarrier (同步)
   │                     └── spirv.Load / spirv.Store (内存操作)
   │
   └── gpu.memcpy ──→ spirv.CopyMemory
   └── gpu.barrier ──→ spirv.ControlBarrier

降阶命令：

mlir-opt input.mlir \
  -gpu-kernel-outlining \         # 将 gpu.launch 折叠为 kernel 函数
  -convert-gpu-to-spirv \         # 转换为 SPIR-V 方言
  -spirv-module-to-binary \       # 转换为 SPIR-V 二进制
  -o output.mlir

SPIR-V 的表示层级使其成为 GPU 后端的通用出口：

MLIR IR (gpu dialect)
   │
   ├──→ (convert-gpu-to-spirv) → SPIR-V 方言 → Vulkan / OpenCL 后端
   ├──→ (convert-gpu-to-nvgpu) → nvgpu 方言  → NVIDIA PTX → cubin
   └──→ (gpu-to-llvm)          → LLVM 方言   → AMDGPU / NVPTX 后端

六、MLIR GPU 工具链的实际应用

6.1 运行示例：GPU 向量加法

完整的 GPU 向量加法在 MLIR 中的表示：

module {
  func.func @vector_add(%A: memref<1024xf32>, %B: memref<1024xf32>,
                         %C: memref<1024xf32>) {
    %c1024 = arith.constant 1024 : index
    gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c4, %grid_y = %c1, %grid_z = %c1)
        threads(%tx, %ty, %tz) in (%block_x = %c256, %block_y = %c1, %block_z = %c1) {
      %idx = arith.addi %bx * %c256, %tx : index
      %a = memref.load %A[%idx] : memref<1024xf32>
      %b = memref.load %B[%idx] : memref<1024xf32>
      %sum = arith.addf %a, %b : f32
      memref.store %sum, %C[%idx] : memref<1024xf32>
      gpu.terminator
    }
    return
  }
}

mlir-opt 将此转换为 gpu.module + device kernel，然后由 SPIR-V 或 NVVM 后端消费。

七、与 Triton 和 IREE 的关系

Triton

Triton 是 OpenAI 维护的 GPU 编程语言和编译器。Triton 有自己的 IR（Triton IR）和代码生成管线，并非基于 MLIR GPU 方言。但两者在 AI 编译链中处于同一位置——将高层计算（尤其是 block-level matmul）映射到 GPU。

Triton 与 MLIR 的关系：Torch-MLIR 可以将 PyTorch 模型翻译为 MLIR 方言栈，然后可选通过 Triton 或 MLIR 的 GPU 管线生成代码。两者是竞争+互补的关系——Triton 在 block-level matmul 的自动调度上有优势；MLIR GPU 方言在通用性和多后端支持上有优势。

IREE

IREE 的 GPU 代码生成完全基于 MLIR 的 GPU 方言和 SPIR-V 路径：

IREE 编译流程 (GPU 部分):
linalg (tiling) → GPU tiling → gpu.launch
    → GPU memory promotion (shared memory)
    → SPIR-V / LLVM (NVVM)
    → 设备二进制

IREE 选择 SPIR-V 作为主 GPU 后端的理由是：SPIR-V 是 Khronos 标准，被 Vulkan、OpenCL、Metal（通过 MoltenVK）广泛支持。一套 MLIR Pass 生成的 SPIR-V 可以同时在 Android（Vulkan）、Windows（Vulkan/DirectX via SPIRV-Cross）、Mac（Metal）、Linux（Vulkan）上运行。

八、本篇后续

GPU 代码生成是 AI 编译链的最后几环之一。但完整的端到端 AI 编译还需要之前的所有篇章——Tensor → Linalg → Affine → SCF → GPU。下一篇讲 AI 框架桥接：PyTorch、TensorFlow、JAX 各自怎么把计算图翻译到 MLIR 方言栈。

参考资料

官方文档（A 级）

• MLIR GPU Dialect — https://mlir.llvm.org/docs/Dialects/GPU/
• MLIR SPIR-V Dialect — https://mlir.llvm.org/docs/Dialects/SPIR-V/

源码（A 级）

• mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
• mlir/lib/Conversion/GPUToSPIRV/
• mlir/lib/Dialect/GPU/Transforms/

社区项目（B 级）

• IREE GPU codegen — https://github.com/iree-org/iree
• Triton — https://github.com/triton-lang/triton