[Mlir-commits] [mlir] [mlir] make transform.foreach_match forward arguments (PR #89920)

[Nicolas Vasilache]

================
@@ -78,3 +78,113 @@ module attributes { transform.with_named_sequence } {
     transform.yield
   }
 }
+
+// -----
+
+// expected-remark @below {{op from within the matcher}}
+module attributes { transform.with_named_sequence } {
+  // expected-remark @below {{returned root}}
+  func.func @foo() {
+    return
+  }
+
+  transform.named_sequence @match_fail(
+      %op: !transform.any_op {transform.readonly},
+      %root: !transform.any_op {transform.readonly},
+      %param: !transform.param<i64> {transform.readonly}) -> (!transform.any_op, !transform.param<i64>) {
+    transform.test_succeed_if_operand_of_op_kind %op, "test.impossible_to_match" : !transform.any_op
+    transform.yield %root, %param : !transform.any_op, !transform.param<i64>
+  }
+
+  transform.named_sequence @return(
+      %root: !transform.any_op {transform.readonly},
+      %param: !transform.param<i64> {transform.readonly}) -> (!transform.param<i64>, !transform.param<i64>, !transform.any_op) {
+    %func = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
+    transform.yield %param, %param, %func : !transform.param<i64>, !transform.param<i64>, !transform.any_op
+  }
+
+  transform.named_sequence @match_succeed(
+      %op: !transform.any_op {transform.readonly},
+      %root: !transform.any_op {transform.readonly},
+      %param: !transform.param<i64> {transform.readonly}) -> (!transform.any_op, !transform.param<i64>) {
+    transform.debug.emit_remark_at %root, "op from within the matcher" : !transform.any_op
+    // expected-remark @below {{param from within the matcher 42}}
+    transform.debug.emit_param_as_remark %param, "param from within the matcher" : !transform.param<i64>
+    transform.yield %root, %param : !transform.any_op, !transform.param<i64>
+  }
+
+  transform.named_sequence @__transform_main(%root: !transform.any_op) {
+    %param = transform.param.constant 42 : i64 -> !transform.param<i64>
+    %func = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
+    %func2, %param1, %param2, %roots = transform.foreach_match restrict_root in %func, %root, %param
+      @match_fail -> @return,
+      @match_succeed -> @return
+      : (!transform.any_op, !transform.any_op, !transform.param<i64>) -> (!transform.any_op, !transform.param<i64>, !transform.param<i64>, !transform.any_op)
+    transform.debug.emit_remark_at %roots, "returned root" : !transform.any_op
+    // expected-remark @below {{42 : i64, 42 : i64}}
+    transform.debug.emit_param_as_remark %param1: !transform.param<i64>
+    %num_roots = transform.num_associations %roots : (!transform.any_op) -> !transform.param<i64>
+    // expected-remark @below {{2 : i64}}
+    transform.debug.emit_param_as_remark %num_roots : !transform.param<i64>
+    transform.yield
+  }
+}
+
+// -----
+
+module attributes { transform.with_named_sequence } {
+  func.func private @foo()
+  func.func private @bar()
+
+  transform.named_sequence @return(
+      %func: !transform.any_op {transform.readonly}) -> (!transform.any_op) {
+    transform.yield %func : !transform.any_op
+  }
+
+  transform.named_sequence @match(
+      %op: !transform.any_op {transform.readonly},
+      %func: !transform.any_op {transform.readonly}) -> (!transform.any_op) {
+    transform.yield %func : !transform.any_op
+  }
+
+  transform.named_sequence @__transform_main(%root: !transform.any_op) {
+    %func = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
+    %func2, %func3 = transform.foreach_match flatten_results restrict_root in %func, %func
+      @match -> @return
+      : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+    %num = transform.num_associations %func3 : (!transform.any_op) -> !transform.param<i64>
+    // 2 funcs are yielded for each of the 2 funcs = 4:
----------------
nicolasvasilache wrote:

Starting to review from the test, I am having trouble understanding why `num_associations %func3 == 4` , I would have expected less due to `%func2` in `%func2, %func3 =` but I can't put the finger on it.
Could you please explain?

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