)
筑基 MLIR深入理解常驻幕后的 ODSOperation Definition Specification在深入 MLIRMulti-Level Intermediate Representation的编译世界时我们频繁提及各种强大的方言Dialect如高层线性代数的MHLO/StableHLO、硬件抽象层的TOSA以及底层的Linalg。但你有没有想过这些方言中成百上千的算子Operations简称 Op究竟是如何被定义出来的如果手头有一款自研的专用 AI 芯片NPU想要为它定制一套专属的方言算子难道需要纯手工去写几万行繁琐、重复且容易出错的 C 类结构、属性校验和 C 接口吗答案是否定的。MLIR 生态引入了一套强悍的元编程基础设施——ODSOperation Definition Specification算子定义规范。ODS 才是 MLIR 拥有无限扩展能力的“幕后造物主”。本文将为你深度拆解 ODS 的设计哲学、核心语法、工作原理以及它是如何将编译器开发从“体力活”解放为“艺术”的。1. 什么是 ODS现代编译器的元编程ODSOperation Definition Specification是 MLIR 中基于TableGen语法的一套声明式Declarative算子定义规范。TableGen 是一种什么工具TableGen (.td文件) 是 LLVM 生度系统中久经沙场的领域特定语言DSL。它的核心思想是用一种结构化的声明方式来描述数据然后通过特定的后端生成器Backend自动把这些描述转化为成千上万行的 C 源代码。在传统编译器中新加一个算子开发者需要手动编写C 类定义与构造函数输入、输出、属性的验证逻辑Verifier防止非法输入算子的解析器Parser与打印器Printer用于支持文本 IR 的读写各种辅助接口Traits / Interfaces的挂载而在 MLIR 中你只需要在.td文件里用 ODS 语法写下一段几十行的声明MLIR 的mlir-tblgen编译器就会在编译期自动帮你喷涌出所有的 C 样板代码Boilerplate Code。2. ODS 的核心设计哲学单源真理Single Source of TruthODS 的核心设计哲学非常明确关于一个算子的所有事实应当且只应当在一处被定义。这种单源真理的设计带来了编译工程上的三大颠覆式优势极低的维护成本如果一个算子的输入类型从“只能是张量”扩展到“可以是张量或标量”你不需要去改 C 的校验函数、不需要去改文本解析逻辑只需要在 ODS 里修改一行类型声明底层几十处关联代码自动同步更新。绝对的行为一致性自动生成的验证器Verifier能确保进入编译管线的每一个算子都百分之百合规从源头上杜绝了因为手写 C 校验逻辑不严谨而导致的编译器崩溃Segmentation Fault。极强的自文档化Self-documentingODS 强制要求在定义算子时附带summary和description。这些文本不仅留在代码里还会被 TableGen 自动导出为精美的 Markdown 开发文档。方言的文档与源码永远保持最新绝不脱节。3. 深度剖析一个标准 ODS 算子定义的拆解我们来看一个典型的 ODS 算子定义案例。假设我们在定制一个属于自己芯片的方言my_npu现在需要定义一个定制的矩阵乘法算子MatMulOp// 定义属于 my_npu 方言的基类 class MyNPU_Opstring mnemonic, listTrait traits [] : OpMyNPU_Dialect, mnemonic, traits; // 正式定义 MatMulOp def MyNPU_MatMulOp : MyNPU_Opmatmul, [Pure, SameOperandsAndResultType] { let summary NPU 专属矩阵乘法算子; let description [{ 该算子在 NPU 的矩阵计算阵列Tensor Core上执行高效的二维矩阵乘法。 它要求两个输入张量的维度必须满足紧密乘积条件且输出与输入的数据类型保持一致。 }]; // 1. 定义算子的输入Arguments包括操作数Operands和属性Attributes let arguments (ins AnyTensor:$lhs, AnyTensor:$rhs, DefaultValuedAttrBoolAttr, false:$transpose_lhs ); // 2. 定义算子的输出/结果Results let results (outs AnyTensor:$output ); // 3. 自定义文本 IR 的打印与解析行为 let assemblyFormat [{ $lhs , $rhs attr-dict : type($output) }]; // 4. 挂载额外的 C 自定义验证逻辑可选 let hasVerifier 1; }四大核心构建块Building Blocks深度解构方言与特征挂载Dialect TraitsMyNPU_Opmatmul, [Pure, SameOperandsAndResultType]声明了算子在 IR 中的完整字面量为my_npu.matmul。同时挂载了两个强大的Trait特征Pure告诉编译器这个算子没有副作用可以放心被死代码消除DCEPass 优化失掉SameOperandsAndResultType则是一种编译约束规定输入输出的元素类型必须完全相同。输入Arguments包含了操作数$lhs/$rhs在运行时流动的数据流以及属性$transpose_lhs在编译期就确定的静态配置如是否转置。ODS 允许你使用AnyTensor、F32Tensor等极丰富的内置约束甚至可以用DefaultValuedAttr直接指定默认值。输出Results指定了计算完成后产出的张量数量和类型。配合 SSA静态单赋值形态MLIR 会在底层为其生成完美的数据流拓扑追踪追踪。装配格式Assembly FormatassemblyFormat是 ODS 的神来之笔。它通过一种极简的声明式 DSL直接定义了该算子在.mlir文本文件里的长相。上文定义的算子打印出来就是%res my_npu.matmul %matrixA, %matrixB {transpose_lhs true} : tensor128x128xf32编译器会自动据此生成相应的 Parser 和 Printer再也不用手写繁琐的字符串解析。4. ODS 的工作流水线从.td到.h/.inc了解 ODS 的运行机制能帮你更好地在工程中配置 CMake 和排查编译错误。当你按下编译快捷键时构建系统在幕后完成了以下流水线[开发者编写的 op.td] ───► (输入) │ ▼ [ mlir-tblgen ] (编译期执行的元编译器) │ ┌──────────────┴──────────────┐ ▼ (指定 -gen-op-decls) ▼ (指定 -gen-op-defs) [ op.h.inc ] [ op.cpp.inc ] (包含 C 类声明、派生类、) (包含 C 函数实现、验证器、) (Getter 访问器等) (自动解析/打印逻辑等) │ │ └──────────────┬──────────────┘ ▼ (由 CMake 隐式 #include 嵌入) [ 最终的编译产物 Dialect ]在你的 C 源码中你只需要通过简单的宏控就能将这些生成的代码无缝注入到你的方言命名空间中#defineGET_OP_CLASSES#includeMyNPU/MyNPUOps.h.inc这种设计让编译器的核心业务逻辑代码保持得极度干净C 文件里只存放你最核心的图优化、降级Lowering算法所有的账目式结构定义通通交由 ODS 托管。5. ODS 与上层AST / 框架及下层硬件的协同在整个大模型算力基础设施AI Infra中ODS 实际上定义了信息流经编译器时的“骨架”。对接上层如 Python AST / Torch当 Torch-MLIR 等前端解析了 Python 端的模型 AST 后它会调用 ODS 自动生成的 COpBuilder接口。由于 ODS 定义了严格的输入类型任何在前端暴露的参数错误都会在进门的第一时间被 ODS 的 Verifier 捕获并拦截。服务中层优化Passes诸如算子融合Operator Fusion、循环平铺Tiling等 Pass高度依赖 ODS 生成的 Trait例如前文提到的Pure或IsolatedFromAbove。Pass 不需要知道my_npu.matmul的底层细节只需要探查它是否携带这些 ODS 特征即可安全地进行全局图变换。驱动底层硬件NPU / GPUODS 定义的属性Attributes最终会伴随算子一路走到最底层的硬件发射期。这些属性将被无缝翻译成 NPU 的 DMA 传输配置、或者是 GPU Kernel 的 Grid/Block 线程维度确保人类的语义被丝毫不差地灌注进晶体管中。总结在 MLIR 生态中如果把各类方言和基础设施比作长城那么ODS 就是铸造每一块标准规格砖瓦的模具。它用 TableGen 的优雅声明消除了传统编译器开发中长久存在的样板代码噩梦使得大模型时代的编译器专家能将 100% 的精力聚焦在算子演进与图级优化本身。