[flang-commits] [flang] 0c265f8 - [flang][cuda] Keep delayed descriptor allocs out of CFG loops (#232992)
via flang-commits
flang-commits at lists.llvm.org
Mon Oct 19 10:55:12 PDT 2026
Author: Zhen Wang
Date: 2026-10-19T10:55:06-07:00
New Revision: 0c265f8418491b5a1026d07ca6146befdd970a80
URL: https://github.com/llvm/llvm-project/commit/0c265f8418491b5a1026d07ca6146befdd970a80
DIFF: https://github.com/llvm/llvm-project/commit/0c265f8418491b5a1026d07ca6146befdd970a80.diff
LOG: [flang][cuda] Keep delayed descriptor allocs out of CFG loops (#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.
Added:
Modified:
flang/include/flang/Optimizer/Transforms/Passes.td
flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp
flang/test/Transforms/CUF/cuf-alloc-delay.fir
Removed:
################################################################################
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..7721f55257d92 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.
//
@@ -19,8 +20,10 @@
#include "flang/Optimizer/Dialect/FIRType.h"
#include "mlir/IR/Block.h"
#include "mlir/IR/Dominance.h"
+#include "mlir/IR/RegionGraphTraits.h"
#include "mlir/Pass/Pass.h"
#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SCCIterator.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"
@@ -48,13 +51,15 @@ 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 one of \p cyclicBlocks,
+/// the blocks 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.
static mlir::Operation *
findDelayTarget(fir::DeclareOp declareOp, mlir::Block *entryBlock,
mlir::DominanceInfo &domInfo,
+ const llvm::SmallPtrSetImpl<mlir::Block *> &cyclicBlocks,
llvm::SmallVectorImpl<fir::StoreOp> &hostAssocStores) {
mlir::Region *funcRegion = entryBlock->getParent();
@@ -108,6 +113,18 @@ 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. The entry block has no
+ // predecessors, so it is never on a cycle and the climb always ends.
+ while (cyclicBlocks.contains(common)) {
+ mlir::DominanceInfoNode *node = domInfo.getNode(common);
+ if (!node || !node->getIDom())
+ return nullptr;
+ common = node->getIDom()->getBlock();
+ }
+
mlir::Operation *earliest = nullptr;
for (mlir::Operation *use : uses)
if (use->getBlock() == common &&
@@ -133,6 +150,15 @@ struct CUFAllocDelay : public fir::impl::CUFAllocDelayBase<CUFAllocDelay> {
if (auto allocOp = mlir::dyn_cast<cuf::AllocOp>(op))
if (mlir::isa<fir::BaseBoxType>(allocOp.getInType()))
boxAllocOps.push_back(allocOp);
+ if (boxAllocOps.empty())
+ return;
+
+ // Blocks on a CFG cycle, natural loop or not: every block of a strongly
+ // connected component with more than one block or with a self edge.
+ llvm::SmallPtrSet<mlir::Block *, 8> cyclicBlocks;
+ for (auto it = llvm::scc_begin(&func.getBody()); !it.isAtEnd(); ++it)
+ if (it.hasCycle())
+ cyclicBlocks.insert(it->begin(), it->end());
for (cuf::AllocOp allocOp : boxAllocOps) {
// Find the fir.declare and fir.store using this cuf.alloc; bail on any
@@ -152,8 +178,8 @@ struct CUFAllocDelay : public fir::impl::CUFAllocDelayBase<CUFAllocDelay> {
continue;
llvm::SmallVector<fir::StoreOp> hostAssocStores;
- mlir::Operation *delayTarget =
- findDelayTarget(declareOp, &entryBlock, domInfo, hostAssocStores);
+ mlir::Operation *delayTarget = findDelayTarget(
+ declareOp, &entryBlock, domInfo, cyclicBlocks, hostAssocStores);
if (!delayTarget)
continue;
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