Homecompile

# 前言

本篇是学习 MLIR 中的 Transform Dialect 的第二篇，主要通过一个具体的 “全连接 + 偏置 + ReLU” 层例子，介绍了如何使用 Transform Dialect 来编写变换脚本。内容来自 MLIR 官方教程 《Chapter 1: Combining Existing Transformations》。

相关链接： LLVM Project ，MLIR 官方文档，MLIR 官网教程， 《Chapter 0: A Primer on “Structured” Linalg Operations》，【MLIR】Transform Dialect【0】保持计算结构。

作为初学者，错误在所难免，还望不吝赐教。

# 简介

Transform 方言允许用户精确地针对 IR 中的特定操作进行转换，并将这些转换串联起来，即对前一个转换生成的操作应用转换。为此，转换以 IR 中其他操作的形式表达。我们称包含这些操作的 IR 为转换（transform）IR，而正在被转换的 IR 称为负载（payload）IR。
转换 IR 的应用始终从一个顶层操作开始。在 C++ API 中，该操作会被传递给 applyTransforms 函数。这个顶层操作用于指定是否应执行其他转换以及如何执行。最常见的顶层操作 transform.named_sequence，会依次应用其体内的列出的其他转换操作，类似于函数或宏。

基本概念：

  • Transform IR：用 IR 操作来表示变换本身。
  • Payload IR：被变换的 IR。
  • 变换操作作用于三类值：操作句柄、值句柄、参数。
  • 转换 IR 的应用始终从一个顶层操作开始，最常见的是 transform.named_sequence，类似函数或宏。

下面以一个简单的变换序列来说明，该序列适用于常见的 “全连接 + 偏置 + ReLU” 机器学习层。其过程可归结为先进行矩阵乘法，然后执行（逐元素）矩阵加法，并对结果取与 0 的逐元素最大值。这可以用以下中间表示（IR）来表达：

func.func @fc_relu(%lhs: tensor<512x512xf32>, %rhs: tensor<512x512xf32>,
                   %bias: tensor<512x512xf32>, %output: tensor<512x512xf32>)
                   -> tensor<512x512xf32> {
  // Matrix-matrix multiplication.
  %matmul = linalg.matmul ins(%lhs, %rhs: tensor<512x512xf32>, tensor<512x512xf32>)
                          outs(%output: tensor<512x512xf32>) -> tensor<512x512xf32>
  // Elementwise addition.
  %biased = linalg.elementwise <add>
    ins(%matmul, %bias : tensor<512x512xf32>, tensor<512x512xf32>)
    outs(%output : tensor<512x512xf32>) -> tensor<512x512xf32>
  // Elementwise max with 0 (ReLU).
  %c0f = arith.constant 0.0 : f32
  %relued = linalg.elementwise <max_signed>
    indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> ()>, affine_map<(d0, d1) -> (d0, d1)>]
    ins(%biased, %c0f : tensor<512x512xf32>, f32)
    outs(%output : tensor<512x512xf32>) -> tensor<512x512xf32>
  func.return %relued : tensor<512x512xf32>
}
# 顶层序列操作 (Top-Level Sequence Operation)


因为性能原因，我们希望对 “全连接 + 偏置 + ReLU” 这一串操作进行分块（tile）和融合（fuse），以利用缓存局部性。这些变换需要按顺序依次执行，所以自然从一个对应的顶层变换操作开始 —— 也就是 transform.named_sequence。


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">):
    transform.yield
  }
}
1）特殊名称 @__transform_main 和第一个参数

 @__transform_main 这个名字以及第一个参数 % arg0 是解释器 pass 强制要求的。

类比 C 程序入口 main 必须有 int (int argc, char** argv) 签名。

% arg0 会关联到顶层 payload 操作，通常就是 pass 所作用的那个操作。

注意：如果通过 applyTransforms 或 applyNamedSequence 以编程方式应用变换，则这些都不是必需的。

（2）其余入口块参数是可选的

