MLIR Progressive Lowering：AI 编译器多层渐进式降级架构深度实战

1 阅读 0 评论 0 点赞

摘要：现代 AI 编译器面临的核心挑战是如何将高层计算图高效映射到异构硬件。MLIR 通过 Dialect 层次化设计实现了渐进式降级（Progressive Lowering），让每一层 IR 只关注一个抽象级别，从而构建出可组合、可扩展的编译器基础设施。

一、为什么需要渐进式降级

传统编译器通常采用"一步到位"的方式，直接从高层 IR 降到 LLVM IR 或 PTX。这种方案在 AI 计算日益复杂的今天暴露了三个致命问题：

问题 1：优化 pass 爆炸。当所有优化都挤在同一层 IR 上时，pass 之间互相依赖、互相干扰。一个简单的算子融合可能触发内存布局变化，进而影响向量化决策，形成多米诺效应。

问题 2：硬件适配碎片化。不同硬件（GPU Nvidia/AMD/Intel、NPU、TPU、自定义 ASIC）各有独特的指令集和内存层次结构。统一的 lowering 路径需要对每种硬件编写独立的 lowering table，代码量呈 O(N×M) 增长。

问题 3：调试困难。从高层计算图直接降到 PTX，中间缺失了多个抽象级别的 checkpoint，出错时难以定位是哪个阶段的优化出了问题。

MLIR 的答案是将编译过程拆分为多个 Dialect（方言），每个 Dialect 代表一个抽象级别，通过逐级 lowering 实现分层治理。

二、MLIR Dialect 层次体系

MLIR 的核心创新在于其 Dialect 机制。每个 Dialect 定义了一组 Operation、Type 和 Attribute，代表特定抽象级别的操作原语。

典型 AI 编译器 Dialect 栈

torch-mlir / topline          ← PyTorch/TensorFlow 导入层
  ↓
linalg                        ← 通用 tensor 算子层
  ↓
vector                        ← 向量化抽象层
  ↓
scf / affine                  ← 控制流与循环优化层
  ↓
llvm / rocdl / nvvm           ← 后端具体指令层
  ↓
LLVM IR / PTX / HSACO         ← 最终目标代码

每一层 Dialect 都有其专用的优化 passes。例如： - linalg 层：算子融合、tiling、padding - vector 层：向量化策略、vector masking、contraction lowering - scf 层：循环展开、循环分块、循环交换

Graph 1：Dialect 抽象级别与优化 pass 映射

| 抽象级别 | Dialect      | 优化 pass 示例                    | 典型变换                 |
|----------|-------------|-----------------------------------|--------------------------|
| L6 最高  | torch/tosa  | shape 推导、类型特化               | 图级 dead code elimination |
| L5       | linalg      | 通用算子融合、tiling-by-1          | elementwise fusion       |
| L4       | linalg-on-tensors | bufferization、内存规划      | buffer placement         |
| L3       | vector      | vectorize、unroll、masking         | contraction → vector     |
| L2       | scf + memref| 循环优化、内存预取                  | loop interchange         |
| L1 最低  | llvm/nvvm   | 指令选择、寄存器分配                 | LLVM DAG selection       |

三、渐进式降级的核心机制

3.1 Conversion Pattern 与 Dialect Conversion

MLIR 通过 ConversionPattern 实现跨 Dialect 的 lowering。每个 pattern 定义了从高层 Op 到低层 Op 的重写规则。

// 示例：将 linalg.matmul 转换为 vector 层操作
struct MatmulVectorizationPattern
    : public OpConversionPattern<linalg::MatmulOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult matchAndRewrite(
      linalg::MatmulOp op,
      OpAdaptor adaptor,
      ConversionRewriter &rewriter) const override {
    auto loc = op.getLoc();
    auto lhs = adaptor.getLhs();
    auto rhs = adaptor.getRhs();
    auto dst = adaptor.getDInit();

    // 创建 vector.contract 表示矩阵乘法
    auto contractor = rewriter.create<vector::ContractionOp>(
        loc, lhs, rhs, dst,
        rewriter.getI64ArrayAttr({1, 2, 0}),  // parallel/reduction iterators
        rewriter.getI64ArrayAttr({MITER, NITER, KITER}));
    rewriter.replaceOp(op, contractor.getResultV());
    return success();
  }
};

3.2 Partial Conversion 与 Legality 检查

渐进式降级的关键特性是部分转换（Partial Conversion）。MLIR 的 Conversion Framework 允许某些 Op 被转换、其他 Op 保持不变，通过 ConversionTarget 精确控制转换范围。

