第 7 课:逐级 Lowering 到 LLVM Dialect

本课目标与完成标准

学完后你应能:设计高层→linalg/scf→cf/memref→LLVM 的 lowering pipeline,并区分 LLVM dialect 与 LLVM IR。

通过必须同时满足:

  • 独立补齐本课唯一的代码填空题,并通过给定检查;
  • 三个问答题均说明因果链,而不是只报术语;
  • 能指出至少一个正确性边界和一个性能取舍;
  • 总分不低于 8/10,且没有一票否决级概念错误。

前置关系

  • 课程前置:编译原理基础、C++ 阅读能力
  • 本课在路线中的作用:MLIR 通常保留多层 IR:高层语义先变为结构化计算,再降控制流/内存,最后到 LLVM dialect 并翻译成 LLVM IR。

核心心智模型

1. 它是什么,解决什么问题

把高层神经网络算子一步编译到底层硬件代码,会让各层的复杂性耦合在一起,优化 Pass 也难以复用。MLIR 采用渐进式多级降低(Progressive Multi-Level Lowering): High-Level Graph (Torch/TOSA)→图融合Linalg on Tensors→分块缓冲SCF / MemRef→控制流扁平化CF / LLVM Dialect→翻译LLVM IR→机器码ELF / PTX \text{High-Level Graph (Torch/TOSA)} \xrightarrow{\text{图融合}} \text{Linalg on Tensors} \xrightarrow{\text{分块缓冲}} \text{SCF / MemRef} \xrightarrow{\text{控制流扁平化}} \text{CF / LLVM Dialect} \xrightarrow{\text{翻译}} \text{LLVM IR} \xrightarrow{\text{机器码}} \text{ELF / PTX} 每一个编译阶段仅剥离一层抽象,消灭一个维度的语义(例如:先消除多维张量代数,再消除结构化循环,再消除结构化控制流,最后消除高级内存描述符)。 本课解决的核心工程问题是:理解渐进降级流水线的强前置依赖图,深刻辨析 MLIR 中的 llvm Dialect 与底层独立 llvm::Module(LLVM IR)的本质分水岭,掌握从 LLVM 方言经过 mlir-translate 导出为机器码的端到端完整链路。

2. 它如何工作

渐进降级拓扑与最终代码生成时序如图 7-1 与图 7-2 所示:

图 7-1:MLIR 渐进式多级降低与代码生成阶段性能加速比

图 7-2:从 LLVM Dialect 到底层硬件二进制目标代码的翻译运行流程

  1. 中端逐步降解(Middle-end Lowering Pipeline):

    • 阶段 1:linalg 算子通过 Tiling 和 Bufferization 降解为 scf.for(结构化循环)和 memref(物理内存读写);
    • 阶段 2:scf.for / scf.if 通过 -convert-scf-to-cf 展开为带条件跳转的基本块控制流(cf.br / cf.cond_br);
    • 阶段 3:memref 通过 -finalize-memref-to-llvm 展开为包含物理基指针、对齐指针、偏移量与多维步长的一维扁平 LLVM 结构体描述符(MemRef Descriptor Struct);
    • 阶段 4:arith 算术指令、cf 分支与 func 函数通过 -convert-arith-to-llvm、-convert-cf-to-llvm 和 -convert-func-to-llvm 降为 llvm Dialect,最后用 -reconcile-unrealized-casts 消掉各转换之间遗留的 unrealized_conversion_cast。
  2. LLVM Dialect 的本质: 此时虽然所有操作的前缀都变成了 llvm.(如 llvm.func、llvm.add、llvm.load、llvm.return),但它们在内存中仍然完完全全是 MLIR 的 Operation 实例。它们仍然遵循 MLIR 的 C++ 内存对象树,并驻留在 mlir::MLIRContext 之中。

  3. 离开 MLIR 世界的物理翻译(mlir-translate): 如图 7-2 所示,只有调用 mlir-translate --mlir-to-llvmir 时,程序才会启动 1:1 的树遍历,在内存中正式构建出真正的 LLVM 原生数据结构 llvm::Module 和 llvm::Instruction。随后移交给标准 LLVM 后端优化器(opt -O3),经过指令选择(ISel)、寄存器分配(RegAlloc)最终编译为目标平台的机器码(x86/ARM ELF 或 NVIDIA PTX/CUBIN)。

