[Mlir-commits] [mlir] [mlir][linalg] Bail out of decomposition when peeled stmt has no results (PR #222618)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


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

@llvm/pr-subscribers-mlir-linalg

Author: Akun (AkunCube)

<details>
<summary>Changes</summary>

Current `DecomposeLinalgOp` peels the first statement of a linalg.generic body into a new generic op and then calls `replaceOpUsesWithinBlock` on the peeled operation to rewire its uses within the residual op.

When the peeled statement produces no results (e.g. a bufferization.dealloc_tensor as the first body statement), replaceOpUsesWithinBlock iterates over zero results and leaves allUsesReplaced vacuously true, tripping the assertion and crashing mlir-opt -test-linalg-decompose-ops.

Reject such ops up front with a notifyMatchFailure instead of letting the decomposition proceed and crash, and add a regression test.

Fixed: #<!-- -->220523

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp (+7) 
- (modified) mlir/test/Dialect/Linalg/decompose-ops.mlir (+36) 


``````````diff
diff --git a/mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp b/mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp
index 3c8d2aadc8c75..2fb5d11f95a37 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp
@@ -266,6 +266,13 @@ DecomposeLinalgOp::matchAndRewrite(GenericOp genericOp,
                                        "operation has less than 3 statements");
   }
 
+  /// The peeled statement is expected to produce at least one result.
+  if (body->getOperations().begin()->getNumResults() == 0) {
+    return rewriter.notifyMatchFailure(
+        &(*body->getOperations().begin()),
+        "peeled statement is expected to have at least one result");
+  }
+
   /// Check that the peeled statement has a scalar element type.
   if (llvm::any_of(body->getOperations().begin()->getResultTypes(),
                    [](Type t) { return !t.isIntOrIndexOrFloat(); })) {
diff --git a/mlir/test/Dialect/Linalg/decompose-ops.mlir b/mlir/test/Dialect/Linalg/decompose-ops.mlir
index d5be02976a663..43bd42528f34f 100644
--- a/mlir/test/Dialect/Linalg/decompose-ops.mlir
+++ b/mlir/test/Dialect/Linalg/decompose-ops.mlir
@@ -416,3 +416,39 @@ func.func @destination_passing_style(
 // CANONICALIZECHECK-NEXT:     %[[S2:.+]] = arith.mulf %[[ARG4]], %[[ARG6]]
 // CANONICALIZECHECK-NEXT:     linalg.yield %[[ARG5]], %[[S2]]
 //      CANONICALIZECHECK:   return %[[GENERIC1]], %[[GENERIC2]]#1
+
+// -----
+
+// The peeled statement is expected to produce at least one result. Here the
+// first statement (`bufferization.dealloc_tensor`) has no results, so the
+// operation is not decomposed.
+#map0 = affine_map<(d0) -> (d0)>
+func.func @first_statement_no_result(%arg0 : tensor<27xi32>) -> tensor<27xi32> {
+  %init = tensor.empty() : tensor<27xi32>
+  %r = linalg.generic {
+    indexing_maps = [#map0, #map0],
+    iterator_types = ["parallel"]}
+    ins(%arg0 : tensor<27xi32>)
+    outs(%init : tensor<27xi32>) {
+    ^bb0(%b0 : i32, %b1 : i32):
+      bufferization.dealloc_tensor %init : tensor<27xi32>
+      %0 = arith.addi %b0, %b1 : i32
+      linalg.yield %0 : i32
+  } -> tensor<27xi32>
+  return %r : tensor<27xi32>
+}
+//  CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0) -> (d0)>
+//      CHECK: func.func @first_statement_no_result(
+//      CHECK:   %[[GENERIC:.+]] = linalg.generic
+//      CHECK:   bufferization.dealloc_tensor
+//      CHECK:   arith.addi
+//      CHECK:   linalg.yield
+//      CHECK:   return %[[GENERIC]]
+
+// CANONICALIZECHECK-DAG: #[[MAP0:.+]] = affine_map<(d0) -> (d0)>
+//      CANONICALIZECHECK: func.func @first_statement_no_result(
+//      CANONICALIZECHECK:   %[[GENERIC:.+]] = linalg.generic
+//      CANONICALIZECHECK:   bufferization.dealloc_tensor
+//      CANONICALIZECHECK:   arith.addi
+//      CANONICALIZECHECK:   linalg.yield
+//      CANONICALIZECHECK:   return %[[GENERIC]]

``````````

</details>


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