# 前言

本篇是学习 MLIR 中的 Transform Dialect 的第五篇,主要讲的是如何在 Transform Dialect 中通过 “匹配操作(Match Operations)” 来自动发现和选择 payload IR 中需要变换的操作。内容来自 MLIR 官方教程 《Chapter 4: Matching Payload with Transform Operations》。
相关链接: LLVM Project ,MLIR 官方文档,MLIR 官网教程。
作为初学者,错误在所难免,还望不吝赐教。

# 简介

到目前为止,前面几章一直假设在调用转换方言解释器时,特定的负载操作由调用者识别。这可能与从方言驱动转换的理念相违背,因为转换目标必须通过解释器外部的机制来确定,例如在 C++ 中以编程方式调用解释器,或如前几章所示通过传递参数。
在前几章中,Transform Dialect 脚本假设 payload 中要变换的操作是由调用者在外部指定的。我们先以第一章的例子回顾一下这个调用者外部手动指定,以下是 Transform 脚本的入口签名:

transform.named_sequence @__transform_main(
    %arg0: !transform.any_op,
    %arg1: !transform.op<"linalg.matmul">,
    %arg2: !transform.op<"linalg.elementwise">) {
  ...
}

这里有三个参数,% arg0 !transform.any_op 第一个参数,通常是 pass root;% arg1 !transform.op<"linalg.matmul"> 第二个参数,期望绑定到 linalg.matmul;% arg2 !transform.op<"linalg.elementwise"> 第三个参数,期望绑定到 linalg.elementwise。
调用者通过指令手动调用, debug-bind-trailing-args=linalg.matmul,linalg.elementwise 。

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

debug-bind-trailing-args=linalg.matmul,linalg.elementwise 把 Transform 脚本 @__transform_main 的尾部参数(trailing arguments)依次绑定到 payload 中匹配给定名字的操作上。
绑定好之后,可以从 state 中取出来,即后续通过 auto payload = state.getPayloadOps(getCall()); 取出。 如果没有手动绑定,state 里就是空的。
总结来说,debug-bind-trailing-args 是外部绑定,把 @__transform_main 的尾部参数关联到 payload 操作;apply 里的 for (Operation *payloadOp : payload) 是消费已经建立好的关联;没有绑定,state.getPayloadOps (...) 返回空,循环不执行。

这算是一种临时的、根据操作名进行的匹配,与 Transform Dialect 的设计理念不符 (变换应该由 dialect 内部驱动,而不是依赖外部机制);增加使用成本 (C++ 交互复杂,认知负担大)。因此 Transform Dialect 引入了一类特殊的操作 ——Match Operations,用于在 payload IR 中自动匹配符合条件的操作。

# Simple match

以下是 Chapter 1 的 “fully connected layer” 例子:

// Original function to optimize.
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>
    ins(%biased, %c0f : tensor<512x512xf32>, f32)
    outs(%output : tensor<512x512xf32>) -> tensor<512x512xf32>
  func.return %relued : tensor<512x512xf32>
}

在这一章节中,需要调用额外的命令行外部参数 bind-first-extra-to-ops=linalg.matmul bind-second-extra-to-ops=linalg.elementwise ,把这两个操作绑给 Transform 变换脚本。但与其外部绑定,不如用 Match Operations 在 payload IR 里自己发现相关操作:

