[Mlir-commits] [mlir] [mlir] [liveness] Conservatively mark operands of return-like op inside non-callable and non-regionbranch op as live (PR #140793)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Nhat Nguyen (nhat-nguyen)

<details>
<summary>Changes</summary>

Currently the liveness analysis always marks operands yielded in regions that aren't classified as `RegionBranchOpInterface` or `CallableOpInterface` as non-live. Examples for these ops include linalg.generic (with `linalg.yield` as terminator) or gpu ops (with `gpu.yield` as terminator).

This in turn makes the `remove-dead-values` pass always incorrectly removes the bodies of these ops, leading to invalid IR. Because these ops define their own semantics, I have conservatively marked all operands of these yield ops to be live.

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


2 Files Affected:

- (modified) mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp (+7-3) 
- (modified) mlir/test/Transforms/remove-dead-values.mlir (+40) 


``````````diff
diff --git a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
index c12149a1a0242..d61cdb143e7dd 100644
--- a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
@@ -51,7 +51,11 @@ ChangeResult Liveness::meet(const AbstractSparseLattice &other) {
 /// A value is considered "live" iff it:
 ///   (1) has memory effects OR
 ///   (2) is returned by a public function OR
-///   (3) is used to compute a value of type (1) or (2).
+///   (3) is used to compute a value of type (1) or (2) OR
+///   (4) is returned by a return-like op whose parent isn't a callable
+///       nor a RegionBranchOpInterface (e.g.: linalg.yield, gpu.yield,...)
+///       These ops have their own semantics, so we conservatively mark the
+///       the yield value as live.
 /// It is also to be noted that a value could be of multiple types (1/2/3) at
 /// the same time.
 ///
@@ -73,8 +77,8 @@ ChangeResult Liveness::meet(const AbstractSparseLattice &other) {
 LogicalResult
 LivenessAnalysis::visitOperation(Operation *op, ArrayRef<Liveness *> operands,
                                  ArrayRef<const Liveness *> results) {
-  // This marks values of type (1.a) liveness as "live".
-  if (!isMemoryEffectFree(op)) {
+  // This marks values of type (1.a) and (4) liveness as "live".
+  if (!isMemoryEffectFree(op) || op->hasTrait<OpTrait::ReturnLike>()) {
     for (auto *operand : operands)
       propagateIfChanged(operand, operand->markLive());
   }
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 21d53b0742e07..87df57df54d7a 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -468,3 +468,43 @@ func.func private @no_block_func_declaration() -> ()
 
 // CHECK: llvm.func @no_block_external_func()
 llvm.func @no_block_external_func() attributes {sym_visibility = "private"}
+
+// -----
+
+// Check that yielded values aren't incorrectly removed in gpu regions
+gpu.module @test_module_3 {
+  gpu.func @gpu_all_reduce_region() {
+    %arg0 = arith.constant 1 : i32
+    %result = gpu.all_reduce %arg0 uniform {
+    ^bb(%lhs : i32, %rhs : i32):
+      %xor = arith.xori %lhs, %rhs : i32
+      "gpu.yield"(%xor) : (i32) -> ()
+    } : (i32) -> (i32)
+    gpu.return
+  }
+}
+
+// CHECK-LABEL: func @gpu_all_reduce_region()
+// CHECK: %[[yield:.*]] = arith.xori %{{.*}}, %{{.*}} : i32
+// CHECK: gpu.yield %[[yield]] : i32
+
+// -----
+
+// Check that yielded values aren't incorrectly removed in linalg regions
+module {
+  func.func @linalg_red_add(%arg0: tensor<?xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
+    %0 = linalg.generic {
+      indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (0)>],
+      iterator_types = ["reduction"]
+    } ins(%arg0 : tensor<?xf32>) outs(%arg1 : tensor<1xf32>) {
+    ^bb0(%in: f32, %out: f32):
+      %1 = arith.addf %in, %out : f32
+      linalg.yield %1 : f32
+    } -> tensor<1xf32>
+    return %0 : tensor<1xf32>
+  }
+}
+
+// CHECK-LABEL: func @linalg_red_add
+// CHECK: %[[yield:.*]] = arith.addf %{{.*}}, %{{.*}} : f32
+// CHECK: linalg.yield %[[yield]] : f32

``````````

</details>


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