剩下的参数（这里是 % arg1、% arg2）是可选的，可以关联到 payload 的属性、操作或值，方便在序列中使用。

这些关联也是在调用 applyTransforms 时指定的。

在本例中，我们关心的是将要被分块和融合的矩阵乘法和逐元素操作。

（3）所有值句柄都有 Transform dialect 类型，这些类型规定了它们所关联的 payload IR 实体的性质：

!transform.any_op：句柄可关联到任意 payload 操作。

!transform.op<"X">：句柄只能关联到类型为 X 的 payload 操作。

这些约束在创建句柄 /payload 关联时会被验证。对于顶层变换操作的入口块参数，这个验证发生在 applyTransforms 函数的早期阶段。如果约束不满足，变换应用会失败，并向用户产生诊断信息。


# 失败传播 (Failure Propagation)


Transform 方言基础设施具有一种专门处理诊断信息的机制，能够支持可恢复的错误。通过考虑一个 “未命名序列” 操作来理解这一点，该操作包含一个强制性属性，用于指定失败传播模式。有两种选择：

“propagate”：如果嵌套中的任一转换失败，则整个序列转换将失败；

“suppress”：即使嵌套中的某一项转换失败，序列仍会成功执行，但不会尝试在失败后继续执行后续的转换。


这里说的是 “未命名序列”（unnamed sequence），即 transform.sequence，它和之前讲的 transform.named_sequence 是两类东西：

transform.named_sequence：有名字，可以作为顶层入口。

transform.sequence：匿名，通常嵌套在别的序列里，用来控制失败传播。


transform.sequence failures(propagate) {
^bb0(%arg0: !transform.any_op,
     %arg1: !transform.op<"linalg.matmul">,
     %arg2: !transform.op<"linalg.elementwise">):
  transform.debug.emit_remark_at %arg1, "matmul"
      : !transform.op<"linalg.matmul">
  transform.debug.emit_remark_at %arg2, "elemwise_binaries"
      : !transform.op<"linalg.elementwise">
  transform.yield
}
    

  • transform.sequence failures(propagate) ：定义一个匿名序列，失败传播模式为  propagate 。
    • 

  • 入口块参数  %arg0 、 %arg1 、 %arg2  是句柄。
    • 

  • transform.debug.emit_remark_at  % arg1, "matmul"：在 % arg1 关联的 payload 操作处发出一个 remark（备注），内容为 "matmul"。
    • 

  • transform.debug.emit_remark_at  % arg2, "elemwise_binaries"：同理，在 % arg2 关联的操作处发出 remark。
    • 

  • transform.yield ：序列结束。
    • 



# Transform 方言解释器 (Transform Dialect Interpreter)


如何在不重新编译编译器的情况下，运行变换脚本。

Transform Dialect 的一个核心优势是：变换脚本是 IR，不是硬编码在 C++ 里的。


这意味着：改变换脚本，不需要重新编译 MLIR。可以把变换脚本当成 “配置” 或 “程序”，交给一个解释器去执行。这个解释器就是 transform-interpreter pass。


$ mlir-opt sequence.mlir --pass-pipeline="
    builtin.module(transform-interpreter{
        debug-bind-trailing-args=linalg.matmul,linalg.elementwise})"
    

  • mlir-opt：MLIR 的通用优化 / 测试工具。
    • 

  • sequence.mlir：输入文件，同时包含 payload IR 和 transform IR，嵌套在同一个 module 里。
    • 

  • --pass-pipeline=...：指定要运行的 pass 流水线。
    • 

  • builtin.module (...)：在 module 级别运行。
    • 

  • transform-interpreter {...}：运行 transform 解释器 pass。
    • 

  • debug-bind-trailing-args=linalg.matmul,linalg.elementwise：把顶层序列的额外参数绑定到所有匹配的 payload 操作上。
    • 



解释器会做哪些东西：

找到 @__transform_main 这个 named sequence。把它应用到 pass 的 anchor operation（锚点操作，通常是 module）。根据 debug-bind-trailing-args 的选项，把额外参数关联到对应的 payload 操作：

