[Mlir-commits] [mlir] [mlir] add a chapter on matchers to the transform dialect tutorial (PR #76725)

[Oleksandr Alex Zinenko]

================
@@ -0,0 +1,578 @@
+# Chapter 4: Matching Payload with Transform Operations
+
+**Check the continuously-tested version of MLIR files under
+[mlir/test/Examples/transform/Ch4](https://github.com/llvm/llvm-project/tree/main/mlir/test/Examples/transform/Ch4).**
+
+Up until now, we were applying transform dialect scripts under the assumption
+that specific payload operations are identified by the caller when the transform
+dialect interpreter is invoked. This may be seen as contrary to the idea of
+driving transformations from a dialect since the transformation targets must be
+identified by the caller in C++. It also adds practical overhead due to
+increased interaction with the interpreter in C++, and cognitive overhead of
+manipulating two interfaces at once. To remedy this, Transform dialect proposes
+a subset of operations for _matching_ payload operations that need to be
+transformed.
+
+_Match_ operations are simply transform operations with some additional
+guarantees. In particular, they are not expected to modify the payload IR and
+are expected to fail if their operands (typically payload operation handles) are
+not associated with payload IR objects having desired properties, such as
+operation names or kinds of arguments. Using simple combinator operations, it
+becomes possible to set up a higher-level match and rewrite infrastructure
+directly within the transform dialect.
+
+
+## Simple match
+
+Let us reconsider the “fully connected layer” example from Chapter 1, reproduced
+below for convenience.
+
+
+```mlir
+// 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.elemwise_binary { fun = #linalg.binary_fn<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.elemwise_binary { fun = #linalg.binary_fn<max_signed> }
+    ins(%biased, %c0f : tensor<512x512xf32>, f32)
+    outs(%output : tensor<512x512xf32>) -> tensor<512x512xf32>
+  func.return %relued : tensor<512x512xf32>
+}
+
+```
+
+
+In Chapter 1, we were calling the test transform interpreter pass with
+additional arguments, `bind-first-extra-to-ops=linalg.matmul
+bind-second-extra-to-ops=linalg.elemwise_binary`, to provide initial
+associations for operation handles. Instead, we can use match operations to
+discover relevant operations in the payload IR. Match operations can be combined
+with “regular” transform operations using, e.g., the
+`transform.collect_matching` combinator operation that leverages the concept of
+named sequences to organize matchers.
+
+
+```mlir
+// 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.elemwise_binary"]
+      : !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.test_print_remark_at_operand
+      %elemwise_binary, "elementwise binary" : !transform.any_op
+    transform.yield
+  }
+  transform.named_sequence @print_matmul(
+      %matmul: !transform.any_op {transform.readonly}) {
+    transform.test_print_remark_at_operand %matmul, "matmul" : !transform.any_op
+    transform.yield
+  }
+}
+
+```
+
+
+This script can be executed using the non-test interpreter pass running on the
+root operation of the translation unit without additional flags: `mlir-opt
+--transform-interpreter`. It will emit corresponding remarks at elementwise and
+matmul operations. In debug builds, the infrastructure provides a convenient
+method to understand the matching process by passing
+`-debug-only=transform-matcher` to `mlir-opt` or a derived tool. It will print
+the silenceable failure messages produced by the match operations into the debug
+stream, for example:
+
+
+```
+<...>
+[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
+<...>
+```
+
+
+This is now sufficient to run the rest of the transform script from Chapter 1,
+substituting `%arg1` with `%matmul` and `%arg2` with `%elemwise`.
+
+
+## Matching Chains of Operations
+
+The matcher above remains naive as it matches _all_ operations of the certain
+kind under the payload root. These operations may or may not be related, and
+may, for example, belong to different functions. Even if they are in a single
+function, if there are multiple groups of such operations, we wouldn’t be able
+to differentiate them with this approach. In reality, we want to match a
+specific group of operations where a `matmul` operation produces a result that
+is used by an elementwise operation, which in turn feeds another elementwise
+operation in a similar way.
+
+This can be achieved using the following matcher sequence.
+
+
+```mlir
+// This is also a matcher sequence. It is similarly given an operation to
+// match and nested operations must succeed in order for a match to be deemed
+// successful. It starts matching from the last operation in the use-def chain
+// and goes back because each operand (use) has exactly one definition.
+transform.named_sequence @match_matmul_elemwise(
+    %last: !transform.any_op {transform.readonly})
+    -> (!transform.any_op, !transform.any_op, !transform.any_op) {
+  // The last operation must be an elementwise binary.
+  transform.match.operation_name %last ["linalg.elemwise_binary"]
+    : !transform.any_op
+  // Its first operand must be defined by another operation, to which we
+  // will get a handle here. We are guaranteed that the first operand exists
+  // because we know the operation is binary, but even in absence of such a
+  // guarantee, this operation would have produced a silenceable failure when
+  // `%last` does not have enough operands.
+  %middle = transform.get_producer_of_operand %last[0]
+    : (!transform.any_op) -> !transform.any_op
+  // The defining operation must itself be an elementwise binary.
+  transform.match.operation_name %middle ["linalg.elemwise_binary"]
+    : !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
+    : !transform.any_op, !transform.any_op, !transform.any_op
+}
+```
+
+This matcher is applicable in presence of other `elemwise` and `matmul`
+operations and will return the triple of _related_ operations rather than
+operations in the order in which they are found. It can be exercised similarly
+to the previous incarnation, as follows.
+
+```mlir
+// Alternative entry point.
+transform.named_sequence @__transform_main(
----------------
ftynse wrote:

I prefer to keep it this way in case somebody wants to copy-paste the script when reading, and will get surprised when nothing happens because the entry point name is different.

https://github.com/llvm/llvm-project/pull/76725