[flang-commits] [flang] [flang][cuda] Defer on_device() folding in host copies of OpenACC routines (PR #208125)
Zhen Wang via flang-commits
flang-commits at lists.llvm.org
Tue Jul 7 17:23:53 PDT 2026
https://github.com/wangzpgi created https://github.com/llvm/llvm-project/pull/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.
>From 5581702108fdbbf3085092a9bbd441cedac0eb7e Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Tue, 7 Jul 2026 10:20:49 -0700
Subject: [PATCH 1/3] Fix on_device() folding to host value in OpenACC routines
under managed/unified memory
---
.../flang/Optimizer/Transforms/Passes.td | 7 +++
.../Transforms/CUDA/CUFFunctionRewrite.cpp | 22 ++++++-
.../test/Fir/CUDA/cuda-function-rewrite.mlir | 63 +++++++++++++++++++
3 files changed, 89 insertions(+), 3 deletions(-)
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 0c02aabfcc935..ed02766e8a4ba 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
@@ -14,6 +14,7 @@
#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"
@@ -45,8 +46,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,
@@ -74,6 +76,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();
@@ -103,6 +117,8 @@ class CallConversion : public OpRewritePattern<fir::CallOp> {
// is recovered independently of external name mangling.
const llvm::StringMap<genFunctionType> genMappings_ = {
{"on_device", &genOnDevice}};
+
+ bool deferAccRoutines_;
};
class CUFFunctionRewrite
@@ -112,7 +128,7 @@ class CUFFunctionRewrite
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.
>From ba3a24856c5444347e4e5a39df3d6a0280239891 Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Tue, 7 Jul 2026 14:55:46 -0700
Subject: [PATCH 2/3] fix
---
flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp | 2 ++
1 file changed, 2 insertions(+)
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
index ed02766e8a4ba..4e709049819af 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
@@ -124,6 +124,8 @@ class CallConversion : public OpRewritePattern<fir::CallOp> {
class CUFFunctionRewrite
: public fir::impl::CUFFunctionRewriteBase<CUFFunctionRewrite> {
public:
+ using CUFFunctionRewriteBase::CUFFunctionRewriteBase;
+
void runOnOperation() override {
auto *ctx = &getContext();
mlir::RewritePatternSet patterns(ctx);
>From f2406b9ae106050cf49061d309ab80aa46a4cbeb Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Tue, 7 Jul 2026 17:19:13 -0700
Subject: [PATCH 3/3] add initializer
---
flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
index 4e709049819af..523b2b6ec347b 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
@@ -118,7 +118,7 @@ class CallConversion : public OpRewritePattern<fir::CallOp> {
const llvm::StringMap<genFunctionType> genMappings_ = {
{"on_device", &genOnDevice}};
- bool deferAccRoutines_;
+ bool deferAccRoutines_ = false;
};
class CUFFunctionRewrite
More information about the flang-commits
mailing list