[flang-commits] [flang] [flang][cuda] Fix cuf-alloc-delay for host-associated allocatables (PR #211103)
Vijay Kandiah via flang-commits
flang-commits at lists.llvm.org
Tue Jul 21 14:46:43 PDT 2026
https://github.com/VijayKandiah updated https://github.com/llvm/llvm-project/pull/211103
>From d3016ce88c829a343f4fb5ecb69b179d7f052f31 Mon Sep 17 00:00:00 2001
From: Vijay Kandiah <vkandiah at nvidia.com>
Date: Tue, 21 Jul 2026 13:24:48 -0700
Subject: [PATCH 1/2] [flang][cuda] Fix cuf-alloc-delay for host-associated
allocatables
---
.../Transforms/CUDA/CUFAllocDelay.cpp | 61 +++++---------
flang/test/Transforms/CUF/cuf-alloc-delay.fir | 81 +++++++++++++++++--
2 files changed, 94 insertions(+), 48 deletions(-)
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp
index 9ed701937b003..d81c64aa06568 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp
@@ -28,34 +28,23 @@ namespace fir {
namespace {
-/// Find the earliest use of any of the declare results, returning the
-/// operation before which the cuf.alloc group should be placed.
-///
-/// Uses inside nested regions (fir.if, fir.do_loop, etc.) are resolved to
-/// the parent op in the entry block. When all real uses reside in a single
-/// successor block of the entry block, the target is placed in that block
-/// (just before the earliest use there) so the alloc is deferred past any
-/// setup code that precedes the branch.
-///
-/// Stores to fir.llvm_ptr destinations (host-association tuple slots) are
-/// skipped as "uses" and collected in \p hostAssocStores so the caller can
-/// move them along with the sunk group.
-static mlir::Operation *
-findDelayTarget(fir::DeclareOp declareOp, mlir::Block *entryBlock,
- llvm::SmallVectorImpl<fir::StoreOp> &hostAssocStores) {
+/// Return the op before which the cuf.alloc group should be placed: the
+/// earliest use of the descriptor. Every use counts, including a store of the
+/// descriptor into a host-association tuple slot (fir.store to fir.llvm_ptr) -
+/// that store reads its own fir.coordinate_of slot, so the group is only sunk
+/// to *before* it and the store itself is never moved. Sinking strictly before
+/// any use keeps dominance valid.
+static mlir::Operation *findDelayTarget(fir::DeclareOp declareOp,
+ mlir::Block *entryBlock) {
mlir::Operation *earliest = nullptr;
mlir::Region *funcRegion = entryBlock->getParent();
- // Track successor-block uses: which blocks, and the earliest op in each.
+ // Uses per successor block, with the earliest op in each.
llvm::SmallDenseMap<mlir::Block *, mlir::Operation *> successorEarliest;
- auto updateEarliest = [&](mlir::Operation *user) {
- if (auto store = mlir::dyn_cast<fir::StoreOp>(user)) {
- if (mlir::isa<fir::LLVMPointerType>(store.getMemref().getType())) {
- hostAssocStores.push_back(store);
- return;
- }
- }
+ // Resolve a use in a nested region or successor block to a target in/after
+ // the entry block.
+ auto recordRealUse = [&](mlir::Operation *user) {
mlir::Operation *target = user;
while (target->getBlock() != entryBlock) {
// User in another block of the same function.
@@ -77,7 +66,7 @@ findDelayTarget(fir::DeclareOp declareOp, mlir::Block *entryBlock,
for (mlir::Value result : declareOp->getResults()) {
for (mlir::Operation *user : result.getUsers())
- updateEarliest(user);
+ recordRealUse(user);
}
if (earliest)
@@ -110,9 +99,8 @@ struct CUFAllocDelay : public fir::impl::CUFAllocDelayBase<CUFAllocDelay> {
boxAllocOps.push_back(allocOp);
for (cuf::AllocOp allocOp : boxAllocOps) {
- // Find the fir.declare and fir.store that use this cuf.alloc.
- // Bail out if the alloc has any unexpected users to avoid breaking
- // dominance for patterns we don't track.
+ // Find the fir.declare and fir.store using this cuf.alloc; bail on any
+ // unexpected user.
fir::DeclareOp declareOp = nullptr;
fir::StoreOp storeOp = nullptr;
bool hasUnknownUser = false;
@@ -127,23 +115,19 @@ struct CUFAllocDelay : public fir::impl::CUFAllocDelayBase<CUFAllocDelay> {
if (!declareOp || hasUnknownUser)
continue;
- llvm::SmallVector<fir::StoreOp> hostAssocStores;
- mlir::Operation *delayTarget =
- findDelayTarget(declareOp, &entryBlock, hostAssocStores);
+ mlir::Operation *delayTarget = findDelayTarget(declareOp, &entryBlock);
if (!delayTarget)
continue;
- // Don't move if target is in the same block and at or before current pos.
+ // Skip if the target is at or before the alloc, or is the declare.
if (delayTarget->getBlock() == allocOp->getBlock() &&
(delayTarget->isBeforeInBlock(allocOp) || delayTarget == allocOp))
continue;
- // Don't move to the declare itself.
if (delayTarget == declareOp)
continue;
- // Move {cuf.alloc, fir.store, fir.declare} before the delay target.
- // The embox/zero_bits/shape/constants stay at their original positions
- // since they still dominate the new locations.
+ // Sink {cuf.alloc, fir.store, fir.declare} before the target; the
+ // embox/shape/constants stay put and still dominate the new position.
allocOp->moveBefore(delayTarget);
if (storeOp)
storeOp->moveAfter(allocOp);
@@ -151,13 +135,6 @@ struct CUFAllocDelay : public fir::impl::CUFAllocDelayBase<CUFAllocDelay> {
declareOp->moveAfter(storeOp);
else
declareOp->moveAfter(allocOp);
-
- // Move host-association tuple stores after the declare.
- mlir::Operation *insertAfter = declareOp;
- for (fir::StoreOp hostStore : hostAssocStores) {
- hostStore->moveAfter(insertAfter);
- insertAfter = hostStore;
- }
}
}
};
diff --git a/flang/test/Transforms/CUF/cuf-alloc-delay.fir b/flang/test/Transforms/CUF/cuf-alloc-delay.fir
index f77c8ba6ad2fb..ae97ffc45bafc 100644
--- a/flang/test/Transforms/CUF/cuf-alloc-delay.fir
+++ b/flang/test/Transforms/CUF/cuf-alloc-delay.fir
@@ -326,8 +326,8 @@ func.func @_QPuse_in_multi_successor(%arg0: !fir.ref<i32>) {
// -----
-// Test 10: Host-association tuple store — fir.store to fir.llvm_ptr
-// should not block delaying and should move along with the group.
+// Test 10: A host-association store (fir.store to fir.llvm_ptr) is a use, so
+// the group sinks to just before it and the store is not moved.
func.func @_QPhost_assoc() {
%tuple = fir.alloca tuple<!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>
%c0_i32 = arith.constant 0 : i32
@@ -349,19 +349,88 @@ func.func @_QPhost_assoc() {
func.func private @_QFPcontained(!fir.ref<tuple<!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>>)
// CHECK-LABEL: func.func @_QPhost_assoc
-// The fir.store to fir.llvm_ptr (host-assoc tuple) does not block delaying.
-// cuf.alloc+store+declare+host-assoc-store all move past %c99 to just
-// before cuf.allocate.
// CHECK: fir.alloca
// CHECK: fir.coordinate_of
// CHECK: fir.zero_bits
// CHECK: fir.embox
-// CHECK: arith.constant 99
// CHECK: cuf.alloc
// CHECK: fir.store {{.*}} : !fir.ref<!fir.box
// CHECK: fir.declare
// CHECK: fir.store {{.*}} : !fir.llvm_ptr
+// CHECK: arith.constant 99
// CHECK: cuf.allocate
// CHECK: fir.call @_QFPcontained
// CHECK: cuf.free
// CHECK: return
+
+// -----
+
+// Test 11: The descriptor's only entry-block use is its host-association store
+// (the allocate happens inside the callee). The group and store stay before the
+// call, so the callee sees an initialized descriptor.
+func.func @_QPhost_assoc_call() {
+ %tuple = fir.alloca tuple<!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>
+ %c0_i32 = arith.constant 0 : i32
+ %slot = fir.coordinate_of %tuple, %c0_i32 : (!fir.ref<tuple<!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>>, i32) -> !fir.llvm_ptr<!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>
+ %0 = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a", data_attr = #cuf.cuda<managed>, uniq_name = "_QFhost_assoc_callEa"} -> !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 = 3 : i32} : (!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 {data_attr = #cuf.cuda<managed>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFhost_assoc_callEa"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ fir.store %4 to %slot : !fir.llvm_ptr<!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>
+ fir.call @_QFPgo(%tuple) : (!fir.ref<tuple<!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>>) -> ()
+ cuf.free %4 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<managed>}
+ return
+}
+func.func private @_QFPgo(!fir.ref<tuple<!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>>)
+
+// CHECK-LABEL: func.func @_QPhost_assoc_call
+// CHECK: cuf.alloc
+// CHECK: fir.store {{.*}} : !fir.ref<!fir.box
+// CHECK: fir.declare
+// CHECK: fir.store {{.*}} : !fir.llvm_ptr
+// CHECK: fir.call @_QFPgo
+// CHECK: cuf.free
+// CHECK: return
+
+// -----
+
+// Test 12: A shared tuple whose slots fill in order, so the descriptor's own
+// slot (index 1) is defined after an earlier coordinate_of of the same tuple
+// (index 0). The group must sink to before its host-association store, not
+// before its own slot.
+func.func @_QPhost_assoc_late_slot() {
+ %tuple = fir.alloca tuple<!fir.ref<i32>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>
+ %c0_i32 = arith.constant 0 : i32
+ %c1_i32 = arith.constant 1 : i32
+ %0 = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a", data_attr = #cuf.cuda<managed>, uniq_name = "_QFhost_assoc_late_slotEa"} -> !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 = 3 : i32} : (!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 {data_attr = #cuf.cuda<managed>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFhost_assoc_late_slotEa"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ %slot0 = fir.coordinate_of %tuple, %c0_i32 : (!fir.ref<tuple<!fir.ref<i32>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>>, i32) -> !fir.llvm_ptr<!fir.ref<i32>>
+ %scalar = fir.alloca i32
+ fir.store %scalar to %slot0 : !fir.llvm_ptr<!fir.ref<i32>>
+ %slot1 = fir.coordinate_of %tuple, %c1_i32 : (!fir.ref<tuple<!fir.ref<i32>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>>, i32) -> !fir.llvm_ptr<!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>
+ fir.store %4 to %slot1 : !fir.llvm_ptr<!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>
+ fir.call @_QFPgo2(%tuple) : (!fir.ref<tuple<!fir.ref<i32>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>>) -> ()
+ cuf.free %4 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<managed>}
+ return
+}
+func.func private @_QFPgo2(!fir.ref<tuple<!fir.ref<i32>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>>)
+
+// CHECK-LABEL: func.func @_QPhost_assoc_late_slot
+// CHECK: fir.coordinate_of {{.*}} -> !fir.llvm_ptr<!fir.ref<i32>>
+// CHECK: fir.store {{.*}} : !fir.llvm_ptr<!fir.ref<i32>>
+// CHECK: fir.coordinate_of {{.*}} -> !fir.llvm_ptr<!fir.ref<!fir.box
+// CHECK: cuf.alloc
+// CHECK: fir.store {{.*}} : !fir.ref<!fir.box
+// CHECK: fir.declare
+// CHECK: fir.store {{.*}} : !fir.llvm_ptr<!fir.ref<!fir.box
+// CHECK: fir.call @_QFPgo2
+// CHECK: cuf.free
+// CHECK: return
>From f0d6665d617078c758966228523f79364da46f51 Mon Sep 17 00:00:00 2001
From: Vijay Kandiah <vkandiah at nvidia.com>
Date: Tue, 21 Jul 2026 14:46:30 -0700
Subject: [PATCH 2/2] [flang][cuda] cuf-alloc-delay update comments
---
flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp | 10 ++++------
flang/test/Transforms/CUF/cuf-alloc-delay.fir | 3 +--
2 files changed, 5 insertions(+), 8 deletions(-)
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp
index d81c64aa06568..f9e62cccdc6e2 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp
@@ -28,12 +28,10 @@ namespace fir {
namespace {
-/// Return the op before which the cuf.alloc group should be placed: the
-/// earliest use of the descriptor. Every use counts, including a store of the
-/// descriptor into a host-association tuple slot (fir.store to fir.llvm_ptr) -
-/// that store reads its own fir.coordinate_of slot, so the group is only sunk
-/// to *before* it and the store itself is never moved. Sinking strictly before
-/// any use keeps dominance valid.
+/// Find the earliest use of the descriptor and return the op before which the
+/// cuf.alloc group should be placed. Uses in nested regions (fir.if,
+/// fir.do_loop, ...) resolve to the enclosing entry-block op; uses confined to
+/// a single successor block resolve to that block.
static mlir::Operation *findDelayTarget(fir::DeclareOp declareOp,
mlir::Block *entryBlock) {
mlir::Operation *earliest = nullptr;
diff --git a/flang/test/Transforms/CUF/cuf-alloc-delay.fir b/flang/test/Transforms/CUF/cuf-alloc-delay.fir
index ae97ffc45bafc..d4e9a6a4479f3 100644
--- a/flang/test/Transforms/CUF/cuf-alloc-delay.fir
+++ b/flang/test/Transforms/CUF/cuf-alloc-delay.fir
@@ -399,8 +399,7 @@ func.func private @_QFPgo(!fir.ref<tuple<!fir.ref<!fir.box<!fir.heap<!fir.array<
// Test 12: A shared tuple whose slots fill in order, so the descriptor's own
// slot (index 1) is defined after an earlier coordinate_of of the same tuple
-// (index 0). The group must sink to before its host-association store, not
-// before its own slot.
+// (index 0). The group must sink to before its host-association store.
func.func @_QPhost_assoc_late_slot() {
%tuple = fir.alloca tuple<!fir.ref<i32>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>
%c0_i32 = arith.constant 0 : i32
More information about the flang-commits
mailing list