# 前言

本篇是学习 MLIR 中的 Transform Dialect 的第三篇，添加新的 Transformation 操作。这一章节主要讲了如何为 MLIR 的 Transform 方言添加一个简单的新变换操作。内容来自 MLIR 官方教程《Chapter 2: Adding a Simple New Transformation Operation》
相关链接： LLVM Project ，MLIR 官方文档，MLIR 官网教程。
作为初学者，错误在所难免，还望不吝赐教。

# 基本简介

在前一章节《Chapter 1: Combining Existing Transformations》中，最后一小节 More Handle Invalidation 提到将融合后的操作继续 tile 到 4×4 大小，得到一个 func.call，调用被 outline 出来的函数。

	%func, %call = transform.loop.outline %outline_target
	{func_name = "outlined"}

这一节没有将 func.call 替换成对微内核（microkernel）或硬件 intrinsic 的调用。Transform 方言本身没有现成的操作能做这件事，所以需要自己定义新的 transform op。
本章节开始做这样一件事，即定义 ChangeCallTargetOp，真正把 func.call 的被调用者改成 "microkernel"。

# Setting Up to Add New Transformations

在真正定义一个新的 Transform 方言操作之前，需要先把 “架子” 搭好。更具体地：

1. 把实现放在哪里 —— 用方言扩展机制
2. 扩展类怎么定义、init () 里要做什么
3. 操作用 ODS 怎么写
4. 用 CMake + TableGen 生成代码，并把扩展注册进去

MLIR 是鼓励把变换操作贡献到上游（upstream）的，也就是把需要拓展的 transform 操作直接添加到 transform 方言中去，丰富该方言的功能。但现在教程中我们要添加的操作是个非常私人的需求，即树外（out-of-tree）方言，所以 MLIR 提供了方言扩展机制：在不修改方言本身的前提下，向方言注入额外的操作。

方言拓展是这样定义的：

	// In MyExtension.cpp.
	#include "mlir/Dialect/Transform/IR/TransformDialect.h"

	// Define a new Transform dialect extension. This uses the CRTP idiom to
	// identify extensions.
	class MyExtension : public ::mlir::transform::TransformDialectExtension<MyExtension> { // 定义一个方言拓展类
	public:
	// The extension must derive the base constructor.
	using Base::Base; // 继承基类的构造函数

	// This function initializes the extension, similarly to `initialize` in
	// dialect definitions. List individual operations and dependent dialects
	// here.
	void init();
	};

	void MyExtension::init() {
	// Similarly to dialects, an extension can declare a dependent dialect. This
	// dialect will be loaded along with the extension and, therefore, along with
	// the Transform dialect. Only declare as dependent the dialects that contain
	// the attributes or types used by transform operations. Do NOT declare as
	// dependent the dialects produced during the transformation.
	// 声明 dependent dialects（依赖方言）
	// declareDependentDialect<MyDialect>();

	// When transformations are applied, they may produce new operations from
	// previously unloaded dialects. Typically, a pass would need to declare
	// itself dependent on the dialects containing such new operations. To avoid
	// confusion with the dialects the extension itself depends on, the Transform
	// dialects differentiates between:
	// - dependent dialects, which are used by the transform operations, and
	// - generated dialects, which contain the entities (attributes, operations,
	// types) that may be produced by applying the transformation even when
	// not present in the original payload IR.
	// In the following chapter, we will be add operations that generate function
	// calls and structured control flow operations, so let's declare the
	// corresponding dialects as generated.
	// 声明 generated dialects（生成方言）
	declareGeneratedDialect<::mlir::scf::SCFDialect>();
	declareGeneratedDialect<::mlir::func::FuncDialect>();

	// Finally, we register the additional transform operations with the dialect.
	registerTransformOps<
	// TODO: list the operation classes. // 注册变换操作，先留空，等 ODS 定义好操作后再填。
	>();
	}

定义的方言拓展类非常简单，继承基类，使用基类的构造函数，以及一个 init () 函数。
init () 函数主要做三件事：声明 dependent dialects（依赖方言）、声明 generated dialects（生成方言）、注册变换操作。
值得注意的是，声明 dependent dialects（依赖方言），指的是如果该拓展还需要依赖其他方言的属性或类型，那么就要增加方言依赖。尽管当前的拓展肯定依赖于 transform 方言，但 TransformDialectExtension 已经指明了依赖，Transform 方言被加载时，这些扩展会被一起加载，所以这里不用额外声明。

