[flang-commits] [flang] [flang][cuda] Keep delayed descriptor allocs out of CFG loops (PR #232992)
Zhen Wang via flang-commits
flang-commits at lists.llvm.org
Fri Oct 16 18:04:51 PDT 2026
https://github.com/wangzpgi created https://github.com/llvm/llvm-project/pull/232992
CUFAllocDelay places a descriptor's cuf.alloc, its initial store and its fir.declare at the nearest common dominator of all their uses. When a DO loop lowers to branches, for example because it contains a RETURN, and the descriptor is used both inside the loop and after it, that block is the loop header. The group then runs on every trip, including the exit trip, and the code after the loop reads a fresh, empty descriptor:
```fortran
do jj = 1, 1
if (.false.) return
allocate(clvars(1))
end do
allocate(cvars(1))
cvars(:) = clvars(:) ! Assign: mismatching element counts (to 1, from 0)
```
After computing the common dominator, climb the dominator tree while the block is on a CFG cycle, tested with Block::isReachable on itself, so irreducible cycles are covered as well. The entry block has no predecessors, so the climb always terminates there at the latest. Structured loops were already safe because uses in nested regions resolve to the enclosing top-level op.
Tests cover a loop with uses in the body and after it, a nested loop, a loop reached through a non-loop successor, all uses inside the body, and an irreducible cycle.
>From 33e3e96fc519a9f0cb51dd5d24d3265b1e120a2b Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Fri, 16 Oct 2026 18:02:11 -0700
Subject: [PATCH] Keep delayed descriptor allocs out of CFG loops
---
.../flang/Optimizer/Transforms/Passes.td | 3 +-
.../Transforms/CUDA/CUFAllocDelay.cpp | 17 +-
flang/test/Transforms/CUF/cuf-alloc-delay.fir | 231 ++++++++++++++++++
3 files changed, 248 insertions(+), 3 deletions(-)
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index 07d3296a17be9..b5642cbdccede 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -589,7 +589,8 @@ def CUFAllocDelay : Pass<"cuf-alloc-delay", "::mlir::func::FuncOp"> {
let summary = "Delay cuf.alloc of descriptor types to just before first use";
let description = [{
Moves cuf.alloc operations for allocatable/pointer descriptors (box types)
- from function entry to just before their first use. This defers
+ from function entry to just before their first use, but never into a block
+ on a CFG cycle, where the alloc would run on every trip. This defers
cudaMallocManaged calls so that users can call cudaSetDevice before any
CUDA context is created.
}];
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp
index 5b3f7afd7f656..89073152bb6e4 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp
@@ -7,7 +7,8 @@
//===----------------------------------------------------------------------===//
//
// Delay cuf.alloc of descriptor (box) types from function entry to just before
-// their first use, possibly in a later block that dominates every use. This
+// their first use, possibly in a later block that dominates every use. The
+// target block is never on a CFG cycle, so the alloc runs at most once. This
// defers cudaMallocManaged calls so that users can call cudaSetDevice before
// any CUDA context is created.
//
@@ -48,7 +49,8 @@ static fir::CoordinateOp getHostAssocTupleSlot(fir::StoreOp storeOp,
/// Find the point before which the cuf.alloc group should be placed: the
/// earliest use in the block that dominates all uses, or that block's
/// terminator if it holds no use itself. Uses in nested regions resolve to
-/// their enclosing top-level op.
+/// their enclosing top-level op. The block is never on a CFG cycle, since an
+/// alloc there would run on every trip.
///
/// Host association stores go to \p hostAssocStores and sink with the group
/// instead of constraining it; the tuple's readers constrain it instead.
@@ -108,6 +110,17 @@ findDelayTarget(fir::DeclareOp declareOp, mlir::Block *entryBlock,
return nullptr;
}
+ // A block on a CFG cycle (e.g. the header of a DO loop that lowered to
+ // branches because it contains a RETURN) executes once per trip, and the
+ // alloc would replace the descriptor each time. Climb the dominator tree to
+ // the nearest block that is not on a cycle. Reachability covers irreducible
+ // cycles too, and the entry block has no predecessors, so it always ends.
+ while (common->isReachable(common)) {
+ mlir::DominanceInfoNode *idom = domInfo.getNode(common)->getIDom();
+ assert(idom && "entry block cannot be on a cycle");
+ common = idom->getBlock();
+ }
+
mlir::Operation *earliest = nullptr;
for (mlir::Operation *use : uses)
if (use->getBlock() == common &&
diff --git a/flang/test/Transforms/CUF/cuf-alloc-delay.fir b/flang/test/Transforms/CUF/cuf-alloc-delay.fir
index 793756067e498..a41479faf6344 100644
--- a/flang/test/Transforms/CUF/cuf-alloc-delay.fir
+++ b/flang/test/Transforms/CUF/cuf-alloc-delay.fir
@@ -621,3 +621,234 @@ func.func private @_QPsetup5()
// CHECK: fir.load
// CHECK: cuf.allocate
// CHECK: cuf.free
+
+// -----
+
+// Test 18: Uses in the body of a CFG loop and after it. Their nearest common
+// dominator is the loop header, which runs once per trip, so the group must
+// stay before the branch into the loop rather than sink into the header.
+func.func @_QPuse_in_cfg_loop() {
+ %0 = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> uniq_name("_QFuse_in_cfg_loopEa") bindc_name("a") data_attr(device) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %1 = fir.zero_bits !fir.heap<!fir.array<?xf32>>
+ %c0 = arith.constant 0 : index
+ %2 = fir.shape %c0 : (index) -> !fir.shape<1>
+ %3 = fir.embox %1(%2) allocator_idx(2) : (!fir.heap<!fir.array<?xf32>>, !fir.shape<1>) -> !fir.box<!fir.heap<!fir.array<?xf32>>>
+ fir.store %3 to %0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %4 = fir.declare %0 uniq_name("_QFuse_in_cfg_loopEa") fortran_attrs<allocatable> data_attr(#cuf.cuda<device>) : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %c1 = arith.constant 1 : index
+ cf.br ^bb1(%c0 : index)
+^bb1(%i: index):
+ %5 = arith.cmpi slt, %i, %c1 : index
+ cf.cond_br %5, ^bb2, ^bb3
+^bb2:
+ %6 = cuf.allocate %4 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> data_attr(device) -> i32
+ %7 = arith.addi %i, %c1 : index
+ cf.br ^bb1(%7 : index)
+^bb3:
+ %8 = fir.load %4 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ cuf.free %4 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> data_attr(device)
+ return
+}
+
+// CHECK-LABEL: func.func @_QPuse_in_cfg_loop
+// CHECK: fir.embox
+// CHECK: arith.constant 1
+// CHECK: cuf.alloc
+// CHECK: fir.store
+// CHECK: fir.declare
+// CHECK: cf.br ^bb1
+// CHECK: ^bb1
+// CHECK-NOT: cuf.alloc !fir.box
+// CHECK: cf.cond_br
+// CHECK: ^bb2:
+// CHECK: cuf.allocate
+// CHECK: ^bb3:
+// CHECK: fir.load
+// CHECK: cuf.free
+
+// -----
+
+// Test 19: Uses in the body and the exit of an inner CFG loop, both inside
+// an outer CFG loop. The nearest common dominator is the inner header, so the
+// climb has to pass the outer header as well and end at the entry block.
+func.func @_QPuse_in_nested_cfg_loop() {
+ %0 = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> uniq_name("_QFuse_in_nested_cfg_loopEa") bindc_name("a") data_attr(device) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %1 = fir.zero_bits !fir.heap<!fir.array<?xf32>>
+ %c0 = arith.constant 0 : index
+ %2 = fir.shape %c0 : (index) -> !fir.shape<1>
+ %3 = fir.embox %1(%2) allocator_idx(2) : (!fir.heap<!fir.array<?xf32>>, !fir.shape<1>) -> !fir.box<!fir.heap<!fir.array<?xf32>>>
+ fir.store %3 to %0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %4 = fir.declare %0 uniq_name("_QFuse_in_nested_cfg_loopEa") fortran_attrs<allocatable> data_attr(#cuf.cuda<device>) : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %c1 = arith.constant 1 : index
+ cf.br ^bb1(%c0 : index)
+^bb1(%i: index):
+ %5 = arith.cmpi slt, %i, %c1 : index
+ cf.cond_br %5, ^bb2(%c0 : index), ^bb5
+^bb2(%j: index):
+ %6 = arith.cmpi slt, %j, %c1 : index
+ cf.cond_br %6, ^bb3, ^bb4
+^bb3:
+ %7 = cuf.allocate %4 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> data_attr(device) -> i32
+ %8 = arith.addi %j, %c1 : index
+ cf.br ^bb2(%8 : index)
+^bb4:
+ cuf.free %4 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> data_attr(device)
+ %9 = arith.addi %i, %c1 : index
+ cf.br ^bb1(%9 : index)
+^bb5:
+ return
+}
+
+// CHECK-LABEL: func.func @_QPuse_in_nested_cfg_loop
+// CHECK: fir.embox
+// CHECK: arith.constant 1
+// CHECK: cuf.alloc
+// CHECK: fir.store
+// CHECK: fir.declare
+// CHECK: cf.br ^bb1
+// CHECK: ^bb1
+// CHECK-NOT: cuf.alloc !fir.box
+// CHECK: ^bb3:
+// CHECK: cuf.allocate
+// CHECK: ^bb4:
+// CHECK: cuf.free
+
+// -----
+
+// Test 20: The loop is only reached through a non-loop successor block. The
+// group still sinks out of the entry block, into that block before the branch
+// into the loop header.
+func.func @_QPcfg_loop_in_successor(%arg0: !fir.ref<i32>) {
+ %0 = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> uniq_name("_QFcfg_loop_in_successorEa") bindc_name("a") data_attr(device) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %1 = fir.zero_bits !fir.heap<!fir.array<?xf32>>
+ %c0 = arith.constant 0 : index
+ %2 = fir.shape %c0 : (index) -> !fir.shape<1>
+ %3 = fir.embox %1(%2) allocator_idx(2) : (!fir.heap<!fir.array<?xf32>>, !fir.shape<1>) -> !fir.box<!fir.heap<!fir.array<?xf32>>>
+ fir.store %3 to %0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %4 = fir.declare %0 uniq_name("_QFcfg_loop_in_successorEa") fortran_attrs<allocatable> data_attr(#cuf.cuda<device>) : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %c1 = arith.constant 1 : index
+ %5 = fir.load %arg0 : !fir.ref<i32>
+ %c0_i32 = arith.constant 0 : i32
+ %6 = arith.cmpi slt, %5, %c0_i32 : i32
+ cf.cond_br %6, ^bb1, ^bb2
+^bb1:
+ fir.call @_FortranAStopStatementText(%c0_i32, %c0_i32) : (i32, i32) -> none
+ fir.unreachable
+^bb2:
+ %c2 = arith.constant 2 : index
+ cf.br ^bb3(%c0 : index)
+^bb3(%i: index):
+ %7 = arith.cmpi slt, %i, %c2 : index
+ cf.cond_br %7, ^bb4, ^bb5
+^bb4:
+ %8 = cuf.allocate %4 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> data_attr(device) -> i32
+ %9 = arith.addi %i, %c1 : index
+ cf.br ^bb3(%9 : index)
+^bb5:
+ cuf.free %4 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> data_attr(device)
+ return
+}
+func.func private @_FortranAStopStatementText(i32, i32) -> none
+
+// CHECK-LABEL: func.func @_QPcfg_loop_in_successor
+// CHECK: fir.embox
+// CHECK: cf.cond_br
+// CHECK: ^bb1:
+// CHECK: fir.unreachable
+// CHECK: ^bb2:
+// CHECK: arith.constant 2
+// CHECK: cuf.alloc
+// CHECK: fir.store
+// CHECK: fir.declare
+// CHECK: cf.br ^bb3
+// CHECK: ^bb3
+// CHECK-NOT: cuf.alloc !fir.box
+// CHECK: ^bb4:
+// CHECK: cuf.allocate
+// CHECK: ^bb5:
+// CHECK: cuf.free
+
+// -----
+
+// Test 21: All uses are in the body block of a CFG loop, so the nearest common
+// dominator is the body itself and the climb takes two steps, body to header
+// to entry.
+func.func @_QPuses_only_in_cfg_loop_body() {
+ %0 = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> uniq_name("_QFuses_only_in_cfg_loop_bodyEa") bindc_name("a") data_attr(device) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %1 = fir.zero_bits !fir.heap<!fir.array<?xf32>>
+ %c0 = arith.constant 0 : index
+ %2 = fir.shape %c0 : (index) -> !fir.shape<1>
+ %3 = fir.embox %1(%2) allocator_idx(2) : (!fir.heap<!fir.array<?xf32>>, !fir.shape<1>) -> !fir.box<!fir.heap<!fir.array<?xf32>>>
+ fir.store %3 to %0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %4 = fir.declare %0 uniq_name("_QFuses_only_in_cfg_loop_bodyEa") fortran_attrs<allocatable> data_attr(#cuf.cuda<device>) : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %c1 = arith.constant 1 : index
+ cf.br ^bb1(%c0 : index)
+^bb1(%i: index):
+ %5 = arith.cmpi slt, %i, %c1 : index
+ cf.cond_br %5, ^bb2, ^bb3
+^bb2:
+ %6 = cuf.allocate %4 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> data_attr(device) -> i32
+ cuf.free %4 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> data_attr(device)
+ %7 = arith.addi %i, %c1 : index
+ cf.br ^bb1(%7 : index)
+^bb3:
+ return
+}
+
+// CHECK-LABEL: func.func @_QPuses_only_in_cfg_loop_body
+// CHECK: fir.embox
+// CHECK: arith.constant 1
+// CHECK: cuf.alloc
+// CHECK: fir.store
+// CHECK: fir.declare
+// CHECK: cf.br ^bb1
+// CHECK: ^bb1
+// CHECK-NOT: cuf.alloc !fir.box
+// CHECK: ^bb2:
+// CHECK: cuf.allocate
+// CHECK: cuf.free
+
+// -----
+
+// Test 22: An irreducible cycle (two entries, as GOTO can build) that natural
+// loop analysis does not report as a loop. The uses are in a block on the
+// cycle and after it, so the group must stay at entry.
+func.func @_QPirreducible_cycle(%arg0: !fir.ref<i32>) {
+ %0 = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> uniq_name("_QFirreducible_cycleEa") bindc_name("a") data_attr(device) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %1 = fir.zero_bits !fir.heap<!fir.array<?xf32>>
+ %c0 = arith.constant 0 : index
+ %2 = fir.shape %c0 : (index) -> !fir.shape<1>
+ %3 = fir.embox %1(%2) allocator_idx(2) : (!fir.heap<!fir.array<?xf32>>, !fir.shape<1>) -> !fir.box<!fir.heap<!fir.array<?xf32>>>
+ fir.store %3 to %0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %4 = fir.declare %0 uniq_name("_QFirreducible_cycleEa") fortran_attrs<allocatable> data_attr(#cuf.cuda<device>) : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %5 = fir.load %arg0 : !fir.ref<i32>
+ %c0_i32 = arith.constant 0 : i32
+ %6 = arith.cmpi slt, %5, %c0_i32 : i32
+ cf.cond_br %6, ^bb1, ^bb2
+^bb1:
+ cf.br ^bb2
+^bb2:
+ %7 = cuf.allocate %4 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> data_attr(device) -> i32
+ %8 = fir.load %arg0 : !fir.ref<i32>
+ %9 = arith.cmpi sgt, %8, %c0_i32 : i32
+ cf.cond_br %9, ^bb1, ^bb3
+^bb3:
+ cuf.free %4 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> data_attr(device)
+ return
+}
+
+// CHECK-LABEL: func.func @_QPirreducible_cycle
+// CHECK: fir.embox
+// CHECK: fir.load
+// CHECK: arith.cmpi
+// CHECK: cuf.alloc
+// CHECK: fir.store
+// CHECK: fir.declare
+// CHECK: cf.cond_br
+// CHECK: ^bb1:
+// CHECK-NOT: cuf.alloc !fir.box
+// CHECK: ^bb2:
+// CHECK-NOT: cuf.alloc !fir.box
+// CHECK: cuf.allocate
+// CHECK: ^bb3:
+// CHECK: cuf.free
More information about the flang-commits
mailing list