[flang] Add a pass to get OpenACC device ptr for CUDA kernel (#212299)

When a CUDA kernel is launched inside an OpenACC data region, it does
not properly get the device pointers and instead uses host data. This PR
adds a pass to get the device pointers set up by the OpenACC data
construct and pass them explicitly to the CUDA kernel.

Note: A possibly non-contiguous array argument is currently not supported. This pass will skip these cases.

---------

Co-authored-by: Yebin Chon <ychon@nvidia.com>

flang/include/flang/Optimizer/OpenACC/Passes.td

@@ -109,4 +109,31 @@ def ACCOptimizeFirstprivateMap
   let dependentDialects = ["mlir::acc::OpenACCDialect", "fir::FIROpsDialect"];
 }
 
+def ACCDevicePtrToCUFKernel
+    : Pass<"acc-device-ptr-to-cuf-kernel", "mlir::ModuleOp"> {
+  let summary = "Pass device addresses to CUDA Fortran kernels launched inside "
+                "OpenACC data regions";
+  let description = [{
+    A CUDA Fortran kernel launched inside an OpenACC data region must receive
+    the device address of any host variable that OpenACC has made present, not
+    the host address. Otherwise the kernel dereferences a host pointer, which is
+    only valid on shared-memory/unified-addressing systems and is illegal on a
+    device with a separate address space.
+
+    For each cuf.kernel_launch whose reference arguments resolve to a variable
+    mapped by an enclosing acc.data region, this pass wraps the launch in an
+    acc.host_data construct with acc.use_device operands for the mapped
+    variables, and rebuilds the launch argument addressing on top of the
+    use_device result so the kernel receives the device pointer.
+
+    Both directly-addressed variables (e.g. static arrays, whose data address is
+    the mapped varPtr) and descriptor-based variables (e.g. allocatables and
+    pointers, whose data address is box_addr(load(<descriptor>))) are handled.
+
+    This must run before the pass that lowers cuf.kernel_launch to
+    gpu.launch_func.
+  }];
+  let dependentDialects = ["mlir::acc::OpenACCDialect"];
+}
+
 #endif // FORTRAN_OPTIMIZER_OPENACC_PASSES

flang/lib/Optimizer/OpenACC/Transforms/ACCDevicePtrToCUFKernel.cpp

