【MLIR】跟着Toy语言学习MLIR【6】支持复合结构类型

Пусть этот камень будет более крепким, чем человек

【MLIR】Transform Dialect【0】保持计算结构【MLIR】跟着Toy语言学习MLIR【5】降级到LLVM及CodeGen 【MLIR】跟着Toy语言学习MLIR【4】部分降级到低层方言【MLIR】跟着Toy语言学习MLIR【3】通过接口实现通用转换【MLIR】跟着Toy语言学习MLIR【2】pattern匹配和重写【MLIR】跟着Toy语言学习MLIR【1】Toy语言和 Toy Dialect 【琐记】你的爱好能坚持多少年【琐记】烟火与尘埃【Triton】Triton实现矩阵乘【LLM推理加速】FlashAttention 【LLM推理加速】PagedAttention 【LLM推理加速】Online Softmax LLM基础知识【1】 Transformer模型【AI编译】LayerGroup Tiling Tile的疑惑和思考【AI编译】深度优先的Tile调度，万事大吉？【AI编译】多级流水线Tile调度策略【CUDA C++】GPU内存使用【3】【AI编译】Cache缓存地址映射【CUDA C++】GPU存储【2】【CUDA C++】GPU基本介绍【1】【00】0序章-不受欢迎的来客【转载】我来了——持续低熵【Halide】调度优化【2】【感想】写作进度报告5 【Halide】调度优化【1】【转载】北大中文男足战报2 【BYOC】TVM切分子图【转载】北大中文男足战报1 【AI编译】张量生命周期管理

XianMu · 2026-09-05 · via Пусть этот камень будет более крепким, чем человек

# 前言

现在开启 MLIR 学习系列。本篇是跟着 Toy 语言学习 MLIR 的第六篇，前述章节已经完整介绍了 toy 语言经 MLIR 降级到 llvm ir，并通过 JIT 进行编译执行的过程，本章节是教程中增添复合类型章节。前述内容请参考【MLIR】跟着 Toy 语言学习 MLIR【1】Toy 语言和 Toy Dialect，【MLIR】跟着 Toy 语言学习 MLIR【2】pattern 匹配和重写，【MLIR】跟着 Toy 语言学习 MLIR【3】通过接口实现通用转换。
相关链接： LLVM Project ，MLIR 官方文档，MLIR 官网教程，【编译器】使用 llvm 编译自定义语言【1】构建 AST ，【MLIR】跟着 Toy 语言学习 MLIR【2】pattern 匹配和重写，【MLIR】跟着 Toy 语言学习 MLIR【3】通过接口实现通用转换。
作为初学者，错误在所难免，还望不吝赐教。

# 基本简介

