014、MLIR的ODS（Operation Definition Specification）详解

MLIR的ODS（Operation Definition Specification）详解从一个让我熬夜到凌晨三点的bug说起去年做AI加速器编译器的时候，我定义了一组自定义算子。按照MLIR官方教程写ODS，编译通过，跑起来也正常。直到某天，一个同事在优化pass里用PatternRewriter替换算子时，程序直接崩了——段错误。查了两天，最后发现是ODS里一个let声明顺序的问题。从那以后，我养成了一个习惯：每次写完ODS，先手动检查生成的C++代码，再跑测试。这个经历让我意识到，ODS不是简单的“声明式定义”，它背后生成的代码逻辑、约束检查、与C++模板的交互，处处是坑。今天这篇笔记，就把我踩过的坑和总结的经验摊开来讲。ODS到底是什么？别被名字唬住ODS全称Operation Definition Specification，本质是一套基于TableGen的声明式语言。TableGen是LLVM项目自带的代码生成工具，MLIR用它来定义Operation、Type、Attribute、Dialect等核心组件。你可能会问：为什么不用C++直接写？因为MLIR的Operation数量动辄成百上千，手写C++类、接口、builder、parser、printer、verifier，重复劳动不说，还容易出错。ODS让你用声明式语法描述“这个操作长什么样”，TableGen自动生成对应的C++代码。但注意，ODS不是万能的。它擅长定义“结构固定”的操作，如果你的操作逻辑极其复杂（比如动态shape推理、自定义约束），ODS生成的代码可能不够用，需要手动写C++扩展。ODS文件的基本骨架一个典型的ODS文件长这样（以MyDialect.td为例）：// 引入MLIR基础定义 include "mlir/IR/OpBase.td" // 定义Dialect def My_Dialect : Dialect { let name = "my"; let cppNamespace = "::my"; let summary = "My custom dialect for demo"; let description = [{ This dialect defines operations for ... }]; } // 定义Operation def My_AddOp : My_Op"add" { let summary = "Addition operation"; let description = [{ Performs element-wise addition of two tensors. }]; let arguments = (ins TensorOf[F32]:$lhs, // 左操作数，F32张量 TensorOf[F32]:$rhs // 右操作数，F32张量 ); let results = (outs TensorOf[F32]:$result ); let assemblyFormat = "$lhs `,` $rhs attr-dict"; }这里有几个关键点：My_Op"add"：My_Op是一个在Dialect定义中生成的基类，"add"是操作的名字，最终会生成mlir::my::AddOp这个C++类。arguments和results：分别定义输入和输出。注意ins和outs后面的括号里，每个参数用类型:$名字的格式。assemblyFormat：控制打印和解析的格式。这里$lhs,$rhs表示用逗号分隔两个操作数，attr-dict表示打印属性字典。踩坑提醒：assemblyFormat里的变量名必须和arguments/results里定义的名字完全一致，大小写敏感。我见过有人把$lhs写成$LHS，编译不报错，但解析时永远失败。参数类型：不只是TensorOfODS支持丰富的类型约束，远不止TensorOf。常用的有：AnyTensor：任意张量类型TensorOf[F32, F16]：指定元素类型的张量RankedTensorOf[F32]：有rank的张量（即不是动态rank）St