第 3 课：ODS/TableGen 与自定义 Operation

本课目标与完成标准

学完后你应能：用 ODS 声明 op 的 operands/results/traits，并解释生成代码与手写语义钩子的边界。

通过必须同时满足：

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

前置关系

  • 课程前置：编译原理基础、C++ 阅读能力
  • 本课在路线中的作用：ODS 用 TableGen 描述 dialect/op 的结构约束并生成 C++ 声明、parser/printer 框架和文档，减少重复样板。

核心心智模型

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

在 C++ 中手动编写一个工业级 MLIR Operation 需要耗费数以百行的样板代码（Boilerplate）：手写类定义、构造函数、访问器（Accessors）、Operand/Result 数量校验、Parser/Printer 文本解析序列化、类型检查 Verifier 以及文档。这极其冗长且极易出错。 为了消除样板代码并实现单一事实来源（Single Source of Truth），MLIR 提供了基于 LLVM TableGen 语言的领域特定元编程框架——ODS（Operation Definition Specification）。 开发者在声明式的 .td 文件中定义算子的输入、输出、属性、Traits 与文本语法格式，由专用代码生成工具 mlir-tblgen 在编译期全自动生成强类型的 C++ 声明与实现。 本课解决的核心工程问题是：掌握基于 ODS 的声明式算子建模生命周期，理解 Trait 的物理语义承诺，以及划分编译期代码生成（ODS）与手动 C++ 语义钩子（Verifier / Fold / Canonicalize）的清晰职责边界。

2. 它如何工作

ODS 声明规范与全自动构建生命周期如图 3-1 与图 3-2 所示：

图 3-1：ODS 声明式算子规范 (Operation Definition Specification) 元模型

图 3-2：基于 mlir-tblgen 的自定义算子全自动构建与接入生命周期

  1. 元模型定义（.td 描述）： 在 TableGen 脚本中，继承基础模板定义算子：
    • arguments：使用 (ins ...) 声明输入操作数与属性；
    • results：使用 (outs ...) 声明返回值；
    • traits：挂载编译期语义特征（如 Pure 表示无副作用、Commutative 表示可交换、SameOperandsAndResultType 表示输入输出类型同构）；
    • assemblyFormat：声明声明式自定义打印与解析规则，免写手写 Parser。
  2. 代码生成（mlir-tblgen 驱动）： 如图 3-2 所示，在项目 CMake 构建阶段：
    • 运行 mlir-tblgen -gen-op-decls 生成 *.h.inc（包含强类型 C++ 包装类、getLhs() / getRhs() 访问器、Trait 继承声明）；
    • 运行 mlir-tblgen -gen-op-defs 生成 *.cpp.inc（包含操作数偏移行内解析、自动类型校验骨架、build() 工厂方法与 Parser/Printer 逻辑）。
  3. C++ 语义钩子接入： 对于超越声明式表达能力的复杂业务逻辑（如多维张量沿某一轴求和的形状推导、常量折叠），开发者在自定义 C++ 文件中实现声明的语义钩子函数：
    • verify()：执行复杂的深层跨字段数值校验；
    • fold()：执行编译期常量计算（Constant Folding）；
    • getCanonicalizationPatterns()：注册代数化简等价重写规则。

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

  • Trait 是严肃的语义承诺，而非非绑定的注释： Trait（如 Pure、Commutative、SameOperandsAndResultType）会直接被下游优化 Pass（如 DCE、CSE、规范化重写）作为无条件成立的前提公理使用。若给实际向显存写入数据或改变全局状态的 Op（如 memref.store）错误标记了 Pure，下游 DCE Pass 会在发现其返回值未被使用时直接将其整条物理消除，导致严重的状态静默损坏。

  • 类型系统与 ODS Constraint 匹配： 若在 arguments 声明了 AnyTensor:$lhs 和 AnyTensor:$rhs，但未在 Trait 中约束类型相等，则 $lhs 可以是 tensor<4xf32> 而 $rhs 可以是 tensor<8xi32>。若要确保逐元素相加合法，必须同时施加 SameOperandsAndResultType 或显式的 AllTypesMatch 约束。

  • 生成代码的包含方式： 生成的 *.h.inc 已按 dialect 的 cppNamespace 自带命名空间，应在全局作用域、先 #define GET_OP_CLASSES 再包含；包含前还要先包含 dialect 头文件和 op 用到的 interface/trait 头文件，否则会出现类型未声明的编译错误。

4. 性能与工程取舍

  • 声明式约束 vs 手写 C++ Verifier： 能用 ODS 声明式表达的约束（如 SameTypeOperands、TypeConstraint）尽量在 .td 中声明。ODS 能自动生成直观易读的报错提示，并能自动导出至 Markdown 语法手册；过度依赖手写 C++ verify() 会降低算子文档的自解释性。

  • 自动 AssemblyFormat vs 手写 Parser： 声明式 assemblyFormat 能覆盖大多数自定义语法格式，并能自动保证 Parser 与 Printer 的完全对称可逆（Round-trip Safety）；仅当涉及极特殊的复杂动态语法解析时才需要开启 let hasCustomAssemblyFormat = 1;。

具体演示

以二元加法算子 My_AddOp 为例：

