[flang-commits] [flang] ec1b3b1 - [flang][cuda] Add cuf.on_device op and lower to it (#227089)

via flang-commits flang-commits at lists.llvm.org
Mon Sep 28 19:57:19 PDT 2026


Author: Valentin Clement (バレンタイン クレメン)
Date: 2026-09-28T19:57:13-07:00
New Revision: ec1b3b10a71c2181b0dfb90f1178e2e739680e02

URL: https://github.com/llvm/llvm-project/commit/ec1b3b10a71c2181b0dfb90f1178e2e739680e02
DIFF: https://github.com/llvm/llvm-project/commit/ec1b3b10a71c2181b0dfb90f1178e2e739680e02.diff

LOG: [flang][cuda] Add cuf.on_device op and lower to it (#227089)

Replace the CUFFunctionRewrite pass, which matched fir.call names and
folded them to constants, with a cuf.on_device operation emitted when
the intrinsic is lowered. CUFOpConversion folds the operation once the
code is in its host or device context, and leaves the host copy of an
OpenACC routine unfolded so the device clone is not baked to false. Any
operation that remains is folded by the late CUF conversion, and the
rewrite pass is dropped from the pipeline.

Added: 
    flang/test/Fir/CUDA/cuda-on-device.mlir
    flang/test/Lower/CUDA/cuda-on-device.cuf

Modified: 
    flang/include/flang/Optimizer/Builder/CUDAIntrinsicCall.h
    flang/include/flang/Optimizer/Builder/CUFCommon.h
    flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
    flang/include/flang/Optimizer/Transforms/Passes.td
    flang/lib/Optimizer/Builder/CUDAIntrinsicCall.cpp
    flang/lib/Optimizer/Builder/CUFCommon.cpp
    flang/lib/Optimizer/Transforms/CMakeLists.txt
    flang/lib/Optimizer/Transforms/CUDA/CUFOpConversion.cpp
    flang/lib/Optimizer/Transforms/CUDA/CUFOpConversionLate.cpp
    flang/test/Fir/cuf.mlir

Removed: 
    flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
    flang/test/Fir/CUDA/cuda-function-rewrite.mlir


################################################################################
diff  --git a/flang/include/flang/Optimizer/Builder/CUDAIntrinsicCall.h b/flang/include/flang/Optimizer/Builder/CUDAIntrinsicCall.h
index 2974d774b2d1a..09eb8b474831c 100644
--- a/flang/include/flang/Optimizer/Builder/CUDAIntrinsicCall.h
+++ b/flang/include/flang/Optimizer/Builder/CUDAIntrinsicCall.h
@@ -75,6 +75,7 @@ struct CUDAIntrinsicLibrary : IntrinsicLibrary {
   mlir::Value genMatchAnySync(mlir::Type, llvm::ArrayRef<mlir::Value>);
   template <typename OpTy>
   mlir::Value genNVVMTime(mlir::Type, llvm::ArrayRef<mlir::Value>);
+  mlir::Value genOnDevice(mlir::Type, llvm::ArrayRef<mlir::Value>);
   void genSyncThreads(llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genSyncThreadsAnd(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genSyncThreadsCount(mlir::Type, llvm::ArrayRef<mlir::Value>);

diff  --git a/flang/include/flang/Optimizer/Builder/CUFCommon.h b/flang/include/flang/Optimizer/Builder/CUFCommon.h
index fc297a8919126..e1f8ae454b68e 100644
--- a/flang/include/flang/Optimizer/Builder/CUFCommon.h
+++ b/flang/include/flang/Optimizer/Builder/CUFCommon.h
@@ -37,6 +37,13 @@ mlir::gpu::GPUModuleOp getOrCreateGPUModule(mlir::ModuleOp mod,
 bool isCUDADeviceContext(mlir::Operation *op);
 bool isCUDADeviceContext(mlir::Region &,
                          bool isDoConcurrentOffloadEnabled = false);
+
+/// True when \p op is compiled for device execution: a CUDA device procedure,
+/// a GPU module or launch, or an OpenACC offload region. Host and host_device
+/// procedures return false. Unlike isCUDADeviceContext, gpu.module and
+/// gpu.launch count even without a CUDA procedure attribute, and host_data
+/// does not.
+bool isExecutingOnDevice(mlir::Operation *op);
 bool isRegisteredDeviceGlobal(fir::GlobalOp op);
 bool isRegisteredDeviceAttr(std::optional<cuf::DataAttribute> attr);
 

diff  --git a/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td b/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
index bae9794d91f91..b2a1888da7f68 100644
--- a/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
+++ b/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
@@ -396,6 +396,21 @@ def cuf_SharedMemoryOp
       CArg<"llvm::ArrayRef<mlir::NamedAttribute>", "{}">:$attributes)>];
 }
 
+def cuf_OnDeviceOp : cuf_Op<"on_device", [Pure]> {
+  let summary = "Query whether execution is on the device";
+
+  let description = [{
+    CUDA Fortran `on_device()` intrinsic. Produces true when this operation
+    is compiled into device code and false on the host. Folding is delayed so
+    a host copy of an OpenACC routine can be cloned for the device before each
+    copy is folded in its own context.
+  }];
+
+  let results = (outs I1:$res);
+
+  let assemblyFormat = "attr-dict `:` type($res)";
+}
+
 def cuf_DeviceIsActiveOp : cuf_Op<"device_is_active", []> {
   let summary = "Query whether the CUDA device context is still active";
 

diff  --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index 51f97efe02885..9bd2a033e1365 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -589,14 +589,15 @@ def CUFAllocationConversion : Pass<"cuf-allocation-convert", "mlir::ModuleOp"> {
 def CUFOpConversion : Pass<"cuf-convert", "mlir::ModuleOp"> {
   let summary = "Convert some CUF operations to runtime calls";
   let dependentDialects = ["fir::FIROpsDialect", "hlfir::hlfirDialect",
-                           "mlir::gpu::GPUDialect", "mlir::DLTIDialect"];
+                           "mlir::arith::ArithDialect", "mlir::gpu::GPUDialect",
+                           "mlir::DLTIDialect"];
   let options = [Option<
       "deferAccRoutineDataTransfers", "defer-acc-routine-data-transfers",
       "bool", /*default=*/"false",
-      "Do not convert cuf.data_transfer operations in the host copy of an "
-      "OpenACC routine. The routine is later specialized for host and device, "
-      "after which each copy can be handled according to its execution "
-      "context.">];
+      "Do not convert cuf.data_transfer or cuf.on_device operations in the "
+      "host copy of an OpenACC routine. The routine is later specialized for "
+      "host and device, after which each copy can be handled according to its "
+      "execution context.">];
 }
 
 def CUFOpConversionLate : Pass<"cuf-convert-late", "mlir::ModuleOp"> {
@@ -714,18 +715,6 @@ def CUFDeviceFuncTransform
               "descriptor I/O">];
 }
 
-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", ""> {
   let summary = "Attach CUDA attribute to CUF kernel generated launch";
   let description = [{

diff  --git a/flang/lib/Optimizer/Builder/CUDAIntrinsicCall.cpp b/flang/lib/Optimizer/Builder/CUDAIntrinsicCall.cpp
index db41d0f5151bd..1877ad8fd0672 100644
--- a/flang/lib/Optimizer/Builder/CUDAIntrinsicCall.cpp
+++ b/flang/lib/Optimizer/Builder/CUDAIntrinsicCall.cpp
@@ -642,16 +642,41 @@ static constexpr IntrinsicHandler cudaHandlers[]{
 };
 static_assert(fir::isSorted(cudaHandlers) && "map must be sorted");
 
-const IntrinsicHandler *findCUDAIntrinsicHandler(llvm::StringRef name,
-                                                 bool isBindcCall) {
-  if (isBindcCall)
-    return nullptr;
+// BIND(C) CUDA Fortran procedures. Kept separate because the other handlers
+// are not BIND(C) and must not match a user procedure with the same name.
+static constexpr IntrinsicHandler cudaBindcHandlers[]{
+    {"on_device",
+     static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genOnDevice),
+     {},
+     /*isElemental=*/false},
+};
+static_assert(fir::isSorted(cudaBindcHandlers) && "map must be sorted");
+
+static const IntrinsicHandler *
+lookupCUDAHandler(llvm::ArrayRef<IntrinsicHandler> handlers,
+                  llvm::StringRef name) {
   auto compare = [](const IntrinsicHandler &cudaHandler, llvm::StringRef name) {
     return name.compare(cudaHandler.name) > 0;
   };
-  auto result = llvm::lower_bound(cudaHandlers, name, compare);
-  return result != std::end(cudaHandlers) && result->name == name ? result
-                                                                  : nullptr;
+  auto result = llvm::lower_bound(handlers, name, compare);
+  return result != handlers.end() && result->name == name ? &*result : nullptr;
+}
+
+const IntrinsicHandler *findCUDAIntrinsicHandler(llvm::StringRef name,
+                                                 bool isBindcCall) {
+  // The tables have 
diff erent lengths, so they cannot share a ternary: that
+  // would decay both arrays to a pointer.
+  if (isBindcCall)
+    return lookupCUDAHandler(cudaBindcHandlers, name);
+  return lookupCUDAHandler(cudaHandlers, name);
+}
+
+mlir::Value
+CUDAIntrinsicLibrary::genOnDevice(mlir::Type resultType,
+                                  llvm::ArrayRef<mlir::Value> args) {
+  assert(args.empty() && "on_device takes no arguments");
+  mlir::Value onDevice = cuf::OnDeviceOp::create(builder, loc);
+  return builder.createConvert(loc, resultType, onDevice);
 }
 
 static mlir::Value convertPtrToNVVMSpace(fir::FirOpBuilder &builder,

diff  --git a/flang/lib/Optimizer/Builder/CUFCommon.cpp b/flang/lib/Optimizer/Builder/CUFCommon.cpp
index 2ba5e445634cf..de5831a05636a 100644
--- a/flang/lib/Optimizer/Builder/CUFCommon.cpp
+++ b/flang/lib/Optimizer/Builder/CUFCommon.cpp
@@ -15,6 +15,7 @@
 #include "flang/Optimizer/Support/AllocationPolicy.h"
 #include "flang/Optimizer/Support/InternalNames.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "mlir/IR/SymbolTable.h"
 #include "llvm/ADT/StringSet.h"
@@ -70,6 +71,28 @@ bool cuf::isCUDADeviceContext(mlir::Region &region,
   return false;
 }
 
+bool cuf::isExecutingOnDevice(mlir::Operation *op) {
+  if (!op)
+    return false;
+  if (op->getParentOfType<cuf::KernelOp>() ||
+      op->getParentOfType<mlir::acc::OffloadRegionOpInterface>() ||
+      op->getParentOfType<mlir::gpu::GPUModuleOp>() ||
+      op->getParentOfType<mlir::gpu::LaunchOp>() ||
+      op->getParentOfType<mlir::gpu::GPUFuncOp>())
+    return true;
+  if (auto funcOp = op->getParentOfType<mlir::func::FuncOp>()) {
+    if (mlir::acc::isSpecializedAccRoutine(funcOp))
+      return true;
+    if (auto cudaProcAttr =
+            funcOp.getOperation()->getAttrOfType<cuf::ProcAttributeAttr>(
+                cuf::getProcAttrName())) {
+      return cudaProcAttr.getValue() != cuf::ProcAttribute::Host &&
+             cudaProcAttr.getValue() != cuf::ProcAttribute::HostDevice;
+    }
+  }
+  return false;
+}
+
 bool cuf::isRegisteredDeviceAttr(std::optional<cuf::DataAttribute> attr) {
   if (attr && (*attr == cuf::DataAttribute::Device ||
                *attr == cuf::DataAttribute::Managed ||

diff  --git a/flang/lib/Optimizer/Transforms/CMakeLists.txt b/flang/lib/Optimizer/Transforms/CMakeLists.txt
index cb70854eec128..0a3d28098054b 100644
--- a/flang/lib/Optimizer/Transforms/CMakeLists.txt
+++ b/flang/lib/Optimizer/Transforms/CMakeLists.txt
@@ -16,7 +16,6 @@ add_flang_library(FIRTransforms
   CUDA/CUFDeviceFuncTransform.cpp
   CUDA/CUFDeviceGlobal.cpp
   CUDA/CUFDuplicateDeviceFunc.cpp
-  CUDA/CUFFunctionRewrite.cpp
   CUDA/CUFGPUToLLVMConversion.cpp
   CUDA/CUFLaunchAttachAttr.cpp
   CUDA/CUFOpConversion.cpp

diff  --git a/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
deleted file mode 100644
index 63c03d86eba3e..0000000000000
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
+++ /dev/null
@@ -1,133 +0,0 @@
-//===-- CUFFunctionRewrite.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/Dialect/FIRDialect.h"
-#include "flang/Optimizer/Dialect/FIROps.h"
-#include "flang/Optimizer/Support/DataLayout.h"
-#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"
-#include "mlir/IR/ValueRange.h"
-#include "mlir/Interfaces/FunctionInterfaces.h"
-#include "mlir/Pass/Pass.h"
-#include "mlir/Support/LogicalResult.h"
-#include "llvm/ADT/SmallVector.h"
-#include "llvm/ADT/StringMap.h"
-#include "llvm/ADT/StringRef.h"
-#include "llvm/ADT/StringSet.h"
-#include "llvm/Support/Debug.h"
-
-#define DEBUG_TYPE "flang-cuf-function-rewrite"
-
-namespace fir {
-#define GEN_PASS_DEF_CUFFUNCTIONREWRITE
-#include "flang/Optimizer/Transforms/Passes.h.inc"
-} // namespace fir
-
-using namespace mlir;
-
-namespace {
-
-using genFunctionType =
-    std::function<mlir::Value(mlir::RewriterBase &, fir::CallOp op)>;
-
-class CallConversion {
-public:
-  explicit CallConversion(bool deferAccRoutines)
-      : deferAccRoutines_(deferAccRoutines) {}
-
-  LogicalResult matchAndRewrite(fir::CallOp op,
-                                mlir::RewriterBase &rewriter) const {
-    auto callee = op.getCallee();
-    if (!callee)
-      return failure();
-
-    // Match on the callee's Fortran leaf name rather than on its symbol name so
-    // this does not depend on the target's name-mangling convention or on where
-    // in the pipeline the pass runs. getPresentableFunctionName restores the
-    // original name saved by external-name conversion and returns the
-    // deconstructed leaf name.
-    auto func = mlir::dyn_cast_or_null<mlir::FunctionOpInterface>(
-        mlir::SymbolTable::lookupNearestSymbolFrom(op, *callee));
-    if (!func)
-      return failure();
-
-    auto fct = genMappings_.find(fir::getPresentableFunctionName(func));
-    if (fct == genMappings_.end())
-      return failure();
-
-    // Only rewrite a compiler-provided declaration, never a user-defined
-    // procedure that happens to share the name.
-    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();
-    }
-
-    rewriter.setInsertionPoint(op);
-    mlir::Value result = fct->second(rewriter, op);
-    if (!result)
-      return failure();
-    rewriter.replaceOp(op, result);
-    return success();
-  }
-
-private:
-  static mlir::Value genOnDevice(mlir::RewriterBase &rewriter, fir::CallOp op) {
-    // Only fold calls that match the intrinsic's shape: no arguments and a
-    // single logical result.
-    if (!op.getArgs().empty() || op.getNumResults() != 1)
-      return {};
-    mlir::Type resTy = op.getResult(0).getType();
-    if (!mlir::isa<fir::LogicalType>(resTy))
-      return {};
-    mlir::Location loc = op.getLoc();
-    unsigned inGPUMod = op->getParentOfType<gpu::GPUModuleOp>() ? 1 : 0;
-    mlir::Type i1Ty = rewriter.getIntegerType(1);
-    mlir::Value t = mlir::arith::ConstantOp::create(
-        rewriter, loc, i1Ty, rewriter.getIntegerAttr(i1Ty, inGPUMod));
-    return fir::ConvertOp::create(rewriter, loc, resTy, t);
-  }
-
-  // Recognized by Fortran leaf name; see matchAndRewrite for how the leaf name
-  // 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 {
-    CallConversion conversion(deferAccRoutines);
-    mlir::IRRewriter rewriter(&getContext());
-    getOperation()->walk([&](fir::CallOp op) {
-      (void)conversion.matchAndRewrite(op, rewriter);
-    });
-  }
-};
-
-} // namespace

diff  --git a/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversion.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversion.cpp
index d537e9442c8e1..596368d94e434 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversion.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversion.cpp
@@ -59,6 +59,36 @@ static bool inDeviceContext(mlir::Operation *op) {
   return false;
 }
 
+static bool keepOnDeviceOp(cuf::OnDeviceOp op, bool deferAccRoutines) {
+  // Leave the host copy of an OpenACC routine alone until it has been cloned
+  // for the device. Folding it to false here would bake the host value into
+  // that clone.
+  if (!deferAccRoutines || cuf::isExecutingOnDevice(op))
+    return false;
+  if (auto funcOp = op->getParentOfType<mlir::func::FuncOp>())
+    return mlir::acc::isAccRoutine(funcOp);
+  return false;
+}
+
+struct CUFOnDeviceOpConversion
+    : public mlir::OpRewritePattern<cuf::OnDeviceOp> {
+  CUFOnDeviceOpConversion(mlir::MLIRContext *context, bool deferAccRoutines)
+      : OpRewritePattern(context), deferAccRoutines(deferAccRoutines) {}
+
+  mlir::LogicalResult
+  matchAndRewrite(cuf::OnDeviceOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+    if (keepOnDeviceOp(op, deferAccRoutines))
+      return mlir::failure();
+    rewriter.replaceOpWithNewOp<mlir::arith::ConstantOp>(
+        op, rewriter.getBoolAttr(cuf::isExecutingOnDevice(op)));
+    return mlir::success();
+  }
+
+private:
+  bool deferAccRoutines;
+};
+
 static mlir::Value createConvertOp(mlir::PatternRewriter &rewriter,
                                    mlir::Location loc, mlir::Type toTy,
                                    mlir::Value val) {
@@ -663,6 +693,10 @@ class CUFOpConversion : public fir::impl::CUFOpConversionBase<CUFOpConversion> {
           return funcOp && mlir::acc::isAccRoutine(funcOp);
         });
     target.addLegalOp<cuf::DeviceIsActiveOp>();
+    target.addDynamicallyLegalOp<cuf::OnDeviceOp>([&](cuf::OnDeviceOp op) {
+      return keepOnDeviceOp(op, deferAccRoutineDataTransfers);
+    });
+    patterns.insert<CUFOnDeviceOpConversion>(ctx, deferAccRoutineDataTransfers);
     cuf::populateCUFToFIRConversionPatterns(typeConverter, *dl, symtab,
                                             patterns);
     if (mlir::failed(mlir::applyPartialConversion(getOperation(), target,

diff  --git a/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversionLate.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversionLate.cpp
index 291e55fb1d3bc..ba460e59e162a 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversionLate.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversionLate.cpp
@@ -110,6 +110,19 @@ struct CUFDeviceAddressOpConversion
   const mlir::SymbolTable &symTab;
 };
 
+struct CUFOnDeviceOpConversion
+    : public mlir::OpRewritePattern<cuf::OnDeviceOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(cuf::OnDeviceOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+    rewriter.replaceOpWithNewOp<mlir::arith::ConstantOp>(
+        op, rewriter.getBoolAttr(cuf::isExecutingOnDevice(op)));
+    return mlir::success();
+  }
+};
+
 struct CUFDeviceIsActiveOpConversion
     : public mlir::OpRewritePattern<cuf::DeviceIsActiveOp> {
   using OpRewritePattern::OpRewritePattern;
@@ -153,9 +166,11 @@ class CUFOpConversionLate
     mlir::SymbolTable symtab(module);
     target.addLegalDialect<fir::FIROpsDialect, mlir::arith::ArithDialect,
                            mlir::gpu::GPUDialect>();
+    target.addIllegalOp<cuf::OnDeviceOp>();
     patterns.insert<CUFDeviceAddressOpConversion>(patterns.getContext(),
                                                   symtab);
-    patterns.insert<CUFDeviceIsActiveOpConversion>(patterns.getContext());
+    patterns.insert<CUFOnDeviceOpConversion, CUFDeviceIsActiveOpConversion>(
+        patterns.getContext());
     if (mlir::failed(mlir::applyPartialConversion(getOperation(), target,
                                                   std::move(patterns)))) {
       mlir::emitError(mlir::UnknownLoc::get(ctx),

diff  --git a/flang/test/Fir/CUDA/cuda-function-rewrite.mlir b/flang/test/Fir/CUDA/cuda-function-rewrite.mlir
deleted file mode 100644
index dbeb89d567dde..0000000000000
--- a/flang/test/Fir/CUDA/cuda-function-rewrite.mlir
+++ /dev/null
@@ -1,241 +0,0 @@
-// 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 {
-  func.func private @on_device() -> !fir.logical<4>
-  func.func @_QMmtestsPdo2(%arg0: !fir.ref<i32> {cuf.data_attr = #cuf.cuda<device>, fir.bindc_name = "c"}, %arg1: !fir.ref<i32> {cuf.data_attr = #cuf.cuda<device>, fir.bindc_name = "i"}) attributes {cuf.proc_attr = #cuf.cuda_proc<host_device>} {
-    %c2_i32 = arith.constant 2 : i32
-    %c1_i32 = arith.constant 1 : i32
-    %0 = fir.dummy_scope : !fir.dscope
-    %5 = fir.declare %arg0 dummy_scope %0 {uniq_name = "_QMmtestsFdo2Ec"} : (!fir.ref<i32>, !fir.dscope) -> !fir.ref<i32>
-    %8 = fir.declare %arg1 dummy_scope %0 {uniq_name = "_QMmtestsFdo2Ei"} : (!fir.ref<i32>, !fir.dscope) -> !fir.ref<i32>
-    %13 = fir.call @on_device() proc_attrs<bind_c> fastmath<contract> : () -> !fir.logical<4>
-    %14 = fir.convert %13 : (!fir.logical<4>) -> i1
-    fir.if %14 {
-      fir.store %c1_i32 to %5 : !fir.ref<i32>
-    } else {
-      fir.store %c2_i32 to %5 : !fir.ref<i32>
-    }
-    return
-  }
-}
-
-// CHECK-LABEL: gpu.module @cuda_device_mod
-// CHECK: func.func @_QMmtestsPdo2
-// CHECK: arith.constant true
-// CHECK-NOT: fir.call @on_device
-
-// -----
-
-// Test the bind(c) name "on_device" on host side.
-func.func private @on_device() -> !fir.logical<4>
-func.func @_QMmtestsPdo3(%arg0: !fir.ref<i32> {cuf.data_attr = #cuf.cuda<device>, fir.bindc_name = "c"}, %arg1: !fir.ref<i32> {cuf.data_attr = #cuf.cuda<device>, fir.bindc_name = "i"}) attributes {cuf.proc_attr = #cuf.cuda_proc<host_device>} {
-  %c2_i32 = arith.constant 2 : i32
-  %c1_i32 = arith.constant 1 : i32
-  %0 = fir.dummy_scope : !fir.dscope
-  %5 = fir.declare %arg0 dummy_scope %0 {uniq_name = "_QMmtestsFdo2Ec"} : (!fir.ref<i32>, !fir.dscope) -> !fir.ref<i32>
-  %8 = fir.declare %arg1 dummy_scope %0 {uniq_name = "_QMmtestsFdo2Ei"} : (!fir.ref<i32>, !fir.dscope) -> !fir.ref<i32>
-  %13 = fir.call @on_device() proc_attrs<bind_c> fastmath<contract> : () -> !fir.logical<4>
-  %14 = fir.convert %13 : (!fir.logical<4>) -> i1
-  fir.if %14 {
-    fir.store %c1_i32 to %5 : !fir.ref<i32>
-  } else {
-    fir.store %c2_i32 to %5 : !fir.ref<i32>
-  }
-  return
-}
-
-// CHECK-LABEL: func.func @_QMmtestsPdo3
-// CHECK: arith.constant false
-// CHECK-NOT: fir.call @on_device
-
-// -----
-
-// Test on_device() with Fortran name mangling (_QPon_device) in device context.
-gpu.module @acc_device_mod {
-  func.func private @_QPon_device() -> !fir.logical<4>
-  func.func @_QMmtestPsub_device() {
-    %c2_i32 = arith.constant 2 : i32
-    %c1_i32 = arith.constant 1 : i32
-    %0 = fir.alloca i32
-    %13 = fir.call @_QPon_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_device_mod
-// CHECK: func.func @_QMmtestPsub_device
-// CHECK: arith.constant true
-// CHECK-NOT: fir.call @_QPon_device
-
-// -----
-
-// Test _QPon_device on host side.
-func.func private @_QPon_device() -> !fir.logical<4>
-func.func @_QMmtestPsub_host() {
-  %c2_i32 = arith.constant 2 : i32
-  %c1_i32 = arith.constant 1 : i32
-  %0 = fir.alloca i32
-  %13 = fir.call @_QPon_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 @_QMmtestPsub_host
-// CHECK: arith.constant false
-// CHECK-NOT: fir.call @_QPon_device
-
-// -----
-
-// Test externally-mangled on_device_ (after ExternalNameConversion) in device
-// context. The original name is recovered from the fir.internal_name attribute.
-gpu.module @acc_extname_device_mod {
-  func.func private @on_device_() -> !fir.logical<4> attributes {fir.internal_name = "_QPon_device"}
-  func.func @_QMmtestPsub_extname_device() {
-    %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_extname_device_mod
-// CHECK: func.func @_QMmtestPsub_extname_device
-// CHECK: arith.constant true
-// CHECK-NOT: fir.call @on_device_
-
-// -----
-
-// Test on_device_ on host side (original name recovered from fir.internal_name).
-func.func private @on_device_() -> !fir.logical<4> attributes {fir.internal_name = "_QPon_device"}
-func.func @_QMmtestPsub_extname_host() {
-  %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 @_QMmtestPsub_extname_host
-// CHECK: arith.constant false
-// CHECK-NOT: fir.call @on_device_
-
-// 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: arith.constant false
-// DEFER-NOT: fir.call @on_device_
-
-// -----
-
-// 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: arith.constant false
-// CHECK-NOT: fir.call @on_device_
-
-// 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: arith.constant true
-// CHECK-NOT: fir.call @on_device_
-
-// DEFER-LABEL: gpu.module @acc_routine_device_mod
-// DEFER: func.func @_QMmtestPaccroutine_device
-// DEFER: arith.constant true
-// DEFER-NOT: fir.call @on_device_
-
-// -----
-
-// A user-defined procedure named on_device (with a body) must not be folded.
-func.func @_QPon_device() -> !fir.logical<4> {
-  %true = arith.constant true
-  %0 = fir.convert %true : (i1) -> !fir.logical<4>
-  return %0 : !fir.logical<4>
-}
-func.func @_QMmtestPsub_userdef() -> !fir.logical<4> {
-  %13 = fir.call @_QPon_device() fastmath<contract> : () -> !fir.logical<4>
-  return %13 : !fir.logical<4>
-}
-
-// CHECK-LABEL: func.func @_QMmtestPsub_userdef
-// CHECK: fir.call @_QPon_device()
-
-// -----
-
-// A call whose signature does not match the intrinsic (extra argument) must not
-// be folded.
-func.func private @_QPon_device(i32) -> !fir.logical<4>
-func.func @_QMmtestPsub_badsig(%arg0: i32) -> !fir.logical<4> {
-  %13 = fir.call @_QPon_device(%arg0) fastmath<contract> : (i32) -> !fir.logical<4>
-  return %13 : !fir.logical<4>
-}
-
-// CHECK-LABEL: func.func @_QMmtestPsub_badsig
-// CHECK: fir.call @_QPon_device(%arg0)

diff  --git a/flang/test/Fir/CUDA/cuda-on-device.mlir b/flang/test/Fir/CUDA/cuda-on-device.mlir
new file mode 100644
index 0000000000000..94d550622f243
--- /dev/null
+++ b/flang/test/Fir/CUDA/cuda-on-device.mlir
@@ -0,0 +1,110 @@
+// RUN: fir-opt --cuf-convert %s | FileCheck %s --check-prefix=FOLD
+// RUN: fir-opt --cuf-convert-late %s | FileCheck %s --check-prefix=FOLD
+// RUN: fir-opt --cuf-convert="defer-acc-routine-data-transfers=true" %s | FileCheck %s --check-prefix=DEFER
+
+module attributes {gpu.container_module, dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<i64, dense<64> : vector<2xi64>>, #dlti.dl_entry<i32, dense<32> : vector<2xi64>>, #dlti.dl_entry<i8, dense<8> : vector<2xi64>>, #dlti.dl_entry<i1, dense<8> : vector<2xi64>>, #dlti.dl_entry<!llvm.ptr, dense<64> : vector<4xi64>>, #dlti.dl_entry<"dlti.endianness", "little">, #dlti.dl_entry<"dlti.stack_alignment", 128 : i64>>} {
+  func.func @host() -> i1 {
+    %0 = cuf.on_device : i1
+    return %0 : i1
+  }
+
+  func.func @host_device() -> i1 attributes {cuf.proc_attr = #cuf.cuda_proc<host_device>} {
+    %0 = cuf.on_device : i1
+    return %0 : i1
+  }
+
+  func.func @device_proc() -> i1 attributes {cuf.proc_attr = #cuf.cuda_proc<device>} {
+    %0 = cuf.on_device : i1
+    return %0 : i1
+  }
+
+  func.func @launch() -> i1 {
+    %c1 = arith.constant 1 : index
+    %0 = arith.constant false
+    gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1, %grid_z = %c1)
+               threads(%tx, %ty, %tz) in (%block_x = %c1, %block_y = %c1, %block_z = %c1) {
+      %1 = cuf.on_device : i1
+      gpu.terminator
+    }
+    return %0 : i1
+  }
+
+  gpu.module @cuda_device_mod {
+    gpu.func @on_device_kernel() -> i1 {
+      %0 = cuf.on_device : i1
+      gpu.return %0 : i1
+    }
+  }
+
+  func.func @acc_routine() -> i1 attributes {acc.routine_info = #acc.routine_info<[@routine]>} {
+    %0 = cuf.on_device : i1
+    return %0 : i1
+  }
+
+  func.func @device_specialized() -> i1 attributes {acc.specialized_routine = #acc.specialized_routine<@routine, <seq>, "device_specialized">} {
+    %0 = acc.compute_region -> i1 {
+      %1 = cuf.on_device : i1
+      acc.yield %1 : i1
+    } <{origin = "acc.routine"}>
+    return %0 : i1
+  }
+}
+
+// FOLD-LABEL: func.func @host()
+// FOLD: %[[FALSE:.*]] = arith.constant false
+// FOLD: return %[[FALSE]] : i1
+// FOLD-NOT: cuf.on_device
+
+// DEFER-LABEL: func.func @host()
+// DEFER: arith.constant false
+// DEFER-NOT: cuf.on_device
+
+// FOLD-LABEL: func.func @host_device()
+// FOLD: arith.constant false
+// FOLD-NOT: cuf.on_device
+
+// DEFER-LABEL: func.func @host_device()
+// DEFER: arith.constant false
+// DEFER-NOT: cuf.on_device
+
+// FOLD-LABEL: func.func @device_proc()
+// FOLD: arith.constant true
+// FOLD-NOT: cuf.on_device
+
+// DEFER-LABEL: func.func @device_proc()
+// DEFER: arith.constant true
+// DEFER-NOT: cuf.on_device
+
+// FOLD-LABEL: func.func @launch()
+// FOLD: gpu.launch
+// FOLD: arith.constant true
+// FOLD-NOT: cuf.on_device
+
+// DEFER-LABEL: func.func @launch()
+// DEFER: gpu.launch
+// DEFER: arith.constant true
+// DEFER-NOT: cuf.on_device
+
+// FOLD-LABEL: gpu.func @on_device_kernel()
+// FOLD: arith.constant true
+// FOLD-NOT: cuf.on_device
+
+// DEFER-LABEL: gpu.func @on_device_kernel()
+// DEFER: arith.constant true
+// DEFER-NOT: cuf.on_device
+
+// FOLD-LABEL: func.func @acc_routine()
+// FOLD: arith.constant false
+// FOLD-NOT: cuf.on_device
+
+// DEFER-LABEL: func.func @acc_routine()
+// DEFER: cuf.on_device : i1
+// DEFER-NOT: arith.constant
+
+// FOLD-LABEL: func.func @device_specialized()
+// FOLD: arith.constant true
+// FOLD-NOT: cuf.on_device
+
+// DEFER-LABEL: func.func @device_specialized()
+// DEFER: arith.constant true
+// DEFER-NOT: cuf.on_device

diff  --git a/flang/test/Fir/cuf.mlir b/flang/test/Fir/cuf.mlir
index 75b9a1d20868b..385212a1fbbda 100644
--- a/flang/test/Fir/cuf.mlir
+++ b/flang/test/Fir/cuf.mlir
@@ -121,3 +121,12 @@ func.func @_QPdevice_is_active() {
 }
 
 // CHECK: %{{.*}} = cuf.device_is_active : i1
+
+// -----
+
+func.func @_QPon_device() {
+  %0 = cuf.on_device : i1
+  return
+}
+
+// CHECK: %{{.*}} = cuf.on_device : i1

diff  --git a/flang/test/Lower/CUDA/cuda-on-device.cuf b/flang/test/Lower/CUDA/cuda-on-device.cuf
new file mode 100644
index 0000000000000..c083f82f631d9
--- /dev/null
+++ b/flang/test/Lower/CUDA/cuda-on-device.cuf
@@ -0,0 +1,23 @@
+! RUN: bbc -emit-hlfir -fcuda %s -o - | FileCheck %s
+
+! on_device() lowers to cuf.on_device. Folding to a host or device constant
+! happens later, once each copy is in its final execution context.
+
+subroutine on_host(r)
+  use cudadevice
+  logical :: r
+  r = on_device()
+end subroutine
+
+attributes(device) subroutine on_dev(r)
+  logical, device :: r
+  r = on_device()
+end subroutine
+
+! CHECK-LABEL: func.func @_QPon_host
+! CHECK: cuf.on_device : i1
+! CHECK-NOT: fir.call @on_device
+
+! CHECK-LABEL: func.func @_QPon_dev
+! CHECK: cuf.on_device : i1
+! CHECK-NOT: fir.call @on_device


        


More information about the flang-commits mailing list