第一个额外参数 → 所有 linalg.matmul 操作

第二个额外参数 → 所有 linalg.elementwise 操作


注意：@__transform_main 的第一个参数 % arg0 自动关联到顶层 payload 操作（这里是 module），不需要在 debug-bind-trailing-args 里指定。debug-bind-trailing-args 只负责后面的额外参数。


运行后，因为脚本里用了 transform.debug.emit_remark_at，所以会输出 remark：


sequence.mlir:5:13: remark: matmul
  %matmul = linalg.matmul ins(%lhs, %rhs: tensor<512x512xf32>, tensor<512x512xf32>)
            ^
sequence.mlir:5:13: note: see current operation: %0 = linalg.matmul ins(%arg0, %arg1 : tensor<512x512xf32>, tensor<512x512xf32>) outs(%arg3 : tensor<512x512xf32>) -> tensor<512x512xf32>
sequence.mlir:9:13: remark: elemwise_binaries
  %biased = linalg.elementwise <add>
            ^
sequence.mlir:9:13: note: see current operation: %1 = linalg.elementwise <add> ins(%0, %arg2 : tensor<512x512xf32>, tensor<512x512xf32>) outs(%arg3 : tensor<512x512xf32>) -> tensor<512x512xf32>
sequence.mlir:15:13: remark: elemwise_binaries
  %relued = linalg.elementwise <max_signed>
            ^
sequence.mlir:15:13: note: see current operation: %2 = linalg.elementwise <max_signed> indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> ()>, affine_map<(d0, d1) -> (d0, d1)>] ins(%1, %cst : tensor<512x512xf32>, f32) outs(%arg3 : tensor<512x512xf32>) -> tensor<512x512xf32>
从以上打印输出中可以看到：

% arg1 成功关联到了 linalg.matmul 操作，打印出 "matmul"。


% arg2 成功关联到了两个 linalg.elementwise 操作（add 和 max_signed），所以打印出两次 "elemwise_binaries"。


# 指定变换 (Specifying Transformations)


真正开始动手做变换。

对 matmul 做分块：


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">) {
    // The actual tiling transformation takes tile sizes as attributes.
    %loop, %tiled = transform.structured.tile_using_forall %arg1
                    tile_sizes [4, 32]
      : (!transform.op<"linalg.matmul">)
     -> (!transform.any_op, !transform.any_op)
    transform.yield
  }
}
其中的核心就一个：


%loop, %tiled = transform.structured.tile_using_forall %arg1 tile_sizes [4, 32]
用 % arg1 这个句柄，找到它指向的 linalg.matmul 操作。

对它进行 tiling（分块），分块大小为 [4, 32]。

它返回两个新的句柄：% tiled：指向分块后产生的 linalg.generic（作用在原始数据的子集上）。% loop：指向包围它的 scf.forall “multi-for” 循环。

这样可以得到分块之后的代码：payload IR 里的 matmul 被包在了一个 scf.forall 循环里。


