[Mlir-commits] [mlir] [mlir][interface] Add getRegionNonForwardedValues API to RegionBranchOpInterface (PR #175212)

[lonely eagle]

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

>From 8150fbc468f7c4287178b8bba172a90733e41bd9 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Wed, 7 Jan 2026 15:41:34 +0000
Subject: [PATCH] add getRegionNonforwardedArguments API.

add getRegionNonForwardedValues API to RegionBranchOpInterface.
---
 .../mlir/Interfaces/ControlFlowInterfaces.td  |  8 ++++++
 mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp | 24 ++++++++----------
 mlir/lib/Interfaces/ControlFlowInterfaces.cpp | 25 +++++++++++++++++++
 3 files changed, 43 insertions(+), 14 deletions(-)

diff --git a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td
index 627eeb14b9023..af582a0c1e101 100644
--- a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td
+++ b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td
@@ -379,6 +379,14 @@ def RegionBranchOpInterface : OpInterface<"RegionBranchOpInterface"> {
     ::mlir::OperandRange getSuccessorOperands(
         ::mlir::RegionBranchPoint src, ::mlir::RegionSuccessor dest);
 
+    /// Return the non-forwarded values from the source branch point to the
+    /// destination region successor.
+    ///
+    /// If the `RegionSuccessor` is a region, it will return non-forwarded
+    /// arguments, if it is a parent, it will return non-forwarded results.
+    ::llvm::SmallVector<Value> getRegionNonForwardedValues(
+        ::mlir::RegionBranchPoint src, ::mlir::RegionSuccessor dest);
+
     /// Build a mapping from successor operands to successor input. Each
     /// successor operand could be forwarded to multiple successor inputs.
     /// Operands that are not forwarded are not added to the map. Unless a
diff --git a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
index 6515e42bb2081..4f420dddf5eb7 100644
--- a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
@@ -329,8 +329,9 @@ void AbstractSparseForwardDataFlowAnalysis::visitRegionSuccessors(
       }
     }
 
-    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));
   }
 }
 
@@ -606,7 +607,6 @@ void AbstractSparseBackwardDataFlowAnalysis::visitRegionSuccessors(
       unaccounted.reset(operand->getOperandNumber());
     }
   }
-
   Operation *op = branch.getOperation();
   SmallVector<RegionSuccessor> successors;
   SmallVector<Attribute> operands(op->getNumOperands(), nullptr);
@@ -614,17 +614,13 @@ void AbstractSparseBackwardDataFlowAnalysis::visitRegionSuccessors(
   for (RegionSuccessor &successor : successors) {
     if (successor.isParent())
       continue;
-    SmallVector<BlockArgument> noControlFlowArguments;
-    MutableArrayRef<BlockArgument> arguments =
-        successor.getSuccessor()->getArguments();
-    ValueRange inputs = branch.getSuccessorInputs(successor);
-    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);
-      }
-    }
+    auto ValueToArgument = [](Value value) {
+      return cast<BlockArgument>(value);
+    };
+    SmallVector<BlockArgument> noControlFlowArguments =
+        llvm::map_to_vector(branch.getRegionNonForwardedValues(
+                                RegionBranchPoint::parent(), successor),
+                            ValueToArgument);
     visitNonControlFlowArguments(successor, noControlFlowArguments);
   }
 
diff --git a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
index ebf78d8bd60ce..da5a17f533695 100644
--- a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
+++ b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
@@ -442,6 +442,31 @@ RegionBranchOpInterface::getSuccessorOperands(RegionBranchPoint src,
   return src.getTerminatorPredecessorOrNull().getSuccessorOperands(dest);
 }
 
+SmallVector<Value>
+RegionBranchOpInterface::getRegionNonForwardedValues(RegionBranchPoint src,
+                                                     RegionSuccessor dest) {
+  SmallVector<Value> nonForwardValues;
+  OperandRange successorOperands = getSuccessorOperands(src, dest);
+  ValueRange successorInputs = getSuccessorInputs(dest);
+  size_t firstIndex = 0;
+  if (!successorOperands.empty()) {
+    firstIndex =
+        dest.isParent()
+            ? cast<OpResult>(successorInputs.front()).getResultNumber()
+            : cast<BlockArgument>(successorInputs.front()).getArgNumber();
+  }
+  ValueRange values;
+  if (dest.isParent()) {
+    values = this->getOperation()->getResults();
+  } else {
+    values = dest.getSuccessor()->getArguments();
+  }
+  nonForwardValues.append(llvm::to_vector(values.take_front(firstIndex)));
+  nonForwardValues.append(llvm::to_vector(
+      values.drop_front(firstIndex + successorOperands.size())));
+  return nonForwardValues;
+}
+
 static MutableArrayRef<OpOperand> operandsToOpOperands(OperandRange &operands) {
   return MutableArrayRef<OpOperand>(operands.getBase(), operands.size());
 }