[Mlir-commits] [mlir] 568110f - [MLIR][Linalg] Fix segfault when an out operand is yielded by a nested op (#216977)

[llvmlistbot at llvm.org]

Author: Chibuoyim (Wilson) Ogbonna
Date: 2026-08-24T10:38:57+01:00
New Revision: 568110f3a8a2baa22f78e552e17d724fcc5c5286

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

LOG: [MLIR][Linalg] Fix segfault when an out operand is yielded by a nested op (#216977)

This closes
[203752](https://github.com/llvm/llvm-project/issues/203752), a smaller
reproducer is:
```mlir
#map = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1) -> (d0)>
#map2 = affine_map<(d0) -> (d0)>

func.func @yield_of_nested_op_keeps_output(%arg0: tensor<1x1xi32>, %arg1: tensor<1xi32>,
                                           %arg2: memref<1xi1>) {
  %c42 = arith.constant 42 : i32
  %c0 = arith.constant 0 : index
  %false = arith.constant false
  %0 = linalg.generic {indexing_maps = [#map, #map1],
                       iterator_types = ["parallel", "reduction"]}
      ins(%arg0 : tensor<1x1xi32>) outs(%arg1 : tensor<1xi32>) {
  ^bb0(%in: i32, %out: i32):
    %1 = linalg.generic {indexing_maps = [#map2], iterator_types = ["parallel"]}
        outs(%arg1 : tensor<1xi32>) {
    ^bb0(%out_0: i32):
      memref.store %false, %arg2[%c0] : memref<1xi1>
      linalg.yield %out : i32
    } -> tensor<1xi32>
    linalg.yield %c42 : i32
  } -> tensor<1xi32>
  return
}
```
Segfaults when run with `mlir-opt --linalg-fuse-elementwise-ops`

Currently, `isResultValueDead` drops the `out` operand because it only
checks that the user of `%out` is a `linalg.yield`, not that it is this
op's own terminator. But here `%out`'s only use is the *nested*
generic's yield, so the operand is dropped while the value is still live
inside that region!

Simple fix is checking that the yield is the terminator of the generic
being rewritten before marking the result dead.

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/Transforms/EraseUnusedOperandsAndResults.cpp
    mlir/test/Dialect/Linalg/erase-unused-operands-and-results.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Transforms/EraseUnusedOperandsAndResults.cpp b/mlir/lib/Dialect/Linalg/Transforms/EraseUnusedOperandsAndResults.cpp
index cbd63576619b6..6800b3042d524 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/EraseUnusedOperandsAndResults.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/EraseUnusedOperandsAndResults.cpp
@@ -40,9 +40,11 @@ static bool isResultValueDead(linalg::GenericOp genericOp, OpResult result) {
   if (!argUserOp->use_empty())
     return false;
 
-  // Check that argUser is a yield.
+  // Check that argUser is this op's own terminator. A nested op's
+  // `linalg.yield` also matches, but leaves the argument live inside that
+  // region.
   auto yieldOp = dyn_cast<linalg::YieldOp>(argUserOp);
-  if (!yieldOp)
+  if (!yieldOp || yieldOp != genericOp.getBody()->getTerminator())
     return false;
 
   // Check outArg data is not being used by other outArgs.

diff  --git a/mlir/test/Dialect/Linalg/erase-unused-operands-and-results.mlir b/mlir/test/Dialect/Linalg/erase-unused-operands-and-results.mlir
index dea3c22be0015..ff2bb0e1d24b7 100644
--- a/mlir/test/Dialect/Linalg/erase-unused-operands-and-results.mlir
+++ b/mlir/test/Dialect/Linalg/erase-unused-operands-and-results.mlir
@@ -520,3 +520,45 @@ func.func @remove_unnecessary_input(%a: tensor<?xf32>, %b: tensor<?xf32>)
   } -> tensor<?xf32>
   return %0 : tensor<?xf32>
 }
+
+// -----
+
+// Keep an out operand whose block argument is yielded by a nested op.
+
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (d0)>
+#map2 = affine_map<(d0) -> (d0)>
+func.func @keep_result_used_by_nested_op(%arg0 : tensor<1x1xi32>, %arg1 : tensor<1xi32>,
+    %arg2 : memref<1xi1>) {
+  %c0 = arith.constant 0 : index
+  %false = arith.constant false
+  %cst = arith.constant 42 : i32
+  %0 = linalg.generic {
+      indexing_maps = [#map0, #map1],
+      iterator_types = ["parallel", "reduction"]}
+      ins(%arg0 : tensor<1x1xi32>)
+      outs(%arg1 : tensor<1xi32>) {
+    ^bb0(%b0 : i32, %b1 : i32) :
+      %1 = linalg.generic {
+          indexing_maps = [#map2],
+          iterator_types = ["parallel"]}
+          outs(%arg1 : tensor<1xi32>) {
+        ^bb0(%b2 : i32) :
+          memref.store %false, %arg2[%c0] : memref<1xi1>
+          linalg.yield %b1 : i32
+        } -> tensor<1xi32>
+      linalg.yield %cst : i32
+    } -> tensor<1xi32>
+  return
+}
+
+//  CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+//  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d0)>
+//      CHECK: func @keep_result_used_by_nested_op(
+// CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: tensor<1x1xi32>
+// CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: tensor<1xi32>
+//      CHECK:   linalg.generic
+// CHECK-SAME:       indexing_maps = [#[[MAP0]], #[[MAP1]]]
+// CHECK-SAME:       outs(%[[ARG1]] :
+// CHECK-NEXT:   ^bb0(%{{.+}}: i32, %[[B1:.+]]: i32)
+//      CHECK:     linalg.yield %[[B1]]