# 前言

本篇是学习 MLIR 中的 Transform Dialect 的第四篇，本章主要讲的是 MLIR Transform Dialect 的进阶用法，核心围绕如何用更规范、更类型安全的方式编写 transform 操作，以及如何定义自己的 transform 类型。内容来自 MLIR 官方教程《Chapter 3: More than Simple Transform Operations》。
相关链接： LLVM Project ，MLIR 官方文档，MLIR 官网教程，《Chapter 0: A Primer on “Structured” Linalg Operations》。
作为初学者，错误在所难免，还望不吝赐教。

# 简介

Transform 操作中的重复模式。
在 Transform Dialect 中，很多 transform 操作都有这样的特点：对 handle 关联的每一个 payload 操作单独应用同一个变换，并且要求这些 payload 操作必须是某种特定类型。例如：把每个 func.call 的被调用者改成另一个符号；给每个 scf.for 加上某个属性；把每个 linalg.matmul 替换成自定义实现。如果每个这样的操作都手写完整的 apply 方法，就会出现大量重复代码：

	// 伪代码：每个操作都要写一遍
	DiagnosedSilenceableFailure MyOp::apply(...) {
	for (Operation *op : state.getPayloadOps(getHandle())) {
	// 1. 验证 op 类型
	auto call = dyn_cast<func::CallOp>(op);
	if (!call) return emitSilenceableError() << "not a func.call";
	// 2. 执行变换
	updateCallee(call, getNewTarget());
	// 3. 处理结果（这里没有结果）
	}
	return success();
	}

这些步骤 —— 验证类型、遍历 payload、合并结果 —— 是完全模式化的，应该由框架自动完成，而不是每个操作都重复写。
MLIR 提出两个机制来解决这个问题。

1. 用 Transform 类型收窄操作数约束

原本 handle 的类型是宽泛的 TransformHandleTypeInterface，它可以指向任意 payload 操作。现在改为更具体的： Transform_ConcreteOpType<"func.call">:$call 。
这表示：这个 handle 只能关联 func.call 类型的 payload 操作。如果实际关联的不是 func.call，在验证阶段就会失败，不需要在 apply 里手动检查。
这样做的好处：类型安全：约束写在 ODS 定义里，而不是散落在 C++ 实现中；提前失败：不满足条件的 payload 在验证阶段就报错，而不是运行到一半才失败；
代码更简洁：applyToOne 的第二个参数可以直接写成 func::CallOp，不需要 dyn_cast。

2. 用 TransformEachOpTrait 自动生成 apply 骨架

TransformEachOpTrait 是一个 trait，它告诉 Transform dialect，这个操作会对每个 payload 操作单独应用一次，你（框架）帮我生成 apply 方法的骨架。
框架自动完成的事情包括：验证：检查操作数 handle 是否有效；遍历：对 handle 关联的每个 payload 操作调用一次 applyToOne；结果合并：把每次 applyToOne 产生的结果按顺序合并到最终的结果 handle 中；
失败处理：如果某次 applyToOne 返回失败，按规则传播。

这是改进后的 ODS 操作定义：

	// In MyExtension.td.

	// Define the new operation. By convention, prefix its name with the name of the dialect extension, "my.". The full operation name will be further prefixed with "transform.".
	def ChangeCallTargetOp : Op<Transform_Dialect, "my.change_call_target",
	// Indicate that the operation implements the required TransformOpInterface and
	// MemoryEffectsOpInterface. Use the TransformEach trait to provide the
	// implementation for TransformOpInterface.
	[TransformOpInterface, TransformEachOpTrait,
	DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
	// Provide a brief and a full description. It is recommended that the latter describes
	// the effects on the operands and how the operation processes various failure modes.
	let summary = "Changes the callee of a call operation to the specified one"; // 描述信息
	let description = [{
	For each `func.call` payload operation associated with the handle, changes its
	callee to be the symbol whose name is provided as an attribute to this operation.

	Generates a silenceable failure if the operand is associated with payload operations
	that are not `func.call`.
	Only reads the operand.
	}];

	// The arguments include the handle to the payload operations and the attribute that
	// specifies the new callee. The handle must implement TransformHandleTypeInterface.
	// We use a string attribute as the symbol may not exist in the transform IR so the
	// verification may fail.
	let arguments = (ins
	Transform_ConcreteOpType<"func.call">:$call,
	StrAttr:$new_target);

	// The results are empty as the transformation does not produce any new payload.
	let results = (outs);

	// Provide nice syntax.
	let assemblyFormat = "$call `,` $new_target attr-dict `:` type($call)";

	// Declare the function implementing the interface for a single payload operation.
	let extraClassDeclaration = [{
	::mlir::DiagnosedSilenceableFailure applyToOne(
	::mlir::transform::TransformRewriter &rewriter,
	::mlir::func::CallOp call,
	::mlir::transform::ApplyToEachResultList &results,
	::mlir::transform::TransformState &state);
	}];
	}

