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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Chibuoyim (Wilson) Ogbonna (bruteforceboy)

<details>
<summary>Changes</summary>

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.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Linalg/Transforms/EraseUnusedOperandsAndResults.cpp (+4-2) 
- (modified) mlir/test/Dialect/Linalg/erase-unused-operands-and-results.mlir (+42) 


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

``````````

</details>


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