[flang-commits] [flang] e2d6498 - [flang][cuda] Defer on_device() folding in host copies of OpenACC routines (#208125)
via flang-commits
flang-commits at lists.llvm.org
Wed Jul 8 09:18:08 PDT 2026
Author: Zhen Wang
Date: 2026-07-08T16:18:03Z
New Revision: e2d6498e36235ecd9955d8d375a0c1a253c2e11a
URL: https://github.com/llvm/llvm-project/commit/e2d6498e36235ecd9955d8d375a0c1a253c2e11a
DIFF: https://github.com/llvm/llvm-project/commit/e2d6498e36235ecd9955d8d375a0c1a253c2e11a.diff
LOG: [flang][cuda] Defer on_device() folding in host copies of OpenACC routines (#208125)
Add a `defer-acc-routines` option to `cuf-function-rewrite`. When set,
`on_device()` is not folded in the host copy of an OpenACC routine (has
`acc.routine_info`, not in a `gpu.module`), because that body is later
cloned into the device routine and would otherwise bake in the host
value (`.false.`). A later run folds each copy in its own context. Calls
already in a `gpu.module` are still folded.
Added:
Modified:
flang/include/flang/Optimizer/Transforms/Passes.td
flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
flang/test/Fir/CUDA/cuda-function-rewrite.mlir
Removed:
################################################################################
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index 5617e282ed5ce..4a042926b7a75 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -588,6 +588,13 @@ def CUFDeviceFuncTransform
def CUFFunctionRewrite : Pass<"cuf-function-rewrite", ""> {
let summary = "Convert some CUDA Fortran specific call";
let dependentDialects = ["fir::FIROpsDialect"];
+ let options = [Option<
+ "deferAccRoutines", "defer-acc-routines", "bool", /*default=*/"false",
+ "Do not fold device-detection intrinsics located in the host copy of an "
+ "OpenACC routine. Such a routine is later cloned to build its device "
+ "counterpart, so folding it to the host value here would bake that value "
+ "into the device clone. Deferred calls are folded by a later run that "
+ "executes after device specialization.">];
}
def CUFLaunchAttachAttr : Pass<"cuf-launch-attach-attr", ""> {
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
index 33c65d09b8aca..aaa40f477088e 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
@@ -12,6 +12,8 @@
#include "flang/Optimizer/Support/Utils.h"
#include "flang/Optimizer/Transforms/Passes.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/SymbolTable.h"
@@ -43,8 +45,9 @@ using genFunctionType =
class CallConversion : public OpRewritePattern<fir::CallOp> {
public:
- CallConversion(MLIRContext *context)
- : OpRewritePattern<fir::CallOp>(context) {}
+ CallConversion(MLIRContext *context, bool deferAccRoutines)
+ : OpRewritePattern<fir::CallOp>(context),
+ deferAccRoutines_(deferAccRoutines) {}
LogicalResult
matchAndRewrite(fir::CallOp op,
@@ -72,6 +75,18 @@ class CallConversion : public OpRewritePattern<fir::CallOp> {
if (!func.isExternal())
return failure();
+ // Defer folding in the host copy of an OpenACC routine. Device
+ // specialization later clones the host body to build the device routine, so
+ // folding it to the host value now would bake that value into the device
+ // clone. A later run (after specialization) folds each copy in its own
+ // host/device context. Calls already inside a gpu.module are device copies
+ // and are always safe to fold.
+ if (deferAccRoutines_ && !op->getParentOfType<gpu::GPUModuleOp>()) {
+ if (auto enclosing = op->getParentOfType<mlir::FunctionOpInterface>())
+ if (mlir::acc::isAccRoutine(enclosing))
+ return failure();
+ }
+
mlir::Value result = fct->second(rewriter, op);
if (!result)
return failure();
@@ -101,16 +116,20 @@ class CallConversion : public OpRewritePattern<fir::CallOp> {
// is recovered independently of external name mangling.
const llvm::StringMap<genFunctionType> genMappings_ = {
{"on_device", &genOnDevice}};
+
+ bool deferAccRoutines_ = false;
};
class CUFFunctionRewrite
: public fir::impl::CUFFunctionRewriteBase<CUFFunctionRewrite> {
public:
+ using CUFFunctionRewriteBase::CUFFunctionRewriteBase;
+
void runOnOperation() override {
auto *ctx = &getContext();
mlir::RewritePatternSet patterns(ctx);
- patterns.insert<CallConversion>(patterns.getContext());
+ patterns.insert<CallConversion>(patterns.getContext(), deferAccRoutines);
if (mlir::failed(
mlir::applyPatternsGreedily(getOperation(), std::move(patterns)))) {
diff --git a/flang/test/Fir/CUDA/cuda-function-rewrite.mlir b/flang/test/Fir/CUDA/cuda-function-rewrite.mlir
index 649064a46f54a..862b3d1ea4960 100644
--- a/flang/test/Fir/CUDA/cuda-function-rewrite.mlir
+++ b/flang/test/Fir/CUDA/cuda-function-rewrite.mlir
@@ -1,4 +1,5 @@
// RUN: fir-opt --split-input-file --cuf-function-rewrite %s | FileCheck %s
+// RUN: fir-opt --split-input-file --cuf-function-rewrite="defer-acc-routines=true" %s | FileCheck %s --check-prefix=DEFER
// Test the bind(c) name "on_device" in device context.
gpu.module @cuda_device_mod {
@@ -138,6 +139,68 @@ func.func @_QMmtestPsub_extname_host() {
// CHECK-LABEL: func.func @_QMmtestPsub_extname_host
// CHECK: fir.if %false
+// A plain host function (not an OpenACC routine) is still folded to .false.
+// even with defer-acc-routines, which only defers OpenACC routine host copies.
+// DEFER-LABEL: func.func @_QMmtestPsub_extname_host
+// DEFER: fir.if %false
+
+// -----
+
+// Host copy of an OpenACC routine. Folded to .false. by default, but with
+// defer-acc-routines the call is left in place so the later device
+// specialization clones an unfolded body (each copy is folded in its own
+// host/device context by a subsequent run).
+func.func private @on_device_() -> !fir.logical<4> attributes {fir.internal_name = "_QPon_device"}
+func.func @_QMmtestPaccroutine_host() attributes {acc.routine_info = #acc.routine_info<[@acc_routine_0]>} {
+ %c2_i32 = arith.constant 2 : i32
+ %c1_i32 = arith.constant 1 : i32
+ %0 = fir.alloca i32
+ %13 = fir.call @on_device_() fastmath<contract> : () -> !fir.logical<4>
+ %14 = fir.convert %13 : (!fir.logical<4>) -> i1
+ fir.if %14 {
+ fir.store %c1_i32 to %0 : !fir.ref<i32>
+ } else {
+ fir.store %c2_i32 to %0 : !fir.ref<i32>
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @_QMmtestPaccroutine_host
+// CHECK: fir.if %false
+
+// DEFER-LABEL: func.func @_QMmtestPaccroutine_host
+// DEFER: fir.call @on_device_()
+
+// -----
+
+// Device copy (inside gpu.module) of an OpenACC routine is always folded to
+// .true., even with defer-acc-routines, because it is already in its final
+// device placement.
+gpu.module @acc_routine_device_mod {
+ func.func private @on_device_() -> !fir.logical<4> attributes {fir.internal_name = "_QPon_device"}
+ func.func @_QMmtestPaccroutine_device() attributes {acc.routine_info = #acc.routine_info<[@acc_routine_0]>} {
+ %c2_i32 = arith.constant 2 : i32
+ %c1_i32 = arith.constant 1 : i32
+ %0 = fir.alloca i32
+ %13 = fir.call @on_device_() fastmath<contract> : () -> !fir.logical<4>
+ %14 = fir.convert %13 : (!fir.logical<4>) -> i1
+ fir.if %14 {
+ fir.store %c1_i32 to %0 : !fir.ref<i32>
+ } else {
+ fir.store %c2_i32 to %0 : !fir.ref<i32>
+ }
+ return
+ }
+}
+
+// CHECK-LABEL: gpu.module @acc_routine_device_mod
+// CHECK: func.func @_QMmtestPaccroutine_device
+// CHECK: fir.if %true
+
+// DEFER-LABEL: gpu.module @acc_routine_device_mod
+// DEFER: func.func @_QMmtestPaccroutine_device
+// DEFER: fir.if %true
+
// -----
// A user-defined procedure named on_device (with a body) must not be folded.
More information about the flang-commits
mailing list