// The module containing named sequences must have an attribute allowing them
// to enable verification.
module @transforms attributes { transform.with_named_sequence } {
  // Entry point. This takes as the only argument the root operation (typically
  // pass root) given to the transform interpreter.
  transform.named_sequence @__transform_main(
      %root: !transform.any_op {transform.readonly}) {
    // Collect operations that match the criteria specified in named sequence.
    // If the named sequence fails with a silenceable failure, silences it (the
    // message is forwarded to the debug stream). If the named sequence
    // succeeds, appends its results to the results of this operation.
    %elemwise = transform.collect_matching @match_elemwise in %root
      : (!transform.any_op) -> !transform.any_op
    %matmul = transform.collect_matching @match_matmul in %root
      : (!transform.any_op) -> !transform.any_op
    transform.include @print_elemwise failures(propagate)  (%elemwise)
      : (!transform.any_op) -> ()
    transform.include @print_matmul failures(propagate)  (%matmul)
      : (!transform.any_op) -> ()
    transform.yield
  }
  // This is a matcher sequence. It is given an operation to match and the
  // match is considered successful unless any nested operation produces a
  // failure. The values yielded by this operation will be forwarded to the
  // rewriter sequence on success.
  transform.named_sequence @match_elemwise(
      %entry: !transform.any_op {transform.readonly}) -> !transform.any_op {
    transform.match.operation_name %entry ["linalg.elementwise"]
      : !transform.any_op
    transform.yield %entry : !transform.any_op
  }
  transform.named_sequence @match_matmul(
      %entry: !transform.any_op {transform.readonly}) -> !transform.any_op {
    transform.match.operation_name %entry ["linalg.matmul"] : !transform.any_op
    transform.yield %entry : !transform.any_op
  }
  // This is a rewriter sequence.
  transform.named_sequence @print_elemwise(
      %elemwise_binary: !transform.any_op {transform.readonly}) {
    transform.debug.emit_remark_at
      %elemwise_binary, "elementwise binary" : !transform.any_op
    transform.yield
  }
  transform.named_sequence @print_matmul(
      %matmul: !transform.any_op {transform.readonly}) {
    transform.debug.emit_remark_at %matmul, "matmul" : !transform.any_op
    transform.yield
  }
}

引入新的核心操作是 transform.collect_matching @match_elemwise in %root , 意思是在 % root 关联的 payload IR 里,遍历所有操作,对每个操作尝试执行 @match_elemwise 这个命名序列。如果匹配成功,就把结果收集起来;如果匹配失败(silenceable failure),就静默忽略,继续试下一个。
与之配套的是 transform.include,用来在匹配成功后调用一个 “重写序列”:

transform.include @print_elemwise failures(propagate) (%elemwise)

意思是:把 % elemwise 作为参数传给 @print_elemwise 执行;如果里面出现失败,就向上传播。
那再来看看具体的匹配序列和重写序列:

transform.named_sequence @match_elemwise(
    %entry: !transform.any_op) -> !transform.any_op {
  transform.match.operation_name %entry ["linalg.elementwise"]
    : !transform.any_op
  transform.yield %entry : !transform.any_op
}

transform.match.operation_name 检查它是不是 linalg.elementwise;如果是,transform.yield 把 % entry 作为匹配结果返回;如果不是, match.operation_name 产生 silenceable failure,整个匹配序列失败, collect_matching 会静默忽略这个候选,继续下一个。

transform.named_sequence @print_elemwise(
      %elemwise_binary: !transform.any_op) {
    transform.debug.emit_remark_at
      %elemwise_binary, "elementwise binary" : !transform.any_op
    transform.yield
  }

重写序列只是进行了打印操作。
在 debug 构建中,增加 -debug-only=transform-matcher 能够得到打印信息:

<...>
[transform-matcher] matching %0 = linalg.matmul ins(%arg0, %arg1 : tensor<512x512xf32>, tensor<512x512xf32>) outs(%arg3 : tensor<512x512xf32>) -> tensor<512x512xf32> @0x5622eee08410
[transform-matcher] matcher match_elemwise failed: wrong operation name
<...>

是因为 解释器尝试用 @match_elemwise 去匹配一个 linalg.matmul;因为操作名不对,匹配失败。

# Matching Chains of Operations

