006、MLIR中的Dialect（方言）体系详解

006、MLIR中的Dialect（方言）体系详解

从一次凌晨三点的崩溃说起

去年冬天，我在调试一个自定义硬件加速器的编译流程。MLIR生成的IR在某个pass之后突然“变形”——原本清晰的linalg.generic操作变成了满屏的scf.for嵌套，中间还夹杂着几个我从未见过的arith.addi。更诡异的是，同样的IR在另一台机器上跑得好好的。折腾到凌晨三点，终于发现是Dialect注册顺序的问题：某个自定义Dialect在arith之前被加载，导致arith的operation被错误地解析成了我的自定义操作。

这个教训让我意识到：Dialect不是MLIR的“插件”，而是它的骨架。不理解Dialect体系，你写的每一个pass都是在雷区蹦迪。

Dialect到底是什么？别被“方言”这个翻译骗了

很多人把Dialect理解成“一种编程语言的方言”，比如C++方言、Python方言。但在MLIR里，Dialect更像是一个操作符的命名空间，外加一套类型系统和行为约定。

想象你正在设计一个芯片的指令集。你有算术指令、访存指令、控制流指令。在MLIR里，这些指令不会混在一起——arith Dialect管加减乘除，memref Dialect管内存操作，scf Dialect管循环和条件分支。每个Dialect就是一个独立的“指令集模块”，有自己的操作符、类型、甚至语法规则。

关键点：Dialect之间可以互相引用，但不能互相依赖对方的内部实现。这就像你的芯片里ALU单元和Cache单元可以互相传数据，但ALU不需要知道Cache的替换策略。

从代码层面看Dialect的“三件套”

每个Dialect在C++层面由三部分组成，缺一不可。我见过太多新手只注册了Operation就以为完事了，结果在IR解析时疯狂报错。

1. Dialect类本身

// 别这样写：把Dialect定义和Operation定义混在一个文件里
// 正确做法：Dialect类单独一个文件，Operation按功能拆分

class MyCustomDialect : public mlir::Dialect {
public:
explicit MyCustomDialect(mlir::MLIRContext *ctx)
: mlir::Dialect(getDialectNamespace(), ctx,
TypeID::get<MyCustomDialect>()) {
// 这里踩过坑：必须按依赖顺序添加Operation和Type
addOperations<
#define GET_OP_LIST
#include "MyCustomOps.cpp.inc"
>();
addTypes<
#define GET_TYPEDEF_LIST
#include "MyCustomTypes.cpp.inc"
>();
}

// 这个静态方法必须实现，MLIR用它来查找Dialect
static llvm::StringRef getDialectNamespace() { return "my_custom"; }
};

注意那个TypeID::get<MyCustomDialect>()——这是MLIR用来区分不同Dialect的“身份证”。如果你不小心把两个Dialect的TypeID写成了同一个，恭喜你，你会在某个深夜遇到“Operation不属于任何已知Dialect”的诡异错误。

2. Operation定义（TableGen才是王道）

手写Operation定义？那是2019年的做法。现在都用TableGen（.td文件），自动生成C++代码。

// MyCustomOps.td
def MyCustom_Dialect : Dialect {
let name = "my_custom";
let summary = "My custom hardware accelerator dialect";
let description = [{
This dialect models operations for a custom AI accelerator.
All operations operate on 16-bit fixed-point data.
}];
let cppNamespace = "::my_custom";
}

// 定义一个加法操作
def AddOp : MyCustom_Op<"add"> {
let summary = "Fixed-point addition";
let arguments = (ins
FixedPoint16:$lhs, // 这里踩过坑：类型必须提前注册
FixedPoint16:$rhs
);
let results = (outs FixedPoint16:$result);

// 别这样写：把验证逻辑写在C++里，然后用TableGen的verifier字段
// 正确做法：用TableGen的let verifier = ? 或者写C++验证函数
let hasVerifier = 1;
}