当前要实现的接口是 applyToOne ：

	::mlir::DiagnosedSilenceableFailure ChangeCallTargetOp::applyToOne(
	::mlir::transform::TransformRewriter &rewriter,
	::mlir::func::CallOp call,
	::mlir::transform::ApplyToEachResultList &results,
	::mlir::transform::TransformState &state) {
	// Call the actual transformation function.
	updateCallee(call, getNewTarget());
	// Indicate success.
	return DiagnosedSilenceableFailure::success();
	}

只需要写单个 payload 操作的变换逻辑；不需要循环、不需要类型检查、不需要结果合并；返回 success () 表示这次应用成功。也就是，trait + 类型约束把验证、遍历、错误处理都自动化了。
如果没有 TransformEachOpTrait 和类型约束，等价的手写实现大概是，也就是前一章的内容：

	DiagnosedSilenceableFailure ChangeCallTargetOp::apply(
	TransformRewriter &rewriter,
	TransformResults &results,
	TransformState &state) {
	for (Operation *op : state.getPayloadOps(getCall())) {
	auto call = dyn_cast<func::CallOp>(op);
	if (!call) {
	auto diag = emitSilenceableError() << "expected func.call";
	diag.attachNote(op->getLoc()) << "offending operation";
	return diag;
	}
	updateCallee(call, getNewTarget());
	}
	return success();
	}

# Defining a Transform Type

上一节用 Transform_ConcreteOpType<"func.call"> 把 handle 约束为具体的 func.call 操作。但有时我们想约束的不是某个具体操作，而是实现了某个接口的一类操作。
当前 ChangeCallTargetOp 只能作用于 func.call；但我们希望它能作用于任何实现了 CallOpInterface 的 payload 操作（比如 func.call、call、自定义调用操作等）；Transform Dialect 目前没有提供一个现成的类型来检查 “payload 操作是否实现了 CallOpInterface”。
所以需要在扩展里自己定义一个 transform 类型。

Transform Dialect 允许扩展定义并注入额外的属性和类型。类型定义也使用 ODS：

	// Transform dialect allows additional types to be defined and injected.
	def CallOpInterfaceHandle
	: TypeDef<Transform_Dialect, "CallOpInterfaceHandle",
	// The type must implement `TransformHandleTypeInterface`.
	[DeclareTypeInterfaceMethods<TransformHandleTypeInterface>]> {

	// The usual components of a type such as description, mnemonic and assembly format
	// should be provided.
	let summary = "handle to payload operations implementing CallOpInterface";
	let mnemonic = "my.call_op_interface";
	let assemblyFormat = "";
	}

定义完类型后，必须实现 TransformHandleTypeInterface 的接口方法，最关键的是 checkPayload，该方法验证即将与这个 handle 关联的 payload 操作是否满足类型约束：

	// In MyExtension.cpp.

	// The interface declares this method to verify constraints this type has on
	// payload operations. It returns the now familiar tri-state result.
	mlir::DiagnosedSilenceableFailure
	mlir::transform::CallOpInterfaceHandleType::checkPayload(
	// Location at which diagnostics should be emitted.
	mlir::Location loc,
	// List of payload operations that are about to be associated with the
	// handle that has this type.
	llvm::ArrayRef<mlir::Operation *> payload) const {

	// All payload operations are expected to implement CallOpInterface, check this.
	for (Operation *op : payload) {
	if (llvm::isa<mlir::CallOpInterface>(op))
	continue;

	// By convention, these verifiers always emit a silenceable failure since they are
	// checking a precondition.
	DiagnosedSilenceableFailure diag = emitSilenceableError(loc)
	<< "expected the payload operation to implement CallOpInterface";
	diag.attachNote(op->getLoc()) << "offending operation";
	return diag;
	}

	// If everything is okay, return success.
	return DiagnosedSilenceableFailure::success();
	}

新类型需要在扩展的 init () 中注册，和操作注册放在一起：

	// In MyExtension.cpp.

	void MyExtension::init() {
	// ...

	registerTypes<
	#define GET_TYPEDEF_LIST
	#include "MyExtensionTypes.cpp.inc"
	>();
	}