上面的匹配器仍然很 naive,它仍然是靠操作名称匹配 payload root 下所有某种类型的操作。这些操作可能彼此无关,比如属于不同函数;即使在同一函数里,如果有多组这样的操作,也无法区分。
也就是说,只按操作名匹配,无法区分操作之间的依赖关系,也无法区分多组同名操作。
具体来说,上一节的 @match_elemwise 会匹配到: 那个 elementwise;<max_signed> 那个 elementwise;但它们只是 “名字都叫 linalg.elementwise”,并没有表达出操作链关系:

matmul -> elementwise <add> -> elementwise <max_signed>

现实中我们想要的匹配是:一个 matmul 产生的结果被一个 elementwise 使用,而这个 elementwise 的结果又被另一个 elementwise 使用。也就是说,要匹配的是一组有 use-def 依赖关系的操作,而不是 “所有同名操作”。
所以需要一个新匹配器 @match_matmul_elemwise

transform.named_sequence @match_matmul_elemwise(
    %last: !transform.any_op {transform.readonly})
    -> (!transform.any_op, !transform.any_op, !transform.any_op) {
  // 1. 最后一个操作必须是 elementwise
  transform.match.operation_name %last ["linalg.elementwise"]
    : !transform.any_op
  // 2. 它的第一个操作数必须由另一个操作产生
  %middle = transform.get_producer_of_operand %last[0]
    : (!transform.any_op) -> !transform.any_op
  // 3. 那个定义操作也必须是 elementwise
  transform.match.operation_name %middle ["linalg.elementwise"]
    : !transform.any_op
  // 4. 那个操作的第一个操作数必须由另一个操作产生
  %matmul = transform.get_producer_of_operand %middle[0]
    : (!transform.any_op) -> !transform.any_op
  // 5. 那个操作必须是 matmul
  transform.match.operation_name %matmul ["linalg.matmul"] : !transform.any_op
  // 6. 返回三个句柄:matmul、中间的 elementwise、最后的 elementwise
  transform.yield %matmul, %middle, %last
    : !transform.any_op, !transform.any_op, !transform.any_op
}

其关键在于从后往前匹配,从 use-def 链的最后一个操作开始,往回找。为什么往前匹配呢:因为一个操作数(use)只有一个定义(definition);但从一个定义出发,可能有多个 use;所以往回找是确定的、唯一的,往前找会分叉。
get_producer_of_operand 本身会检查操作数是否存在,不存在就失败,匹配器整体失败。
匹配器逐级检查,最后返回三个句柄。即使 payload 里还有其他 elementwise 和 matmul,这个匹配器也只返回有依赖关系的那一组;它返回的是三元组(matmul, middle, last),而不是按发现顺序返回一堆操作。
入口里使用方法:

// Alternative entry point.
transform.named_sequence @__transform_main(
    %root: !transform.any_op {transform.readonly}) {
  // Collect groups of operations that match the criteria specified in the
  // named sequence.
  %matmul, %el1, %el2 = transform.collect_matching @match_matmul_elemwise in %root
    : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
  %elemwise = transform.merge_handles %el1, %el2 : !transform.any_op
  transform.include @print_elemwise failures(propagate)  (%elemwise)
    : (!transform.any_op) -> ()
  transform.include @print_matmul failures(propagate)  (%matmul)
    : (!transform.any_op) -> ()
  transform.yield
}

需要注意的是:上一节的 collect_matching 返回一个句柄,这一节返回三个。
这个匹配器比上一节精确,但是仍然存在不足。
1. 它硬编码了 “第一个操作数”:

transform.get_producer_of_operand %last[0]

这要求 elementwise 的第一个操作数必须由另一个操作产生。但很多二元操作是可交换的,比如:

max(add(matmul(...), bias), 0)
max(0, add(matmul(...), bias))

这两种写法语义相同,但第一种的第一个操作数是 add 的结果,第二种的第一个操作数是 0。硬编码 [0] 只能匹配其中一种。
2. 它还要求 elementwise 的 operand 位置固定,不够通用。

# Defining Match Operations

