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

[Chibuoyim Ogbonna]

https://github.com/bruteforceboy created https://github.com/llvm/llvm-project/pull/216977

This closes [203752](https://github.com/llvm/llvm-project/issues/203752), a smaller reproducer is:
```
#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.

>From 32d88ef24dfba3824c52e59f5e8995b600889098 Mon Sep 17 00:00:00 2001
From: bruteforceboy <chibuoyim.faith.ogbonna at huawei.com>
Date: Tue, 18 Aug 2026 10:56:38 +0100
Subject: [PATCH] [MLIR][Linalg] Fix segfault when an out operand is yielded by
 a nested op

---
 .../EraseUnusedOperandsAndResults.cpp         |  6 ++-
 .../erase-unused-operands-and-results.mlir    | 42 +++++++++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)

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