这里有个血泪教训：TableGen里定义的Operation，其参数类型必须在同一个Dialect中注册过。如果你用了memref的类型，但你的Dialect没有依赖memref，编译能过，运行时直接崩溃。

3. 类型系统

Dialect可以定义自己的类型。比如我的自定义Dialect需要FixedPoint16类型：

def FixedPoint16 : MyCustom_Type<"FixedPoint16"> {
let summary = "16-bit fixed-point number";
let description = [{
Q7.8 format: 1 sign bit, 7 integer bits, 8 fractional bits
}];
let parameters = (ins
"int":$integerBits, // 这里踩过坑：参数类型必须是MLIR支持的
"int":$fractionalBits
);
let assemblyFormat = "`<` $integerBits `,` $fractionalBits `>`";
}

注意那个assemblyFormat——它决定了你的类型在IR里长什么样。比如fixed_point<7, 8>。如果你不定义这个，MLIR会用默认的打印格式，那玩意儿可读性极差。

Dialect之间的“爱恨情仇”：依赖、转换与降级

依赖关系：谁依赖谁，谁不能依赖谁

MLIR的Dialect依赖关系是有向无环图。你不能让Dialect A依赖Dialect B，同时Dialect B又依赖Dialect A——MLIR会直接拒绝加载。

实际项目中常见的依赖链：

tosa -> linalg -> scf + arith -> llvm

tosa（Tensor Operator Set Architecture）依赖linalg（线性代数）做算子分解，linalg依赖scf（结构化控制流）和arith（算术）做循环展开和标量计算，最后全部降级到llvm Dialect。

经验之谈：设计自定义Dialect时，尽量让它只依赖builtin和arith这两个最基础的Dialect。依赖越多，后续维护越痛苦。

转换（Conversion）：从高级Dialect到低级Dialect

转换是MLIR的核心能力。比如把tosa.add转换成linalg.generic，再转换成scf.for+arith.addi。

// 一个简单的转换模式：把自定义AddOp转换成arith.addi
struct AddOpConversion : public mlir::OpRewritePattern<my_custom::AddOp> {
using OpRewritePattern::OpRewritePattern;

mlir::LogicalResult matchAndRewrite(
my_custom::AddOp op, mlir::PatternRewriter &rewriter) const override {

// 这里踩过坑：必须先检查操作数类型是否匹配
auto lhs = op.getLhs();
auto rhs = op.getRhs();

// 别这样写：直接创建arith.addi而不检查类型
// 正确做法：先转换类型
auto newLhs = rewriter.create<arith::ExtSIOp>(
op.getLoc(), rewriter.getIntegerType(32), lhs);
auto newRhs = rewriter.create<arith::ExtSIOp>(
op.getLoc(), rewriter.getIntegerType(32), rhs);

auto addOp = rewriter.create<arith::AddIOp>(
op.getLoc(), newLhs, newRhs);

// 把结果截断回16位
auto truncOp = rewriter.create<arith::TruncIOp>(
op.getLoc(), rewriter.getIntegerType(16), addOp);

rewriter.replaceOp(op, truncOp.getResult());
return mlir::success();
}
};

这个转换看起来简单，但实际项目中80%的bug都出在类型转换上。永远不要假设操作数的类型和目标Dialect的类型完全匹配。

降级（Lowering）：从抽象到具体

降级是转换的一种特殊形式——把高级语义的操作变成低级语义的操作组合。比如把linalg.generic降级成scf.for循环。

降级过程中最头疼的是中间状态的合法性。你可能会遇到这种情况：降级到一半，IR里同时存在linalg和scf的操作，但某些pass只认识linalg，某些pass只认识scf，结果IR卡在中间状态无法继续。

我的做法：每次降级只做一件事，降级完成后立即运行验证（mlir-opt –verify-each），确保IR在任何时刻都是合法的。

