在 MLIR 中,如果你要自定义一个带有参数的复合类型(比如前面提到的!toy.struct<tensor<f64>, tensor<f64>>),你就必须要和TypeStorage(类型存储层)打交道。而StructTypeStorage就是专门用来在内存中存储这个结构体成员信息的“底层打工人”。
在编写编译器时,这部分逻辑往往是最容易让人绕晕的。我们可以用最通俗的语言,把它的核心原理和必须实现的三个硬核函数拆解清楚。
1. 为什么不能直接定义一个 C++ 类,而是需要 Storage?
在 MLIR 的世界里,有一个至高无上的铁律:所有的Type(类型)对象都是不可变的(Immutable),并且在全局是唯一化的(Unique)。
这意味着,如果在代码的不同地方出现了两个!toy.struct<tensor<f64>>,它们在内存里其实共享同一个地址。这样做能让编译器在做类型对比时(比如检查函数参数是否匹配),只需要极其高效地对比两个指针地址是否相等,而不用去深拷贝或者递归遍历整个结构体树。
为了实现这个机制,MLIR 采用了值对象模式(Flyweight Pattern,享元模式),把类型拆成了两层:
Toy::StructType(外壳):这是一个轻量级的包装类,里面其实只包含一个指向存储层的指针。它就像是一个“句柄”,暴露给用户使用,可以被随意复制。StructTypeStorage(内层存储):它才是真正躺在 MLIR 统一管理(MLIRContext)的堆内存里、存着具体数据的实体。
2.StructTypeStorage的硬核解构
一个标准的StructTypeStorage通常需要继承TypeStorage,它的核心任务就是存下结构体里有哪些子类型(即一个Type的数组ArrayRef<Type>)。
它的 C++ 核心模板结构和必须实现的三个魔法函数如下:
① 核心数据成员
structStructTypeStorage:publicTypeStorage{// 真正存数据的地方:一个分配在 Context 内存池中的类型数组ArrayRef<Type>elementTypes;// 构造函数StructTypeStorage(ArrayRef<Type>elementTypes):elementTypes(elementTypes){}};② 关键函数一:作为唯一标识的 Key
因为类型要全局唯一,MLIR 需要一个“Key”来去重。对于结构体来说,决定它唯一性的就是它的成员类型列表。所以它的 Key 直接用ArrayRef<Type>即可:
// 用来作为哈希表的键usingKeyTy=ArrayRef<Type>;// 检查当前存储的内容和新来的 Key 是否是一回事booloperator==(constKeyTy&key)const{returnkey==elementTypes;}③ 关键函数二:计算哈希值(Hash Key)
MLIR 底层是用一个类似于DenseMap的哈希表来做全局唯一化的。所以必须提供一个哈希函数,把结构体里所有成员的类型指针混合计算出一个哈希值:
staticllvm::hash_codehashKey(constKeyTy&key){// 混合计算数组中所有 Type 指针的哈希值returnllvm::hash_combine_range(key.begin(),key.end();}④ 关键函数三:内存分配与构造(Construct)
这是最神奇的一步。因为ArrayRef本身只是一个只读引用,它不拥有数据。当哈希表里找不到当前的结构体,需要创建新的 Storage 时,我们必须把传入的子类型数组,深拷贝到MLIRContext专门的内存分配器(Allocator)中。
staticStructTypeStorage*construct(TypeStorageAllocator&allocator,constKeyTy&key){// 1. 在 Context 的专有内存池里为 Type 数组占坑并拷贝数据ArrayRef<Type>elementTypes=allocator.copyInto(key);// 2. 在分配好的内存上原地构造(Placement New)Storage 对象returnnew(allocator.allocate<StructTypeStorage>())StructTypeStorage(elementTypes);}3. 现在的解放者:TableGen (ODS)
在早期(或者深入学习底层原理时),编译器开发者必须手动手写上面这一堆啰嗦的 C++ 样板代码(Boilerplate)。
不过,现代 MLIR 引入了强悍的TableGen (TypeDefs)机制。现在你只需要在.td文件里写几行配置,MLIR 的后台工具就会自动帮你生成上述一整套StructTypeStorage的 C++ 代码:
def Toy_StructType : TypeDef<Toy_Dialect, "Struct"> { let mnemonic = "struct"; // 你只需要声明参数,底层的 Storage、Key、Construct 统统由工具自动生成! let parameters = (ins ArrayRefParameter<"Type">:$elementTypes); }💡 总结
在编译器架构中,StructTypeStorage是实现底层类型唯一化(Type Uniquing)的功臣。它通过把类型数据“收拢”并深拷贝到专有的内存池中,配合哈希去重,保证了整个编译生命周期中,相同结构的复合类型在内存里独一无二。这也是为什么 MLIR 能够从容应对各种动辄包含成千上万个算子和复杂类型的大型计算图优化的底层基石。