MLIR 多层次中间表示：构建领域专用编译器的深度实践

传统编译器优化框架（如 LLVM IR）在处理深度学习框架多样化的计算图时，面临抽象层级单一、领域语义丢失的困境。MLIR（Multi-Level Intermediate Representation）通过"方言（Dialect）"机制，让编译器基础设施能够同时保留高层张量语义和底层硬件特性，实现从 AI 模型到芯片指令的无缝 lowering。本文深入剖析 MLIR 的核心架构，并通过构建一个完整的张量计算 Dialect，展示如何用 C++ 和 TableGen 实现多级 IR 转换管线。

一、编译器碎片化困境与 MLIR 的诞生

深度学习编译器生态长期面临"碎片化"问题。TensorFlow 的 XLA 使用 HLO IR，TVM 使用 Tensor Halide IR，PyTorch 的 TorchScript 有专属的中间表示，而硬件厂商（NVIDIA、Intel、AMD）又各自维护一套私有编译器栈。每种 IR 都试图在"表达力"和"可优化性"之间寻找平衡，但没有任何一种能同时满足从算子融合到指令调度的全链路需求。

MLIR 的核心洞察是：编译器不需要一种 IR，而是一套可扩展的 IR 定义框架。它引入三个关键抽象：

• Operation（操作）：IR 的基本单元，有零个或多个操作数、结果、属性和区域
• Type（类型）：可扩展的类型系统，支持自定义张量类型、量化类型、内存引用类型
• Attribute（属性）：编译时常量元数据，如布局信息、编译提示、硬件亲和性标注

这些抽象通过 Dialect（方言） 分组管理，每个 Dialect 定义一组语义相关的 Operation、Type 和 Attribute。Dialect 之间通过显式的 conversion pass 相互转换，形成多层 lowering 链。

1.1 与传统 IR 的关键差异

二、MLIR 核心架构解析

2.1 Operation 的解剖结构

每个 MLIR Operation 是一个 SSA 值生产者，其结构如下：

// 一个典型的 MLIR Operation 结构
%result = "matmul_inspect"(%lhs, %rhs) {
  tile_sizes = [32, 32, 32],
  pipeline_depth = 2
} : (tensor<256x256xf32>, tensor<256x256xf32>) -> tensor<256x256xf32>

其中 %lhs 和 %rhs 是 SSA 操作数，"matmul_inspect" 是 Operation 名称，{ tile_sizes = ... } 是属性字典，最后的类型标注定义输入输出类型签名。

2.2 Region 与嵌套结构

与 LLVM IR 的基本块不同，MLIR Operation 可以包含 Region——一个嵌套的 Operation 列表，形成层级化的控制流：

func.func @forward(%input : tensor<1x3x224x224xf32>) -> tensor<1x1000xf32> {
  %conv_out = "tensor.conv"(%input, %weight) ^
    padding = {0, 0, 0, 0}, strides = {2, 2}
  } : (...) -> tensor<1x64x112x112xf32>
  %relu_out = "tensor.relu"(%conv_out) : (...) -> ...
  return %relu_out : ...
}

这种结构化表示使得高层优化 pass 可以直接在 tensor 级别操作，无需先 lowering 到标量循环。

2.3 类型系统的可扩展性

MLIR 的类型系统允许 Dialect 定义语义化类型：

// DLTI（数据类型接口）描述硬件内存布局
struct MatrixType : public Type::TypeBase<MatrixType, Type, TypeStorage> {
  using Base::Base;

  // 工厂方法
  static MatrixType get(MLIRContext *ctx, unsigned rows, unsigned cols,
                        StringRef layout);

  // 访问器
  unsigned getRows() const;
  unsigned getCols() const;
  StringRef getLayout() const;  // "row_major" 或 "col_major"
};

三、构建自定义 Dialect：张量计算方言实战

下面构建一个名为 tensor_opt 的 Dialect，支持基本张量运算和硬件亲和性标注。

3.1 TableGen 定义

MLIR 使用 TableGen 描述 Dialect 结构，自动生成 C++ 代码：

// TENSOR_OPT.td
ifndef TENSOR_OPT_TD

include "mlir/IR/OpBase.td"

