[flang-commits] [flang] [flang][cuda] Delay descriptor alloc when addressed reused on host/device (PR #220534)

[Scott Manley via flang-commits]

https://github.com/rscottmanley created https://github.com/llvm/llvm-project/pull/220534

CSE can share one fir.coordinate_of between the host-association capture store and a later fir.load. Treating that coordinate_of as a real use made cuf-alloc-delay think the movable group depended on an operand at the sink point, so the device descriptor stayed at function entry and cudaMallocManaged ran before cudaSetDevice.

Count only users of the slot address that actually read it. Stores that populate the tuple still sink with the allocation group.

>From a8e745227e1bd159725445b2a3757f954dedf391 Mon Sep 17 00:00:00 2001
From: Scott Manley <scmanley at nvidia.com>
Date: Wed, 2 Sep 2026 02:46:12 -0700
Subject: [PATCH] [flang][cuda] Delay descriptor alloc when host-assoc slot
 address is reused

CSE can share one fir.coordinate_of between the host-association capture
store and a later fir.load. Treating that coordinate_of as a real use
made cuf-alloc-delay think the movable group depended on an operand at
the sink point, so the device descriptor stayed at function entry and
cudaMallocManaged ran before cudaSetDevice.

Count only users of the slot address that actually read it. Stores that
populate the tuple still sink with the allocation group.
---
 .../Transforms/CUDA/CUFAllocDelay.cpp         | 22 +++++------
 flang/test/Transforms/CUF/cuf-alloc-delay.fir | 39 ++++++++++++++++++-
 2 files changed, 49 insertions(+), 12 deletions(-)

diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp
index 1ee7f96bfc9f3..5b3f7afd7f656 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp
@@ -45,15 +45,6 @@ static fir::CoordinateOp getHostAssocTupleSlot(fir::StoreOp storeOp,
   return coord;
 }
 
-/// Return true if \p coord's result is only stored into, so it writes the tuple
-/// rather than reading it.
-static bool onlyPopulatesSlot(fir::CoordinateOp coord) {
-  return llvm::all_of(coord->getUsers(), [&](mlir::Operation *user) {
-    auto storeOp = mlir::dyn_cast<fir::StoreOp>(user);
-    return storeOp && storeOp.getMemref() == coord.getResult();
-  });
-}
-
 /// 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
@@ -91,9 +82,18 @@ findDelayTarget(fir::DeclareOp declareOp, mlir::Block *entryBlock,
       hostAssocStores.push_back(storeOp);
       for (mlir::Operation *tupleUser : slot.getRef().getUsers()) {
         auto coord = mlir::dyn_cast<fir::CoordinateOp>(tupleUser);
-        if (coord && onlyPopulatesSlot(coord))
+        if (!coord) {
+          recordRealUse(tupleUser);
           continue;
-        recordRealUse(tupleUser);
+        }
+        // A coordinate_of only computes the slot address. Stores through that
+        // address populate the tuple; other users actually consume the slot.
+        for (mlir::Operation *coordUser : coord->getUsers()) {
+          auto slotStore = mlir::dyn_cast<fir::StoreOp>(coordUser);
+          if (slotStore && slotStore.getMemref() == coord.getResult())
+            continue;
+          recordRealUse(coordUser);
+        }
       }
     }
   }
diff --git a/flang/test/Transforms/CUF/cuf-alloc-delay.fir b/flang/test/Transforms/CUF/cuf-alloc-delay.fir
index 99e8f184b96c4..f676938a45bc2 100644
--- a/flang/test/Transforms/CUF/cuf-alloc-delay.fir
+++ b/flang/test/Transforms/CUF/cuf-alloc-delay.fir
@@ -550,7 +550,8 @@ func.func private @_QPsetup3()
 // -----
 
 // Test 16: The tuple slot is read back in the host, so the group sinks past the
-// call but stays before the coordinate_of feeding that read.
+// call. A coordinate_of only computes the slot address, so the group stays
+// before the fir.load that consumes it.
 func.func @_QPtuple_slot_read() {
   %tuple = fir.alloca tuple<!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>
   %c0_i32 = arith.constant 0 : i32
@@ -576,11 +577,47 @@ func.func private @_QPsetup4()
 // CHECK:         fir.coordinate_of
 // CHECK:         fir.embox
 // CHECK:         fir.call @_QPsetup4
+// CHECK:         fir.coordinate_of
 // CHECK:         cuf.alloc
 // CHECK:         fir.store {{.*}} : !fir.ref<!fir.box
 // CHECK:         fir.declare
 // CHECK:         fir.store {{.*}} : !fir.llvm_ptr
+// CHECK:         fir.load
+// CHECK:         cuf.allocate
+// CHECK:         cuf.free
+
+// -----
+
+// Test 17: The same coordinate_of feeds both the host-association capture store
+// and a later load (as CSE can produce). The coordinate_of is not a real use.
+func.func @_QPshared_slot_addr() {
+  %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<device>, uniq_name = "_QFshared_slot_addrEa"} -> !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 : 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<device>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFshared_slot_addrEa"} : (!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 @_QPsetup5() : () -> ()
+  %5 = fir.load %slot : !fir.llvm_ptr<!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>
+  %6 = cuf.allocate %5 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>} -> i32
+  cuf.free %4 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>}
+  return
+}
+func.func private @_QPsetup5()
+
+// CHECK-LABEL: func.func @_QPshared_slot_addr
 // CHECK:         fir.coordinate_of
+// CHECK:         fir.embox
+// CHECK:         fir.call @_QPsetup5
+// CHECK:         cuf.alloc
+// CHECK:         fir.store {{.*}} : !fir.ref<!fir.box
+// CHECK:         fir.declare
+// CHECK:         fir.store {{.*}} : !fir.llvm_ptr
 // CHECK:         fir.load
 // CHECK:         cuf.allocate
 // CHECK:         cuf.free