[flang][cuda] Delay box cuf.alloc past host association captures (#214347) CUFAllocDelay treated the store of a descriptor into a host association tuple as a use, so a device allocatable captured by an internal procedure kept its descriptor allocation in the prologue. That allocates managed memory before the program can call cudaSetDevice, binding a CUDA context to the wrong device. The store now sinks together with the allocation, constrained by the tuple's readers, and the group is placed at the nearest common dominator of all uses so it can sink into a later block. GitOrigin-RevId: 239cbf5baf092fc738e59443d11948b26d8db386
diff --git a/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp b/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp index f9e62cc..1ee7f96 100644 --- a/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp +++ b/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp
@@ -7,8 +7,9 @@ //===----------------------------------------------------------------------===// // // Delay cuf.alloc of descriptor (box) types from function entry to just before -// their first use. This defers cudaMallocManaged calls so that users can call -// cudaSetDevice before any CUDA context is created. +// their first use, possibly in a later block that dominates every use. This +// defers cudaMallocManaged calls so that users can call cudaSetDevice before +// any CUDA context is created. // //===----------------------------------------------------------------------===// @@ -17,8 +18,10 @@ #include "flang/Optimizer/Dialect/FIROps.h" #include "flang/Optimizer/Dialect/FIRType.h" #include "mlir/IR/Block.h" +#include "mlir/IR/Dominance.h" #include "mlir/Pass/Pass.h" #include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SmallVector.h" namespace fir { @@ -28,56 +31,90 @@ namespace { -/// 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; +/// Return the coordinate_of producing the host association tuple slot that +/// \p storeOp writes \p descriptor into, or null if this is not such a capture. +static fir::CoordinateOp getHostAssocTupleSlot(fir::StoreOp storeOp, + mlir::Value descriptor) { + if (storeOp.getValue() != descriptor || + !mlir::isa<fir::LLVMPointerType>(storeOp.getMemref().getType())) + return nullptr; + auto coord = storeOp.getMemref().getDefiningOp<fir::CoordinateOp>(); + if (!coord || + !mlir::isa<mlir::TupleType>(fir::unwrapRefType(coord.getRef().getType()))) + return nullptr; + 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 +/// their enclosing top-level op. +/// +/// 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, + llvm::SmallVectorImpl<fir::StoreOp> &hostAssocStores) { mlir::Region *funcRegion = entryBlock->getParent(); - // Uses per successor block, with the earliest op in each. - llvm::SmallDenseMap<mlir::Block *, mlir::Operation *> successorEarliest; + // Uses resolved to an op that sits directly in a block of the function. + llvm::SmallVector<mlir::Operation *> uses; - // 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. - if (target->getBlock() && target->getBlock()->getParent() == funcRegion) { - mlir::Block *blk = target->getBlock(); - auto it = successorEarliest.find(blk); - if (it == successorEarliest.end() || - target->isBeforeInBlock(it->second)) - successorEarliest[blk] = target; - return; - } - target = target->getParentOp(); - if (!target) - return; - } - if (!earliest || target->isBeforeInBlock(earliest)) - earliest = target; + mlir::Operation *op = user; + while (op && op->getBlock() && op->getBlock()->getParent() != funcRegion) + op = op->getParentOp(); + if (op && op->getBlock()) + uses.push_back(op); }; for (mlir::Value result : declareOp->getResults()) { - for (mlir::Operation *user : result.getUsers()) - recordRealUse(user); + for (mlir::Operation *user : result.getUsers()) { + auto storeOp = mlir::dyn_cast<fir::StoreOp>(user); + fir::CoordinateOp slot = + storeOp ? getHostAssocTupleSlot(storeOp, result) : nullptr; + if (!slot) { + recordRealUse(user); + continue; + } + // Whoever reads the tuple must still see a populated slot. + hostAssocStores.push_back(storeOp); + for (mlir::Operation *tupleUser : slot.getRef().getUsers()) { + auto coord = mlir::dyn_cast<fir::CoordinateOp>(tupleUser); + if (coord && onlyPopulatesSlot(coord)) + continue; + recordRealUse(tupleUser); + } + } } - if (earliest) - return earliest; + if (uses.empty()) + return nullptr; - // No entry-block uses. If all successor uses are in a single block, - // delay directly into that block (before the earliest use there). - // Otherwise fall back to the entry block's terminator. - if (successorEarliest.size() == 1) - return successorEarliest.begin()->second; - if (!successorEarliest.empty()) - return entryBlock->getTerminator(); - return nullptr; + mlir::Block *common = uses.front()->getBlock(); + for (mlir::Operation *use : uses) { + common = domInfo.findNearestCommonDominator(common, use->getBlock()); + if (!common) + return nullptr; + } + + mlir::Operation *earliest = nullptr; + for (mlir::Operation *use : uses) + if (use->getBlock() == common && + (!earliest || use->isBeforeInBlock(earliest))) + earliest = use; + + return earliest ? earliest : common->getTerminator(); } struct CUFAllocDelay : public fir::impl::CUFAllocDelayBase<CUFAllocDelay> { @@ -88,6 +125,7 @@ return; mlir::Block &entryBlock = func.front(); + mlir::DominanceInfo domInfo(func); // Collect box-type cuf.alloc ops in the entry block. llvm::SmallVector<cuf::AllocOp> boxAllocOps; @@ -113,7 +151,9 @@ if (!declareOp || hasUnknownUser) continue; - mlir::Operation *delayTarget = findDelayTarget(declareOp, &entryBlock); + llvm::SmallVector<fir::StoreOp> hostAssocStores; + mlir::Operation *delayTarget = + findDelayTarget(declareOp, &entryBlock, domInfo, hostAssocStores); if (!delayTarget) continue; @@ -124,15 +164,36 @@ if (delayTarget == declareOp) continue; - // 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); + // Ops that move together, keeping their relative order. + llvm::SmallVector<mlir::Operation *> group; + group.push_back(allocOp); if (storeOp) - storeOp->moveAfter(allocOp); - if (storeOp) - declareOp->moveAfter(storeOp); - else - declareOp->moveAfter(allocOp); + group.push_back(storeOp); + group.push_back(declareOp); + for (fir::StoreOp hostAssocStore : hostAssocStores) + group.push_back(hostAssocStore); + + // Whatever the group reads from outside itself stays put, so it must + // already dominate the new position. + llvm::SmallPtrSet<mlir::Operation *, 8> groupSet(group.begin(), + group.end()); + auto readsDominateTarget = [&](mlir::Operation *op) { + return llvm::all_of(op->getOperands(), [&](mlir::Value operand) { + mlir::Operation *def = operand.getDefiningOp(); + return (def && groupSet.contains(def)) || + domInfo.properlyDominates(operand, delayTarget); + }); + }; + if (!llvm::all_of(group, readsDominateTarget)) + continue; + + // Sink the group before the target, preserving its relative order. + group.front()->moveBefore(delayTarget); + mlir::Operation *last = group.front(); + for (mlir::Operation *op : llvm::drop_begin(group)) { + op->moveAfter(last); + last = op; + } } } };
diff --git a/test/Transforms/CUF/cuf-alloc-delay.fir b/test/Transforms/CUF/cuf-alloc-delay.fir index d4e9a6a..99e8f18 100644 --- a/test/Transforms/CUF/cuf-alloc-delay.fir +++ b/test/Transforms/CUF/cuf-alloc-delay.fir
@@ -326,8 +326,8 @@ // ----- -// 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. +// Test 10: A host association store sinks with the group, so the allocation is +// still delayed to the first real use. func.func @_QPhost_assoc() { %tuple = fir.alloca tuple<!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>> %c0_i32 = arith.constant 0 : i32 @@ -353,11 +353,11 @@ // 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 @@ -365,9 +365,8 @@ // ----- -// 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. +// Test 11: The descriptor is only captured, never used in the host, so the group +// and its store must still land before the call. func.func @_QPhost_assoc_call() { %tuple = fir.alloca tuple<!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>> %c0_i32 = arith.constant 0 : i32 @@ -397,9 +396,8 @@ // ----- -// 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. +// Test 12: A shared tuple filled in slot order, so the descriptor's own slot is +// defined after an earlier coordinate_of of the same tuple. 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 @@ -433,3 +431,156 @@ // CHECK: fir.call @_QFPgo2 // CHECK: cuf.free // CHECK: return + +// ----- + +// Test 13: The tuple is populated in the prologue but the descriptor is not used +// until later, so the group and its store sink past the intervening call. +func.func @_QPhost_assoc_prologue() { + %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 = "_QFhost_assoc_prologueEa"} -> !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 = "_QFhost_assoc_prologueEa"} : (!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 @_QPsetup() : () -> () + %5 = cuf.allocate %4 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>} -> i32 + fir.call @_QFPcontained2(%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<device>} + return +} +func.func private @_QPsetup() +func.func private @_QFPcontained2(!fir.ref<tuple<!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>>) + +// CHECK-LABEL: func.func @_QPhost_assoc_prologue +// CHECK: fir.coordinate_of +// CHECK: fir.zero_bits +// CHECK: fir.embox +// CHECK: fir.call @_QPsetup +// CHECK: cuf.alloc +// CHECK: fir.store {{.*}} : !fir.ref<!fir.box +// CHECK: fir.declare +// CHECK: fir.store {{.*}} : !fir.llvm_ptr +// CHECK: cuf.allocate +// CHECK: fir.call @_QFPcontained2 +// CHECK: cuf.free +// CHECK: return + +// ----- + +// Test 14: Uses span several successor blocks, but ^bb2 dominates them all, so +// the group sinks into ^bb2 instead of stopping at the entry terminator. +func.func @_QPuse_in_dominating_successor(%arg0: !fir.ref<i32>) { + %0 = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a", data_attr = #cuf.cuda<device>, uniq_name = "_QFuse_in_dominating_successorEa"} -> !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 = "_QFuse_in_dominating_successorEa"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> + %c0_i32 = arith.constant 0 : i32 + %5 = fir.load %arg0 : !fir.ref<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: + fir.call @_QPinit_runtime() : () -> () + %7 = cuf.allocate %4 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>} -> i32 + %8 = fir.load %arg0 : !fir.ref<i32> + %9 = arith.cmpi slt, %8, %c0_i32 : i32 + cf.cond_br %9, ^bb3, ^bb4 +^bb3: + cuf.free %4 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>} + return +^bb4: + cuf.free %4 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>} + return +} +func.func private @_FortranAStopStatementText(i32, i32) -> none +func.func private @_QPinit_runtime() + +// CHECK-LABEL: func.func @_QPuse_in_dominating_successor +// CHECK: cf.cond_br +// CHECK: ^bb1: +// CHECK: fir.unreachable +// CHECK: ^bb2: +// The allocation happens after the call, not before the branch. +// CHECK: fir.call @_QPinit_runtime +// CHECK: cuf.alloc +// CHECK: fir.store +// CHECK: fir.declare +// CHECK: cuf.allocate + +// ----- + +// Test 15: A store into an !fir.llvm_ptr that is not a tuple slot goes to an +// unknown destination, so it is an ordinary use and pins the group. +func.func @_QPllvm_ptr_not_tuple(%arg0: !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 = "_QFllvm_ptr_not_tupleEa"} -> !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 = "_QFllvm_ptr_not_tupleEa"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> + %opaque = fir.convert %arg0 : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> !fir.llvm_ptr<!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>> + fir.store %4 to %opaque : !fir.llvm_ptr<!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>> + fir.call @_QPsetup3() : () -> () + cuf.free %4 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>} + return +} +func.func private @_QPsetup3() + +// CHECK-LABEL: func.func @_QPllvm_ptr_not_tuple +// CHECK: fir.convert +// CHECK: cuf.alloc +// CHECK: fir.store {{.*}} : !fir.ref<!fir.box +// CHECK: fir.declare +// CHECK: fir.store {{.*}} : !fir.llvm_ptr +// CHECK: fir.call @_QPsetup3 +// CHECK: cuf.free + +// ----- + +// 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. +func.func @_QPtuple_slot_read() { + %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 = "_QFtuple_slot_readEa"} -> !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 = "_QFtuple_slot_readEa"} : (!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 @_QPsetup4() : () -> () + %reread = 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>>>>> + %5 = fir.load %reread : !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 @_QPsetup4() + +// CHECK-LABEL: func.func @_QPtuple_slot_read +// CHECK: fir.coordinate_of +// CHECK: fir.embox +// CHECK: fir.call @_QPsetup4 +// CHECK: cuf.alloc +// CHECK: fir.store {{.*}} : !fir.ref<!fir.box +// CHECK: fir.declare +// CHECK: fir.store {{.*}} : !fir.llvm_ptr +// CHECK: fir.coordinate_of +// CHECK: fir.load +// CHECK: cuf.allocate +// CHECK: cuf.free