[mlir] Fix liveness analysis for block arguments (#88848)

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.

---------

Signed-off-by: ikulagin <[email protected]>
Co-authored-by: ikulagin <[email protected]>

Author: ikulagin

mlir/lib/Analysis/Liveness.cpp

@@ -72,6 +72,10 @@ struct BlockInfoBuilder {
         defValues.insert(result);
       for (Value operand : op->getOperands())
         useValues.insert(operand);
+      for (Region &region : op->getRegions())
+        for (Block &child : region.getBlocks())
+          for (BlockArgument arg : child.getArguments())
+            defValues.insert(arg);
     });
     llvm::set_subtract(useValues, defValues);
   }

mlir/test/Analysis/test-liveness.mlir

@@ -493,3 +493,27 @@ 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: {{^// +}}[[VAL3:[a-z0-9_]+]]{{ *:}}
+  // CHECK: {{^// +}}[[VAL4:[a-z0-9_]+]]{{ *:}}
+  %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: Block: 1
+    // CHECK-NEXT: LiveIn: [[VAL4]]{{ *$}}
+    // CHECK-NEXT: LiveOut:{{ *$}}
+    %1 = arith.addi %arg4, %c1_i32 : i32
+    scf.yield %1 : i32
+  }
+  return
+}