// 仅将 matmul 和 generic 转换为 vector representation
// elementwise op 保持不变，可能在后续 pass 中单独处理
ConversionTarget target(*context);
target.addDynamicallyLegalOp<linalg::MatmulOp>(
    [](linalg::MatmulOp op) { return false; });  // 强制全部转换
target.addDynamicallyLegalOp<linalg::GenericOp>(
    [](linalg::GenericOp op) { return false; });
target.addLegalDialect<vector::VectorDialect>();  // 转换后是 legal 的
target.addLegalDialect<arith::ArithDialect>();

// 应用 partial conversion
RewritePatternSet patterns(context);
patterns.add<MatmulVectorizationPattern>(context);
patterns.add<ElementwiseVectorizationPattern>(context);

if (failed(applyPartialConversion(op, target, std::move(patterns))))
    return failure();

3.3 Lowering Pass 管道编排

实际编译器中，渐进式降级是一个精心编排的 pass pipeline：

mlir-pass-pipeline = {
  // Stage 1: L6 → L5 (Tensor → Linalg)
  func.func(torch-to-linalg),

  // Stage 2: L5 算子优化 (Linalg 层内)
  func.func(linalg-fuse-elementwise-ops),
  func.func(linalg-generalize-named-ops),
  func.func(linalg-fuse-packed-ops),

  // Stage 3: Bufferization (Tensor → Memref)
  empty-tensor-to-alloc-tensor,
  one-shot-bufferize,

  // Stage 4: L5 → L3 (Linalg → Vector/SCCF)
  func.func(linalg-vectorize),
  func.func(convert-vector-to-scf),
  func.func(lower-affine),

  // Stage 5: L3 → L1 (到 LLVM IR)
  func.func(convert-vector-to-llvm),
  func.func(convert-scf-to-cf),
  func.func(convert-cf-to-llvm),
  func.func(convert-arith-to-llvm),
  func.func(convert-memref-to-llvm),
  func.func(reconcile-unrealized-casts),

  // Stage 6: LLVM 后端优化
  llvm-request-c-wrappers,
}

四、实战案例：构建自定义 Dialect

让我们通过一个端到端的示例，展示如何为自定义 AI 加速器设计 Dialect 并完成渐进式降级。

4.1 场景设定

假设我们有一个自定义 NPU 加速器，支持以下原生操作： - npu.matmul：矩阵乘，支持 INT8/FP16 - npu.conv2d：卷积，支持 stride/padding/dilation - npu.activation：激活函数（ReLU/GELU/SiLU） - npu.memcpy：DMA 搬运，支持跨内存层次（L1/L2/DRAM）

4.2 定义自定义 Dialect

标题: NPU Dialect Definition

// NPUOps.td — TableGen 定义
def NPU_Dialect : Dialect {
  let name = "npu";
  let summary = "Custom AI accelerator dialect";
  let description = [{
    Operations representing native NPU instructions
    for matrix multiply, convolution, and activation.
  }];
}

// 矩阵乘操作
class NPU_Op<string mnemonic, list<Trait> traits = []> :
    Op<NPU_Dialect, mnemonic, traits>;

def NPU_MatmulOp : NPU_Op<"matmul", [Pure]> {
  let arguments = (ins
    AnyRankedTensor:$lhs,
    AnyRankedTensor:$rhs,
    AnyTensor:$output,
    OptionalAttr<I64Attr>:$tile_m,
    OptionalAttr<I64Attr>:$tile_n,
    OptionalAttr<I64Attr>:$tile_k
  );
  let results = (outs AnyTensor:$result);
}

def NPU_ActivationOp : NPU_Op<"activation", [Pure]> {
  let arguments = (ins
    AnyRankedTensor:$input,
    I32Attr:$kind  // 0=ReLU, 1=GELU, 2=SiLU
  );
  let results = (outs AnyRankedTensor:$result);
}

4.3 渐进式降级实现

第一层降级：从 Linalg 识别 NPU 可加速的算子并替换为 npu Dialect：

