ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

Mojo 编译器 Post-Parser 调试完全指南:用 kgen 系列工具剖析 MLIR 与 LLVM IR

Mojo 编译器 Post-Parser 调试完全指南:用 kgen 系列工具剖析 MLIR 与 LLVM IR Mojo 编译器 Post-Parser 调试完全指南用 kgen 系列工具剖析 MLIR 与 LLVM IR【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo导读本指南以 Mojo 编译器仓库中的 PostParserDebugging.md 为骨架系统讲解 Mojo 编译流水线在 Elaboration 之后的调试方法。你将掌握kgen、kgen-opt、kgen-llvm-opt、llvm-module-split等内置工具的使用方式与各自定位学会如何定位肇事 Pass、转储中间表示MLIR/LLVM IR/bitcode/汇编、借助KGEN_OPTIONS环境变量注入 Pass 级调优参数以及如何用lldb/gdb调试编译器本身。文章还深入覆盖 Apple GPU 后端AIR 文件的专项调试流程包括 invalid bitcode、错误输出与Failed to create compute pipeline state三类典型问题的排查方法。为什么需要Post-Parser 调试Mojo 编译器的前端把 Mojo 源码解析、语义检查并 Elaborate展开成 MLIR 之后剩下的大量工作——从 MLIR 到 LLVM IR 的降级lowering、自定义 LLVM 优化管线、后端汇编/对象生成——都发生在解析器Parser之后。当生成的 MLIR 或 LLVM IR 不正确时问题往往出在某个具体 Pass 上此时就需要一套面向中间表示的调试工具链。从仓库源码看这套工具集中位于 Mojo/tools 目录下构建目标由各子目录的BUILD.bazel定义如 Mojo/tools/kgen/BUILD.bazel、Mojo/tools/kgen-opt/BUILD.bazel 等。它们共同构成了 Mojo 编译器解析之后的观察窗口。调试工具全家福工具用途源码位置kgen把 Mojo 程序编译到由命令行选项控制的某个阶段并输出结果Mojo/tools/kgen/kgen.cppkgen-opt类似mlir-opt可对给定 MLIR 单独运行某个 PassMojo/tools/kgen-opt/kgen-opt.cppkgen-llvm-opt测试 Mojo 自定义 LLVM 优化管线类似optMojo/tools/kgen-llvm-opt/kgen-llvm-opt.cppllvm-module-split测试把 LLVM IR 拆分到不同模块的 splitterMojo/tools/llvm-module-splitkgen-reduce归约reduceMLIR 以便调试仓库中标注 NOT TESTED较少使用Mojo/tools/kgen-reduce其中kgen与kgen-opt是日常调试的主力下面分别展开。kgen按阶段截停编译kgen的核心逻辑在 Mojo/tools/kgen/kgen.cpp 的runToolPipeline函数中它按命令行指定的Command决定在哪一步停下来并输出中间产物。文档给出的选项如下-elaborate在 Elaboration宿主侧结束后立即停止打印结果 MLIR-emitllvm在 MLIR 降级为 LLVM IR宿主侧后停止打印 LLVM IR-emitllvm-bitcode同上但输出 LLVM bitcode-emitllvm-opt跑完 Mojo 自定义 LLVM 优化管线后停止打印结果 LLVM IR-emitllvm-opt-bitcode同上输出 bitcode-emitasm跑完自定义优化管线并完成后端汇编输出汇编-emitasm-verbose同asm但内联更多信息-emitobject输出二进制目标文件其他选项见--help调试中较少使用。从源码看kgen的main中注册了registerMLIRContextCLOptions、registerAsmPrinterCLOptions、registerDefaultTimingManagerCLOptions、registerPassManagerCLOptions以及KGEN::KGENPassCLOptions::registerOptions()见 Mojo/tools/kgen/kgen.cpp这意味着它同时支持上游mlir-opt的常见旗标# 打印某个 Pass 运行后的 IR-all 表示所有 Pass kgen -elaborate test.mojo --mlir-print-ir-after-all # 打印单个 Pass 后的 IR kgen -elaborate test.mojo --mlir-print-ir-afterinline-parametric # 打印 Pass 统计信息 kgen -elaborate test.mojo --mlir-pass-statistics另外两个值得注意的选项是--save-temps与--temps-dirdir/prefix二者必须配合使用作用是把编译到某些阶段之后的 IR 保存成文件。其底层实现在编译选项中体现为saveTempsPrefixkgen把--temps-dir的值写入options.saveTempsPrefix见 Mojo/tools/kgen/kgen.cpp随后由 Mojo/include/Mojo/Compiler/SaveAsmOutput.h 中的writeTempModule等辅助函数按saveTempsPrefixphase.hashfileExt的命名模式落盘。CompilationOptions中对应的设置接口是setSaveTemps(std::string prefix)见 Mojo/include/Mojo/ToolCommon/CompilationOptions.h。⚠️ 关于--mlir-print-ir-[before|after]的关键限制该选项只对主PassManager生效即用于构建主流水线的那个 PassManager。Mojo 编译器以编译速度为核心目标许多 Pass 会通过构造自己的嵌套PassManager来做激进并行而这些嵌套 PassManager 并不继承主 PassManager 的选项例如不经过applyPassManagerCLOptions。因此--mlir-print-ir-[before|after]不会为这些嵌套 PassManager 打印 IR同样地它也不会为 offload例如 GPU目标打印 MLIR。这意味着在调试 GPU 相关问题时不能依赖--mlir-print-ir-after观察 offload 端的中间表示见下文Apple GPU 调试。kgen-opt单 Pass 实验台kgen-opt相当于 Mojo 编译器自己的mlir-opt可对给定的 MLIR 输入运行单个 Pass。其实现位于 Mojo/tools/kgen-opt/kgen-opt.cpp本质上包装了上游 MLIR 的MlirOptMain并额外注册了 KGEN 方言与全部 KGEN PassKGEN::registerDefaultKGENPasses(kgen-opt)还注册了DebugInfo::registerTransformsPasses()。用法示例# 对某个 MLIR 文件单独运行 inline-parametric Pass kgen-opt -passesinline-parametric dump.mlir -o out.mlir # 通过 --asyncrt-single-thread 强制单线程kgen-opt 特有选项 kgen-opt --asyncrt-single-thread -passesinline-parametric dump.mlir从源码看kgen-opt还内置了一个调试辅助 Passtest-always-fail一个总是失败的 Pass用于测试 crash reproducer可用来验证故障复现流程见 Mojo/tools/kgen-opt/kgen-opt.cpp。⚠️ 注意对于会自行构建嵌套PassManager与子管线的 Pass可能需要给其对应的 CL 选项追加额外参数才能生效。kgen-llvm-opt自定义 LLVM 优化管线测试kgen-llvm-opt用于测试 Mojo 的自定义 LLVM 优化管线。与 LLVM 的opt不同它目前不支持单独测试某个 Pass但支持-passes语法见下。从 Mojo/tools/kgen-llvm-opt/kgen-llvm-opt.cpp 的头部注释看它有两种工作模式完整 KGEN 优化管线通过-O0/-O1/-O2/-O3运行对应优化级别的完整 Mojo 编译管线单 Pass 模式通过-passes...指定 LLVM pass 管线语法与opt -passes...相同例如kgen-llvm-opt -passeskgen-llvmir-downgrade in.bckgen-llvmir-downgrade用于把 IR 降级以兼容更老的 LLVM 后端。此外还支持--disable-optimization-passes禁用优化 Pass 并打印输入模块用于观察未经优化的 IR。这在 Apple GPU 调试中尤为重要见下文。llvm-module-split与kgen-reducellvm-module-split允许测试把 LLVM IR 拆分到不同模块的 splitterMojo 编译器在并行编译时会用到模块拆分见Tips一节kgen-reduce仓库文档标注NOT TESTED其目标是把 MLIR 归约成更小、更易调试的形式但当前似乎无人使用。调试方法论定位肇事 Pass与 LLVM 不同Mojo 编译器目前文档撰写时还没有类似-opt-bisect-limit的自动二分定位选项。因此定位问题 Pass 主要靠人工操作先找出导致问题的 Pass。由于缺少 bisect 选项可能需要手动摆弄编译管线例如通过kgen的--mlir-print-ir-afterpass逐步缩小范围找到第一个产出异常 IR 的 Pass。警惕嵌套 PassManager 及其中的 Pass嵌套 PassManager 中的 Pass 不会被--mlir-print-ir-before/after覆盖见上节遇到这类 Pass 需要另想办法例如临时加op.dump()。确定肇事 Pass 后用kgen转储 IR。用kgen-opt、opt或其他上述工具对 IR 做进一步实验找出根因。⚠️ offload 编译的特殊情况如果需要观察KGEN - LLVMPass 之前的 MLIR即 offload 目标端尚未降级的中间表示目前只能手动在 Pass 里加上op.dump()之类的调试输出并建议用-workqueuesingle-thread关闭多线程以保证输出顺序与稳定性。实战 Tips 与构建配置文档给出了几条实战经验均可在仓库中找到对应的实现依据。1. 运行时失败优先用release编译器复现如果用户程序在运行时失败先用release配置构建的编译器复现br //:install --configrelease原因productionbr //:install --configproduction构建会关闭所有验证与断言assert因此可能生成无效 MLIR 而不报错。release构建保留了验证与断言更早暴露问题。这一结论与源码中KGENPassCLOptions在MODULAR_PRODUCTION宏下把可选项全部退化为默认值的行为相互印证见 Mojo/include/Mojo/ToolCommon/CLOptions.h生产构建去掉了运行期选项解析等开销同时也去掉了大量防御性检查。2. 不要用debug构建跑kgen除非正在调试 Mojo Parser 本身否则不要用debug构建br //:install --configdebug-modular或--configdebug-everything来跑kgen工具。原因debug构建非常慢。用于 Parser 调试之外只会浪费时间。3. 通过KGEN_OPTIONS注入 Pass 级选项部分 Pass 有专属的 CL 选项可以经KGEN_OPTIONS环境变量传给mojo工具KGEN_OPTIONS-kgen-parametric-inline-threshold35 mojo build test.mojo这要求先在mojo源码中定义KGEN_ENABLE_PASS_OPTIONS并重新构建。支持的选项定义在KGENPassCLOptions中见 Mojo/include/Mojo/ToolCommon/CLOptions.h。从该类的实现可以提取出当前支持的选项清单含默认值选项类型说明默认值kgen-automatic-inline-thresholduint64自动内联函数的阈值优先级高于 Pass 自身启发式无需显式指定kgen-parametric-inline-thresholduint64内联 parametric 函数的阈值优先级高于 Pass 自身启发式无需显式指定kgen-parametric-inline-avg-loop-trip-countuint64InlineParametricPass 启发式中估计的循环平均迭代次数4kgen-stack-reuse-promote-to-global-thresholdsize_t只读栈分配被提升为全局变量的字节阈值1024kgen-verifier-max-errorssize_tKGENVerifierPass 单次最多输出的错误数10这些选项的消费者散落在 Mojo/lib/Transforms/InlineParametric.cpp、Mojo/lib/Transforms/AutomaticInline.cpp、Mojo/lib/Transforms/StackReuse.cpp 与 Mojo/lib/Transforms/KGENVerifier.cpp 等 Pass 实现中kgen、kgen-opt、mojo-build、mojo-run等工具入口都会调用KGENPassCLOptions::registerOptions()完成注册。使用时按调试需要调整阈值即可例如调低 parametric 内联阈值以观察内联决策对 IR 的影响。4. 怀疑 LLVM 模块拆分器时关闭并行编译偶尔会出现与 LLVM 模块拆分器splitter相关的编译期或运行期问题。如果怀疑是它尝试禁用并行编译例如用-workqueuesingle-thread再复现。5. 只构建调试所需工具在全新工作区上仍需要先跑一次br //:install --config....来正确设置符号链接之后通常只需增量重建需要的工具即可。例如br //:kgen-tool //:kgen-opt该命令只构建kgen与kgen-opt两个工具避免全量构建带来的时间开销。用调试器调试 Mojo 编译器本身当问题出在编译器自身逻辑而非生成的 IR时需要用调试器单步跟踪。macOS 上lldb是最简单的选择Linux 上可用gdb或lldb。有三种启动调试器的方式对 Bazel 测试目标bd [--gdb] [--config...] //bazel-test对工具加参数--之后是传给被调试工具的选项bd --configdebug-modular --gdb KGEN/tools/mojo -- build \ --target-acceleratoramdgpu:gfx942 test.mojo直接调用lldb或gdb。注意两点debuginfo 包含相对路径因此在 Mojo 工作区之外运行调试器将无法显示符号如果想保留工作区之外的测试目录可以使用下面的 shell 别名把文件参数解析为绝对路径后再进入MODULAR_PATH目录启动调试器alias blldbblldb_() { if [[ -z ${MODULAR_PATH} ]]; then echo ${MODULAR_PATH} must be set fi args(); for arg in $; do if [[ -f ${arg} || -d ${arg} ]]; then args($(readlink -f ${arg})) else args(${arg}) fi done echo (cd ${MODULAR_PATH}; lldb ${args}) (cd ${MODULAR_PATH}; lldb ${args}) unset -f blldb_; } blldb_ alias bgdbbgdb_() { if [[ -z ${MODULAR_PATH} ]]; then echo ${MODULAR_PATH} must be set fi args(); for arg in $; do if [[ -f ${arg} || -d ${arg} ]]; then args($(readlink -f ${arg})) else args(${arg}) fi done echo (cd ${MODULAR_PATH}; gdb ${args}) (cd ${MODULAR_PATH}; gdb ${args}) unset -f bgdb_; } bgdb_这两个别名存在一些显而易见的边界情况例如路径含空格时但在把测试文件放在工作区之外的日常场景下非常实用。Apple GPU 调试专章Apple GPU 没有 LLVM 上游支持Mojo 对它的支持来自逆向工程reverse engineering因此其调试流程非常特殊。整体思路是Mojo 编译器先把 LLVM IR 转换成AIRApple Intermediate Representation兼容的 LLVM IR再交给 Metal 编译器通过air-*工具链生成最终内核。核心命令组合文档推荐的最佳命令组合如下# 生成转成 AIR 之前的 LLVM IR 文件以及生成的 AIR 文件 kgen -elaborate --save-temps --temps-dir # 观察 Metal 编译器对给定 AIR 文件的行为 xcrun -sdk macosx air-* # 得到 AIR 转换之后的 LLVM IR kgen-llvm-opt -O3 -S # 不做任何 LLVM IR - AIR 兼容 IR 的变换直接从 LLVM IR 产出 AIR 文件 kgen-llvm-opt --disable-optimization-passes -O3若想直接在运行时测试某个手工修改过的 AIR 文件KGEN_OPTIONS-kgen-object-compiler-use-custom-airyour-air-file mojo \ build/run test.mojo配合llvm-reduce可以把 LLVM IR 归约到最小复现规模。为 Apple GPU 添加 NVidia/AMDGPU 的某个特性 X这类工作的第一步是理解 Metal 如何支持该特性。可以借助 LLM 或 Metal Shading Language 规范写一个 metal kernel然后编译它并转储 LLVM IR观察特性 X 是如何实现的/Applications/Xcode.app/Contents/Developer/Toolchains/XcodeDefault.xctoolchain/usr/bin/metal \ -mllvm -print-after-all test.metal dump之后再把对应实现移植进编译器和/或 stdlib。问题一Invalid bitcodemetallib compilation failed with codeMetal 编译器仍使用 LLVM 17.0因为它依赖 typed pointers 与 5.0 bitcode reader。Mojo 编译器自带一个应当与 5.0 兼容的 bitcode writer但两者之间可能存在差异。如果看到metallib compilation failed with code说明生成的 bitcode 对 Metal 无效。排查步骤kgen -elaborate test.mojo --save-temps --temps-dir/tmp/apple_gpu xcrun -sdk macosx air-objdump --disassemble /tmp/apple_gpu.*.airair-objdump --disassemble会打印 Metal 编译器不接受的 opcode。如果仍不明确问题所在用llvm-reduce配一个归约脚本缩小测试用例! kgen-llvm-opt --disable-optimization-passes -O3 $1 -o /tmp/kernel_$$.air || exit 1 xcrun -sdk macosx air-objdump --disassemble /tmp/kernel_$$.air \ -o /tmp/metalllib$$.metallib /tmp/reduce_$$.log grep -q the error you care about /tmp/reduce_$$.log || exit 1然后执行归约kgen-llvm-opt -O3 -S /tmp/apple_gpu.pre-split*.ll -o /tmp/apple_gpu.kernel.ll llvm-reduce --testreduce.sh /tmp/apple_gpu.kernel.ll问题二Incorrect output错误输出这是最耗时的调试问题。目前已知会导致错误输出的两类根因load或store上的 address space地址空间不正确某个MTL::Buffer没有被传给 encoder。最佳策略是手动摆弄 AIR 文件找出具体是哪条或哪些指令不对kgen -elaborate test.mojo --save-temps --temps-dir/tmp/apple_gpu kgen-llvm-opt -O3 /tmp/apple_gpu.pre-split.*.ll -S -o /tmp/kernel.ll此时/tmp/kernel.ll包含的是已转换为 AIR 兼容形式的 LLVM IR。接着在运行时验证它kgen-llvm-opt --disable-optimization-passes -O3 /tmp/kernel.ll \ -o /tmp/apple_gpu.air clear-cache KGEN_OPTIONS-kgen-object-compiler-use-custom-air/tmp/apple_gpu.air \ mojo run test.mojo随后就可以逐条修改 AIR定位到底是指令、属性还是 metadata 出错。问题三Failed to create compute pipeline state这是目前最不明确的问题。它发生在 AsyncRT 试图获取需要传入MTL::Buffer的参数个数时——也就是说 Metal 编译器认为生成的 AIR 文件没问题但反射reflection阶段仍然出错。目前的最佳实践是组合 Invalid bitcode 与 Incorrect output 两节的方法并为llvm-reduce编写更复杂的归约脚本。⚠️ 注意由于这类问题必须实际运行程序才能复现llvm-reduce归约过程中某些测试可能出现无限递归或死循环务必给运行二进制加上 timeout。总结Mojo 编译器的 Post-Parser 调试本质上是在不同编译阶段截停并检查中间表示的过程用kgen按阶段截停编译并导出 MLIR / LLVM IR / bitcode / 汇编用kgen-opt对单个 Pass 做最小化实验用kgen-llvm-opt验证自定义 LLVM 优化管线含 AIR 转换用llvm-module-split/llvm-reduce处理模块拆分与用例归约问题通过KGEN_OPTIONS注入 Pass 级选项、通过bd/blldb/bgdb在调试器里跟踪编译器本身。需要特别记住的三条经验--mlir-print-ir-before/after对嵌套 PassManager 与 offload 目标无效生产构建会关闭验证与断言调试时优先用release构建Apple GPU 的调试必须理解 LLVM IR → AIR → Metal 工具链这条特殊路径。掌握了这套方法论无论是定位错误降级、内联策略问题还是排查 GPU 后端的 bitcode 兼容性都能做到有迹可循。延伸阅读MojoCompilerWalkthrough.md编译器整体流程讲解Compiler.mdSupport 库中的编译期支持设施Mojo 编译流水线测试kgen相关集成测试kgen-llvm-opt 工具说明两种工作模式与 Pass 列表KGENPassCLOptions 定义可注入的 Pass 级选项清单【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进