[Mlir-commits] [mlir] [mlir] [liveness] Conservatively mark operands of return-like op inside non-callable and non-regionbranch op as live (PR #140793)
Nhat Nguyen
llvmlistbot at llvm.org
Tue May 20 13:24:28 PDT 2025
https://github.com/nhat-nguyen created https://github.com/llvm/llvm-project/pull/140793
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.
>From b0032a933ad9c845b7a00fbd87fd0a0f8c211be5 Mon Sep 17 00:00:00 2001
From: Nhat Nguyen <hoangnhat2911 at gmail.com>
Date: Tue, 20 May 2025 16:00:37 -0400
Subject: [PATCH 1/4] Add
---
.../Analysis/DataFlow/LivenessAnalysis.cpp | 13 ++++---
mlir/test/Transforms/remove-dead-values.mlir | 37 +++++++++++++++++++
2 files changed, 45 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
index c12149a1a0242..4d581d4b1da09 100644
--- a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
+++ b/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());
}
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 21d53b0742e07..3a00083177151 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/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
>From 1a5666ed487ee7b616d722da50891a6a77937c40 Mon Sep 17 00:00:00 2001
From: Nhat Nguyen <hoangnhat2911 at gmail.com>
Date: Tue, 20 May 2025 16:06:15 -0400
Subject: [PATCH 2/4] Mark all yield values as live
---
mlir/test/Transforms/remove-dead-values.mlir | 7 +++++--
1 file changed, 5 insertions(+), 2 deletions(-)
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 3a00083177151..87df57df54d7a 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -486,14 +486,17 @@ gpu.module @test_module_3 {
// CHECK-LABEL: func @gpu_all_reduce_region()
// CHECK: %[[yield:.*]] = arith.xori %{{.*}}, %{{.*}} : i32
-// CHECK: "gpu.yield"(%[[yield]]) : (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>) {
+ %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
>From ab3164716377106413c305c85411130f1ba97b6d Mon Sep 17 00:00:00 2001
From: Nhat Nguyen <hoangnhat2911 at gmail.com>
Date: Tue, 20 May 2025 16:07:39 -0400
Subject: [PATCH 3/4] Mark all yield values as live
---
mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
index 4d581d4b1da09..a1091f85287bd 100644
--- a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
@@ -49,8 +49,8 @@ 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
-/// (2) is returned by a public function
+/// (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) 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
>From 22c265f38c7ccca81df570327663be2d5005ef6b Mon Sep 17 00:00:00 2001
From: Nhat Nguyen <hoangnhat2911 at gmail.com>
Date: Tue, 20 May 2025 16:10:53 -0400
Subject: [PATCH 4/4] Update comment
---
mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp | 5 +++--
1 file changed, 3 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
index a1091f85287bd..d61cdb143e7dd 100644
--- a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp
@@ -53,8 +53,9 @@ ChangeResult Liveness::meet(const AbstractSparseLattice &other) {
/// (2) is returned by a public function OR
/// (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.
+/// 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.
///
More information about the Mlir-commits
mailing list