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

[via flang-commits]

Author: Scott Manley
Date: 2026-09-02T17:58:31-04:00
New Revision: 81bf35d75a2925ff57a470b7a54c5d3730ade165

URL: https://github.com/llvm/llvm-project/commit/81bf35d75a2925ff57a470b7a54c5d3730ade165
DIFF: https://github.com/llvm/llvm-project/commit/81bf35d75a2925ff57a470b7a54c5d3730ade165.diff

LOG: [flang][cuda] Delay descriptor alloc when addressed reused on host/device (#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.

Added: 
    

Modified: 
    flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp
    flang/test/Transforms/CUF/cuf-alloc-delay.fir

Removed: 
    


################################################################################
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