#map = affine_map<(d0) -> (d0 * 4)>
#map1 = affine_map<(d0) -> (d0 * 32)>
#map2 = affine_map<(d0, d1) -> (d0, d1)>
#map3 = affine_map<(d0, d1) -> ()>
func.func @fc_relu(%arg0: tensor<512x512xf32>,
                   %arg1: tensor<512x512xf32>,
                   %arg2: tensor<512x512xf32>,
                   %arg3: tensor<512x512xf32>) -> tensor<512x512xf32> {
  %0 = scf.forall (%arg4, %arg5) in (128, 16) shared_outs(%arg6 = %arg3) -> (tensor<512x512xf32>) {
    %3 = affine.apply #map(%arg4)
    %4 = affine.apply #map1(%arg5)
    %extracted_slice = tensor.extract_slice %arg0[%3, 0] [4, 512] [1, 1]
                     : tensor<512x512xf32> to tensor<4x512xf32>
    %extracted_slice_0 = tensor.extract_slice %arg1[0, %4] [512, 32] [1, 1]
                     : tensor<512x512xf32> to tensor<512x32xf32>
    %extracted_slice_1 = tensor.extract_slice %arg6[%3, %4] [4, 32] [1, 1]
                     : tensor<512x512xf32> to tensor<4x32xf32>
    %5 = linalg.matmul
         ins(%extracted_slice, %extracted_slice_0
            : tensor<4x512xf32>, tensor<512x32xf32>)
         outs(%extracted_slice_1 : tensor<4x32xf32>) -> tensor<4x32xf32>
    scf.forall.in_parallel {
      tensor.parallel_insert_slice %5 into %arg6[%3, %4] [4, 32] [1, 1]
           : tensor<4x32xf32> into tensor<512x512xf32>
    }
  }
  %1 = linalg.elementwise <add>
     ins(%0, %arg2 : tensor<512x512xf32>, tensor<512x512xf32>)
     outs(%arg3 : tensor<512x512xf32>) -> tensor<512x512xf32>
  %cst = arith.constant 0.000000e+00 : f32
  %2 = linalg.elementwise <max_signed>
    indexing_maps = [#map2, #map3, #map2]
    ins(%1, %cst : tensor<512x512xf32>, f32)
    outs(%arg3 : tensor<512x512xf32>) -> tensor<512x512xf32>
  return %2 : tensor<512x512xf32>
}
还有两个关键概念：消耗（consume）与失效（invalidate）

教程中说 “Besides producing new handles, the tiling transform operation consumes the operand handle. This means that the handle is invalidated after this operation, and is no longer supposed to be used.”

 意思是：transform.structured.tile_using_forall 不仅产生新句柄 % loop 和 % tiled，还消耗了输入句柄 % arg1。被消耗后，% arg1 失效，不能再使用。再用会导致未定义行为（解释器会用昂贵检查报错）。

为什么会被消耗？因为：分块会擦除原来的 matmul 操作，并创建新的小 matmul（结构相似但不同）。句柄 % arg1 原本指向被擦除的那个操作，操作没了，句柄就成了悬空引用（dangling reference）。

教程里写：“Transform operations are required to mark all their operands as either consumed or readonly.”

 每个变换操作必须明确声明它的每个操作数是：consumed（消耗）：操作数关联的 payload 操作会被擦除或重建，句柄失效。readonly（只读）：操作数关联的 payload 操作不会被改动，句柄仍然有效。

例如：


    

  • 

    transform.structured.tile_using_forall：会擦除并重建操作 → consumed。
    

    • 

  • 

    transform.cast：只是创建一个别名句柄，不擦除操作 → readonly。
    

    • 

  • 

    transform.debug.emit_remark_at：只是打印信息 → readonly。
    

    • 



# 句柄失效与昂贵检查模式 (Handle Invalidation and Expensive Checks Mode)


我们很难处理未定义行为，因此 Transform 语言方言解释器默认会执行一系列额外的、可能代价高昂的检查，以检测转换 IR 中的大多数未定义行为。例如，如果我们想在句柄 % arg1 被消费后使用它，这会导致未定义行为，该行为在调试模式下表现为断言失败，在发布模式下则可能导致段错误。

上一节说， tile_using_forall  会消耗  %arg1 ，导致它失效。如果你在失效后还用它，会发生什么？


transform.named_sequence @__transform_main(
     %arg0: !transform.any_op,
     %arg1: !transform.op<"linalg.matmul">,
     %arg2: !transform.op<"linalg.elementwise">) {
  %loop, %tiled = transform.structured.tile_using_forall %arg1 tile_sizes [4, 32]
      : (!transform.op<"linalg.matmul">) -> (!transform.any_op, !transform.any_op)
  // 错误：在 % arg1 失效后还使用它
  transform.debug.emit_remark_at %arg1, "remark" : !transform.op<"linalg.matmul">
  transform.yield
}
这就是未定义行为（undefined behavior）：

在 debug 构建里：可能触发断言（assertion）失败，程序崩溃。

在 release 构建里：可能直接段错误（segmentation fault），因为句柄指向的内存已经被释放。


为了避免这种难以调试的未定义行为，Transform Dialect 解释器默认开启一组额外的、可能昂贵的检查，用来检测大多数未定义行为。在这种情况下，会输出清晰的诊断。


sequence.mlir:28:3: error: op uses a handle invalidated by a previously executed transform op
  transform.debug.emit_remark_at %mm, "elemwise_binaries" : !transform.any_op
  ^
sequence.mlir:26:9: note: handle to invalidated ops
  %mm = transform.cast %matmul : !transform.op<"linalg.matmul"> to !transform.any_op
        ^
sequence.mlir:27:19: note: invalidated by this transform op that consumes its operand #0 and invalidates all handles to payload IR entities associated with this operand and entities nested in them
  %loop, %tiled = transform.structured.tile_using_forall %mm tile_sizes [4, 32]
诊断包含三个部分：

主错误：error: op uses a handle invalidated by a previously executed transform op，指出哪一行用了失效句柄。


附注 1：note: handle to invalidated ops，指出这个句柄是怎么来的（这里是通过 transform.cast 从 % matmul 转来的）。


附注 2：note: invalidated by this transform op that consumes its operand #0 ...，指出是哪个变换操作、消耗了哪个操作数导致失效。


当编译时性能是关注点，且转换序列足够稳定时，可以通过向 pass 传递 disable-expensive-checks 选项，或在 applyTransforms 中传入的 TransformOptions 中设置相应标志，来禁用解释器中的昂贵检查，从而提升性能。


一个更隐蔽的失效场景：别名句柄


%casted = transform.cast %arg1 : !transform.op<"linalg.matmul"> to !transform.any_op
%loop, %tiled = transform.structured.tile_using_forall %arg1 tile_sizes [4, 32]
    : (!transform.op<"linalg.matmul">) -> (!transform.any_op, !transform.any_op)
// 错误：% casted 也失效了
transform.debug.emit_remark_at %casted, "remark" : !transform.any_op
这里面：

% arg1 和 % casted 都指向同一个 payload 操作（那个 matmul）。transform.cast 不消耗操作数，所以 % casted 本身没被消耗。但是，tile_using_forall % arg1 消耗了 % arg1，从而擦除了那个 matmul 操作。操作没了，所有指向它的句柄都失效了，包括 % casted。就像两个指针指向同一块内存，内存释放后，两个指针都悬空。


递归失效：嵌套操作也会失效。

如果句柄 A 指向操作 X；操作 X 被擦除；那么不仅 A 失效，所有指向 X 内部嵌套操作的句柄也失效。因为擦除 X 会连带擦除它的 region 和里面所有操作。例如，如果一个句柄指向某个循环内部的 matmul，而外层循环被 outline 消耗，那么这个内部 matmul 的句柄也会失效。


# 用句柄链式组合变换 (Chaining Transformations with Handles)


讲的是：如何把多个变换串联起来，逐步实现分块 + 融合（tile + fuse）的优化。

方法是，不要逐个分块，而是分块最后一个再融合前面。

之前的做法是：直接对 linalg.matmul 做分块。但教程指出，在结构化操作（structured operations）范式里，典型做法是：找到一条无环数据流图（acyclic dataflow graph）中的最后一个操作。对它做分块（tiling）。

然后逐步把产生它操作数的前面那些操作融合（fuse）进循环里。

这样操作的好处：不需要显式地对每个操作都做分块。融合时，循环会自动调整被融合操作的大小（adapt their sizes）。如果需要，还可以注入重计算（inject recomputation）。


当前的例子是： matmul → add → max_signed

 最后一个操作是 max_signed（即 % max）。所以我们：先对 % max 分块；再把 % add 融合进去；再把 % matmul 融合进去。


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_max, %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, %loop_0 =
        transform.structured.fuse_into_containing_op %add into %loop
          : (!transform.any_op, !transform.any_op)
            -> (!transform.any_op, !transform.any_op)
    %matmul_fused, %loop_1 =
        transform.structured.fuse_into_containing_op %arg1 into %loop_0
          : (!transform.op<"linalg.matmul">, !transform.any_op)
            -> (!transform.any_op, !transform.any_op)
    transform.yield
  }
}
第一步：拆分句柄


