mlir

用途

本技能用于指导智能体使用 MLIR（Multi-Level IR，多级中间表示），涵盖算子（ops）、区域（regions）、块（blocks）与值（values）；内置方言（arith、func、memref、affine、linalg）；使用 ODS 编写自定义方言；使用 ConversionPattern 编写降级 pass；mlir-opt 命令行工具；以及 ML 编译器用例（Torch-MLIR、IREE）。

适用场景

构建领域特定编译器 IR（图形、ML、硬件 DSL）
将高层算子降级到 LLVM 或 GPU 方言
编写渐进式降级流水线（linalg → loops → LLVM）
集成 IREE 或 Torch-MLIR 进行 ML 部署
创建跨方言的可复用转换 pass
在合适的抽象层级上原型化编译器优化

MLIR 结构

Module
└── func.func @main()
    └── region
        └── block ^bb0:
            └── operations (ops) producing SSA values

关键概念：

Operation — 类似指令的节点（arith.addi、memref.load）
Region — 块的容器（函数、控制流）
Block — 包含有序算子的 CFG 节点
Value — 算子或块参数的 SSA 结果

内置方言

方言	用途
arith	整数/浮点算术
func	函数定义与调用
memref	带形状/步长的缓冲区抽象
affine	仿射循环嵌套、map/set 约束
linalg	结构化线性代数算子
scf	结构化控制流（for、if）
llvm	用于最终降级的 LLVM IR 方言
gpu	GPU 内核启动

示例：

// example.mlir
func.func @add(%a: memref<4xf32>, %b: memref<4xf32>, %c: memref<4xf32>) {
  %c0 = arith.constant 0 : index
  %c4 = arith.constant 4 : index
  scf.for %i = %c0 to %c4 step %c1 {
    %av = memref.load %a[%i] : memref<4xf32>
    %bv = memref.load %b[%i] : memref<4xf32>
    %sum = arith.addf %av, %bv : f32
    memref.store %sum, %c[%i] : memref<4xf32>
  }
  return
}

mlir-opt CLI

# Parse and print
mlir-opt example.mlir

# Run canonicalization
mlir-opt example.mlir -canonicalize

# Lower affine to scf
mlir-opt affine.mlir -lower-affine

# Full pipeline toward LLVM
mlir-opt input.mlir \
  --linalg-bufferize \
  --convert-linalg-to-loops \
  --convert-scf-to-cf \
  --convert-arith-to-llvm \
  --convert-memref-to-llvm \
  --convert-func-to-llvm \
  -o llvm.mlir

ODS — 算子定义规范

// MyOps.td
include "mlir/IR/OpBase.td"

def My_Dialect : Dialect {
    let name = "my";
    let summary = "My custom dialect";
}

class My_Op<string mnemonic, list<Trait> traits = []> :
    Op<My_Dialect, mnemonic, traits>;

def AddOp : My_Op<"add", [Pure]> {
    let summary = "Add two values";
    let arguments = (ins AnyType:$lhs, AnyType:$rhs);
    let results = (outs AnyType:$result);
    let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($result)";
}

从 TableGen 生成 C++：

mlir-tblgen -gen-op-defs MyOps.td -I include/ -o MyOps.cpp.inc

自定义方言 C++ 实现

#include "mlir/IR/DialectImplementation.h"
#include "MyDialect.h"

#include "MyOps.cpp.inc"

void MyDialect::initialize() {
    addOperations<
#define GET_OP_LIST
#include "MyOps.cpp.inc"
    >();
}

#define GET_OP_CLASSES
#include "MyOps.cpp.inc"

降级 pass

#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
#include "mlir/Transforms/DialectConversion.h"

struct AddOpLowering : OpConversionPattern<my::AddOp> {
    using OpConversionPattern::OpConversionPattern;

    LogicalResult matchAndRewrite(my::AddOp op, OpAdaptor adaptor,
                                  ConversionPatternRewriter &rewriter) const override {
        rewriter.replaceOpWithNewOp<arith::AddIOp>(op, adaptor.getLhs(), adaptor.getRhs());
        return success();
    }
};

void populateLoweringPatterns(RewritePatternSet &patterns) {
    patterns.add<AddOpLowering>(patterns.getContext());
}

// In pass:
mlir::ConversionTarget target(*context);
target.addIllegalDialect<my::MyDialect>();
target.addLegalDialect<arith::ArithDialect>();

if (failed(applyPartialConversion(module, target, std::move(patterns))))
    signalPassFailure();

用于 ML 编译器的 linalg

%0 = linalg.matmul ins(%A, %B : tensor<128x256xf32>, tensor<256x64xf32>)
                   outs(%C : tensor<128x64xf32>) -> tensor<128x64xf32>

降级路径：linalg → scf loops → affine → llvm

Torch-MLIR 与 IREE

# Torch-MLIR: PyTorch → MLIR
python -m torch_mlir.tools.import-onnx --onnx-model model.onnx -o model.mlir

# IREE: MLIR → GPU/CPU executable
iree-compile --iree-hal-target-backends=llvm-cpu model.mlir -o model.vmfb
iree-run-module --module=model.vmfb --function=main

常见问题

症状	原因	修复
方言未注册	缺少 registerDialect	在工具/pass 初始化中注册
ODS 构建失败	TableGen 包含路径	检查 mlir/IR/OpBase.td 的 -I 路径
降级不完整	残留非法算子	使用 --mlir-print-ir-after-failure 调试
模式中类型不匹配	adaptor 类型错误	使用 OpAdaptor 类型化访问器
mlir-opt 崩溃	无效 IR	运行 -verify-each
降级后函数为空	所有算子非法且未转换	添加缺失的模式

技能介绍

用途