def My_AddOp : My_Op<"add", [Pure, SameOperandsAndResultType]> {
  let summary = "elementwise add";
  let arguments = (ins AnyTensor:$lhs, AnyTensor:$rhs);
  let results = (outs AnyTensor:$result);
}
  • arguments 指定必须输入两个张量 $lhs 与 $rhs；
  • results 指定输出张量 $result；
  • [Pure] 承诺无任何内存副作用，供 DCE 安全消除；
  • [SameOperandsAndResultType] 自动生成校验代码，若用户传入 %a: tensor<2xf32> 与 %b: tensor<3xf32>，解析后的 Verifier 阶段会报错。

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

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

补齐 ODS 中两个输入和一个输出。

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

Code 
%%writefile /tmp/MyOps.td
def My_AddOp : My_Op<"add", [Pure, SameOperandsAndResultType]> {
  let summary = "elementwise add";
  let arguments = (ins AnyTensor:$lhs, ______:$rhs);
  let results = (outs ______:$result);
}

检查方法

有源码构建环境时交给 mlir-tblgen；否则检查约束是否足以拒绝标量和不同类型 operands。

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

Q1

不要背定义：请从输入、状态变化和输出三个阶段解释“ODS/TableGen 与自定义 Operation”的工作机制。

你的答案：

Q2

错误地给有副作用 op 标 Pure，DCE 可能做出什么错误？

你的答案：

Q3

shape 相等这种跨 operand 约束应写 trait、ODS constraint 还是 verifier？

你的答案：

评分与通过规则

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

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

参考答案（仅 answer 分支）

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

Code 
%%writefile /tmp/MyOps.td
def My_AddOp : My_Op<"add", [Pure, SameOperandsAndResultType]> {
  let summary = "elementwise add";
  let arguments = (ins AnyTensor:$lhs, AnyTensor:$rhs);
  let results = (outs AnyTensor:$result);
}

Q1 参考答案

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

  1. 输入阶段（声明式元模型定义）：开发者在 TableGen 脚本（.td 文件）中以极简 DSL 声明 Operation 的核心规格，包括输入操作数类型（arguments）、返回值类型（results）、编译期语义公理（traits，如 Pure）以及文本汇编语法（assemblyFormat）。

  2. 状态变化阶段（编译期代码生成与 C++ 虚实结合）：构建系统调用 mlir-tblgen 驱动前端解析 TableGen AST，自动生成包含静态类型安全访问器、build() 工厂重载、结构化校验逻辑和语法反序列化器的 C++ 源码片段（*.h.inc 与 *.cpp.inc）。开发者补充手写高级语义钩子（如 verify() 边界判断与 fold() 常量折叠），共同编译为本地二进制。

  3. 输出阶段（方言模块装配与注册）：在 Dialect 的 initialize() 中调用 addOperations<...>() 注册生成的 Op 类（模板参数列表通过 #define GET_OP_LIST 后包含 .cpp.inc 得到）；工具通过 DialectRegistry 加载该 Dialect 到 MLIRContext 后，为整个编译器优化管线提供强类型、可被验证的语法树节点。

Q2 参考答案

在 MLIR 优化体系中，Pure Trait（对应内部的 MemoryEffectOpInterface 且没有任何内存副作用）是编译器用于判定死代码消除（DCE: Dead Code Elimination）的核心公理： 若错误地将一个实际上具有物理副作用的操作（例如向全局显存写数据的 store 操作、打印日志的 print 操作、或引发硬件同步的 barrier 操作）标记为 Pure： 当该 Operation 的返回值没有被下游任何其它 Operation 消费（或者该 Operation 本身没有返回值，如单向写回）时： -canonicalize、-cse 以及贪心重写驱动中的 trivially-dead 清理会基于“没有结果被使用且无副作用的操作等价于无用计算”的强假设，直接将该 Operation 从控制流图中斩除删除。 这会导致预期的显存写入或设备交互完全没有被发射，产生极端隐蔽的静默逻辑丢失。

Q3 参考答案

在设计算子约束时，三者的选型与架构分层原则如下：

  1. 优先选择 Trait（如 SameOperandsShape / SameOperandsAndResultShape）：
    • 适用场景：当形状约束是编译器中广泛通用的公理模式时。
    • 优势：Trait 不仅在解析时自动生成校验逻辑，更重要的是为下游 Pass 提供了可查询的 C++ Trait 概念，能被优化算法批量识别（例如代数重写 Pass 可直接针对所有 SameOperandsShape 的算子执行跨节点广播下推）。
  2. 次选 ODS 声明式 Constraint（如在 arguments/results 中使用 Predicate 谓词）：
    • 适用场景：当约束模式较为特殊（例如输入维度必须与输出属性满足算术倍数），但仍能用纯逻辑公式表达时。
    • 优势：保持在 .td 中声明式维护，能自动生成到导出的算子文档中，维护成本低。
  3. 最后选择自定义 C++ verify() 钩子：
    • 适用场景：涉及动态上下文、高阶张量广播（Broadcasting）兼容性检查、或依赖外部硬件属性的极复杂条件逻辑。
    • 优势：拥有 C++ 图遍历的全部图灵完备能力，能构造极度精准详尽的局部诊断错误信息（op->emitOpError(...)）。

参考资料

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

← 上一课
第 2 课：Dialect、Type、Att… 下一课 →
第 4 课：Pass、PatternRewr…