尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

Mojo 生成式内核编译器(kgen)任务清单全解析:从元编程 IR 到内核搜索与生成

Mojo 生成式内核编译器(kgen)任务清单全解析:从元编程 IR 到内核搜索与生成 Mojo 生成式内核编译器kgen任务清单全解析从元编程 IR 到内核搜索与生成【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo导读本文基于 Modular 仓库中一份历史性技术文档《Generative Kernel Compiler Task List》展开。该文档是 kgenkernel generator生成式内核编译器项目早期的工作分解清单以待实现 / 待规划 / 已完成三段式记录了构建一套用编译器生成高性能内核体系所需的全部技术组件——从可带参数的kgen.generator元编程 IR到meta.buffer参数化类型、generator 接口与约束系统、elaboration展开算法再到 JIT 执行与性能测量基础设施。文档中标注为已完成的每一项都能在当前仓库中找到真实的源码落点KGENDialect、Elaborator、POPDialect、LowerLIT 等目录。读完本文你将理解 Mojo 编译器如何描述内核的元程序、如何穷举展开所有内核变体、如何用约束剪枝搜索空间以及这些设计与仓库中具体源码文件的对应关系。一、文档背景一份生成式内核编译器的工程化路线图这份位于 Mojo/docs/compiler/attic/TaskList.md 的文档自述是 Generative Kernel Compiler Language 设计文档的实现任务清单用于把宏大设计拆解成可执行的工程粒度。它被放在attic阁楼目录下从标题中的 (OLD) 可以看出它已属于历史参考材料——文档中Tasks to be implemented一节里还标注为未完成的任务如目标机器建模、memset/erf 生成器、搜索缓存在后来的演进中大多已通过不同方式落地。整份文档的核心思想可以概括为一条主线让内核kernel本身变成程序生成程序的过程——开发者编写带参数的 generator生成器编译器elaborator根据目标硬件与参数约束自动展开并生成所有合法变体再通过实测选择最优实现。这与传统手写汇编/SIMD 内核、或仅靠 LLVM 自动向量化的思路都不同kgen 把选择哪个内核变体从编译期决策升级为可搜索、可测量、可缓存的显式过程。二、文档结构总览三段式任务编排文档采用三段式结构组织任务这种编排本身就是项目治理方式的体现段落状态标记含义Tasks to be implemented❌已有基本范围划分但尚未完整实现To be scoped✅ / ⚠️已确认✅或待定⚠️的远期方向尚未细化Tasks completed✅已完成按自底向上的顺序给出技术组件总览文档还明确了工作流转规则当某个大方向完成时应把任务移到文档末尾的 Tasks Completed 一节。这种未完成→已完成的移动机制使文档成为一份可演进的活文档——从仓库现状看绝大多数 ✅ 任务都已在 Mojo/lib 与 Mojo/include/Mojo 中找到了对应实现。三、核心概念与 IR 设计已完成部分Tasks completed一节按自底向上顺序列出了主要技术组件其中最关键的是元编程参数基础设施——这是整个 kgen 体系的地基。3.1 新库/框架kgen 的命名由来文档记载项目的第一步是在 Modular 仓库中新建一个遵循 include/lib/tools/test 标准结构的顶层目录。命名过程颇为有趣corn/popcorn因涉及 kernel 有 emoji 可用但听起来很蠢被否决Abdul 建议 kgenkernel generator 的缩写Tatiana 建议 kirkernel intermediate representation。最终敲定kgen且这个名字同时被用作一个方言dialect的名字。从仓库看这一决策完整落地存在 Mojo/lib/KGENDialect19 个 .cpp 实现文件、Mojo/include/Mojo/KGENDialectOps/Attrs/Types 的 TableGen 定义以及独立的 Mojo/tools/kgen/kgen.cpp 工具。3.2 元编程参数基础设施kgen.generator 与参数表达式文档指出kgen 的核心需求是在同一个 IR 中同时描述程序和元程序。为此需要一系列 IR 构件文档逐条列出均已完成kgen.generator操作类似函数的操作但允许携带参数。参考先例是 CIRCT 的hw.module。该操作在 KGENOps.td 中以def KGEN_GeneratorOp : KGENOpgenerator, ...定义。参数表达式以简单的二元表达式和参数声明引用起步例如#kgen.param.expradd, ...需要实现代数规范化。仓库中的 KGENParameters.cpp 与 ParameterEvaluator.cpp 承担了参数声明与求值职责。漂亮的打印/解析语法文档作者明确不满意 CIRCT 参数属性的冗长打印如#hw.param.expr.add...套娃要求在大量测试用例落地前把语法糖设计好。当前仓库的语法已经简化为前缀调用形式例如测试文件 kgen-param-exprs.mlir 中的%0 kgen.param.constant add(p1, 42, mul(p2, p2)) %3 kgen.param.constant: i1 to_builtin(:scalarbool eq(42, p1)) %6 kgen.param.constant: i1 to_builtin(:scalarbool identical(:dtype type, f32))该测试通过kgen-opt -verify-parameters校验参数表达式并做常量折叠例如eq(41, 42)折叠为0、identical(bf16, f16)折叠为0。 4.kgen.param.constant操作把参数投影为 SSA 值使元程序中的值能被程序使用。 5.参数返回能力允许 generator 返回参数如白皮书中的panelDotInner类似hw.output但允许参数列表作为属性。 6.参数类型系统决策Issue #983允许任意类型的参数未指定类型时默认是带符号解释的 index 类型——这保证投影到 SSA 域后与架构无关并能与 SCF 方言良好配合。语法上两种写法等价kgen.generator foop1, p2, p3(...) kgen.generator foop1: index, p2: index, p3: index(...)kgen.callIssue #960支持 generator 调用其他 generator并带校验。局部参数声明与绑定Issue #966让 IR 更具表达力和可读性便于内核规模增长后的维护。文档特别强调表达式规范化、pretty print/parse 逻辑、以及尽早建立的 verification 逻辑是让后续所有工作变轻松的关键。仓库中 KGENAttrsFolders.cpp、KGENOpsFolders.cpp 正是表达式折叠/规范化的落地实现。3.3 DType 参数的类型系统!kgen.dtype参数列表里 dtype 极为常见但 MLIR 参数系统允许任意类型如何处理 dtype 就成了问题。文档否决了用带魔数的i8参数表示 DType这种简单但难读的方案改为引入新的!kgen.dtype类型对应 KGENTypes.td 中的def KGEN_DTypeType : KGENTypeDType, dtype, ...允许参数列表中出现 0..255 整数初始化器、直接对应 DType 枚举值的属性让printParamValue对f32这类参数名加引号打印f32使其像 MLIR 关键字一样不能作为裸词为f32等常见类型引入特殊关键字在解析时映射到对应 DType 值。最终效果是下面这种可读性极佳的语法kgen.generator foot1: !kgen.dtype(... kgen.generate foot1: !kgen.dtype f32(...进一步Issue #980结合 kgen 类型上下文还可以省略!kgen.前缀kgen.generator foot1: dtype(... kgen.generate foot1: dtype f32(...3.4 新参数化类型meta.scalar / meta.buffer / meta.simd / pop.pointer有了参数化系统下一步是构建参数化类型参考 CIRCT 的hw.Int、hw.Array类型含义对应 Issuemeta.scalarx具体标量类型x 为 dtype#978meta.buffersize, x缓冲区memref 的类比size 为 index、x 为 dtype#1009meta.simdsize, tySIMD 类型size 可为任意参数表达式—pop.pointerT无界指针运算#1663配套基础设施包括参数作用域校验#979、与具体类型互转的 cast 操作#1221、meta.buffer.address操作#1249。文档特别强调了一条重要设计原则通常不需要为 meta.scalar / meta.simd 类型定义算术操作——这些操作本身就应该是内核kernel。也就是说加法的实现路径是写一个生成加法的 generator而不是在 IR 里内置一个特权fadd指令。3.5 动态 shape 与 dtypebuffer除编译期已知的参数化类型外meta.buffer还必须支持动态大小与动态元素类型Issue #1082。MLIR 中可以用 null 参数表示未知例如kgen.generator fillWithOnes(%dest: !buffer?, ?) {} kgen.generator knownFloatAlgo(%dest: !buffer?, f32) {} kgen.generator fixedSizeUnknownAlgo(%dest: !buffer1024, ?) {}该能力只适用于 buffer 类型——文档明确反对在 scalar/simd 上支持动态 dtype也反对 simd 支持动态长度理由指向仓库中的 Mojo/docs/compiler/manual/Rationale.md。配套的运行时提取操作Issue #1083/#1084可以把参数提取为 SSA 值并做 buffer 类型间互转kgen.generator algo(%dest: !buffer?, ?, %other : !buffer1024, f32) { %dtype meta.buffer.dtype %dest: !buffer?,? // returns !kgen.dtype %size meta.buffer.size %dest: !buffer?,? // returns index %buf1 meta.buffer.cast %dest: !buffer?,? to !buffer?, f32 %buf2 meta.buffer.cast %other: !buffer1024xf32 to !buffer?, f32 }其中meta.buffer.dtype返回对应 DType 枚举的 i8 值meta.buffer.size返回对应 size_t 的 index 类型当参数为已知常量时应提供fold()折叠。3.6 Generator 接口声明与实例设计的另一个关键点是一个 generator 接口可以有多个实现。因此需要把接口声明与实现分离kgen.generator.interface声明接口Issue #1086generator声明需标明自己实现哪个接口并有类型检查保证声明/实现兼容调用点把具体信息如具体 dtype绑定到接口声明侧的泛型参数上通过传递参数 meta.buffer.cast重化reify具体类型。文档给出了如下示意kgen.generator.interface thingdtype?(%dest: !meta.buffer?xdtype) ... %dstCast meta.buffer.cast %dest to buffer?xi32 meta.call thingi32(%dstCast)3.7 约束系统用布尔表达式剪枝生成generator 是以参数为输入产生内核的参数化函数但它可能不是全函数——某些输入对它是非法的。约束constraints机制允许定义布尔条件来控制 generator 对其输入参数的有效性从而让 elaborator 在展开早期就把不满足约束的 generator 剪掉避免浪费搜索时间Issue #1087constraints in(vecLen, 2,4,8,16,32), notequals(type, bf16)文档预告将来表达式语法可迁移到中缀形式Issue #1395如constraints vecLen ∈ {2,4,8,16,32}, type ! bf16。generator 接口同样支持约束用于统一约束其全部实现。配套能力还有elaborator 使用约束剪枝生成Issue #1396生成失败时的展开栈回溯错误报告机制重构诊断与追踪kgen.param.assert针对任意参数表达式包括非直接 generator 参数的广义断言。四、编译器 elaboration 算法穷举展开所有变体有了描述能力多实现、基本参数、约束还需要执行展开的编译器基础设施见 Mojo/lib/Elaborator/Elaborator.cpp 及其辅助文件。文档明确初期不引入搜索而是穷举生成所有可能的内核实现。展开算法的步骤草图实现新的kgen-generate工具读取含生成式内核的 MLIR 文件 库文件并加载校验源 MLIR 文件与库文件中的内核接口在签名等细节上一致处理内核体内的调用与其他操作折叠参数表达式以SCC 顺序自底向上遍历 generator 调用树把环检测为错误库文件中的实现视为同一图的一部分自底向上生成完全特化的内核实现写入目标 MLIR 文件库文件保持不被修改结果应完全特化、参数全部消除为每个内核跟踪绑定确保接口调用点的多向展开方向一致。从仓库看kgen-generate工具如今演化为功能更完整的 Mojo/tools/kgen/kgen.cpp其命令行支持--elaborate仅展开并输出 MLIR、--emit-llvm、--emit-assembly、--emit输出 .o、--execute等命令并在runToolPipeline中调用compiler.runKGENPipeline(*theModule, target)。展开后的kgen.func即已被展开的 kgen.generatorKGENOps.td 中对kgen.func的注释如此说明。五、内核执行与性能测量基础设施展开出大量变体后必须回答哪个是好的生成内核。文档规划了两条执行路径Issue #1610 的 JIT、Issue #1611 的编译为 .o 再执行接入 LLVM 生成 .o 文件或 JIT 到缓冲区后执行运行需要输入生成基础设施与预期 dtype元数据收集真实数据进行对比例如 mmperf 的 matmul 维度列表、memset/memcpy 的长度直方图数据集。在 Mojo/tools/kgen/kgen.cpp 中可以确认这条链路的实现ObjectCompilerMojo/lib/Compiler/ObjectCompiler 相关负责编译ExecutionEngineMojo/lib/ExecutionEngine负责 JIT 执行且支持通过-func指定要执行的函数、通过-ignore-failure容忍失败——这些正是穷举生成所有变体 → 逐个运行测量工作流所需要的 CLI 支撑。六、方言分层lit 与 pop随着内核越写越多直接在 kgen 层编写会变得痛苦。文档规划了两层方言来支撑6.1 高层 lit 方言kgen 层应保持简单、便于分析与变换elaborator、静态分析工具都工作在类型检查过的元程序上因此需要一个会被脱糖desugar到 kgen的高层方言。具体任务包括lit.fn操作Issue #1410generator 实现接口时可以用比接口更特化的签名并据此推断约束模块系统与kgen.include操作涉及分离编译怎么做导入的东西是 kgen 还是 lit 抽象层次等问题。仓库中 Mojo/lib/LowerLITLowerLIT.cpp、ConcreteBindings.cpp、CheckLifetimes.cpp 等正是该层级的实现Mojo/lib/LITDialect 定义了 lit 方言本身。6.2 参数化操作方言 pop由于 LLVM 方言不支持参数化类型为每个整数宽度手写重载如 fadd 示例对人是巨大负担且拖慢编译。解决方案是新增parametric operationspop方言允许这样写kgen.generator fadddt: dtype(%lhs: !kgen.scalardt, %rhs: !kgen.scalardt) - !kgen.scalardt constraints type f32, f64 { %res pop.add %lhs, %rhs : !kgen.scalardt kgen.return %res }文档同时指出每个要用的方言都要适配工作是烦人的代价因此通用特化general specialization方法很重要使我们能与 MLIR 生态中的任何方言对话LLVM 方言足够特殊和重要值得投入适配。另外文档还建议借此机会重新考虑 signless 设计——pop 方言可采纳带符号类型signful types让类型泛型内核更好写不必为 divs/divu 各写一份。仓库中的 Mojo/lib/POPDialectPOPOps.cpp、POPTypes.cpp、POPLLVMFolder.cpp 等即该方言的实现且存在pop.dtype.to_ui8/pop.dtype.from_ui8POPOps.td这类连接 kgen 与底层整数世界的桥接操作。6.3 全类型参数化Issue #2504有了以上支撑可以写对标量与向量元素类型都泛型的内核但还不能写对标量与向量同时泛型的内核——这是另一种参数化形态需要把 MLIR 类型本身作为参数值传入。这是文档中最后一条已完成任务对应 KGENAttrs.td 中大量以!kgen.type承载 MLIR 类型参数的属性设计测试文件 kgen-param-exprs.mlir 中即有mlirType: type与#kgen.type!kgen.param:type mlirType的示例。七、标准库自举基础构件即内核文档还提出一个哲学层面的要求整个系统应定义在自身之内减少特权技术。基础标量与 SIMD 算子如 fadd应该用语言本身定义为内核而不是特权指令。文档给出了用 LLVM 方言 cast 操作实现 fadd 的雏形kgen.generator fadddtypef32(%lhs: scalarf32, %rhs: scalarf32) - scalarf32 { %lhsc meta.cast_to_builtin %lhs: scalarf32 to f32 %rhsc meta.cast_to_builtin %lhs: scalarf32 to f32 %resc llvm.fadd %lhsc, %rhsc : f32 %res meta.cast_from_builtin %resc : f32 to scalarf32 kgen.return %res }选择 LLVM 方言的好处是可直接使用目标特定 intrinsics 与内联汇编是否引入 arith 等更高层方言则需要评估。相关 Issue标量与 SIMD 加法#1088、buffer 加法#1089而fillWithOnes#1090在当时被标记为 ❌ 未完成——这也印证了文档中待实现与已完成是动态流动的。八、待实现与待规划方向文档保留了三条未完成任务Tasks to be implemented和四条远期方向To be scoped它们体现了 kgen 路线的完整野心8.1 未完成任务目标机器建模❌generator 应允许因与目标硬件无关而失败例如用了 VNNI 特性但目标芯片不支持或向量长度不匹配。文档作者明确表示不应让 LLVM 替我们合法化不支持的向量长度即使 LLVM 能做。更有意思的生成器❌实现 memset 与 erf 的真实生成器这需要实现数学库、memset 所需的由搜索得到的决策树、基于 dtype 的动态 switch、scf.for、buffer 分区操作等。用动态规划缓存❌当系统成型后搜索执行时间会成为问题。需要缓存内核、利用层级结构这要求能把微内核从大内核中隔离出来单独评估可能把输入数据约束如长度直方图下推给微内核也允许用户自定义度量如针对预期维度的 FLOPS。8.2 待规划方向多维张量✅ 已确认延后实现1D 的 buffer 类型并未硬编码进系统它只是类型 操作的组合一旦证明该模式可行就能为表格、树等更复杂的数据类型添加新类型与操作。量化类型与生成器支持⚠️可扩展平台允许重新讨论量化类型及其 codegen 方式用简单方案重实现 Ruy/XNNPACK 内核是好的起点但不在此前的里程碑范围内。花哨的搜索算法⚠️基础设施成熟后可探索多种搜索与搜索空间剪枝、黑盒优化、ML 辅助搜索当前阶段优先暴力搜索以集中建设其他基础设施最终希望支持多个可插拔/模块化搜索算法。高阶生成器⚠️希望 generator 能接收 region例如一个vectorize操作接收 region 与宽度并尝试向量化它类似 ISPC 对任意代码做 SPMD 化。这样可以把算法如 erf写成标量代码直接从标量规范生成 SIMD 版本scf.for之类的 lowering 原则上也能这样做。九、仓库中的实现证据与延伸阅读想深入源码的读者可以沿着以下路径继续探索均以仓库根目录为起点kgen 方言定义Mojo/include/Mojo/KGENDialect/KGENOps.td、KGENAttrs.td、KGENTypes.td以及实现 Mojo/lib/KGENDialect参数求值见 ParameterEvaluator.cpp展开算法Mojo/lib/Elaborator/Elaborator.cpp 与 ParametricElaborator.cppkgen 工具Mojo/tools/kgen/kgen.cpp支持 elaborate / emit-llvm / emit-assembly / emit / execute 等命令lit 与 pop 方言Mojo/lib/LowerLIT、Mojo/lib/LITDialect、Mojo/lib/POPDialect测试样例Mojo/test/kgen/kgen-dialect/kgen-param-exprs.mlir参数表达式解析/折叠、kgen-types.mlir、kgen-verify-errors.mlir校验错误路径设计澄清Mojo/docs/compiler/manual/Rationale.md动态 shape 的设计理由。结语从这份历史任务清单到仓库中的实现kgen 的路线图展示了一条清晰的技术演化路径先用参数化 IR 让内核的元程序可描述再用 elaboration 算法穷举展开并配合约束剪枝随后建立执行与测量基础设施来挑选最优变体最后通过 lit/pop 两层方言让内核编写与类型泛型化变得可持续。文档中❌ 未完成的三项目标机建模、memset/erf 生成器、搜索缓存以及⚠️ 待规划的搜索算法与高阶生成器则勾勒出 kgen 更远的野心从穷举 实测走向可插拔搜索 由标量规范直接生成 SIMD。对于想理解 Mojo 编译器内核生成机制的读者这份文档与上述源码目录互为印证是最佳的起点。【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表