操作定义和普通方言操作完全一样，用 TableGen / ODS：

	// In MyExtension.td
	#ifndef MY_EXTENSION
	#define MY_EXTENSION

	include "mlir/Dialect/Transform/IR/TransformDialect.td"
	include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.td"
	include "mlir/IR/OpBase.td"
	include "mlir/Interfaces/SideEffectInterfaces.td"

	def MyOp : Op<Transform_Dialect, "transform.my.op", [
	// TODO: interfaces and traits here.
	]> {
	let summary = "my transform op";
	// TODO: define the operation properties.
	}

	#endif // MY_EXTENSION

用 CMake + TableGen 生成代码：生成 .h.inc/.cpp.inc

	# In CMakeLists.txt next to MyExtension.td.

	# Tell Tablegen to use MyExtension.td as input.
	set(LLVM_TARGET_DEFINITIONS MyExtension.td)

	# Ask Tablegen to generate op declarations and definitions from ODS.
	mlir_tablegen(MyExtension.h.inc -gen-op-decls)
	mlir_tablegen(MyExtension.cpp.inc -gen-op-defs)

	# Add a CMakeTarget we can depend on to ensure the generation happens before the compilation.
	add_public_tablegen_target(MyExtensionIncGen)

	# Don't forget to generate the documentation, this will produce a MyExtension.md under
	# Dialects.
	add_mlir_doc(MyExtension MyExtension Dialects/ -gen-op-doc)

建库并连接：

	# In CMakeLists.txt next to MyExtension.cpp
	add_mlir_library(
	# Library called MyExtension.
	MyExtension

	# Built from the following source files.
	MyExtension.cpp

	# Make sure ODS declaration and definitions are generated before compiling
	# this.
	DEPENDS
	MyExtensionIncGen

	# Link in the transform dialect, and all generated dialects.
	LINK_LIBS PUBLIC
	MLIRTransformDialect
	MLIRFuncDialect
	MLIRSCFDialect
	)

库名 MyExtension；源文件 MyExtension.cpp；DEPENDS MyExtensionIncGen 确保生成；链接 Transform 方言和所有 generated dialects
在 .h/.cpp 里 include 生成文件。

	// MyExtension.h
	#include "mlir/Dialect/Transform/IR/TransformDialect.h"
	#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
	#define GET_OP_CLASSES
	#include "MyExtension.h.inc"


	// MyExtension.cpp
	#include "MyExtension.h"
	#define GET_OP_CLASSES
	#include "MyExtension.cpp.inc"

最后在 init () 函数中补全注册：

	void MyExtension::init() {
	// …

	registerTransformOps<
	#define GET_OP_LIST
	#include "MyExtension.cpp.inc"
	>();
	}

GET_OP_LIST 会展开成所有由 ODS 生成的操作列表。registerTransformOps 还会做额外检查：
操作是否实现了 Transform 方言解释器要求的接口（如 TransformOpInterface）；
是否实现了内存效果接口（MemoryEffectsOpInterface）；
不满足会直接 assert。

# Defining a Transform Operation

前述架子搭好之后，定义并实现一个具体的 transform 操作。这个操作叫 ChangeCallTargetOp，功能是把 func.call 的被调用者改成指定的符号名。
内容分成三部分：

用 ODS 定义操作（MyExtension.td）
实现 TransformOpInterface::apply（MyExtension.cpp）
实现 MemoryEffectsOpInterface::getEffects（MyExtension.cpp）

用 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.
	[DeclareOpInterfaceMethods<TransformOpInterface>, // Transform 方言要求实现这个接口
	DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> { // 用来声明操作对操作数和 payload IR 的副作用（读还是写、消费还是只读）
	// 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 TransformHandleTypeInterface:$call, // 类型约束是 TransformHandleTypeInterface，指向要修改的 payload 操作
	StrAttr:$new_target); // "microkernel" 这个符号不一定在当前的 transform IR 里存在，用 symbol 引用会导致验证失败，所以用 StrAttr

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

	// Provide nice syntax. 自定义语法，生成类似语句 “transform.my.change_call_target % call, "microkernel" : !transform.any_op”
	let assemblyFormat = "$call `,` $new_target attr-dict `:` type($call)";
	}

