系列文章目录章节1什么是MLIR章节2MLIR的引入章节3MLIR完成一个端到端的流程章节4编写一个pass章节5MLIR相关语法章节6量化相关章节7训练相关章节8图优化文章目录一、头文件怎么写完整的代码修改Writing our first pass by j2kun · Pull Request #4 · j2kun/mlir-tutorial · GitHub首先定义一个头文件// lib/Transform/Affine/AffineFullUnroll.h class AffineFullUnrollPass : public PassWrapperAffineFullUnrollPass, OperationPassmlir::func::FuncOp { private: void runOnOperation() override; // implemented in AffineFullUnroll.cpp StringRef getArgument() const final { return affine-full-unroll; } StringRef getDescription() const final { return Fully unroll all affine loops; } };这段代码定义了一个 MLIR Pass编译过程中的一个处理步骤名为AffineFullUnrollPass作用是将所有的 Affine 循环完全展开。class AffineFullUnrollPass定义了一个名为AffineFullUnrollPass的 C 类。PassWrapper...这是 MLIR 提供的一个模板基类用来“包裹”我们的 Pass。它使用了前面提到的CRTP奇异递归模板模式用来实现静态多态。第一个模板参数AffineFullUnrollPass把自己传给了基类。这样PassWrapper在编译期就知道了派生类长什么样从而能自动帮你实现clone()等样板方法不需要你手动写。第二个模板参数OperationPassmlir::func::FuncOp说明这是一个操作级 PassOperation Pass。mlir::func::FuncOp把这个 Pass锚定anchor到函数上。这意味着 Pass 管理器只会把这个 Pass 应用到func.func操作上每次处理一个函数。mlir::func::FuncOp是一个具体的 C 类类型Class Type更准确地说它是 MLIR 中用来表示func.func操作Operation的强类型包装类Wrapper Class。void runOnOperation() override;runOnOperation()这是 Pass 的核心业务逻辑入口。当 Pass 管理器把 Pass 应用到某个FuncOp时就会调用这个方法。override告诉编译器这是在重写基类OperationPass里的虚函数。方法体在哪里注释明确写了implemented in AffineFullUnroll.cpp。这里只做了声明具体“怎么展开循环”的代码在对应的.cpp文件里。暗藏玄机虽然方法体不在这里但在.cpp里实现时你可以直接调用getOperation()来获取当前正在处理的FuncOp。这正是OperationPassmlir::func::FuncOp提供给你的便利。StringRef getArgument() const finalMLIR 的mlir-opt工具需要通过命令行参数来指定运行哪个 Pass。这个方法返回的就是命令行里的参数名affine-full-unroll。这意味着你在终端里可以这样用mlir-opt --affine-full-unroll input.mlir。final禁止派生类继续重写这个方法。StringRef getDescription() const final当你运行mlir-opt --help时每个 Pass 都会有一行简短的描述。这里返回的Fully unroll all affine loops就是显示在帮助信息里的那句话。看过头文件说明我们需要实现的三个函数是runOnOperation执行 pass 逻辑的函数。getArgument在类似mlir-opt的工具中使用的命令行参数。getDescription在类似mlir-opt的工具上运行--help时显示的命令行描述。二、注册pass// tools/tutorial-opt.cpp #include lib/Transform/Affine/AffineFullUnroll.h #include mlir/include/mlir/InitAllDialects.h #include mlir/include/mlir/Pass/PassManager.h #include mlir/include/mlir/Pass/PassRegistry.h #include mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h int main(int argc, char **argv) { mlir::DialectRegistry registry; mlir::registerAllDialects(registry); mlir::PassRegistrationmlir::tutorial::AffineFullUnrollPass(); return mlir::asMainReturnCode( mlir::MlirOptMain(argc, argv, Tutorial Pass Driver, registry)); }DialectRegistry是一个“方言清单”。MLIR 在解析 IR 时需要知道文本里出现的arith.addi、func.func、affine.for等操作分别属于哪个方言才能正确构建 IR。registerAllDialects(registry)把所有内置方言一次性登记进去这样你的工具就能解析包含这些方言的.mlir文件。如果你只关心affine和func两个方言也可以只注册它们registry.insertmlir::func::FuncDialect, mlir::affine::AffineDialect()但注册全部方言更省事也和mlir-opt行为一致。把AffineFullUnrollPass这个类登记到 MLIR 的全局 Pass 注册表。注册时MLIR 会调用你 Pass 类里的getArgument()返回affine-full-unroll和getDescription()返回Fully unroll all affine loops。注册成功后命令行工具就能识别--affine-full-unroll这个参数并在--help里显示它的描述。这解释了为什么你的 Pass 类里需要实现getArgument()和getDescription()——它们就是给PassRegistration用的。MlirOptMain是一个高度封装的工具构造函数它帮你实现了几乎所有的命令行工具逻辑包括它自动处理的事情说明解析命令行参数-o、--split-input-file、--verify-diagnostics、--pass-pipeline等读取输入文件从argv中拿到.mlir文件路径或从 stdin 读解析 MLIR 文本用你传入的registry构建 IR构建 PassManager根据命令行传入的 Pass 名字如--affine-full-unroll从全局注册表里找到对应 Pass 并加入 PassManager运行 Pass按顺序执行所有指定的 Pass输出结果写到-o指定的文件或 stdout错误处理解析失败、Pass 失败时给出诊断信息参数含义argc, argv命令行参数。Tutorial Pass Driver工具的描述字符串会显示在--help开头。registry之前填充好的方言注册表告诉MlirOptMain支持哪些方言。MlirOptMain返回的是mlir::LogicalResultsuccess或failure而main需要返回int。asMainReturnCode做了一层转换success→0进程正常退出failure→1进程异常退出shell 能感知到失败这样在脚本里就能用if ./my-opt input.mlir; then ...判断 Pass 是否成功。三、实现pass主体#include lib/Transform/Affine/AffineFullUnroll.h #include mlir/Dialect/Affine/IR/AffineOps.h #include mlir/Dialect/Affine/LoopUtils.h #include mlir/include/mlir/Pass/Pass.h using mlir::affine::AffineForOp; using mlir::affine::loopUnrollFull; void AffineFullUnrollPass::runOnOperation() { getOperation().walk([](AffineForOp op) { if (failed(loopUnrollFull(op))) { op.emitError(unrolling failed); signalPassFailure(); } }); }void AffineFullUnrollPass::runOnOperation()这是 Pass 的入口。当 PassManager 把这个 Pass 调度到某个func.func上时就会执行这个函数。在函数内部你可以通过getOperation()拿到当前正在处理的FuncOp。3.1遍历 IR 树walkgetOperation().walk([](AffineForOp op)getOperation()返回当前的mlir::func::FuncOp因为我们在头文件里把 Pass 锚定到了FuncOp上。.walk(...)这是 MLIR 中极其常用的访问者模式Visitor Pattern。它会深度优先遍历当前操作内部的所有子操作。[](AffineForOp op)Lambda 表达式。walk会检查遍历到的每个操作只有当操作类型匹配AffineForOp时才会触发这个 Lambda。作用自动找出函数体内所有的affine.for循环包括嵌套在里面的交给 Lambda 处理。3.2 完全展开循环loopUnrollFullif (failed(loopUnrollFull(op)))loopUnrollFull(op)这是 MLIR 的LoopUtils库提供的核心算法。它的作用是分析这个affine.for循环的上下界、步长如果边界是编译期常量且循环次数在合理范围内就把这个循环原地替换成 N 份重复的循环体代码并把循环变量替换成具体的常量。failed(...)loopUnrollFull返回的是LogicalResult成功或失败。如果循环边界是动态的无法在编译期确定次数或者展开后代码量爆炸超过阈值它会返回failure。3.3 错误处理与 Pass 失败op.emitError(unrolling failed); signalPassFailure();op.emitError(...)在当前的操作那个无法展开的affine.for上附加一条错误诊断信息。这会告诉用户“嘿这个循环我没法完全展开原因在这里”。signalPassFailure()这是 MLIR Pass 机制中非常关键的一步。它告诉 PassManager“这个 Pass 执行失败了”。为什么不能直接return因为walk可能已经展开了一部分循环直接return会导致 Pass 处于一个部分完成、未知状态。调用signalPassFailure()会让 PassManager 标记这个 Pass 失败并停止后续的 Pass 流水线避免在错误的状态上继续编译。注意调用signalPassFailure()后代码不会立刻中断walk还会继续遍历剩下的操作但 PassManager 最终会判定整个 Pass 失败。代码作用getOperation().walk(...)遍历当前函数内的所有操作筛选出AffineForOploopUnrollFull(op)尝试将找到的仿射循环完全展开为直线代码failed(...)检查展开是否成功动态边界或代码膨胀会失败op.emitError(...)在失败的具体操作上生成诊断信息signalPassFailure()向 PassManager 报告此 Pass 失败终止后续编译流水线这段代码只有短短几行但完美体现了 MLIR Pass 的编写范式类型锚定只处理FuncOp、IR 遍历walk、调用官方工具库loopUnrollFull、严谨的错误处理signalPassFailure。四、改成用Rewrite Patternstruct AffineFullUnrollPattern : public OpRewritePatternAffineForOp { AffineFullUnrollPattern(mlir::MLIRContext *context) : OpRewritePatternAffineForOp(context, /*benefit*/1) {} LogicalResult matchAndRewrite(AffineForOp op, PatternRewriter rewriter) const override { return loopUnrollFull(op); } };struct AffineFullUnrollPattern : public OpRewritePatternAffineForOp {OpRewritePatternAffineForOp这是 MLIR 提供的模板基类专门用于基于操作Operation的模式重写。模板参数AffineForOp告诉 MLIR 重写引擎“这个模式只对affine.for操作感兴趣”。当引擎遍历 IR 时只有在遇到AffineForOp时才会把该操作交给这个模式去尝试匹配。CRTP 的体现与PassWrapper类似OpRewritePattern内部也使用了 CRTP 等机制但它不需要派生类把自己传进去因为AffineForOp已经指明了目标类型。AffineFullUnrollPattern(mlir::MLIRContext *context) : OpRewritePatternAffineForOp(context, /*benefit*/1) {}contextMLIR 的上下文对象包含方言、类型等全局信息。/*benefit*/1这是重写模式的收益值Benefit。当多个模式同时匹配同一个操作时MLIR 的 Greedy Rewrite Driver 会优先执行benefit值更高的模式。这里设为 1表示这是一个基础优先级的模式。LogicalResult matchAndRewrite(AffineForOp op, PatternRewriter rewriter) const override { return loopUnrollFull(op); } };这是所有 Rewrite Pattern 必须实现的纯虚函数包含两个核心步骤Match匹配检查op是否满足重写条件。Rewrite重写如果满足使用rewriter修改 IR。在这段代码中匹配和重写被合二为一了它直接调用loopUnrollFull(op)。这个函数内部会先检查循环边界是否静态、是否满足展开条件相当于 Match如果满足就修改 IR相当于 Rewrite。return loopUnrollFull(op);直接返回LogicalResult。如果返回success()重写引擎会认为修改成功并继续处理其他操作。如果返回failure()例如循环次数动态未知无法展开重写引擎会认为这个模式没有匹配成功不会对 IR 做任何修改并尝试其他模式。PatternRewriter rewriter这是 MLIR 提供的安全修改 IR 的工具。虽然这段代码没有直接用它因为loopUnrollFull内部自己处理了重写但通常情况下你需要通过rewriter.replaceOp()、rewriter.eraseOp()等方法来修改 IR绝对不要直接删除或修改操作否则会导致引擎崩溃。如何使用这个 Pattern定义好这个 Pattern 后它不会自己运行。你通常需要在一个 Pass 内部使用applyPatternsAndFoldGreedily贪心应用模式并折叠来驱动它。例如你可以写一个AffineFullUnrollPatternPassvoid runOnOperation() override { RewritePatternSet patterns(getContext()); patterns.addAffineFullUnrollPattern(getContext()); // 启动重写引擎不断应用模式直到 IR 不再变化 if (failed(applyPatternsAndFoldGreedily(getOperation(), std::move(patterns)))) { signalPassFailure(); } }一句话总结用 Pass当你需要一个独立的、有明确边界的编译阶段比如全局分析、跨函数变换、需要按特定顺序执行的流水线。用 Pattern当你需要定义局部的、基于匹配的重写规则比如常量折叠、操作替换、死代码消除并希望它们被自动、反复、智能地组合执行。核心对比对比维度Pass (例如AffineFullUnrollPass)Pattern (例如AffineFullUnrollPattern)作用域宏观。可以是Module、Func、甚至特定的 Op。微观。总是锚定到某个或某几个具体的 Op如AffineForOp。执行方式主动。自己写walk或遍历逻辑控制执行流程。被动。由重写引擎如GreedyPatternRewriteDriver遍历并调用。状态与通信可以有成员变量可以调用AnalysisManager获取全局分析结果如支配树、数据流。通常无状态。不直接访问全局分析结果需要的话得通过 Operation 的 Attr/Type 或特殊机制传递。失败处理失败即中断整条 Pass 流水线signalPassFailure()。匹配失败直接跳过不影响其他 Pattern 和整体编译。命令行暴露通过PassRegistration注册可以直接在mlir-opt里用--my-pass调用。不直接暴露。必须作为某个 Pass如--canonicalize的一部分去运行。组合与迭代按顺序执行一个 Pass 跑完才跑下一个。高度可组合。几十个 Pattern 放在一起引擎会不断尝试直到 IR 稳定不动点。你的需求选哪个“我要在终端里跑mlir-opt --my-opt”Pass“我要写一个规则op A可以安全地变成op B”Pattern“我要获取整个函数的支配树然后基于它重排指令”Pass“我要让x * 1自动变成x并且尽可能多地化简”Pattern“我要分析整个 Module 的符号表然后死代码消除”Pass“我写了一堆局部化简规则希望它们能自动组合生效”Pattern放进RewritePatternSet由 Pass 驱动在真实的 MLIR 开发中最强大的做法是在一个 Pass 内部去驱动 Pattern。参考https://zhuanlan.zhihu.com/p/2065755079806919146
