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

[llvmlistbot at llvm.org]

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

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

>From 1312b506d8159dfdb02d6f3179f06a4c8d7ad424 Mon Sep 17 00:00:00 2001
From: cubehan3 <cubehan3 at gmail.com>
Date: Mon, 7 Sep 2026 11:27:15 +0000
Subject: [PATCH] [mlir][linalg] Bail out of decomposition when peeled stmt has
 no results

---
 .../Linalg/Transforms/DecomposeLinalgOps.cpp  |  7 ++++
 mlir/test/Dialect/Linalg/decompose-ops.mlir   | 36 +++++++++++++++++++
 2 files changed, 43 insertions(+)

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]]