[flang-commits] [flang] [flang][cuda] Defer on_device() folding in host copies of OpenACC routines (PR #208125)

Tue Jul 7 17:28:49 PDT 2026

llvmorg-github-actions[bot] wrote:




@llvm/pr-subscribers-flang-fir-hlfir

Author: Zhen Wang (wangzpgi)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/208125.diff


3 Files Affected:

- (modified) flang/include/flang/Optimizer/Transforms/Passes.td (+7) 
- (modified) flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp (+21-3) 
- (modified) flang/test/Fir/CUDA/cuda-function-rewrite.mlir (+63) 


``````````diff

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..523b2b6ec347b 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,16 +117,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.

``````````

</details>


https://github.com/llvm/llvm-project/pull/208125