官网 Toy 教程展示了如何将自定义语言 Toy 借助 MLIR 一步步编译为可执行机器码的过程。Toy 是一种简单的自定义语言，为了简便，其所有数据类型定义为 fp64 类型的 Tensor，支持 +/* 操作和 transpose 等有限的操作。

以下是整个编译降级的过程：

Toy txt -> Toy AST -> Toy Dialect -> Affine Dialect -> llvm Dialect -> llvm IR -> 机器码（通过 JIT 编译）

编译流程图

# 在 Toy 中定义 Struct 类型

Struct 在 Toy 语言中的定义非常简单，和 C++ 语言中的定义与初始化差不多：

	# A struct is defined by using the `struct` keyword followed by a name.
	struct MyStruct {
	# Inside of the struct is a list of variable declarations without initializers
	# or shapes, which may also be other previously defined structs.
	var a;
	var b;
	}

这是给的 toy 语言例子，包含了 struct 的定义初始化和调用。

	struct Struct {
	var a;
	var b;
	}

	# User defined generic function may operate on struct types as well.
	def multiply_transpose(Struct value) {
	# We can access the elements of a struct via the '.' operator.
	return transpose(value.a) * transpose(value.b);
	}

	def main() {
	# We initialize struct values using a composite initializer.
	Struct value = {[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]};

	# We pass these arguments to functions like we do with variables.
	var c = multiply_transpose(value);
	print(c);
	}

# 在 MLIR 中定义 struct

在 MLIR 中，我们也需要为结构体类型定义一种表示方式。MLIR 本身没有提供完全满足我们需求的类型，因此我们需要自行定义。
可以简单地将结构体定义为一组元素类型的无名容器。

# 定义 Type Class

当我们需要定义一个新的复杂类型（如 struct）时，需要提供自定义的存储类（Storage Class）
MLIR 中的类型 Type 是 “值类型”，Type 对象本身是一个轻量级的包装器，它不存储实际的数据，而是指向一个内部的存储对象。


	class Type {
	TypeStorage *impl;

	};

举例来说，String 在 Java 中也是值类型，在 Java 中，String 对象是不可变的，多个 String 变量可以共享同一个底层字符数组。
Type 类本身只是 TypeStorage 对象的一个简单包装器，而这个 TypeStorage 对象是在 MLIRContext 中被唯一化的。“唯一化” 意味着：同一种类型在同一个上下文中只会有一个存储实例。这种设计，节省内存：同一种类型只存一份，避免重复；快速比较：比较两个类型只需要比较指针地址，而不是逐字段比较。

Defining the Storage Class

什么时候需要定义存储类？当需要定义一个包含参数化数据的新类型（如 struct 类型，需要存储元素类型列表）时，你需要提供一个派生的存储类。而那些没有额外数据的单例类型（如 index 类型，i32 类型，f64 类型等）则不需要自定义存储类，直接使用默认的 TypeStorage。
类型存储对象包含构建和唯一化类型实例所需的所有数据。派生的存储类必须继承自基础的 mlir::TypeStorage，并提供一组别名和钩子，供 MLIRContext 用于唯一化处理。以下是结构体类型 struct 的存储实例定义，可以在 Dialect.cpp 中找到：


	struct StructTypeStorage : public mlir::TypeStorage {




	using 核心概念 = llvm::ArrayRef<mlir::Type>;


	StructTypeStorage(llvm::ArrayRef<mlir::Type> elementTypes)
	: elementTypes(elementTypes) {}




	bool operator==(const KeyTy &key) const { return key == elementTypes; }





	static llvm::hash_code hashKey(const KeyTy &key) {
	return llvm::hash_value(key);
	}






	static KeyTy getKey(llvm::ArrayRef<mlir::Type> elementTypes) {
	return KeyTy(elementTypes);
	}





	static StructTypeStorage *construct(mlir::TypeStorageAllocator &allocator,
	const KeyTy &key) {

	llvm::ArrayRef<mlir::Type> elementTypes = allocator.copyInto(key);


	return new (allocator.allocate<StructTypeStorage>())
	StructTypeStorage(elementTypes);
	}


	llvm::ArrayRef<mlir::Type> elementTypes;
	};

这个 struct 的存储实例，是为了支持 toy 语言中的 struct 语法。有了它，toy 语言就可以定义各种结构的 struct。
keyTy 是 MLIR 的核心概念，它定义了用于唯一化存储实例（一个存储类型）的键类型。对于 StructType，唯一化的依据是它的元素类型列表，当 MLIRContext 在创建类型时，会检查是否已经存在相同键的存储实例。如果有，直接返回已存在的实例，否则创建新的。
构造函数接收元素类型列表，并将其存储到 elementTypes 成员变量中。元素类型类型列表决定了当前定义的 struct 的唯一性。
相等性比较，看起来好像是比较两个定义实例是否相同的，防止重复定义。
哈希函数：为键计算哈希值，用于在哈希表中快速定位存储实例。
键构建函数 getKey 将键列表构建为 KeyTy ，当然 KeyTy 本身就是键列表。
存储构造函数：MLIR 框架要求提供的一个 “工厂方法”，职责是 “构造存储实例”（Construct a Storage Instance）。

Defining the Type Class
定义类型。
在定义好存储类后，我们可以添加用户可见的 StructType 类的定义。这个类是我们实际进行交互的对象。






	class StructType : public mlir::Type::TypeBase<StructType, mlir::Type, StructTypeStorage> {
	public:

	using Base::Base;



	static StructType get(llvm::ArrayRef<mlir::Type> elementTypes) {
	assert(!elementTypes.empty() && "expected at least 1 element type");




	mlir::MLIRContext *ctx = elementTypes.front().getContext();
	return Base::get(ctx, elementTypes);
	}


	llvm::ArrayRef<mlir::Type> getElementTypes() {

	return getImpl()->elementTypes;
	}


	size_t getNumElementTypes() { return getElementTypes().size(); }
	};

Type 类定义遵循了 MLIR 中 CRTP（奇异递归模板模式，Curiously Recurring Template Pattern）的设计模式。它帮助实现了一些功能，如唯一化（Uniquing）机制，类型转换（cast、dyn_cast）支持，上下文管理等。
StructType 对象是轻量级的，它们只包含一个指向 StructTypeStorage 的指针。

多个 StructType 对象可以共享同一个 StructTypeStorage 实例（如果它们的 KeyTy 相同）。

然后在 toy dialect 中注册这个自定义类型：

	void ToyDialect::initialize() {
	addTypes<StructType>();
	}

Exposing to ODS
当定义了一个新的类型，我们需要确保 ODS 框架能够 “看得到” 这个新类型，以便在操作定义和方言中的自动生成工具中使用它。以下是一个简单的示例：




	def Toy_StructType :
	DialectType<Toy_Dialect, CPred<"isa<StructType>($_self)">,
	"Toy struct type">;


	def Toy_Type : AnyTypeOf<[F64Tensor, Toy_StructType]>;

它是如何绑定到 C 代码中的自定义类型 StructType 的呢？ DialectType<Toy_Dialect, ...> 指明这是一个属于 Toy_Dialect 方言的类型。CPred 是连接 ODS 和 C 的 “桥梁”，通过 CPred<"isa<StructType>($_self)"> 这个谓词（Predicate）与 C++ 中的 StructType 类绑定在一起的。 isa<StructType>($_self) 是一个 C++ 表达式，用于检查一个类型是否是 StructType 。

Parsing and Printing
此时，可以在 MLIR 生成和转换过程中使用 StructType，但无法输出或解析 .mlir 文件。为此，我们需要为 StructType 的实例添加解析和打印支持。这可以通过重写 ToyDialect 中的 parseType 和 printType 方法来实现。

	class ToyDialect : public mlir::Dialect {
	public:

	mlir::Type parseType(mlir::DialectAsmParser &parser) const override;


	void printType(mlir::Type type,
	mlir::DialectAsmPrinter &printer) const override;
	};

根据 MLIR 语言参考文档所述，方言类型通常表示为：! dialect-namespace <type-data>，在某些情况下可提供更美观的格式。所以教程将 struct 类型的打印定义为如下形式：

struct-type ::= `struct` `<` type (`,` type)* `>`

那么解析（Parser）就可以定义为：


	mlir::Type ToyDialect::parseType(mlir::DialectAsmParser &parser) const {









	if (parser.parseKeyword("struct") \|\| parser.parseLess())
	return Type();


	SmallVector<mlir::Type, 1> elementTypes;
	do {

	SMLoc typeLoc = parser.getCurrentLocation();
	mlir::Type elementType;
	if (parser.parseType(elementType))
	return nullptr;


	if (!isa<mlir::TensorType, StructType>(elementType)) {
	parser.emitError(typeLoc, "element type for a struct must either "
	"be a TensorType or a StructType, got: ")
	<< elementType;
	return Type();
	}
	elementTypes.push_back(elementType);


	} while (succeeded(parser.parseOptionalComma()));


	if (parser.parseGreater())
	return Type();
	return StructType::get(elementTypes);
	}

解析类，帮助从文本中反序列化回内存对象。
一个打印（Printer）的实现可以这样定义：


	void ToyDialect::printType(mlir::Type type, mlir::DialectAsmPrinter &printer) const {

	StructType structType = type.cast<StructType>();


	printer << "struct<";
	llvm::interleaveComma(structType.getElementTypes(), printer);
	printer << '>';
	}

让 MLIR 能够将 StructType 序列化为文本（如 .mlir 文件）

这是 toy 语言中的 strcut 定义：

	struct Struct {
	var a;
	var b;
	}

	def multiply_transpose(Struct value) {
	}

它将会生成：

	module {
	toy.func @multiply_transpose(%arg0: !toy.struct<tensor<xf64>, tensor<xf64>>) {
	toy.return
	}
	}

它也支持从 mlir 文本到内存对象的双向转换。

Operating on StructType
现在已经定义好了 strcut 这个类型，现在要做的是让我们 toy 方言中已有的操作支持 strcut 这个新的类型。
例如我们已经存在的操作 ReturnOp ，需要进行如下更新:

	def ReturnOp : Toy_Op<"return", [Terminator, HasParent<"FuncOp">]> {
	...
	let arguments = (ins Variadic<Toy_Type>:$input);
	...
	}

将原来的类型约束从 F64Tensor 变为了 Toy_Type ，而 Toy_Type 是个复合类型约束，它表示 F64Tensor 或 StructType 。
Adding New Toy Operations
除了现有操作外，还需要添加一些新的操作 Operation 来更好的处理 truct 。
struct 常量操作 toy.struct_constant。

	%0 = toy.struct_constant [
	dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]> : tensor<2x3xf64>
	] : !toy.struct<tensor<*xf64>>

其实在这里我有个疑问，为什么 "教程选择创建一个新操作 struct_constant，而不是扩展现有的 ConstantOp (复用)，将其支持 struct" 呢？
问 llm，得到的答复是：解析 / 打印变得复杂、验证逻辑变得复杂、类型推导变得复杂、可读性下降。所以没有复用原来的常量操作。

struct 访问结构体成员操作 toy.struct_access


	%1 = toy.struct_access %0[0] : !toy.struct<tensor<xf64>> -> tensor<xf64>

这个是为了表示 toy 语言中 Value.a， Value.b 这种访问数值的操作。
这时候新的 toy 语言示例：

	struct Struct {
	var a;
	var b;
	}

	# User defined generic function may operate on struct types as well.
	def multiply_transpose(Struct value) {
	# We can access the elements of a struct via the '.' operator.
	return transpose(value.a) * transpose(value.b);
	}

	def main() {
	# We initialize struct values using a composite initializer.
	Struct value = {[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]};

	# We pass these arguments to functions like we do with variables.
	var c = multiply_transpose(value);
	print(c);
	}

能够得到 mlir 打印输出：

	module {
	toy.func @multiply_transpose(%arg0: !toy.struct<tensor<xf64>, tensor<xf64>>) -> tensor<*xf64> {
	%0 = toy.struct_access %arg0[0] : !toy.struct<tensor<xf64>, tensor<xf64>> -> tensor<*xf64>
	%1 = toy.transpose(%0 : tensor<xf64>) to tensor<xf64>
	%2 = toy.struct_access %arg0[1] : !toy.struct<tensor<xf64>, tensor<xf64>> -> tensor<*xf64>
	%3 = toy.transpose(%2 : tensor<xf64>) to tensor<xf64>
	%4 = toy.mul %1, %3 : tensor<*xf64>
	toy.return %4 : tensor<*xf64>
	}
	toy.func @main() {
	%0 = toy.struct_constant [
	dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>,
	dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>
	] : !toy.struct<tensor<xf64>, tensor<xf64>>
	%1 = toy.generic_call @multiply_transpose(%0) : (!toy.struct<tensor<xf64>, tensor<xf64>>) -> tensor<*xf64>
	toy.print %1 : tensor<*xf64>
	toy.return
	}
	}

# 操作在 StructType 方面的优化

现在我们已经有一些操作支持了 StructType 类型，也有了一些新的进行常量折叠的机会。
在内联 inline 之后，当前的 mlir 输出会是如下这个样子：

	module {
	toy.func @main() {
	%0 = toy.struct_constant [
	dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>,
	dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>
	] : !toy.struct<tensor<xf64>, tensor<xf64>>
	%1 = toy.struct_access %0[0] : !toy.struct<tensor<xf64>, tensor<xf64>> -> tensor<*xf64>
	%2 = toy.transpose(%1 : tensor<xf64>) to tensor<xf64>
	%3 = toy.struct_access %0[1] : !toy.struct<tensor<xf64>, tensor<xf64>> -> tensor<*xf64>
	%4 = toy.transpose(%3 : tensor<xf64>) to tensor<xf64>
	%5 = toy.mul %2, %4 : tensor<*xf64>
	toy.print %5 : tensor<*xf64>
	toy.return
	}
	}

上述 mlir 中有很多针对 toy.struct_constant 的 toy.struct_access ，如果没有常量折叠，值需要在运行时通过 struct_access 操作获取；而有了常量折叠，编译器可以直接将其替换为常量。
在 CH7 章节 ToyCombine.cpp 能找到以下代码。


	OpFoldResult ConstantOp::fold(FoldAdaptor adaptor) { return value(); }


	OpFoldResult StructConstantOp::fold(FoldAdaptor adaptor) {
	return value();
	}


	OpFoldResult StructAccessOp::fold(FoldAdaptor adaptor) {
	auto structAttr = dyn_cast_or_null<mlir::ArrayAttr>(adaptor.getInput());
	if (!structAttr)
	return nullptr;

	size_t elementIndex = index().getZExtValue();
	return structAttr[elementIndex];
	}

OPS 时定义指定该操作可以常量折叠。





	def ConstantOp : Toy_Op<"constant",
	[ConstantLike, Pure,
	DeclareOpInterfaceMethods<ShapeInferenceOpInterface>]> {
	...

	let hasFolder = 1;
	}

常量物化（Constant Materialization）

常量物化是指在折叠过程中，当需要创建一个新的常量时，如何生成对应的常量操作。也就是折叠 struct_access 后，我们需要创建一个新的常量来表示结果。但是，这个结果可能是一个张量（用 toy.constant）或者一个结构体（用 toy.struct_constant）。MLIR 框架不知道如何为 Toy 方言创建常量，因为它不知道 toy.constant 和 toy.struct_constant 的区别。

	mlir::Operation *ToyDialect::materializeConstant(mlir::OpBuilder &builder,
	mlir::Attribute value,
	mlir::Type type,
	mlir::Location loc) {
	if (isa<StructType>(type))
	return StructConstantOp::create(builder, loc, type,
	cast<mlir::ArrayAttr>(value));
	return ConstantOp::create(builder, loc, type,
	cast<mlir::DenseElementsAttr>(value));
	}

此后生成的 mlir :

	module {
	toy.func @main() {
	%0 = toy.constant dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>
	%1 = toy.transpose(%0 : tensor<2x3xf64>) to tensor<3x2xf64>
	%2 = toy.mul %1, %1 : tensor<3x2xf64>
	toy.print %2 : tensor<3x2xf64>
	toy.return
	}
	}

# 后记

	日落是免费的，春夏秋冬也是、不要觉得人生是那么无望，希望你快乐。
	我们总是为许多遥不可及的事情奔波，却错过了路边的花开，傍晚落在身上的夕阳。
	忙着生活的同时记得去感受生活中的小细节，生活除了琐碎与平淡，还有可口的美食和无数盛开的花朵。
	晚风吹人醒、万事藏于心、我没说不公平、也没说苦、我说我知道了。

本博客目前以及可预期的将来都不会支持评论功能。各位大侠如若有指教和问题，可以在我的 github 项目或随便一个项目下提出 issue，并指明哪一篇博客，看到一定及时回复！

此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。原文来自 — 版权归原作者所有。

推荐订阅源

Пусть этот камень будет более крепким, чем человек

# 前言

# 基本简介

# 在 Toy 中定义 Struct 类型

# 在 MLIR 中定义 struct

# 定义 Type Class

# 操作在 StructType 方面的优化

# 后记