[Libomptarget] Support mapping indirect host calls to device functions

[jhuber6]

Summary

The changes in D157738 allowed for us to emit stub globals on the device
in the offloading entry section. These globals contain the addresses of
device functions and allow us to map host functions to their
corresponding device equivalent. This patch provides the initial support
required to build a table on the device to lookup the associated value.
This is done by finding these entries and creating a global table on the
device that can be searched with a simple binary search.

This requires an allocation, which supposedly should be automatically
freed at plugin shutdown. This includes a basic test which looks up device
pointers via a host pointer using the added function. This will need to be built
upon to provide full support for these calls in the runtime.

To support reverse offloading it would also be useful to provide a reverse table
that allows us to get host functions from device stubs.

Depends on D157738

[RaviNarayanaswamy]

openmp/libomptarget/DeviceRTL/src/Misc.cpp

84

Do you want to check if HstPtr is null then return nullptr

openmp/libomptarget/include/omptarget.h

95

Should this be 0x08

[jhuber6]

openmp/libomptarget/DeviceRTL/src/Misc.cpp

84

Yeah that's a good idea.

openmp/libomptarget/include/omptarget.h

95

I didn't design it this way, but the frontend uses different flags for globals and kernels. All the other flags overlap so why should this one be the exception without breaking ABI and redesigning it?

[RaviNarayanaswamy]

openmp/libomptarget/include/omptarget.h

95

You want the lookup table to be small. Does TARGET_DTOR get added to the indirect lookup table?

[RaviNarayanaswamy]

openmp/libomptarget/DeviceRTL/src/Misc.cpp

84

Also if the Table does not exist or size zero return the same pointer since it might already be a device pointer.

[jhuber6]

Return if null host pointer.

[jdoerfert]

openmp/libomptarget/DeviceRTL/src/Misc.cpp

86

Check if the HstPtr is outside the table range. This should allows us to avoid the binary search for basically all host pointers in practice.

openmp/libomptarget/include/omptarget.h

95

0x08 plz

[jhuber6]

openmp/libomptarget/DeviceRTL/src/Misc.cpp

86

We'd need to look up the beginning and the end right? That's still two loads to check.

openmp/libomptarget/include/omptarget.h

95

This comes from the previous patch, Clang for whatever reason uses a different set of enums for kernels and globals. All the others overlap so I don't know why this should be the exception. The case where this was important was the registration of the global constructors and destructors, which shouldn't have been looking at globals anyway.

[jdoerfert]

openmp/libomptarget/DeviceRTL/src/Misc.cpp

86

Yes, up to two loads/checks.
Below you check log(N) + 1 values.
With Size >= 3 it's a win, assuming device ptrs.

openmp/libomptarget/include/omptarget.h

95

All the others overlap so I don't know why this should be the exception

I don't see this as a valid argument. If we started fresh, we would not make them overlap to help debugging efforts. Why would one "stick to it" for no good reason?

[jplehr]

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

310

Can this error?

[jhuber6]

Adjust table lookups.

[jhuber6]

Change flag to 8.

[RaviNarayanaswamy]

LGTM

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

enum OMPTargetGlobalVarEntryKind : uint32_t {
     OMPTargetGlobalVarEntryTo = 0x0,
     /// Mark the entry as a to declare target link.
     OMPTargetGlobalVarEntryLink = 0x1,
     /// Mark the entry as a declare target enter.
     OMPTargetGlobalVarEntryEnter = 0x2,
     /// Mark the entry as having no declare target entry kind.
     OMPTargetGlobalVarEntryNone = 0x3,
     /// Mark the entry as a declare target indirect global.
-    OMPTargetGlobalVarEntryIndirect = 0x4,
+    OMPTargetGlobalVarEntryIndirect = 0x8,
   };
 
   /// Kind of device clause for declare target variables
   /// and functions
   /// NOTE: Currently not used as a part of a variable entry
   /// used for Flang and Clang to interface with the variable
   /// related registration functions
   enum OMPTargetDeviceClauseKind : uint32_t {

openmp/libomptarget/DeviceRTL/include/Configuration.h

 uint32_t getDebugKind();
 
 /// Return the amount of dynamic shared memory that was allocated at launch.
 uint64_t getDynamicMemorySize();
 
 /// Returns the cycles per second of the device's fixed frequency clock.
 uint64_t getClockFrequency();
 
+/// Returns the pointer to the beginning of the indirect call table.
+void *getIndirectCallTablePtr();
+
+/// Returns the size of the indirect call table.
+uint64_t getIndirectCallTableSize();
+
 /// Return if debugging is enabled for the given debug kind.
 bool isDebugMode(DebugKind Level);
 
 /// Indicates if this kernel may require thread-specific states, or if it was
 /// explicitly disabled by the user.
 bool mayUseThreadStates();
 
 /// Indicates if this kernel may require data environments for nested
 /// parallelism, or if it was explicitly disabled by the user.
 bool mayUseNestedParallelism();
 
 } // namespace config
 } // namespace ompx
 
 #endif