def TensorOpt_Dialect : Dialect {
  let name = "tensor_opt";
  let summary = "Tensor optimization dialect for accelerator targeting";
  let description = [{
    This dialect provides operations for tensor computations with 
    hardware affinity annotations, enabling multi-stage lowering 
    from algorithmic representation to hardware-specific code.
  }];
  let cppNamespace = "::mlir::tensor_opt";
}

// 基础张量运算操作
def TensorOpt_MatMulOp : Op<TensorOpt_Dialect, "matmul",
    [Pure, AttrSizedOperandSegments]> {
  let summary = "Matrix multiplication with tiling hints";
  let description = [{
    Performs matrix multiplication with optional tiling annotations
    for accelerator-friendly code generation.
  }];

  let arguments = (ins
    AnyRankedTensor:$lhs,
    AnyRankedTensor:$rhs,
    OptionalAttr<I64ArrayAttr>:$tile_sizes,
    OptionalAttr<BoolAttr>:$transpose_lhs,
    OptionalAttr<StrAttr>:$preferred_device
  );

  let results = (outs AnyRankedTensor:$output);

  let assemblyFormat = [{
    attr-dict `(` $lhs `,` $rhs `) type($lhs) `->` type($output)
  }];
}

// ReLU 激活操作
def TensorOpt_ReluOp : Op<TensorOpt_Dialect, "relu", [Pure]> {
  let summary = "ReLU activation function";

  let arguments = (ins AnyRankedTensor:$input);
  let results = (outs AnyRankedTensor:$output);

  let assemblyFormat = [{
    attr-dict `(` $input `) ` type($input) `->` type($output)
  }];
}

endif // TENSOR_OPT_TD

3.2 C++ 注册与实现

// TensorOptDialect.cpp
#include "mlir/IR/DialectImplementation.h"
#include "mlir/Interfaces/InferTypeOpInterface.h"

using namespace mlir;

void TensorOptDialect::initialize() {
  addOperations<
#define GET_OP_LIST
#include "TensorOptOps.inc"
  >();
}

// MatMulOp 的折叠优化：零矩阵短路
LogicalResult MatMulOp::fold(FoldAdaptor adaptor,
                              SmallVectorImpl<OpFoldResult> &results) {
  // 如果左侧或右侧是常量零张量，直接返回零
  auto lhsConst = adaptor.getLhs().dyn_cast<DenseElementsAttr>();
  auto rhsConst = adaptor.getRhs().dyn_cast<DenseElementsAttr>();

  if (lhsConst && lhsAttr.isa<SplatElementsAttr>() &&
      lhsConst.getSplatValue<float>() == 0.0f) {
    results.push_back(createZeroInitTensor(getResult().getType()));
    return success();
  }
  // ... 类似处理 rhs ...
  return failure();
}

// 类型推断接口
LogicalResult MatMulOp::inferReturnTypes(
    MLIRContext *context, std::optional<Location> location,
    ValueRange operands, DictionaryAttr properties,
    OpaqueProperties operandsProperties, RegionRange regions,
    SmallVectorImpl<Type> &inferredReturnTypes) {

  auto lhsType = operands[0].getType().cast<RankedTensorType>();
  auto rhsType = operands[1].getType().cast<RankedTensorType>();

  // 简化的形状推导：MxK * KNx -> MxN
  auto M = lhsType.getDimSize(0);
  auto N = rhsType.getDimSize(1);
  auto elementType = lhsType.getElementType();

  auto resultType = RankedTensorType::get({M, N}, elementType);
  inferredReturnTypes.push_back(resultType);
  return success();
}

四、Pass Manager 与多级 Lowering 管线

MLIR 的 Pass Manager 支持在任意抽象层级执行转换，并通过 ConversionTarget 管理合法化要求。

4.1 Pass 定义

// TensorOptLoweringPass.cpp
struct TensorOptToAffinePass
    : public PassWrapper<TensorOptToAffinePass, OperationPass<ModuleOp>> {

  void runOnOperation() override {
    // 定义转换目标：所有 tensor_opt.* 操作转换后消失
    ConversionTarget target(getContext());
    target.addIllegalDialect<tensor_opt::TensorOptDialect>();
    target.addLegalDialect<affine::AffineDialect, arith::ArithDialect>();

    // 定义 lowering 模式
    TensorOptToAffineLowering lowering;

    // 执行完整转换
    if (failed(applyFullConversion(getOperation(), target, lowering))) {
      signalPassFailure();
    }
  }

  // Pass 配置选项
  Option<int> maxTileSize{*this, "max-tile-size", 
                          cl::desc("Maximum tile size for tiling"),
                          cl::init(64)};
};

// 注册 Pass
static PassRegistration<TensorOptToAffinePass> pass(
    "convert-tensor-opt-to-affine",
    "Convert tensor_opt operations to affine loops");

4.2 Lowering 模式实现

class TensorOptToAffineLowering : public ConversionPatternRewritePattern {
public:
  using ConversionPatternRewritePattern::ConversionPatternRewritePattern;

  LogicalResult matchAndRewrite(
      tensor_opt::MatMulOp op, OpAdaptor adaptor,
      PatternRewriter &rewriter) const override {

    auto loc = op.getLoc();
    auto lhs = adaptor.getLhs();
    auto rhs = adaptor.getRhs();
    auto resultType = op.getResult().getType();

    // 创建累加器（零初始化）
    auto initVal = rewriter.create<arith::ConstantOp>(
        loc, rewriter.getF32FloatAttr(0.0));
    auto initTensor = rewriter.create<linalg::FillOp>(
        loc, initVal, resultType);

    // 生成 affine 循环嵌套
    auto M = rewriter create<arith::ConstantIndexOp>(loc, /* lhs.rows */);
    auto K = rewriter.create<arith::ConstantIndexOp>(loc, /* lhs.cols */);
    auto N = rewriter.create<arith::ConstantIndexOp>(loc, /* rhs.cols */);

    // 创建累加循环
    auto loop = scf::buildLoopNest(
        rewriter, loc, {M, K, N},
        [&](OpBuilder &b, Location loc, ValueRange ivs) {
      auto lhsElem = b.create<affine::LoadOp>(loc, lhs, ivs[0], ivs[1]);
      auto rhsElem = b.create<affine::LoadOp>(loc, rhs, ivs[1], ivs[2]);
      auto mul = b.create<arith::MulFOp>(loc, lhsElem, rhsElem);
      auto acc = b.create<affine::LoadOp>(loc, initTensor, ivs[0], ivs[2]);
      auto updated = b.create<arith::AddFOp>(loc, acc, mul);
      b.create<affine::StoreOp>(loc, updated, initTensor, ivs[0], ivs[2]);
    });

    // 替换原操作
    rewriter.replaceOp(op, loop.getResult(0));
    return success();
  }
};

4.3 完整 Lowering 管线

一个完整的 MLIR 编译器管线通常包含以下阶段：

tensor_opt.dialect (高层张量语义)
    │
    ▼ convert-tensor-opt-to-linalg
linalg.dialect (结构化运算)
    │
    ▼ linalg-fuse-operators (算子融合)
    ▼ linalg-tile (分块)
    ▼ linalg-vectorize (向量化)
    │
    ▼ convert-linalg-to-affine
affine.dialect + scf.dialect (循环嵌套)
    │
    ▼ affine-loop-fusion (循环融合)
    ▼ affine-scalar-replacement (标量替换)
    │
    ▼ convert-scf-to-cf + convert-affine-to-llvm
cf.dialect + llvm.dialect (LLVM IR)
    │
    ▼ LLVM 后端优化
machine code (目标指令)

每一层都保留特定的领域语义：高层做算法级优化（算子融合、内存分配），中层做循环变换（分块、展开、向量化），底层做指令级优化（寄存器分配、指令调度）。

五、MLIR 在 AI 编译器中的工业实践

5.1 IREE（Intermediate Representation Execution Environment）

Google 的 IREE 项目使用 MLIR 作为端到端 AI 编译器基础设施，其核心管线包括：

• StableHLO → IREE Input Dialect：从 TensorFlow/XLA 的前端 IR 导入
• IREE Transform Dialect：执行基于脚本的转换（类似 eBPF 的轻量级 pass）
• IREE Stream Dialect：显式数据流分区与异步执行调度
• SPIR-V / LLVM / WebGPU 后端：跨硬件编译

IREE 的关键创新是将"调度（scheduling）"与"执行（execution）"分离：Stream Dialect 显式标记张量数据的生产-消费者依赖，然后由 Partition Pass 将数据流分配到不同计算单元（CPU/GPU/NPU）。

5.2 TVM Ansor 与 Auto-Scheduling

虽然 TVM 使用自己的 TIR IR，但其自动调度器 Ansor 的设计理念与 MLIR 的 Transform Dialect 高度一致：

• 搜索空间设计：定义循环展开、向量化、分块的参数空间
• 代价模型：基于 XGBoost 的性能预测模型
• 进化搜索：使用遗传算法在程序变换空间中搜索最优配置

MLIR 的 Transform Dialect 将这种搜索过程本身表达为 IR 变换，使得 Auto-Scheduling 可以复用编译器的 pass 基础设施。

5.3 Triton：Python DSL 到 GPU 代码的桥梁

OpenAI Triton 使用 MLIR 构建的张量级编译器，允许用户用 Python DSL 编写自定义 GPU kernel：

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, N, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < N
    x = tl.load(x_ptr + offs, mask=mask)
    y = tl.load(y_ptr + offs, mask=mask)
    tl.store(out_ptr + offs, x + y, mask=mask)

Triton 编译器通过 MLIR 的 Triton Dialect 将 Python AST 转换为 SSA IR，然后经过 Triton → LLVM → PTX 的 lowering 管线生成 GPU 代码，避免了手写 CUDA 的繁琐。

六、工程实践建议与权衡

6.1 何时使用 MLIR

MLIR 适合以下场景：

1. 多目标编译：需要将同一 ML 模型部署到 CPU、GPU、NPU 等多种硬件
2. 领域特定优化：张量运算、量子计算、密码学等需要保留高层语义的领域
3. 硬件协同设计：新型加速器需要从指令集到编译器栈的垂直整合
4. 程序分析工具：需要在不同抽象层级进行静态分析或安全审计

6.2 学习曲线与团队配置

MLIR 的学习曲线主要来自：

• MLIR 概念密度：Dialect、Operation、Attribute、Type、Region、Pass、Pattern、Rewriter、Legalization 等术语需要较长时间消化
• C++ 要求：虽然 MLIR 提供 Python Bindings，但生产级 Dialect 仍需 C++ 实现
• TableGen 语法：MLIR 大量使用 TableGen 自动生成代码，是额外的学习成本

建议团队配置：至少 1 名 LLVM/MLIR 专家 + 2-3 名领域专家。初期投入回报周期约 3-6 个月。

6.3 与直接写 LLVM IR 的取舍

对于简单场景（如单一硬件平台、少量算子），直接使用 LLVM IR 可能更快。MLIR 的价值在以下情况才显现：

• 需要跨多个抽象层级共享优化逻辑
• 需要为不同硬件编写可复用的 lowering pass
• 需要与多种前端框架对接
• 需要进行复杂的程序变换（如微分、并行化、内存规划）

七、总结与展望

MLIR 代表了编译器基础设施从"单一 IR 统治一切"到"可扩展多层 IR 协作"的范式转变。它的核心价值不在于某一种特定的 IR 格式，而在于提供了一套可扩展的类型系统、转换框架和 pass 管理基础设施，让编译器开发者能够聚焦于领域语义的表达和优化。

随着 AI 芯片异构化趋势加速（CPU + GPU + NPU + DPU 混合部署），MLIR 作为"编译器中的 Linux"，其生态地位将进一步巩固。未来的方向包括：

• MLIR 与 ML 框架的深度集成：PyTorch 2.0 的 torch.compile 已经开始探索 MLIR 后端
• 安全关键领域的应用：IEC 61508 标准要求编译器的可验证性，MLIR 的形式化语义基础有望满足此需求
• 与 WebAssembly 的融合：将 MLIR IR 编译为 WASM 字节码，实现跨平台的 AI 推理部署

对于系统工程师而言，掌握 MLIR 不仅意味着理解一种新的编译器框架，更意味着拥有一种在硬件抽象与算法语义之间架设桥梁的能力——这在 AI 硬件加速日益重要的今天，正变得愈发关键。

参考资源： - MLIR 官方文档 - MLIR GitHub Repository - IREE 架构文档 - Triton: An Intermediate Language and Compiler for Tiled Neural Network Computation - Chris Lattner et al., "MLIR: A Compiler Infrastructure for the End of Moore's Law"