# 前言
本篇是学习 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 会匹配到:
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,并指明哪一篇博客,看到一定及时回复!