[flang-commits] [flang] cd1b0d9 - [flang][cuda] Route dynamic autos through malloc_unified/free_unified (#212965)
via flang-commits
flang-commits at lists.llvm.org
Thu Aug 20 09:03:15 PDT 2026
Author: Matsu
Date: 2026-08-20T09:03:09-07:00
New Revision: cd1b0d94cd58b21224bffe86d77fd6c6fdc5e371
URL: https://github.com/llvm/llvm-project/commit/cd1b0d94cd58b21224bffe86d77fd6c6fdc5e371
DIFF: https://github.com/llvm/llvm-project/commit/cd1b0d94cd58b21224bffe86d77fd6c6fdc5e371.diff
LOG: [flang][cuda] Route dynamic autos through malloc_unified/free_unified (#212965)
Example:
```fortran
subroutine work(n)
integer :: n
real :: a(n)
call compute(a)
end subroutine
```
In this code, `a` is an automatic array placed on the stack. Under
`-gpu=mem:unified|managed` it must come from the unified/managed
allocator
entry points instead. Renaming every host malloc/free in the module
would be
unsafe: `fir.freemem` also releases buffers the Fortran runtime
allocated with
libc malloc (transformational intrinsic results, polymorphic
temporaries), and
inline ALLOCATE memory can be released by the runtime.
Fix: record the mode on the module at lowering; in the allocation
placement
passes, move named dynamic-size locals of host functions to
`fir.allocmem`/`fir.freemem` pairs marked with that mode; lower only
marked
pairs to those entry points, so runtime-allocated memory keeps libc
free.
Device code, compiler temporaries, and unmarked allocations are
unaffected.
Added:
flang/lib/Optimizer/Transforms/CudaHeapAllocPromotion.cpp
flang/test/Driver/cuda-heap-alloc-promotion-pipeline.f90
flang/test/Fir/CUDA/cuda-heap-alloc-managed.fir
flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
flang/test/Lower/CUDA/cuda-gpu-unified-automatic-array.f90
Modified:
flang/include/flang/Optimizer/CodeGen/CGPasses.td
flang/include/flang/Optimizer/CodeGen/CodeGen.h
flang/include/flang/Optimizer/Dialect/Support/FIRContext.h
flang/include/flang/Optimizer/Transforms/MemoryUtils.h
flang/include/flang/Optimizer/Transforms/Passes.td
flang/lib/Lower/Bridge.cpp
flang/lib/Optimizer/CodeGen/CodeGen.cpp
flang/lib/Optimizer/Dialect/Support/FIRContext.cpp
flang/lib/Optimizer/Passes/Pipelines.cpp
flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
flang/lib/Optimizer/Transforms/CMakeLists.txt
flang/lib/Optimizer/Transforms/MemoryAllocation.cpp
flang/lib/Optimizer/Transforms/MemoryUtils.cpp
flang/test/Driver/bbc-mlir-pass-pipeline.f90
flang/test/Driver/mlir-debug-pass-pipeline.f90
flang/test/Driver/mlir-pass-pipeline.f90
flang/test/Fir/basic-program.fir
Removed:
################################################################################
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/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..61cabd1df584e 100644
--- a/flang/include/flang/Optimizer/Transforms/MemoryUtils.h
+++ b/flang/include/flang/Optimizer/Transforms/MemoryUtils.h
@@ -57,6 +57,20 @@ 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. This is what the cuda-heap-alloc-promotion pass runs.
+bool promoteDynamicVariableAllocasToCudaHeap(mlir::RewriterBase &rewriter,
+ mlir::Operation *func);
+
} // namespace fir
#endif // FORTRAN_OPTIMIZER_TRANSFORMS_MEMORYUTILS_H
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/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index f9becabaa1fe2..9aa11faf1bc14 100644
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -6901,6 +6901,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..7696f5f900c5e 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,43 @@ getMallocInModule(ModuleOp mod, fir::AllocMemOp op,
return mlir::SymbolRefAttr::get(mallocDecl);
}
+/// 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.str();
+ switch (fir::getCudaHeapAllocMode(op)) {
+ case fir::CudaHeapAllocMode::Unified:
+ return (plain + options.unifiedHeapAllocSuffix).str();
+ case fir::CudaHeapAllocMode::Managed:
+ return (plain + options.managedHeapAllocSuffix).str();
+ case fir::CudaHeapAllocMode::None:
+ return plain.str();
+ }
+ llvm_unreachable("unexpected CudaHeapAllocMode");
+}
+
/// Return the LLVMFuncOp corresponding to the standard malloc call.
static mlir::SymbolRefAttr getMalloc(fir::AllocMemOp op,
mlir::ConversionPatternRewriter &rewriter,
- mlir::Type indexType) {
+ 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);
+ 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);
@@ -1345,19 +1366,19 @@ getAlignedAllocInModule(ModuleOp mod, fir::AllocMemOp op,
static mlir::SymbolRefAttr
getAlignedAlloc(fir::AllocMemOp op, mlir::ConversionPatternRewriter &rewriter,
- mlir::Type indexType) {
+ 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);
+ 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);
@@ -1378,11 +1399,13 @@ getPosixMemalignInModule(ModuleOp mod, fir::AllocMemOp op,
static mlir::SymbolRefAttr
getPosixMemalign(fir::AllocMemOp op, mlir::ConversionPatternRewriter &rewriter,
- mlir::Type indexType) {
+ 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);
+ 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
@@ -1465,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{
@@ -1487,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},
@@ -1496,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));
@@ -1519,8 +1544,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))
@@ -1540,11 +1565,13 @@ getFreeInModule(ModuleOp mod, fir::FreeMemOp op,
}
static mlir::SymbolRefAttr getFree(fir::FreeMemOp op,
- mlir::ConversionPatternRewriter &rewriter) {
+ 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);
+ 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) {
@@ -1566,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()},
@@ -4747,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/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/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 c8ad34f7f60ad..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;
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 fd1d566ca2825..db1df9874cdf2 100644
--- a/flang/lib/Optimizer/Transforms/MemoryAllocation.cpp
+++ b/flang/lib/Optimizer/Transforms/MemoryAllocation.cpp
@@ -58,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;
diff --git a/flang/lib/Optimizer/Transforms/MemoryUtils.cpp b/flang/lib/Optimizer/Transforms/MemoryUtils.cpp
index d1b27d9872ea1..d1c457b4d904e 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,84 @@ bool fir::replaceAllocas(mlir::RewriterBase &rewriter,
rewriter.restoreInsertionPoint(insertPoint);
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.
+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()))
+ // 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 =
+ llvm::dyn_cast<mlir::omp::OffloadModuleInterface>(mod.getOperation()))
+ return offloadMod.getIsTargetDevice();
+ return false;
+}
+
+bool fir::promoteDynamicVariableAllocasToCudaHeap(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;
+ // 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 uniqName && !uniqName->empty();
+ };
+ auto genAllocmem = [&](mlir::OpBuilder &builder, fir::AllocaOp alloca,
+ bool) -> mlir::Value {
+ 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.
+ 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/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
new file mode 100644
index 0000000000000..7f3220a66c922
--- /dev/null
+++ b/flang/test/Fir/CUDA/cuda-heap-alloc-managed.fir
@@ -0,0 +1,44 @@
+// 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.
+
+// 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..c5c20565c1c73
--- /dev/null
+++ b/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
@@ -0,0 +1,120 @@
+// 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
+
+// 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
+// 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>>
+// 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
+ %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
+}
+
+// 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(
+// 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(
+// 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>>
+ 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/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
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..c3c15c98ceed0
--- /dev/null
+++ b/flang/test/Lower/CUDA/cuda-gpu-unified-automatic-array.f90
@@ -0,0 +1,54 @@
+! 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, 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-NOT: cuf.alloc
+! CHECK-NOT: data_attr = #cuf.cuda
+! CHECK: fir.alloca !fir.array<?xf32>
+subroutine vla(n)
+ integer :: n
+ real :: a(n)
+ a(1) = 1.0
+end subroutine
+
+! CHECK-LABEL: func.func @_QPadjustable(
+! 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 remain ordinary stack allocations.
+! CHECK-LABEL: func.func @_QPfixed(
+! CHECK-NOT: cuf.alloc
+! CHECK: fir.alloca !fir.array<128xf32>
+subroutine fixed
+ real :: a(128)
+ a(1) = 1.0
+end subroutine
+
+! Dummy adjustable arrays are caller-allocated.
+! CHECK-LABEL: func.func @_QPdummy_adj(
+! CHECK-NOT: cuf.alloc
+! CHECK-NOT: fir.alloca !fir.array<?xf32>
+subroutine dummy_adj(a, n)
+ integer :: n
+ real :: a(n)
+ a(1) = 1.0
+end subroutine
More information about the flang-commits
mailing list