// LinalgToNPU Pass
struct LinalgMatmulToNPUPattern
    : public OpRewritePattern<linalg::MatmulOp> {
  OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(
      linalg::MatmulOp op, PatternRewriter &rewriter) const override {

    // 检查硬件约束
    auto lhsType = op.getLhs().getType().cast<ShapedType>();
    auto rhsType = op.getRhs().getType().cast<ShapedType>();
    auto M = lhsType.getShape()[0];
    auto K = lhsType.getShape()[1];
    auto N = rhsType.getShape()[1];

    // NPU 仅支持特定数据类型和形状约束
    if (!isSupportedElementType(lhsType.getElementType()))
      return failure();
    if (M < 16 || N < 16 || K < 16)  // 太小不值得 DMA 搬运
      return failure();

    // 确定 tiling 策略
    auto tileM = std::min(M, static_cast<int64_t>(128));
    auto tileN = std::min(N, static_cast<int64_t>(128));
    auto tileK = std::min(K, static_cast<int64_t>(256));

    // 替换为 npu.matmul
    auto npuMatmul = rewriter.create<npu::MatmulOp>(
        op.getLoc(), op.getLhs(), op.getRhs(), op.getDInit(),
        rewriter.getI64IntegerAttr(tileM),
        rewriter.getI64IntegerAttr(tileN),
        rewriter.getI64IntegerAttr(tileK));

    rewriter.replaceOp(op, npuMatmul.getResult());
    return success();
  }
};

第二层降级：从 npu Dialect 到具体的 runtime call / DMA 指令序列：

// NPUToRuntime Call Pattern
struct MatmulToRuntimePattern : OpRewritePattern<npu::MatmulOp> {
  LogicalResult matchAndRewrite(
      npu::MatmulOp op, PatternRewriter &rewriter) const override {

    // 分配 L1 buffer
    auto lhsL1 = rewriter.create<npu::DmaMemcpyOp>(
        op.getLoc(), op.getLhs(), DmaDirection::HOST_TO_L1);
    auto rhsL1 = rewriter.create<npu::DmaMemcpyOp>(
        op.getLoc(), op.getRhs(), DmaDirection::HOST_TO_L1);

    // 生成双缓冲 pipeline
    auto output = rewriter.create<npu::ComputeMatmulOp>(
        op.getLoc(), lhsL1, rhsL1,
        op.getTileM(), op.getTileN(), op.getTileK());

    // 结果写回 DRAM
    rewriter.create<npu::DmaMemcpyOp>(
        op.getLoc(), output, DmaDirection::L1_TO_HOST);
    rewriter.eraseOp(op);
    return success();
  }
};

五、Bufferization：内存规划的桥梁

渐进式降级中最关键的环节之一是 Bufferization（缓冲化）——将基于 value semantics 的 tensor op 转换为基于 memory semantics 的 memref op。

传统方案的问题

早期 IREE 和 TensorFlow/XLA 采用提前规划（eager allocation）：在 lowering 之前一次性分配所有 buffer，然后在后续 pass 中反复 in-place 分析。这导致：

Bufferization pass 需要理解后续所有优化
in-place 分析是 NP-hard 问题
优化 pass 改变计算顺序后可能使原有 buffer 分配失效

One-Shot Bufferize 方案

MLIR 引入的 one-shot-bufferize pass 将 bufferization 与 partial conversion 统一：

输入: linalg op on tensors (value semantics)
      ↓
  [in-place 可用性分析]
      ↓
  判断: 能否 in-place？ ─┬─ 是 → 直接重用已有 buffer
                        └─ 否 → 分配新 buffer (memref.alloc)
      ↓
输出: linalg/scf op on memrefs (reference semantics)

核心算法基于别名分析（Alias Analysis）和冲突读分析（Conflicting Read Analysis）：

// 伪代码：in-place 分析逻辑
bool canBufferInPlace(OpResult bufferizedResult, Value operand,
                      BufferizationState &state) {
  // 检查 1: operand 是否在该 use 之后被其他 use 读取
  auto hasConflictingRead = state.hasReadAfterWrite(
      operand, bufferizedResult.getDefiningOp());
  if (hasConflictingRead) return false;

  // 检查 2: operand 和 result buffer 是否有别名
  auto aliasInfo = analyzeAliasing(operand, bufferizedResult);
  if (mayAlias(aliasInfo)) return false;

  // 检查 3: 写入是否覆盖整个 shape（否则不能 in-place）
  if (!writesFullTensor(bufferizedResult)) return false;

  return true; // 可以 in-place
}

六、硬件后端的渐进式降级策略

6.1 GPU PTX Backend

NVIDIA GPU 的渐进式降级栈：

vector/arith/scf
  ↓
convert-vector-to-gpu      (→ gpu.subgroup_mma / gpu.barrier)
  ↓
convert-gpu-to-nvvm        (→ nvvm.mma.sync / nvvm.shfl.sync)
  ↓
