[flang-commits] [flang] [flang] Add RegionBranchOpInterface to fir.do_loop (PR #202418)
Susan Tan ス-ザン タン via flang-commits
flang-commits at lists.llvm.org
Tue Jun 9 08:27:48 PDT 2026
https://github.com/SusanTan updated https://github.com/llvm/llvm-project/pull/202418
>From cc10439817bcb01d69ed23e2b914059d1e4f6570 Mon Sep 17 00:00:00 2001
From: Susan Tan <zujunt at nvidia.com>
Date: Mon, 8 Jun 2026 12:20:49 -0700
Subject: [PATCH 1/2] add RegionBranchOpInterface to fir.do_loop
---
.../include/flang/Optimizer/Dialect/FIROps.td | 8 +-
flang/lib/Optimizer/Dialect/FIROps.cpp | 59 ++++++++++++++
.../alias-analysis-regionbranch.mlir | 78 ++++++++++++++++++-
3 files changed, 140 insertions(+), 5 deletions(-)
diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index ec42855761bcb..83ec441febf44 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -2337,7 +2337,9 @@ def fir_LenParamIndexOp : fir_OneResultOp<"len_param_index", [NoMemoryEffect]> {
def fir_ResultOp
: fir_Op<"result", [Pure, ReturnLike, Terminator,
- ParentOneOf<["IfOp", "DoLoopOp", "IterWhileOp"]>]> {
+ ParentOneOf<["IfOp", "DoLoopOp", "IterWhileOp"]>,
+ DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
+ ["getMutableSuccessorOperands"]>]> {
let summary = "special terminator for use in fir region operations";
let description = [{
@@ -2367,7 +2369,9 @@ class region_Op<string mnemonic, list<Trait> traits = []> :
def fir_DoLoopOp : region_Op<"do_loop", [AttrSizedOperandSegments,
DeclareOpInterfaceMethods<LoopLikeOpInterface,
- ["getYieldedValuesMutable"]>]> {
+ ["getYieldedValuesMutable"]>,
+ DeclareOpInterfaceMethods<RegionBranchOpInterface,
+ ["getEntrySuccessorOperands", "getSuccessorInputs"]>]> {
let summary = "generalized loop operation";
let description = [{
Generalized high-level looping construct. This operation is similar to
diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index c3c6397c16fca..e57288f392b1d 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -3973,6 +3973,65 @@ fir::DoLoopOp::getYieldedValuesMutable() {
: term->getOpOperands();
}
+void fir::DoLoopOp::getSuccessorRegions(
+ mlir::RegionBranchPoint point,
+ llvm::SmallVectorImpl<mlir::RegionSuccessor> ®ions) {
+ // Entry (parent → body): loop executes ≥1 iterations, or is skipped (0-trip).
+ // Loop-back (body → body): another iteration follows.
+ regions.push_back(mlir::RegionSuccessor(&getRegion()));
+ // Exit (body → parent): last iteration completes.
+ regions.push_back(mlir::RegionSuccessor::parent());
+}
+
+// Note: when finalValue is set, result[0] tracks the IV's final value.
+mlir::OperandRange
+fir::DoLoopOp::getEntrySuccessorOperands(mlir::RegionSuccessor successor) {
+ // Initial iter-arg values, forwarded into the body or (zero-trip) directly
+ // to the parent's iter results.
+ return getInitArgs();
+}
+
+mlir::ValueRange
+fir::DoLoopOp::getSuccessorInputs(mlir::RegionSuccessor successor) {
+ // Returns the receiving slots for values flowing into `successor`.
+ //
+ // For a loop with finalValue and one iter arg:
+ // %r:2 = fir.do_loop ... iter_args(%acc = %init) -> (index, i32)
+ //
+ // body successor → receiving slots = [%acc] (regionIterArgs)
+ // parent successor → receiving slots = [%r#1] (results, drop %r#0)
+ //
+ // %r#0 (finalValue IV) is excluded: it has no symmetric iter-arg slot in
+ // the body, so it cannot participate in the 4-edge count check.
+ if (successor.isParent())
+ return getResults().drop_front(getFinalValue() ? 1 : 0);
+ return getRegionIterArgs();
+}
+
+mlir::MutableOperandRange
+fir::ResultOp::getMutableSuccessorOperands(mlir::RegionSuccessor successor) {
+ // called on fir.result to determine
+ // which of the operands is sending to the specific successor
+ //
+ // Without finalValue (N iter args, N results):
+ // fir.result %a0, %a1 : T0, T1 → send all N ops to body/%acc or results
+ //
+ // With finalValue (N iter args, N+1 results):
+ // fir.result %iv_next, %a0 : index, T0
+ // ops[0]: IV final value — only meaningful on the last exit,
+ // has no iter-arg slot in the body; excluded here.
+ // ops[1..]: iter-carried values — sent on both loop-back and
+ // exit.
+ //
+ // For fir.if / fir.iter_while the parent cast fails; all ops are returned.
+ if (auto doLoop =
+ mlir::dyn_cast<fir::DoLoopOp>(getOperation()->getParentOp()))
+ if (doLoop.getFinalValue() && getNumOperands() > 0)
+ return mlir::MutableOperandRange(getOperation(), /*start=*/1,
+ getNumOperands() - 1);
+ return mlir::MutableOperandRange(getOperation());
+}
+
//===----------------------------------------------------------------------===//
// DTEntryOp
//===----------------------------------------------------------------------===//
diff --git a/flang/test/Analysis/AliasAnalysis/alias-analysis-regionbranch.mlir b/flang/test/Analysis/AliasAnalysis/alias-analysis-regionbranch.mlir
index ea0dd5d99570f..f721e6a853439 100644
--- a/flang/test/Analysis/AliasAnalysis/alias-analysis-regionbranch.mlir
+++ b/flang/test/Analysis/AliasAnalysis/alias-analysis-regionbranch.mlir
@@ -263,11 +263,11 @@ func.func @test_rb_optional_ref_outside_target(%arg0: !fir.ref<i32> {fir.bindc_n
// -----
-// Add loop test whose result is based on BlockArguments to ensure
-// no cycle.
+// Case: no finalValue, with iter args.
+// The iter arg is passed through unchanged; the result MustAliases the init.
// CHECK-LABEL: Testing : "test_rb_do_loop_iter_carry_ref"
// CHECK-DAG: carry_init#0 <-> outside_loop#0: NoAlias
-// CHECK-DAG: carry_init#0 <-> loop_join_ref#0: MayAlias
+// CHECK-DAG: carry_init#0 <-> loop_join_ref#0: MustAlias
// CHECK-DAG: outside_loop#0 <-> loop_join_ref#0: MayAlias
func.func @test_rb_do_loop_iter_carry_ref() {
@@ -284,3 +284,75 @@ func.func @test_rb_do_loop_iter_carry_ref() {
%join = fir.convert %loop_res {test.ptr = "loop_join_ref"} : (!fir.ref<f32>) -> !fir.ref<f32>
return
}
+
+// -----
+
+// Case: no finalValue, no iter args.
+// No value flow through the interface; two unrelated allocas remain NoAlias.
+// CHECK-LABEL: Testing : "test_rb_do_loop_no_iter_args"
+// CHECK-DAG: no_iter_a#0 <-> no_iter_b#0: NoAlias
+
+func.func @test_rb_do_loop_no_iter_args() {
+ %lb = arith.constant 1 : index
+ %ub = arith.constant 2 : index
+ %st = arith.constant 1 : index
+ %a = fir.alloca f32 {uniq_name = "_QFEa"}
+ %b = fir.alloca f32 {uniq_name = "_QFEb"}
+ %da = fir.declare %a {uniq_name = "_QFEa", test.ptr = "no_iter_a"} : (!fir.ref<f32>) -> !fir.ref<f32>
+ %db = fir.declare %b {uniq_name = "_QFEb", test.ptr = "no_iter_b"} : (!fir.ref<f32>) -> !fir.ref<f32>
+ fir.do_loop %iv = %lb to %ub step %st {
+ fir.store %iv to %da : !fir.ref<index>
+ fir.result
+ }
+ return
+}
+
+// -----
+
+// Case: finalValue only, no iter args.
+// The loop result is index (the IV's final value), not a reference.
+// Unrelated allocas inside and outside the loop remain NoAlias.
+// CHECK-LABEL: Testing : "test_rb_do_loop_final_value_only"
+// CHECK-DAG: fv_only_a#0 <-> fv_only_b#0: NoAlias
+
+func.func @test_rb_do_loop_final_value_only() {
+ %lb = arith.constant 1 : index
+ %ub = arith.constant 2 : index
+ %st = arith.constant 1 : index
+ %a = fir.alloca f32 {uniq_name = "_QFEfv_a"}
+ %b = fir.alloca f32 {uniq_name = "_QFEfv_b"}
+ %da = fir.declare %a {uniq_name = "_QFEfv_a", test.ptr = "fv_only_a"} : (!fir.ref<f32>) -> !fir.ref<f32>
+ %db = fir.declare %b {uniq_name = "_QFEfv_b", test.ptr = "fv_only_b"} : (!fir.ref<f32>) -> !fir.ref<f32>
+ %iv_final:1 = fir.do_loop %iv = %lb to %ub step %st -> index {
+ fir.result %iv : index
+ }
+ return
+}
+
+// -----
+
+// Case: finalValue with iter args.
+// result#0 is the IV's final value (index); result#1 is the iter result.
+// The iter arg is passed through unchanged; result#1 MustAliases the init.
+// CHECK-LABEL: Testing : "test_rb_do_loop_final_value_and_iter_carry_ref"
+// CHECK-DAG: fv_carry_init#0 <-> fv_outside_loop#0: NoAlias
+// CHECK-DAG: fv_carry_init#0 <-> fv_loop_join_ref#0: MustAlias
+// CHECK-DAG: fv_outside_loop#0 <-> fv_loop_join_ref#0: MayAlias
+
+func.func @test_rb_do_loop_final_value_and_iter_carry_ref() {
+ %lb = arith.constant 1 : index
+ %ub = arith.constant 2 : index
+ %st = arith.constant 1 : index
+ %a_carry = fir.alloca f32 {uniq_name = "_QFEfv_carry"}
+ %d_carry = fir.declare %a_carry {uniq_name = "_QFEfv_carry", test.ptr = "fv_carry_init"} : (!fir.ref<f32>) -> !fir.ref<f32>
+ %a_out = fir.alloca f32 {uniq_name = "_QFEfv_outside"}
+ %d_out = fir.declare %a_out {uniq_name = "_QFEfv_outside", test.ptr = "fv_outside_loop"} : (!fir.ref<f32>) -> !fir.ref<f32>
+ // -> (index, !fir.ref<f32>): result#0 = IV final value, result#1 = iter result
+ %loop_res:2 = fir.do_loop %iv = %lb to %ub step %st
+ iter_args(%carry = %d_carry) -> (index, !fir.ref<f32>) {
+ %next_iv = arith.addi %iv, %st : index
+ fir.result %next_iv, %carry : index, !fir.ref<f32>
+ }
+ %join = fir.convert %loop_res#1 {test.ptr = "fv_loop_join_ref"} : (!fir.ref<f32>) -> !fir.ref<f32>
+ return
+}
>From ae8421a81f60e051da9874d5c6d29b269d67f3ae Mon Sep 17 00:00:00 2001
From: Susan Tan <zujunt at nvidia.com>
Date: Mon, 8 Jun 2026 12:20:49 -0700
Subject: [PATCH 2/2] add RegionBranchOpInterface to fir.do_loop
---
.../include/flang/Optimizer/Dialect/FIROps.td | 8 +-
flang/lib/Optimizer/Dialect/FIROps.cpp | 59 ++++++++++++++
.../alias-analysis-regionbranch.mlir | 79 ++++++++++++++++++-
3 files changed, 141 insertions(+), 5 deletions(-)
diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index ec42855761bcb..83ec441febf44 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -2337,7 +2337,9 @@ def fir_LenParamIndexOp : fir_OneResultOp<"len_param_index", [NoMemoryEffect]> {
def fir_ResultOp
: fir_Op<"result", [Pure, ReturnLike, Terminator,
- ParentOneOf<["IfOp", "DoLoopOp", "IterWhileOp"]>]> {
+ ParentOneOf<["IfOp", "DoLoopOp", "IterWhileOp"]>,
+ DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
+ ["getMutableSuccessorOperands"]>]> {
let summary = "special terminator for use in fir region operations";
let description = [{
@@ -2367,7 +2369,9 @@ class region_Op<string mnemonic, list<Trait> traits = []> :
def fir_DoLoopOp : region_Op<"do_loop", [AttrSizedOperandSegments,
DeclareOpInterfaceMethods<LoopLikeOpInterface,
- ["getYieldedValuesMutable"]>]> {
+ ["getYieldedValuesMutable"]>,
+ DeclareOpInterfaceMethods<RegionBranchOpInterface,
+ ["getEntrySuccessorOperands", "getSuccessorInputs"]>]> {
let summary = "generalized loop operation";
let description = [{
Generalized high-level looping construct. This operation is similar to
diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index c3c6397c16fca..e57288f392b1d 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -3973,6 +3973,65 @@ fir::DoLoopOp::getYieldedValuesMutable() {
: term->getOpOperands();
}
+void fir::DoLoopOp::getSuccessorRegions(
+ mlir::RegionBranchPoint point,
+ llvm::SmallVectorImpl<mlir::RegionSuccessor> ®ions) {
+ // Entry (parent → body): loop executes ≥1 iterations, or is skipped (0-trip).
+ // Loop-back (body → body): another iteration follows.
+ regions.push_back(mlir::RegionSuccessor(&getRegion()));
+ // Exit (body → parent): last iteration completes.
+ regions.push_back(mlir::RegionSuccessor::parent());
+}
+
+// Note: when finalValue is set, result[0] tracks the IV's final value.
+mlir::OperandRange
+fir::DoLoopOp::getEntrySuccessorOperands(mlir::RegionSuccessor successor) {
+ // Initial iter-arg values, forwarded into the body or (zero-trip) directly
+ // to the parent's iter results.
+ return getInitArgs();
+}
+
+mlir::ValueRange
+fir::DoLoopOp::getSuccessorInputs(mlir::RegionSuccessor successor) {
+ // Returns the receiving slots for values flowing into `successor`.
+ //
+ // For a loop with finalValue and one iter arg:
+ // %r:2 = fir.do_loop ... iter_args(%acc = %init) -> (index, i32)
+ //
+ // body successor → receiving slots = [%acc] (regionIterArgs)
+ // parent successor → receiving slots = [%r#1] (results, drop %r#0)
+ //
+ // %r#0 (finalValue IV) is excluded: it has no symmetric iter-arg slot in
+ // the body, so it cannot participate in the 4-edge count check.
+ if (successor.isParent())
+ return getResults().drop_front(getFinalValue() ? 1 : 0);
+ return getRegionIterArgs();
+}
+
+mlir::MutableOperandRange
+fir::ResultOp::getMutableSuccessorOperands(mlir::RegionSuccessor successor) {
+ // called on fir.result to determine
+ // which of the operands is sending to the specific successor
+ //
+ // Without finalValue (N iter args, N results):
+ // fir.result %a0, %a1 : T0, T1 → send all N ops to body/%acc or results
+ //
+ // With finalValue (N iter args, N+1 results):
+ // fir.result %iv_next, %a0 : index, T0
+ // ops[0]: IV final value — only meaningful on the last exit,
+ // has no iter-arg slot in the body; excluded here.
+ // ops[1..]: iter-carried values — sent on both loop-back and
+ // exit.
+ //
+ // For fir.if / fir.iter_while the parent cast fails; all ops are returned.
+ if (auto doLoop =
+ mlir::dyn_cast<fir::DoLoopOp>(getOperation()->getParentOp()))
+ if (doLoop.getFinalValue() && getNumOperands() > 0)
+ return mlir::MutableOperandRange(getOperation(), /*start=*/1,
+ getNumOperands() - 1);
+ return mlir::MutableOperandRange(getOperation());
+}
+
//===----------------------------------------------------------------------===//
// DTEntryOp
//===----------------------------------------------------------------------===//
diff --git a/flang/test/Analysis/AliasAnalysis/alias-analysis-regionbranch.mlir b/flang/test/Analysis/AliasAnalysis/alias-analysis-regionbranch.mlir
index ea0dd5d99570f..cc2374471186e 100644
--- a/flang/test/Analysis/AliasAnalysis/alias-analysis-regionbranch.mlir
+++ b/flang/test/Analysis/AliasAnalysis/alias-analysis-regionbranch.mlir
@@ -263,11 +263,11 @@ func.func @test_rb_optional_ref_outside_target(%arg0: !fir.ref<i32> {fir.bindc_n
// -----
-// Add loop test whose result is based on BlockArguments to ensure
-// no cycle.
+// Case: no finalValue, with iter args.
+// The iter arg is passed through unchanged; the result MustAliases the init.
// CHECK-LABEL: Testing : "test_rb_do_loop_iter_carry_ref"
// CHECK-DAG: carry_init#0 <-> outside_loop#0: NoAlias
-// CHECK-DAG: carry_init#0 <-> loop_join_ref#0: MayAlias
+// CHECK-DAG: carry_init#0 <-> loop_join_ref#0: MustAlias
// CHECK-DAG: outside_loop#0 <-> loop_join_ref#0: MayAlias
func.func @test_rb_do_loop_iter_carry_ref() {
@@ -284,3 +284,76 @@ func.func @test_rb_do_loop_iter_carry_ref() {
%join = fir.convert %loop_res {test.ptr = "loop_join_ref"} : (!fir.ref<f32>) -> !fir.ref<f32>
return
}
+
+// -----
+
+// Case: no finalValue, no iter args.
+// No value flow through the interface; two unrelated allocas remain NoAlias.
+// CHECK-LABEL: Testing : "test_rb_do_loop_no_iter_args"
+// CHECK-DAG: no_iter_a#0 <-> no_iter_b#0: NoAlias
+
+func.func @test_rb_do_loop_no_iter_args() {
+ %lb = arith.constant 1 : index
+ %ub = arith.constant 2 : index
+ %st = arith.constant 1 : index
+ %a = fir.alloca f32 {uniq_name = "_QFEa"}
+ %b = fir.alloca f32 {uniq_name = "_QFEb"}
+ %da = fir.declare %a {uniq_name = "_QFEa", test.ptr = "no_iter_a"} : (!fir.ref<f32>) -> !fir.ref<f32>
+ %db = fir.declare %b {uniq_name = "_QFEb", test.ptr = "no_iter_b"} : (!fir.ref<f32>) -> !fir.ref<f32>
+ fir.do_loop %iv = %lb to %ub step %st {
+ %fval = fir.convert %iv : (index) -> f32
+ fir.store %fval to %da : !fir.ref<f32>
+ fir.result
+ }
+ return
+}
+
+// -----
+
+// Case: finalValue only, no iter args.
+// The loop result is index (the IV's final value), not a reference.
+// Unrelated allocas inside and outside the loop remain NoAlias.
+// CHECK-LABEL: Testing : "test_rb_do_loop_final_value_only"
+// CHECK-DAG: fv_only_a#0 <-> fv_only_b#0: NoAlias
+
+func.func @test_rb_do_loop_final_value_only() {
+ %lb = arith.constant 1 : index
+ %ub = arith.constant 2 : index
+ %st = arith.constant 1 : index
+ %a = fir.alloca f32 {uniq_name = "_QFEfv_a"}
+ %b = fir.alloca f32 {uniq_name = "_QFEfv_b"}
+ %da = fir.declare %a {uniq_name = "_QFEfv_a", test.ptr = "fv_only_a"} : (!fir.ref<f32>) -> !fir.ref<f32>
+ %db = fir.declare %b {uniq_name = "_QFEfv_b", test.ptr = "fv_only_b"} : (!fir.ref<f32>) -> !fir.ref<f32>
+ %iv_final:1 = fir.do_loop %iv = %lb to %ub step %st -> index {
+ fir.result %iv : index
+ }
+ return
+}
+
+// -----
+
+// Case: finalValue with iter args.
+// result#0 is the IV's final value (index); result#1 is the iter result.
+// The iter arg is passed through unchanged; result#1 MustAliases the init.
+// CHECK-LABEL: Testing : "test_rb_do_loop_final_value_and_iter_carry_ref"
+// CHECK-DAG: fv_carry_init#0 <-> fv_outside_loop#0: NoAlias
+// CHECK-DAG: fv_carry_init#0 <-> fv_loop_join_ref#0: MustAlias
+// CHECK-DAG: fv_outside_loop#0 <-> fv_loop_join_ref#0: MayAlias
+
+func.func @test_rb_do_loop_final_value_and_iter_carry_ref() {
+ %lb = arith.constant 1 : index
+ %ub = arith.constant 2 : index
+ %st = arith.constant 1 : index
+ %a_carry = fir.alloca f32 {uniq_name = "_QFEfv_carry"}
+ %d_carry = fir.declare %a_carry {uniq_name = "_QFEfv_carry", test.ptr = "fv_carry_init"} : (!fir.ref<f32>) -> !fir.ref<f32>
+ %a_out = fir.alloca f32 {uniq_name = "_QFEfv_outside"}
+ %d_out = fir.declare %a_out {uniq_name = "_QFEfv_outside", test.ptr = "fv_outside_loop"} : (!fir.ref<f32>) -> !fir.ref<f32>
+ // -> (index, !fir.ref<f32>): result#0 = IV final value, result#1 = iter result
+ %loop_res:2 = fir.do_loop %iv = %lb to %ub step %st
+ iter_args(%carry = %d_carry) -> (index, !fir.ref<f32>) {
+ %next_iv = arith.addi %iv, %st : index
+ fir.result %next_iv, %carry : index, !fir.ref<f32>
+ }
+ %join = fir.convert %loop_res#1 {test.ptr = "fv_loop_join_ref"} : (!fir.ref<f32>) -> !fir.ref<f32>
+ return
+}
More information about the flang-commits
mailing list