[flang-commits] [flang] [flang][cuda] do not honor -fstack-arrays in device procedures (PR #223040)
Zhen Wang via flang-commits
flang-commits at lists.llvm.org
Fri Sep 11 12:54:20 PDT 2026
https://github.com/wangzpgi created https://github.com/llvm/llvm-project/pull/223040
The device stack is far smaller than the host one, so an automatic array that fits the host stack easily overflows it. With `-fstack-arrays` (implied by `-Ofast`) automatic arrays in CUDA Fortran `global` and `device` procedures were kept on the device stack and overflowed it at runtime.
Lowering now records a `fir.allocation_policy` with `stack_arrays` disabled on those procedures, and `getAllocationPolicy` honors a policy on the enclosing function instead of only the one on the module. The size based part of the policy is unchanged, so small constant size arrays still go on the device stack. Host procedures, including the host copy of `attributes(host,device)`, are unaffected.
>From 11559969f9a7218860c5850cf22bf12966b99cd0 Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Fri, 11 Sep 2026 12:49:34 -0700
Subject: [PATCH] [flang][cuda] do not honor -fstack-arrays in device
procedures
---
.../Optimizer/Support/AllocationPolicy.h | 16 ++++--
flang/lib/Lower/CallInterface.cpp | 21 ++++++--
.../Optimizer/Support/AllocationPolicy.cpp | 37 ++++++++++----
.../lib/Optimizer/Transforms/StackArrays.cpp | 9 ++++
.../Lower/CUDA/cuda-allocation-policy.cuf | 50 +++++++++++++++++++
.../allocation-policy-precedence.fir | 47 +++++++++++++++++
.../stack-arrays-policy-opt-out.fir | 35 +++++++++++++
7 files changed, 198 insertions(+), 17 deletions(-)
create mode 100644 flang/test/Lower/CUDA/cuda-allocation-policy.cuf
create mode 100644 flang/test/Transforms/stack-arrays-policy-opt-out.fir
diff --git a/flang/include/flang/Optimizer/Support/AllocationPolicy.h b/flang/include/flang/Optimizer/Support/AllocationPolicy.h
index cb4e57628744f..f43dc9d5b55c2 100644
--- a/flang/include/flang/Optimizer/Support/AllocationPolicy.h
+++ b/flang/include/flang/Optimizer/Support/AllocationPolicy.h
@@ -133,12 +133,22 @@ AllocationPolicy getCommandLineAllocationPolicy(bool stackArrays);
/// any policy already recorded there.
void setAllocationPolicy(mlir::ModuleOp mod, const AllocationPolicy &policy);
+/// Record \p policy on \p op, which is meant to be a function that needs a
+/// policy of its own, narrower than the module one.
+void setAllocationPolicy(mlir::Operation *op, const AllocationPolicy &policy);
+
/// Get the policy recorded on \p mod, or the defaults if none was recorded.
AllocationPolicy getAllocationPolicy(mlir::ModuleOp mod);
-/// Get the policy in effect for \p op, which is the one recorded on its
-/// enclosing ModuleOp. Returns the defaults if \p op is not inside a module or
-/// if no policy was recorded.
+/// Get the policy recorded directly on \p op, without looking at its parents.
+/// Use this to tell "this operation opted out" from "nothing was recorded
+/// anywhere", which the defaults cannot express.
+std::optional<AllocationPolicy> getLocalAllocationPolicy(mlir::Operation *op);
+
+/// Get the policy in effect for \p op: the one recorded on the innermost
+/// enclosing operation that carries one, usually the enclosing function if it
+/// has its own policy, otherwise the ModuleOp. Returns the defaults if no
+/// policy was recorded.
AllocationPolicy getAllocationPolicy(mlir::Operation *op);
} // namespace fir
diff --git a/flang/lib/Lower/CallInterface.cpp b/flang/lib/Lower/CallInterface.cpp
index 64960b5e09a98..0709986f692a0 100644
--- a/flang/lib/Lower/CallInterface.cpp
+++ b/flang/lib/Lower/CallInterface.cpp
@@ -20,6 +20,7 @@
#include "flang/Optimizer/Builder/FIRBuilder.h"
#include "flang/Optimizer/Builder/Todo.h"
#include "flang/Optimizer/Dialect/FIROpsSupport.h"
+#include "flang/Optimizer/Support/AllocationPolicy.h"
#include "flang/Optimizer/Support/InternalNames.h"
#include "flang/Optimizer/Support/Utils.h"
#include "flang/Semantics/symbol.h"
@@ -695,10 +696,22 @@ setCUDAAttributes(mlir::func::FuncOp func,
std::optional<Fortran::evaluate::characteristics::Procedure>
characteristic) {
if (characteristic && characteristic->cudaSubprogramAttrs) {
- func.getOperation()->setAttr(
- cuf::getProcAttrName(),
- cuf::getProcAttribute(func.getContext(),
- *characteristic->cudaSubprogramAttrs));
+ auto procAttr = cuf::getProcAttribute(func.getContext(),
+ *characteristic->cudaSubprogramAttrs);
+ func.getOperation()->setAttr(cuf::getProcAttrName(), procAttr);
+ // -fstack-arrays cannot be honored in device code: the device stack is
+ // orders of magnitude smaller, and an automatic array that fits the host
+ // stack easily overflows it. host_device is the host copy of the routine.
+ cuf::ProcAttribute proc = procAttr.getValue();
+ if (proc != cuf::ProcAttribute::Host &&
+ proc != cuf::ProcAttribute::HostDevice) {
+ fir::AllocationPolicy policy =
+ fir::getAllocationPolicy(func.getOperation());
+ if (policy.stackArrays) {
+ policy.stackArrays = false;
+ fir::setAllocationPolicy(func.getOperation(), policy);
+ }
+ }
}
if (sym) {
diff --git a/flang/lib/Optimizer/Support/AllocationPolicy.cpp b/flang/lib/Optimizer/Support/AllocationPolicy.cpp
index ad8eca26f8473..fa0036db9758b 100644
--- a/flang/lib/Optimizer/Support/AllocationPolicy.cpp
+++ b/flang/lib/Optimizer/Support/AllocationPolicy.cpp
@@ -91,12 +91,17 @@ fir::AllocationPolicy fir::getCommandLineAllocationPolicy(bool stackArrays) {
return policy;
}
+void fir::setAllocationPolicy(mlir::Operation *op,
+ const fir::AllocationPolicy &policy) {
+ op->setAttr(allocationPolicyName, fir::AllocationPolicyAttr::get(
+ op->getContext(), policy.stackArrays,
+ policy.smallArrayThresholdBytes,
+ policy.totalStackLimitBytes));
+}
+
void fir::setAllocationPolicy(mlir::ModuleOp mod,
const fir::AllocationPolicy &policy) {
- mod->setAttr(allocationPolicyName, fir::AllocationPolicyAttr::get(
- mod.getContext(), policy.stackArrays,
- policy.smallArrayThresholdBytes,
- policy.totalStackLimitBytes));
+ setAllocationPolicy(mod.getOperation(), policy);
}
fir::AllocationPolicy fir::getAllocationPolicy(mlir::ModuleOp mod) {
@@ -109,11 +114,23 @@ fir::AllocationPolicy fir::getAllocationPolicy(mlir::ModuleOp mod) {
attr.getTotalStackLimit()};
}
+std::optional<fir::AllocationPolicy>
+fir::getLocalAllocationPolicy(mlir::Operation *op) {
+ auto attr =
+ op->getAttrOfType<fir::AllocationPolicyAttr>(allocationPolicyName);
+ if (!attr)
+ return std::nullopt;
+ return fir::AllocationPolicy{attr.getStackArrays(),
+ attr.getSmallArrayThreshold(),
+ attr.getTotalStackLimit()};
+}
+
fir::AllocationPolicy fir::getAllocationPolicy(mlir::Operation *op) {
- auto mod = mlir::dyn_cast<mlir::ModuleOp>(op);
- if (!mod)
- mod = op->getParentOfType<mlir::ModuleOp>();
- if (!mod)
- return fir::AllocationPolicy{};
- return getAllocationPolicy(mod);
+ // The innermost policy wins, so that a function can narrow the module one
+ // (e.g. device code, where the stack is a scarce resource).
+ for (mlir::Operation *cur = op; cur; cur = cur->getParentOp())
+ if (std::optional<fir::AllocationPolicy> policy =
+ fir::getLocalAllocationPolicy(cur))
+ return *policy;
+ return fir::AllocationPolicy{};
}
diff --git a/flang/lib/Optimizer/Transforms/StackArrays.cpp b/flang/lib/Optimizer/Transforms/StackArrays.cpp
index 8c2b39c40762c..7d4e6d49642d3 100644
--- a/flang/lib/Optimizer/Transforms/StackArrays.cpp
+++ b/flang/lib/Optimizer/Transforms/StackArrays.cpp
@@ -14,6 +14,7 @@
#include "flang/Optimizer/Dialect/FIROps.h"
#include "flang/Optimizer/Dialect/FIRType.h"
#include "flang/Optimizer/Dialect/Support/FIRContext.h"
+#include "flang/Optimizer/Support/AllocationPolicy.h"
#include "flang/Optimizer/Support/DataLayout.h"
#include "flang/Optimizer/Transforms/Passes.h"
#include "mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h"
@@ -774,6 +775,14 @@ llvm::StringRef StackArraysPass::getDescription() const {
void StackArraysPass::runOnOperation() {
mlir::func::FuncOp func = getOperation();
+ // This pass only runs under -fstack-arrays, so honor a function that opted
+ // out in its own policy (device code, where the stack is tiny). Functions
+ // without a policy of their own are left to the module setting.
+ if (std::optional<fir::AllocationPolicy> policy =
+ fir::getLocalAllocationPolicy(func))
+ if (!policy->stackArrays)
+ return;
+
auto &analysis = getAnalysis<fir::StackArraysAnalysisWrapper>();
const fir::StackArraysAnalysisWrapper::AllocMemMap *candidateOps =
analysis.getCandidateOps(func);
diff --git a/flang/test/Lower/CUDA/cuda-allocation-policy.cuf b/flang/test/Lower/CUDA/cuda-allocation-policy.cuf
new file mode 100644
index 0000000000000..8749d6b78e7da
--- /dev/null
+++ b/flang/test/Lower/CUDA/cuda-allocation-policy.cuf
@@ -0,0 +1,50 @@
+! Test that -fstack-arrays is not applied to device code. The device stack is
+! far smaller than the host one, so lowering records a policy of its own on
+! device and global procedures to keep their automatic arrays on the heap.
+
+! RUN: %flang_fc1 -emit-fir %s -o - | FileCheck %s --check-prefix=DEFAULT
+! RUN: %flang_fc1 -emit-fir -fstack-arrays %s -o - | FileCheck %s
+
+module m
+contains
+ attributes(global) subroutine kernel(a, n)
+ integer, device :: a(*)
+ integer, value :: n
+ integer :: auto(n)
+ do i = 1, n
+ auto(i) = i
+ end do
+ a(1) = sum(auto(1:n))
+ end subroutine
+
+ attributes(device) subroutine devsub(a, n)
+ integer, device :: a(*)
+ integer, value :: n
+ integer :: auto(n)
+ do i = 1, n
+ auto(i) = i
+ end do
+ a(1) = sum(auto(1:n))
+ end subroutine
+
+ attributes(host) subroutine hostsub(a, n)
+ integer :: a(*)
+ integer :: n
+ integer :: auto(n)
+ do i = 1, n
+ auto(i) = i
+ end do
+ a(1) = sum(auto(1:n))
+ end subroutine
+end module
+
+! Without -fstack-arrays the module policy already keeps automatic arrays on the
+! heap, so no procedure needs a policy of its own.
+! DEFAULT-NOT: func.func{{.*}}fir.allocation_policy
+
+! CHECK: func.func @_QMmPkernel({{.*}}) attributes {cuf.proc_attr = #cuf.cuda_proc<global>, fir.allocation_policy = #fir.allocation_policy<stack_arrays = false,
+! CHECK: func.func @_QMmPdevsub({{.*}}) attributes {cuf.proc_attr = #cuf.cuda_proc<device>, fir.allocation_policy = #fir.allocation_policy<stack_arrays = false,
+
+! The host procedure keeps the module policy, so -fstack-arrays still applies
+! to it.
+! CHECK: func.func @_QMmPhostsub({{.*}}) attributes {cuf.proc_attr = #cuf.cuda_proc<host>}
diff --git a/flang/test/Transforms/allocation-policy-precedence.fir b/flang/test/Transforms/allocation-policy-precedence.fir
index 7f3a9803b8758..10e99cd6ac7de 100644
--- a/flang/test/Transforms/allocation-policy-precedence.fir
+++ b/flang/test/Transforms/allocation-policy-precedence.fir
@@ -60,3 +60,50 @@ func.func @without_policy_attribute() {
return
}
}
+
+// -----
+
+// A policy on the function narrows the module one: the function policy wins for
+// the allocations inside it, the module policy still applies elsewhere.
+//
+// <2000xi32> is 8000 bytes, above the option threshold too, so the function
+// sends it to the heap in both runs.
+
+// CHECK-LABEL: func.func @function_policy
+// CHECK: fir.allocmem !fir.array<2000xi32>
+// CHECK-LABEL: func.func @module_policy
+// CHECK: fir.alloca !fir.array<2000xi32>
+// OPTION-LABEL: func.func @function_policy
+// OPTION: fir.allocmem !fir.array<2000xi32>
+// OPTION-LABEL: func.func @module_policy
+// OPTION: fir.alloca !fir.array<2000xi32>
+module attributes {fir.allocation_policy =
+ #fir.allocation_policy<stack_arrays = true,
+ small_array_threshold = 1024,
+ total_stack_limit = 4194304>,
+ fir.defaultkind = "a1c4d8i4l4r4", fir.kindmap = "",
+ llvm.data_layout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"} {
+func.func @function_policy() attributes {fir.allocation_policy =
+ #fir.allocation_policy<stack_arrays = false,
+ small_array_threshold = 64,
+ total_stack_limit = 4194304>} {
+ %0 = fir.allocmem !fir.array<2000xi32>
+ %c0 = arith.constant 0 : index
+ %v = arith.constant 0 : i32
+ %r = fir.convert %0 : (!fir.heap<!fir.array<2000xi32>>) -> !fir.ref<!fir.array<2000xi32>>
+ %e = fir.coordinate_of %r, %c0 : (!fir.ref<!fir.array<2000xi32>>, index) -> !fir.ref<i32>
+ fir.store %v to %e : !fir.ref<i32>
+ fir.freemem %0 : !fir.heap<!fir.array<2000xi32>>
+ return
+}
+func.func @module_policy() {
+ %0 = fir.allocmem !fir.array<2000xi32>
+ %c0 = arith.constant 0 : index
+ %v = arith.constant 0 : i32
+ %r = fir.convert %0 : (!fir.heap<!fir.array<2000xi32>>) -> !fir.ref<!fir.array<2000xi32>>
+ %e = fir.coordinate_of %r, %c0 : (!fir.ref<!fir.array<2000xi32>>, index) -> !fir.ref<i32>
+ fir.store %v to %e : !fir.ref<i32>
+ fir.freemem %0 : !fir.heap<!fir.array<2000xi32>>
+ return
+}
+}
diff --git a/flang/test/Transforms/stack-arrays-policy-opt-out.fir b/flang/test/Transforms/stack-arrays-policy-opt-out.fir
new file mode 100644
index 0000000000000..52f75a7da398a
--- /dev/null
+++ b/flang/test/Transforms/stack-arrays-policy-opt-out.fir
@@ -0,0 +1,35 @@
+// Test that a function opting out of -fstack-arrays in its own policy keeps its
+// heap allocations, while a function without a policy of its own is unaffected.
+
+// RUN: fir-opt --stack-arrays %s | FileCheck %s
+
+// CHECK-LABEL: func.func @opted_out
+// CHECK: fir.allocmem !fir.array<42xi32>
+// CHECK-NOT: fir.alloca
+func.func @opted_out() attributes {fir.allocation_policy =
+ #fir.allocation_policy<stack_arrays = false,
+ small_array_threshold = 1024,
+ total_stack_limit = 4194304>} {
+ %0 = fir.allocmem !fir.array<42xi32>
+ %c0 = arith.constant 0 : index
+ %v = arith.constant 0 : i32
+ %r = fir.convert %0 : (!fir.heap<!fir.array<42xi32>>) -> !fir.ref<!fir.array<42xi32>>
+ %e = fir.coordinate_of %r, %c0 : (!fir.ref<!fir.array<42xi32>>, index) -> !fir.ref<i32>
+ fir.store %v to %e : !fir.ref<i32>
+ fir.freemem %0 : !fir.heap<!fir.array<42xi32>>
+ return
+}
+
+// CHECK-LABEL: func.func @no_policy
+// CHECK: fir.alloca !fir.array<42xi32>
+// CHECK-NOT: fir.allocmem
+func.func @no_policy() {
+ %0 = fir.allocmem !fir.array<42xi32>
+ %c0 = arith.constant 0 : index
+ %v = arith.constant 0 : i32
+ %r = fir.convert %0 : (!fir.heap<!fir.array<42xi32>>) -> !fir.ref<!fir.array<42xi32>>
+ %e = fir.coordinate_of %r, %c0 : (!fir.ref<!fir.array<42xi32>>, index) -> !fir.ref<i32>
+ fir.store %v to %e : !fir.ref<i32>
+ fir.freemem %0 : !fir.heap<!fir.array<42xi32>>
+ return
+}
More information about the flang-commits
mailing list