注册后，这个类型就正式进入 Transform Dialect，可以在操作定义和 IR 中使用了。
之前 $call 的类型必须是 Transform_ConcreteOpType<"func.call">。现在可以用 AnyTypeOf 允许两种类型：

	def ChangeCallTargetOp : ... {
	let arguments = (ins
	// Allow the handle to be to concrete `func.call` ops as well as any op implementing
	// the `CallOpInterface`.
	AnyTypeOf<[Transform_ConcreteOpType<"func.call">, CallOpInterfaceHandle]>:$call,
	StrAttr:$new_target);
	}

这样 ChangeCallTargetOp 就从 “只作用于 func.call” 泛化为 “作用于任何调用类操作”。
把前面定义的 Transform 类型和 Transform 操作，实际写进一个 MLIR 文件里，用 Transform Dialect 的 interpreter 去运行，看看效果。

	// Cast to our new type.
	%casted = transform.cast %call : !transform.any_op to !transform.my.call_op_interface
	// Using our new operation.
	transform.my.change_call_target %casted, "microkernel" : !transform.my.call_op_interface

# Operand Consumption

当 transform 操作会 “消费” 操作数 handle 时，如何正确定义和实现这个操作呢？
前面 ChangeCallTargetOp 只是修改已有的 func.call（改被调用者），操作数 handle 之后还能继续用，所以它是只读操作数。如果我们现在将原有的 func.call 重写成自定义操作 my.mm4；原来的 func.call 会被 erase 掉，新的 my.mm4 被插入。
这种情况下，操作数 handle 指向的原操作已经不存在了，不能再被后续 transform 使用；所以这个 handle 被消费（consumed）了；同时变换产生了新的 payload 操作，需要返回一个新的 handle 指向它。
接下来定义一个新的 Transform 操作，CallToOp，相对于 ChangeCallTargetOp “修改已有 func.call 的被调用者符号，操作数只读，之后还能继续用”，CallToOp “把 func.call 替换成自定义操作 my.mm4，操作数被消费。”

	// In MyExtension.td.

	// Define another transform operation.
	def CallToOp : Op<Transform_Dialect, "my.call_to_op",
	// Indicate that the operation implements the required TransformOpInterface and
	// MemoryEffectsOpInterface. Use the TransformEach trait to provide the
	// implementation for TransformOpInterface.
	[TransformOpInterface, TransformEachOpTrait,
	DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
	// Summary and description omitted for brevity.

	// The argument is the handle to the payload operations.
	let arguments = (ins CallOpInterfaceHandle:$call); // 操作数 $call 类型是上一节定义的 CallOpInterfaceHandle，即任何实现 CallOpInterface 的操作。

	// The result is the handle to the payload operations produced during the
	// transformation.
	let results = (outs TransformHandleTypeInterface:$transformed); // 返回一个 handle $transformed，指向变换后新产生的 payload 操作。

	// Provide nice syntax.
	let assemblyFormat = "$call attr-dict `:` functional-type(operands, results)";

	// Declare the function implementing the interface for a single payload operation.
	let extraClassDeclaration = [{
	::mlir::DiagnosedSilenceableFailure applyToOne(
	::mlir::transform::TransformRewriter &rewriter,
	::mlir::CallOpInterface call,
	::mlir::transform::ApplyToEachResultList &results,
	::mlir::transform::TransformState &state);
	}];
	}

% transformed 是一个新的 handle，它指向 payload IR 里新产生的 my.mm4 操作。那么 CallToOp 这个 Transform 操作为什么需要一个指向 my.mm4 操作的句柄呢？
假设 CallToOp 不返回任何结果，只是把 func.call 替换成 my.mm4，然后就结束了。那么：payload IR 里确实多了一个 my.mm4；但 transform 脚本里没有任何东西 “知道” 这个新操作的存在；后续的 transform 操作无法再对新产生的 my.mm4 做任何事。
也就是说，这个新操作就 “失联” 了，你只能靠它自己存在于 payload IR 里，但 transform 流程无法继续操作它。
因此，CallToOp 必须返回一个新 handle % transformed，指向新产生的 my.mm4。

和前面小节一样，上述操作需要实现 applyToOne () 函数。
DeclareOpInterfaceMethods：声明内存效果接口方法，需要手动实现 getEffects。

applyToOne 函数

	// In MyExtension.cpp.

	::mlir::DiagnosedSilenceableFailure CallToOp::applyToOne(
	::mlir::transform::TransformRewriter &rewriter,
	::mlir::CallOpInterface call,
	::mlir::transform::ApplyToEachResultList &results,
	::mlir::transform::TransformState &state) {
	// Call the actual rewrite.
	Operation *rewritten = rewriteToOp(call);

	// Report an error if the rewriter produced a null pointer. Note that it may have
	// irreversibly modified the payload IR, so we produce a definite failure.
	if (!rewritten) {
	return emitDefiniteError() << "failed to rewrite call to operation";
	}

	// On success, push the resulting operation into the result list. The list is expected
	// to contain exactly one entity per result and per application. The handles will be
	// associated with lists of the respective values produced by each application.
	results.push_back(rewritten);

	// If everything is fine, return success.
	return DiagnosedSilenceableFailure::success();
	}

