[Mlir-commits] [mlir] [mlir] transform.structured.match loop-like flag (PR #65336)

[Oleksandr Alex Zinenko]

https://github.com/ftynse created https://github.com/llvm/llvm-project/pull/65336:

Add an enum option to `transform.structured.match` operation to match payload operations implementing LoopLikeOpInterface.

>From 6cf663b62174e4f28dbb130275883f2ed7028ebd Mon Sep 17 00:00:00 2001
From: Alex Zinenko <zinenko at google.com>
Date: Tue, 5 Sep 2023 15:16:16 +0000
Subject: [PATCH] [mlir] transform.structured.match loop-like flag

Add an enum option to `transform.structured.match` operation to match
payload operations implementing LoopLikeOpInterface.
---
 .../TransformOps/LinalgTransformEnums.td      |  3 +-
 .../TransformOps/LinalgTransformOps.cpp       |  3 ++
 .../Dialect/Linalg/transform-op-match.mlir    | 29 +++++++++++++++++++
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformEnums.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformEnums.td
index bf997ac2b068c8a..e4cf4e194b7abf6 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformEnums.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformEnums.td
@@ -3,7 +3,8 @@ include "mlir/IR/EnumAttr.td"
 def MatchInterfaceEnum : I32EnumAttr<"MatchInterfaceEnum", "An interface to match",
     [
       I32EnumAttrCase<"LinalgOp", 0>,
-      I32EnumAttrCase<"TilingInterface", 1>
+      I32EnumAttrCase<"TilingInterface", 1>,
+      I32EnumAttrCase<"LoopLikeInterface", 2>,
     ]>{
   let cppNamespace = "mlir::transform";
 }
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 6549c27b0d0dfb4..81ca4286e9b6bb4 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -1147,6 +1147,9 @@ transform::MatchOp::apply(transform::TransformRewriter &rewriter,
       if (iface == transform::MatchInterfaceEnum::TilingInterface &&
           isa<TilingInterface>(op))
         return;
+      if (iface == transform::MatchInterfaceEnum::LoopLikeInterface &&
+          !isa<LoopLikeOpInterface>(op))
+        return;
     }
 
     // Check if all specified attributes match.
diff --git a/mlir/test/Dialect/Linalg/transform-op-match.mlir b/mlir/test/Dialect/Linalg/transform-op-match.mlir
index c6c177c3e86d099..cbc5504a66fc619 100644
--- a/mlir/test/Dialect/Linalg/transform-op-match.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-match.mlir
@@ -76,3 +76,32 @@ transform.sequence failures(propagate) {
 // expected-remark @below {{0}}
   transform.test_print_number_of_associated_payload_ir_ops %no_match : !transform.any_op
 }
+
+// -----
+
+func.func private @callee()
+
+func.func @foo(%lb: index, %ub: index, %step: index) {
+  // expected-remark @below {{loop-like}}
+  scf.for %i = %lb to %ub step %step {
+    func.call @callee() : () -> ()
+    scf.yield
+  }
+  // expected-remark @below {{loop-like}}
+  scf.parallel (%i) = (%lb) to (%ub) step (%step) {
+    func.call @callee() : () -> ()
+    scf.yield
+  }
+  // expected-remark @below {{loop-like}}
+  scf.forall (%i) in (%ub) {
+    func.call @callee() : () -> ()
+  }
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+  %matched = transform.structured.match interface{LoopLikeInterface} in %arg0 : (!transform.any_op) -> !transform.any_op
+  transform.test_print_remark_at_operand %matched, "loop-like" : !transform.any_op
+  transform.yield
+}