openmp/libomptarget/DeviceRTL/src/Configuration.cpp

 uint64_t config::getDynamicMemorySize() {
   return __omp_rtl_device_environment.DynamicMemSize;
 }
 
 uint64_t config::getClockFrequency() {
   return __omp_rtl_device_environment.ClockFrequency;
 }
 
+void *config::getIndirectCallTablePtr() {
+  return reinterpret_cast<void *>(
+      __omp_rtl_device_environment.IndirectCallTable);
+}
+
+uint64_t config::getIndirectCallTableSize() {
+  return __omp_rtl_device_environment.IndirectCallTableSize;
+}
+
 bool config::isDebugMode(config::DebugKind Kind) {
   return config::getDebugKind() & Kind;
 }
 
 bool config::mayUseThreadStates() { return !__omp_rtl_assume_no_thread_state; }
 
 bool config::mayUseNestedParallelism() {
   if (__omp_rtl_assume_no_nested_parallelism)
     return false;
   return state::getKernelEnvironment().Configuration.MayUseNestedParallelism;
 }
 
 #pragma omp end declare target

openmp/libomptarget/DeviceRTL/src/Misc.cpp

 double getWTime() {
   unsigned long long nsecs;
   asm volatile("mov.u64  %0, %%globaltimer;" : "=l"(nsecs));
   return (double)nsecs * getWTick();
 }
 
 #pragma omp end declare variant
 
+/// Lookup a device-side function using a host pointer /p HstPtr using the table
+/// provided by the device plugin. The table is an ordered pair of host and
+/// device pointers sorted on the value of the host pointer.
+void *indirectCallLookup(void *HstPtr) {
+  if (!HstPtr)
+    return nullptr;
+
+  struct IndirectCallTable {
+    void *HstPtr;
+    void *DevPtr;
+  };
+  IndirectCallTable *Table =
+      reinterpret_cast<IndirectCallTable *>(config::getIndirectCallTablePtr());
+  uint64_t TableSize = config::getIndirectCallTableSize();
+
+  // If the table is empty we assume this is device pointer.
+  if (!Table || !TableSize)
+    return HstPtr;
+
+  uint32_t Left = 0;
+  uint32_t Right = TableSize;
+
+  // If the pointer is definitely not contained in the table we exit early.
+  if (HstPtr < Table[Left].HstPtr || HstPtr > Table[Right - 1].HstPtr)
+    return HstPtr;
+
+  while (Left != Right) {
+    uint32_t Current = Left + (Right - Left) / 2;
+    if (Table[Current].HstPtr == HstPtr)
+      return Table[Current].DevPtr;
+
+    if (HstPtr < Table[Current].HstPtr)
+      Right = Current;
+    else
+      Left = Current;
+  }
+
+  // If we searched the whole table and found nothing this is a device pointer.
+  return HstPtr;
+}
+
 } // namespace impl
 } // namespace ompx
 
 /// Interfaces
 ///
 ///{
 
 extern "C" {
 int32_t __kmpc_cancellationpoint(IdentTy *, int32_t, int32_t) { return 0; }
 
 int32_t __kmpc_cancel(IdentTy *, int32_t, int32_t) { return 0; }
 
 double omp_get_wtick(void) { return ompx::impl::getWTick(); }
 
 double omp_get_wtime(void) { return ompx::impl::getWTime(); }
+
+void *__llvm_omp_indirect_call_lookup(void *HstPtr) {
+  return ompx::impl::indirectCallLookup(HstPtr);
+}
 }
 
 ///}
 #pragma omp end declare target

