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

[Andrzej Warzyński]

================
@@ -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,
----------------
banach-space wrote:

```suggestion
This is now sufficient to run the rest of [the transform script from Chapter 1](https://mlir.llvm.org/docs/Tutorials/transform/Ch1/#chaining-transformations-with-handles),
```
? Also, note that that example uses `!transform.op<"linalg.matmul">` rather than `!transform.any_op` (same for `linalg.elemwise_binary`).

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