[Mlir-commits] [mlir] [mlir] Fix liveness analysis (PR #88848)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Ivan Kulagin (ikulagin)

<details>
<summary>Changes</summary>

The current implementation does not take into account definitions created by arguments of nested blocks. This leads to an incorrect construction of the live-in set of an outer block. Arguments of nested blocks are added to the live-in set of an outer block.

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


2 Files Affected:

- (modified) mlir/lib/Analysis/Liveness.cpp (+12-5) 
- (modified) mlir/test/Analysis/test-liveness.mlir (+56) 


``````````diff
diff --git a/mlir/lib/Analysis/Liveness.cpp b/mlir/lib/Analysis/Liveness.cpp
index a8e0daeabf4061..50ac7129dab87f 100644
--- a/mlir/lib/Analysis/Liveness.cpp
+++ b/mlir/lib/Analysis/Liveness.cpp
@@ -67,11 +67,18 @@ struct BlockInfoBuilder {
     // Mark all nested operation results as defined, and nested operation
     // operands as used. All defined value will be removed from the used set
     // at the end.
-    block->walk([&](Operation *op) {
-      for (Value result : op->getResults())
-        defValues.insert(result);
-      for (Value operand : op->getOperands())
-        useValues.insert(operand);
+    block->walk([&](Block *nestedBlock) {
+      if (block != nestedBlock) {
+        for (BlockArgument arg : nestedBlock->getArguments()) {
+          defValues.insert(arg);
+        }
+      }
+      for (Operation &op : *nestedBlock) {
+        for (Value result : op.getResults())
+          defValues.insert(result);
+        for (Value operand : op.getOperands())
+          useValues.insert(operand);
+      }
     });
     llvm::set_subtract(useValues, defValues);
   }
diff --git a/mlir/test/Analysis/test-liveness.mlir b/mlir/test/Analysis/test-liveness.mlir
index 8ae3d09a6cd122..17b5e7d5ae8990 100644
--- a/mlir/test/Analysis/test-liveness.mlir
+++ b/mlir/test/Analysis/test-liveness.mlir
@@ -493,3 +493,59 @@ func.func @nested_region3(
   }
   return %1 : i32
 }
+
+// -----
+
+// CHECK-LABEL: Testing : nested_region4
+
+func.func @nested_region4(%arg0: index, %arg1: index, %arg2: index) {
+  // CHECK: Block: 0
+  // CHECK-NEXT: LiveIn:
+  // CHECK-NEXT: LiveOut:
+  // CHECK-NEXT: BeginLivenessIntervals
+  // CHECK-NEXT: val_3
+  // CHECK-NEXT: %c0_i32 = arith.constant 0
+  // CHECK-NEXT: %c1_i32 = arith.constant 1
+  // CHECK-NEXT: %0 = scf.for
+  // COM: Skipping the body of the scf.for...
+  // CHECK:      val_4
+  // CHECK-NEXT: %c1_i32 = arith.constant 1
+  // CHECK-NEXT: %0 = scf.for
+  // COM: Skipping the body of the scf.for...
+  // CHECK:      // %1 = arith.addi
+  // CHECK-NEXT: val_5
+  // CHECK-NEXT: %0 = scf.for
+  // COM: Skipping the body of the scf.for...
+  // CHECK:      EndLivenessIntervals
+  // CHECK-NEXT: BeginCurrentlyLive
+  // CHECK-NEXT: %c0_i32 = arith.constant 0
+  // CHECK-SAME: arg0 at 0 arg1 at 0 arg2 at 0 val_3
+  // CHECK-NEXT: %c1_i32 = arith.constant 1
+  // CHECK-SAME: arg0 at 0 arg1 at 0 arg2 at 0 val_3 val_4
+  // CHECK-NEXT: %0 = scf.for
+  // COM: Skipping the body of the scf.for...
+  // CHECK:      arg0 at 0 arg1 at 0 arg2 at 0 val_3 val_4 val_5
+  // CHECK-NEXT: EndCurrentlyLive
+  %c0_i32 = arith.constant 0 : i32
+  %c1_i32 = arith.constant 1 : i32
+
+  %0 = scf.for %arg3 = %arg0 to %arg1 step %arg2 iter_args(%arg4 = %c0_i32) -> (i32) {
+    // CHECK-NEXT: Block: 1
+    // CHECK-NEXT: LiveIn: val_4
+    // CHECK-NEXT: LiveOut:
+    // CHECK-NEXT: BeginLivenessIntervals
+    // CHECK-NEXT: val_8
+    // CHECK-NEXT: %1 = arith.addi
+    // CHECK-NEXT: scf.yield %1
+    // CHECK-NEXT: EndLivenessIntervals
+    // CHECK-NEXT: BeginCurrentlyLive
+    // CHECK-NEXT: %1 = arith.addi
+    // CHECK-SAME: val_4 arg0 at 1 arg1 at 1 val_8
+    // CHECK-NEXT: scf.yield %1
+    // CHECK-SAME: val_8
+    // CHECK-NEXT: EndCurrentlyLive
+    %1 = arith.addi %arg4, %c1_i32 : i32
+    scf.yield %1 : i32
+  }
+  return
+}

``````````

</details>


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