上一小节中提到,匹配器硬编码了操作数的位置。
这两种写法语义相同,但第一种的第一个操作数是 add 的结果,第二种的第一个操作数是 0。硬编码 [0] 只能匹配其中一种。
2. 它还要求 elementwise 的 operand 位置固定,不够通用。
Defining Match Operations

上一小节中提到,匹配器硬编码了操作数的位置。

%middle = transform.get_producer_of_operand %last[0]
%matmul = transform.get_producer_of_operand %middle[0]

这要求:最后一个 elementwise 的第一个操作数必须由前一个 elementwise 产生;中间 elementwise 的第一个操作数必须由 matmul 产生。
但很多二元操作是可交换的,比如:

max(add(matmul(...), bias), 0)
max(0, add(matmul(...), bias))

这两种写法语义相同,但第二种的第一个操作数是 0,不是 add 的结果。硬编码 [0] 只能匹配第一种。
所以需要一个更通用的匹配操作,如果任意一个操作数满足嵌套匹配条件,就算成功;同时返回匹配到的操作数位置。

一个新的 Transform 操作,HasOperandSatisfyingOp:

// Define the new operation. By convention, prefix its name with `match`
// followed by the name of the dialect extension.
def HasOperandSatisfyingOp : TransformDialectOp<"match.my.has_operand_satisfying",
    [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
     DeclareOpInterfaceMethods<TransformOpInterface>,
     // Indicate that the operation implements MatchOpInterface in addition to
     // the TransformOpInterface. This interface is only used as a tag at this
     // point and has no methods that are mandatory to implement.
     MatchOpInterface,
     SingleBlockImplicitTerminator<"::mlir::transform::YieldOp">]> {
  let summary = "Succeed if any of the operands matches all nested criteria";
  let arguments = (ins TransformHandleTypeInterface:$op);
  let results = (outs TransformParamTypeInterface:$position,
                      Variadic<Transform_AnyHandleOrParamType>:$results);
  // Match operations can be arbitrarily complex, e.g., containing regions.
  let regions = (region SizedRegion<1>:$body);
  let hasVerifier = 1;
  let assemblyFormat = [{
    $op `:` functional-type($op, results) attr-dict-with-keyword $body
  }];
}

MatchOpInterface 表示标记这是 Match Operation,不修改 payload。
这个接口没有额外方法,只是一个验证契约,表示这个操作是用于匹配的,不会试图变换 payload。
核心逻辑在 apply 里,思路是:遍历 payload 操作的所有操作数,对每个操作数,把嵌套匹配体执行一遍;如果所有嵌套匹配都成功,就认为这个操作数匹配成功,返回它的位置和 yield 出来的值。