convert-nvvm-to-llvm       (→ llvm.call @nvvm intrinsics)
  ↓
LLVM NVPTX Backend         (→ PTX → cubin → SASS)

关键的 gpu.subgroup_mma 利用 Tensor Core 的 MMA 指令：```mlir // 使用 Tensor Core 的矩阵乘片段 %descA = gpu.mma_create async [%lhs] tile:[128, 128, 16] ... %descB = gpu.mma_create async [%rhs] tile:[128, 128, 16] ... %accum = gpu.mma_compute async %descA, %descB, %acc

### 6.2 AMD ROCm Backend

AMD GPU 的降级差异在于使用 `rocdl` Dialect 映射到 ROCm Device Library：

vector/arith/scf ↓ convert-vector-to-rocdl (→ rocdl.mfma / rocdl.ds.permute) ↓ convert-rocdl-to-llvm (→ llvm.call @llvm.amdgcn.mfma) ↓ LLVM AMDGPU Backend (→ AMDGPU ISA)

### 6.3 CPU SIMD Backend

CPU 目标的渐进式降级重点在于自动向量化：

linalg (on tensors) ↓ linalg-vectorize (→ vector.contract / vector.transfer_read) ↓ vector-unroll (展开到硬件向量宽度) ↓ convert-vector-to-llvm (→ llvm.intr.masked.load / <4 x float>) ↓ LLVM X86/AArch64 Backend (→ AVX-512 / NEON / SVE)

## 七、生产级 MLIR 编译器的工程实践

### 7.1 Pass 依赖与缓存

生产级编译器通过 `mlir-cpu-runner` 或 JIT 引擎管理 pass pipeline：

```cpp
// 编译 session 管理
class CompilationSession {
public:
  LogicalResult compile(func::FuncOp func) {
    // 阶段 1：高层优化（缓存友好）
    {
      PassManager pm(getContext());
      addHighLevelPasses(pm);
      if (failed(pm.run(func))) return failure();
    }

    // 阶段 2：Bufferization（不可逆变换）
    {
      PassManager pm(getContext());
      pm.addPass(createOneShotBufferizePass());
      if (failed(pm.run(func))) return failure();
    }

    // 阶段 3：向量化（依赖 buffer 分配结果）
    {
      PassManager pm(getContext());
      pm.addPass(createLinalgVectorizePass());
      pm.addPass(createCanonicalizerPass());  // 清理残留
      if (failed(pm.run(func))) return failure();
    }

    // 阶段 4：后端 lowering
    {
      PassManager pm(getContext());
      addBackendPasses(pm);
      if (failed(pm.run(func))) return failure();
    }

    // 阶段 5：Native lowering
    return lowerToLLVMAndCodegen(func);
  }
};

7.2 调试与诊断

渐进式降级的一个核心优势是中间 checkpoint 的可用性：

# 在每个 lowering pass 后 dump IR
mlir-opt input.mlir \
  --torch-to-linalg \
  --mlir-print-ir-after-all \        # 每步后打印 IR
  --mlir-print-ir-after-change \     # 仅变化时打印
  --mlir-print-ir-on-failure \       # 失败时打印当前 IR
  --verify-each \                    # 每步后做验证
  -o output.mlir

# 使用 reproducer 捕获失败场景
mlir-opt broken.mlir \
  --linalg-vectorize \
  --mlir-print-ir-on-failure 2>&1 | tee reproducer.mlir

7.3 TableGen 驱动的重写规则

MLIR 推荐使用 TableGen 定义 declarative rewrite rules（DRR），减少样板代码：

// 将 linalg.transpose 直接降级为 vector.transpose
def TransposeToVector : Pattern<
  (Linalg_TransposeOp $input, $output, $permutation),
  (Vector_TransposeOp 
     (Vector_TransferReadOp $input, ...),  // 读取
     $permutation),                         // 转置
  [(IsStaticShape $input),                 // 约束：静态 shape
   (IsVectorizableType $elementType)]     // 约束：可向量化元素
>;

八、渐进式降级的性能收益

渐进式降级带来的性能收益主要来自三个维度：

维度 1：分层优化效果更优

每一层 IR 只关注一个抽象级别，pass 之间的干扰大幅减少。据 IREE 团队的实验数据，渐进式降级相比一步式 lowering，torch-mlir → linalg → vector → LLVM 三步走，在 NVIDIA A100 上平均获得 1.3x-2.1x 的端到端性能提升。

维度 2：硬件适配成本更低

