[Mlir-commits] [mlir] [mlir][linalg] Canonicalize non-identity `linalg.generic` ops (PR #101430)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-linalg

Author: Ian Wood (IanWood1)

<details>
<summary>Changes</summary>

Extend `linalg.generic`'s canonicalization patterns to be able to erase ops with non-identity indexing maps but that are still noops.

---
Full diff: https://github.com/llvm/llvm-project/pull/101430.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp (+6-3) 
- (modified) mlir/test/Dialect/Linalg/canonicalize.mlir (+20) 


``````````diff
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 99b625d99fec2..34f403330621f 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -1217,9 +1217,12 @@ struct EraseIdentityLinalgOp : public OpRewritePattern<OpTy> {
 
   LogicalResult matchAndRewrite(OpTy linalgOp,
                                 PatternRewriter &rewriter) const override {
-    // Check all indexing maps are identity.
-    if (llvm::any_of(linalgOp.getIndexingMapsArray(),
-                     [](AffineMap map) { return !map.isIdentity(); }))
+    // All indexing maps must be equal permutations
+    auto indexingMaps = linalgOp.getIndexingMapsArray();
+    if (!llvm::all_equal(indexingMaps))
+      return failure();
+
+    if (!indexingMaps.empty() && !indexingMaps.front().isPermutation())
       return failure();
 
     // Check that the body of the linalg operation is just a linalg.yield
diff --git a/mlir/test/Dialect/Linalg/canonicalize.mlir b/mlir/test/Dialect/Linalg/canonicalize.mlir
index b1212b8863a5d..a50fbb0fc3b86 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -911,6 +911,26 @@ func.func @identity_buffer(%arg0 : memref<?xf32>, %arg1: memref<?xf32>) {
 
 // -----
 
+#map = affine_map<(d0, d1) -> (d1, d0)>
+func.func @erase_non_identity_noop(%arg0 : tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %0 = linalg.generic {
+    indexing_maps = [#map, #map],
+    iterator_types = ["parallel", "parallel"]
+  } ins(%arg0 : tensor<?x?xf32>)
+    outs(%arg1 : tensor<?x?xf32>) {
+  ^bb0(%in: f32, %out: f32):
+    linalg.yield %in: f32
+  } -> tensor<?x?xf32>
+  return %0 : tensor<?x?xf32> 
+}
+
+// Do not erase ops with buffer semantics.
+// CHECK-LABEL: func @erase_non_identity_noop
+//  CHECK-SAME:   (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: tensor<?x?xf32>)
+//       CHECK:   return %[[ARG0]] : tensor<?x?xf32>
+
+// -----
+
 // Just make sure that we don't crash.
 
 // CHECK-LABEL: func @dedeplicate_regression_test

``````````

</details>


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