// Matcher ops implement `apply` similarly to other transform ops. They are not
// expected to modify payload, but use the tri-state result to signal failure or
// success to match, as well as potential irrecoverable errors.
mlir::DiagnosedSilenceableFailure
mlir::transform::HasOperandSatisfyingOp::apply(
    mlir::transform::TransformRewriter &rewriter,
    mlir::transform::TransformResults &results,
    mlir::transform::TransformState &state) {
  // For simplicity, only handle a single payload op. Actual implementations
  // can use `SingleOpMatcher` trait to simplify implementation and document
  // this expectation.
  auto payloadOps = state.getPayloadOps(getOp()); 
  if (!llvm::hasSingleElement(payloadOps))  // 简单起见,就处理了单个 payload 操作
    return emitSilenceableError() << "expected single payload";
  // Iterate over all operands of the payload op to see if they can be matched
  // using the body of this op.
  Operation *payload = *payloadOps.begin();
  for (OpOperand &operand : payload->getOpOperands()) {  // 遍历操作数
    // Create a scope for transform values defined in the body. This corresponds
    // to the syntactic scope of the region attached to this op. Any values
    // associated with payloads from now on will be automatically dissociated
    // when this object is destroyed, i.e. at the end of the iteration.
    // Associate the block argument handle with the operand.
    auto matchScope = state.make_region_scope(getBody());
    if (failed(state.mapBlockArgument(getBody().getArgument(0),{operand.get()}))) {
      return DiagnosedSilenceableFailure::definiteFailure();
    }
    // Iterate over all nested matchers with the current mapping and see if they succeed.
    bool matchSucceeded = true;  
    for (Operation &matcher : getBody().front().without_terminator()) {  // 执行嵌套匹配
      // Matcher ops are applied similarly to any other transform op.
      DiagnosedSilenceableFailure diag =
          state.applyTransform(cast<TransformOpInterface>(matcher));
      // Definite failures are immediately propagated as they are irrecoverable.
      if (diag.isDefiniteFailure())
        return diag;
      // On success, keep checking the remaining conditions.
      if (diag.succeeded())
        continue;
      // Report failure-to-match for debugging purposes and stop matching this operand.
      assert(diag.isSilenceableFailure());
      DEBUG_MATCHER(DBGS_MATCHER()
                    << "failed to match operand #" << operand.getOperandNumber()
                    << ": " << diag.getMessage());
      (void)diag.silence();
      matchSucceeded = false;
      break;
    }
    // If failed to match this operand, try other operands.
    if (!matchSucceeded)
      continue;
    // If we reached this point, the matching succeeded for the current operand.
    // Remap the values associated with terminator operands to be associated
    // with op results, and also map the parameter result to the operand's
    // position. Note that it is safe to do here despite the end of the scope
    // as `results` are integrated into `state` by the interpreter after `apply`
    // returns rather than immediately.
    SmallVector<SmallVector<MappedValue>> yieldedMappings;
    transform::detail::prepareValueMappings(
        yieldedMappings, getBody().front().getTerminator()->getOperands(),
        state);
    results.setParams(cast<OpResult>(getPosition()),
                      {rewriter.getI32IntegerAttr(operand.getOperandNumber())});
    for (auto &&[result, mapping] : llvm::zip(getResults(), yieldedMappings))
      results.setMappedValues(result, mapping);
    return DiagnosedSilenceableFailure::success();
  }
  // If we reached this point, none of the operands succeeded the match.
  return emitSilenceableError()
         << "none of the operands satisfied the conditions";
}

根据惯例,实现 MatchOpInterface 的操作不得修改有效载荷 IR,因此必须明确指出其仅读取操作数句柄和有效载荷作为其效果。

void transform::CollectMatchingOp::getEffects(
    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
  onlyReadsHandle(getRoot(), effects);
  producesHandle(getResults(), effects);
  onlyReadsPayload(effects);
}

此操作现在可以被包含在转换方言扩展中,加载并用于我们的匹配器。具体来说,我们将使用它来表示示例中 “max” 逐元素操作的任一操作数,均可由前一个逐元素操作生成。为简化起见,前一个操作仍需依赖矩阵乘法(matmul)来生成第一个操作数。更新后的匹配器序列如下所示。

transform.named_sequence @match_matmul_elemwise(
    %last: !transform.any_op {transform.readonly})
    -> (!transform.any_op, !transform.any_op, !transform.any_op,
        !transform.param<i32>) {
  // The last operation must be an elementwise binary.
  transform.match.operation_name %last ["linalg.elementwise"]
    : !transform.any_op
  // One of its operands must be defined by another operation, to which we
  // will get a handle here. This is achieved thanks to a newly defined
  // operation that tries to match operands one by one using the match
  // operations nested in its region.
  %pos, %middle = transform.match.my.has_operand_satisfying %last
      : (!transform.any_op) -> (!transform.param<i32>, !transform.any_op) {
  ^bb0(%operand: !transform.any_value):
    // The operand must be defined by an operation.
    %def = transform.get_defining_op %operand
      : (!transform.any_value) -> !transform.any_op
    // The defining operation must itself be an elementwise binary.
    transform.match.operation_name %def ["linalg.elementwise"]
      : !transform.any_op
    transform.yield %def : !transform.any_op
  }
  // And the first operand of that operation must be defined by yet another
  // operation.
  %matmul = transform.get_producer_of_operand %middle[0]
    : (!transform.any_op) -> !transform.any_op
  // And that operation is a matmul.
  transform.match.operation_name %matmul ["linalg.matmul"] : !transform.any_op
  // We will yield the handles to the matmul and the two elementwise
  // operations separately.
  transform.yield %matmul, %middle, %last, %pos
    : !transform.any_op, !transform.any_op, !transform.any_op,
      !transform.param<i32>
}