DeclareOpInterfaceMethods 表明这里只是声明了接口，具体实现在.cpp 文件中。

为了完成转换操作的定义，我们需要实现接口方法。目前 TransformOpInterface 只需要一个方法 ——apply，该方法执行实际的转换操作。将方法体限制在对 Transform 语法结构的处理上，并将实际的转换逻辑实现为独立函数，以便在代码其他地方调用，是一种良好的做法。与重写模式类似，所有 IR 都必须通过提供的重写器进行修改。

实现 TransformOpInterface::apply， apply 是 TransformOpInterface 目前唯一要求的方法，负责执行真正的变换。

	// In MyExtension.cpp

	// Implementation of our Transform dialect operation.
	// This operation returns a tri-state result that can be one of:
	// - success when the transformation succeeded;
	// - definite failure when the transformation failed in such a way that
	// following transformations are impossible or undesirable, typically it could
	// have left payload IR in an invalid state; it is expected that a diagnostic
	// is emitted immediately before returning the definite error;
	// - silenceable failure when the transformation failed but following
	// transformations are still applicable, typically this means a precondition
	// for the transformation is not satisfied and the payload IR has not been
	// modified. The silenceable failure additionally carries a Diagnostic that
	// can be emitted to the user.
	::mlir::DiagnosedSilenceableFailure mlir::transform::ChangeCallTargetOp::apply(
	// The rewriter that should be used when modifying IR.
	::mlir::transform::TransformRewriter &rewriter, // 类似 rewrite pattern 里的 rewriter
	// The list of payload IR entities that will be associated with the
	// transform IR values defined by this transform operation. In this case, it
	// can remain empty as there are no results.
	::mlir::transform::TransformResults &results,
	// The transform application state. This object can be used to query the
	// current associations between transform IR values and payload IR entities.
	// It can also carry additional user-defined state.
	::mlir::transform::TransformState &state) {

	// First, we need to obtain the list of payload operations that are associated
	// with the operand handle.
	auto payload = state.getPayloadOps(getCall());

	// Then, we iterate over the list of operands and call the actual IR-mutating
	// function. We also check the preconditions here.
	for (Operation *payloadOp : payload) {
	auto call = dyn_cast<::mlir::func::CallOp>(payloadOp);
	if (!call) {
	DiagnosedSilenceableFailure diag = emitSilenceableError()
	<< "only applies to func.call payloads";
	diag.attachNote(payloadOp->getLoc()) << "offending payload";
	return diag;
	}

	updateCallee(call, getNewTarget());
	}

	// If everything went well, return success.
	return DiagnosedSilenceableFailure::success(); // 返回成功
	}

函数签名有三个：

rewriter：修改 IR 必须用它，类似 rewrite pattern 里的 rewriter；
results：用来填充这个 transform 操作的结果，这里没有结果，所以不用管；
state：变换应用状态，可以查询 transform IR 值和 payload IR 实体之间的关联。

实现逻辑：用 state.getPayloadOps (getCall ()) 拿到 $call handle 关联的所有 payload 操作；遍历它们；用 dyn_castfunc::CallOp 检查是不是 func.call；不是，就发 silenceable error，附上出错位置，返回；是，就调用 updateCallee (call, getNewTarget ()) 做实际修改；全部成功，返回 success。
文档里说：“It is a good practice to limit the body of the method to manipulation of the Transform dialect constructs and have the actual transformation implemented as a standalone function so it can be used from other places in the code.”
也就是说，apply 里尽量只做 Transform 方言层面的调度和检查；真正的 IR 修改（这里是 updateCallee）抽成独立函数，方便复用。

实现 MemoryEffectsOpInterface::getEffects

	// In MyExtension.cpp

	void ChangeCallTargetOp::getEffects(
	::llvm::SmallVectorImpl<::mlir::MemoryEffects::EffectInstance> &effects) {
	// Indicate that the `call` handle is only read by this operation because the
	// associated operation is not erased but rather modified in-place, so the
	// reference to it remains valid.
	onlyReadsHandle(getCall(), effects);

	// Indicate that the payload is modified by this operation.
	modifiesPayload(effects);
	}

