[Mlir-commits] [mlir] [MLIR] [Transforms] Let lowerToLoopsUsingSCFForOp delete target op, fixes #83252 (PR #83256)

[llvmlistbot at llvm.org]

https://github.com/lhunloh updated https://github.com/llvm/llvm-project/pull/83256

>From 2f1ef500a3295a3763174b6f1b4718a5a16a38f6 Mon Sep 17 00:00:00 2001
From: Lars <larshunloh at uni-muenster.de>
Date: Wed, 28 Feb 2024 13:15:14 +0100
Subject: [PATCH 1/3] Let lowerToLoopsUsingSCFForOp delete target op.

The function mlir::scf::lowerToLoopsUsingSCFForOp was not deleting its
(structured) target op, resulting in IR with the expected loop nest in
front of the still remaining (structured) op, e.g. a linalg.matmul.
---
 mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
index 1a84a59ddb69df..e1e9be858b251e 100644
--- a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
@@ -1133,5 +1133,6 @@ mlir::scf::lowerToLoopsUsingSCFForOp(RewriterBase &rewriter,
   if (failed(op.generateScalarImplementation(rewriter, op.getLoc(), ivs))) {
     return failure();
   }
+  rewriter.eraseOp(op);
   return loops;
 }

>From 4f6bd46c53778f39f194e5b96c5c9a6f61e31790 Mon Sep 17 00:00:00 2001
From: Lars <larshunloh at uni-muenster.de>
Date: Wed, 28 Feb 2024 19:18:11 +0100
Subject: [PATCH 2/3] Moved target deletion of lowerToLoopsUsingSCFForOp to
 caller

The caller of function `lowerToLoopsUsingSCFForOp`, `ConvertToLoopsOp` was not deleting its
(structured) target op, resulting in IR with the expected loop nest in
front of the still remaining (structured) op, e.g. a linalg.matmul.
---
 mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp | 1 +
 mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp      | 1 -
 2 files changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 299965bcfc3ab3..d542da43049b5c 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -2121,6 +2121,7 @@ DiagnosedSilenceableFailure transform::ConvertToLoopsOp::applyToOne(
       scf::lowerToLoopsUsingSCFForOp(rewriter, target);
   if (failed(loops))
     return emitDefaultDefiniteFailure(target);
+  rewriter.eraseOp(target);
   return DiagnosedSilenceableFailure::success();
 }
 
diff --git a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
index e1e9be858b251e..1a84a59ddb69df 100644
--- a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
@@ -1133,6 +1133,5 @@ mlir::scf::lowerToLoopsUsingSCFForOp(RewriterBase &rewriter,
   if (failed(op.generateScalarImplementation(rewriter, op.getLoc(), ivs))) {
     return failure();
   }
-  rewriter.eraseOp(op);
   return loops;
 }

>From 06b97faebe2c8403667a77e4da2d8a3b6c165ac0 Mon Sep 17 00:00:00 2001
From: Lars <larshunloh at uni-muenster.de>
Date: Thu, 29 Feb 2024 10:40:45 +0100
Subject: [PATCH 3/3] Test deletion of structured op via
 transform.structured.convert_to_loops

This tests that the op is indeed deleted by
transform.structured.convert_to_loops.
---
 .../lower-to-loops-using-interface.mlir       | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/mlir/test/Interfaces/TilingInterface/lower-to-loops-using-interface.mlir b/mlir/test/Interfaces/TilingInterface/lower-to-loops-using-interface.mlir
index 7969de0d456bb6..9f472b607912b9 100644
--- a/mlir/test/Interfaces/TilingInterface/lower-to-loops-using-interface.mlir
+++ b/mlir/test/Interfaces/TilingInterface/lower-to-loops-using-interface.mlir
@@ -36,6 +36,26 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
+func.func @gemm(%arg0 : memref<?x?xf32>, %arg1 : memref<?x?xf32>,
+  %arg2 : memref<?x?xf32>) {
+  linalg.matmul ins(%arg0, %arg1 : memref<?x?xf32>, memref<?x?xf32>)
+      outs(%arg2 : memref<?x?xf32>)
+  return
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {
+    %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1
+      : (!transform.any_op) -> !transform.any_op
+    transform.structured.convert_to_loops %matmul : !transform.any_op
+    transform.yield
+  }
+}
+// CHECK-LABEL: func @gemm
+//   CHECK-NOT:   linalg.matmul ins({{.*}})
+
+// -----
+
 func.func @indexed_generic(%arg0 : memref<200x300xi32>, %arg1 : memref<300xi16>,
     %arg2 : memref<200xi8>, %arg3 : memref<300x200xi64>) {
   linalg.generic {