[Mlir-commits] [mlir] 60a108b - [mlir][interface] Add getNonSuccessorInputs API to RegionBranchOpInterface (#175212)

[llvmlistbot at llvm.org]

Author: lonely eagle
Date: 2026-01-24T00:25:12+08:00
New Revision: 60a108bcc8a47f6b4091a6291636fbe4dec9615f

URL: https://github.com/llvm/llvm-project/commit/60a108bcc8a47f6b4091a6291636fbe4dec9615f
DIFF: https://github.com/llvm/llvm-project/commit/60a108bcc8a47f6b4091a6291636fbe4dec9615f.diff

LOG: [mlir][interface] Add getNonSuccessorInputs API to RegionBranchOpInterface (#175212)

Add getNonSuccessorInputs API to RegionBranchOpInterface.It is used to
return the non-forwarded arguments or results.

Added: 
    

Modified: 
    mlir/include/mlir/Interfaces/ControlFlowInterfaces.td
    mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
    mlir/lib/Interfaces/ControlFlowInterfaces.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td
index 627eeb14b9023..8975b1235a7e3 100644
--- a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td
+++ b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td
@@ -379,6 +379,13 @@ def RegionBranchOpInterface : OpInterface<"RegionBranchOpInterface"> {
     ::mlir::OperandRange getSuccessorOperands(
         ::mlir::RegionBranchPoint src, ::mlir::RegionSuccessor dest);
 
+    /// Return all successor inputs for the given region successor.
+    ///
+    /// If the "successor" is a region, it will return non-forwarded arguments,
+    /// if it is a "parent", it will return non-forwarded results.
+    ::llvm::SmallVector<Value> getNonSuccessorInputs(
+        ::mlir::RegionSuccessor successor);
+
     /// 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..7dad9676e7e53 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,11 @@ 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.getNonSuccessorInputs(successor), valueToArgument);
     visitNonControlFlowArguments(successor, noControlFlowArguments);
   }
 

diff  --git a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
index cd0a1a8371aac..2f95531455b2b 100644
--- a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
+++ b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
@@ -443,6 +443,24 @@ RegionBranchOpInterface::getSuccessorOperands(RegionBranchPoint src,
   return src.getTerminatorPredecessorOrNull().getSuccessorOperands(dest);
 }
 
+SmallVector<Value>
+RegionBranchOpInterface::getNonSuccessorInputs(RegionSuccessor successor) {
+  SmallVector<Value> results = llvm::to_vector(
+      successor.isParent()
+          ? ValueRange(getOperation()->getResults())
+          : ValueRange(successor.getSuccessor()->getArguments()));
+  ValueRange successorInputs = getSuccessorInputs(successor);
+  if (!successorInputs.empty()) {
+    unsigned inputBegin =
+        successor.isParent()
+            ? cast<OpResult>(successorInputs.front()).getResultNumber()
+            : cast<BlockArgument>(successorInputs.front()).getArgNumber();
+    results.erase(results.begin() + inputBegin,
+                  results.begin() + inputBegin + successorInputs.size());
+  }
+  return results;
+}
+
 static MutableArrayRef<OpOperand> operandsToOpOperands(OperandRange &operands) {
   return MutableArrayRef<OpOperand>(operands.getBase(), operands.size());
 }