Mojo 编译器内幕：Common Types and Tools——MLIR 节点、值包装器与辅助工具的全面指南

Mojo 编译器内幕Common Types and Tools——MLIR 节点、值包装器与辅助工具的全面指南【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo本篇指南以 Mojo/docs/compiler/manual/CommonTypesAndTools.md 为骨架结合 Mojo 编译器源码TableGen 定义、生成的 C、MojoParser 实现展开带你认识在 Mojo 编译器代码库中反复出现的三大类类型——MLIR 节点、数据包装器Data Wrappers与辅助工具Helpers掌握如何阅读.td定义、理解值分类体系并学会在编译器开发中正确创建与操作各类对象。读完本文你将具备为 Mojo 编译器添加新特性所需的大部分工具与心智模型。在深入一个陌生的编译器之前必须做好充分的准备。本页结束时你应该拥有为 Mojo 编译器添加一个有趣特性所需的大部分工具。建议先阅读 Passes and Intermediate Representations 和 Terminology 以建立前置知识。代码库中的三类类型Nodes、Data Wrappers 与 HelpersMojo 编译器代码库中你会看到三种类型的结构理解它们的区别是入门的第一步MLIR 节点MLIR Nodes包括属性attributes和操作operations。如果你在找 Mojo 的 AST 或 IR找的就是它们。我们用 TableGen 向 MLIR 定义这些节点然后 MLIR 为我们生成 C 代码。例如Mojo/include/Mojo/LITDialect/LITTypes.td中的def LIT_StructType { ... }块告诉 MLIR 如何生成LIT::StructType。数据包装器Data Wrappers围绕上述 MLIR 事物的轻量包装。例如ASTType是mlir::Type的包装器CValue、LValue、MLValue、DLValue、SBValue、MBValue、MBPValue、PValue、SRValue、MRValue是对我们定义的各种 MLIR 属性的包装。辅助工具Helpers如 builder、transformer、算法等帮助我们操作前两类对象。例如IREmitter包装了一个mlir::OpBuilder帮助我们输出表达式OverloadSet表示一次调用的合法候选集并提供方法帮助缩小范围并完成调用ParameterInferenceState在进行参数推断时保存所有暂定结论/绑定。当把Spaceship[42]与struct Spaceship[N: Int]匹配时正是ParameterInferenceState推算出N 42。下文将分别深入这三类类型。MLIR 节点一切皆定义在.td文件中先回顾一下对下面的main.mojo运行br //Mojo/tools/kgen-translate -- -import-mojo main.mojodef foo(arg: Int): pass def main(): foo(5)会产出如下 MLIRlit.fn ”main()”() - !kgen.none attributes {sourceName “main”, specialFnKind 0 : i8} { %0 kgen.param.constant: !Int {5} %1 lit.call main::”foo(::Int)”(%0) : !lit.generator(“arg”: !Int) - !kgen.none %none kgen.param.constant: none #kgen.none lit.return %none : !kgen.none lit.end_fn }上面出现的所有东西都是 MLIR 节点。本小节将教你如何认识它们。经验法则每个 MLIR 节点都定义在某个.td文件中。例如lit.return由Mojo/include/Mojo/LITDialect/LITOps.td中的def LIT_ReturnOp定义lit.call由Mojo/include/Mojo/LITDialect/LITOps.td中的def LIT_CallOp定义lit.fn由Mojo/include/Mojo/LITDialect/LITOps.td中的def LIT_FnOp定义kgen.param.constant由Mojo/include/Mojo/KGENDialect/KGENOps.td中的def KGEN_ParamConstantOp定义。你通常可以通过全局搜索其def找到某事物定义所在的.td文件搜索模式如下^def.*lit.*return^def.*lit.*call^def.*lit.*fn^def.*kgen.*param.*constant以上都是 MLIR 操作operation我们倾向于把操作放在各个 pass 的操作定义文件中如LITOps.td、KGENOps.td、POPOps.td、HLCFOps.td等。那么操作之外的东西呢比如值、类型和编译期数据回忆 Passes and IR 中的基础 MLIR 记号规则%name— 运行时值一个 MLIR SSA 值某个 MLIR 操作的结果。name— 一个 SymbolRefAttr。!name— 一个 MLIR 类型。#expr或{expr}— 编译期数据 / 属性。其余任何东西都是 MLIR 操作。类型定义在LITTypes.td、KGENTypes.td、POPTypes.td等文件中编译期数据/属性定义在LITAttrs.td、KGENAttrs.td、POPAttrs.td等文件中。例如上面片段里的!kgen.none由KGENTypes.td中的def KGEN_NoneType定义#kgen.none由KGENAttrs.td中的def KGEN_NoneAttr定义。解剖一个 TableGen 定义以LIT_ReturnOp为例来看LITOps.td中的一个定义def LIT_ReturnOp : LITOpreturn, [] { let summary Lexical return statement.; ... let arguments (ins VariadicAnyType:$operands); let assemblyFormat $operands attr-dict (: type($operands)^)?; let hasVerifier 1; }这是一份 Operation Definition Specification即通常所说的 tablegen 文件。最重要的几个部分顶部的return这就是为什么它在 MLIR 中以lit.return出现let arguments 列出该操作的属性/字段。return只有一个字段名为operands它表示 Mojo 代码return 42, True中的42和Truelet assemblyFormat 描述该操作在 MLIR 中如何打印与解析。这就是它显示为lit.return %none : !kgen.none的原因——你能在assemblyFormat中看到$operands、:与type($operands)的对应关系。在源码中该定义位于 Mojo/include/Mojo/LITDialect/LITOps.td其description还详细说明了lit.return并非终结符terminator而是对应 Mojo 中的 return 语句用于建模函数可以有多个返回、且这些返回不必是代码块最后一条语句的事实。从 TableGen 到生成的 CLIT_CallOp的完整旅程编译器自身的构建系统会把.td文件转换成 C。以LIT_CallOp为例源码中的定义见 Mojo/include/Mojo/LITDialect/LITOps.tddef LIT_CallOp : LITOpcall, [ DeclareOpInterfaceMethodsKGEN_CallOpInterface] { let summary call a function; ... let arguments (ins AttrOfTypeLIT_FnTypeGeneratorType:$callee, KGEN_ParameterExprArrayAttr:$implicitOrigins, VariadicAnyType:$operands ); let results (outs VariadicAnyType:$results); let hasVerifier 1; let assemblyFormat [{ customCallOp($callee, $implicitOrigins, $operands, type($operands), type($results)) attr-dict }]; let extraClassDeclaration [{ /// Get the direct callee symbol if this is a direct call. SymbolRefAttr getDirectCallee(); /// Get the callee signature type. FnTypeGeneratorType getCalleeType() { return castFnTypeGeneratorType(getCallee().getType()); } }]; }它会被转换成生成文件LIT.h.inc中的 C 代码大致如下class CallOp : public ::mlir::OpCallOp, ::mlir::OpTrait::ZeroRegions, ... { public: ... static constexpr ::llvm::StringLiteral getOperationName() { return ::llvm::StringLiteral(lit.call); } ... ::mlir::TypedAttr getCallee(); ::llvm::ArrayRefTypedAttr getImplicitOrigins(); ... static void build(::mlir::OpBuilder odsBuilder, ::mlir::OperationState odsState, ::mlir::TypeRange results, ::mlir::TypedAttr callee, ::M::KGEN::ParameterExprArrayAttr implicitOrigins, ::mlir::ValueRange operands); ... ::llvm::LogicalResult verify(); static ::mlir::ParseResult parse(::mlir::OpAsmParser parser, ::mlir::OperationState result); void print(::mlir::OpAsmPrinter _odsPrinter); ... /// Get the direct callee symbol if this is a direct call. SymbolRefAttr getDirectCallee(); /// Get the callee signature type. FnTypeGeneratorType getCalleeType() { return castFnTypeGeneratorType(getCallee().getType()); } };注意两者的对应关系arguments中的各个字段callee、implicitOrigins、operands各自得到对应的辅助方法比如上面的getCalleeextraClassDeclaration中的文本getDirectCallee、getCalleeType也原样出现在最终 C 中。此外MLIR 还为我们生成了一些有用的辅助方法如verify、parse和print。实际定义中LIT_CallOp还带有KGEN_CallOpInterface接口声明与可选的tailKind属性DefaultValuedAttrKGEN_TailKindAttr, {}对应 Mojo 的尾调用支持。编译器自身如何与这些节点交互我们自己的 C 代码可以与这些节点交互。例如SharedState.cppinclude 了LITOps.h而后者又 include 生成的LIT.h.inc。SharedState.cpp中有这样一段使用call.getCallee()的代码if (auto call dyn_castLIT::CallOp(op)) { SmallVectorTypedAttr calleeOperands; calleeOperands.push_back(evaluator.getReboundAttribute(call.getCallee())); ...我们也可以在 MLIR 节点上添加自己的方法。例如上面的LIT_CallOp定义了extraClassDeclarationlet extraClassDeclaration [{ /// Get the direct callee symbol if this is a direct call. SymbolRefAttr getDirectCallee(); /// Get the callee signature type. FnTypeGeneratorType getCalleeType() { return castFnTypeGeneratorType(getCallee().getType()); } }];而我们自己在LITOps.cpp中定义getDirectCalleeSymbolRefAttr LIT::CallOp::getDirectCallee() { if (auto symbolCst dyn_castSymbolConstantAttr(getCallee())) return symbolCst.getSymbol(); return {}; }为厘清涉及的文件关系我们编写LITOps.tdLIT.h.inc由 MLIR 根据LITOps.td、LITAttrs.td等自动生成我们编写LITOps.h并在其中#includeLIT.h.inc我们编写LITOps.cpp我们编写编译器的其余部分如SharedState.cpp并通过 includeLITOps.h使用以上所有内容。Parameter Operator Code编译期查询编译器的机制Parameter Operator CodePOC是一组具名操作它们是POC枚举的变体。POC定义了#kgen.param.exprop, args...MLIR 属性所支持的操作名称。这一机制为 Mojo 代码提供了一种在编译期向编译器查询值的方式。它的用途多种多样从简单的操作如sizeof()到创新性的用例如compile_assembly一种把 Mojo 函数编译为汇编、并将该汇编嵌入最终二进制的机制。在KGENOps.td中定义如下节选def KGEN_POCAttr : I32EnumAttrPOC, Parameter Operator Code, [ /// Fully associative variadic expressions. I32EnumAttrCaseAdd, 0, add, I32EnumAttrCaseMul, 1, mul, I32EnumAttrCaseMulNoWrap, 2, mul_no_wrap, I32EnumAttrCaseAnd, 3, and, I32EnumAttrCaseOr, 4, or, I32EnumAttrCaseXor, 5, xor, I32EnumAttrCaseMax, 6, max, I32EnumAttrCaseMin, 7, min,这些枚举项以I32EnumAttr的形式定义定义位于Mojo/include/Mojo/KGENDialect/KGENOps.td每个I32EnumAttrCase名称, 编号, 字符串符号三元组把枚举值与其在 MLIR 文本中的符号形式关联起来供#kgen.param.expradd, ...这类属性在汇编文本中引用。数据包装器Data Wrappers值Values在解析器中几乎每个操作运行时表达式和参数表达式都会产生一个值。在基于 LLVM 的编译器中你可以用llvm::ValueRef引用该表达式而在基于 MLIR 的编译器中我们用mlir::Value引用它。……不过只要可能我们尽量避免直接用mlir::Value。当解析器知道自己处理的是哪种数据时我们的代码会更稳健。例如LIT::CallOp的操作数必须是运行时值而它的 callee 的参数必须是编译期值。来看处理 store如x y的代码。我们要确保x是一个 l-value即一个var而不是aliasy是一个具体值concrete value即具有已知大小。如果不检查这些编译器很容易产生导致崩溃的代码。于是我们这样做制作包装类LValue它只持有一个我们已知是 l-value的mlir::Value制作包装类CValue它只持有一个我们已知具有已知大小的mlir::Value让处理 store 的代码IREmitter::emitStoreToLValue只接受LValue和CValue如下CValue IREmitter::emitStoreToLValue(ASTExprAndCValue value, LValue destLV, ExprContext context) {现在调用方被迫确保源是具体值、目标是 l-value。LValue和CValue并不是仅有的两种值。SRValue只持有可寄存器传递register-passable的类型原始类型如int64、float32以及任何标记了register_passable已弃用改用RegisterPassabletrait或register_passable(trivial)已弃用改用TrivialRegisterPassabletrait的结构体参见 Life of Mojo reg-passable argumentsMValue只持有内存类型不可寄存器传递的东西如大多数结构体。事实上这里存在一整套分类体系。以下是你会遇到的各种值出自Mojo/include/Mojo/MojoParser/IRValues.h其头部注释中记录了完整层级AnyValue - Expr emitted to MLIR... UValue - unresolved value that cannot be materialized OverloadSetUValue - with an unresolved overload set InitializerUValue - constructor operands for an unknown type CValue - Concrete value: something with a known type. LValue - mutable reference to storage MLValue - value is in memory with a mutable reference DLValue - with dynamic get/set accessors BValue - with a borrowed value SBValue - value is register-passable and in an SSA register MBValue - value is in memory with a reference (may be mutable) MBPValue - reference with parametric mutability PValue - value is a parameter expression. RValue - with an owned value SRValue - with a register-passable value in an SSA register MRValue - value is in memory with a mutable reference PValue - with a parameter value说明这是文档写作时记录的分类。从当前仓库源码Mojo/include/Mojo/MojoParser/IRValues.h看该层级还在演化UValue下多了InferredBaseAttrRefUValue推断的 base 属性引用如.f64MBValue下多了PMBValuecomptime 内存借用值MRValue下多了PMRValuecomptime 内存右值LValue下多了RLValue指向ref x绑定以进行初始化的 LValue。此外源码注释还强调SRValue与内存专用类型不兼容但MRValue可以持有任何类型包括RegisterPassable类型。例如在下面这段 Mojo 代码中struct Spaceship: var hp: Int64 def launch(imm ship: Spaceship): var x: Int64 ship.hp我们的解析器会把ship视为MBValue因为它是内存类型并且我们以imm引用的方式借用它。以上内容主要适用于解析器。在后续 pass 中我们更多直接与 MLIR 节点以及 MLIR 内建类如mlir::Value和mlir::Type交互。ASTDeclMojo 解析器在很大程度上并不以 AST 的视角思考它处理的是 IR解析器直接消费文本惰性地在行进中 lex产出半扁平化的 IR一棵结构化控制流if、loop、try等的树其中包含 SSA 语句lit.call等的列表。然而尽管解析器并不_产出_ AST它确实会创建一个临时的 AST。对于每个作用域fn、struct、if、loop、try任何继承自ASTDeclInterface的东西我们都有一个ASTDecl为它做簿记。它的主要用途是跟踪成员直接拥有的子成员、从父 trait 继承的子成员以及任何从该作用域内部以某种方式可见的东西。然后解析器中的各个地方可以使用ASTDecl::lookupInCurrentScope查找匹配某个名字的成员。它还持有一个游标cursor用于惰性、渐进式地解析自身参见ASTDecl::getCursor。一个ASTDecl通常由一个操作operation支撑但也可以由一个CValue支撑。有一次我们观察到一个包含子成员 T 的struct MyStructT: AnyType其 T 是一个由ParamDeclRefAttr(T)PValue支撑的ASTDecl。因此我们最好的推测是当想存储预计算的引用ParamDeclRefAttr而不是指向实际声明的指针ParamDeclAttr时ASTDecl可以是CValue——这大概是出于性能考虑。Types解析器通常不直接与mlir::Type交互我们改用ASTType它带有很多有用的方法。HelpersBuilders、Transformers、Algorithms在编译器中你会在很多地方看到大量有用的工具。Walkers / Replacersmlir::AttrTypeWalker让我们遍历一个属性参数值及其间接包含的所有属性。mlir::AttrTypeReplacer一个可以边遍历边做替换的AttrTypeWalker。IndexParameterReplacer和ParameterReplacer更高级的AttrTypeReplacer它们是深度感知的知道当前参数引用节点引用的是父作用域中的东西还是别的东西。ParameterEvaluator一个带有替换列表的ParameterReplacer。ParserParameterEvaluator带一些便捷构造函数的ParameterEvaluator。函数调用器Function CallersOverloadSet帮助缩小你想调用的确切函数的类。如非必要尽量不要直接使用它。emitGetterSetterAccess用于进行字段访问如ship.hp或下标访问如my_list[42]。Witness Tables 与 ConformanceConformance 指我们检查一个结构体或 trait是否正确满足父 trait 的所有要求。这是通过doesNominalTypeConformTo完成的它还有一个副作用首先为结构体 trait配对或 trait trait 配对填好特定的 witness table即 vtable也叫ConformanceOp参见 CALROC。主要参与方doesNominalTypeConformTo调用DeclResolver::resolveBody(ConformanceOp, ASTDecl )后者又调用verifyConformance。更多细节参见 Conformance.md。分配Allocation我们很少使用new、make_unique或make_shared。经验法则分配一个 MLIR 操作时通过OpBuilder::create或ImplicitLocOpBuilder::create创建例如b.createLIT::ReturnOp(results)分配一个属性时我们常用它的::get静态方法例如BoolAttr::get(context, value)分配一个mlir::Type时我们通常调用其::get静态方法例如LIT::StructType::get(symbol, unbound, sig)。::get方法通常在我们的 tablegen.td文件中通过let builders [指令声明其实现定义在对应的.cpp文件中。如果你不确定某个Something是如何创建的可以搜索(alloc\w*|create)Something|Something::get解析器特有创建ASTType时它是一个轻量包装器、一个值类型给定mlir::Type可以直接构造有些可以从SharedState获取例如shared.lookupBuiltinType(Bool, declScope, loc)StructDeclOp::bindReference从StructDeclOp创建一个ASTType。分配不应逃逸出解析器调用的内存时使用SharedState::allocPersistent或ExprParser::alloc后者会替你调用allocPersistent。我们为表达式节点如allocBoolLiteralNode(startTok.getLoc(), false)和ASTDecl这样做。一个实用经验如果在解析语句/表达式时反复用到某个模式不妨把它提取成公共辅助函数。小结为 Mojo 编译器添加特性的工具箱回顾本文一个 Mojo 编译器开发者最常用的工具箱包括定位节点在Mojo/include/Mojo/*/下按 dialect 找到对应.td文件LITOps.td、KGENOps.td、POPOps.td、HLCFOps.td、LITTypes.td、KGENTypes.td、LITAttrs.td、KGENAttrs.td等通过^def.*全局搜索确认定义位置理解生成代码阅读生成的*.h.inc如LIT.h.inc理解arguments/results/assemblyFormat/extraClassDeclaration如何映射为 C 方法getCallee、verify、parse、print等并利用extraClassDeclaration在LITOps.cpp等文件中为节点添加自定义方法值分类用IRValues.h中的AnyValue→UValue/CValue→LValue/BValue/RValue层级来约束解析器对值语义的理解通过LValue/CValue等强类型包装把类型错误限制在编译期利用辅助工具IREmitter输出表达式、OverloadSet处理重载决议、AttrTypeWalker/ParameterReplacer族做深度感知的属性替换、doesNominalTypeConformTo处理 conformance 与 witness table遵循分配规范MLIR 操作用OpBuilder::create、属性与类型用::get、解析器内短生命周期内存用allocPersistent/ExprParser::alloc避免裸new/make_shared造成生命周期与所有权混乱。带着这套知识当你再次面对lit.fn、kgen.param.constant或某个MBValue时你不仅知道它们是什么还知道它们从哪里来、如何被创建、如何被操作——这正是不熟悉的编译器向你打开大门的方式。【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考