openmp/libomptarget/include/Environment.h

 #include "llvm/Frontend/OpenMP/OMPDeviceConstants.h"
 
 struct DeviceEnvironmentTy {
   uint32_t DebugKind;
   uint32_t NumDevices;
   uint32_t DeviceNum;
   uint32_t DynamicMemSize;
   uint64_t ClockFrequency;
+  uintptr_t IndirectCallTable;
+  uint64_t IndirectCallTableSize;
 };
 
 // NOTE: Please don't change the order of those members as their indices are
 // used in the middle end. Always add the new data member at the end.
 // Different from KernelEnvironmentTy below, this structure contains members
 // that might be modified at runtime.
 struct DynamicEnvironmentTy {
   /// Current indentation level for the function trace. Only accessed by thread

openmp/libomptarget/include/omptarget.h

enum tgt_map_type {
   // This is an OpenMP extension for the sake of OpenACC support.
   OMP_TGT_MAPTYPE_OMPX_HOLD       = 0x2000,
   // descriptor for non-contiguous target-update
   OMP_TGT_MAPTYPE_NON_CONTIG      = 0x100000000000,
   // member of struct, member given by [16 MSBs] - 1
   OMP_TGT_MAPTYPE_MEMBER_OF       = 0xffff000000000000
 };
 
+/// Flags for offload entries.
 enum OpenMPOffloadingDeclareTargetFlags {
-  /// Mark the entry as having a 'link' attribute.
+  /// Mark the entry global as having a 'link' attribute.
   OMP_DECLARE_TARGET_LINK = 0x01,
-  /// Mark the entry as being a global constructor.
+  /// Mark the entry kernel as being a global constructor.
   OMP_DECLARE_TARGET_CTOR = 0x02,
-  /// Mark the entry as being a global destructor.
-  OMP_DECLARE_TARGET_DTOR = 0x04
+  /// Mark the entry kernel as being a global destructor.
+  OMP_DECLARE_TARGET_DTOR = 0x04,
+  /// Mark the entry global as being an indirectly callable function.
+  OMP_DECLARE_TARGET_INDIRECT = 0x08
 };
 
 enum OpenMPOffloadingRequiresDirFlags {
   /// flag undefined.
   OMP_REQ_UNDEFINED               = 0x000,
   /// no requires directive present.
   OMP_REQ_NONE                    = 0x001,
   /// reverse_offload clause.

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

Error init(GenericDeviceTy *Device, uint64_t MemSize, RRStatusTy Status,
 
     return Plugin::success();
   }
 
   void deinit() { Device->free(MemoryStart); }
 
 } RecordReplay;
 
+// Extract the mapping of host function pointers to device function pointers
+// from the entry table. Functions marked as 'indirect' in OpenMP will have
+// offloading entries generated for them which map the host's function pointer
+// to a global containing the corresponding function pointer on the device.
+static Expected<std::pair<void *, uint64_t>>
+setupIndirectCallTable(GenericPluginTy &Plugin, GenericDeviceTy &Device,
+                       DeviceImageTy &Image) {
+  GenericGlobalHandlerTy &Handler = Plugin.getGlobalHandler();
+
+  llvm::ArrayRef<__tgt_offload_entry> Entries(Image.getTgtImage()->EntriesBegin,
+                                              Image.getTgtImage()->EntriesEnd);
+  llvm::SmallVector<std::pair<void *, void *>> IndirectCallTable;
+  for (const auto &Entry : Entries) {
+    if (Entry.size == 0 || !(Entry.flags & OMP_DECLARE_TARGET_INDIRECT))
+      continue;
+
+    assert(Entry.size == sizeof(void *) && "Global not a function pointer?");
+    auto &[HstPtr, DevPtr] = IndirectCallTable.emplace_back();
+
+    GlobalTy DeviceGlobal(Entry.name, Entry.size);
+    if (auto Err =
+            Handler.getGlobalMetadataFromDevice(Device, Image, DeviceGlobal))
+      return std::move(Err);
+
+    HstPtr = Entry.addr;
+    if (auto Err = Device.dataRetrieve(&DevPtr, DeviceGlobal.getPtr(),
+                                       Entry.size, nullptr))
+      return std::move(Err);
+  }
+
+  // If we do not have any indirect globals we exit early.
+  if (IndirectCallTable.empty())
+    return std::pair{nullptr, 0};
+
+  // Sort the array to allow for more efficient lookup of device pointers.
+  llvm::sort(IndirectCallTable,
+             [](const auto &x, const auto &y) { return x.first < y.first; });
+
+  uint64_t TableSize =
+      IndirectCallTable.size() * sizeof(std::pair<void *, void *>);
+  void *DevicePtr = Device.allocate(TableSize, nullptr, TARGET_ALLOC_DEVICE);
+  if (auto Err = Device.dataSubmit(DevicePtr, IndirectCallTable.data(),
+                                   TableSize, nullptr))
+    return std::move(Err);
+  return std::pair<void *, uint64_t>(DevicePtr, IndirectCallTable.size());
+}
+
 AsyncInfoWrapperTy::AsyncInfoWrapperTy(GenericDeviceTy &Device,
                                        __tgt_async_info *AsyncInfoPtr)
     : Device(Device),
       AsyncInfoPtr(AsyncInfoPtr ? AsyncInfoPtr : &LocalAsyncInfo) {}
 
 void AsyncInfoWrapperTy::finalize(Error &Err) {
   assert(AsyncInfoPtr && "AsyncInfoWrapperTy already finalized");
 

 }
 
 Error GenericDeviceTy::setupDeviceEnvironment(GenericPluginTy &Plugin,
                                               DeviceImageTy &Image) {
   // There are some plugins that do not need this step.
   if (!shouldSetupDeviceEnvironment())
     return Plugin::success();
 
+  // Obtain a table mapping host function pointers to device function pointers.
+  auto CallTablePairOrErr = setupIndirectCallTable(Plugin, *this, Image);
+  if (!CallTablePairOrErr)
+    return CallTablePairOrErr.takeError();
+
   DeviceEnvironmentTy DeviceEnvironment;
   DeviceEnvironment.DebugKind = OMPX_DebugKind;
   DeviceEnvironment.NumDevices = Plugin.getNumDevices();
   // TODO: The device ID used here is not the real device ID used by OpenMP.
   DeviceEnvironment.DeviceNum = DeviceId;
   DeviceEnvironment.DynamicMemSize = OMPX_SharedMemorySize;
   DeviceEnvironment.ClockFrequency = getClockFrequency();
+  DeviceEnvironment.IndirectCallTable =
+      reinterpret_cast<uintptr_t>(CallTablePairOrErr->first);
+  DeviceEnvironment.IndirectCallTableSize = CallTablePairOrErr->second;
 
   // Create the metainfo of the device environment global.
   GlobalTy DevEnvGlobal("__omp_rtl_device_environment",
                         sizeof(DeviceEnvironmentTy), &DeviceEnvironment);
 
   // Write device environment values to the device.
   GenericGlobalHandlerTy &GHandler = Plugin.getGlobalHandler();
   if (auto Err = GHandler.writeGlobalToDevice(*this, Image, DevEnvGlobal)) {

openmp/libomptarget/src/rtl.cpp

static void registerGlobalCtorsDtorsForImage(__tgt_bin_desc *Desc,
                                              RTLInfoTy *RTL) {
 
   for (int32_t I = 0; I < RTL->NumberOfDevices; ++I) {
     DeviceTy &Device = *PM->Devices[RTL->Idx + I];
     Device.PendingGlobalsMtx.lock();
     Device.HasPendingGlobals = true;
     for (__tgt_offload_entry *Entry = Img->EntriesBegin;
          Entry != Img->EntriesEnd; ++Entry) {
+      // Globals are not callable and use a different set of flags.
+      if (Entry->size != 0)
+        continue;
+
       if (Entry->flags & OMP_DECLARE_TARGET_CTOR) {
         DP("Adding ctor " DPxMOD " to the pending list.\n",
            DPxPTR(Entry->addr));
         Device.PendingCtorsDtors[Desc].PendingCtors.push_back(Entry->addr);
       } else if (Entry->flags & OMP_DECLARE_TARGET_DTOR) {
         // Dtors are pushed in reverse order so they are executed from end
         // to beginning when unregistering the library!
         DP("Adding dtor " DPxMOD " to the pending list.\n",

openmp/libomptarget/test/api/omp_indirect_call.c

+// RUN: %libomptarget-compile-run-and-check-generic
+
+#include <assert.h>
+#include <stdio.h>
+
+#pragma omp begin declare variant match(device = {kind(gpu)})
+// Provided by the runtime.
+void *__llvm_omp_indirect_call_lookup(void *host_ptr);
+#pragma omp declare target to(__llvm_omp_indirect_call_lookup)                 \
+    device_type(nohost)
+#pragma omp end declare variant
+
+#pragma omp begin declare variant match(device = {kind(cpu)})
+// We assume unified addressing on the CPU target.
+void *__llvm_omp_indirect_call_lookup(void *host_ptr) { return host_ptr; }
+#pragma omp end declare variant
+
+#pragma omp begin declare target indirect
+void foo(int *x) { *x = *x + 1; }
+void bar(int *x) { *x = *x + 1; }
+void baz(int *x) { *x = *x + 1; }
+#pragma omp end declare target
+
+int main() {
+  void *foo_ptr = foo;
+  void *bar_ptr = bar;
+  void *baz_ptr = baz;
+
+  int count = 0;
+  void *foo_res;
+  void *bar_res;
+  void *baz_res;
+#pragma omp target map(to : foo_ptr, bar_ptr, baz_ptr) map(tofrom : count)
+  {
+    foo_res = __llvm_omp_indirect_call_lookup(foo_ptr);
+    ((void (*)(int *))foo_res)(&count);
+    bar_res = __llvm_omp_indirect_call_lookup(bar_ptr);
+    ((void (*)(int *))bar_res)(&count);
+    baz_res = __llvm_omp_indirect_call_lookup(baz_ptr);
+    ((void (*)(int *))baz_res)(&count);
+  }
+
+  assert(count == 3 && "Calling failed");
+
+  // CHECK: PASS
+  printf("PASS\n");
+}