[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
Wed Jul 1 15:52:39 PDT 2026


https://github.com/wangzpgi updated https://github.com/llvm/llvm-project/pull/206882

>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 1/3] 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

>From b2d61f5cb319d905ec067d5108c6775d37cb24e1 Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Tue, 30 Jun 2026 21:15:10 -0700
Subject: [PATCH 2/3] format

---
 flang/lib/Optimizer/Transforms/CUDA/CUFAllocSinking.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFAllocSinking.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFAllocSinking.cpp
index 0e18f272dead3..1b2b409f12658 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFAllocSinking.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFAllocSinking.cpp
@@ -59,8 +59,7 @@ findSinkTarget(hlfir::DeclareOp declareOp, mlir::Block *entryBlock,
     mlir::Operation *target = user;
     while (target->getBlock() != entryBlock) {
       // User in another block of the same function.
-      if (target->getBlock() &&
-          target->getBlock()->getParent() == funcRegion) {
+      if (target->getBlock() && target->getBlock()->getParent() == funcRegion) {
         mlir::Block *blk = target->getBlock();
         auto it = successorEarliest.find(blk);
         if (it == successorEarliest.end() ||

>From ff42432710fec68990a9b406636fbd1e54229aad Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Wed, 1 Jul 2026 14:35:47 -0700
Subject: [PATCH 3/3] Rename CUFAllocSinking to CUFAllocDelay, remove
 nvshmemx_init_attr reference, drop upstream pipeline insertion

---
 .../flang/Optimizer/Transforms/Passes.td      |  8 ++---
 flang/lib/Frontend/FrontendActions.cpp        |  3 --
 flang/lib/Optimizer/Transforms/CMakeLists.txt |  2 +-
 ...{CUFAllocSinking.cpp => CUFAllocDelay.cpp} | 36 +++++++++----------
 ...-alloc-sinking.fir => cuf-alloc-delay.fir} |  2 +-
 5 files changed, 24 insertions(+), 27 deletions(-)
 rename flang/lib/Optimizer/Transforms/CUDA/{CUFAllocSinking.cpp => CUFAllocDelay.cpp} (83%)
 rename flang/test/Transforms/CUF/{cuf-alloc-sinking.fir => cuf-alloc-delay.fir} (99%)

diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index 44be25e8412d4..5617e282ed5ce 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -492,13 +492,13 @@ 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";
+def CUFAllocDelay : Pass<"cuf-alloc-delay", "::mlir::func::FuncOp"> {
+  let summary = "Delay 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.
+    cudaMallocManaged calls so that users can call cudaSetDevice before any
+    CUDA context is created.
   }];
   let dependentDialects = ["fir::FIROpsDialect", "cuf::CUFDialect"];
 }
diff --git a/flang/lib/Frontend/FrontendActions.cpp b/flang/lib/Frontend/FrontendActions.cpp
index c82df49d5d4c6..5fe876595d5c0 100644
--- a/flang/lib/Frontend/FrontendActions.cpp
+++ b/flang/lib/Frontend/FrontendActions.cpp
@@ -636,9 +636,6 @@ 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 69d72cd42d58c..2d41770a40d71 100644
--- a/flang/lib/Optimizer/Transforms/CMakeLists.txt
+++ b/flang/lib/Optimizer/Transforms/CMakeLists.txt
@@ -12,7 +12,7 @@ add_flang_library(FIRTransforms
   CUDA/CUFAddConstructor.cpp
   CUDA/CUFAllocationConversion.cpp
   CUDA/CUFAllocationConversion.cpp
-  CUDA/CUFAllocSinking.cpp
+  CUDA/CUFAllocDelay.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/CUFAllocDelay.cpp
similarity index 83%
rename from flang/lib/Optimizer/Transforms/CUDA/CUFAllocSinking.cpp
rename to flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp
index 1b2b409f12658..31f261f373070 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFAllocSinking.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFAllocDelay.cpp
@@ -1,4 +1,4 @@
-//===-- CUFAllocSinking.cpp -----------------------------------------------===//
+//===-- CUFAllocDelay.cpp -------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -6,9 +6,9 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// Sink cuf.alloc of descriptor (box) types from function entry to just before
+// 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 or nvshmemx_init_attr before any CUDA context is created.
+// cudaSetDevice before any CUDA context is created.
 //
 //===----------------------------------------------------------------------===//
 
@@ -22,7 +22,7 @@
 #include "llvm/ADT/SmallVector.h"
 
 namespace fir {
-#define GEN_PASS_DEF_CUFALLOCSINKING
+#define GEN_PASS_DEF_CUFALLOCDELAY
 #include "flang/Optimizer/Transforms/Passes.h.inc"
 } // namespace fir
 
@@ -41,8 +41,8 @@ namespace {
 /// 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) {
+findDelayTarget(hlfir::DeclareOp declareOp, mlir::Block *entryBlock,
+                llvm::SmallVectorImpl<fir::StoreOp> &hostAssocStores) {
   mlir::Operation *earliest = nullptr;
   mlir::Region *funcRegion = entryBlock->getParent();
 
@@ -84,7 +84,7 @@ findSinkTarget(hlfir::DeclareOp declareOp, mlir::Block *entryBlock,
     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).
+  // 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;
@@ -93,8 +93,8 @@ findSinkTarget(hlfir::DeclareOp declareOp, mlir::Block *entryBlock,
   return nullptr;
 }
 
-struct CUFAllocSinking
-    : public fir::impl::CUFAllocSinkingBase<CUFAllocSinking> {
+struct CUFAllocDelay
+    : public fir::impl::CUFAllocDelayBase<CUFAllocDelay> {
 
   void runOnOperation() override {
     mlir::func::FuncOp func = getOperation();
@@ -129,23 +129,23 @@ struct CUFAllocSinking
         continue;
 
       llvm::SmallVector<fir::StoreOp> hostAssocStores;
-      mlir::Operation *sinkTarget =
-          findSinkTarget(declareOp, &entryBlock, hostAssocStores);
-      if (!sinkTarget)
+      mlir::Operation *delayTarget =
+          findDelayTarget(declareOp, &entryBlock, hostAssocStores);
+      if (!delayTarget)
         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))
+      if (delayTarget->getBlock() == allocOp->getBlock() &&
+          (delayTarget->isBeforeInBlock(allocOp) || delayTarget == allocOp))
         continue;
-      // Don't sink to the declare itself.
-      if (sinkTarget == declareOp)
+      // Don't move to the declare itself.
+      if (delayTarget == declareOp)
         continue;
 
-      // Move {cuf.alloc, fir.store, hlfir.declare} before the sink target.
+      // Move {cuf.alloc, fir.store, hlfir.declare} before the delay target.
       // The embox/zero_bits/shape/constants stay at their original positions
       // since they still dominate the new locations.
-      allocOp->moveBefore(sinkTarget);
+      allocOp->moveBefore(delayTarget);
       if (storeOp)
         storeOp->moveAfter(allocOp);
       if (storeOp)
diff --git a/flang/test/Transforms/CUF/cuf-alloc-sinking.fir b/flang/test/Transforms/CUF/cuf-alloc-delay.fir
similarity index 99%
rename from flang/test/Transforms/CUF/cuf-alloc-sinking.fir
rename to flang/test/Transforms/CUF/cuf-alloc-delay.fir
index 935e302544be6..92b8dcd79171e 100644
--- a/flang/test/Transforms/CUF/cuf-alloc-sinking.fir
+++ b/flang/test/Transforms/CUF/cuf-alloc-delay.fir
@@ -1,4 +1,4 @@
-// RUN: fir-opt -cuf-alloc-sinking --split-input-file %s | FileCheck %s
+// RUN: fir-opt -cuf-alloc-delay --split-input-file %s | FileCheck %s
 
 // Test 1: Basic sinking — cuf.alloc of box type is sunk to just before
 // the first use (AllocatableSetBounds).



More information about the flang-commits mailing list