%add, %max = transform.split_handle %arg2
    : (!transform.op<"linalg.elementwise">)
    -> (!transform.any_op, !transform.any_op)
% arg2 关联了两个 linalg.elementwise 操作（add 和 max_signed）。我们只想对最后一个（max_signed）做分块，所以用 transform.split_handle 把它拆成两个句柄：% add：指向 linalg.elementwise    % max：指向 linalg.elementwise <max_signed>

 第二步：对最后一个操作分块


%tiled_max, %loop =
    transform.structured.tile_using_forall %max tile_sizes [8, 32]
      : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
对 % max 做分块，大小 [8, 32]，得到：% tiled_max：分块后的 max 操作（作用在子集上）。% loop：包围它的 scf.forall 循环。


第三步：逐步融合前面的操作


%add_fused, %loop_0 =
    transform.structured.fuse_into_containing_op %add into %loop
      : (!transform.any_op, !transform.any_op)
        -> (!transform.any_op, !transform.any_op)
%matmul_fused, %loop_1 =
    transform.structured.fuse_into_containing_op %arg1 into %loop_0
      : (!transform.op<"linalg.matmul">, !transform.any_op)
        -> (!transform.any_op, !transform.any_op)
这里做了两次融合：

把 % add 融合进 % loop：% add_fused：融合后的 add 操作； % loop_0：融合后的循环。