3. 正确性条件与常见误区

  • Pass 管道的严格拓扑偏序(Pass Order Invariant): Pass 列表有顺序约束。例如 scf 必须先经 -convert-scf-to-cf 降为 cf,后面的 -convert-cf-to-llvm 才有东西可转;-reconcile-unrealized-casts 必须放在所有 *-to-llvm 转换之后,过早运行会留下无法消除的 cast。注意 -convert-func-to-llvm 使用 LLVMTypeConverter,函数签名中的 memref 会被它直接转换为描述符结构体,不要求 memref 先降级。

  • 混淆 LLVM Dialect 与 LLVM IR: 最常见的新手认知误区是以为打印出 llvm.add 就等于生成了汇编或 LLVM IR。llvm.add 只是 MLIR 世界里对方言操作的一个镜像抽象,尚未逃离 MLIR 运行时的作用域,无法直接送入汇编器或链接器。

  • Index 类型的目标机器位宽依赖(DataLayout Constraint): MLIR 的 index 类型的位宽由目标决定。降级到 LLVM 方言时,LLVMTypeConverter 按 index-bitwidth 选项或 Module 上的数据布局(DataLayout)决定 index 映射为 i32 还是 i64,未指定时默认 64 位。为 32 位目标编译却沿用默认值,会得到与目标指针宽度不一致的代码。

4. 性能与工程取舍

  • 多层降级的收益倍增: 如图 7-1 所示,高层 Pass 负责做全局性、大粒度的算子融合与并行瓦片切分(耗时几十毫秒,带来数倍吞吐提升);中层 Pass 负责消除冗余循环;底层 LLVM 负责精细的指令调度、SIMD 向量展开和寄存器染色分配。各层 Pass 各司其职,避免了底层编译器在杂乱无章的扁平代码中尝试逆向工程高阶意图。

  • 异构目标(GPU/CPU)的分叉时机(Lowering Forking Point): 在异构计算编译中,Host(CPU)与 Device(GPU)的代码必须在降级流水线中找准分叉点。过早分叉会导致公共优化逻辑无法复用;过晚分叉则无法充分利用 GPU 独有的硬件层次(如 Warp/Block 与 Shared Memory)。

具体演示

以一个极简整数加法函数的降级终端形态为例:

module {
  llvm.func @add(%a: i32, %b: i32) -> i32 {
    %sum = llvm.add %a, %b : i32
    llvm.return %sum : i32
  }
}
  • 输入:经由 arith.addi 与 func.func 降级而来的终极 MLIR 形式。

  • 操作:llvm.add 与 llvm.return 严格映射到 LLVM 标量指令集。

  • 终检验证:运行 mlir-translate --mlir-to-llvmir,该 MLIR 会被无缝翻译成纯正的 LLVM IR:

    define i32 @add(i32 %0, i32 %1) {
      %3 = add i32 %0, %1
      ret i32 %3
    }

    此时方可移交 clang 或 llc 编译为 .o 目标二进制机器码。

请在阅读后先合上这一节,用自己的语言复述“输入状态 → 中间状态 → 输出状态”,再做练习。

实践任务:唯一代码填空题

补齐 LLVM dialect 的整数加法与返回。

规则:只能修改 TODO/______ 所在位置;不要删除断言或放宽误差。代码注释说明了每个边界条件。

Code
%%writefile /tmp/lesson07.mlir
module {
  llvm.func @add(%a: i32, %b: i32) -> i32 {
    %sum = ______ %a, %b : i32
    ______ %sum : i32
  }
}

检查方法

用 mlir-opt 验证,再用 mlir-translate --mlir-to-llvmir 检查得到 LLVM IR;无工具时标记待验证。

提交时请给出:补齐后的代码、实际运行输出(环境不可用时注明“仅静态审查”)以及对失败用例的解释。

Q1

不要背定义:请从输入、状态变化和输出三个阶段解释“逐级 Lowering 到 LLVM Dialect”的工作机制。

你的答案:

Q2

为什么看到 llvm.* op 还不能声称已经生成机器码?

你的答案:

Q3

若目标是 GPU/NVVM,应在哪一层分叉 pipeline?

你的答案:

评分与通过规则

  • 代码 4 分:正常输入 2 分,边界输入 1 分,解释实现 1 分;
  • Q1~Q3 各 2 分;
  • 一票否决:结果碰巧正确但核心因果链错误、删除边界检查、把未运行结果说成实测。

需要提示时按四级机制请求:概念区域 → 具体方向 → 关键局部 → 完整答案。

参考答案(仅 answer 分支)

先完成题目再核对。即使代码一致,也要能解释关键步骤,并尝试更换一个输入规模。

Code
%%writefile /tmp/lesson07.mlir
module {
  llvm.func @add(%a: i32, %b: i32) -> i32 {
    %sum = llvm.add %a, %b : i32
    llvm.return %sum : i32
  }
}

