[flang-commits] [flang] [flang][cuda] Add CUFAllocSinking pass to defer device descriptor allocation (PR #206882)
Zhen Wang via flang-commits
flang-commits at lists.llvm.org
Tue Jun 30 21:12:09 PDT 2026
https://github.com/wangzpgi created https://github.com/llvm/llvm-project/pull/206882
Add a new MLIR pass that sinks `cuf.alloc` operations for allocatable descriptor (box) types from function entry to just before their first use. This defers `cudaMallocManaged` calls so that programs can call `cudaSetDevice` or initialize communication libraries before any CUDA context is created.
>From eba97cd31c04813c245a6c0f121679d8bf605375 Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Tue, 30 Jun 2026 15:14:19 -0700
Subject: [PATCH] Add CUFAllocSinking pass to defer device descriptor
allocation to first use
---
.../flang/Optimizer/Transforms/Passes.td | 11 +
flang/lib/Frontend/FrontendActions.cpp | 3 +
flang/lib/Optimizer/Transforms/CMakeLists.txt | 1 +
.../Transforms/CUDA/CUFAllocSinking.cpp | 167 ++++++++
.../test/Transforms/CUF/cuf-alloc-sinking.fir | 367 ++++++++++++++++++
5 files changed, 549 insertions(+)
create mode 100644 flang/lib/Optimizer/Transforms/CUDA/CUFAllocSinking.cpp
create mode 100644 flang/test/Transforms/CUF/cuf-alloc-sinking.fir
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index e7a83f64acef9..44be25e8412d4 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -492,6 +492,17 @@ def AssumedRankOpConversion : Pass<"fir-assumed-rank-op", "mlir::ModuleOp"> {
];
}
+def CUFAllocSinking : Pass<"cuf-alloc-sinking", "::mlir::func::FuncOp"> {
+ let summary = "Sink 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
+ cudaMallocManaged calls so that users can call cudaSetDevice or
+ nvshmemx_init_attr before any CUDA context is created.
+ }];
+ let dependentDialects = ["fir::FIROpsDialect", "cuf::CUFDialect"];
+}
+
def CUFAllocationConversion : Pass<"cuf-allocation-convert", "mlir::ModuleOp"> {
let summary = "Convert allocation related CUF operations to runtime calls";
let dependentDialects = ["fir::FIROpsDialect"];
diff --git a/flang/lib/Frontend/FrontendActions.cpp b/flang/lib/Frontend/FrontendActions.cpp
index 5fe876595d5c0..c82df49d5d4c6 100644
--- a/flang/lib/Frontend/FrontendActions.cpp
+++ b/flang/lib/Frontend/FrontendActions.cpp
@@ -636,6 +636,9 @@ void CodeGenAction::lowerHLFIRToFIR() {
if (ci.getInvocation().getLangOpts().OpenMPIsTargetDevice)
config.EnableOpenMPIsTargetDevice = true;
// Create the pass pipeline
+ if (ci.getInvocation().getFortranOpts().features.IsEnabled(
+ Fortran::common::LanguageFeature::CUDA))
+ pm.addNestedPass<mlir::func::FuncOp>(fir::createCUFAllocSinking());
fir::createHLFIRToFIRPassPipeline(pm, enableOpenMP, config);
(void)mlir::applyPassManagerCLOptions(pm);
diff --git a/flang/lib/Optimizer/Transforms/CMakeLists.txt b/flang/lib/Optimizer/Transforms/CMakeLists.txt
index 5a3059ebbd97f..69d72cd42d58c 100644
--- a/flang/lib/Optimizer/Transforms/CMakeLists.txt
+++ b/flang/lib/Optimizer/Transforms/CMakeLists.txt
@@ -12,6 +12,7 @@ add_flang_library(FIRTransforms
CUDA/CUFAddConstructor.cpp
CUDA/CUFAllocationConversion.cpp
CUDA/CUFAllocationConversion.cpp
+ CUDA/CUFAllocSinking.cpp
CUDA/CUFComputeSharedMemoryOffsetsAndSize.cpp
CUDA/CUFDeviceFuncTransform.cpp
CUDA/CUFDeviceGlobal.cpp
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFAllocSinking.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFAllocSinking.cpp
new file mode 100644
index 0000000000000..0e18f272dead3
--- /dev/null
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFAllocSinking.cpp
@@ -0,0 +1,167 @@
+//===-- CUFAllocSinking.cpp -----------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Sink 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 or nvshmemx_init_attr before any CUDA context is created.
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Dialect/CUF/CUFOps.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/HLFIR/HLFIROps.h"
+#include "mlir/IR/Block.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace fir {
+#define GEN_PASS_DEF_CUFALLOCSINKING
+#include "flang/Optimizer/Transforms/Passes.h.inc"
+} // 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 *
+findSinkTarget(hlfir::DeclareOp declareOp, mlir::Block *entryBlock,
+ llvm::SmallVectorImpl<fir::StoreOp> &hostAssocStores) {
+ mlir::Operation *earliest = nullptr;
+ mlir::Region *funcRegion = entryBlock->getParent();
+
+ // Track successor-block uses: which blocks, and 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;
+ }
+ }
+ 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;
+ };
+
+ for (mlir::Value result : declareOp->getResults()) {
+ for (mlir::Operation *user : result.getUsers())
+ updateEarliest(user);
+ }
+
+ if (earliest)
+ return earliest;
+
+ // No entry-block uses. If all successor uses are in a single block,
+ // sink 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;
+}
+
+struct CUFAllocSinking
+ : public fir::impl::CUFAllocSinkingBase<CUFAllocSinking> {
+
+ void runOnOperation() override {
+ mlir::func::FuncOp func = getOperation();
+ if (func.empty())
+ return;
+
+ mlir::Block &entryBlock = func.front();
+
+ // Collect box-type cuf.alloc ops in the entry block.
+ llvm::SmallVector<cuf::AllocOp> boxAllocOps;
+ for (auto &op : entryBlock)
+ if (auto allocOp = mlir::dyn_cast<cuf::AllocOp>(op))
+ if (mlir::isa<fir::BaseBoxType>(allocOp.getInType()))
+ boxAllocOps.push_back(allocOp);
+
+ for (cuf::AllocOp allocOp : boxAllocOps) {
+ // Find the hlfir.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.
+ hlfir::DeclareOp declareOp = nullptr;
+ fir::StoreOp storeOp = nullptr;
+ bool hasUnknownUser = false;
+ for (mlir::Operation *user : allocOp->getUsers()) {
+ if (auto decl = mlir::dyn_cast<hlfir::DeclareOp>(user))
+ declareOp = decl;
+ else if (auto store = mlir::dyn_cast<fir::StoreOp>(user))
+ storeOp = store;
+ else
+ hasUnknownUser = true;
+ }
+ if (!declareOp || hasUnknownUser)
+ continue;
+
+ llvm::SmallVector<fir::StoreOp> hostAssocStores;
+ mlir::Operation *sinkTarget =
+ findSinkTarget(declareOp, &entryBlock, hostAssocStores);
+ if (!sinkTarget)
+ continue;
+
+ // Don't move if target is in the same block and at or before current pos.
+ if (sinkTarget->getBlock() == allocOp->getBlock() &&
+ (sinkTarget->isBeforeInBlock(allocOp) || sinkTarget == allocOp))
+ continue;
+ // Don't sink to the declare itself.
+ if (sinkTarget == declareOp)
+ continue;
+
+ // Move {cuf.alloc, fir.store, hlfir.declare} before the sink target.
+ // The embox/zero_bits/shape/constants stay at their original positions
+ // since they still dominate the new locations.
+ allocOp->moveBefore(sinkTarget);
+ if (storeOp)
+ storeOp->moveAfter(allocOp);
+ if (storeOp)
+ 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;
+ }
+ }
+ }
+};
+
+} // end anonymous namespace
diff --git a/flang/test/Transforms/CUF/cuf-alloc-sinking.fir b/flang/test/Transforms/CUF/cuf-alloc-sinking.fir
new file mode 100644
index 0000000000000..935e302544be6
--- /dev/null
+++ b/flang/test/Transforms/CUF/cuf-alloc-sinking.fir
@@ -0,0 +1,367 @@
+// RUN: fir-opt -cuf-alloc-sinking --split-input-file %s | FileCheck %s
+
+// Test 1: Basic sinking — cuf.alloc of box type is sunk to just before
+// the first use (AllocatableSetBounds).
+func.func @_QPbasic() {
+ %0 = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a", data_attr = #cuf.cuda<device>, uniq_name = "_QFbasicEa"} -> !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:2 = hlfir.declare %0 {data_attr = #cuf.cuda<device>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFbasicEa"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>)
+ %c10 = arith.constant 10 : index
+ %c1 = arith.constant 1 : index
+ %5 = fir.convert %c1 : (index) -> i32
+ %6 = fir.convert %c10 : (index) -> i64
+ fir.call @_FortranAAllocatableSetBounds(%4#0, %5, %6, %6) : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, i32, i64, i64) -> ()
+ %7 = cuf.allocate %4#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>} -> i32
+ return
+}
+func.func private @_FortranAAllocatableSetBounds(!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, i32, i64, i64)
+
+// CHECK-LABEL: func.func @_QPbasic()
+// The embox/zero_bits/shape stay at entry; cuf.alloc+store+declare sink to
+// just before the first use (AllocatableSetBounds).
+// CHECK: fir.zero_bits
+// CHECK: fir.shape
+// CHECK: fir.embox
+// CHECK: arith.constant 10
+// CHECK: arith.constant 1
+// CHECK: cuf.alloc
+// CHECK: fir.store
+// CHECK: hlfir.declare
+// CHECK: fir.call @_FortranAAllocatableSetBounds
+// CHECK: cuf.allocate
+
+// -----
+
+// Test 2: Non-box cuf.alloc (fixed-size device array) should NOT be sunk.
+func.func @_QPfixed_array() {
+ %0 = cuf.alloc !fir.array<10xf32> {bindc_name = "a", data_attr = #cuf.cuda<device>, uniq_name = "_QFfixed_arrayEa"} -> !fir.ref<!fir.array<10xf32>>
+ %c10 = arith.constant 10 : index
+ %1 = fir.shape %c10 : (index) -> !fir.shape<1>
+ %2:2 = hlfir.declare %0(%1) {data_attr = #cuf.cuda<device>, uniq_name = "_QFfixed_arrayEa"} : (!fir.ref<!fir.array<10xf32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xf32>>, !fir.ref<!fir.array<10xf32>>)
+ %c42 = arith.constant 42 : index
+ return
+}
+
+// CHECK-LABEL: func.func @_QPfixed_array()
+// CHECK: cuf.alloc !fir.array<10xf32>
+// CHECK: arith.constant 10
+// CHECK: fir.shape
+// CHECK: hlfir.declare
+// CHECK: arith.constant 42
+
+// -----
+
+// Test 3: First use inside fir.if — cuf.alloc should be sunk before the
+// fir.if, not inside it.
+func.func @_QPuse_in_if(%arg0: !fir.ref<!fir.logical<4>>) {
+ %0 = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a", data_attr = #cuf.cuda<device>, uniq_name = "_QFuse_in_ifEa"} -> !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:2 = hlfir.declare %0 {data_attr = #cuf.cuda<device>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFuse_in_ifEa"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>)
+ %c42 = arith.constant 42 : index
+ %5 = fir.load %arg0 : !fir.ref<!fir.logical<4>>
+ %6 = fir.convert %5 : (!fir.logical<4>) -> i1
+ fir.if %6 {
+ %7 = cuf.allocate %4#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>} -> i32
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @_QPuse_in_if
+// cuf.alloc+store+declare sink to just before fir.if (not inside it).
+// CHECK: fir.zero_bits
+// CHECK: fir.embox
+// CHECK: arith.constant 42
+// CHECK: fir.load
+// CHECK: fir.convert
+// CHECK: cuf.alloc
+// CHECK: fir.store
+// CHECK: hlfir.declare
+// CHECK: fir.if
+// CHECK: cuf.allocate
+
+// -----
+
+// Test 4: cuf.free stays at function end after sinking.
+func.func @_QPwith_free() {
+ %0 = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a", data_attr = #cuf.cuda<device>, uniq_name = "_QFwith_freeEa"} -> !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:2 = hlfir.declare %0 {data_attr = #cuf.cuda<device>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFwith_freeEa"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>)
+ %c10 = arith.constant 10 : index
+ %c1 = arith.constant 1 : index
+ %5 = fir.convert %c1 : (index) -> i32
+ %6 = fir.convert %c10 : (index) -> i64
+ fir.call @_FortranAAllocatableSetBounds(%4#0, %5, %6, %6) : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, i32, i64, i64) -> ()
+ %7 = cuf.allocate %4#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>} -> i32
+ cuf.free %4#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>}
+ return
+}
+func.func private @_FortranAAllocatableSetBounds(!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, i32, i64, i64)
+
+// CHECK-LABEL: func.func @_QPwith_free()
+// CHECK: fir.zero_bits
+// CHECK: fir.embox
+// CHECK: cuf.alloc
+// CHECK: fir.store
+// CHECK: hlfir.declare
+// CHECK: fir.call @_FortranAAllocatableSetBounds
+// CHECK: cuf.allocate
+// CHECK: cuf.free
+// CHECK: return
+
+// -----
+
+// Test 5: Multiple allocatable descriptors sink independently.
+func.func @_QPmultiple() {
+ %0 = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a", data_attr = #cuf.cuda<device>, uniq_name = "_QFmultipleEa"} -> !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:2 = hlfir.declare %0 {data_attr = #cuf.cuda<device>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFmultipleEa"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>)
+ %10 = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf64>>> {bindc_name = "b", data_attr = #cuf.cuda<device>, uniq_name = "_QFmultipleEb"} -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>
+ %11 = fir.zero_bits !fir.heap<!fir.array<?xf64>>
+ %12 = fir.shape %c0 : (index) -> !fir.shape<1>
+ %13 = fir.embox %11(%12) {allocator_idx = 2 : i32} : (!fir.heap<!fir.array<?xf64>>, !fir.shape<1>) -> !fir.box<!fir.heap<!fir.array<?xf64>>>
+ fir.store %13 to %10 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>
+ %14:2 = hlfir.declare %10 {data_attr = #cuf.cuda<device>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFmultipleEb"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>)
+ %c42 = arith.constant 42 : index
+ %20 = cuf.allocate %14#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>> {data_attr = #cuf.cuda<device>} -> i32
+ %c99 = arith.constant 99 : index
+ %21 = cuf.allocate %4#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>} -> i32
+ cuf.free %14#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>> {data_attr = #cuf.cuda<device>}
+ cuf.free %4#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>}
+ return
+}
+
+// CHECK-LABEL: func.func @_QPmultiple()
+// Both allocs sink past the constants to just before their first real use.
+// "b" is used first (cuf.allocate %b), "a" is used later (cuf.allocate %a).
+// CHECK: fir.zero_bits
+// CHECK: fir.embox
+// CHECK: fir.zero_bits
+// CHECK: fir.embox
+// CHECK: arith.constant 42
+// CHECK: cuf.alloc {{.*}}uniq_name = "_QFmultipleEb"
+// CHECK: fir.store
+// CHECK: hlfir.declare {{.*}}uniq_name = "_QFmultipleEb"
+// CHECK: cuf.allocate
+// CHECK: arith.constant 99
+// CHECK: cuf.alloc {{.*}}uniq_name = "_QFmultipleEa"
+// CHECK: fir.store
+// CHECK: hlfir.declare {{.*}}uniq_name = "_QFmultipleEa"
+// CHECK: cuf.allocate
+// CHECK: cuf.free
+// CHECK: cuf.free
+// CHECK: return
+
+// -----
+
+// Test 6: Dead descriptor (declare with only cuf.free user) — still sinks
+// to just before the free.
+func.func @_QPdead_var() {
+ %0 = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a", data_attr = #cuf.cuda<device>, uniq_name = "_QFdead_varEa"} -> !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:2 = hlfir.declare %0 {data_attr = #cuf.cuda<device>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFdead_varEa"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>)
+ %c42 = arith.constant 42 : index
+ cuf.free %4#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>}
+ return
+}
+
+// CHECK-LABEL: func.func @_QPdead_var()
+// CHECK: fir.zero_bits
+// CHECK: fir.embox
+// CHECK: arith.constant 42
+// CHECK: cuf.alloc
+// CHECK: fir.store
+// CHECK: hlfir.declare
+// CHECK: cuf.free
+// CHECK: return
+
+// -----
+
+// Test 7: First use inside fir.do_loop — cuf.alloc sinks before the loop.
+func.func @_QPuse_in_loop() {
+ %0 = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a", data_attr = #cuf.cuda<device>, uniq_name = "_QFuse_in_loopEa"} -> !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:2 = hlfir.declare %0 {data_attr = #cuf.cuda<device>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFuse_in_loopEa"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>)
+ %c99 = arith.constant 99 : index
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ fir.do_loop %i = %c1 to %c10 step %c1 {
+ %5 = fir.load %4#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ }
+ cuf.free %4#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>}
+ return
+}
+
+// CHECK-LABEL: func.func @_QPuse_in_loop()
+// CHECK: fir.zero_bits
+// CHECK: fir.embox
+// CHECK: arith.constant 99
+// CHECK: arith.constant 1
+// CHECK: arith.constant 10
+// CHECK: cuf.alloc
+// CHECK: fir.store
+// CHECK: hlfir.declare
+// CHECK: fir.do_loop
+// CHECK: fir.load
+// CHECK: cuf.free
+// CHECK: return
+
+// -----
+
+// Test 8: All uses in a successor block (e.g. after a branch from
+// stop-in-if). The cuf.alloc should sink to just before the entry
+// block's terminator so it still dominates the successor.
+func.func @_QPuse_in_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_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:2 = hlfir.declare %0 {data_attr = #cuf.cuda<device>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFuse_in_successorEa"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>)
+ %c42 = arith.constant 42 : 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:
+ %7 = cuf.allocate %4#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>} -> i32
+ cuf.free %4#0 : !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 @_FortranAAllocatableSetBounds(!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, i32, i64, i64)
+
+// CHECK-LABEL: func.func @_QPuse_in_successor
+// All uses are in ^bb2 only, so cuf.alloc+store+declare sink directly
+// into ^bb2 before the first use.
+// CHECK: fir.zero_bits
+// CHECK: fir.embox
+// CHECK: arith.constant 42
+// CHECK: fir.load
+// CHECK: arith.cmpi
+// CHECK: cf.cond_br
+// CHECK: ^bb1:
+// CHECK: fir.unreachable
+// CHECK: ^bb2:
+// CHECK: cuf.alloc
+// CHECK: fir.store
+// CHECK: hlfir.declare
+// CHECK: cuf.allocate
+// CHECK: cuf.free
+// CHECK: return
+
+// -----
+
+// Test 9: Uses in multiple successor blocks — cuf.alloc sinks to just
+// before the entry block's terminator (the branch) so it dominates both.
+func.func @_QPuse_in_multi_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_multi_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:2 = hlfir.declare %0 {data_attr = #cuf.cuda<device>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFuse_in_multi_successorEa"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>)
+ %c42 = arith.constant 42 : 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:
+ %7 = fir.load %4#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+ cuf.free %4#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>}
+ return
+^bb2:
+ %8 = cuf.allocate %4#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>} -> i32
+ cuf.free %4#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>}
+ return
+}
+
+// CHECK-LABEL: func.func @_QPuse_in_multi_successor
+// Uses span ^bb1 and ^bb2 — alloc sinks to before the entry block's
+// terminator (cf.cond_br).
+// CHECK: fir.zero_bits
+// CHECK: fir.embox
+// CHECK: arith.constant 42
+// CHECK: fir.load
+// CHECK: arith.constant 0
+// CHECK: arith.cmpi
+// CHECK: cuf.alloc
+// CHECK: fir.store
+// CHECK: hlfir.declare
+// CHECK: cf.cond_br
+// CHECK: ^bb1:
+// CHECK: fir.load
+// CHECK: cuf.free
+// CHECK: ^bb2:
+// CHECK: cuf.allocate
+// CHECK: cuf.free
+
+// -----
+
+// Test 10: Host-association tuple store — fir.store to fir.llvm_ptr
+// should not block sinking and should move along with the sunk group.
+func.func @_QPhost_assoc() {
+ %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_assocEa"} -> !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:2 = hlfir.declare %0 {data_attr = #cuf.cuda<device>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFhost_assocEa"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>)
+ fir.store %4#0 to %slot : !fir.llvm_ptr<!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>
+ %c99 = arith.constant 99 : index
+ %5 = cuf.allocate %4#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>} -> i32
+ fir.call @_QFPcontained(%tuple) : (!fir.ref<tuple<!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>>>) -> ()
+ cuf.free %4#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>}
+ return
+}
+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 sinking.
+// cuf.alloc+store+declare+host-assoc-store all sink 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: hlfir.declare
+// CHECK: fir.store {{.*}} : !fir.llvm_ptr
+// CHECK: cuf.allocate
+// CHECK: fir.call @_QFPcontained
+// CHECK: cuf.free
+// CHECK: return
More information about the flang-commits
mailing list