@@ -0,0 +1,188 @@
+//===- ACCDevicePtrToCUFKernel.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
+//
+//===----------------------------------------------------------------------===//
+//
+// A CUDA Fortran kernel launched inside an OpenACC data
+// region must receive the device address of any host variable that OpenACC has
+// made present, not the host address. This pass wraps each
+// cuf.kernel_launch in an acc.host_data construct with acc.use_device operands
+// for the mapped host variables, and rebuilds the launch's argument addressing
+// on top of the use_device result. The host_data/use_device lowering then
+// materializes the present-table device pointer, and any array-section
+// addressing is recomputed on the device pointer.
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Dialect/CUF/CUFOps.h"
+#include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/HLFIR/HLFIROps.h"
+#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/Dialect/OpenACC/OpenACCUtils.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace fir {
+namespace acc {
+#define GEN_PASS_DEF_ACCDEVICEPTRTOCUFKERNEL
+#include "flang/Optimizer/OpenACC/Passes.h.inc"
+} // namespace acc
+} // namespace fir
+
+using namespace mlir;
+
+namespace {
+
+/// Walk down an addressing chain to the underlying variable that OpenACC maps
+/// as a data-clause pointer
+static Value getMappedVar(Value value) {
+  while (true) {
+    Operation *def = value.getDefiningOp();
+    // Stop at the variable OpenACC maps as the data-clause pointer.
+    if (isa_and_nonnull<fir::DeclareOp, hlfir::DeclareOp>(def))
+      return value;
+    if (auto view = dyn_cast_or_null<fir::FortranObjectViewOpInterface>(def)) {
+      value = view.getViewSource(cast<OpResult>(value));
+      continue;
+    }
+    // TODO: drop this case once (if) fir.emboxchar implements the
+    // FortranObjectViewOpInterface.
+    if (auto embox = dyn_cast_or_null<fir::EmboxCharOp>(def)) {
+      value = embox.getMemref();
+      continue;
+    }
+    // Descriptor-based (allocatable/pointer) variables: the data address is
+    // extracted from the descriptor via box_addr(load(<descriptor ref>)). Peel
+    // both so the walk reaches the descriptor variable, which is what OpenACC
+    // maps as the data-clause varPtr for such variables.
+    if (auto load = value.getDefiningOp<fir::LoadOp>()) {
+      // Only a load that produces a descriptor is part of the addressing
+      // chain; scalar loads are ordinary values, not addressing steps.
+      if (mlir::isa<fir::BaseBoxType>(load.getType())) {
+        value = load.getMemref();
+        continue;
+      }
+    }
+    return {};
+  }
+}
+
+/// Collects the host variables made present by an
+/// OpenACC data directive that dominates `launch`. The result is
+/// computed once per launch, then queried per argument.
+static llvm::DenseSet<Value>
+collectPresentAccVars(cuf::KernelLaunchOp launch, DominanceInfo &domInfo,
+                      PostDominanceInfo &postDomInfo) {
+  llvm::DenseSet<Value> presentVars;
+  for (Value dataClause :
+       acc::getDominatingDataClauses(launch, domInfo, postDomInfo))
+    if (Value hostVar = acc::getVar(dataClause.getDefiningOp()))
+      if (Value mappedVar = getMappedVar(hostVar))
+        presentVars.insert(mappedVar);
+  return presentVars;
+}
+
+class ACCDevicePtrToCUFKernel
+    : public fir::acc::impl::ACCDevicePtrToCUFKernelBase<
+          ACCDevicePtrToCUFKernel> {
+public:
+  using fir::acc::impl::ACCDevicePtrToCUFKernelBase<
+      ACCDevicePtrToCUFKernel>::ACCDevicePtrToCUFKernelBase;
+
+  void runOnOperation() override {
+    llvm::SmallVector<cuf::KernelLaunchOp> launches;
+    getOperation().walk(
+        [&](cuf::KernelLaunchOp launch) { launches.push_back(launch); });
+
+    for (cuf::KernelLaunchOp launch : launches)
+      rewriteLaunch(launch);
+  }
+
+private:
+  void rewriteLaunch(cuf::KernelLaunchOp launch) {
+    // Collect kernel arguments that are references to a host variable made
+    // present by an enclosing acc.data region.
+    llvm::SmallVector<OpOperand *> deviceOperands;
+    llvm::SetVector<Value> deviceArgs;
+
+    // Special case: a boxchar argument is not pointer-like, so acc.use_device
+    // cannot take it. Translate its base address and rebuild the boxchar on the
+    // device address. Keyed by the launch operand so the general path below is
+    // untouched. TODO: remove once fir.emboxchar is a
+    // FortranObjectViewOpInterface.
+    llvm::DenseMap<OpOperand *, fir::EmboxCharOp> boxCharArgs;
+
+    DominanceInfo domInfo;
+    PostDominanceInfo postDomInfo;
+    llvm::DenseSet<Value> presentAccVars =
+        collectPresentAccVars(launch, domInfo, postDomInfo);
+
+    for (OpOperand &operand : launch.getArgsMutable()) {
+      Value arg = operand.get();
+      Value mappedVar = getMappedVar(arg);
+      // Check if mappedVar is present due to an enclosing OpenACC data region.
+      if (!mappedVar || !presentAccVars.contains(mappedVar))
+        continue;
+      deviceOperands.push_back(&operand);
+      if (auto embox = arg.getDefiningOp<fir::EmboxCharOp>()) {
+        boxCharArgs[&operand] = embox;
+        deviceArgs.insert(embox.getMemref()); // translate the base pointer
+      } else {
+        deviceArgs.insert(arg); // general case
+      }
+    }
+
+    if (deviceOperands.empty())
+      return;
+
+    OpBuilder builder(launch);
+    Location loc = launch.getLoc();
+
+    // One acc.use_device per distinct kernel argument, emitted before the
+    // launch so it dominates the host_data region.
+    llvm::DenseMap<Value, Value> deviceArgsMap;
+    llvm::SmallVector<Value> hostDataOperands;
+    for (Value arg : deviceArgs) {
+      Value deviceVar = acc::UseDeviceOp::create(builder, loc, arg,
+                                                 /*structured=*/true,
+                                                 /*implicit=*/false)
+                            .getAccVar();
+      deviceArgsMap[arg] = deviceVar;
+      hostDataOperands.push_back(deviceVar);
+    }
+
+    // Wrap the launch in an acc.host_data region.
+    auto hostData = acc::HostDataOp::create(builder, loc, /*ifCond=*/Value{},
+                                            hostDataOperands);
+    hostData.setIfPresent(true);
+    Block *body = builder.createBlock(&hostData.getRegion());
+    builder.setInsertionPointToStart(body);
+    Operation *terminator = acc::TerminatorOp::create(builder, loc);
+    launch->moveBefore(terminator);
+
+    for (OpOperand *operand : deviceOperands) {
+      if (fir::EmboxCharOp embox = boxCharArgs.lookup(operand)) {
+        // Char case: rebuild the boxchar on the device base address. Must be
+        // emitted before the launch so the rebuilt value dominates its use.
+        builder.setInsertionPoint(launch);
+        operand->assign(fir::EmboxCharOp::create(
+            builder, embox.getLoc(), embox.getType(),
+            deviceArgsMap[embox.getMemref()], embox.getLen()));
+      } else {
+        // General case
+        operand->assign(deviceArgsMap[operand->get()]);
+      }
+    }
+  }
+};
+
+} // namespace