新增硬件后端只需实现从 vector ops 到 native Dialect 的 lowering pattern，无需关心高层的图优化。从 IREE 的实际经验看，支持一个新硬件后端（如 CUDA → ROCm）的工作量约 2-3 人月（含基础 kernel library）。

维度 3：编译时间更可控

渐进式降级天然支持增量编译和并行化。不同函数可以独立进入不同的 lowering 阶段，形成 pipeline parallelism。在 IREE 中，大型模型（如 7B LLM）的编译时间从 45 分钟优化到 8 分钟。

九、总结：渐进式降级的工程哲学

MLIR 的渐进式降级体现了软件工程中"单一职责"与"关注点分离"的核心原则：

每个 Dialect 只解决一个问题：linalg 负责算子语义，vector 负责 SIMD 抽象，LLVM Dialect 负责具体指令
每个 Pass 只做一件事：operator fusion pass 不涉及内存布局，bufferization pass 不涉及向量化
可组合优于继承：通过组合不同的 pass pipeline 适配不同硬件，无需修改核心 IR

MLIR 证明了编译器基础设施不一定是大一统的黑箱，通过精细的层次化设计，可以构建出既灵活又高效的 AI 编译栈。对于正在构建自定义 AI 加速器的团队，从 vector Dialect 开始设计自己的 lowering 路径是最低成本、最高灵活度的选择。

参考资源

点赞(0) 打赏

本文分类：操作系统
本文标签：MLIR 编译器 AI编译 LLVM 渐进式降级 Hardware Backend Dialect Pass Pipeline
浏览次数：1 次浏览
发布日期：2026-10-07 17:11:14
本文链接：https://www.ybb.press/cat_操作系统_bf19/15670.html

上一篇 > Linux 内核网络栈深度解析：从 Socket 到驱动层
下一篇 > Linux 内核块设备与文件系统深度实战：从 VFS 到 io_uring

MLIR Progressive Lowering：AI 编译器多层渐进式降级架构深度实战

MLIR Progressive Lowering：AI 编译器多层渐进式降级架构深度实战

一、为什么需要渐进式降级

二、MLIR Dialect 层次体系

典型 AI 编译器 Dialect 栈

Graph 1：Dialect 抽象级别与优化 pass 映射

三、渐进式降级的核心机制

3.1 Conversion Pattern 与 Dialect Conversion

3.2 Partial Conversion 与 Legality 检查

3.3 Lowering Pass 管道编排

四、实战案例：构建自定义 Dialect

4.1 场景设定

4.2 定义自定义 Dialect

4.3 渐进式降级实现

五、Bufferization：内存规划的桥梁

传统方案的问题

One-Shot Bufferize 方案

六、硬件后端的渐进式降级策略

6.1 GPU PTX Backend

7.2 调试与诊断

7.3 TableGen 驱动的重写规则

八、渐进式降级的性能收益

九、总结：渐进式降级的工程哲学

参考资源

评论列表共有 0 条评论

发表评论取消回复

MLIR Progressive Lowering：AI 编译器多层渐进式降级架构深度实战

MLIR Progressive Lowering：AI 编译器多层渐进式降级架构深度实战

一、为什么需要渐进式降级

二、MLIR Dialect 层次体系

典型 AI 编译器 Dialect 栈

Graph 1：Dialect 抽象级别与优化 pass 映射

三、渐进式降级的核心机制

3.1 Conversion Pattern 与 Dialect Conversion

3.2 Partial Conversion 与 Legality 检查

3.3 Lowering Pass 管道编排

四、实战案例：构建自定义 Dialect

4.1 场景设定

4.2 定义自定义 Dialect

4.3 渐进式降级实现

五、Bufferization：内存规划的桥梁

传统方案的问题

One-Shot Bufferize 方案

六、硬件后端的渐进式降级策略

6.1 GPU PTX Backend

7.2 调试与诊断

7.3 TableGen 驱动的重写规则

八、渐进式降级的性能收益

九、总结：渐进式降级的工程哲学

参考资源

MLIR Progressive Lowering：AI 编译器多层渐进式降级架构深度实战

V8 Turbofan JIT 编译器内部机制：从字节码到机器码的优化流水线深度解析

MLIR 编译器基础设施：多级中间表示与 AI/ML 系统编译的全栈工程

PyTorch 2 torch.compile 深度工程实战：Dynamo 追踪与 Inductor 后端优化全链路剖析

评论列表 共有 0 条评论

发表评论 取消回复

评论列表共有 0 条评论

发表评论取消回复