Add

Author: nhat-nguyen

mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp

@@ -49,9 +49,12 @@ ChangeResult Liveness::meet(const AbstractSparseLattice &other) {
 /// For every value, liveness analysis determines whether or not it is "live".
 ///
 /// 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).
+///   (1) has memory effects
+///   (2) is returned by a public function
+///   (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
+///       (e.g.: linalg.yield, gpu.yield,...) These ops have their own
+///       semantics, so we conservatively mark the 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 +76,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());
   }

mlir/test/Transforms/remove-dead-values.mlir

@@ -468,3 +468,40 @@ 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 = [#map, #map1], 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