flang/lib/Optimizer/OpenACC/Transforms/CMakeLists.txt

@@ -1,5 +1,6 @@
 add_flang_library(FIROpenACCTransforms
   ACCDeclareActionConversion.cpp
+  ACCDevicePtrToCUFKernel.cpp
   ACCEmitNYIFlang.cpp
   ACCInitializeFIRAnalyses.cpp
   ACCPipeline.cpp
@@ -11,6 +12,7 @@ add_flang_library(FIROpenACCTransforms
   FIROpenACCPassesIncGen
 
   LINK_LIBS
+  CUFDialect
   FIRAnalysis
   FIRBuilder
   FIRDialect

flang/test/Fir/OpenACC/device-ptr-to-cuf-kernel.mlir

@@ -0,0 +1,135 @@
+// RUN: fir-opt %s --acc-device-ptr-to-cuf-kernel -split-input-file | FileCheck %s
+
+// A CUF kernel launched inside an acc.data region that maps a directly-addressed
+// (static) array. The kernel argument is the mapped variable itself, so it is
+// wrapped in acc.use_device and substituted directly.
+func.func @static_array() {
+  %c1 = arith.constant 1 : i32
+  %c100 = arith.constant 100 : index
+  %0 = fir.alloca !fir.array<100xi32> {bindc_name = "a", uniq_name = "_QFEa"}
+  %sh = fir.shape %c100 : (index) -> !fir.shape<1>
+  %1 = fir.declare %0(%sh) {uniq_name = "_QFEa"} : (!fir.ref<!fir.array<100xi32>>, !fir.shape<1>) -> !fir.ref<!fir.array<100xi32>>
+  %2 = acc.copyin varPtr(%1 : !fir.ref<!fir.array<100xi32>>) -> !fir.ref<!fir.array<100xi32>> {name = "a"}
+  acc.data dataOperands(%2 : !fir.ref<!fir.array<100xi32>>) {
+    cuf.kernel_launch @kernel<<<%c1, %c1, %c1, %c1, %c1, %c1>>>(%1) : (!fir.ref<!fir.array<100xi32>>)
+    acc.terminator
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @static_array
+// CHECK: %[[DECL:.*]] = fir.declare
+// CHECK: acc.data
+// CHECK: %[[DEV:.*]] = acc.use_device varPtr(%[[DECL]] : !fir.ref<!fir.array<100xi32>>)
+// CHECK: acc.host_data dataOperands(%[[DEV]]
+// CHECK: cuf.kernel_launch @kernel<<<{{.*}}>>>(%[[DEV]]) : (!fir.ref<!fir.array<100xi32>>)
+// CHECK: attributes {ifPresent}
+
+// -----
+
+// A CUF kernel launched with an array section. The launch argument is already an
+// interior pointer (&a(3)) computed on the host base. The pass wraps that
+// interior pointer directly in acc.use_device and substitutes it; the section
+// addressing is left untouched (not rebuilt on a device base).
+func.func @array_section() {
+  %c1 = arith.constant 1 : i32
+  %c2_i32 = arith.constant 2 : i32
+  %c100 = arith.constant 100 : index
+  %c3 = arith.constant 3 : index
+  %0 = fir.alloca !fir.array<100xi32> {uniq_name = "_QFEa"}
+  %sh = fir.shape %c100 : (index) -> !fir.shape<1>
+  %1 = fir.declare %0(%sh) {uniq_name = "_QFEa"} : (!fir.ref<!fir.array<100xi32>>, !fir.shape<1>) -> !fir.ref<!fir.array<100xi32>>
+  %2 = acc.copyin varPtr(%1 : !fir.ref<!fir.array<100xi32>>) -> !fir.ref<!fir.array<100xi32>> {name = "a"}
+  acc.data dataOperands(%2 : !fir.ref<!fir.array<100xi32>>) {
+    %3 = fir.array_coor %1(%sh) %c3 : (!fir.ref<!fir.array<100xi32>>, !fir.shape<1>, index) -> !fir.ref<i32>
+    %4 = fir.convert %3 : (!fir.ref<i32>) -> !fir.ref<!fir.array<?xi32>>
+    cuf.kernel_launch @kernel<<<%c1, %c1, %c1, %c1, %c1, %c1>>>(%4, %c2_i32) : (!fir.ref<!fir.array<?xi32>>, i32)
+    acc.terminator
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @array_section
+// CHECK: %[[DECL:.*]] = fir.declare
+// CHECK: acc.data
+// CHECK: %[[COOR:.*]] = fir.array_coor %[[DECL]]
+// CHECK: %[[CONV:.*]] = fir.convert %[[COOR]] : (!fir.ref<i32>) -> !fir.ref<!fir.array<?xi32>>
+// CHECK: %[[DEV:.*]] = acc.use_device varPtr(%[[CONV]] : !fir.ref<!fir.array<?xi32>>)
+// CHECK: acc.host_data dataOperands(%[[DEV]]
+// CHECK: cuf.kernel_launch @kernel<<<{{.*}}>>>(%[[DEV]], {{.*}}) : (!fir.ref<!fir.array<?xi32>>, i32)
+// CHECK: attributes {ifPresent}
+
+// -----
+
+// A CUF kernel launched inside an acc.data region that maps a descriptor-based
+// (allocatable) variable. OpenACC maps the descriptor, but the launch already
+// extracts the data address on the host via box_addr(load(<descriptor>)). The
+// pass wraps that data pointer in acc.use_device and substitutes it directly;
+// the descriptor addressing is left on the host descriptor and not rebuilt.
+func.func @descriptor_array() {
+  %c1 = arith.constant 1 : i32
+  %0 = fir.alloca !fir.box<!fir.heap<!fir.array<?xi32>>> {bindc_name = "h", uniq_name = "_QFEh"}
+  %1 = fir.declare %0 {fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFEh"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+  %2 = acc.create varPtr(%1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>> {name = "h"}
+  acc.data dataOperands(%2 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>) {
+    %3 = fir.load %1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+    %4 = fir.box_addr %3 : (!fir.box<!fir.heap<!fir.array<?xi32>>>) -> !fir.heap<!fir.array<?xi32>>
+    %5 = fir.convert %4 : (!fir.heap<!fir.array<?xi32>>) -> !fir.ref<!fir.array<?xi32>>
+    cuf.kernel_launch @kernel<<<%c1, %c1, %c1, %c1, %c1, %c1>>>(%5) : (!fir.ref<!fir.array<?xi32>>)
+    acc.terminator
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @descriptor_array
+// CHECK: %[[DECL:.*]] = fir.declare
+// CHECK: acc.data
+// CHECK: %[[LOAD:.*]] = fir.load %[[DECL]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+// CHECK: %[[ADDR:.*]] = fir.box_addr %[[LOAD]] : (!fir.box<!fir.heap<!fir.array<?xi32>>>) -> !fir.heap<!fir.array<?xi32>>
+// CHECK: %[[CONV:.*]] = fir.convert %[[ADDR]] : (!fir.heap<!fir.array<?xi32>>) -> !fir.ref<!fir.array<?xi32>>
+// CHECK: %[[DEV:.*]] = acc.use_device varPtr(%[[CONV]] : !fir.ref<!fir.array<?xi32>>)
+// CHECK: acc.host_data dataOperands(%[[DEV]]
+// CHECK: cuf.kernel_launch @kernel<<<{{.*}}>>>(%[[DEV]]) : (!fir.ref<!fir.array<?xi32>>)
+// CHECK: attributes {ifPresent}
+
+// -----
+
+// No enclosing acc.data region: the launch must be left untouched.
+func.func @no_enclosing_data() {
+  %c1 = arith.constant 1 : i32
+  %c100 = arith.constant 100 : index
+  %0 = fir.alloca !fir.array<100xi32> {uniq_name = "_QFEa"}
+  %sh = fir.shape %c100 : (index) -> !fir.shape<1>
+  %1 = fir.declare %0(%sh) {uniq_name = "_QFEa"} : (!fir.ref<!fir.array<100xi32>>, !fir.shape<1>) -> !fir.ref<!fir.array<100xi32>>
+  cuf.kernel_launch @kernel<<<%c1, %c1, %c1, %c1, %c1, %c1>>>(%1) : (!fir.ref<!fir.array<100xi32>>)
+  return
+}
+
+// CHECK-LABEL: func.func @no_enclosing_data
+// CHECK-NOT: acc.use_device
+// CHECK-NOT: acc.host_data
+// CHECK: cuf.kernel_launch @kernel
+
+// -----
+
+// The enclosing acc.data maps a different variable than the one launched: no
+// rewrite.
+func.func @unmapped_arg() {
+  %c1 = arith.constant 1 : i32
+  %c100 = arith.constant 100 : index
+  %sh = fir.shape %c100 : (index) -> !fir.shape<1>
+  %0 = fir.alloca !fir.array<100xi32> {uniq_name = "_QFEa"}
+  %1 = fir.declare %0(%sh) {uniq_name = "_QFEa"} : (!fir.ref<!fir.array<100xi32>>, !fir.shape<1>) -> !fir.ref<!fir.array<100xi32>>
+  %2 = fir.alloca !fir.array<100xi32> {uniq_name = "_QFEb"}
+  %3 = fir.declare %2(%sh) {uniq_name = "_QFEb"} : (!fir.ref<!fir.array<100xi32>>, !fir.shape<1>) -> !fir.ref<!fir.array<100xi32>>
+  %4 = acc.copyin varPtr(%3 : !fir.ref<!fir.array<100xi32>>) -> !fir.ref<!fir.array<100xi32>> {name = "b"}
+  acc.data dataOperands(%4 : !fir.ref<!fir.array<100xi32>>) {
+    cuf.kernel_launch @kernel<<<%c1, %c1, %c1, %c1, %c1, %c1>>>(%1) : (!fir.ref<!fir.array<100xi32>>)
+    acc.terminator
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @unmapped_arg
+// CHECK-NOT: acc.use_device
+// CHECK-NOT: acc.host_data