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

[llvmlistbot at llvm.org]

Author: Oleksandr "Alex" Zinenko
Date: 2023-09-05T17:56:00+02:00
New Revision: c17735053b22a83a9d0d727775da9339d260ac7d

URL: https://github.com/llvm/llvm-project/commit/c17735053b22a83a9d0d727775da9339d260ac7d
DIFF: https://github.com/llvm/llvm-project/commit/c17735053b22a83a9d0d727775da9339d260ac7d.diff

LOG: [mlir] transform.structured.match loop-like flag (#65336)

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformEnums.td
    mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
    mlir/test/Dialect/Linalg/transform-op-match.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformEnums.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformEnums.td
index bf997ac2b068c8..e4cf4e194b7abf 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 6549c27b0d0dfb..81ca4286e9b6bb 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 c6c177c3e86d09..cbc5504a66fc61 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
+}