[flang-commits] [flang] [flang][cuda] duplicate host_device procedures early and outline the copies under their original name (PR #224400)
Zhen Wang via flang-commits
flang-commits at lists.llvm.org
Thu Sep 17 12:57:28 PDT 2026
https://github.com/wangzpgi created https://github.com/llvm/llvm-project/pull/224400
An `attributes(host,device)` procedure is lowered as one `func.func`, which `cuf-transform-device-func` clones into the GPU module at the end of the FIR pipeline. By then the FIR optimizer has already run on the single body with the host policy, so under `-fstack-arrays` an automatic array stays a `fir.alloca` and overflows the device stack at run time. Procedures reached from device code only through such a procedure, and device bodies inlined into a kernel, have the same problem.
This adds `cuf-duplicate-device-func`, meant to run before the FIR optimizer. It gives every `host_device` procedure, and every procedure without a device attribute that device code reaches through calls or procedure references, a device copy as a `func.func` carrying `cuf.proc_attr = device`, the device allocation policy and a `cuf.device_copy_of` marker naming the original, and redirects references inside device code to the copies. Declarations are not copied. `cuf-transform-device-func` then outlines a copy under its original name, since device symbol names are the cross-unit ABI, restores the original names inside the outlined code and `cuf.kernel` regions, and erases the copies. Without markers it behaves as before.
The device code gathering is factored out of `cuf-transform-device-func` into `cuf::collectDeviceCode` and shared by both passes. `fir::getPresentedFunction` and `fir::getPresentedCallee` let diagnostics and remarks report a copy as the procedure it copies, and `fir::getPresentableFunctionName` uses them. Lowering uses the same `cuf::setDeviceAllocationPolicy` helper as the copies.
>From 9468c243d27563e26225f73753d9391910bdaaa4 Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Thu, 17 Sep 2026 09:22:32 -0700
Subject: [PATCH] [flang][cuda] duplicate host_device procedures early and
outline the copies by their original name
---
.../flang/Optimizer/Builder/CUFCommon.h | 44 ++++++
.../Dialect/CUF/Attributes/CUFAttr.h | 18 +++
flang/include/flang/Optimizer/Support/Utils.h | 13 +-
.../flang/Optimizer/Transforms/Passes.td | 16 +++
flang/lib/Lower/CallInterface.cpp | 7 +-
flang/lib/Optimizer/Builder/CUFCommon.cpp | 105 ++++++++++++++
flang/lib/Optimizer/Support/Utils.cpp | 24 ++++
flang/lib/Optimizer/Transforms/CMakeLists.txt | 1 +
.../CUDA/CUFDeviceFuncTransform.cpp | 132 ++++++------------
.../CUDA/CUFDuplicateDeviceFunc.cpp | 107 ++++++++++++++
.../CUDA/cuda-device-func-transform-aio.mlir | 5 +
.../cuda-device-func-transform-copies.mlir | 81 +++++++++++
.../Fir/CUDA/cuda-duplicate-device-func.mlir | 111 +++++++++++++++
13 files changed, 572 insertions(+), 92 deletions(-)
create mode 100644 flang/lib/Optimizer/Transforms/CUDA/CUFDuplicateDeviceFunc.cpp
create mode 100644 flang/test/Fir/CUDA/cuda-device-func-transform-copies.mlir
create mode 100644 flang/test/Fir/CUDA/cuda-duplicate-device-func.mlir
diff --git a/flang/include/flang/Optimizer/Builder/CUFCommon.h b/flang/include/flang/Optimizer/Builder/CUFCommon.h
index 736f90123969c..fc297a8919126 100644
--- a/flang/include/flang/Optimizer/Builder/CUFCommon.h
+++ b/flang/include/flang/Optimizer/Builder/CUFCommon.h
@@ -10,11 +10,18 @@
#define FORTRAN_OPTIMIZER_TRANSFORMS_CUFCOMMON_H_
#include "flang/Optimizer/Dialect/FIROps.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/IR/BuiltinOps.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SetVector.h"
static constexpr llvm::StringRef cudaDeviceModuleName = "cuda_device_mod";
static constexpr llvm::StringRef cudaSharedMemSuffix = "__shared_mem__";
+/// Appended to the name of the device copy of a procedure while it lives in
+/// the host module, where it cannot share the symbol of the original. The dot
+/// keeps it clear of any Fortran or C identifier.
+static constexpr llvm::StringRef cudaDeviceCopySuffix = ".device";
namespace fir {
class FirOpBuilder;
@@ -33,6 +40,43 @@ bool isCUDADeviceContext(mlir::Region &,
bool isRegisteredDeviceGlobal(fir::GlobalOp op);
bool isRegisteredDeviceAttr(std::optional<cuf::DataAttribute> attr);
+/// True for procedures that have a device side: attributes(device), (global),
+/// (grid_global) and (host,device). Unlike isCUDADeviceContext, host_device
+/// counts, since its body is compiled for the device as well.
+bool isDeviceProcedure(mlir::func::FuncOp funcOp);
+
+/// The device code of a module, as gathered by collectDeviceCode.
+struct DeviceCodeSet {
+ /// Procedures with a device proc attribute (see isDeviceProcedure).
+ llvm::SetVector<mlir::func::FuncOp> deviceFuncs;
+ /// Procedures without one that device code reaches, through calls or
+ /// procedure references, directly or through other such procedures. OpenACC
+ /// routines are excluded, the OpenACC pipeline moves those to the device
+ /// itself. Declarations are included: device code needs them too.
+ llvm::SetVector<mlir::func::FuncOp> calledFromDevice;
+ /// Device procedures referenced from a derived-type binding table, whose
+ /// host symbol must be kept for the table to verify.
+ llvm::SetVector<mlir::func::FuncOp> keepInModule;
+};
+
+/// Gather the device code of \p mod. A host_device procedure that already has
+/// a device copy is host code and is not walked for callees. With
+/// \p rejectDynamicDispatch, a type-bound call with dynamic dispatch in device
+/// code is reported as not yet implemented.
+DeviceCodeSet collectDeviceCode(mlir::ModuleOp mod, mlir::SymbolTable &symTab,
+ bool rejectDynamicDispatch = false);
+
+/// Point every fir.call and fir.address_of under \p root whose symbol is a key
+/// of \p map at the mapped symbol instead.
+void remapProcedureSymbols(
+ mlir::Operation *root,
+ const llvm::DenseMap<mlir::StringAttr, mlir::FlatSymbolRefAttr> &map);
+
+/// Record on \p func the allocation policy of device code: the policy in
+/// effect for it with stack arrays disabled, since the device stack is far
+/// smaller than the host one.
+void setDeviceAllocationPolicy(mlir::Operation *func);
+
void genPointerSync(const mlir::Value box, fir::FirOpBuilder &builder);
int computeElementByteSize(mlir::Location loc, mlir::Type type,
diff --git a/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h b/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h
index 88a47749b8808..6ce3c5ef71636 100644
--- a/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h
+++ b/flang/include/flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h
@@ -15,6 +15,8 @@
#include "flang/Support/Fortran.h"
#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Operation.h"
+#include <optional>
namespace llvm {
class StringRef;
@@ -35,6 +37,22 @@ namespace cuf {
static constexpr llvm::StringRef dataAttrName = "data_attr";
static constexpr llvm::StringRef getDataAttrName() { return "cuf.data_attr"; }
static constexpr llvm::StringRef getProcAttrName() { return "cuf.proc_attr"; }
+/// On the device copy of a procedure, the symbol of the procedure it copies.
+static constexpr llvm::StringRef getDeviceCopyOfAttrName() {
+ return "cuf.device_copy_of";
+}
+/// Mark \p copy as the device copy of the procedure named \p original.
+inline void setDeviceCopyOf(mlir::Operation *copy, llvm::StringRef original) {
+ copy->setAttr(getDeviceCopyOfAttrName(),
+ mlir::FlatSymbolRefAttr::get(copy->getContext(), original));
+}
+/// The name of the procedure \p op is the device copy of, if it is one.
+inline std::optional<llvm::StringRef> getDeviceCopyOf(mlir::Operation *op) {
+ if (auto ref =
+ op->getAttrOfType<mlir::FlatSymbolRefAttr>(getDeviceCopyOfAttrName()))
+ return ref.getValue();
+ return std::nullopt;
+}
/// Attribute to carry CUDA launch_bounds values.
static constexpr llvm::StringRef getLaunchBoundsAttrName() {
diff --git a/flang/include/flang/Optimizer/Support/Utils.h b/flang/include/flang/Optimizer/Support/Utils.h
index d2f0be15d7dd9..ef7bff4e719db 100644
--- a/flang/include/flang/Optimizer/Support/Utils.h
+++ b/flang/include/flang/Optimizer/Support/Utils.h
@@ -255,8 +255,19 @@ mlir::Value integerCast(const fir::LLVMTypeConverter &converter,
/// otherwise it returns std::nullopt.
std::optional<bool> isNewAllocationResult(mlir::OpResult result);
+/// The procedure \p func stands for in diagnostics and remarks: itself, or the
+/// procedure it is a compiler-made copy of (the device copy of a CUDA Fortran
+/// host_device procedure) when that one can be found.
+mlir::FunctionOpInterface getPresentedFunction(mlir::FunctionOpInterface func);
+
+/// Same for the callee of \p call named by \p callee; null if it does not
+/// resolve to a function.
+mlir::FunctionOpInterface getPresentedCallee(mlir::Operation *call,
+ mlir::SymbolRefAttr callee);
+
/// Used to obtain user-facing function name that can be used in
-/// diagnostics and remarks without mangling or underscores.
+/// diagnostics and remarks without mangling or underscores. Compiler-made
+/// copies are reported under the name of the procedure they copy.
std::string getPresentableFunctionName(mlir::FunctionOpInterface func);
} // namespace fir
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index be2cbe31811e8..653f490715707 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -661,6 +661,22 @@ def CUFComputeSharedMemoryOffsetsAndSize
];
}
+def CUFDuplicateDeviceFunc
+ : Pass<"cuf-duplicate-device-func", "::mlir::ModuleOp"> {
+ let summary = "Give host_device procedures a device copy before optimization";
+ let description = [{
+ Creates the device copy of every attributes(host,device) procedure, and of
+ every procedure without a device attribute that device code calls, as a
+ func.func in the host module. The copy carries the device proc attribute,
+ the device allocation policy and a cuf.device_copy_of marker naming the
+ original, and device code is redirected to it while host code keeps the
+ original. Running before the FIR optimizer lets each copy be optimized for
+ its own target; cuf-transform-device-func later moves the copies into the
+ GPU module under their original name.
+ }];
+ let dependentDialects = ["cuf::CUFDialect", "fir::FIROpsDialect"];
+}
+
def CUFDeviceFuncTransform
: Pass<"cuf-transform-device-func", "::mlir::ModuleOp"> {
let summary = "Transform device function to GPU func";
diff --git a/flang/lib/Lower/CallInterface.cpp b/flang/lib/Lower/CallInterface.cpp
index 9885345bed3a5..2bfe94559cd88 100644
--- a/flang/lib/Lower/CallInterface.cpp
+++ b/flang/lib/Lower/CallInterface.cpp
@@ -16,11 +16,11 @@
#include "flang/Lower/PFTBuilder.h"
#include "flang/Lower/StatementContext.h"
#include "flang/Lower/Support/Utils.h"
+#include "flang/Optimizer/Builder/CUFCommon.h"
#include "flang/Optimizer/Builder/Character.h"
#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"
@@ -706,10 +706,7 @@ setCUDAAttributes(mlir::func::FuncOp func,
cuf::ProcAttribute proc = procAttr.getValue();
if (proc != cuf::ProcAttribute::Host &&
proc != cuf::ProcAttribute::HostDevice) {
- fir::AllocationPolicy policy =
- fir::getAllocationPolicy(func.getOperation());
- policy.stackArrays = false;
- fir::setAllocationPolicy(func.getOperation(), policy);
+ cuf::setDeviceAllocationPolicy(func.getOperation());
}
}
diff --git a/flang/lib/Optimizer/Builder/CUFCommon.cpp b/flang/lib/Optimizer/Builder/CUFCommon.cpp
index 601f9fd6e2b31..96b3247ec77bd 100644
--- a/flang/lib/Optimizer/Builder/CUFCommon.cpp
+++ b/flang/lib/Optimizer/Builder/CUFCommon.cpp
@@ -8,11 +8,15 @@
#include "flang/Optimizer/Builder/CUFCommon.h"
#include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "flang/Optimizer/Builder/Todo.h"
#include "flang/Optimizer/Dialect/CUF/CUFOps.h"
#include "flang/Optimizer/Dialect/Support/KindMapping.h"
#include "flang/Optimizer/HLFIR/HLFIROps.h"
+#include "flang/Optimizer/Support/AllocationPolicy.h"
+#include "flang/Optimizer/Support/InternalNames.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "llvm/ADT/StringSet.h"
/// Retrieve or create the CUDA Fortran GPU module in the give in \p mod.
mlir::gpu::GPUModuleOp cuf::getOrCreateGPUModule(mlir::ModuleOp mod,
@@ -156,3 +160,104 @@ mlir::Value cuf::computeElementCount(mlir::PatternRewriter &rewriter,
}
return mlir::Value();
}
+
+bool cuf::isDeviceProcedure(mlir::func::FuncOp funcOp) {
+ auto procAttr =
+ funcOp->getAttrOfType<cuf::ProcAttributeAttr>(cuf::getProcAttrName());
+ if (!procAttr)
+ return false;
+ switch (procAttr.getValue()) {
+ case cuf::ProcAttribute::Device:
+ case cuf::ProcAttribute::Global:
+ case cuf::ProcAttribute::GridGlobal:
+ case cuf::ProcAttribute::HostDevice:
+ return true;
+ case cuf::ProcAttribute::Host:
+ return false;
+ }
+ return false;
+}
+
+cuf::DeviceCodeSet cuf::collectDeviceCode(mlir::ModuleOp mod,
+ mlir::SymbolTable &symTab,
+ bool rejectDynamicDispatch) {
+ cuf::DeviceCodeSet code;
+ llvm::StringSet<> hasDeviceCopy;
+ mod.walk([&](mlir::func::FuncOp funcOp) {
+ if (cuf::isDeviceProcedure(funcOp)) {
+ code.deviceFuncs.insert(funcOp);
+ if (std::optional<llvm::StringRef> original =
+ cuf::getDeviceCopyOf(funcOp))
+ hasDeviceCopy.insert(*original);
+ }
+ });
+
+ // Everything device code reaches without a device attribute of its own,
+ // following calls and procedure references through the procedures found.
+ llvm::SmallVector<mlir::Operation *> worklist;
+ auto found = [&](mlir::StringAttr name) {
+ auto callee = symTab.lookup<mlir::func::FuncOp>(name);
+ if (!callee || mlir::acc::isAccRoutine(callee) ||
+ code.deviceFuncs.count(callee))
+ return;
+ if (code.calledFromDevice.insert(callee) && !callee.isDeclaration())
+ worklist.push_back(callee);
+ };
+ auto scan = [&](mlir::Operation *root) {
+ root->walk([&](mlir::Operation *op) {
+ if (auto call = mlir::dyn_cast<fir::CallOp>(op)) {
+ if (mlir::SymbolRefAttr callee = call.getCalleeAttr())
+ found(callee.getLeafReference());
+ } else if (auto addrOf = mlir::dyn_cast<fir::AddrOfOp>(op)) {
+ found(addrOf.getSymbol().getLeafReference());
+ } else if (rejectDynamicDispatch && mlir::isa<fir::DispatchOp>(op)) {
+ TODO(op->getLoc(),
+ "type-bound procedure call with dynamic dispatch in device code");
+ }
+ });
+ };
+ for (mlir::func::FuncOp funcOp : code.deviceFuncs)
+ if (!hasDeviceCopy.contains(funcOp.getSymName()))
+ worklist.push_back(funcOp);
+ mod.walk([&](cuf::KernelOp kernelOp) { worklist.push_back(kernelOp); });
+ while (!worklist.empty())
+ scan(worklist.pop_back_val());
+
+ // A device procedure referenced from a binding table keeps a host symbol,
+ // or the table fails to verify once lowered to the LLVM dialect.
+ for (fir::GlobalOp globalOp : mod.getOps<fir::GlobalOp>()) {
+ if (!globalOp.getName().contains(fir::kBindingTableSeparator))
+ continue;
+ globalOp.walk([&](fir::AddrOfOp addrOfOp) {
+ auto funcOp = symTab.lookup<mlir::func::FuncOp>(
+ addrOfOp.getSymbol().getLeafReference());
+ if (funcOp && code.deviceFuncs.count(funcOp))
+ code.keepInModule.insert(funcOp);
+ });
+ }
+ return code;
+}
+
+void cuf::remapProcedureSymbols(
+ mlir::Operation *root,
+ const llvm::DenseMap<mlir::StringAttr, mlir::FlatSymbolRefAttr> &map) {
+ if (map.empty())
+ return;
+ root->walk([&](mlir::Operation *op) {
+ if (auto call = mlir::dyn_cast<fir::CallOp>(op)) {
+ if (mlir::SymbolRefAttr callee = call.getCalleeAttr())
+ if (auto it = map.find(callee.getLeafReference()); it != map.end())
+ call.setCalleeAttr(it->second);
+ } else if (auto addrOf = mlir::dyn_cast<fir::AddrOfOp>(op)) {
+ if (auto it = map.find(addrOf.getSymbol().getLeafReference());
+ it != map.end())
+ addrOf.setSymbolAttr(it->second);
+ }
+ });
+}
+
+void cuf::setDeviceAllocationPolicy(mlir::Operation *func) {
+ fir::AllocationPolicy policy = fir::getAllocationPolicy(func);
+ policy.stackArrays = false;
+ fir::setAllocationPolicy(func, policy);
+}
diff --git a/flang/lib/Optimizer/Support/Utils.cpp b/flang/lib/Optimizer/Support/Utils.cpp
index a9590d3073ff2..93dc0c569a97d 100644
--- a/flang/lib/Optimizer/Support/Utils.cpp
+++ b/flang/lib/Optimizer/Support/Utils.cpp
@@ -14,6 +14,7 @@
#include "flang/Optimizer/Dialect/FIROps.h"
#include "flang/Optimizer/Dialect/FIRType.h"
#include "flang/Optimizer/Support/InternalNames.h"
+#include "mlir/IR/SymbolTable.h"
fir::TypeInfoOp fir::lookupTypeInfoOp(fir::RecordType recordType,
mlir::ModuleOp module,
@@ -149,7 +150,30 @@ std::optional<bool> fir::isNewAllocationResult(mlir::OpResult result) {
return false;
}
+mlir::FunctionOpInterface
+fir::getPresentedFunction(mlir::FunctionOpInterface func) {
+ if (std::optional<llvm::StringRef> original =
+ cuf::getDeviceCopyOf(func.getOperation()))
+ if (auto originalFunc = mlir::SymbolTable::lookupNearestSymbolFrom<
+ mlir::FunctionOpInterface>(
+ func.getOperation(),
+ mlir::StringAttr::get(func.getContext(), *original)))
+ return originalFunc;
+ return func;
+}
+
+mlir::FunctionOpInterface fir::getPresentedCallee(mlir::Operation *call,
+ mlir::SymbolRefAttr callee) {
+ if (!callee)
+ return nullptr;
+ auto func =
+ mlir::SymbolTable::lookupNearestSymbolFrom<mlir::FunctionOpInterface>(
+ call, callee);
+ return func ? getPresentedFunction(func) : nullptr;
+}
+
std::string fir::getPresentableFunctionName(mlir::FunctionOpInterface func) {
+ func = getPresentedFunction(func);
if (func.getName() == fir::NameUniquer::doProgramEntry()) {
// Main program entry is all uppercase - to avoid name conflicts. But
// from a reporting perspective, keep it lowercase for consistency with
diff --git a/flang/lib/Optimizer/Transforms/CMakeLists.txt b/flang/lib/Optimizer/Transforms/CMakeLists.txt
index 9b25e590764ab..7fb77ff1ed224 100644
--- a/flang/lib/Optimizer/Transforms/CMakeLists.txt
+++ b/flang/lib/Optimizer/Transforms/CMakeLists.txt
@@ -15,6 +15,7 @@ add_flang_library(FIRTransforms
CUDA/CUFComputeSharedMemoryOffsetsAndSize.cpp
CUDA/CUFDeviceFuncTransform.cpp
CUDA/CUFDeviceGlobal.cpp
+ CUDA/CUFDuplicateDeviceFunc.cpp
CUDA/CUFFunctionRewrite.cpp
CUDA/CUFGPUToLLVMConversion.cpp
CUDA/CUFLaunchAttachAttr.cpp
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFDeviceFuncTransform.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFDeviceFuncTransform.cpp
index a0303a181775e..f5d605fbd4094 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFDeviceFuncTransform.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFDeviceFuncTransform.cpp
@@ -7,20 +7,18 @@
//===----------------------------------------------------------------------===//
#include "flang/Optimizer/Builder/CUFCommon.h"
-#include "flang/Optimizer/Builder/Todo.h"
#include "flang/Optimizer/Dialect/CUF/CUFOps.h"
#include "flang/Optimizer/Dialect/FIRAttr.h"
#include "flang/Optimizer/Dialect/FIRDialect.h"
#include "flang/Optimizer/Dialect/FIRType.h"
-#include "flang/Optimizer/Support/InternalNames.h"
#include "flang/Optimizer/Transforms/Passes.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
-#include "mlir/Dialect/OpenACC/OpenACC.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/StringSet.h"
@@ -85,7 +83,8 @@ class CUFDeviceFuncTransform
}
static gpu::GPUFuncOp createGPUFuncOp(mlir::func::FuncOp funcOp,
- bool isGlobal, int computeCap) {
+ llvm::StringRef name, bool isGlobal,
+ int computeCap) {
mlir::OpBuilder builder(funcOp.getContext());
mlir::Region &funcOpBody = funcOp.getBody();
@@ -107,9 +106,8 @@ class CUFDeviceFuncTransform
mlir::FunctionType type = mlir::FunctionType::get(
funcOp.getContext(), funcOperandTypes, funcResultTypes);
- auto deviceFuncOp =
- gpu::GPUFuncOp::create(builder, loc, funcOp.getName(), type,
- mlir::TypeRange{}, mlir::TypeRange{});
+ auto deviceFuncOp = gpu::GPUFuncOp::create(
+ builder, loc, name, type, mlir::TypeRange{}, mlir::TypeRange{});
if (mlir::ArrayAttr argAttrs = funcOp.getAllArgAttrs())
deviceFuncOp.setAllArgAttrs(argAttrs);
setIntentInKernelArgAttrs(funcOp, deviceFuncOp);
@@ -183,18 +181,6 @@ class CUFDeviceFuncTransform
symTab.insert(emptyStub);
}
- static bool isDeviceFunc(mlir::func::FuncOp funcOp) {
- if (auto cudaProcAttr =
- funcOp.getOperation()->getAttrOfType<cuf::ProcAttributeAttr>(
- cuf::getProcAttrName()))
- if (cudaProcAttr.getValue() == cuf::ProcAttribute::Device ||
- cudaProcAttr.getValue() == cuf::ProcAttribute::Global ||
- cudaProcAttr.getValue() == cuf::ProcAttribute::GridGlobal ||
- cudaProcAttr.getValue() == cuf::ProcAttribute::HostDevice)
- return true;
- return false;
- }
-
void runOnOperation() override {
// Working on Module operation because inserting/removing function from the
// module is not thread-safe.
@@ -207,64 +193,8 @@ class CUFDeviceFuncTransform
gpu::GPUModuleOp gpuMod = cuf::getOrCreateGPUModule(mod, symbolTable);
mlir::SymbolTable gpuModSymTab(gpuMod);
- llvm::SetVector<mlir::func::FuncOp> funcsToClone;
- llvm::SetVector<mlir::func::FuncOp> deviceFuncs;
- llvm::SetVector<mlir::func::FuncOp> keepInModule;
- llvm::StringSet<> deviceFuncNames;
-
- // Look for all function to migrate to the GPU module.
- mod.walk([&](mlir::func::FuncOp op) {
- if (isDeviceFunc(op)) {
- deviceFuncs.insert(op);
- deviceFuncNames.insert(op.getSymName());
- }
- });
-
- auto processCallOp = [&](fir::CallOp op) {
- if (op.getCallee()) {
- auto func = symbolTable.lookup<mlir::func::FuncOp>(
- op.getCallee()->getLeafReference());
- // ACCRoutineToGPUFunc moves the materialized specialized routine into
- // the GPU module later in the pipeline.
- if (mlir::acc::isAccRoutine(func))
- return;
- if (deviceFuncs.count(func) == 0)
- funcsToClone.insert(func);
- }
- };
-
- // Gather all function called by device functions.
- for (auto funcOp : deviceFuncs) {
- funcOp.walk([&](fir::CallOp op) { processCallOp(op); });
- funcOp.walk([&](fir::DispatchOp op) {
- TODO(op.getLoc(), "type-bound procedure call with dynamic dispatch "
- "in device procedure");
- });
- }
-
- // Functions that are referenced in a derived-type binding table must be
- // kept in the host module to avoid LLVM dialect verification errors.
- for (auto globalOp : mod.getOps<fir::GlobalOp>()) {
- if (globalOp.getName().contains(fir::kBindingTableSeparator)) {
- globalOp.walk([&](fir::AddrOfOp addrOfOp) {
- if (deviceFuncNames.contains(addrOfOp.getSymbol().getLeafReference()))
- keepInModule.insert(
- *llvm::find_if(deviceFuncs, [&](mlir::func::FuncOp f) {
- return f.getSymName() ==
- addrOfOp.getSymbol().getLeafReference();
- }));
- });
- }
- }
-
- // Gather all functions called by CUF kernels.
- mod.walk([&](cuf::KernelOp kernelOp) {
- kernelOp.walk([&](fir::CallOp op) { processCallOp(op); });
- kernelOp.walk([&](fir::DispatchOp op) {
- TODO(op.getLoc(),
- "type-bound procedure call with dynamic dispatch in cuf kernel");
- });
- });
+ cuf::DeviceCodeSet code = cuf::collectDeviceCode(
+ mod, symbolTable, /*rejectDynamicDispatch=*/true);
// Optionally report an error when device code calls the runtime function
// _FortranAioOutputDescriptor, which is not supported on the device.
@@ -278,35 +208,65 @@ class CUFDeviceFuncTransform
signalPassFailure();
}
};
- for (auto funcOp : deviceFuncs)
+ for (auto funcOp : code.deviceFuncs)
funcOp.walk(checkForAioOutputDescriptor);
mod.walk([&](cuf::KernelOp kernelOp) {
kernelOp.walk(checkForAioOutputDescriptor);
});
}
- for (auto funcOp : funcsToClone)
- gpuModSymTab.insert(funcOp->clone());
+ // Device copies made by cuf-duplicate-device-func carry the symbol of the
+ // procedure they copy. They take that name back on the device, where
+ // cross-unit references resolve by it, and their originals stay host code.
+ llvm::DenseMap<mlir::StringAttr, mlir::FlatSymbolRefAttr> originalOf;
+ llvm::StringSet<> hasDeviceCopy;
+ for (mlir::func::FuncOp funcOp : code.deviceFuncs)
+ if (std::optional<llvm::StringRef> original =
+ cuf::getDeviceCopyOf(funcOp)) {
+ originalOf[funcOp.getSymNameAttr()] =
+ mlir::FlatSymbolRefAttr::get(ctx, *original);
+ hasDeviceCopy.insert(*original);
+ }
+ // cuf.kernel regions stay in host code until they are lowered, and are
+ // matched to the device by the original names.
+ mod.walk([&](cuf::KernelOp kernelOp) {
+ cuf::remapProcedureSymbols(kernelOp, originalOf);
+ });
+
+ for (auto funcOp : code.calledFromDevice)
+ if (!hasDeviceCopy.contains(funcOp.getSymName()))
+ gpuModSymTab.insert(funcOp->clone());
- for (auto funcOp : deviceFuncs) {
+ for (auto funcOp : code.deviceFuncs) {
auto cudaProcAttr =
funcOp.getOperation()->getAttrOfType<cuf::ProcAttributeAttr>(
cuf::getProcAttrName());
auto isGlobal = cudaProcAttr.getValue() == cuf::ProcAttribute::Global ||
cudaProcAttr.getValue() == cuf::ProcAttribute::GridGlobal;
+ std::optional<llvm::StringRef> copyOf = cuf::getDeviceCopyOf(funcOp);
+ // A host_device original whose device copy exists stays host code.
+ if (!copyOf && hasDeviceCopy.contains(funcOp.getSymName()))
+ continue;
+ llvm::StringRef deviceName = copyOf ? *copyOf : funcOp.getSymName();
if (funcOp.isDeclaration()) {
- mlir::Operation *clonedFuncOp = funcOp->clone();
+ auto clonedFuncOp = mlir::cast<func::FuncOp>(funcOp->clone());
+ if (copyOf) {
+ clonedFuncOp.setSymName(deviceName);
+ clonedFuncOp->removeAttr(cuf::getDeviceCopyOfAttrName());
+ }
if (isGlobal) {
clonedFuncOp->setAttr(gpu::GPUDialect::getKernelFuncAttrName(),
builder.getUnitAttr());
clonedFuncOp->removeAttr(cuf::getProcAttrName());
- if (auto funcOp = mlir::dyn_cast<func::FuncOp>(clonedFuncOp))
- funcOp.setNested();
+ clonedFuncOp.setNested();
}
gpuModSymTab.insert(clonedFuncOp);
+ if (copyOf)
+ funcOp.erase();
} else {
gpu::GPUFuncOp deviceFuncOp =
- createGPUFuncOp(funcOp, isGlobal, computeCap);
+ createGPUFuncOp(funcOp, deviceName, isGlobal, computeCap);
+ cuf::remapProcedureSymbols(deviceFuncOp, originalOf);
gpuModSymTab.insert(deviceFuncOp);
if (cudaProcAttr.getValue() != cuf::ProcAttribute::HostDevice) {
@@ -314,7 +274,7 @@ class CUFDeviceFuncTransform
// declaration for the kernel registration. Currently we just
// erase its body but in the future, the body should be rewritten
// to be able to launch CUDA Fortran kernel from C code.
- if (isGlobal || keepInModule.contains(funcOp))
+ if (isGlobal || code.keepInModule.contains(funcOp))
createHostStub(funcOp, symbolTable, mod);
else
funcOp.erase();
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFDuplicateDeviceFunc.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFDuplicateDeviceFunc.cpp
new file mode 100644
index 0000000000000..d11b9f611c45d
--- /dev/null
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFDuplicateDeviceFunc.cpp
@@ -0,0 +1,107 @@
+//===-- CUFDuplicateDeviceFunc.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/Builder/CUFCommon.h"
+#include "flang/Optimizer/Dialect/CUF/CUFDialect.h"
+#include "flang/Optimizer/Dialect/CUF/CUFOps.h"
+#include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SetVector.h"
+
+namespace fir {
+#define GEN_PASS_DEF_CUFDUPLICATEDEVICEFUNC
+#include "flang/Optimizer/Transforms/Passes.h.inc"
+} // namespace fir
+
+namespace {
+
+class CUFDuplicateDeviceFunc
+ : public fir::impl::CUFDuplicateDeviceFuncBase<CUFDuplicateDeviceFunc> {
+ using CUFDuplicateDeviceFuncBase<
+ CUFDuplicateDeviceFunc>::CUFDuplicateDeviceFuncBase;
+
+ static bool isHostDevice(mlir::func::FuncOp funcOp) {
+ auto procAttr =
+ funcOp->getAttrOfType<cuf::ProcAttributeAttr>(cuf::getProcAttrName());
+ return procAttr && procAttr.getValue() == cuf::ProcAttribute::HostDevice;
+ }
+
+ /// Create the device copy of \p funcOp: same body, device proc attribute, the
+ /// device allocation policy, and a marker naming the original.
+ static mlir::func::FuncOp createDeviceCopy(mlir::func::FuncOp funcOp,
+ llvm::StringRef copyName,
+ mlir::SymbolTable &symTab) {
+ auto copy = mlir::cast<mlir::func::FuncOp>(funcOp->clone());
+ copy.setSymName(copyName);
+ copy->setAttr(cuf::getProcAttrName(),
+ cuf::ProcAttributeAttr::get(funcOp.getContext(),
+ cuf::ProcAttribute::Device));
+ cuf::setDeviceCopyOf(copy, funcOp.getSymName());
+ cuf::setDeviceAllocationPolicy(copy.getOperation());
+ symTab.insert(copy,
+ std::next(mlir::Block::iterator(funcOp.getOperation())));
+ return copy;
+ }
+
+ void runOnOperation() override {
+ mlir::ModuleOp mod = getOperation();
+ mlir::SymbolTable symTab(mod);
+ cuf::DeviceCodeSet code = cuf::collectDeviceCode(mod, symTab);
+
+ // Procedures that need a copy of their own on the device: host_device ones,
+ // whose original stays host code, and the procedures without a device
+ // attribute that device code reaches. Declarations have nothing to
+ // optimize; cuf-transform-device-func clones them by name as before.
+ llvm::SetVector<mlir::func::FuncOp> toCopy;
+ for (mlir::func::FuncOp funcOp : code.deviceFuncs)
+ if (isHostDevice(funcOp) && !funcOp.isDeclaration())
+ toCopy.insert(funcOp);
+ for (mlir::func::FuncOp funcOp : code.calledFromDevice)
+ if (!funcOp.isDeclaration())
+ toCopy.insert(funcOp);
+ if (toCopy.empty())
+ return;
+
+ llvm::DenseMap<mlir::StringAttr, mlir::FlatSymbolRefAttr> copyOf;
+ llvm::SmallVector<mlir::func::FuncOp> copies;
+ for (mlir::func::FuncOp funcOp : toCopy) {
+ std::string copyName = (funcOp.getSymName() + cudaDeviceCopySuffix).str();
+ auto copy = symTab.lookup<mlir::func::FuncOp>(copyName);
+ if (copy && cuf::getDeviceCopyOf(copy) != funcOp.getSymName()) {
+ funcOp.emitError("cannot create the device copy of this procedure: "
+ "the symbol '")
+ << copyName << "' is already taken";
+ signalPassFailure();
+ return;
+ }
+ if (!copy)
+ copy = createDeviceCopy(funcOp, copyName, symTab);
+ copies.push_back(copy);
+ copyOf[funcOp.getSymNameAttr()] =
+ mlir::FlatSymbolRefAttr::get(copy.getSymNameAttr());
+ }
+
+ // Device code refers to the copies, host code keeps the originals. Device
+ // code is every device procedure except the host_device originals, the
+ // copies themselves, and the cuf.kernel regions.
+ for (mlir::func::FuncOp funcOp : code.deviceFuncs)
+ if (!isHostDevice(funcOp))
+ cuf::remapProcedureSymbols(funcOp, copyOf);
+ for (mlir::func::FuncOp copy : copies)
+ cuf::remapProcedureSymbols(copy, copyOf);
+ mod.walk([&](cuf::KernelOp kernelOp) {
+ cuf::remapProcedureSymbols(kernelOp, copyOf);
+ });
+ }
+};
+
+} // end anonymous namespace
diff --git a/flang/test/Fir/CUDA/cuda-device-func-transform-aio.mlir b/flang/test/Fir/CUDA/cuda-device-func-transform-aio.mlir
index 11b283bb5f1c7..d3892f3bdf1d2 100644
--- a/flang/test/Fir/CUDA/cuda-device-func-transform-aio.mlir
+++ b/flang/test/Fir/CUDA/cuda-device-func-transform-aio.mlir
@@ -6,6 +6,11 @@
// RUN: fir-opt --split-input-file --cuf-transform-device-func="check-io-output-descriptor=true" \
// RUN: --verify-diagnostics %s
+// The check must also hold when cuf-duplicate-device-func has run first, as in
+// the driver pipeline: runtime declarations are not given device copies.
+// RUN: fir-opt --split-input-file --cuf-duplicate-device-func \
+// RUN: --cuf-transform-device-func="check-io-output-descriptor=true" --verify-diagnostics %s
+
func.func private @_FortranAioOutputDescriptor(!fir.ref<i8>, !fir.box<none>) -> i1
func.func @_QPsub_aio_device(%arg0: !fir.ref<i8>, %arg1: !fir.box<none>) attributes {cuf.proc_attr = #cuf.cuda_proc<device>} {
diff --git a/flang/test/Fir/CUDA/cuda-device-func-transform-copies.mlir b/flang/test/Fir/CUDA/cuda-device-func-transform-copies.mlir
new file mode 100644
index 0000000000000..3a47dc0cf7676
--- /dev/null
+++ b/flang/test/Fir/CUDA/cuda-device-func-transform-copies.mlir
@@ -0,0 +1,81 @@
+// Device copies made by cuf-duplicate-device-func take their original name back
+// on the device; their host_device originals stay in the host module.
+
+// RUN: fir-opt --cuf-transform-device-func %s | FileCheck %s
+// RUN: fir-opt --cuf-duplicate-device-func --cuf-transform-device-func %s | FileCheck %s
+
+module attributes {fir.defaultkind = "a1c4d8i4l4r4", fir.kindmap = ""} {
+
+func.func @_QPhostdev(%arg0: !fir.ref<i32>) attributes {cuf.proc_attr = #cuf.cuda_proc<host_device>} {
+ %0 = fir.load %arg0 : !fir.ref<i32>
+ return
+}
+func.func @_QPhostdev.device(%arg0: !fir.ref<i32>) attributes {cuf.device_copy_of = @_QPhostdev, cuf.proc_attr = #cuf.cuda_proc<device>} {
+ %0 = fir.load %arg0 : !fir.ref<i32>
+ return
+}
+
+func.func @host_used_in_device() {
+ return
+}
+func.func @host_used_in_device.device() attributes {cuf.device_copy_of = @host_used_in_device, cuf.proc_attr = #cuf.cuda_proc<device>} {
+ return
+}
+
+func.func private @_QMotherPdecl() attributes {cuf.proc_attr = #cuf.cuda_proc<host_device>}
+
+func.func @_QPkernel(%arg0: !fir.ref<i32>) attributes {cuf.proc_attr = #cuf.cuda_proc<global>} {
+ fir.call @_QPhostdev.device(%arg0) : (!fir.ref<i32>) -> ()
+ fir.call @host_used_in_device.device() : () -> ()
+ fir.call @_QMotherPdecl() : () -> ()
+ %0 = fir.address_of(@host_used_in_device.device) : () -> ()
+ return
+}
+
+func.func @_QPhostcaller(%arg0: !fir.ref<i32>) {
+ fir.call @_QPhostdev(%arg0) : (!fir.ref<i32>) -> ()
+ fir.call @host_used_in_device() : () -> ()
+ %c1_i32 = arith.constant 1 : i32
+ %c1 = arith.constant 1 : index
+ cuf.kernel<<<%c1_i32, %c1_i32>>> (%i : index) = (%c1 : index) to (%c1 : index) step (%c1 : index) {
+ fir.call @_QPhostdev.device(%arg0) : (!fir.ref<i32>) -> ()
+ "fir.end"() : () -> ()
+ }
+ return
+}
+
+}
+
+// Host module: originals kept, copies gone.
+// CHECK-LABEL: func.func @_QPhostdev(
+// CHECK-SAME: cuf.proc_attr = #cuf.cuda_proc<host_device>
+// CHECK: fir.load
+// CHECK-NOT: func.func @_QPhostdev.device
+// CHECK: func.func @host_used_in_device()
+// CHECK-NOT: func.func @host_used_in_device.device
+// CHECK: func.func private @_QMotherPdecl()
+
+// The cuf.kernel region refers to the original name again.
+// CHECK-LABEL: func.func @_QPhostcaller(
+// CHECK: fir.call @_QPhostdev(
+// CHECK: fir.call @host_used_in_device()
+// CHECK: cuf.kernel
+// CHECK: fir.call @_QPhostdev(
+
+// GPU module: one function per original name, calls restored.
+// CHECK: gpu.module @cuda_device_mod
+// CHECK: gpu.func @_QPhostdev(
+// CHECK: gpu.func @host_used_in_device()
+// CHECK: func.func private @_QMotherPdecl()
+// CHECK: gpu.func @_QPkernel(
+// CHECK-SAME: kernel
+// CHECK: fir.call @_QPhostdev(
+// CHECK: fir.call @host_used_in_device()
+// CHECK: fir.call @_QMotherPdecl()
+// CHECK: fir.address_of(@host_used_in_device)
+
+// The kernel keeps a host stub for registration, re-inserted after the module.
+// CHECK: func.func @_QPkernel(
+// CHECK-NEXT: return
+// CHECK-NOT: .device
+// CHECK-NOT: cuf.device_copy_of
diff --git a/flang/test/Fir/CUDA/cuda-duplicate-device-func.mlir b/flang/test/Fir/CUDA/cuda-duplicate-device-func.mlir
new file mode 100644
index 0000000000000..a8ce3cc6365c1
--- /dev/null
+++ b/flang/test/Fir/CUDA/cuda-duplicate-device-func.mlir
@@ -0,0 +1,111 @@
+// RUN: fir-opt --cuf-duplicate-device-func %s | FileCheck %s
+
+module attributes {fir.allocation_policy = #fir.allocation_policy<stack_arrays = true, small_array_threshold = 1024, total_stack_limit = 4194304>, fir.defaultkind = "a1c4d8i4l4r4", fir.kindmap = ""} {
+
+// A host_device procedure: the original stays host code, the copy is device
+// code and opts out of -fstack-arrays.
+func.func @_QPhostdev(%arg0: !fir.ref<i32>) attributes {cuf.proc_attr = #cuf.cuda_proc<host_device>} {
+ %0 = fir.load %arg0 : !fir.ref<i32>
+ fir.call @deeper() : () -> ()
+ return
+}
+
+// Reached only through the host_device procedure above: copied as well, and
+// the copy of the caller refers to it.
+func.func @deeper() {
+ return
+}
+
+// Reached only through a procedure reference in device code.
+func.func @by_address() {
+ return
+}
+
+// Called directly by device code.
+func.func @host_used_in_device() {
+ return
+}
+
+// Declarations are not copied; device code keeps calling them by name.
+func.func private @_QMotherPdecl() attributes {cuf.proc_attr = #cuf.cuda_proc<host_device>}
+func.func private @_FortranAioOutputDescriptor(!fir.ref<i8>, !fir.box<none>) -> i1 attributes {fir.runtime}
+
+// Not duplicated: already device code.
+func.func @_QPdevonly() attributes {cuf.proc_attr = #cuf.cuda_proc<device>} {
+ return
+}
+
+// Not duplicated: OpenACC routines are moved by the OpenACC pipeline.
+func.func @acc_routine() attributes {acc.routine_info = #acc.routine_info<[@acc_routine_info]>} {
+ return
+}
+
+// Device code: every reference to a copied procedure is redirected.
+func.func @_QPkernel(%arg0: !fir.ref<i32>, %arg1: !fir.ref<i8>, %arg2: !fir.box<none>) attributes {cuf.proc_attr = #cuf.cuda_proc<global>} {
+ fir.call @_QPhostdev(%arg0) : (!fir.ref<i32>) -> ()
+ fir.call @host_used_in_device() : () -> ()
+ %0 = fir.address_of(@by_address) : () -> ()
+ fir.call @_QMotherPdecl() : () -> ()
+ %1 = fir.call @_FortranAioOutputDescriptor(%arg1, %arg2) : (!fir.ref<i8>, !fir.box<none>) -> i1
+ fir.call @_QPdevonly() : () -> ()
+ fir.call @acc_routine() : () -> ()
+ return
+}
+
+// Host code: references keep the originals.
+func.func @_QPhostcaller(%arg0: !fir.ref<i32>) {
+ fir.call @_QPhostdev(%arg0) : (!fir.ref<i32>) -> ()
+ fir.call @host_used_in_device() : () -> ()
+ %0 = fir.address_of(@by_address) : () -> ()
+ return
+}
+
+}
+
+// CHECK-LABEL: func.func @_QPhostdev(
+// CHECK-SAME: attributes {cuf.proc_attr = #cuf.cuda_proc<host_device>}
+// CHECK: fir.call @deeper()
+
+// CHECK: func.func @_QPhostdev.device(
+// CHECK-SAME: cuf.device_copy_of = @_QPhostdev
+// CHECK-SAME: cuf.proc_attr = #cuf.cuda_proc<device>
+// CHECK-SAME: fir.allocation_policy = #fir.allocation_policy<stack_arrays = false
+// CHECK: fir.load
+// CHECK: fir.call @deeper.device()
+
+// CHECK: func.func @deeper()
+// CHECK-NOT: cuf.device_copy_of
+// CHECK: func.func @deeper.device()
+// CHECK-SAME: cuf.device_copy_of = @deeper
+
+// CHECK: func.func @by_address()
+// CHECK: func.func @by_address.device()
+// CHECK-SAME: cuf.device_copy_of = @by_address
+
+// CHECK: func.func @host_used_in_device()
+// CHECK: func.func @host_used_in_device.device()
+// CHECK-SAME: cuf.device_copy_of = @host_used_in_device
+
+// CHECK: func.func private @_QMotherPdecl()
+// CHECK-NOT: @_QMotherPdecl.device
+// CHECK: func.func private @_FortranAioOutputDescriptor(
+// CHECK-NOT: @_FortranAioOutputDescriptor.device
+
+// CHECK: func.func @_QPdevonly()
+// CHECK-NOT: @_QPdevonly.device
+// CHECK: func.func @acc_routine()
+// CHECK-NOT: @acc_routine.device
+
+// CHECK-LABEL: func.func @_QPkernel(
+// CHECK: fir.call @_QPhostdev.device(
+// CHECK: fir.call @host_used_in_device.device()
+// CHECK: fir.address_of(@by_address.device)
+// CHECK: fir.call @_QMotherPdecl()
+// CHECK: fir.call @_FortranAioOutputDescriptor(
+// CHECK: fir.call @_QPdevonly()
+// CHECK: fir.call @acc_routine()
+
+// CHECK-LABEL: func.func @_QPhostcaller(
+// CHECK: fir.call @_QPhostdev(
+// CHECK: fir.call @host_used_in_device()
+// CHECK: fir.address_of(@by_address)
More information about the flang-commits
mailing list