Mojo 编译器泛型处理内幕解析器中泛型的四个非显而易见机制IRAIDAI / DCRTODS / PSTIAIRAID / STCHDDDOS【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo本文聚焦 Mojo 编译器源码仓库中Mojo/docs/compiler/arcana/Generics.md所记录的泛型解析核心机制IndexRefAttrInterface基于索引的参数引用、深度depth与索引index语义、ParameterScopeTypeInterface对深度的作用以及同型不同深度带来的相等性比较陷阱。读完本文你将理解 Mojo 编译器如何在解析阶段用相对位置代替参数名来描述泛型引用掌握编写或阅读编译器泛型相关代码如IndexParameterReplacer、ParameterReplacer、ParameterEvaluator时必备的深度感知depth-aware思维并清楚为何 MLIR 内置的.walk不能用于此类遍历。引言为什么泛型处理在解析器中如此非显而易见Mojo 是一门支持强参数多态parametric polymorphism的语言其泛型参数如T: AnyType、N: Int在编译器前端解析、KGEN 方言生成、求值等各阶段都被大量引用。而泛型处理的难点在于同一个类型或表达式里引用的参数可能来自不同层级的嵌套签名且这些引用需要在不依赖参数名的情况下保持语义一致、可判等。Generics.md是 Mojo 编译器秘辛arcana系列文档之一它用四个缩略词概括了解析器处理泛型时的四个非显而易见之处IRAIDAIIndexRefAttrInterface、深度与索引Depths and IndexesDCRTODS深度不能引用 Op 声明的签名Depths Cannot Refer To Op-Declared SignaturesPSTIAIRAIDParameterScopeTypeInterface影响IndexRefAttrInterface的深度STCHDDDOS同一类型在不同作用域可以有不同的深度Same Type Can Have Different Depths Depending On Scope本文将以这四个机制为骨架结合仓库中 KGENDialect 与 MojoParser 的源码实现Mojo/include/Mojo/KGENDialect/KGENAttrInterfaces.td、Mojo/include/Mojo/KGENDialect/KGENAttrs.td、Mojo/include/Mojo/KGENDialect/KGENTypeInterfaces.td、Mojo/include/Mojo/KGENDialect/ParameterReplacer.h逐一展开。IRAIDAIIndexRefAttrInterface、深度与索引两个等价的别名考虑以下两个 Mojo 别名alias它们的类型是完全相等的alias A: defT: AnyType-None ... alias B: defY: AnyType-None ...直觉上我们知道A和B是同一个函数类型——参数名T与Y只是占位符实际语义没有差别。但要让编译器知道它们相等一个简洁的办法是把名字擦除改用相对位置来描述参数引用。于是上面的两个别名可以写成alias A: def_: AnyType)-None ... alias B: def_: AnyType)-None ...这里*(0,0)就是一个索引式参数引用index-based parameter reference其中第一个分量是深度depth第二个分量是索引index。IndexRefAttrInterface的两个字段在 KGENDialect 的 TableGen 接口定义中Mojo/include/Mojo/KGENDialect/KGENAttrInterfaces.tdIndexRefAttrInterface的官方描述指出这是一种用一对整数、完全不涉及名字的相对参数引用方案其价值在于即便参数名不同也能判定两个类型是否相等Index-based parameter references are a relative parameter referencing scheme that uses a pair of integers to reference parameters in a way that doesnt involve names. This is useful for later knowing if two types are equal, even if they have different parameter names.IndexRefAttrInterface的所有实现都包含两个字段字段含义取值范围depth目标参数声明位于第几层外层签名ParameterScopeTypeInterface/ParameterScopeAttrInterface中。0表示最近的包含签名1表示再往上一层非负整数index该参数声明在目标签名中的序号第几个输入参数非负整数接口还暴露了三个方法对应 TableGen 中的InterfaceMethodgetDepth()返回深度、getIndex()返回索引、replace()在不改变被引用参数语义的前提下修改 depth/index 并携带新的子元素。对应到具体属性实现Mojo/include/Mojo/KGENDialect/KGENAttrs.td#kgen.param.index.ref的 MLIR 文本形式即#kgen.param.index.refdepth, index : type例如// 最近签名中的第二个输入参数 #kgen.param.index.ref0, 1 : index // 下一层外层签名的第一个输入参数 #kgen.param.index.ref1, 0 : !lit.structInt非零深度示例SIMD[N, D]中的N文档给出了一个体现非零深度的典型例子alias bar: def D: DType, N: Int, f: def[Y: AnyType-None ](...) ...这里SIMD[N, D]中的N是一个#kgen.param.index.ref1, 0 : !lit.structInt有时写成*(1,0)。原因在于N并没有引用最近的包含签名即f所在的那个def[Y: AnyType]而是引用它的外层签名bar自身的签名。在参数引用与参数声明之间隔着1 层签名所以其深度为1索引为0N是bar签名的第 2 个参数从 0 起算。原文档特别提醒计算depth时要格外小心很容易算错。并非所有参数引用都用索引需要澄清的是不是所有参数引用都采用索引与深度。还有一种传统的ParamDeclRefAttr#kgen.param.decl.ref它仍然按名字引用参数。其定义见Mojo/include/Mojo/KGENDialect/KGENAttrs.td// 对名为 p、类型为 i1 的参数的引用 #kgen.param.decl.refp : i1ParamDeclRefAttr是一个带类型的属性包含被引用参数的类型和名字其类型必须始终与被引用参数的类型一致。至于什么时候必须用哪种引用就是下一节 DCRTODS 的主题。DCRTODS深度不能引用 Op 声明的签名一个反直觉的例子考虑下面的代码def fooX: AnyType: alias bar: defY: AnyType ...按照 IRAIDAI 的直觉你可能会认为bar的类型是defAnyType, *(0,0))——因为类型一般不会包含ParamDeclRefAttr。但这个推断是错误的。bar的真实类型是defAnyType)也就是说Y的引用被替换成了*(0,0)而X的引用保持为名字X不变。为什么规则看参数声明由谁发出关键规则只有一条当引用的是某个 Op 的参数声明时使用ParamDeclRefAttr按名字。当引用的是某个 Type 或 Attr 的参数声明时使用ParamIndexRefAttr按深度索引。回到例子把涉及的参数声明逐一归类参数声明由谁声明引用方式Xdef fooOpParamDeclRefAttr保留名字Xbaralias barOpParamDeclRefAttrY: AnyTypefnTypeParamIndexRefAttr*(0,0)X和bar都是由 Op 声明的因此用ParamDeclRefAttr按名字引用Y是由函数类型声明的因此用ParamIndexRefAttr按索引引用。在Mojo/include/Mojo/KGENDialect/KGENAttrs.td中ParamIndexRefAttr的注释也给出了同样的反例comptime zork版本并明确标注了禁止出现的写法def fooX: AnyType: comptime zork: def...( # 不允许 #kgen.param.index.ref1, 0 : !lit.structInt )-None ...经验法则原文档给出的经验法则非常简洁深度depth只能指向包含它的GeneratorType和GeneratorAttr对于其他一切比如X这样的 Op 参数使用ParamDeclRefAttr。这条规则也是后续所有深度处理代码PSTIAIRAID、STCHDDDOS的前提。PSTIAIRAIDParameterScopeTypeInterface影响IndexRefAttrInterface的深度深度的两种情形回顾 IRAIDAIdepth 0表示引用最近的外层签名例如alias A: def[T: AnyType-None ...其中x: T的T就是*(0,0)深度 0索引 0。depth 1表示引用再外层一层的签名例如前面SIMD[N, D]中的Nalias bar: def D: DType, N: Int, f: def[Y: AnyType-None ](...) ...机制所有签名都继承ParameterScopeTypeInterface为了管理深度编译器让所有会声明参数的签名类型统一继承ParameterScopeTypeInterfaceMojo/include/Mojo/KGENDialect/KGENTypeInterfaces.td代码在处理深度时专门监视该接口。该接口的官方描述用一个更完整的例子说明了其语义def foo[T: AnyType](): comptime bork: def T: AnyType, inner_f: def[Y: AnyType - None ] - None ...bork:之后的def是一个kgen.generator它实现了ParameterScopeTypeInterface。描述中还明确指出ParameterScopeTypeInterface也会导致其内部即便是间接包含的ParamIndexRefAttr/ImplicitOriginRefAttr等的depth字段变大见 PSTIAIRAID。也就是说每进入一层签名内部引用的深度就加 1——这正是深度感知遍历的基础。同理属性一侧对应ParameterScopeAttrInterfaceMojo/include/Mojo/KGENDialect/KGENAttrInterfaces.td它让属性内部的ParamIndexRefAttr可以引用所在作用域声明的参数。深度感知搜索depth-aware searching文档给出的例子alias bar: def[ T: AnyType, L: List[T], f: defY: AnyType-None ](...) ...它被解释为alias bar: def[ _: AnyType, _: List[*(0,0)], _: def_: AnyType, List[*(1,0)])-None ](...) ...如果我们要查找代码中所有对T的引用不能简单地在整棵树上找*(0,0)——那样会同时命中第一个T正确和内层签名中Y的引用错误。正确做法是在递归搜索过程中记录当前处于第几层签名。在外层签名中找*(0,0)一旦搜索下潜进内层签名就改为找*(1,0)。这正是文档所说的深度感知搜索depth-aware searching它在处理任何可能哪怕间接包含签名的对象时都至关重要。结论慎用AttrTypeReplacer与AttrTypeWalker提防 MLIR 内置.walk文档给出两条极具实战价值的警告三思而后用AttrTypeReplacer或AttrTypeWalker——它们不是深度感知的。应当优先考虑IndexParameterReplacer、ParameterReplacer、ParameterEvaluator、ParserParameterEvaluator、IndexRefRemapper等深度感知工具。MLIR 的内置.walk方法不是深度感知的原文用了加粗强调——直接用它遍历带签名嵌套的结构会算错深度。以IndexParameterReplacer为例Mojo/include/Mojo/KGENDialect/ParameterReplacer.h其类注释与实现印证了上述设计把这个depth传给replaceImpl是这个类的核心目的它让replaceImpl的实现知道在递归遍历中目前深入签名作用域多少层。例如它们可以检查depth 0判断是否在原始作用域或检查indexRef.getDepth() depth判断该 indexRef 是否引用原始作用域的参数声明。其doReplace的递归实现展示了深度递增的关键逻辑// Increment depth when looking inside signatures, see PSTIAIRAID. if constexpr (std::is_base_of_vAttribute, T) if (isaParameterScopeAttrInterface(value)) depth; if constexpr (std::is_base_of_vType, T) if (isaParameterScopeTypeInterface(value)) depth;即每当递归进入一个实现了ParameterScopeAttrInterface/ParameterScopeTypeInterface的子对象就把当前深度加一再向下传递。这与 PSTIAIRAID 描述的语义完全吻合也解释了为什么这些 replacer 能正确处理嵌套签名而 MLIR 原生的.walk不能。STCHDDDOS同一类型在不同作用域可有不同深度同型不同 depth继续沿用 PSTIAIRAID 的例子alias bar: def[ T: AnyType, L: List[T], f: defY: AnyType-None ](...) ...解释后alias bar: def[ _: AnyType, _: List[*(0,0)], _: def_: AnyType, List[*(1,0)])-None ](...) ...注意List[T]现在出现了两个不同的类型表示List[*(0,0)]外层签名中的TList[*(1,0)]内层签名中引用的外层T尽管它们语义上是同一个类型。结论相等的类型不总是指针可比较的文档指出这意味着相等的类型并不总是指针可比较的。这颇具讽刺意味——因为引入深度的初衷按 IRAIDAI恰恰是为了便于判定两个类型是否相等源码接口描述中的 This is useful for later knowing if two types are equal。于是编译器需要在多处对这种同型不同深度做特殊处理。比较前的两种调整策略当需要比较List[*(0,0)]与List[*(1,0)]是否相等时必须先统一深度二选一递减更深签名里引用的深度把List[*(1,0)]中的深度从 1 减为 0递增外层签名里引用的深度把List[*(0,0)]中的深度从 0 增为 1。只有完成这样的深度统一之后才能安全地比较两个类型是否相等。典型触发场景实例化参数化值STCHDDDOS-A / STCHDDDOS-B文档给出了两个必须警惕的关键场景场景 A实例化一个参数化生成器值generator value。例如给定一个生成器值index *(0,0) 1对任意索引类型参数返回其加 1可以把外层作用域的索引引用绑定进去如bind_params(index *(0,0) 1, *(0,1))。对bind_params算子做折叠会得到一个新的、输入参数为 0 个的生成器值 *(1,1) 1。注意它仍然是一个生成器值必须再执行一次显式的实例化步骤来解开生成器而这次解包需要对生成器体内的所有索引引用做一次深度递减——因为此时生成器体处于少一层的生成器作用域之下。场景 Bapply算子作用于已实例化的参数化函数。当被调方是一个已实例化的参数化函数时产生的函数类型同样需要对其中的索引引用做深度递减文档中称之为up-binding。这个操作与场景 A 是同一类问题在另一个算子上的体现。这两个场景的共同点在于任何跨越签名边界解包/折叠参数化结构的操作都必须同步调整内部索引引用的深度否则会产生悬挂的、语义错误的深度值。实战总结处理泛型深度时需要记住的清单综合原文档与仓库实现处理 Mojo 解析器 / KGEN 方言中的泛型时建议按以下清单自查先分清引用种类引用 Op 声明的参数用ParamDeclRefAttr按名字引用 Type/Attr 声明的参数用ParamIndexRefAttr按depthindex。深度只能指向包含它的GeneratorType/GeneratorAttrDCRTODS。算深度时数清签名层数0是最近的签名每向外一层 1计算时容易出错IRAIDAI。遍历优先用深度感知工具如IndexParameterReplacer、ParameterReplacer、ParameterEvaluator、ParserParameterEvaluator、IndexRefRemapper避免AttrTypeReplacer/AttrTypeWalker绝不要用 MLIR 内置.walkPSTIAIRAID。比较相等性前先统一深度要么递减深层引用要么递增外层引用然后再判等要留意bind_params折叠、apply实例化等会改变深度层次的操作STCHDDDOS-A/B。这些机制的最终目标是在不依赖参数名的前提下让形异义同的泛型类型可以判等、替换与求值——理解depth/index这对相对坐标是读懂 Mojo 编译器泛型前端MojoParser 与 KGENDialect代码的一把钥匙。感兴趣的读者可以继续在仓库中阅读Mojo/include/Mojo/KGENDialect/ParameterEvaluator.h其中第 2 条规则专门描述ParamIndexRefAttr的替换规则、Mojo/include/Mojo/KGENDialect/KGENParameters.h递归遍历并替换ParamIndexRefAttr、以及ParamIndexRefAttrFinder检测器以及Mojo/test/mojo-parser/parameters.mojo中对应的解析器测试用例来加深理解。【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考