[flang-commits] [clang] [flang] [mlir] [Flang][OpenMP] Lower DECLARE TARGET INDIRECT clause (PR #208387)

[via flang-commits]

================
@@ -0,0 +1,94 @@
+//===- IndirectCallLookup.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
+//
+//===----------------------------------------------------------------------===//
+//
+// On a GPU target device, rewrites indirect fir.call ops so the callee (a host
+// function address held in a procedure pointer) is resolved to the device
+// address via the `__llvm_omp_indirect_call_lookup` runtime function.
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/OpenMP/Passes.h"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
+#include "mlir/Dialect/OpenMP/OpenMPInterfaces.h"
+#include "mlir/IR/Builders.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace flangomp {
+#define GEN_PASS_DEF_INDIRECTCALLLOOKUPPASS
+#include "flang/Optimizer/OpenMP/Passes.h.inc"
+} // namespace flangomp
+
+using namespace mlir;
+
+/// Runtime function that maps a host function address to the device address.
+static constexpr llvm::StringRef indirectCallLookupName =
+    "__llvm_omp_indirect_call_lookup";
+
+namespace {
+class IndirectCallLookupPass
+    : public flangomp::impl::IndirectCallLookupPassBase<
+          IndirectCallLookupPass> {
+public:
+  void runOnOperation() override {
+    mlir::ModuleOp module = getOperation();
+    auto offloadMod = mlir::dyn_cast<mlir::omp::OffloadModuleInterface>(
+        module.getOperation());
+
+    // Only a GPU target device needs host-to-device address translation.
+    if (!offloadMod || !offloadMod.getIsTargetDevice() ||
+        !offloadMod.getIsGPU())
+      return;
+
+    // An indirect fir.call has no callee symbol; operand 0 is the callee value.
+    llvm::SmallVector<fir::CallOp> indirectCalls;
+    module.walk([&](fir::CallOp call) {
+      if (!call.getCallee())
+        indirectCalls.push_back(call);
+    });
+    if (indirectCalls.empty())
+      return;
+
+    mlir::MLIRContext *ctx = &getContext();
+    mlir::OpBuilder builder(ctx);
+
+    // A function value lowers to a pointer, so an opaque `() -> ()` type
+    // matches the runtime function's ptr argument and result.
+    auto opaqueFnTy = mlir::FunctionType::get(ctx, {}, {});
+
+    // Declare the runtime lookup function once.
+    auto lookupFn =
+        module.lookupSymbol<mlir::func::FuncOp>(indirectCallLookupName);
+    if (!lookupFn) {
+      builder.setInsertionPointToStart(module.getBody());
+      lookupFn = mlir::func::FuncOp::create(
+          builder, module.getLoc(), indirectCallLookupName,
+          mlir::FunctionType::get(ctx, {opaqueFnTy}, {opaqueFnTy}));
----------------
chichunchen wrote:

The lookup seems to lose address space. Could `mlir::FunctionType` be lowered using the module's address space? (need test for address space as well)

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