[Mlir-commits] [mlir] [mlir][interface] Add getSuccessorForwardOperands to BranchOpInterface (PR #190735)

[lonely eagle]

https://github.com/linuxlonelyeagle created https://github.com/llvm/llvm-project/pull/190735

This PR fixes a limitation where I could get the BranchOp edges but was unable to access the corresponding SuccessorForwardOperands. It adds getSuccessorForwardOperands to BranchOpInterface. It allows passing a BranchOp edge (Block*) as an argument and returns the corresponding SuccessorForwardOperands. Part of https://github.com/llvm/llvm-project/pull/189253.

>From 0dd938c326ee4605fa554cbce124854b3e79e17d Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Tue, 7 Apr 2026 05:47:28 +0000
Subject: [PATCH] add getSuccessorForwardOperands to BranchOpInterface.

---
 .../Dialect/ControlFlow/IR/ControlFlowOps.td  | 10 ++++----
 .../mlir/Interfaces/ControlFlowInterfaces.td  |  9 +++++++
 .../Dialect/ControlFlow/IR/ControlFlowOps.cpp | 24 +++++++++++++++++++
 3 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td b/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td
index a441fd82546e3..fab6c1e754ed5 100644
--- a/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td
+++ b/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td
@@ -65,7 +65,8 @@ def AssertOp : CF_Op<"assert",
 //===----------------------------------------------------------------------===//
 
 def BranchOp : CF_Op<"br", [
-    DeclareOpInterfaceMethods<BranchOpInterface, ["getSuccessorForOperands"]>,
+    DeclareOpInterfaceMethods<BranchOpInterface,
+    ["getSuccessorForOperands", "getSuccessorForwardOperands"]>,
     Pure, Terminator
   ]> {
   let summary = "Branch operation";
@@ -114,8 +115,8 @@ def BranchOp : CF_Op<"br", [
 
 def CondBranchOp
     : CF_Op<"cond_br", [AttrSizedOperandSegments,
-                        DeclareOpInterfaceMethods<
-                            BranchOpInterface, ["getSuccessorForOperands"]>,
+                        DeclareOpInterfaceMethods< BranchOpInterface,
+                        ["getSuccessorForOperands", "getSuccessorForwardOperands"]>,
                         WeightedBranchOpInterface, Pure, Terminator]> {
   let summary = "Conditional branch operation";
   let description = [{
@@ -241,7 +242,8 @@ def CondBranchOp
 
 def SwitchOp : CF_Op<"switch",
     [AttrSizedOperandSegments,
-     DeclareOpInterfaceMethods<BranchOpInterface, ["getSuccessorForOperands"]>,
+     DeclareOpInterfaceMethods<BranchOpInterface,
+     ["getSuccessorForOperands", "getSuccessorForwardOperands"]>,
      Pure, Terminator]> {
   let summary = "Switch operation";
   let description = [{
diff --git a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td
index 06fa724e05fab..d32be0c63acc7 100644
--- a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td
+++ b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td
@@ -98,6 +98,15 @@ def BranchOpInterface : OpInterface<"BranchOpInterface"> {
       (ins "::mlir::Type":$lhs, "::mlir::Type":$rhs), [{}],
        [{ return lhs == rhs; }]
     >,
+    InterfaceMethod<[{
+      This method is called to returns the operands of this operation that
+      are passed to the specified successor's block arguments. If the successor
+      is not valid for this operation, or no operands are forwarded, an empty
+      ValueRange is returned.
+      }],
+      "ValueRange", "getSuccessorForwardOperands",
+      (ins "Block *":$successor), [{}],[{ return {};}]
+    >,
   ];
 
   let verify = [{
diff --git a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
index 435c37bc95aac..f6eb0f05911b8 100644
--- a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
+++ b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
@@ -296,6 +296,12 @@ Block *BranchOp::getSuccessorForOperands(ArrayRef<Attribute>) {
   return getDest();
 }
 
+ValueRange BranchOp::getSuccessorForwardOperands(Block *successor) {
+  if (successor == getDest())
+    return getDestOperands();
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // CondBranchOp
 //===----------------------------------------------------------------------===//
@@ -583,6 +589,14 @@ Block *CondBranchOp::getSuccessorForOperands(ArrayRef<Attribute> operands) {
   return nullptr;
 }
 
+ValueRange CondBranchOp::getSuccessorForwardOperands(Block *successor) {
+  if (successor == getTrueDest())
+    return getTrueOperands();
+  else if (successor == getFalseDest())
+    return getFalseOperands();
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // SwitchOp
 //===----------------------------------------------------------------------===//
@@ -1034,6 +1048,16 @@ void SwitchOp::getCanonicalizationPatterns(RewritePatternSet &results,
       .add<SimplifyUniformBlockArguments>(context);
 }
 
+ValueRange SwitchOp::getSuccessorForwardOperands(Block *successor) {
+  if (successor == getDefaultDestination())
+    return getDefaultOperands();
+  SuccessorRange caseDests = getCaseDestinations();
+  auto it = llvm::find(caseDests, successor);
+  if (it == caseDests.end())
+    return {};
+  return getCaseOperands(std::distance(caseDests.begin(), it));
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//