关键变化点:用 has_operand_satisfying 替代硬编码的 % last [0];body 里检查当前操作数是否由 linalg.elementwise 产生;成功时返回 % pos(操作数位置)和 % middle(匹配到的 elementwise);后面的 matmul 检查仍然用 % middle [0],为简单起见保留。
这实现了预期的效果,并使 max (add (matmul (...), bias), 0) 和 max (0, add (matmul (...), bias)) 的结果保持一致。因为 has_operand_satisfying 不关心是第几个操作数,只关心有没有一个操作数满足条件。
% pos 值是转换方言中的参数,用于存储在转换应用过程中已知为常量的实体列表。通常情况下,参数为数值,但一般可以是任何 MLIR 属性。这里 % pos 是!transform.param,存的是匹配到的操作数位置。后续可以根据它区分不同的匹配组。

用 transform.foreach_match 演示独立匹配。

// Traverses the payload IR associated with the operand handle, invoking
// @match_matmul_elemwise on each of the operations. If the named sequence
// succeeds, i.e., if none of the nested match (transform) operations
// produced a silenceable failure, invokes @print_matmul_elemwise and
// forwards the values yielded as arguments of the new invocation. If the
// named sequence fails with a silenceable failure, silences it (the message
// is forwarded to the debug stream). Definite failures are propagated
// immediately and unconditionally, as usual.
transform.foreach_match in %root
  @match_matmul_elemwise -> @print_matmul_elemwise
  : (!transform.any_op) -> !transform.any_op

遍历 % root 关联的 payload 操作;对每个操作执行 @match_matmul_elemwise ;如果匹配成功(没有 silenceable failure),调用 @print_matmul_elemwise ,把 yield 出来的值作为参数传入;如果匹配失败(silenceable failure),静默它;definite failure 照常立即传播。
@print_matmul_elemwise 会用 % pos 参数区分两组匹配。这样就能演示:不同组的操作可以独立匹配、独立处理。

# Matchers for Inferred Features

前几节匹配器仍有局限,它们比较基础,主要依赖操作名和 use-def 链。用 API 或声明式重写规则也能实现类似效果,有时甚至更简洁。
Transform Dialect 真正的威力在于:可以匹配 payload 的 “推断特征”—— 那些不是操作属性、也不是 IR 组件之间简单关系的性质。这些特征无法通过声明式方式表达,需要实际分析操作的结构、属性、body 等。
用一个例子说明:
linalg.generic 表示的矩阵乘法,前几节用的是 %matmul = linalg.matmul ... ,按操作名匹配 linalg.matmul 很简单。但如果矩阵乘法是用 linalg.generic 写的:

%matmul = linalg.generic {
  indexing_maps = [...],
  iterator_types = ["parallel", "parallel", "reduction"]
} ins(...) outs(...) {
  ^bb0(...):
    %mul = arith.mulf ...
    %add = arith.addf ...
    linalg.yield %add
}

这时:操作名是 linalg.generic,不是 linalg.matmul;前几节的 match.operation_name % last ["linalg.matmul"] 就完全不适用了。
但语义上它仍然是矩阵乘法,而且这种表示在变换过程中很常见,比如高维 contraction 被分解成围绕矩阵乘法的循环。
所以需要一个能推断 “这是不是矩阵乘法” 的匹配器。
那么,什么样的特征才是一个矩阵乘法呢?教程列出矩阵乘法在 linalg.generic 上应具备的特征:

Total rank of 3.
Two inputs accessed as projected permutation of iteration dimensions.
One output accessed as projected permutation of iteration dimensions.
Iteration dimensions can be subdivided into LHS parallel, RHS parallel and reduction dimensions.
The body block consists of a multiplication and an addition.

这些特征:不是操作的直接属性;不能只靠检查一个 attribute 或 type 得到;需要分析 indexing_maps 、 iterator_types 、 body 等。

这正是 Transform Dialect 能做的事:在 matcher op 的 apply 里实现这些分析,并复用。

新的匹配序列:

transform.named_sequence @match_generic_matmul(
    %candidate: !transform.any_op {transform.readonly}) -> !transform.any_op {
  // Match a structured linear algebra operation.
  transform.match.structured %candidate : !transform.any_op {
  ^bb0(%c: !transform.any_op):
    // With a rank equal to 3.
    %rank = transform.match.structured.rank %c
      : (!transform.any_op) -> !transform.param<i64>
    %c3 = transform.param.constant 3 : i64 -> !transform.param<i64>
    transform.match.param.cmpi eq %rank, %c3 : !transform.param<i64>
    // With 2 inputs.
    %n_ins = transform.match.structured.num_inputs %c
      : (!transform.any_op) -> !transform.param<i64>
    %c2 = transform.param.constant 2 : i64 -> !transform.param<i64>
    transform.match.param.cmpi eq %n_ins, %c2 : !transform.param<i64>
    // With 1 output (note that structured ops in destination passing style
    // has as many inits as outputs).
    %n_inits = transform.match.structured.num_inits %c
      : (!transform.any_op) -> !transform.param<i64>
    %c1 = transform.param.constant 1 : i64 -> !transform.param<i64>
    transform.match.param.cmpi eq %n_inits, %c1 : !transform.param<i64>
    // All inputs and inits are accessed with a projected permutation.
    transform.match.structured.input %c[all] {projected_permutation}
      : !transform.any_op
    transform.match.structured.init %c[0] {projected_permutation}
      : !transform.any_op
    // The body is a mulf/addf contraction with appropriate dimensions.
    transform.match.structured.body %c
      { contraction = ["arith.mulf", "arith.addf"] } : !transform.any_op
    %batch, %lhs, %rhs, %reduction =
    transform.match.structured.classify_contraction_dims %c
      : (!transform.any_op)
      -> (!transform.param<i64>, !transform.param<i64>, !transform.param<i64>,
          !transform.param<i64>)
    // There is one of lhs, rhs and reduction dimensions and zero batch
    // dimensions.
    %n_batch = transform.num_associations %batch
      : (!transform.param<i64>) -> !transform.param<i64>
    %n_lhs = transform.num_associations %lhs
      : (!transform.param<i64>) -> !transform.param<i64>
    %n_rhs = transform.num_associations %rhs
      : (!transform.param<i64>) -> !transform.param<i64>
    %n_reduction = transform.num_associations %reduction
      : (!transform.param<i64>) -> !transform.param<i64>
    %c0 = transform.param.constant 0 : i64 -> !transform.param<i64>
    transform.match.param.cmpi eq %n_batch, %c0 : !transform.param<i64>
    transform.match.param.cmpi eq %n_lhs, %c1 : !transform.param<i64>
    transform.match.param.cmpi eq %n_rhs, %c1 : !transform.param<i64>
    transform.match.param.cmpi eq %n_reduction, %c1 : !transform.param<i64>
  }
  transform.yield %candidate : !transform.any_op
}

该匹配序列借助 transform.match.structured 依次检查了上述条件。

# 后记

春风若有怜花意,可否许我再少年。
                                        —— 《咏苍髯》

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

Edited on

Give me a cup of [coffee]~( ̄▽ ̄)~*

XianMu WeChat Pay

WeChat Pay

XianMu Alipay

Alipay