这个接口声明操作对操作数和 payload 的副作用。必须实现，否则 Transform 方言的 verifier 会在 debug 构建里 assert。

# Registration and Usage

注册扩展：提供一个注册钩子，让项目在启动时把 MyExtension 加进去；

使用新操作：在 Transform 序列里实际调用 transform.my.change_call_target 。

	// In TransformDialect.cpp (don't forget a declaration in TransformDialect.h);

	void registerMyExtension(::mlir::DialectRegistry &registry) {
	registry.addExtensions<MyExtension>();
	}

这是项目的入口钩子，通常在 main 或初始化代码里调用；DialectRegistry 是方言和扩展的注册表；addExtensions() 把前面定义的扩展注册进去；
注册后，Transform 方言加载时，MyExtension 也会被加载，init () 被调用，里面的 declareGeneratedDialect 和 registerTransformOps 生效。

mlir/test/Examples/transform/Ch2/sequence.mlir 这里包含了测试例子，以及微内核 microkernel 的实现。
如以此来，第一章提到的 Transform 转换变成下面这个样子。

	module attributes {transform.with_named_sequence} {
	transform.named_sequence @__transform_main(
	%arg0: !transform.any_op,
	%arg1: !transform.op<"linalg.matmul">,
	%arg2: !transform.op<"linalg.elementwise">) {
	// Since the %arg2 handle is associated with both elementwise operations,
	// we need to split it into two handles so we can target only the second
	// elementwise operation.
	%add, %max = transform.split_handle %arg2
	: (!transform.op<"linalg.elementwise">)
	-> (!transform.any_op, !transform.any_op)

	// The actual tiling transformation takes tile sizes as attributes. It
	// produces a handle to the loop generated during tiling.
	%tiled, %loop = transform.structured.tile_using_forall %max
	tile_sizes [8, 32]
	: (!transform.any_op) -> (!transform.any_op, !transform.any_op)

	// We can now fuse the other operations into the loop. Here, we fuse
	// operations one-by-one. This requires the operation that is being fused
	// to define the value used within the loop, so the order of such fusions
	// is important. We could also use "transform.merge_handles" to obtain
	// a single handle to all operations and give it to
	// `fuse_into_containing_op` that would take care of the ordering in this
	// case.
	%add_fused, %loop2 = transform.structured.fuse_into_containing_op %add into %loop
	: (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
	%matmul_fused, %loop3 = transform.structured.fuse_into_containing_op %arg1
	into %loop2
	: (!transform.op<"linalg.matmul">, !transform.any_op)
	-> (!transform.any_op, !transform.any_op)

	// Tile again to get the desired size. Note that this time this tiles the
	// "add" operation and fuses matmul into the loop, but doesn't affect the
	// "max" operation. This illustrates the precise targeting with the
	// transform dialect. Otherwise, it is difficult to differentiate "add" and
	// "max", both of which having the same kind.
	%tiled_second, %loop_second = transform.structured.tile_using_forall %add_fused
	tile_sizes [4, 4]
	: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
	%matmul_fused_2, %loop_second_2 = transform.structured.fuse_into_containing_op %matmul_fused
	into %loop_second
	: (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)

	// Since outlining is currently only implemented for region-holding
	// operations such as loops, use tiling to size 1 to materialize the outer
	// loop that is going to be outlined.
	%_0, %loop_third = transform.structured.tile_using_forall %tiled_second tile_sizes [1]
	: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
	%_1, %outline_target = transform.structured.fuse_into_containing_op %matmul_fused_2 into %loop_third
	: (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
	%func, %call = transform.loop.outline %outline_target
	{func_name = "outlined"}
	: (!transform.any_op) -> (!transform.any_op, !transform.any_op)

	// Rewrite the call target.
	transform.my.change_call_target %call, "microkernel" : !transform.any_op

	transform.yield
	}
	}

# 后记

	看着这四十万亿千米外的银色墓场，维纳尔感慨万千。
	“其实吧，从科学角度讲，毁灭一词并不准确，没有真正毁掉什么，更没有灭掉什么。
	物质总量一点不少都还在，角动量也还在，只是物质的组合方式变了变。
	像一副扑克牌，仅仅重洗而已……可生命是一手同花顺，一洗什么都没了。”
	—— 《三体》

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