Q1 参考答案

从三个阶段拆解其运行机制:

  1. 输入阶段(高阶结构化张量准备):以已完成内存分析的结构化计算 IR(如 linalg + scf + memref)为起点,确保所有高阶数据流已经确定性降级为带物理布局和循环边界的中间形态。

  2. 状态变化阶段(洋葱式逐级抽象剥离与拓扑降级):

    • 第一级(控制流降级):运行 -convert-scf-to-cf 将高层 scf.for/if 剥离为基础分支控制流 cf.br/cond_br;
    • 第二级(内存描述符展开):运行 -finalize-memref-to-llvm 将抽象多维 memref 展开为包含基地址指针、对齐指针与动态维度的 LLVM 扁平 Struct 描述符;
    • 第三级(算术与函数降级):运行 -convert-arith-to-llvm 和 -convert-func-to-llvm,将所有算术操作与函数符号整体替换为 llvm. 前缀的镜像算子,再运行 -reconcile-unrealized-casts,此时全模块只含 LLVM Dialect。
  3. 输出阶段(跨框架物理翻译与机器码生成):调用 mlir-translate --mlir-to-llvmir 走出 MLIR 上下文,在内存中正式具象化原生的 llvm::Module C++ 对象树。随后调用 LLVM 原生 TargetMachine 发射本地机器汇编或二进制可执行文件(ELF / PTX)。

Q2 参考答案

在 MLIR 文本中看到 llvm.* 操作(如 llvm.add、llvm.load),绝对不能声称已经生成了物理机器码,原因在于:

  1. 依然属于 MLIR 内存树节点:llvm.* 仅仅是 MLIR 官方提供的众多 Dialect 之一(LLVMDialect)。它本质上与 arith、linalg 没有高下之分,依然是由 mlir::Operation 构成的抽象语法树,驻留在 MLIRContext 堆内存中,其数据结构并不是物理机器指令,甚至还不是原生的 llvm::Instruction。

  2. 距离机器码还有巨大的物理跨度:

    • 必须通过 mlir-translate 序列化并转换成原生的 llvm::Module(LLVM IR);
    • 真正的机器码生成(如 x86 汇编或 PTX 汇编)是由 LLVM 后端(Target-specific Code Generator)完成的。这期间必须经历指令选择(Instruction Selection: SelectionDAG 或 GlobalISel)、机器指令调度(Machine Instruction Scheduling)、物理寄存器分配(Register Allocation)、栈帧布局(Prologue/Epilogue Insertion)等极其沉重复杂的硬件映射过程。当前代码甚至尚未脱离纯文本 IR 阶段。

Q3 参考答案

在编译异构计算程序(如包含 CPU 宿主代码与 GPU 内核的神经网络)时,Pipeline 通常在 Bufferize 之后、并行循环映射为 gpu.launch 并完成 kernel outlining(gpu.launch_func + gpu.module)的这一层分叉:

  1. 分叉点位置(gpu / nvvm 方言分界线): 流水线在高层图优化、Linalg 算子融合和 Bufferization 阶段保持统一。当计算被分块并映射到 GPU 的 Grid / Block 并行层次后(例如 -convert-linalg-to-parallel-loops → -gpu-map-parallel-loops → -convert-parallel-loops-to-gpu 生成 gpu.launch,再用 -gpu-kernel-outlining 拆出 kernel),流水线正式分叉:
    • Device(GPU)分支:outlining 后的 kernel 位于嵌套的 gpu.module 中,用 -convert-gpu-to-nvvm(或 -convert-gpu-to-rocdl)降为 nvvm/rocdl + llvm 方言,再由 -nvvm-attach-target 与 -gpu-module-to-binary 序列化为 PTX/CUBIN,作为二进制字面量嵌入 IR。
    • Host(CPU)分支:-gpu-to-llvm 将外层控制流、gpu.alloc、gpu.launch_func 等降为 llvm 方言中对 mgpuMemAlloc、mgpuLaunchKernel 等运行时包装函数的调用,这些包装函数(如 libmlir_cuda_runtime)再调用 CUDA Driver API。
  2. 过早与过晚分叉的工程坏味道:过早分叉(如在 Tensor 层)会导致 Host 和 Device 无法共享图优化和常量折叠;过晚分叉(如降到底层标量 LLVM 后)会导致 GPU 的多维网格、Warp 洗牌和共享内存语义被不可逆地打散,再也无法生成高性能的 GPU 汇编。

参考资料

资料用于建立事实基线;面试回答仍需用自己的语言组织。