[flang-commits] [flang] [flang][cuda] Route dynamic autos through malloc_unified/free_unified (PR #212965)
via flang-commits
flang-commits at lists.llvm.org
Wed Aug 19 17:06:32 PDT 2026
https://github.com/khaki3 updated https://github.com/llvm/llvm-project/pull/212965
>From 05d1737bf2bb3dbf382b56f09b0031de0ae77263 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Thu, 30 Jul 2026 01:29:13 -0700
Subject: [PATCH 1/7] [flang][cuda] Allocate adjustable autos in
unified/managed memory
Under -gpu=mem:unified|managed, tag dynamic-extent automatic arrays so
lowering uses cuf.alloc/cuf.free; fixed-size automatic arrays stay on
the stack.
---
flang/lib/Semantics/resolve-names.cpp | 31 ++++++++++
.../CUDA/cuda-gpu-unified-automatic-array.f90 | 58 +++++++++++++++++++
2 files changed, 89 insertions(+)
create mode 100644 flang/test/Lower/CUDA/cuda-gpu-unified-automatic-array.f90
diff --git a/flang/lib/Semantics/resolve-names.cpp b/flang/lib/Semantics/resolve-names.cpp
index 27c4e96d269aa..d188f778a31bd 100644
--- a/flang/lib/Semantics/resolve-names.cpp
+++ b/flang/lib/Semantics/resolve-names.cpp
@@ -10775,6 +10775,37 @@ void ResolveNamesVisitor::FinishSpecificationPart(
context().languageFeatures().IsEnabled(
common::LanguageFeature::CudaPinned))
object->set_cudaDataAttr(common::CUDADataAttr::Pinned);
+ } else if (!object->cudaDataAttr() && !IsDummy(symbol) &&
+ !IsAllocatable(symbol) && !IsPointer(symbol) && !IsSaved(symbol) &&
+ !IsCUDADeviceContext(&symbol.owner()) &&
+ object->shape().IsExplicitShape()) {
+ // Under -gpu=mem:unified|managed, allocate adjustable / VLA automatic
+ // arrays in CUDA unified/managed memory (fixed-size automatic arrays
+ // stay on the stack). Tag those locals so lowering uses
+ // cuf.alloc/cuf.free. Unlike the allocatable managed tagging above,
+ // this does not require -fcuda: OpenACC + -gpu=mem:unified relies on
+ // it, and cuf.alloc does not go through the CUDA Fortran
+ // managed-descriptor pipeline that motivated the -fcuda gate.
+ auto boundIsNonConstant{[](const Bound &b) {
+ return !b.isExplicit() || !b.GetExplicit() ||
+ !evaluate::IsConstantExpr(*b.GetExplicit());
+ }};
+ bool hasDynamicExtent{false};
+ for (const ShapeSpec &ss : object->shape()) {
+ if (boundIsNonConstant(ss.lbound()) ||
+ boundIsNonConstant(ss.ubound())) {
+ hasDynamicExtent = true;
+ break;
+ }
+ }
+ if (hasDynamicExtent) {
+ if (context().languageFeatures().IsEnabled(
+ common::LanguageFeature::CudaUnified))
+ object->set_cudaDataAttr(common::CUDADataAttr::Unified);
+ else if (context().languageFeatures().IsEnabled(
+ common::LanguageFeature::CudaManaged))
+ object->set_cudaDataAttr(common::CUDADataAttr::Managed);
+ }
}
}
}
diff --git a/flang/test/Lower/CUDA/cuda-gpu-unified-automatic-array.f90 b/flang/test/Lower/CUDA/cuda-gpu-unified-automatic-array.f90
new file mode 100644
index 0000000000000..07638b0ca8dfa
--- /dev/null
+++ b/flang/test/Lower/CUDA/cuda-gpu-unified-automatic-array.f90
@@ -0,0 +1,58 @@
+! RUN: bbc -emit-hlfir -gpu=unified %s -o - | FileCheck %s
+! RUN: bbc -emit-hlfir -gpu=managed %s -o - | FileCheck %s --check-prefix=MANAGED
+
+! Under -gpu=mem:unified|managed, allocate adjustable / VLA automatic arrays
+! in CUDA unified/managed memory. Fixed-size automatic arrays stay on the
+! stack so host pointers remain shared under unified memory (e.g. OpenACC).
+
+module m_adj
+ integer :: nx = 32
+end module
+
+! CHECK-LABEL: func.func @_QPvla(
+! CHECK: %[[ALLOC:.*]] = cuf.alloc !fir.array<?xf32>, %{{.*}} : index {bindc_name = "a", data_attr = #cuf.cuda<unified>, uniq_name = "_QFvlaEa"}
+! CHECK: %[[DECL:.*]]:2 = hlfir.declare %[[ALLOC]](%{{.*}}) {data_attr = #cuf.cuda<unified>, uniq_name = "_QFvlaEa"}
+! CHECK: cuf.free %[[DECL]]#1 : !fir.ref<!fir.array<?xf32>> {data_attr = #cuf.cuda<unified>}
+! MANAGED-LABEL: func.func @_QPvla(
+! MANAGED: cuf.alloc !fir.array<?xf32>, %{{.*}} : index {{{.*}}data_attr = #cuf.cuda<managed>
+! MANAGED: cuf.free %{{.*}} : !fir.ref<!fir.array<?xf32>> {data_attr = #cuf.cuda<managed>}
+subroutine vla(n)
+ integer :: n
+ real :: a(n)
+ a(1) = 1.0
+end subroutine
+
+! CHECK-LABEL: func.func @_QPadjustable(
+! CHECK: cuf.alloc !fir.array<?xf32>, %{{.*}} : index {{{.*}}data_attr = #cuf.cuda<unified>
+! CHECK: cuf.free %{{.*}} {data_attr = #cuf.cuda<unified>}
+! MANAGED-LABEL: func.func @_QPadjustable(
+! MANAGED: cuf.alloc !fir.array<?xf32>, %{{.*}} : index {{{.*}}data_attr = #cuf.cuda<managed>
+subroutine adjustable
+ use m_adj
+ real :: a(0:(nx+1)/2)
+ a(0) = 0.0
+end subroutine
+
+! Fixed-size automatic arrays must remain ordinary stack allocations.
+! CHECK-LABEL: func.func @_QPfixed(
+! CHECK-NOT: cuf.alloc
+! CHECK: fir.alloca !fir.array<128xf32>
+! CHECK-NOT: cuf.free
+! MANAGED-LABEL: func.func @_QPfixed(
+! MANAGED-NOT: cuf.alloc
+! MANAGED: fir.alloca !fir.array<128xf32>
+subroutine fixed
+ real :: a(128)
+ a(1) = 1.0
+end subroutine
+
+! Dummy adjustable arrays are caller-allocated; do not retag them.
+! CHECK-LABEL: func.func @_QPdummy_adj(
+! CHECK-NOT: cuf.alloc
+! MANAGED-LABEL: func.func @_QPdummy_adj(
+! MANAGED-NOT: cuf.alloc
+subroutine dummy_adj(a, n)
+ integer :: n
+ real :: a(n)
+ a(1) = 1.0
+end subroutine
>From 1335398ee485b250ebcf4bfc4a71d18e0119b16b Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Thu, 30 Jul 2026 01:38:51 -0700
Subject: [PATCH 2/7] [flang][cuda] Check unified/managed features before shape
walk
---
flang/lib/Semantics/resolve-names.cpp | 36 +++++++++++++--------------
1 file changed, 18 insertions(+), 18 deletions(-)
diff --git a/flang/lib/Semantics/resolve-names.cpp b/flang/lib/Semantics/resolve-names.cpp
index d188f778a31bd..dd6a5401a8250 100644
--- a/flang/lib/Semantics/resolve-names.cpp
+++ b/flang/lib/Semantics/resolve-names.cpp
@@ -10786,26 +10786,26 @@ void ResolveNamesVisitor::FinishSpecificationPart(
// this does not require -fcuda: OpenACC + -gpu=mem:unified relies on
// it, and cuf.alloc does not go through the CUDA Fortran
// managed-descriptor pipeline that motivated the -fcuda gate.
- auto boundIsNonConstant{[](const Bound &b) {
- return !b.isExplicit() || !b.GetExplicit() ||
- !evaluate::IsConstantExpr(*b.GetExplicit());
- }};
- bool hasDynamicExtent{false};
- for (const ShapeSpec &ss : object->shape()) {
- if (boundIsNonConstant(ss.lbound()) ||
- boundIsNonConstant(ss.ubound())) {
- hasDynamicExtent = true;
- break;
+ std::optional<common::CUDADataAttr> attr;
+ if (context().languageFeatures().IsEnabled(
+ common::LanguageFeature::CudaUnified))
+ attr = common::CUDADataAttr::Unified;
+ else if (context().languageFeatures().IsEnabled(
+ common::LanguageFeature::CudaManaged))
+ attr = common::CUDADataAttr::Managed;
+ if (attr) {
+ auto boundIsNonConstant{[](const Bound &b) {
+ return !b.isExplicit() || !b.GetExplicit() ||
+ !evaluate::IsConstantExpr(*b.GetExplicit());
+ }};
+ for (const ShapeSpec &ss : object->shape()) {
+ if (boundIsNonConstant(ss.lbound()) ||
+ boundIsNonConstant(ss.ubound())) {
+ object->set_cudaDataAttr(*attr);
+ break;
+ }
}
}
- if (hasDynamicExtent) {
- if (context().languageFeatures().IsEnabled(
- common::LanguageFeature::CudaUnified))
- object->set_cudaDataAttr(common::CUDADataAttr::Unified);
- else if (context().languageFeatures().IsEnabled(
- common::LanguageFeature::CudaManaged))
- object->set_cudaDataAttr(common::CUDADataAttr::Managed);
- }
}
}
}
>From 4a978ae8ffb384edff72c3742256a69a9288e19d Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Thu, 30 Jul 2026 13:48:45 -0700
Subject: [PATCH 3/7] [flang][cuda] Route dynamic autos through
malloc_unified/free_unified
Drop implicit cudaDataAttr tagging for adjustable/VLA automatic arrays.
Under -gpu=mem:unified|managed, promote those locals to fir.allocmem and
lower host heap allocations via the indirect unified/managed allocators.
---
.../Optimizer/Dialect/Support/FIRContext.h | 14 +++
.../flang/Optimizer/Transforms/MemoryUtils.h | 7 ++
flang/lib/Lower/Bridge.cpp | 7 ++
flang/lib/Optimizer/CodeGen/CodeGen.cpp | 84 +++++++++++++-----
.../Optimizer/Dialect/Support/FIRContext.cpp | 41 +++++++++
.../Transforms/AllocationPlacement.cpp | 5 ++
.../Optimizer/Transforms/MemoryAllocation.cpp | 3 +
.../lib/Optimizer/Transforms/MemoryUtils.cpp | 65 ++++++++++++++
.../lib/Optimizer/Transforms/StackArrays.cpp | 7 ++
flang/lib/Semantics/resolve-names.cpp | 31 -------
.../test/Fir/CUDA/cuda-heap-alloc-managed.fir | 44 ++++++++++
.../test/Fir/CUDA/cuda-heap-alloc-unified.fir | 86 +++++++++++++++++++
.../CUDA/cuda-gpu-unified-automatic-array.f90 | 42 ++++-----
13 files changed, 360 insertions(+), 76 deletions(-)
create mode 100644 flang/test/Fir/CUDA/cuda-heap-alloc-managed.fir
create mode 100644 flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
diff --git a/flang/include/flang/Optimizer/Dialect/Support/FIRContext.h b/flang/include/flang/Optimizer/Dialect/Support/FIRContext.h
index 79337584d6d67..27eee3a101e29 100644
--- a/flang/include/flang/Optimizer/Dialect/Support/FIRContext.h
+++ b/flang/include/flang/Optimizer/Dialect/Support/FIRContext.h
@@ -121,6 +121,20 @@ void setIsPIE(mlir::ModuleOp mod, bool value);
/// Get whether the module is compiled as a position-independent executable.
bool getIsPIE(mlir::ModuleOp mod);
+/// Host heap allocator selected under -gpu=mem:unified|managed, recorded on the
+/// module by lowering and consumed by the allocation placement passes.
+enum class CudaHeapAllocMode { None, Unified, Managed };
+
+void setCudaHeapAllocMode(mlir::ModuleOp mod, CudaHeapAllocMode mode);
+CudaHeapAllocMode getCudaHeapAllocMode(mlir::ModuleOp mod);
+
+/// Same attribute on a fir.allocmem/fir.freemem pair: this allocation uses the
+/// indirect runtime entry points (`malloc_unified`/`free_unified`, ...) instead
+/// of libc. Only pairs created together may be marked, since the allocator and
+/// the deallocator must match.
+void setCudaHeapAllocMode(mlir::Operation *op, CudaHeapAllocMode mode);
+CudaHeapAllocMode getCudaHeapAllocMode(mlir::Operation *op);
+
/// Helper for determining the target from the host, etc. Tools may use this
/// function to provide a consistent interpretation of the `--target=<string>`
/// command-line option.
diff --git a/flang/include/flang/Optimizer/Transforms/MemoryUtils.h b/flang/include/flang/Optimizer/Transforms/MemoryUtils.h
index 92a519cd0c838..62438c7c44dd3 100644
--- a/flang/include/flang/Optimizer/Transforms/MemoryUtils.h
+++ b/flang/include/flang/Optimizer/Transforms/MemoryUtils.h
@@ -57,6 +57,13 @@ bool replaceAllocas(mlir::RewriterBase &rewriter, mlir::Operation *parentOp,
MustRewriteCallBack, AllocaRewriterCallBack,
DeallocCallBack);
+/// Under -gpu=mem:unified|managed, move the dynamically sized fir.alloca of
+/// \p func to fir.allocmem/fir.freemem pairs marked for the unified/managed
+/// allocator. Does nothing for device code, which keeps its stack allocations.
+/// Returns true if the function was modified.
+bool promoteDynamicAllocasToCudaHeap(mlir::RewriterBase &rewriter,
+ mlir::Operation *func);
+
} // namespace fir
#endif // FORTRAN_OPTIMIZER_TRANSFORMS_MEMORYUTILS_H
diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index c52b8bd67a111..febcea8a63dac 100644
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -6898,6 +6898,13 @@ Fortran::lower::LoweringBridge::LoweringBridge(
fir::setIsPIE(*module, cgOpts.IsPIE);
if (cgOpts.RecordCommandLine)
fir::setCommandline(*module, *cgOpts.RecordCommandLine);
+ // Under -gpu=mem:unified|managed, host heap allocations use the matching
+ // indirect runtime allocators (malloc_unified / malloc_managed).
+ if (languageFeatures.IsEnabled(Fortran::common::LanguageFeature::CudaUnified))
+ fir::setCudaHeapAllocMode(*module, fir::CudaHeapAllocMode::Unified);
+ else if (languageFeatures.IsEnabled(
+ Fortran::common::LanguageFeature::CudaManaged))
+ fir::setCudaHeapAllocMode(*module, fir::CudaHeapAllocMode::Managed);
}
Fortran::lower::LoweringBridge::~LoweringBridge() {
diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
index 4e33dc008e53a..8108b19a4981c 100644
--- a/flang/lib/Optimizer/CodeGen/CodeGen.cpp
+++ b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -22,6 +22,7 @@
#include "flang/Optimizer/Dialect/FIRDialect.h"
#include "flang/Optimizer/Dialect/FIROps.h"
#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/Dialect/Support/FIRContext.h"
#include "flang/Optimizer/Support/DataLayout.h"
#include "flang/Optimizer/Support/InternalNames.h"
#include "flang/Optimizer/Support/TypeCode.h"
@@ -1293,8 +1294,7 @@ template <typename ModuleOp>
static mlir::SymbolRefAttr
getMallocInModule(ModuleOp mod, fir::AllocMemOp op,
mlir::ConversionPatternRewriter &rewriter,
- mlir::Type indexType) {
- static constexpr char mallocName[] = "malloc";
+ mlir::Type indexType, llvm::StringRef mallocName) {
if (auto mallocFunc =
mod.template lookupSymbol<mlir::LLVM::LLVMFuncOp>(mallocName))
return mlir::SymbolRefAttr::get(mallocFunc);
@@ -1311,22 +1311,61 @@ getMallocInModule(ModuleOp mod, fir::AllocMemOp op,
return mlir::SymbolRefAttr::get(mallocDecl);
}
+/// Allocator entry points for an allocation marked by the allocation placement
+/// passes under -gpu=mem:unified|managed. Only marked fir.allocmem/fir.freemem
+/// pairs are routed: memory that the Fortran runtime allocated must keep being
+/// released by libc free, and vice versa.
+static llvm::StringRef getHeapAllocName(mlir::Operation *op,
+ llvm::StringRef plain,
+ llvm::StringRef unified,
+ llvm::StringRef managed) {
+ // Device modules keep libc names; the indirect entry points are host-side.
+ if (op->getParentOfType<mlir::gpu::GPUModuleOp>())
+ return plain;
+ switch (fir::getCudaHeapAllocMode(op)) {
+ case fir::CudaHeapAllocMode::Unified:
+ return unified;
+ case fir::CudaHeapAllocMode::Managed:
+ return managed;
+ case fir::CudaHeapAllocMode::None:
+ return plain;
+ }
+ llvm_unreachable("unexpected CudaHeapAllocMode");
+}
+
+static llvm::StringRef getHostHeapMallocName(mlir::Operation *op) {
+ return getHeapAllocName(op, "malloc", "malloc_unified", "malloc_managed");
+}
+
+static llvm::StringRef getHostHeapFreeName(mlir::Operation *op) {
+ return getHeapAllocName(op, "free", "free_unified", "free_managed");
+}
+
+static llvm::StringRef getHostHeapAlignedAllocName(mlir::Operation *op) {
+ return getHeapAllocName(op, "aligned_alloc", "aligned_alloc_unified",
+ "aligned_alloc_managed");
+}
+
+static llvm::StringRef getHostHeapPosixMemalignName(mlir::Operation *op) {
+ return getHeapAllocName(op, "posix_memalign", "posix_memalign_unified",
+ "posix_memalign_managed");
+}
+
/// Return the LLVMFuncOp corresponding to the standard malloc call.
static mlir::SymbolRefAttr getMalloc(fir::AllocMemOp op,
mlir::ConversionPatternRewriter &rewriter,
mlir::Type indexType) {
+ llvm::StringRef name = getHostHeapMallocName(op);
if (auto mod = op->getParentOfType<mlir::gpu::GPUModuleOp>())
- return getMallocInModule(mod, op, rewriter, indexType);
+ return getMallocInModule(mod, op, rewriter, indexType, name);
auto mod = op->getParentOfType<mlir::ModuleOp>();
- return getMallocInModule(mod, op, rewriter, indexType);
+ return getMallocInModule(mod, op, rewriter, indexType, name);
}
template <typename ModuleOp>
-static mlir::SymbolRefAttr
-getAlignedAllocInModule(ModuleOp mod, fir::AllocMemOp op,
- mlir::ConversionPatternRewriter &rewriter,
- mlir::Type indexType) {
- static constexpr char alignedAllocName[] = "aligned_alloc";
+static mlir::SymbolRefAttr getAlignedAllocInModule(
+ ModuleOp mod, fir::AllocMemOp op, mlir::ConversionPatternRewriter &rewriter,
+ mlir::Type indexType, llvm::StringRef alignedAllocName) {
if (auto func =
mod.template lookupSymbol<mlir::LLVM::LLVMFuncOp>(alignedAllocName))
return mlir::SymbolRefAttr::get(func);
@@ -1346,18 +1385,17 @@ getAlignedAllocInModule(ModuleOp mod, fir::AllocMemOp op,
static mlir::SymbolRefAttr
getAlignedAlloc(fir::AllocMemOp op, mlir::ConversionPatternRewriter &rewriter,
mlir::Type indexType) {
+ llvm::StringRef name = getHostHeapAlignedAllocName(op);
if (auto mod = op->getParentOfType<mlir::gpu::GPUModuleOp>())
- return getAlignedAllocInModule(mod, op, rewriter, indexType);
+ return getAlignedAllocInModule(mod, op, rewriter, indexType, name);
auto mod = op->getParentOfType<mlir::ModuleOp>();
- return getAlignedAllocInModule(mod, op, rewriter, indexType);
+ return getAlignedAllocInModule(mod, op, rewriter, indexType, name);
}
template <typename ModuleOp>
-static mlir::SymbolRefAttr
-getPosixMemalignInModule(ModuleOp mod, fir::AllocMemOp op,
- mlir::ConversionPatternRewriter &rewriter,
- mlir::Type indexType) {
- static constexpr char posixMemalignName[] = "posix_memalign";
+static mlir::SymbolRefAttr getPosixMemalignInModule(
+ ModuleOp mod, fir::AllocMemOp op, mlir::ConversionPatternRewriter &rewriter,
+ mlir::Type indexType, llvm::StringRef posixMemalignName) {
if (auto func =
mod.template lookupSymbol<mlir::LLVM::LLVMFuncOp>(posixMemalignName))
return mlir::SymbolRefAttr::get(func);
@@ -1379,10 +1417,11 @@ getPosixMemalignInModule(ModuleOp mod, fir::AllocMemOp op,
static mlir::SymbolRefAttr
getPosixMemalign(fir::AllocMemOp op, mlir::ConversionPatternRewriter &rewriter,
mlir::Type indexType) {
+ llvm::StringRef name = getHostHeapPosixMemalignName(op);
if (auto mod = op->getParentOfType<mlir::gpu::GPUModuleOp>())
- return getPosixMemalignInModule(mod, op, rewriter, indexType);
+ return getPosixMemalignInModule(mod, op, rewriter, indexType, name);
auto mod = op->getParentOfType<mlir::ModuleOp>();
- return getPosixMemalignInModule(mod, op, rewriter, indexType);
+ return getPosixMemalignInModule(mod, op, rewriter, indexType, name);
}
/// Return value of the stride in bytes between adjacent elements
@@ -1519,8 +1558,8 @@ struct AllocMemOpConversion : public fir::FIROpConversion<fir::AllocMemOp> {
template <typename ModuleOp>
static mlir::SymbolRefAttr
getFreeInModule(ModuleOp mod, fir::FreeMemOp op,
- mlir::ConversionPatternRewriter &rewriter) {
- static constexpr char freeName[] = "free";
+ mlir::ConversionPatternRewriter &rewriter,
+ llvm::StringRef freeName) {
// Check if free already defined in the module.
if (auto freeFunc =
mod.template lookupSymbol<mlir::LLVM::LLVMFuncOp>(freeName))
@@ -1541,10 +1580,11 @@ getFreeInModule(ModuleOp mod, fir::FreeMemOp op,
static mlir::SymbolRefAttr getFree(fir::FreeMemOp op,
mlir::ConversionPatternRewriter &rewriter) {
+ llvm::StringRef name = getHostHeapFreeName(op);
if (auto mod = op->getParentOfType<mlir::gpu::GPUModuleOp>())
- return getFreeInModule(mod, op, rewriter);
+ return getFreeInModule(mod, op, rewriter, name);
auto mod = op->getParentOfType<mlir::ModuleOp>();
- return getFreeInModule(mod, op, rewriter);
+ return getFreeInModule(mod, op, rewriter, name);
}
static unsigned getDimension(mlir::LLVM::LLVMArrayType ty) {
diff --git a/flang/lib/Optimizer/Dialect/Support/FIRContext.cpp b/flang/lib/Optimizer/Dialect/Support/FIRContext.cpp
index 16757f934c8d7..66ae5c4c653ac 100644
--- a/flang/lib/Optimizer/Dialect/Support/FIRContext.cpp
+++ b/flang/lib/Optimizer/Dialect/Support/FIRContext.cpp
@@ -248,6 +248,47 @@ void fir::setIsPIE(mlir::ModuleOp mod, bool value) {
bool fir::getIsPIE(mlir::ModuleOp mod) { return mod->hasAttr(isPIEName); }
+static constexpr const char *cudaHeapAllocModeName = "fir.cuda_heap_alloc";
+
+static void setCudaHeapAllocModeOn(mlir::Operation *op,
+ fir::CudaHeapAllocMode mode) {
+ if (mode == fir::CudaHeapAllocMode::None) {
+ if (op->hasAttr(cudaHeapAllocModeName))
+ op->removeAttr(cudaHeapAllocModeName);
+ return;
+ }
+ llvm::StringRef value =
+ mode == fir::CudaHeapAllocMode::Unified ? "unified" : "managed";
+ op->setAttr(cudaHeapAllocModeName,
+ mlir::StringAttr::get(op->getContext(), value));
+}
+
+static fir::CudaHeapAllocMode getCudaHeapAllocModeOf(mlir::Operation *op) {
+ if (auto attr = op->getAttrOfType<mlir::StringAttr>(cudaHeapAllocModeName)) {
+ if (attr.getValue() == "unified")
+ return fir::CudaHeapAllocMode::Unified;
+ if (attr.getValue() == "managed")
+ return fir::CudaHeapAllocMode::Managed;
+ }
+ return fir::CudaHeapAllocMode::None;
+}
+
+void fir::setCudaHeapAllocMode(mlir::ModuleOp mod, CudaHeapAllocMode mode) {
+ setCudaHeapAllocModeOn(mod.getOperation(), mode);
+}
+
+fir::CudaHeapAllocMode fir::getCudaHeapAllocMode(mlir::ModuleOp mod) {
+ return getCudaHeapAllocModeOf(mod.getOperation());
+}
+
+void fir::setCudaHeapAllocMode(mlir::Operation *op, CudaHeapAllocMode mode) {
+ setCudaHeapAllocModeOn(op, mode);
+}
+
+fir::CudaHeapAllocMode fir::getCudaHeapAllocMode(mlir::Operation *op) {
+ return getCudaHeapAllocModeOf(op);
+}
+
std::string fir::determineTargetTriple(llvm::StringRef triple) {
// Treat "" or "default" as stand-ins for the default machine.
if (triple.empty() || triple == "default")
diff --git a/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp b/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
index c8ad34f7f60ad..231d381649f95 100644
--- a/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
+++ b/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
@@ -199,6 +199,11 @@ void AllocationPlacementPass::runOnOperation() {
if (func.empty())
return;
+ // Done first: the pairs it creates are marked fir.must_be_heap, so the
+ // placement decisions below leave them alone.
+ mlir::IRRewriter cudaHeapRewriter(&getContext());
+ fir::promoteDynamicAllocasToCudaHeap(cudaHeapRewriter, func.getOperation());
+
fir::AllocationPlacementThresholds baseThresholds;
baseThresholds.stackArrays = stackArrays;
baseThresholds.smallArrayThresholdBytes = smallArrayThresholdBytes;
diff --git a/flang/lib/Optimizer/Transforms/MemoryAllocation.cpp b/flang/lib/Optimizer/Transforms/MemoryAllocation.cpp
index fd1d566ca2825..6c4342b141050 100644
--- a/flang/lib/Optimizer/Transforms/MemoryAllocation.cpp
+++ b/flang/lib/Optimizer/Transforms/MemoryAllocation.cpp
@@ -9,6 +9,7 @@
#include "flang/Optimizer/Dialect/FIRDialect.h"
#include "flang/Optimizer/Dialect/FIROps.h"
#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/Dialect/Support/FIRContext.h"
#include "flang/Optimizer/Transforms/MemoryUtils.h"
#include "flang/Optimizer/Transforms/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
@@ -131,6 +132,8 @@ class MemoryAllocationOpt
// If func is a declaration, skip it.
if (func.empty())
return;
+ mlir::IRRewriter cudaHeapRewriter(context);
+ fir::promoteDynamicAllocasToCudaHeap(cudaHeapRewriter, func.getOperation());
auto tryReplacing = [&](fir::AllocaOp alloca) {
bool res = !keepStackAllocation(alloca, options);
if (res) {
diff --git a/flang/lib/Optimizer/Transforms/MemoryUtils.cpp b/flang/lib/Optimizer/Transforms/MemoryUtils.cpp
index d1b27d9872ea1..0e5155dd0a361 100644
--- a/flang/lib/Optimizer/Transforms/MemoryUtils.cpp
+++ b/flang/lib/Optimizer/Transforms/MemoryUtils.cpp
@@ -8,7 +8,12 @@
#include "flang/Optimizer/Transforms/MemoryUtils.h"
#include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h"
+#include "flang/Optimizer/Dialect/FIRAttr.h"
+#include "flang/Optimizer/Dialect/Support/FIRContext.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Dominance.h"
#include "llvm/ADT/STLExtras.h"
@@ -309,3 +314,63 @@ bool fir::replaceAllocas(mlir::RewriterBase &rewriter,
rewriter.restoreInsertionPoint(insertPoint);
return replacedAllRequestedAlloca;
}
+
+/// Device code keeps its stack allocations: the unified/managed entry points
+/// are host-only, and a kernel-side heap allocation would be a large
+/// regression over a device stack array.
+static bool isDeviceCode(mlir::Operation *func, mlir::ModuleOp mod) {
+ if (func->getParentOfType<mlir::gpu::GPUModuleOp>())
+ return true;
+ if (auto procAttr =
+ func->getAttrOfType<cuf::ProcAttributeAttr>(cuf::getProcAttrName()))
+ return procAttr.getValue() != cuf::ProcAttribute::Host;
+ if (mlir::acc::isAccRoutine(func))
+ return true;
+ if (auto offloadMod =
+ llvm::dyn_cast<mlir::omp::OffloadModuleInterface>(mod.getOperation()))
+ return offloadMod.getIsTargetDevice();
+ return false;
+}
+
+bool fir::promoteDynamicAllocasToCudaHeap(mlir::RewriterBase &rewriter,
+ mlir::Operation *func) {
+ auto mod = func->getParentOfType<mlir::ModuleOp>();
+ if (!mod)
+ return false;
+ fir::CudaHeapAllocMode mode = fir::getCudaHeapAllocMode(mod);
+ if (mode == fir::CudaHeapAllocMode::None || isDeviceCode(func, mod))
+ return false;
+
+ bool changed = false;
+ // Named locals only: automatic arrays and automatic character. Compiler
+ // temporaries do not need unified memory and would turn a stack save/restore
+ // into a malloc/free pair, possibly per loop iteration.
+ auto mustReplace = [](fir::AllocaOp alloca) {
+ std::optional<llvm::StringRef> uniqName = alloca.getUniqName();
+ return alloca.isDynamic() && uniqName && !uniqName->empty();
+ };
+ auto genAllocmem = [&](mlir::OpBuilder &builder, fir::AllocaOp alloca,
+ bool) -> mlir::Value {
+ auto name = [](std::optional<llvm::StringRef> opt) -> llvm::StringRef {
+ return opt ? *opt : llvm::StringRef{};
+ };
+ auto heap = fir::AllocMemOp::create(
+ builder, alloca.getLoc(), alloca.getInType(),
+ name(alloca.getUniqName()), name(alloca.getBindcName()),
+ alloca.getTypeparams(), alloca.getShape());
+ fir::setCudaHeapAllocMode(heap.getOperation(), mode);
+ // Keep the placement passes from sinking it back to the stack: the
+ // allocator is chosen here and the matching free is emitted below.
+ heap->setAttr(fir::MustBeHeapAttr::getAttrName(),
+ fir::MustBeHeapAttr::get(builder.getContext(), true));
+ changed = true;
+ return heap;
+ };
+ auto genFreemem = [&](mlir::Location loc, mlir::OpBuilder &builder,
+ mlir::Value allocmem) {
+ auto free = fir::FreeMemOp::create(builder, loc, allocmem);
+ fir::setCudaHeapAllocMode(free.getOperation(), mode);
+ };
+ fir::replaceAllocas(rewriter, func, mustReplace, genAllocmem, genFreemem);
+ return changed;
+}
diff --git a/flang/lib/Optimizer/Transforms/StackArrays.cpp b/flang/lib/Optimizer/Transforms/StackArrays.cpp
index 77861e67a07b1..dbd4f5f08b701 100644
--- a/flang/lib/Optimizer/Transforms/StackArrays.cpp
+++ b/flang/lib/Optimizer/Transforms/StackArrays.cpp
@@ -15,6 +15,7 @@
#include "flang/Optimizer/Dialect/FIRType.h"
#include "flang/Optimizer/Dialect/Support/FIRContext.h"
#include "flang/Optimizer/Support/DataLayout.h"
+#include "flang/Optimizer/Transforms/MemoryUtils.h"
#include "flang/Optimizer/Transforms/Passes.h"
#include "mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h"
#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
@@ -746,6 +747,12 @@ llvm::StringRef StackArraysPass::getDescription() const {
void StackArraysPass::runOnOperation() {
mlir::func::FuncOp func = getOperation();
+ // -fstack-arrays does not apply to the automatic arrays that
+ // -gpu=mem:unified|managed must place in unified/managed memory. Done before
+ // the analysis below, which skips the fir.must_be_heap pairs it creates.
+ mlir::IRRewriter cudaHeapRewriter(&getContext());
+ fir::promoteDynamicAllocasToCudaHeap(cudaHeapRewriter, func.getOperation());
+
auto &analysis = getAnalysis<fir::StackArraysAnalysisWrapper>();
const fir::StackArraysAnalysisWrapper::AllocMemMap *candidateOps =
analysis.getCandidateOps(func);
diff --git a/flang/lib/Semantics/resolve-names.cpp b/flang/lib/Semantics/resolve-names.cpp
index dd6a5401a8250..27c4e96d269aa 100644
--- a/flang/lib/Semantics/resolve-names.cpp
+++ b/flang/lib/Semantics/resolve-names.cpp
@@ -10775,37 +10775,6 @@ void ResolveNamesVisitor::FinishSpecificationPart(
context().languageFeatures().IsEnabled(
common::LanguageFeature::CudaPinned))
object->set_cudaDataAttr(common::CUDADataAttr::Pinned);
- } else if (!object->cudaDataAttr() && !IsDummy(symbol) &&
- !IsAllocatable(symbol) && !IsPointer(symbol) && !IsSaved(symbol) &&
- !IsCUDADeviceContext(&symbol.owner()) &&
- object->shape().IsExplicitShape()) {
- // Under -gpu=mem:unified|managed, allocate adjustable / VLA automatic
- // arrays in CUDA unified/managed memory (fixed-size automatic arrays
- // stay on the stack). Tag those locals so lowering uses
- // cuf.alloc/cuf.free. Unlike the allocatable managed tagging above,
- // this does not require -fcuda: OpenACC + -gpu=mem:unified relies on
- // it, and cuf.alloc does not go through the CUDA Fortran
- // managed-descriptor pipeline that motivated the -fcuda gate.
- std::optional<common::CUDADataAttr> attr;
- if (context().languageFeatures().IsEnabled(
- common::LanguageFeature::CudaUnified))
- attr = common::CUDADataAttr::Unified;
- else if (context().languageFeatures().IsEnabled(
- common::LanguageFeature::CudaManaged))
- attr = common::CUDADataAttr::Managed;
- if (attr) {
- auto boundIsNonConstant{[](const Bound &b) {
- return !b.isExplicit() || !b.GetExplicit() ||
- !evaluate::IsConstantExpr(*b.GetExplicit());
- }};
- for (const ShapeSpec &ss : object->shape()) {
- if (boundIsNonConstant(ss.lbound()) ||
- boundIsNonConstant(ss.ubound())) {
- object->set_cudaDataAttr(*attr);
- break;
- }
- }
- }
}
}
}
diff --git a/flang/test/Fir/CUDA/cuda-heap-alloc-managed.fir b/flang/test/Fir/CUDA/cuda-heap-alloc-managed.fir
new file mode 100644
index 0000000000000..c6fc559ce5088
--- /dev/null
+++ b/flang/test/Fir/CUDA/cuda-heap-alloc-managed.fir
@@ -0,0 +1,44 @@
+// RUN: fir-opt --memory-allocation-opt %s | FileCheck %s --check-prefix=HEAP
+// RUN: fir-opt --fir-to-llvm-ir %s | FileCheck %s --check-prefix=LLVM
+
+// Same routing as cuda-heap-alloc-unified.fir, with the managed entry points.
+
+// Declarations are emitted at the top of the module, before any function.
+// LLVM-DAG: llvm.func @malloc_managed(i64) -> !llvm.ptr
+// LLVM-DAG: llvm.func @free_managed(!llvm.ptr)
+// LLVM-DAG: llvm.func @malloc(i64) -> !llvm.ptr
+// LLVM-DAG: llvm.func @free(!llvm.ptr)
+
+module attributes {fir.cuda_heap_alloc = "managed"} {
+
+// HEAP-LABEL: func.func @vla(
+// HEAP: %[[MEM:.*]] = fir.allocmem !fir.array<?xf32>, %{{.*}} {bindc_name = "a", fir.cuda_heap_alloc = "managed", fir.must_be_heap = true, uniq_name = "_QFvlaEa"}
+// HEAP: fir.freemem %[[MEM]] {fir.cuda_heap_alloc = "managed"} : !fir.heap<!fir.array<?xf32>>
+func.func @vla(%arg0: !fir.ref<i32>) {
+ %0 = fir.load %arg0 : !fir.ref<i32>
+ %1 = fir.convert %0 : (i32) -> index
+ %2 = fir.alloca !fir.array<?xf32>, %1 {bindc_name = "a", uniq_name = "_QFvlaEa"}
+ return
+}
+
+// LLVM-LABEL: llvm.func @marked_heap(
+// LLVM: llvm.call @malloc_managed(
+// LLVM: llvm.call @free_managed(
+func.func @marked_heap(%n: index) {
+ %0 = fir.allocmem !fir.array<?xf32>, %n {fir.cuda_heap_alloc = "managed"}
+ fir.freemem %0 {fir.cuda_heap_alloc = "managed"} : !fir.heap<!fir.array<?xf32>>
+ return
+}
+
+// LLVM-LABEL: llvm.func @unmarked_heap(
+// LLVM-NOT: llvm.call @malloc_managed(
+// LLVM: llvm.call @malloc(
+// LLVM-NOT: llvm.call @free_managed(
+// LLVM: llvm.call @free(
+func.func @unmarked_heap(%n: index) {
+ %0 = fir.allocmem !fir.array<?xf32>, %n
+ fir.freemem %0 : !fir.heap<!fir.array<?xf32>>
+ return
+}
+
+}
diff --git a/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir b/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
new file mode 100644
index 0000000000000..7a043cf9737a4
--- /dev/null
+++ b/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
@@ -0,0 +1,86 @@
+// RUN: fir-opt --memory-allocation-opt %s | FileCheck %s --check-prefix=HEAP
+// RUN: fir-opt --stack-arrays %s | FileCheck %s --check-prefix=HEAP
+// RUN: fir-opt --fir-to-llvm-ir %s | FileCheck %s --check-prefix=LLVM
+
+// Under fir.cuda_heap_alloc = "unified", named automatic arrays move to the
+// heap and are marked. Only marked allocations use malloc_unified: memory the
+// Fortran runtime allocated must keep being released by libc free.
+
+// Declarations are emitted at the top of the module, before any function.
+// LLVM-DAG: llvm.func @malloc_unified(i64) -> !llvm.ptr
+// LLVM-DAG: llvm.func @free_unified(!llvm.ptr)
+// LLVM-DAG: llvm.func @malloc(i64) -> !llvm.ptr
+// LLVM-DAG: llvm.func @free(!llvm.ptr)
+
+module attributes {fir.cuda_heap_alloc = "unified"} {
+
+// HEAP-LABEL: func.func @vla(
+// HEAP: %[[MEM:.*]] = fir.allocmem !fir.array<?xf32>, %{{.*}} {bindc_name = "a", fir.cuda_heap_alloc = "unified", fir.must_be_heap = true, uniq_name = "_QFvlaEa"}
+// HEAP: fir.freemem %[[MEM]] {fir.cuda_heap_alloc = "unified"} : !fir.heap<!fir.array<?xf32>>
+func.func @vla(%arg0: !fir.ref<i32>) {
+ %0 = fir.load %arg0 : !fir.ref<i32>
+ %1 = fir.convert %0 : (i32) -> index
+ %2 = fir.alloca !fir.array<?xf32>, %1 {bindc_name = "a", uniq_name = "_QFvlaEa"}
+ return
+}
+
+// Automatic character is an automatic too.
+// HEAP-LABEL: func.func @autochar(
+// HEAP: fir.allocmem !fir.char<1,?>(%{{.*}} : index) {{{.*}}fir.cuda_heap_alloc = "unified"
+func.func @autochar(%arg0: index) {
+ %0 = fir.alloca !fir.char<1,?>(%arg0 : index) {bindc_name = "s", uniq_name = "_QFautocharEs"}
+ return
+}
+
+// Fixed-size automatics stay on the stack.
+// HEAP-LABEL: func.func @fixed(
+// HEAP: fir.alloca !fir.array<128xf32>
+// HEAP-NOT: fir.allocmem
+func.func @fixed() {
+ %0 = fir.alloca !fir.array<128xf32> {bindc_name = "a", uniq_name = "_QFfixedEa"}
+ return
+}
+
+// Compiler temporaries are not automatics: turning them into malloc/free would
+// cost an allocation per loop iteration.
+// HEAP-LABEL: func.func @anon_temp(
+// HEAP: fir.alloca !fir.array<?xf32>
+// HEAP-NOT: fir.allocmem
+func.func @anon_temp(%arg0: index) {
+ %0 = fir.alloca !fir.array<?xf32>, %arg0
+ return
+}
+
+// Device code keeps its stack allocation: the entry points are host-only.
+// HEAP-LABEL: func.func @device_vla(
+// HEAP: fir.alloca !fir.array<?xf32>
+// HEAP-NOT: fir.allocmem
+func.func @device_vla(%arg0: index) attributes {cuf.proc_attr = #cuf.cuda_proc<global>} {
+ %0 = fir.alloca !fir.array<?xf32>, %arg0 {bindc_name = "a", uniq_name = "_QFdevice_vlaEa"}
+ return
+}
+
+// HEAP-LABEL: func.func @marked_heap(
+// LLVM-LABEL: llvm.func @marked_heap(
+// LLVM: llvm.call @malloc_unified(
+// LLVM: llvm.call @free_unified(
+func.func @marked_heap(%n: index) {
+ %0 = fir.allocmem !fir.array<?xf32>, %n {fir.cuda_heap_alloc = "unified"}
+ fir.freemem %0 {fir.cuda_heap_alloc = "unified"} : !fir.heap<!fir.array<?xf32>>
+ return
+}
+
+// An unmarked pair belongs to libc: this is the shape of a buffer the Fortran
+// runtime allocated and lowered code releases.
+// LLVM-LABEL: llvm.func @unmarked_heap(
+// LLVM-NOT: llvm.call @malloc_unified(
+// LLVM: llvm.call @malloc(
+// LLVM-NOT: llvm.call @free_unified(
+// LLVM: llvm.call @free(
+func.func @unmarked_heap(%n: index) {
+ %0 = fir.allocmem !fir.array<?xf32>, %n
+ fir.freemem %0 : !fir.heap<!fir.array<?xf32>>
+ return
+}
+
+}
diff --git a/flang/test/Lower/CUDA/cuda-gpu-unified-automatic-array.f90 b/flang/test/Lower/CUDA/cuda-gpu-unified-automatic-array.f90
index 07638b0ca8dfa..c3c15c98ceed0 100644
--- a/flang/test/Lower/CUDA/cuda-gpu-unified-automatic-array.f90
+++ b/flang/test/Lower/CUDA/cuda-gpu-unified-automatic-array.f90
@@ -1,21 +1,23 @@
-! RUN: bbc -emit-hlfir -gpu=unified %s -o - | FileCheck %s
-! RUN: bbc -emit-hlfir -gpu=managed %s -o - | FileCheck %s --check-prefix=MANAGED
+! RUN: bbc -emit-hlfir -gpu=unified %s -o - | FileCheck %s --check-prefixes=CHECK,UNIFIED
+! RUN: bbc -emit-hlfir -gpu=managed %s -o - | FileCheck %s --check-prefixes=CHECK,MANAGED
+! RUN: bbc -emit-hlfir %s -o - | FileCheck %s --check-prefixes=CHECK,NOFLAG
-! Under -gpu=mem:unified|managed, allocate adjustable / VLA automatic arrays
-! in CUDA unified/managed memory. Fixed-size automatic arrays stay on the
-! stack so host pointers remain shared under unified memory (e.g. OpenACC).
+! Under -gpu=mem:unified|managed, dynamic automatic arrays are later moved to
+! the heap and allocated with malloc_unified / malloc_managed. Lowering only
+! records the mode on the module; symbols stay unmarked (no cudaDataAttr).
+
+! UNIFIED: module attributes {{{.*}}fir.cuda_heap_alloc = "unified"
+! MANAGED: module attributes {{{.*}}fir.cuda_heap_alloc = "managed"
+! NOFLAG-NOT: fir.cuda_heap_alloc
module m_adj
integer :: nx = 32
end module
! CHECK-LABEL: func.func @_QPvla(
-! CHECK: %[[ALLOC:.*]] = cuf.alloc !fir.array<?xf32>, %{{.*}} : index {bindc_name = "a", data_attr = #cuf.cuda<unified>, uniq_name = "_QFvlaEa"}
-! CHECK: %[[DECL:.*]]:2 = hlfir.declare %[[ALLOC]](%{{.*}}) {data_attr = #cuf.cuda<unified>, uniq_name = "_QFvlaEa"}
-! CHECK: cuf.free %[[DECL]]#1 : !fir.ref<!fir.array<?xf32>> {data_attr = #cuf.cuda<unified>}
-! MANAGED-LABEL: func.func @_QPvla(
-! MANAGED: cuf.alloc !fir.array<?xf32>, %{{.*}} : index {{{.*}}data_attr = #cuf.cuda<managed>
-! MANAGED: cuf.free %{{.*}} : !fir.ref<!fir.array<?xf32>> {data_attr = #cuf.cuda<managed>}
+! CHECK-NOT: cuf.alloc
+! CHECK-NOT: data_attr = #cuf.cuda
+! CHECK: fir.alloca !fir.array<?xf32>
subroutine vla(n)
integer :: n
real :: a(n)
@@ -23,34 +25,28 @@ subroutine vla(n)
end subroutine
! CHECK-LABEL: func.func @_QPadjustable(
-! CHECK: cuf.alloc !fir.array<?xf32>, %{{.*}} : index {{{.*}}data_attr = #cuf.cuda<unified>
-! CHECK: cuf.free %{{.*}} {data_attr = #cuf.cuda<unified>}
-! MANAGED-LABEL: func.func @_QPadjustable(
-! MANAGED: cuf.alloc !fir.array<?xf32>, %{{.*}} : index {{{.*}}data_attr = #cuf.cuda<managed>
+! CHECK-NOT: cuf.alloc
+! CHECK-NOT: data_attr = #cuf.cuda
+! CHECK: fir.alloca !fir.array<?xf32>
subroutine adjustable
use m_adj
real :: a(0:(nx+1)/2)
a(0) = 0.0
end subroutine
-! Fixed-size automatic arrays must remain ordinary stack allocations.
+! Fixed-size automatic arrays remain ordinary stack allocations.
! CHECK-LABEL: func.func @_QPfixed(
! CHECK-NOT: cuf.alloc
! CHECK: fir.alloca !fir.array<128xf32>
-! CHECK-NOT: cuf.free
-! MANAGED-LABEL: func.func @_QPfixed(
-! MANAGED-NOT: cuf.alloc
-! MANAGED: fir.alloca !fir.array<128xf32>
subroutine fixed
real :: a(128)
a(1) = 1.0
end subroutine
-! Dummy adjustable arrays are caller-allocated; do not retag them.
+! Dummy adjustable arrays are caller-allocated.
! CHECK-LABEL: func.func @_QPdummy_adj(
! CHECK-NOT: cuf.alloc
-! MANAGED-LABEL: func.func @_QPdummy_adj(
-! MANAGED-NOT: cuf.alloc
+! CHECK-NOT: fir.alloca !fir.array<?xf32>
subroutine dummy_adj(a, n)
integer :: n
real :: a(n)
>From 47f3856e457e31f356870c3688e0f1eacab7d9aa Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Thu, 30 Jul 2026 14:57:26 -0700
Subject: [PATCH 4/7] [flang][cuda] Make the heap allocator entry point names
configurable
Derive them as the libc name plus a per-mode suffix set through
fir-to-llvm-ir options, instead of hardcoding the names in CodeGen.
Co-Authored-By: Claude Opus 5 <noreply at anthropic.com>
---
.../flang/Optimizer/CodeGen/CGPasses.td | 8 +-
.../include/flang/Optimizer/CodeGen/CodeGen.h | 6 ++
flang/lib/Optimizer/CodeGen/CodeGen.cpp | 78 +++++++++----------
.../test/Fir/CUDA/cuda-heap-alloc-unified.fir | 5 ++
4 files changed, 53 insertions(+), 44 deletions(-)
diff --git a/flang/include/flang/Optimizer/CodeGen/CGPasses.td b/flang/include/flang/Optimizer/CodeGen/CGPasses.td
index 2741e0206dfec..1163aff79e171 100644
--- a/flang/include/flang/Optimizer/CodeGen/CGPasses.td
+++ b/flang/include/flang/Optimizer/CodeGen/CGPasses.td
@@ -45,7 +45,13 @@ def FIRToLLVMLowering : Pass<"fir-to-llvm-ir", "mlir::ModuleOp"> {
"std::string", /*default=*/"",
"Name of the function to call to allocate CUDA Fortran descriptors. "
"Must have the same signature as CUFAllocDescriptor. "
- "Defaults to CUFAllocDescriptor.">
+ "Defaults to CUFAllocDescriptor.">,
+ Option<"unifiedHeapAllocSuffix", "unified-heap-alloc-suffix", "std::string",
+ /*default=*/"", "Suffix of the allocator entry points used for "
+ "allocations marked with the unified heap allocation mode.">,
+ Option<"managedHeapAllocSuffix", "managed-heap-alloc-suffix", "std::string",
+ /*default=*/"", "Suffix of the allocator entry points used for "
+ "allocations marked with the managed heap allocation mode.">
];
}
diff --git a/flang/include/flang/Optimizer/CodeGen/CodeGen.h b/flang/include/flang/Optimizer/CodeGen/CodeGen.h
index 948c240967c5b..1d36788fb84f9 100644
--- a/flang/include/flang/Optimizer/CodeGen/CodeGen.h
+++ b/flang/include/flang/Optimizer/CodeGen/CodeGen.h
@@ -71,6 +71,12 @@ struct FIRToLLVMPassOptions {
// Conversion pass of the MLIR complex dialect.
Fortran::frontend::CodeGenOptions::ComplexRangeKind ComplexRange =
Fortran::frontend::CodeGenOptions::ComplexRangeKind::CX_Full;
+
+ // Suffix appended to the libc allocator name (malloc, free, aligned_alloc,
+ // posix_memalign) for allocations marked with a heap allocation mode, e.g.
+ // malloc -> malloc_unified. Lets a runtime name its entry points otherwise.
+ std::string unifiedHeapAllocSuffix = "_unified";
+ std::string managedHeapAllocSuffix = "_managed";
};
/// Convert FIR to the LLVM IR dialect with default options.
diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
index 8108b19a4981c..7696f5f900c5e 100644
--- a/flang/lib/Optimizer/CodeGen/CodeGen.cpp
+++ b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -1311,51 +1311,33 @@ getMallocInModule(ModuleOp mod, fir::AllocMemOp op,
return mlir::SymbolRefAttr::get(mallocDecl);
}
-/// Allocator entry points for an allocation marked by the allocation placement
-/// passes under -gpu=mem:unified|managed. Only marked fir.allocmem/fir.freemem
-/// pairs are routed: memory that the Fortran runtime allocated must keep being
-/// released by libc free, and vice versa.
-static llvm::StringRef getHeapAllocName(mlir::Operation *op,
- llvm::StringRef plain,
- llvm::StringRef unified,
- llvm::StringRef managed) {
- // Device modules keep libc names; the indirect entry points are host-side.
+/// Allocator entry point for an allocation marked by the allocation placement
+/// passes with a heap allocation mode: the libc name plus the mode suffix from
+/// the pass options, e.g. malloc -> malloc_unified. Only marked
+/// fir.allocmem/fir.freemem pairs are routed, since memory the Fortran runtime
+/// allocated must keep being released by libc free, and vice versa.
+static std::string getHeapAllocName(mlir::Operation *op, llvm::StringRef plain,
+ const fir::FIRToLLVMPassOptions &options) {
+ // Device modules keep libc names; the mode entry points are host-side.
if (op->getParentOfType<mlir::gpu::GPUModuleOp>())
- return plain;
+ return plain.str();
switch (fir::getCudaHeapAllocMode(op)) {
case fir::CudaHeapAllocMode::Unified:
- return unified;
+ return (plain + options.unifiedHeapAllocSuffix).str();
case fir::CudaHeapAllocMode::Managed:
- return managed;
+ return (plain + options.managedHeapAllocSuffix).str();
case fir::CudaHeapAllocMode::None:
- return plain;
+ return plain.str();
}
llvm_unreachable("unexpected CudaHeapAllocMode");
}
-static llvm::StringRef getHostHeapMallocName(mlir::Operation *op) {
- return getHeapAllocName(op, "malloc", "malloc_unified", "malloc_managed");
-}
-
-static llvm::StringRef getHostHeapFreeName(mlir::Operation *op) {
- return getHeapAllocName(op, "free", "free_unified", "free_managed");
-}
-
-static llvm::StringRef getHostHeapAlignedAllocName(mlir::Operation *op) {
- return getHeapAllocName(op, "aligned_alloc", "aligned_alloc_unified",
- "aligned_alloc_managed");
-}
-
-static llvm::StringRef getHostHeapPosixMemalignName(mlir::Operation *op) {
- return getHeapAllocName(op, "posix_memalign", "posix_memalign_unified",
- "posix_memalign_managed");
-}
-
/// Return the LLVMFuncOp corresponding to the standard malloc call.
static mlir::SymbolRefAttr getMalloc(fir::AllocMemOp op,
mlir::ConversionPatternRewriter &rewriter,
- mlir::Type indexType) {
- llvm::StringRef name = getHostHeapMallocName(op);
+ mlir::Type indexType,
+ const fir::FIRToLLVMPassOptions &options) {
+ std::string name = getHeapAllocName(op, "malloc", options);
if (auto mod = op->getParentOfType<mlir::gpu::GPUModuleOp>())
return getMallocInModule(mod, op, rewriter, indexType, name);
auto mod = op->getParentOfType<mlir::ModuleOp>();
@@ -1384,8 +1366,9 @@ static mlir::SymbolRefAttr getAlignedAllocInModule(
static mlir::SymbolRefAttr
getAlignedAlloc(fir::AllocMemOp op, mlir::ConversionPatternRewriter &rewriter,
- mlir::Type indexType) {
- llvm::StringRef name = getHostHeapAlignedAllocName(op);
+ mlir::Type indexType,
+ const fir::FIRToLLVMPassOptions &options) {
+ std::string name = getHeapAllocName(op, "aligned_alloc", options);
if (auto mod = op->getParentOfType<mlir::gpu::GPUModuleOp>())
return getAlignedAllocInModule(mod, op, rewriter, indexType, name);
auto mod = op->getParentOfType<mlir::ModuleOp>();
@@ -1416,8 +1399,9 @@ static mlir::SymbolRefAttr getPosixMemalignInModule(
static mlir::SymbolRefAttr
getPosixMemalign(fir::AllocMemOp op, mlir::ConversionPatternRewriter &rewriter,
- mlir::Type indexType) {
- llvm::StringRef name = getHostHeapPosixMemalignName(op);
+ mlir::Type indexType,
+ const fir::FIRToLLVMPassOptions &options) {
+ std::string name = getHeapAllocName(op, "posix_memalign", options);
if (auto mod = op->getParentOfType<mlir::gpu::GPUModuleOp>())
return getPosixMemalignInModule(mod, op, rewriter, indexType, name);
auto mod = op->getParentOfType<mlir::ModuleOp>();
@@ -1504,7 +1488,8 @@ struct AllocMemOpConversion : public fir::FIROpConversion<fir::AllocMemOp> {
mlir::Value nullPtr =
mlir::LLVM::ZeroOp::create(rewriter, loc, ptrTy);
mlir::LLVM::StoreOp::create(rewriter, loc, nullPtr, memptr);
- heap->setAttr("callee", getPosixMemalign(heap, rewriter, mallocTy));
+ heap->setAttr("callee", getPosixMemalign(heap, rewriter, mallocTy,
+ this->options));
mlir::LLVM::CallOp::create(
rewriter, loc,
mlir::TypeRange{
@@ -1526,7 +1511,8 @@ struct AllocMemOpConversion : public fir::FIROpConversion<fir::AllocMemOp> {
~static_cast<std::int64_t>(*alignment - 1));
mlir::Value roundedSize = mlir::LLVM::AndOp::create(
rewriter, loc, mallocTy, sizePlus, notAlignMinusOne);
- heap->setAttr("callee", getAlignedAlloc(heap, rewriter, mallocTy));
+ heap->setAttr("callee",
+ getAlignedAlloc(heap, rewriter, mallocTy, this->options));
rewriter.replaceOpWithNewOp<mlir::LLVM::CallOp>(
heap, ::getLlvmPtrType(heap.getContext()),
mlir::ValueRange{alignVal, roundedSize},
@@ -1535,7 +1521,7 @@ struct AllocMemOpConversion : public fir::FIROpConversion<fir::AllocMemOp> {
}
}
- heap->setAttr("callee", getMalloc(heap, rewriter, mallocTy));
+ heap->setAttr("callee", getMalloc(heap, rewriter, mallocTy, this->options));
rewriter.replaceOpWithNewOp<mlir::LLVM::CallOp>(
heap, ::getLlvmPtrType(heap.getContext()), size,
addLLVMOpBundleAttrs(rewriter, heap->getAttrs(), 1));
@@ -1579,8 +1565,9 @@ getFreeInModule(ModuleOp mod, fir::FreeMemOp op,
}
static mlir::SymbolRefAttr getFree(fir::FreeMemOp op,
- mlir::ConversionPatternRewriter &rewriter) {
- llvm::StringRef name = getHostHeapFreeName(op);
+ mlir::ConversionPatternRewriter &rewriter,
+ const fir::FIRToLLVMPassOptions &options) {
+ std::string name = getHeapAllocName(op, "free", options);
if (auto mod = op->getParentOfType<mlir::gpu::GPUModuleOp>())
return getFreeInModule(mod, op, rewriter, name);
auto mod = op->getParentOfType<mlir::ModuleOp>();
@@ -1606,7 +1593,7 @@ struct FreeMemOpConversion : public fir::FIROpConversion<fir::FreeMemOp> {
matchAndRewrite(fir::FreeMemOp freemem, OpAdaptor adaptor,
mlir::ConversionPatternRewriter &rewriter) const override {
mlir::Location loc = freemem.getLoc();
- freemem->setAttr("callee", getFree(freemem, rewriter));
+ freemem->setAttr("callee", getFree(freemem, rewriter, this->options));
mlir::LLVM::CallOp::create(
rewriter, loc, mlir::TypeRange{},
mlir::ValueRange{adaptor.getHeapref()},
@@ -4787,6 +4774,11 @@ class FIRToLLVMLowering
if (!cudaDescriptorAllocFunction.empty())
options.cudaDescriptorAllocFunction = cudaDescriptorAllocFunction;
+ if (!unifiedHeapAllocSuffix.empty())
+ options.unifiedHeapAllocSuffix = unifiedHeapAllocSuffix;
+ if (!managedHeapAllocSuffix.empty())
+ options.managedHeapAllocSuffix = managedHeapAllocSuffix;
+
// Run dynamic pass pipeline for converting Math dialect
// operations into other dialects (llvm, func, etc.).
// Some conversions of Math operations cannot be done
diff --git a/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir b/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
index 7a043cf9737a4..a98f4121271cb 100644
--- a/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
+++ b/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
@@ -1,6 +1,7 @@
// RUN: fir-opt --memory-allocation-opt %s | FileCheck %s --check-prefix=HEAP
// RUN: fir-opt --stack-arrays %s | FileCheck %s --check-prefix=HEAP
// RUN: fir-opt --fir-to-llvm-ir %s | FileCheck %s --check-prefix=LLVM
+// RUN: fir-opt --fir-to-llvm-ir=unified-heap-alloc-suffix=_pool %s | FileCheck %s --check-prefix=SUFFIX
// Under fir.cuda_heap_alloc = "unified", named automatic arrays move to the
// heap and are marked. Only marked allocations use malloc_unified: memory the
@@ -64,6 +65,10 @@ func.func @device_vla(%arg0: index) attributes {cuf.proc_attr = #cuf.cuda_proc<g
// LLVM-LABEL: llvm.func @marked_heap(
// LLVM: llvm.call @malloc_unified(
// LLVM: llvm.call @free_unified(
+// The entry point names are the libc name plus a configurable suffix.
+// SUFFIX-LABEL: llvm.func @marked_heap(
+// SUFFIX: llvm.call @malloc_pool(
+// SUFFIX: llvm.call @free_pool(
func.func @marked_heap(%n: index) {
%0 = fir.allocmem !fir.array<?xf32>, %n {fir.cuda_heap_alloc = "unified"}
fir.freemem %0 {fir.cuda_heap_alloc = "unified"} : !fir.heap<!fir.array<?xf32>>
>From 4ed105f92affc417bc1308060ee1efbbad0313a8 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Mon, 17 Aug 2026 04:26:28 -0700
Subject: [PATCH 5/7] [flang][cuda] Only promote non-pinned user variables to
unified memory
Rename promoteDynamicAllocasToCudaHeap to
promoteDynamicVariableAllocasToCudaHeap and say in its description that it
only rewrites user variables, and skip the allocas marked fir.must_be_stack.
An array function result carries that attribute and has its storage replaced
by the caller buffer, so promoting it only left a dead malloc/free pair in a
function returning an array.
---
.../flang/Optimizer/Transforms/MemoryUtils.h | 14 +++++++-----
.../Transforms/AllocationPlacement.cpp | 3 ++-
.../Optimizer/Transforms/MemoryAllocation.cpp | 3 ++-
.../lib/Optimizer/Transforms/MemoryUtils.cpp | 22 ++++++++++++++-----
.../lib/Optimizer/Transforms/StackArrays.cpp | 3 ++-
.../test/Fir/CUDA/cuda-heap-alloc-unified.fir | 10 +++++++++
6 files changed, 40 insertions(+), 15 deletions(-)
diff --git a/flang/include/flang/Optimizer/Transforms/MemoryUtils.h b/flang/include/flang/Optimizer/Transforms/MemoryUtils.h
index 62438c7c44dd3..6cf27adc36221 100644
--- a/flang/include/flang/Optimizer/Transforms/MemoryUtils.h
+++ b/flang/include/flang/Optimizer/Transforms/MemoryUtils.h
@@ -57,12 +57,14 @@ bool replaceAllocas(mlir::RewriterBase &rewriter, mlir::Operation *parentOp,
MustRewriteCallBack, AllocaRewriterCallBack,
DeallocCallBack);
-/// Under -gpu=mem:unified|managed, move the dynamically sized fir.alloca of
-/// \p func to fir.allocmem/fir.freemem pairs marked for the unified/managed
-/// allocator. Does nothing for device code, which keeps its stack allocations.
-/// Returns true if the function was modified.
-bool promoteDynamicAllocasToCudaHeap(mlir::RewriterBase &rewriter,
- mlir::Operation *func);
+/// Under -gpu=mem:unified|managed, move the dynamically sized fir.alloca of the
+/// user variables of \p func (automatic arrays and automatic character) to
+/// fir.allocmem/fir.freemem pairs marked for the unified/managed allocator.
+/// Compiler temporaries, fir.must_be_stack allocations, and device code, which
+/// keeps its stack allocations, are left alone. Returns true if the function
+/// was modified.
+bool promoteDynamicVariableAllocasToCudaHeap(mlir::RewriterBase &rewriter,
+ mlir::Operation *func);
} // namespace fir
diff --git a/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp b/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
index 231d381649f95..7a4aa3c6b7bb5 100644
--- a/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
+++ b/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
@@ -202,7 +202,8 @@ void AllocationPlacementPass::runOnOperation() {
// Done first: the pairs it creates are marked fir.must_be_heap, so the
// placement decisions below leave them alone.
mlir::IRRewriter cudaHeapRewriter(&getContext());
- fir::promoteDynamicAllocasToCudaHeap(cudaHeapRewriter, func.getOperation());
+ fir::promoteDynamicVariableAllocasToCudaHeap(cudaHeapRewriter,
+ func.getOperation());
fir::AllocationPlacementThresholds baseThresholds;
baseThresholds.stackArrays = stackArrays;
diff --git a/flang/lib/Optimizer/Transforms/MemoryAllocation.cpp b/flang/lib/Optimizer/Transforms/MemoryAllocation.cpp
index 6c4342b141050..058c64cb68395 100644
--- a/flang/lib/Optimizer/Transforms/MemoryAllocation.cpp
+++ b/flang/lib/Optimizer/Transforms/MemoryAllocation.cpp
@@ -133,7 +133,8 @@ class MemoryAllocationOpt
if (func.empty())
return;
mlir::IRRewriter cudaHeapRewriter(context);
- fir::promoteDynamicAllocasToCudaHeap(cudaHeapRewriter, func.getOperation());
+ fir::promoteDynamicVariableAllocasToCudaHeap(cudaHeapRewriter,
+ func.getOperation());
auto tryReplacing = [&](fir::AllocaOp alloca) {
bool res = !keepStackAllocation(alloca, options);
if (res) {
diff --git a/flang/lib/Optimizer/Transforms/MemoryUtils.cpp b/flang/lib/Optimizer/Transforms/MemoryUtils.cpp
index 0e5155dd0a361..4ba436805fd27 100644
--- a/flang/lib/Optimizer/Transforms/MemoryUtils.cpp
+++ b/flang/lib/Optimizer/Transforms/MemoryUtils.cpp
@@ -332,8 +332,8 @@ static bool isDeviceCode(mlir::Operation *func, mlir::ModuleOp mod) {
return false;
}
-bool fir::promoteDynamicAllocasToCudaHeap(mlir::RewriterBase &rewriter,
- mlir::Operation *func) {
+bool fir::promoteDynamicVariableAllocasToCudaHeap(mlir::RewriterBase &rewriter,
+ mlir::Operation *func) {
auto mod = func->getParentOfType<mlir::ModuleOp>();
if (!mod)
return false;
@@ -342,12 +342,22 @@ bool fir::promoteDynamicAllocasToCudaHeap(mlir::RewriterBase &rewriter,
return false;
bool changed = false;
- // Named locals only: automatic arrays and automatic character. Compiler
- // temporaries do not need unified memory and would turn a stack save/restore
- // into a malloc/free pair, possibly per loop iteration.
+ // User variables only: automatic arrays and automatic character, which are
+ // the ones carrying a uniqued name. Compiler temporaries do not need unified
+ // memory and would turn a stack save/restore into a malloc/free pair,
+ // possibly per loop iteration.
auto mustReplace = [](fir::AllocaOp alloca) {
+ if (!alloca.isDynamic())
+ return false;
+ // An alloca pinned to the stack (e.g. an array function result, whose
+ // storage the abstract-result pass replaces by the caller buffer) would
+ // only be left with a dead malloc/free pair.
+ if (auto attr = alloca->getAttrOfType<fir::MustBeStackAttr>(
+ fir::MustBeStackAttr::getAttrName()))
+ if (attr.getValue())
+ return false;
std::optional<llvm::StringRef> uniqName = alloca.getUniqName();
- return alloca.isDynamic() && uniqName && !uniqName->empty();
+ return uniqName && !uniqName->empty();
};
auto genAllocmem = [&](mlir::OpBuilder &builder, fir::AllocaOp alloca,
bool) -> mlir::Value {
diff --git a/flang/lib/Optimizer/Transforms/StackArrays.cpp b/flang/lib/Optimizer/Transforms/StackArrays.cpp
index dbd4f5f08b701..915af1a8d74b9 100644
--- a/flang/lib/Optimizer/Transforms/StackArrays.cpp
+++ b/flang/lib/Optimizer/Transforms/StackArrays.cpp
@@ -751,7 +751,8 @@ void StackArraysPass::runOnOperation() {
// -gpu=mem:unified|managed must place in unified/managed memory. Done before
// the analysis below, which skips the fir.must_be_heap pairs it creates.
mlir::IRRewriter cudaHeapRewriter(&getContext());
- fir::promoteDynamicAllocasToCudaHeap(cudaHeapRewriter, func.getOperation());
+ fir::promoteDynamicVariableAllocasToCudaHeap(cudaHeapRewriter,
+ func.getOperation());
auto &analysis = getAnalysis<fir::StackArraysAnalysisWrapper>();
const fir::StackArraysAnalysisWrapper::AllocMemMap *candidateOps =
diff --git a/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir b/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
index a98f4121271cb..ca18116d41ab6 100644
--- a/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
+++ b/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
@@ -61,6 +61,16 @@ func.func @device_vla(%arg0: index) attributes {cuf.proc_attr = #cuf.cuda_proc<g
return
}
+// An alloca pinned to the stack stays there: the array function result below is
+// replaced by the caller buffer, so a heap pair would only be dead code.
+// HEAP-LABEL: func.func @array_result(
+// HEAP: fir.alloca !fir.array<?xf32>
+// HEAP-NOT: fir.allocmem
+func.func @array_result(%arg0: index) {
+ %0 = fir.alloca !fir.array<?xf32>, %arg0 {bindc_name = "res", fir.must_be_stack = true, uniq_name = "_QFarray_resultEres"}
+ return
+}
+
// HEAP-LABEL: func.func @marked_heap(
// LLVM-LABEL: llvm.func @marked_heap(
// LLVM: llvm.call @malloc_unified(
>From 09c87554351ed873135234b94ea59960003d755a Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Wed, 19 Aug 2026 16:00:42 -0700
Subject: [PATCH 6/7] [flang][cuda] Promote the unified/managed automatic
variables in its own pass
Allocating the dynamically sized automatic variables in unified or managed
memory is a correctness requirement of -gpu=mem:unified|managed rather than a
placement heuristic, so do it in a cuda-heap-alloc-promotion pass of its own
instead of from whichever array allocation pass the pipeline happens to select.
It used to be skipped entirely with -disable-memory-allocation-opt.
Share the alloca to allocmem construction with the memory-allocation-opt and
allocation-placement passes as well, and stop counting attributes(host,device)
as device code, which matches the inDeviceContext helpers of the CUF passes.
---
.../flang/Optimizer/Transforms/MemoryUtils.h | 7 +++-
.../flang/Optimizer/Transforms/Passes.td | 19 ++++++++++
flang/lib/Optimizer/Passes/Pipelines.cpp | 6 +++
.../Transforms/AllocationPlacement.cpp | 20 +---------
flang/lib/Optimizer/Transforms/CMakeLists.txt | 1 +
.../Transforms/CudaHeapAllocPromotion.cpp | 38 +++++++++++++++++++
.../Optimizer/Transforms/MemoryAllocation.cpp | 16 +-------
.../lib/Optimizer/Transforms/MemoryUtils.cpp | 27 +++++++++----
.../lib/Optimizer/Transforms/StackArrays.cpp | 8 ----
flang/test/Driver/bbc-mlir-pass-pipeline.f90 | 1 +
.../cuda-heap-alloc-promotion-pipeline.f90 | 13 +++++++
.../test/Driver/mlir-debug-pass-pipeline.f90 | 1 +
flang/test/Driver/mlir-pass-pipeline.f90 | 1 +
.../test/Fir/CUDA/cuda-heap-alloc-managed.fir | 2 +-
.../test/Fir/CUDA/cuda-heap-alloc-unified.fir | 20 +++++++++-
flang/test/Fir/basic-program.fir | 1 +
16 files changed, 128 insertions(+), 53 deletions(-)
create mode 100644 flang/lib/Optimizer/Transforms/CudaHeapAllocPromotion.cpp
create mode 100644 flang/test/Driver/cuda-heap-alloc-promotion-pipeline.f90
diff --git a/flang/include/flang/Optimizer/Transforms/MemoryUtils.h b/flang/include/flang/Optimizer/Transforms/MemoryUtils.h
index 6cf27adc36221..61cabd1df584e 100644
--- a/flang/include/flang/Optimizer/Transforms/MemoryUtils.h
+++ b/flang/include/flang/Optimizer/Transforms/MemoryUtils.h
@@ -57,12 +57,17 @@ bool replaceAllocas(mlir::RewriterBase &rewriter, mlir::Operation *parentOp,
MustRewriteCallBack, AllocaRewriterCallBack,
DeallocCallBack);
+/// Create the fir.allocmem that replaces \p alloca: same allocated type, names,
+/// type parameters and shape. Any extra attribute is left to the caller.
+fir::AllocMemOp createAllocMemFromAlloca(mlir::OpBuilder &builder,
+ fir::AllocaOp alloca);
+
/// Under -gpu=mem:unified|managed, move the dynamically sized fir.alloca of the
/// user variables of \p func (automatic arrays and automatic character) to
/// fir.allocmem/fir.freemem pairs marked for the unified/managed allocator.
/// Compiler temporaries, fir.must_be_stack allocations, and device code, which
/// keeps its stack allocations, are left alone. Returns true if the function
-/// was modified.
+/// was modified. This is what the cuda-heap-alloc-promotion pass runs.
bool promoteDynamicVariableAllocasToCudaHeap(mlir::RewriterBase &rewriter,
mlir::Operation *func);
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index 98090fefeeedc..891c60eff97c3 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -285,6 +285,25 @@ def SimplifyIntrinsics : Pass<"simplify-intrinsics", "mlir::ModuleOp"> {
];
}
+def CudaHeapAllocPromotion
+ : Pass<"cuda-heap-alloc-promotion", "mlir::func::FuncOp"> {
+ let summary = "Allocate dynamically sized automatic variables in CUDA "
+ "unified or managed memory.";
+ let description = [{
+ Under -gpu=mem:unified|managed, which sets fir.cuda_heap_alloc on the
+ module, rewrite the dynamically sized fir.alloca of the user variables into
+ fir.allocmem/fir.freemem pairs marked with that mode, so that codegen calls
+ the matching allocator entry points instead of libc malloc and free. The
+ host pointer of such a variable has to be device accessible, so this is a
+ correctness requirement of those modes rather than a placement heuristic:
+ the pass runs on its own instead of being part of whichever array
+ allocation pass the pipeline happens to select. The pairs it creates are
+ marked fir.must_be_heap, which keeps those passes from moving them back to
+ the stack. Without the module attribute the pass does nothing.
+ }];
+ let dependentDialects = ["fir::FIROpsDialect"];
+}
+
def MemoryAllocationOpt : Pass<"memory-allocation-opt", "mlir::func::FuncOp"> {
let summary = "Convert stack to heap allocations and vice versa.";
let description = [{
diff --git a/flang/lib/Optimizer/Passes/Pipelines.cpp b/flang/lib/Optimizer/Passes/Pipelines.cpp
index cc61237760178..1fe255a3afd4c 100644
--- a/flang/lib/Optimizer/Passes/Pipelines.cpp
+++ b/flang/lib/Optimizer/Passes/Pipelines.cpp
@@ -202,6 +202,12 @@ void createDefaultFIROptimizerPassPipeline(mlir::PassManager &pm,
pm.addPass(mlir::createCSEPass());
+ // Unconditional and ahead of the array allocation placement below: under
+ // -gpu=mem:unified|managed the unified/managed allocators are required for
+ // correctness, so this must not depend on which placement pass is selected
+ // or on -disable-memory-allocation-opt.
+ pm.addPass(fir::createCudaHeapAllocPromotion());
+
if (enableAllocationPlacement)
fir::addAllocationPlacement(pm, pc.StackArrays);
else if (pc.StackArrays)
diff --git a/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp b/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
index 7a4aa3c6b7bb5..1c4e4fadde39d 100644
--- a/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
+++ b/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
@@ -126,20 +126,10 @@ getConstantByteSize(mlir::Operation *op,
}
/// Replacement generator used for stack-to-heap conversions (fir.alloca ->
-/// fir.allocmem). Mirrors the MemoryAllocation pass.
+/// fir.allocmem).
static mlir::Value genAllocmem(mlir::OpBuilder &builder, fir::AllocaOp alloca,
bool /*deallocPointsDominateAlloc*/) {
- mlir::Type varTy = alloca.getInType();
- auto unpackName = [](std::optional<llvm::StringRef> opt) -> llvm::StringRef {
- if (opt)
- return *opt;
- return {};
- };
- llvm::StringRef uniqName = unpackName(alloca.getUniqName());
- llvm::StringRef bindcName = unpackName(alloca.getBindcName());
- auto heap = fir::AllocMemOp::create(builder, alloca.getLoc(), varTy, uniqName,
- bindcName, alloca.getTypeparams(),
- alloca.getShape());
+ fir::AllocMemOp heap = fir::createAllocMemFromAlloca(builder, alloca);
LLVM_DEBUG(llvm::dbgs() << "allocation placement: replaced " << alloca
<< " with " << heap << '\n');
return heap;
@@ -199,12 +189,6 @@ void AllocationPlacementPass::runOnOperation() {
if (func.empty())
return;
- // Done first: the pairs it creates are marked fir.must_be_heap, so the
- // placement decisions below leave them alone.
- mlir::IRRewriter cudaHeapRewriter(&getContext());
- fir::promoteDynamicVariableAllocasToCudaHeap(cudaHeapRewriter,
- func.getOperation());
-
fir::AllocationPlacementThresholds baseThresholds;
baseThresholds.stackArrays = stackArrays;
baseThresholds.smallArrayThresholdBytes = smallArrayThresholdBytes;
diff --git a/flang/lib/Optimizer/Transforms/CMakeLists.txt b/flang/lib/Optimizer/Transforms/CMakeLists.txt
index f26f8c5c64bb0..9b25e590764ab 100644
--- a/flang/lib/Optimizer/Transforms/CMakeLists.txt
+++ b/flang/lib/Optimizer/Transforms/CMakeLists.txt
@@ -26,6 +26,7 @@ add_flang_library(FIRTransforms
ConstantArgumentGlobalisation.cpp
ControlFlowConverter.cpp
ConvertComplexPow.cpp
+ CudaHeapAllocPromotion.cpp
DebugTypeGenerator.cpp
EmitMIFGlobalCtors.cpp
ExternalNameConversion.cpp
diff --git a/flang/lib/Optimizer/Transforms/CudaHeapAllocPromotion.cpp b/flang/lib/Optimizer/Transforms/CudaHeapAllocPromotion.cpp
new file mode 100644
index 0000000000000..da4b15c43ba44
--- /dev/null
+++ b/flang/lib/Optimizer/Transforms/CudaHeapAllocPromotion.cpp
@@ -0,0 +1,38 @@
+//===- CudaHeapAllocPromotion.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
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Transforms/MemoryUtils.h"
+#include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+
+namespace fir {
+#define GEN_PASS_DEF_CUDAHEAPALLOCPROMOTION
+#include "flang/Optimizer/Transforms/Passes.h.inc"
+} // namespace fir
+
+#define DEBUG_TYPE "cuda-heap-alloc-promotion"
+
+namespace {
+class CudaHeapAllocPromotion
+ : public fir::impl::CudaHeapAllocPromotionBase<CudaHeapAllocPromotion> {
+public:
+ using CudaHeapAllocPromotionBase<
+ CudaHeapAllocPromotion>::CudaHeapAllocPromotionBase;
+
+ void runOnOperation() override {
+ mlir::func::FuncOp func = getOperation();
+ if (func.empty())
+ return;
+ mlir::IRRewriter rewriter(&getContext());
+ fir::promoteDynamicVariableAllocasToCudaHeap(rewriter, func.getOperation());
+ }
+};
+} // namespace
diff --git a/flang/lib/Optimizer/Transforms/MemoryAllocation.cpp b/flang/lib/Optimizer/Transforms/MemoryAllocation.cpp
index 058c64cb68395..db1df9874cdf2 100644
--- a/flang/lib/Optimizer/Transforms/MemoryAllocation.cpp
+++ b/flang/lib/Optimizer/Transforms/MemoryAllocation.cpp
@@ -9,7 +9,6 @@
#include "flang/Optimizer/Dialect/FIRDialect.h"
#include "flang/Optimizer/Dialect/FIROps.h"
#include "flang/Optimizer/Dialect/FIRType.h"
-#include "flang/Optimizer/Dialect/Support/FIRContext.h"
#include "flang/Optimizer/Transforms/MemoryUtils.h"
#include "flang/Optimizer/Transforms/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
@@ -59,17 +58,7 @@ keepStackAllocation(fir::AllocaOp alloca,
static mlir::Value genAllocmem(mlir::OpBuilder &builder, fir::AllocaOp alloca,
bool deallocPointsDominateAlloc) {
- mlir::Type varTy = alloca.getInType();
- auto unpackName = [](std::optional<llvm::StringRef> opt) -> llvm::StringRef {
- if (opt)
- return *opt;
- return {};
- };
- llvm::StringRef uniqName = unpackName(alloca.getUniqName());
- llvm::StringRef bindcName = unpackName(alloca.getBindcName());
- auto heap = fir::AllocMemOp::create(builder, alloca.getLoc(), varTy, uniqName,
- bindcName, alloca.getTypeparams(),
- alloca.getShape());
+ fir::AllocMemOp heap = fir::createAllocMemFromAlloca(builder, alloca);
LLVM_DEBUG(llvm::dbgs() << "memory allocation opt: replaced " << alloca
<< " with " << heap << '\n');
return heap;
@@ -132,9 +121,6 @@ class MemoryAllocationOpt
// If func is a declaration, skip it.
if (func.empty())
return;
- mlir::IRRewriter cudaHeapRewriter(context);
- fir::promoteDynamicVariableAllocasToCudaHeap(cudaHeapRewriter,
- func.getOperation());
auto tryReplacing = [&](fir::AllocaOp alloca) {
bool res = !keepStackAllocation(alloca, options);
if (res) {
diff --git a/flang/lib/Optimizer/Transforms/MemoryUtils.cpp b/flang/lib/Optimizer/Transforms/MemoryUtils.cpp
index 4ba436805fd27..d1c457b4d904e 100644
--- a/flang/lib/Optimizer/Transforms/MemoryUtils.cpp
+++ b/flang/lib/Optimizer/Transforms/MemoryUtils.cpp
@@ -315,6 +315,19 @@ bool fir::replaceAllocas(mlir::RewriterBase &rewriter,
return replacedAllRequestedAlloca;
}
+fir::AllocMemOp fir::createAllocMemFromAlloca(mlir::OpBuilder &builder,
+ fir::AllocaOp alloca) {
+ auto unpackName = [](std::optional<llvm::StringRef> opt) -> llvm::StringRef {
+ if (opt)
+ return *opt;
+ return {};
+ };
+ return fir::AllocMemOp::create(builder, alloca.getLoc(), alloca.getInType(),
+ unpackName(alloca.getUniqName()),
+ unpackName(alloca.getBindcName()),
+ alloca.getTypeparams(), alloca.getShape());
+}
+
/// Device code keeps its stack allocations: the unified/managed entry points
/// are host-only, and a kernel-side heap allocation would be a large
/// regression over a device stack array.
@@ -323,7 +336,11 @@ static bool isDeviceCode(mlir::Operation *func, mlir::ModuleOp mod) {
return true;
if (auto procAttr =
func->getAttrOfType<cuf::ProcAttributeAttr>(cuf::getProcAttrName()))
- return procAttr.getValue() != cuf::ProcAttribute::Host;
+ // As in the inDeviceContext helpers of the CUF passes, attributes(host,
+ // device) is not device code here: this is the host copy of the routine,
+ // and its device copy is in the gpu.module handled above.
+ return procAttr.getValue() != cuf::ProcAttribute::Host &&
+ procAttr.getValue() != cuf::ProcAttribute::HostDevice;
if (mlir::acc::isAccRoutine(func))
return true;
if (auto offloadMod =
@@ -361,13 +378,7 @@ bool fir::promoteDynamicVariableAllocasToCudaHeap(mlir::RewriterBase &rewriter,
};
auto genAllocmem = [&](mlir::OpBuilder &builder, fir::AllocaOp alloca,
bool) -> mlir::Value {
- auto name = [](std::optional<llvm::StringRef> opt) -> llvm::StringRef {
- return opt ? *opt : llvm::StringRef{};
- };
- auto heap = fir::AllocMemOp::create(
- builder, alloca.getLoc(), alloca.getInType(),
- name(alloca.getUniqName()), name(alloca.getBindcName()),
- alloca.getTypeparams(), alloca.getShape());
+ fir::AllocMemOp heap = fir::createAllocMemFromAlloca(builder, alloca);
fir::setCudaHeapAllocMode(heap.getOperation(), mode);
// Keep the placement passes from sinking it back to the stack: the
// allocator is chosen here and the matching free is emitted below.
diff --git a/flang/lib/Optimizer/Transforms/StackArrays.cpp b/flang/lib/Optimizer/Transforms/StackArrays.cpp
index 915af1a8d74b9..77861e67a07b1 100644
--- a/flang/lib/Optimizer/Transforms/StackArrays.cpp
+++ b/flang/lib/Optimizer/Transforms/StackArrays.cpp
@@ -15,7 +15,6 @@
#include "flang/Optimizer/Dialect/FIRType.h"
#include "flang/Optimizer/Dialect/Support/FIRContext.h"
#include "flang/Optimizer/Support/DataLayout.h"
-#include "flang/Optimizer/Transforms/MemoryUtils.h"
#include "flang/Optimizer/Transforms/Passes.h"
#include "mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h"
#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
@@ -747,13 +746,6 @@ llvm::StringRef StackArraysPass::getDescription() const {
void StackArraysPass::runOnOperation() {
mlir::func::FuncOp func = getOperation();
- // -fstack-arrays does not apply to the automatic arrays that
- // -gpu=mem:unified|managed must place in unified/managed memory. Done before
- // the analysis below, which skips the fir.must_be_heap pairs it creates.
- mlir::IRRewriter cudaHeapRewriter(&getContext());
- fir::promoteDynamicVariableAllocasToCudaHeap(cudaHeapRewriter,
- func.getOperation());
-
auto &analysis = getAnalysis<fir::StackArraysAnalysisWrapper>();
const fir::StackArraysAnalysisWrapper::AllocMemMap *candidateOps =
analysis.getCandidateOps(func);
diff --git a/flang/test/Driver/bbc-mlir-pass-pipeline.f90 b/flang/test/Driver/bbc-mlir-pass-pipeline.f90
index ae1f5d3c01de4..80328c84794cb 100644
--- a/flang/test/Driver/bbc-mlir-pass-pipeline.f90
+++ b/flang/test/Driver/bbc-mlir-pass-pipeline.f90
@@ -38,6 +38,7 @@
! CHECK-NEXT: (S) 0 num-dce'd - Number of operations DCE'd
! CHECK-NEXT: 'func.func' Pipeline
+! CHECK-NEXT: CudaHeapAllocPromotion
! CHECK-NEXT: MemoryAllocationOpt
! CHECK-NEXT: Inliner
diff --git a/flang/test/Driver/cuda-heap-alloc-promotion-pipeline.f90 b/flang/test/Driver/cuda-heap-alloc-promotion-pipeline.f90
new file mode 100644
index 0000000000000..4e96603fee593
--- /dev/null
+++ b/flang/test/Driver/cuda-heap-alloc-promotion-pipeline.f90
@@ -0,0 +1,13 @@
+! Allocating the dynamically sized automatic variables in unified or managed
+! memory is a correctness requirement of -gpu=mem:unified|managed, so the pass
+! doing it stays in the pipeline even where the array allocation optimization
+! is disabled.
+
+! RUN: %flang_fc1 -S -mmlir --mlir-pass-statistics -mmlir --mlir-pass-statistics-display=pipeline -mmlir -disable-memory-allocation-opt -o /dev/null %s 2>&1 | FileCheck %s
+
+! REQUIRES: asserts
+
+end program
+
+! CHECK: CudaHeapAllocPromotion
+! CHECK-NOT: MemoryAllocationOpt
diff --git a/flang/test/Driver/mlir-debug-pass-pipeline.f90 b/flang/test/Driver/mlir-debug-pass-pipeline.f90
index c103600a22412..75173939ab5df 100644
--- a/flang/test/Driver/mlir-debug-pass-pipeline.f90
+++ b/flang/test/Driver/mlir-debug-pass-pipeline.f90
@@ -75,6 +75,7 @@
! ALL-NEXT: (S) 0 num-dce'd - Number of operations DCE'd
! ALL-NEXT: 'func.func' Pipeline
+! ALL-NEXT: CudaHeapAllocPromotion
! ALL-NEXT: MemoryAllocationOpt
! ALL-NEXT: Inliner
diff --git a/flang/test/Driver/mlir-pass-pipeline.f90 b/flang/test/Driver/mlir-pass-pipeline.f90
index ccf9aa8922040..13910af836186 100644
--- a/flang/test/Driver/mlir-pass-pipeline.f90
+++ b/flang/test/Driver/mlir-pass-pipeline.f90
@@ -127,6 +127,7 @@
! ALL-NEXT: (S) 0 num-dce'd - Number of operations DCE'd
! ALL-NEXT: 'func.func' Pipeline
+! ALL-NEXT: CudaHeapAllocPromotion
! ALL-NEXT: MemoryAllocationOpt
! ALL-NEXT: Inliner
diff --git a/flang/test/Fir/CUDA/cuda-heap-alloc-managed.fir b/flang/test/Fir/CUDA/cuda-heap-alloc-managed.fir
index c6fc559ce5088..7f3220a66c922 100644
--- a/flang/test/Fir/CUDA/cuda-heap-alloc-managed.fir
+++ b/flang/test/Fir/CUDA/cuda-heap-alloc-managed.fir
@@ -1,4 +1,4 @@
-// RUN: fir-opt --memory-allocation-opt %s | FileCheck %s --check-prefix=HEAP
+// RUN: fir-opt --cuda-heap-alloc-promotion %s | FileCheck %s --check-prefix=HEAP
// RUN: fir-opt --fir-to-llvm-ir %s | FileCheck %s --check-prefix=LLVM
// Same routing as cuda-heap-alloc-unified.fir, with the managed entry points.
diff --git a/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir b/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
index ca18116d41ab6..9a50f13cfa192 100644
--- a/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
+++ b/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
@@ -1,8 +1,14 @@
-// RUN: fir-opt --memory-allocation-opt %s | FileCheck %s --check-prefix=HEAP
-// RUN: fir-opt --stack-arrays %s | FileCheck %s --check-prefix=HEAP
+// RUN: fir-opt --cuda-heap-alloc-promotion %s | FileCheck %s --check-prefix=HEAP
// RUN: fir-opt --fir-to-llvm-ir %s | FileCheck %s --check-prefix=LLVM
// RUN: fir-opt --fir-to-llvm-ir=unified-heap-alloc-suffix=_pool %s | FileCheck %s --check-prefix=SUFFIX
+// The promotion is a correctness requirement of the mode, so it does not depend
+// on the array placement pass that follows it: what it produces is marked
+// fir.must_be_heap, which both of them leave alone.
+// RUN: fir-opt --cuda-heap-alloc-promotion --memory-allocation-opt %s | FileCheck %s --check-prefix=HEAP
+// RUN: fir-opt --cuda-heap-alloc-promotion --stack-arrays %s | FileCheck %s --check-prefix=HEAP
+// RUN: fir-opt --cuda-heap-alloc-promotion --allocation-placement %s | FileCheck %s --check-prefix=HEAP
+
// Under fir.cuda_heap_alloc = "unified", named automatic arrays move to the
// heap and are marked. Only marked allocations use malloc_unified: memory the
// Fortran runtime allocated must keep being released by libc free.
@@ -61,6 +67,16 @@ func.func @device_vla(%arg0: index) attributes {cuf.proc_attr = #cuf.cuda_proc<g
return
}
+// attributes(host,device) is the host copy of the routine here, so it needs the
+// unified memory just like any other host code. The device copy of it lives in
+// the gpu.module.
+// HEAP-LABEL: func.func @host_device_vla(
+// HEAP: fir.allocmem !fir.array<?xf32>, %{{.*}} {{{.*}}fir.cuda_heap_alloc = "unified"
+func.func @host_device_vla(%arg0: index) attributes {cuf.proc_attr = #cuf.cuda_proc<host_device>} {
+ %0 = fir.alloca !fir.array<?xf32>, %arg0 {bindc_name = "a", uniq_name = "_QFhost_device_vlaEa"}
+ return
+}
+
// An alloca pinned to the stack stays there: the array function result below is
// replaced by the caller buffer, so a heap pair would only be dead code.
// HEAP-LABEL: func.func @array_result(
diff --git a/flang/test/Fir/basic-program.fir b/flang/test/Fir/basic-program.fir
index 536963920bdb7..fa8f666b7b891 100644
--- a/flang/test/Fir/basic-program.fir
+++ b/flang/test/Fir/basic-program.fir
@@ -110,6 +110,7 @@ func.func @_QQmain() {
// PASSES-NEXT: (S) 0 num-dce'd - Number of operations DCE'd
// PASSES-NEXT: 'func.func' Pipeline
+// PASSES-NEXT: CudaHeapAllocPromotion
// PASSES-NEXT: MemoryAllocationOpt
// PASSES-NEXT: Inliner
>From e405af34b4b11485213959baea31ce4125c5f781 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Wed, 19 Aug 2026 17:05:07 -0700
Subject: [PATCH 7/7] [flang][cuda] Check the promoted pairs survive the array
placement passes
The placement passes have policies of their own for the allocations the
promotion does not touch, so give the combined runs a check prefix that only
asserts what they must preserve.
---
flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir | 15 +++++++++------
1 file changed, 9 insertions(+), 6 deletions(-)
diff --git a/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir b/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
index 9a50f13cfa192..c5c20565c1c73 100644
--- a/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
+++ b/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
@@ -2,12 +2,12 @@
// RUN: fir-opt --fir-to-llvm-ir %s | FileCheck %s --check-prefix=LLVM
// RUN: fir-opt --fir-to-llvm-ir=unified-heap-alloc-suffix=_pool %s | FileCheck %s --check-prefix=SUFFIX
-// The promotion is a correctness requirement of the mode, so it does not depend
-// on the array placement pass that follows it: what it produces is marked
-// fir.must_be_heap, which both of them leave alone.
-// RUN: fir-opt --cuda-heap-alloc-promotion --memory-allocation-opt %s | FileCheck %s --check-prefix=HEAP
-// RUN: fir-opt --cuda-heap-alloc-promotion --stack-arrays %s | FileCheck %s --check-prefix=HEAP
-// RUN: fir-opt --cuda-heap-alloc-promotion --allocation-placement %s | FileCheck %s --check-prefix=HEAP
+// Whichever array placement pass follows, the pairs keep their allocator: they
+// are marked fir.must_be_heap, so none of those passes moves them back to the
+// stack.
+// RUN: fir-opt --cuda-heap-alloc-promotion --memory-allocation-opt %s | FileCheck %s --check-prefix=KEEP
+// RUN: fir-opt --cuda-heap-alloc-promotion --stack-arrays %s | FileCheck %s --check-prefix=KEEP
+// RUN: fir-opt --cuda-heap-alloc-promotion --allocation-placement %s | FileCheck %s --check-prefix=KEEP
// Under fir.cuda_heap_alloc = "unified", named automatic arrays move to the
// heap and are marked. Only marked allocations use malloc_unified: memory the
@@ -24,6 +24,9 @@ module attributes {fir.cuda_heap_alloc = "unified"} {
// HEAP-LABEL: func.func @vla(
// HEAP: %[[MEM:.*]] = fir.allocmem !fir.array<?xf32>, %{{.*}} {bindc_name = "a", fir.cuda_heap_alloc = "unified", fir.must_be_heap = true, uniq_name = "_QFvlaEa"}
// HEAP: fir.freemem %[[MEM]] {fir.cuda_heap_alloc = "unified"} : !fir.heap<!fir.array<?xf32>>
+// KEEP-LABEL: func.func @vla(
+// KEEP: %[[KMEM:.*]] = fir.allocmem !fir.array<?xf32>, %{{.*}} {bindc_name = "a", fir.cuda_heap_alloc = "unified", fir.must_be_heap = true, uniq_name = "_QFvlaEa"}
+// KEEP: fir.freemem %[[KMEM]] {fir.cuda_heap_alloc = "unified"} : !fir.heap<!fir.array<?xf32>>
func.func @vla(%arg0: !fir.ref<i32>) {
%0 = fir.load %arg0 : !fir.ref<i32>
%1 = fir.convert %0 : (i32) -> index
More information about the flang-commits
mailing list