把 % arg1（matmul）融合进 % loop_0：% matmul_fused：融合后的 matmul 操作； % loop_1：融合后的循环。


# 更多句柄失效 (More Handle Invalidation)


假设存在一个高效的微内核（microkernel），或者一条以 intrinsic 函数表示的硬件指令，专门做 4×4 矩阵乘法。我们需要将融合操作按所需大小进行 Tile 划分，然后结果函数可以被 4×4 矩阵乘法调用所替换。


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.
    // 拆分 % arg2，对 % max 分块成 8×32，再把 % add 和 % matmul 依次融合进去
    %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, %loop_0 =
        transform.structured.fuse_into_containing_op %add into %loop
          : (!transform.any_op, !transform.any_op)
            -> (!transform.any_op, !transform.any_op)
    %matmul_fused, %loop_1 =
        transform.structured.fuse_into_containing_op %arg1 into %loop_0
          : (!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_2, %loop_2 =
        transform.structured.tile_using_forall %add_fused tile_sizes [4, 4]
          : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
    %matmul_fused_2, %loop_3 =
        transform.structured.fuse_into_containing_op %matmul_fused into %loop_2
          : (!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.
    %_, %outline_target =
        transform.structured.tile_using_forall %tiled_2 tile_sizes [1]
          : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
    transform.structured.fuse_into_containing_op %matmul_fused_2
        into %outline_target
          : (!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.op<"func.call">)
    transform.yield
  }
}
以文章开头的 “全连接 + 偏置 + ReLU” 转换序列为例，查看当前 transform 过程。

原始数据流 ： lhs, rhs → matmul → add → max → return


变换前


func @fc_relu(%lhs, %rhs, %bias, %output) {
  %matmul = linalg.matmul(%lhs, %rhs)        // 512x512
  %biased = linalg.elementwise <add>(%matmul, %bias)   // 512x512
  %relued = linalg.elementwise <max>(%biased, 0)       // 512x512
  return %relued
}
句柄状态：


%arg1 → %matmul
%arg2 → %biased 和 %relued（两个 elementwise）
第 1 步：split_handle % arg2 → % add, % max


