[Mlir-commits] [mlir] [mlir] add getRegionNonforwardedArguments API to RegionSuccessor (PR #175212)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: lonely eagle (linuxlonelyeagle)

<details>
<summary>Changes</summary>

The getSuccessorInputs function is used to access successor inputs, but the region's arguments include arguments that are not forwarded, so I added an API to access them.

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


2 Files Affected:

- (modified) mlir/include/mlir/Interfaces/ControlFlowInterfaces.h (+18) 
- (modified) mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp (+14-31) 


``````````diff
diff --git a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
index b76c2891fad5a..76aa20f11d84b 100644
--- a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
+++ b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
@@ -228,6 +228,24 @@ class RegionSuccessor {
   /// the current region.
   ValueRange getSuccessorInputs() const { return inputs; }
 
+  ValueRange getRegionNonforwardedArguments() const {
+    if (isParent())
+      return {};
+    SmallVector<Value> nonForwardArguments;
+    MutableArrayRef<BlockArgument> arguments = getSuccessor()->getArguments();
+    if (arguments.empty())
+      return {};
+    ValueRange inputs = getSuccessorInputs();
+    if (inputs.empty())
+      return arguments;
+    for (BlockArgument argument : arguments) {
+      if (!llvm::is_contained(inputs, argument)) {
+        nonForwardArguments.push_back(argument);
+      }
+    }
+    return nonForwardArguments;
+  }
+
   bool operator==(RegionSuccessor rhs) const {
     return successor == rhs.successor && inputs == rhs.inputs;
   }
diff --git a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
index b9d861830dd38..97889415da31a 100644
--- a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
@@ -307,29 +307,20 @@ void AbstractSparseForwardDataFlowAnalysis::visitRegionSuccessors(
            "expected the same number of successor inputs as operands");
 
     unsigned firstIndex = 0;
-    if (inputs.size() != lattices.size()) {
-      if (!point->isBlockStart()) {
-        if (!inputs.empty())
-          firstIndex = cast<OpResult>(inputs.front()).getResultNumber();
-        visitNonControlFlowArgumentsImpl(
-            branch,
-            RegionSuccessor(
-                branch, branch->getResults().slice(firstIndex, inputs.size())),
-            lattices, firstIndex);
-      } else {
-        if (!inputs.empty())
-          firstIndex = cast<BlockArgument>(inputs.front()).getArgNumber();
-        Region *region = point->getBlock()->getParent();
-        visitNonControlFlowArgumentsImpl(
-            branch,
-            RegionSuccessor(region, region->getArguments().slice(
-                                        firstIndex, inputs.size())),
-            lattices, firstIndex);
-      }
+    if (inputs.size() != lattices.size() && point->isBlockStart()) {
+      if (!inputs.empty())
+        firstIndex = cast<BlockArgument>(inputs.front()).getArgNumber();
+      Region *region = point->getBlock()->getParent();
+      visitNonControlFlowArgumentsImpl(
+          branch,
+          RegionSuccessor(
+              region, region->getArguments().slice(firstIndex, inputs.size())),
+          lattices, firstIndex);
     }
 
-    for (auto it : llvm::zip(*operands, lattices.drop_front(firstIndex)))
-      join(std::get<1>(it), *getLatticeElementFor(point, std::get<0>(it)));
+    for (auto [lattice, operand] :
+         llvm::zip(lattices.drop_front(firstIndex), *operands))
+      join(lattice, *getLatticeElementFor(point, operand));
   }
 }
 
@@ -612,16 +603,8 @@ void AbstractSparseBackwardDataFlowAnalysis::visitRegionSuccessors(
     if (successor.isParent())
       continue;
     SmallVector<BlockArgument> noControlFlowArguments;
-    MutableArrayRef<BlockArgument> arguments =
-        successor.getSuccessor()->getArguments();
-    ValueRange inputs = successor.getSuccessorInputs();
-    for (BlockArgument argument : arguments) {
-      // Visit blockArgument of RegionBranchOp which isn't "control
-      // flow block arguments". For example, the IV of a loop.
-      if (!llvm::is_contained(inputs, argument)) {
-        noControlFlowArguments.push_back(argument);
-      }
-    }
+    for (Value arg : successor.getRegionNonforwardedArguments())
+      noControlFlowArguments.push_back(cast<BlockArgument>(arg));
     visitNonControlFlowArguments(successor, noControlFlowArguments);
   }
 

``````````

</details>


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