调用实际的改写函数 rewriteToOp (call)，把 func.call 替换成 my.mm4，返回新操作指针。

失败处理：如果 rewritten 为空，说明改写失败；此时 payload IR 可能已经被不可逆地修改了，所以用 emitDefiniteError () 发出确定性失败（definite failure），而不是 silenceable failure。
结果关联：把新产生的操作 rewritten push 到 results 列表；applyToOne 要求每个结果 handle 在每次应用中恰好对应一个 payload 实体；trait 会把所有应用的结果合并，关联到最终的结果 handle $transformed。

getEffects 方法

	void CallToOp::getEffects(
	::llvm::SmallVectorImpl<::mlir::MemoryEffects::EffectInstance> &effects) {
	// Indicate using side effects that the operand handle is consumed, and the
	// result handle is produced.
	consumesHandle(getCall(), effects);
	producesHandle(getTransformed(), effects);

	// Indicate that the payload IR is modified.
	modifiesPayload(effects);
	}

consumesHandle (getCall (), effects) : 操作数 handle $call 被消费，之后不能再使用
producesHandle (getTransformed (), effects) : 结果 handle $transformed 被产生
modifiesPayload (effects) : payload IR 被修改
这些内存效果信息对 interpreter 很重要：它知道哪些 handle 在操作后失效，哪些是新产生的；可以做正确性检查，比如防止使用已消费的 handle。

最后，教程提到结果关联规则。
1. 成功时必须为所有结果指定实体
成功返回时，每个结果 handle 都必须有对应的实体列表；即使列表为空，也要显式设置；否则会触发断言（assertion）。
2. 失败时解释器自动填充空列表
如果操作失败，尚未定义的结果 handle 会被自动关联到空列表；这样后续使用这些 handle 时不会出错，只是关联不到任何 payload。

那么，当一个 transform 操作有结果 handle，但这次变换不想让这个结果指向任何 payload 操作时，该怎么办？
applyToOne 的语义是：对每个 payload 操作应用一次，每次应用为每个结果 handle 产生恰好一个 payload 实体。那想让结果 handle 关联空列表呢？
假设有一个 transform 操作，它接收一个非空的 handle，但变换后不产生任何新的 payload 操作，结果 handle 应该关联空列表。这时候 applyToOne 就不能使用了，需要直接实现 apply 方法。
教程给了个这样的例子：

	::mlir::DiagnosedSilenceableFailure SomeOtherOp::apply(
	::mlir::transform::TransformRewriter &rewriter,
	::mlir::transform::TransformResults &results,
	::mlir::transform::TransformState &state) {
	// ...

	// Associate the result `transformed` with an empty list of payload operations.
	results.set(cast<OpResult>(getTransformed()), {});
	return DiagnosedSilenceableFailure::success();
	}

# Memory Effects Traits

用 trait 来替代手写 MemoryEffectsOpInterface 的方法，从而减少样板代码。
前述 ODS 中是这样声明的内存效果的：

	[TransformOpInterface, TransformEachOpTrait,
	DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]

然后在 C++ 里手写：

	void CallToOp::getEffects(
	llvm::SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
	consumesHandle(getCall(), effects);
	producesHandle(getTransformed(), effects);
	modifiesPayload(effects);
	}

这三行其实是非常模式化的，手写太重复。
MLIR 提供了两个内存效果 trait。
FunctionalStyleTransformOpTrait：所有 handle 类型操作数都被消费；所有结果都被产生；payload IR 被修改。
NavigationTransformOpTrait ：所有 handle 类型操作数只读；所有结果都被产生；payload IR 只读。
用 trait 的简单声明方式（这种情况下就不需要再声明或定义 MemoryEffectsOpInterface 的方法了）：

	[TransformOpInterface, TransformEachOpTrait,
	MemoryEffectsOpInterface, FunctionalStyleTransformOpTrait]

# 后记

	i wish that i could say i had a good childhood
	i guess i have got some memories of running through the woods
	and i believe i even built myself a tree house
	but i cant find a photo where i am smiling with my teeth out
	i wish i had listened more to what my grandpa said
	bu i guess i got the gist of it, cause now and then i hear him in my head
	and what he says is i should aim high
	cause 99 percent of people never really live before they die
	wish i could do it all again
	but meet you when i am ten
	be so much better with my best friend

	—— 《A good childhood》

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