%add, %max = transform.split_handle %arg2
只是把句柄拆开，payload IR 不变。


句柄状态：


%add → %biased（add）
%max → %relued（max）
第 2 步：对 % max 分块，tile_sizes [8, 32]


%tiled, %loop = tile_using_forall %max tile_sizes [8, 32]
% max（% relued）被分块，% relued 被擦除，重建为：


一个 scf.forall 循环 % loop，迭代空间 (64, 16)（512/8=64, 512/32=16）


循环体内有一个分块后的 max 操作 % tiled，处理 8×32 的小块

句柄状态：


%add → %biased（仍然有效，add 没被动过）
%arg1 → %matmul（仍然有效）
%tiled → 循环里的 max（8×32）
%loop → scf.forall 循环
第 3 步：把 % add 融合进 % loop


%add_fused, %loop_0 = fuse_into_containing_op %add into %loop
% biased（add）被融合进循环，add 操作被移动 / 重建到循环体内，原来的 % biased 消失。

句柄状态：


%arg1 → %matmul（仍有效）
%add_fused → 循环里的 add（8×32）
%loop_0 → 新的 scf.forall
第 4 步：把 % arg1（matmul）融合进 % loop_0


%matmul_fused, %loop_1 = fuse_into_containing_op %arg1 into %loop_0
% matmul 被融合进循环，matmul 被移动 / 重建到循环体内。

句柄状态：


%arg1 → 失效（matmul 被融合重建）
%matmul_fused → 循环里的 matmul
%loop_1 → 新的 scf.forall
第 5 步：对 % add_fused 二次分块，tile_sizes [4, 4]


%tiled_2, %loop_2 = tile_using_forall %add_fused tile_sizes [4, 4]
% add_fused 被擦除，重建为：


一个新的内层 scf.forall 循环 % loop_2，迭代空间在 8×32 块内是 (2, 8)（8/4=2, 32/4=8）


循环体内有分块后的 add 操作 % tiled_2，处理 4×4 小块


句柄状态：


%add_fused → 失效
%tiled_2 → 内层循环里的 add（4×4）
%loop_2 → 内层 scf.forall
第 6 步：把 % matmul_fused 融合进 % loop_2


%matmul_fused_2, %loop_3 = fuse_into_containing_op %matmul_fused into %loop_2
matmul 被融合进内层 4×4 循环。

句柄状态：


%matmul_fused → 失效
%matmul_fused_2 → 内层循环里的 matmul
%loop_3 → 内层 scf.forall
第 7 步：用 size 1 分块 materialize 外层循环


%_, %outline_target = tile_using_forall %tiled_2 tile_sizes [1]
% tiled_2（内层循环里的 add）被再次分块，size 1，materialize 出一个新的最内层循环 % outline_target。

句柄状态：


%tiled_2 → 失效
%outline_target → 新造的最内层循环（size 1）
第 8 步：把 % matmul_fused_2 融合进 % outline_target


fuse_into_containing_op %matmul_fused_2 into %outline_target
matmul 被融合进 % outline_target。

句柄状态：


%matmul_fused_2 → 失效
%outline_target → 仍有效（还没被 outline 消耗）
第 9 步：outline % outline_target


%func, %call = transform.loop.outline %outline_target {func_name = "outlined"}
把 % outline_target 循环（连同内部的 matmul 和 add）整体搬到一个独立函数 @outlined，原位置留下 func.call @outlined。

句柄状态：


%outline_target → 失效
%matmul_fused_2、%tiled_2 → 失效（嵌套失效）
%func → 新函数 @outlined
%call → func.call @outlined
第 10 步：目标 —— 替换成微内核调用（未实现）


下一步应该要将 4*4 矩阵乘，替换为微内核调用，可能会在后续章节中提到。


# 后记


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


Edited on
Give me a cup of [coffee]~(￣▽￣)~*
XianMu WeChat Pay
WeChat Pay
XianMu Alipay
Alipay