[flang-commits] [flang] [flang] Let -fstack-arrays win over CudaHeapAllocPromotion under mem:unified (PR #218976)
via flang-commits
flang-commits at lists.llvm.org
Wed Aug 26 09:54:37 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-flang-fir-hlfir
Author: Matsu (khaki3)
<details>
<summary>Changes</summary>
Example:
```fortran
subroutine foo(n)
integer, intent(in) :: n
real :: tmp(n)
```
Compiled with -gpu=mem:unified -fstack-arrays, tmp still ends up in malloc_unified: CudaHeapAllocPromotion marks the allocation fir.must_be_heap, which StackArrays then skips, so -fstack-arrays is silently dropped. Under mem:unified the stack is device accessible, so the promotion is a placement choice there. Under mem:managed only the managed allocator is, so it stays a correctness requirement.
Fix: add a stack-arrays option to the pass and skip the promotion when it is set and the mode is unified.
---
Full diff: https://github.com/llvm/llvm-project/pull/218976.diff
7 Files Affected:
- (modified) flang/include/flang/Optimizer/Transforms/MemoryUtils.h (+6-3)
- (modified) flang/include/flang/Optimizer/Transforms/Passes.td (+10)
- (modified) flang/lib/Optimizer/Passes/Pipelines.cpp (+2-1)
- (modified) flang/lib/Optimizer/Transforms/CudaHeapAllocPromotion.cpp (+2-1)
- (modified) flang/lib/Optimizer/Transforms/MemoryUtils.cpp (+6-1)
- (modified) flang/test/Fir/CUDA/cuda-heap-alloc-managed.fir (+4)
- (modified) flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir (+6)
``````````diff
diff --git a/flang/include/flang/Optimizer/Transforms/MemoryUtils.h b/flang/include/flang/Optimizer/Transforms/MemoryUtils.h
index 61cabd1df584e..a8d928a4aa9c7 100644
--- a/flang/include/flang/Optimizer/Transforms/MemoryUtils.h
+++ b/flang/include/flang/Optimizer/Transforms/MemoryUtils.h
@@ -66,10 +66,13 @@ fir::AllocMemOp createAllocMemFromAlloca(mlir::OpBuilder &builder,
/// 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.
+/// keeps its stack allocations, are left alone. With \p stackArrays,
+/// mem:unified keeps them on the stack too, which mem:managed cannot do.
+/// 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);
+ mlir::Operation *func,
+ bool stackArrays = false);
} // namespace fir
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index f9567c2c91e43..ef18cb33d4b8a 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -301,8 +301,18 @@ def CudaHeapAllocPromotion
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.
+
+ Under mem:unified every host address is device accessible, including the
+ stack, so there the promotion is a placement choice that -fstack-arrays
+ opts out of. Under mem:managed only what the managed allocator hands out
+ is, so the promotion always happens.
}];
let dependentDialects = ["fir::FIROpsDialect"];
+ let options = [
+ Option<"stackArrays", "stack-arrays", "bool", /*default=*/"false",
+ "Keep the automatic arrays of mem:unified on the stack, as "
+ "requested by -fstack-arrays.">
+ ];
}
def MemoryAllocationOpt : Pass<"memory-allocation-opt", "mlir::func::FuncOp"> {
diff --git a/flang/lib/Optimizer/Passes/Pipelines.cpp b/flang/lib/Optimizer/Passes/Pipelines.cpp
index 232b27a148ab2..f39a798c1bc0e 100644
--- a/flang/lib/Optimizer/Passes/Pipelines.cpp
+++ b/flang/lib/Optimizer/Passes/Pipelines.cpp
@@ -208,7 +208,8 @@ void createDefaultFIROptimizerPassPipeline(mlir::PassManager &pm,
// -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());
+ pm.addPass(fir::createCudaHeapAllocPromotion(
+ fir::CudaHeapAllocPromotionOptions{pc.StackArrays}));
if (enableAllocationPlacement)
fir::addAllocationPlacement(pm, pc.StackArrays);
diff --git a/flang/lib/Optimizer/Transforms/CudaHeapAllocPromotion.cpp b/flang/lib/Optimizer/Transforms/CudaHeapAllocPromotion.cpp
index da4b15c43ba44..bcf1afa4b870a 100644
--- a/flang/lib/Optimizer/Transforms/CudaHeapAllocPromotion.cpp
+++ b/flang/lib/Optimizer/Transforms/CudaHeapAllocPromotion.cpp
@@ -32,7 +32,8 @@ class CudaHeapAllocPromotion
if (func.empty())
return;
mlir::IRRewriter rewriter(&getContext());
- fir::promoteDynamicVariableAllocasToCudaHeap(rewriter, func.getOperation());
+ fir::promoteDynamicVariableAllocasToCudaHeap(rewriter, func.getOperation(),
+ stackArrays);
}
};
} // namespace
diff --git a/flang/lib/Optimizer/Transforms/MemoryUtils.cpp b/flang/lib/Optimizer/Transforms/MemoryUtils.cpp
index d1c457b4d904e..00a25a7326f2d 100644
--- a/flang/lib/Optimizer/Transforms/MemoryUtils.cpp
+++ b/flang/lib/Optimizer/Transforms/MemoryUtils.cpp
@@ -350,13 +350,18 @@ static bool isDeviceCode(mlir::Operation *func, mlir::ModuleOp mod) {
}
bool fir::promoteDynamicVariableAllocasToCudaHeap(mlir::RewriterBase &rewriter,
- mlir::Operation *func) {
+ mlir::Operation *func,
+ bool stackArrays) {
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;
+ // The stack is device accessible under unified memory, so -fstack-arrays can
+ // be honored there. Under managed memory only the allocator can be.
+ if (stackArrays && mode == fir::CudaHeapAllocMode::Unified)
+ return false;
bool changed = false;
// User variables only: automatic arrays and automatic character, which are
diff --git a/flang/test/Fir/CUDA/cuda-heap-alloc-managed.fir b/flang/test/Fir/CUDA/cuda-heap-alloc-managed.fir
index 7f3220a66c922..49e7467bc4e11 100644
--- a/flang/test/Fir/CUDA/cuda-heap-alloc-managed.fir
+++ b/flang/test/Fir/CUDA/cuda-heap-alloc-managed.fir
@@ -1,6 +1,10 @@
// 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
+// Only the managed allocator is device accessible here, so -fstack-arrays
+// cannot keep the automatic arrays on the stack.
+// RUN: fir-opt --cuda-heap-alloc-promotion=stack-arrays=true %s | FileCheck %s --check-prefix=HEAP
+
// 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.
diff --git a/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir b/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
index c5c20565c1c73..602217482c16d 100644
--- a/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
+++ b/flang/test/Fir/CUDA/cuda-heap-alloc-unified.fir
@@ -9,6 +9,12 @@
// 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
+// The stack is device accessible here, so -fstack-arrays can keep the automatic
+// arrays on it.
+// RUN: fir-opt --cuda-heap-alloc-promotion=stack-arrays=true %s | FileCheck %s --check-prefix=STACK
+// STACK-NOT: fir.allocmem !fir.array<?xf32>, %{{.*}} {bindc_name
+// STACK-NOT: fir.allocmem !fir.char
+
// 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.
``````````
</details>
https://github.com/llvm/llvm-project/pull/218976
More information about the flang-commits
mailing list