实战：设计一个自定义Dialect的完整流程

假设我们要为一个支持“向量化加法”的硬件设计Dialect。这个硬件可以一次性对两个128位的向量做加法，每个向量包含8个16位整数。

第一步：定义Dialect和Operation

// VectorAdd.td
def VectorAdd_Dialect : Dialect {
let name = "vector_add";
let cppNamespace = "::vector_add";
}

def VecAddOp : VectorAdd_Op<"vec_add"> {
let summary = "Vectorized addition for 8×16-bit integers";
let arguments = (ins
VectorOfLengthAndType<[8], [I16]>:$lhs,
VectorOfLengthAndType<[8], [I16]>:$rhs
);
let results = (outs
VectorOfLengthAndType<[8], [I16]>:$result
);
let assemblyFormat = "$lhs `,` $rhs attr-dict";
}

第二步：实现验证逻辑

mlir::LogicalResult VecAddOp::verify() {
// 这里踩过坑：必须检查向量长度和元素类型
auto lhsType = getLhs().getType().dyn_cast<mlir::VectorType>();
auto rhsType = getRhs().getType().dyn_cast<mlir::VectorType>();

if (!lhsType || !rhsType) {
return emitOpError("operands must be vector types");
}

if (lhsType.getShape()[0] != 8 || rhsType.getShape()[0] != 8) {
return emitOpError("vectors must have exactly 8 elements");
}

if (!lhsType.getElementType().isInteger(16) ||
!rhsType.getElementType().isInteger(16)) {
return emitOpError("element type must be i16");
}

return mlir::success();
}

第三步：实现降级到arith

struct VecAddOpLowering : public mlir::OpRewritePattern<VecAddOp> {
using OpRewritePattern::OpRewritePattern;

mlir::LogicalResult matchAndRewrite(
VecAddOp op, mlir::PatternRewriter &rewriter) const override {

auto loc = op.getLoc();
auto lhs = op.getLhs();
auto rhs = op.getRhs();

// 别这样写：直接展开成8个独立的arith.addi
// 正确做法：用scf.for循环，保持代码简洁
auto zero = rewriter.create<arith::ConstantOp>(
loc, rewriter.getIntegerAttr(rewriter.getI64Type(), 0));
auto eight = rewriter.create<arith::ConstantOp>(
loc, rewriter.getIntegerAttr(rewriter.getI64Type(), 8));
auto one = rewriter.create<arith::ConstantOp>(
loc, rewriter.getIntegerAttr(rewriter.getI64Type(), 1));

// 创建结果向量
auto resultType = op.getResult().getType();
auto result = rewriter.create<scf::ForOp>(
loc, zero, eight, one, llvm::SmallVector<mlir::Value>{},
[&](mlir::OpBuilder &b, mlir::Location loc, mlir::Value iv,
mlir::ValueRange args) {
// 提取元素
auto lhsElem = b.create<vector::ExtractElementOp>(
loc, lhs, iv);
auto rhsElem = b.create<vector::ExtractElementOp>(
loc, rhs, iv);
// 加法
auto sum = b.create<arith::AddIOp>(loc, lhsElem, rhsElem);
// 插入结果向量
auto updated = b.create<vector::InsertElementOp>(
loc, sum, result, iv);
b.create<scf::YieldOp>(loc, updated);
});

rewriter.replaceOp(op, result.getResult(0));
return mlir::success();
}
};

个人经验：Dialect设计的五个“不要”

写在最后

Dialect体系是MLIR最强大的特性，也是最容易踩坑的地方。每次我遇到奇怪的IR解析错误，第一反应就是检查Dialect注册顺序和依赖关系。记住：MLIR的IR是Dialect的IR，不是你的IR。尊重Dialect的边界，你的编译流程才会稳定。

下次遇到“Operation not registered”的错误，别急着怀疑MLIR的bug——先检查你的Dialect注册代码。大概率是那里出了问题。