[OpenMP][NFCI] Unify registration of target entries across plugins [WIP]

[jdoerfert]

Summary

All plugins need to read the target entries and "register" them to some
degree. Almost all of that code is common across plugins and as we
already have multiple now it's time to extract the code.

[jdoerfert]

Make it compile, also extrace device environment handling

openmp/libomptarget/plugins/amdgpu/src/rtl.cpp

 #include "omptargetplugin.h"
 #include "print_tracing.h"
 
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 #include "llvm/Frontend/OpenMP/OMPGridValues.h"
 
 using namespace llvm::omp::plugin;
 
+struct GenericPluginAPIPayloadTy {
+  SymbolInfoTableTy *SymbolInfoTable = nullptr;
+  KernelInfoTableTy *KernelInfoTable = nullptr;
+  GenericPluginAPIPayloadTy(SymbolInfoTableTy &SIT) : SymbolInfoTable(&SIT) {}
+  GenericPluginAPIPayloadTy(KernelInfoTableTy &KIT) : KernelInfoTable(&KIT) {}
+  GenericPluginAPIPayloadTy(SymbolInfoTableTy &SIT, KernelInfoTableTy &KIT)
+      : SymbolInfoTable(&SIT), KernelInfoTable(&KIT) {}
+};
+
 int32_t llvm::omp::plugin::GlobalHandlerTy::getGlobalMetadataFromDevice(
-    int32_t DeviceId, GlobalTy &DeviceGlobal, void *SymbolInfoTablePtr) {
+    int32_t DeviceId, GlobalTy &DeviceGlobal, void *Payload) {
   SymbolInfoTableTy &SymbolInfoTable =
-      *static_cast<SymbolInfoTableTy *>(SymbolInfoTablePtr);
+      *static_cast<GenericPluginAPIPayloadTy *>(Payload)->SymbolInfoTable;
   void *DevPtr;
   unsigned DevSize;
   const char *Name = DeviceGlobal.getName().c_str();
   hsa_status_t Err = interop_hsa_get_symbol_info(SymbolInfoTable, DeviceId,
                                                  Name, &DevPtr, &DevSize);
 
   if (Err != HSA_STATUS_SUCCESS) {
     // Inform the user what symbol prevented offloading
KernelTy(llvm::omp::OMPTgtExecModeFlags _ExecutionMode, int16_t _ConstWGSize,
     if (KernelArgPoolMap.find(N) == KernelArgPoolMap.end()) {
       KernelArgPoolMap.insert(
           std::make_pair(N, std::unique_ptr<KernelArgPool>(new KernelArgPool(
                                 _kernarg_segment_size, KernArgMemoryPool))));
     }
   }
 };
 
-/// List that contains all the kernels.
-/// FIXME: we may need this to be per device and per library.
-std::list<KernelTy> KernelsList;
-
 template <typename Callback> static hsa_status_t FindAgents(Callback CB) {
 
   hsa_status_t err =
       hsa::iterate_agents([&](hsa_agent_t agent) -> hsa_status_t {
         hsa_device_type_t device_type;
         // get_info fails iff HSA runtime not yet initialized
         hsa_status_t err =
             hsa_agent_get_info(agent, HSA_AGENT_INFO_DEVICE, &device_type);
if (S == HSA_STATUS_SUCCESS) {
         // Can't call into HSA to get a string from the integer
         DP("Shutting down HSA failed: %d\n", Err);
       }
     }
   }
 };
 
 /// Class containing all the device information
-class RTLDeviceInfoTy {
+class RTLDeviceInfoTy final : public DeviceInterfaceTy {
+  std::vector<std::list<KernelTy>> KernelLists;
+
   HSALifetime HSA; // First field => constructed first and destructed last
   std::vector<std::list<FuncOrGblEntryTy>> FuncGblEntries;
 
   struct QueueDeleter {
     void operator()(hsa_queue_t *Q) {
       if (Q) {
         hsa_status_t Err = hsa_queue_destroy(Q);
         if (Err != HSA_STATUS_SUCCESS) {
public:
   }
 
   hsa_status_t freesignalpool_memcpy_h2d(void *dest, const void *src,
                                          size_t size, int32_t deviceId) {
     return freesignalpool_memcpy(dest, src, size, impl_memcpy_h2d, deviceId);
   }
 
   // Record entry point associated with device
-  void addOffloadEntry(int32_t device_id, __tgt_offload_entry entry) {
+  void addOffloadEntry(int32_t device_id,
+                       const __tgt_offload_entry entry) override {
     assert(device_id < (int32_t)FuncGblEntries.size() &&
            "Unexpected device id!");
     FuncOrGblEntryTy &E = FuncGblEntries[device_id].back();
 
     E.Entries.push_back(entry);
   }
 
+  GlobalHandlerTy *getGlobalHandler() override { return &GlobalHandler; }
+
+  int64_t getRequiresFlags() override { return RequiresFlags; }
+
+  void *constructKernelEntry(int32_t DeviceId, const __tgt_device_image *Image,
+                             const __tgt_offload_entry *KernelEntry,
+                             void *Payload) override {
+    GenericPluginAPIPayloadTy &GPPayload =
+        *static_cast<GenericPluginAPIPayloadTy *>(Payload);
+    KernelInfoTableTy &KernelInfoTable = *GPPayload.KernelInfoTable;
+
+    uint32_t kernarg_segment_size;
+    hsa_status_t err = interop_hsa_get_kernel_info(
+        KernelInfoTable, DeviceId, KernelEntry->name,
+        HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_KERNARG_SEGMENT_SIZE,
+        &kernarg_segment_size);
+    (void)err;
+
+    // get flat group size if present, else Default_WG_Size
+    int16_t WGSizeVal = RTLDeviceInfoTy::Default_WG_Size;
+
+    // get Kernel Descriptor if present.
+    // Keep struct in sync wih getTgtAttributeStructQTy in CGOpenMPRuntime.cpp
+    struct KernDescValType {
+      uint16_t Version;
+      uint16_t TSize;
+      uint16_t WG_Size;
+    };
+
+    const size_t ImageSize =
+        (char *)Image->ImageEnd - (char *)Image->ImageStart;
+
+    // Read the kernel description global from the binary.
+    StaticGlobalTy<KernDescValType> KernelDescGlobal(KernelEntry->name,
+                                                     "_kern_desc");
+    if (!GlobalHandler.readGlobalFromImage(
+            DeviceId, KernelDescGlobal, (char *)Image->ImageStart, ImageSize)) {
+      // No kernel description available, fallback to work group size global:
+      // Read work group size global from the binary.
+      StaticGlobalTy<int16_t> WGSizeGlobal(KernelEntry->name, "_wg_size");
+      if (!GlobalHandler.readGlobalFromImage(
+              DeviceId, WGSizeGlobal, (char *)Image->ImageStart, ImageSize)) {
+        WGSizeGlobal.setValue(WGSizeVal);
+        INFO(OMP_INFOTYPE_DATA_TRANSFER, DeviceId,
+             "Failed to work group size for %s, defaulting to %i.",
+             KernelEntry->name, WGSizeVal);
+      }
+      WGSizeVal = WGSizeGlobal.getValue();
+    } else if (KernelDescGlobal.getValue().WG_Size) {
+      if (sizeof(KernDescValType) != KernelDescGlobal.getValue().TSize)
+        DP("KernDescVal size %lu does not match advertized size %d for '%s'\n",
+           sizeof(KernDescValType), KernelDescGlobal.getValue().TSize,
+           KernelDescGlobal.getName().c_str());
+      WGSizeVal = KernelDescGlobal.getValue().WG_Size;
+    }
+
+    // Read execution mode global from the binary
+    StaticGlobalTy<llvm::omp::OMPTgtExecModeFlags> ExecModeGlobal(
+        KernelEntry->name, "_exec_mode");
+    if (!GlobalHandler.readGlobalFromImage(
+            DeviceId, ExecModeGlobal, (char *)Image->ImageStart, ImageSize)) {
+      INFO(OMP_INFOTYPE_DATA_TRANSFER, DeviceId,
+           "Failed to read execution mode for %s, defaulting to SPMD.",
+           KernelEntry->name);
+      ExecModeGlobal.setValue(llvm::omp::OMP_TGT_EXEC_MODE_SPMD);
+    }
+
+    auto &KernelList = KernelLists[DeviceId];
+    void *CallStackAddr = nullptr;
+    KernelList.push_back(KernelTy(ExecModeGlobal.getValue(), WGSizeVal,
+                                  DeviceId, CallStackAddr, KernelEntry->name,
+                                  kernarg_segment_size, KernArgPool));
+    return &KernelList.back();
+  }
+
   // Return true if the entry is associated with device
   bool findOffloadEntry(int32_t device_id, void *addr) {
     assert(device_id < (int32_t)FuncGblEntries.size() &&
            "Unexpected device id!");
     FuncOrGblEntryTy &E = FuncGblEntries[device_id].back();
 
     for (auto &it : E.Entries) {
       if (it.addr == addr)
RTLDeviceInfoTy() {
 
     if (NumberOfDevices == 0) {
       DP("There are no devices supporting HSA.\n");
       return;
     } else {
       DP("There are %d devices supporting HSA.\n", NumberOfDevices);
     }
 
+    initDeviceInterface(NumberOfDevices);
+
     // Init the device info
     HSAQueues.resize(NumberOfDevices);
     FuncGblEntries.resize(NumberOfDevices);
     ThreadsPerGroup.resize(NumberOfDevices);
     ComputeUnits.resize(NumberOfDevices);
     GPUName.resize(NumberOfDevices);
     GroupsPerDevice.resize(NumberOfDevices);
     WarpSize.resize(NumberOfDevices);
__tgt_target_table *__tgt_rtl_load_binary_locked(int32_t device_id,
   //
   // - Allocate a large array per-gpu (could be moved to init_device)
   //   - Read a uint64_t at symbol omptarget_nvptx_device_State_size
   //   - Allocate at least that many bytes of gpu memory
   //   - Zero initialize it
   //   - Write the pointer to the symbol omptarget_nvptx_device_State
   //
   // - Pulls some per-kernel information together from various sources and
-  //   records it in the KernelsList for quicker access later
+  //   records it in the KernelList for quicker access later
   //
   // The initialization can be done before or after loading the image onto the
   // gpu. This function presently does a mixture. Using the hsa api to get/set
   // the information is simpler to implement, in exchange for more complicated
   // runtime behaviour. E.g. launching a kernel or using dma to get eight bytes
   // back from the gpu vs a hashtable lookup on the host.
 
   const size_t img_size = (char *)image->ImageEnd - (char *)image->ImageStart;
 
   DeviceInfo.clearOffloadEntriesTable(device_id);
 
   // We do not need to set the ELF version because the caller of this function
   // had to do that to decide the right runtime to use
 
   if (!elf_machine_id_is_amdgcn(image)) {
     return NULL;
   }
 
-  // send device environment data to the device
-  {
-    uint32_t DynamicMemorySize = 0;
-    uint32_t DebugKind = 0;
-    if (const char *EnvStr = getenv("LIBOMPTARGET_DEVICE_RTL_DEBUG"))
-      DebugKind = std::stoi(EnvStr);
-
-    // TODO: The device ID used here is not the real device ID used by OpenMP.
-    StaticGlobalTy<DeviceEnvironmentTy> DeviceEnvGlobal(
-        "omptarget_device_environment", DebugKind,
-        static_cast<uint32_t>(DeviceInfo.NumberOfDevices),
-        static_cast<uint32_t>(device_id),
-        static_cast<uint32_t>(DynamicMemorySize));
-    auto &SymbolInfoTable = DeviceInfo.SymbolInfoTables[device_id];
-    // TODO: Implement "writeGlobalToImage" in the GlobalHandler.
-    if (!DeviceInfo.GlobalHandler.writeGlobalToDevice(
-            device_id, DeviceEnvGlobal, &SymbolInfoTable)) {
-      INFO(OMP_INFOTYPE_PLUGIN_KERNEL, device_id,
-           "Failed to write device environment, abort.");
-      // TODO: Check the device gfx name against the image gfx name.
-      return nullptr;
-    }
-  }
+  GenericPluginAPIPayloadTy Payload(DeviceInfo.SymbolInfoTables[device_id],
+                                    DeviceInfo.KernelInfoTables[device_id]);
+  if (!DeviceInfo.setupDeviceEnvironment(device_id, &Payload))
 
   DP("AMDGPU module successfully loaded!\n");
 
   {
     // the device_State array is either large value in bss or a void* that
     // needs to be assigned to a pointer to an array of size device_state_bytes
     // If absent, it has been deadstripped and needs no setup.
     StaticGlobalTy<void *> DeviceStateGlobal("omptarget_nvptx_device_State");
-    auto &SymbolInfoMap = DeviceInfo.SymbolInfoTables[device_id];
     if (!DeviceInfo.GlobalHandler.getGlobalMetadataFromDevice(
-            device_id, DeviceStateGlobal, &SymbolInfoMap)) {
+            device_id, DeviceStateGlobal, &Payload)) {
       DP("No device_state pointer symbol found, skipping initialization\n");
     } else {
       StaticGlobalTy<uint64_t> DeviceStateSizeGlobal(
           "omptarget_nvptx_device_State_size");
       if (!DeviceInfo.GlobalHandler.readGlobalFromImage(
               device_id, DeviceStateGlobal, (char *)image->ImageStart,
               img_size)) {
         DP("Can't initialize device_State, missing size information\n");
if (!DeviceInfo.GlobalHandler.getGlobalMetadataFromDevice(
       if (DeviceStateSize != DSS.second) {
         DP("Inconsistent sizes of device_State unsupported\n");
         return NULL;
       }
 
       // write ptr to device memory so it can be used by later kernels
       DeviceStateGlobal.setValue(DSS.first.get());
       if (!DeviceInfo.GlobalHandler.writeGlobalToDevice(
-              device_id, DeviceStateGlobal, &SymbolInfoMap)) {
+              device_id, DeviceStateGlobal, &Payload)) {
         DP("memcpy install of state_ptr failed\n");
         return NULL;
       }
     }
   }
 
-  // Here, we take advantage of the data that is appended after img_end to get
-  // the symbols' name we need to load. This data consist of the host entries
-  // begin and end as well as the target name (see the offloading linker script
-  // creation in clang compiler).
-
-  // Find the symbols in the module by name. The name can be obtain by
-  // concatenating the host entry name with the target name
-
-  __tgt_offload_entry *HostBegin = image->EntriesBegin;
-  __tgt_offload_entry *HostEnd = image->EntriesEnd;
-
-  // TODO: This is basically the same in the AMDGPU and CUDA plugin,
-  //       refactor.
-  for (__tgt_offload_entry *e = HostBegin; e != HostEnd; ++e) {
-
-    if (!e->addr) {
-      // The host should have always something in the address to
-      // uniquely identify the target region.
-      DP("Analyzing host entry '<null>' (size = %lld)...\n",
-         (unsigned long long)e->size);
-      return NULL;
-    }
-
-    if (e->size) {
-      __tgt_offload_entry entry = *e;
-
-      StaticGlobalTy<void *> Global(e->name);
-      auto &SymbolInfoMap = DeviceInfo.SymbolInfoTables[device_id];
-      if (!DeviceInfo.GlobalHandler.getGlobalMetadataFromDevice(
-              device_id, Global, &SymbolInfoMap))
-        return nullptr;
-      entry.addr = Global.getPtr();
-
-      DeviceInfo.addOffloadEntry(device_id, entry);
-
-      if (DeviceInfo.RequiresFlags & OMP_REQ_UNIFIED_SHARED_MEMORY &&
-          e->flags & OMP_DECLARE_TARGET_LINK) {
-        // If unified memory is present any target link variables
-        // can access host addresses directly. There is no longer a
-        // need for device copies.
-        Global.setValue(e->addr);
-        if (!DeviceInfo.GlobalHandler.writeGlobalToDevice(device_id, Global,
-                                                          &SymbolInfoMap))
-          return nullptr;
-      }
-
-      continue;
-    }
-
-    DP("to find the kernel name: %s size: %lu\n", e->name, strlen(e->name));
-
-    uint32_t kernarg_segment_size;
-    auto &KernelInfoMap = DeviceInfo.KernelInfoTables[device_id];
-    hsa_status_t err = interop_hsa_get_kernel_info(
-        KernelInfoMap, device_id, e->name,
-        HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_KERNARG_SEGMENT_SIZE,
-        &kernarg_segment_size);
-    (void)err;
-
-    // each arg is a void * in this openmp implementation
-    uint32_t arg_num = kernarg_segment_size / sizeof(void *);
-    std::vector<size_t> arg_sizes(arg_num);
-    for (std::vector<size_t>::iterator it = arg_sizes.begin();
-         it != arg_sizes.end(); it++) {
-      *it = sizeof(void *);
-    }
-
-    // get flat group size if present, else Default_WG_Size
-    int16_t WGSizeVal = RTLDeviceInfoTy::Default_WG_Size;
-
-    // get Kernel Descriptor if present.
-    // Keep struct in sync wih getTgtAttributeStructQTy in CGOpenMPRuntime.cpp
-    struct KernDescValType {
-      uint16_t Version;
-      uint16_t TSize;
-      uint16_t WG_Size;
-    };
-
-    // Read the kernel description global from the binary.
-    StaticGlobalTy<KernDescValType> KernelDescGlobal(e->name, "_kern_desc");
-    if (!DeviceInfo.GlobalHandler.readGlobalFromImage(
-            device_id, KernelDescGlobal, (char *)image->ImageStart, img_size)) {
-      // No kernel description available, fallback to work group size global:
-      // Read work group size global from the binary.
-      StaticGlobalTy<int16_t> WGSizeGlobal(e->name, "_wg_size");
-      if (!DeviceInfo.GlobalHandler.readGlobalFromImage(
-              device_id, WGSizeGlobal, (char *)image->ImageStart, img_size)) {
-        WGSizeGlobal.setValue(WGSizeVal);
-        INFO(OMP_INFOTYPE_DATA_TRANSFER, device_id,
-             "Failed to work group size for %s, defaulting to %i.", e->name,
-             WGSizeVal);
-      }
-      WGSizeVal = WGSizeGlobal.getValue();
-    } else if (KernelDescGlobal.getValue().WG_Size) {
-      if (sizeof(KernDescValType) != KernelDescGlobal.getValue().TSize)
-        DP("KernDescVal size %lu does not match advertized size %d for '%s'\n",
-           sizeof(KernDescValType), KernelDescGlobal.getValue().TSize,
-           KernelDescGlobal.getName().c_str());
-      WGSizeVal = KernelDescGlobal.getValue().WG_Size;
-    }
-
-    // Read execution mode global from the binary
-    StaticGlobalTy<llvm::omp::OMPTgtExecModeFlags> ExecModeGlobal(e->name,
-                                                                  "_exec_mode");
-    if (!DeviceInfo.GlobalHandler.readGlobalFromImage(
-            device_id, ExecModeGlobal, (char *)image->ImageStart, img_size)) {
-      INFO(OMP_INFOTYPE_DATA_TRANSFER, device_id,
-           "Failed to read execution mode for %s, defaulting to SPMD.",
-           e->name);
-      ExecModeGlobal.setValue(llvm::omp::OMP_TGT_EXEC_MODE_SPMD);
-    }
-
-    void *CallStackAddr = nullptr;
-    KernelsList.push_back(
-        KernelTy(ExecModeGlobal.getValue(), WGSizeVal, device_id, CallStackAddr,
-                 e->name, kernarg_segment_size, DeviceInfo.KernArgPool));
-    __tgt_offload_entry entry = *e;
-    entry.addr = (void *)&KernelsList.back();
-    DeviceInfo.addOffloadEntry(device_id, entry);
-    DP("Entry point %ld maps to %s\n", e - HostBegin, e->name);
-  }
+  if (!DeviceInfo.registerOffloadEntries(device_id, image, &Payload))
 
   return DeviceInfo.getOffloadEntriesTable(device_id);
 }
 
 void *__tgt_rtl_data_alloc(int device_id, int64_t size, void *, int32_t kind) {
   void *ptr = NULL;
   assert(device_id < DeviceInfo.NumberOfDevices && "Device ID too large");
 

openmp/libomptarget/plugins/common/CMakeLists.txt

 ##===----------------------------------------------------------------------===##
 #
 # 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
 #
 ##===----------------------------------------------------------------------===##
 #
 # Common parts which can be used by all plugins
 #
 ##===----------------------------------------------------------------------===##
 
 add_subdirectory(DeviceInterface)
 add_subdirectory(elf_common)
 add_subdirectory(MemoryManager)
 add_subdirectory(GlobalHandler)
+add_subdirectory(KernelHandler)

openmp/libomptarget/plugins/common/DeviceInterface/CMakeLists.txt

 ##===----------------------------------------------------------------------===##
 #
 # 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
 #
 ##===----------------------------------------------------------------------===##
 
-add_library(DeviceInterface INTERFACE)
+add_library(DeviceInterface OBJECT DeviceInterface.cpp)
+
+# Build DeviceInterface with PIC to be able to link it with plugin shared libraries.
+set_property(TARGET DeviceInterface PROPERTY POSITION_INDEPENDENT_CODE ON)
+llvm_update_compile_flags(DeviceInterface)
+set(LINK_LLVM_LIBS LLVMSupport LLVMFrontendOpenMP)
+if (LLVM_LINK_LLVM_DYLIB)
+  set(LINK_LLVM_LIBS LLVM)
+endif()
+target_link_libraries(DeviceInterface PRIVATE GlobalHandler ${LINK_LLVM_LIBS})
+include_directories(${LIBOMPTARGET_LLVM_INCLUDE_DIRS})
+add_dependencies(DeviceInterface ${LINK_LLVM_LIBS})
 
 target_include_directories(DeviceInterface INTERFACE ${CMAKE_CURRENT_SOURCE_DIR})

openmp/libomptarget/plugins/common/DeviceInterface/DeviceInterface.h

 //===- DeviceInterface.h - Target independent plugin device interface -----===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_OPENMP_LIBOMPTARGET_PLUGINS_COMMON_DEVICEINTERFACE_DEVICEINTERFACE_H
 #define LLVM_OPENMP_LIBOMPTARGET_PLUGINS_COMMON_DEVICEINTERFACE_DEVICEINTERFACE_H
 
-#include "omptarget.h"
+#include <list>
+#include <vector>
 
-extern "C" {
-int32_t __tgt_rtl_data_submit(int32_t ID, void *TargetPtr, void *HostPtr,
-                              int64_t Size);
-int32_t __tgt_rtl_data_submit_async(int32_t ID, void *TargetPtr, void *HostPtr,
-                                    int64_t Size, __tgt_async_info *AsyncInfo);
-int32_t __tgt_rtl_data_retrieve(int32_t ID, void *HostPtr, void *TargetPtr,
-                                int64_t Size);
-int32_t __tgt_rtl_data_retrieve_async(int32_t ID, void *HostPtr,
-                                      void *TargetPtr, int64_t Size,
-                                      __tgt_async_info *AsyncInfo);
-}
+#include "Debug.h"
+#include "DeviceEnvironment.h"
+#include "omptarget.h"
 
 namespace llvm {
 namespace omp {
 namespace plugin {
 
+int32_t memcpyDtoH(int32_t DeviceId, void *Dst, const void *Src, int32_t Size);
+int32_t memcpyHtoD(int32_t DeviceId, void *Dst, const void *Src, int32_t Size);
+
+class GlobalHandlerTy;
+
 struct DeviceInterfaceTy {
-  static int32_t memcpyDtoH(int32_t DeviceId, void *Dst, const void *Src,
-                            int32_t Size) {
-    return __tgt_rtl_data_retrieve(DeviceId, Dst, const_cast<void *>(Src),
-                                   Size);
+  DeviceEnvironmentTy DeviceEnvironment;
+
+  virtual ~DeviceInterfaceTy() {}
+
+  int32_t setupDeviceEnvironment(int32_t DeviceId, void *Payload);
+
+  int32_t registerOffloadEntries(int32_t DeviceId,
+                                 const __tgt_device_image *Image,
+                                 void *Payload);
+
+protected:
+  void initDeviceInterface(int32_t NumDevices) {
+    DeviceEnvironment.NumDevices = NumDevices;
+
+    if (const char *EnvStr = getenv("LIBOMPTARGET_DEVICE_RTL_DEBUG")) {
+      DeviceEnvironment.DebugKind = std::stoi(EnvStr);
+    }
+    if (const char *EnvStr = getenv("LIBOMPTARGET_SHARED_MEMORY_SIZE")) {
+      DeviceEnvironment.DynamicMemSize = std::stoi(EnvStr);
+      DP("Parsed LIBOMPTARGET_SHARED_MEMORY_SIZE = %" PRIu64 "\n",
+         DeviceEnvironment.DynamicMemSize);
     }
-  static int32_t memcpyHtoD(int32_t DeviceId, void *Dst, const void *Src,
-                            int32_t Size) {
-    return __tgt_rtl_data_submit(DeviceId, Dst, const_cast<void *>(Src), Size);
   }
+
+  virtual void addOffloadEntry(int32_t DeviceId,
+                               const __tgt_offload_entry Entry) = 0;
+  virtual GlobalHandlerTy *getGlobalHandler() = 0;
+
+  virtual int64_t getRequiresFlags() = 0;
+
+  virtual void *constructKernelEntry(int32_t DeviceId,
+                                     const __tgt_device_image *Image,
+                                     const __tgt_offload_entry *KernelEntry,
+                                     void *Payload) = 0;
+
+private:
+  int32_t registerGlobalOffloadEntry(int32_t DeviceId,
+                                     const __tgt_device_image *Image,
+                                     const __tgt_offload_entry *GlobalEntry,
+                                     void *Payload);
+  int32_t registerKernelOffloadEntry(int32_t DeviceId,
+                                     const __tgt_device_image *Image,
+                                     const __tgt_offload_entry *KernelEntry,
+                                     void *Payload);
 };
 
 } // namespace plugin
 } // namespace omp
 } // namespace llvm
 
 #endif // LLVM_OPENMP_LIBOMPTARGET_PLUGINS_COMMON_DEVICEINTERFACE_DEVICEINTERFACE_H

openmp/libomptarget/plugins/common/DeviceInterface/DeviceInterface.cpp

+//===- DeviceInterface.cpp - Target independent plugin device interface ---===//
+//
+// 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 "DeviceInterface.h"
+#include "Debug.h"
+#include "GlobalHandler.h"
+#include "omptarget.h"
+
+extern "C" {
+int32_t __tgt_rtl_data_submit(int32_t ID, void *TargetPtr, void *HostPtr,
+                              int64_t Size);
+int32_t __tgt_rtl_data_submit_async(int32_t ID, void *TargetPtr, void *HostPtr,
+                                    int64_t Size, __tgt_async_info *AsyncInfo);
+int32_t __tgt_rtl_data_retrieve(int32_t ID, void *HostPtr, void *TargetPtr,
+                                int64_t Size);
+int32_t __tgt_rtl_data_retrieve_async(int32_t ID, void *HostPtr,
+                                      void *TargetPtr, int64_t Size,
+                                      __tgt_async_info *AsyncInfo);
+}
+
+using namespace llvm;
+using namespace omp;
+using namespace plugin;
+
+int32_t DeviceInterfaceTy::setupDeviceEnvironment(int32_t DeviceId,
+                                                  void *Payload) {
+
+  // TODO: The device ID used here is not the real device ID used by OpenMP.
+  DeviceEnvironment.DeviceNum = DeviceId;
+
+  GlobalHandlerTy &GlobalHandler = *getGlobalHandler();
+  GlobalTy DeviceEnvGlobal("omptarget_device_environment",
+                           sizeof(DeviceEnvironmentTy), &DeviceEnvironment);
+  // TODO: Implement "writeGlobalToImage" in the GlobalHandler.
+  if (!GlobalHandler.writeGlobalToDevice(DeviceId, DeviceEnvGlobal, Payload)) {
+    INFO(OMP_INFOTYPE_PLUGIN_KERNEL, DeviceId,
+         "Failed to write device environment, abort.");
+    // TODO: Check the device gfx name against the image gfx name.
+    return OFFLOAD_FAIL;
+  }
+  return OFFLOAD_SUCCESS;
+}
+
+int32_t DeviceInterfaceTy::registerOffloadEntries(
+    int32_t DeviceId, const __tgt_device_image *Image, void *Payload) {
+  const __tgt_offload_entry *Begin = Image->EntriesBegin;
+  const __tgt_offload_entry *End = Image->EntriesEnd;
+  for (const __tgt_offload_entry *Entry = Begin; Entry != End; ++Entry) {
+    if (!Entry->addr) {
+      // The host should have always something in the address to
+      // uniquely identify the entry.
+      INFO(OMP_INFOTYPE_ALL, DeviceId,
+           "Unexpected host entry without address (size: %ld), abort!\n",
+           Entry->size);
+      return OFFLOAD_FAIL;
+    }
+
+    if (Entry->size) {
+      if (!registerGlobalOffloadEntry(DeviceId, Image, Entry, Payload))
+        return OFFLOAD_FAIL;
+    } else {
+      if (!registerKernelOffloadEntry(DeviceId, Image, Entry, Payload))
+        return OFFLOAD_FAIL;
+    }
+  }
+  return OFFLOAD_SUCCESS;
+}
+
+int32_t DeviceInterfaceTy::registerGlobalOffloadEntry(
+    int32_t DeviceId, const __tgt_device_image *Image,
+    const __tgt_offload_entry *GlobalEntry, void *Payload) {
+
+  GlobalHandlerTy &GlobalHandler = *getGlobalHandler();
+  __tgt_offload_entry DeviceEntry = *GlobalEntry;
+  StaticGlobalTy<void *> Global(GlobalEntry->name);
+  if (!GlobalHandler.getGlobalMetadataFromDevice(DeviceId, Global, Payload))
+    return OFFLOAD_FAIL;
+  DeviceEntry.addr = Global.getPtr();
+
+  // Note: In the current implementation declare target variables
+  // can either be link or to. This means that once unified
+  // memory is activated via the requires directive, the variable
+  // can be used directly from the host in both cases.
+  // TODO: when variables types other than to or link are added,
+  // the below condition should be changed to explicitly
+  // check for to and link variables types.
+  if (getRequiresFlags() & OMP_REQ_UNIFIED_SHARED_MEMORY &&
+      (GlobalEntry->flags & OMP_DECLARE_TARGET_LINK)) {
+    // If unified memory is present any target link or to variables
+    // can access host addresses directly. There is no longer a
+    // need for device copies.
+    Global.setValue(GlobalEntry->addr);
+    if (!GlobalHandler.writeGlobalToDevice(DeviceId, Global, Payload))
+      return OFFLOAD_FAIL;
+  }
+
+  addOffloadEntry(DeviceId, DeviceEntry);
+  return OFFLOAD_SUCCESS;
+}
+
+int32_t DeviceInterfaceTy::registerKernelOffloadEntry(
+    int32_t DeviceId, const __tgt_device_image *Image,
+    const __tgt_offload_entry *KernelEntry, void *Payload) {
+  __tgt_offload_entry DeviceEntry = *KernelEntry;
+  void *Kernel = constructKernelEntry(DeviceId, Image, KernelEntry, Payload);
+  if (!Kernel)
+    return OFFLOAD_FAIL;
+  DeviceEntry.addr = (void *)Kernel;
+  addOffloadEntry(DeviceId, DeviceEntry);
+  return OFFLOAD_SUCCESS;
+}
+
+int32_t llvm::omp::plugin::memcpyDtoH(int32_t DeviceId, void *Dst,
+                                      const void *Src, int32_t Size) {
+  return __tgt_rtl_data_retrieve(DeviceId, Dst, const_cast<void *>(Src), Size);
+}
+int32_t llvm::omp::plugin::memcpyHtoD(int32_t DeviceId, void *Dst,
+                                      const void *Src, int32_t Size) {
+  return __tgt_rtl_data_submit(DeviceId, Dst, const_cast<void *>(Src), Size);
+}

openmp/libomptarget/plugins/common/GlobalHandler/CMakeLists.txt

 ##===----------------------------------------------------------------------===##
 #
 # 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
 #
 ##===----------------------------------------------------------------------===##
 
 add_library(GlobalHandler INTERFACE)
 
+target_link_libraries(GlobalHandler INTERFACE DeviceInterface)
 target_include_directories(GlobalHandler INTERFACE ${CMAKE_CURRENT_SOURCE_DIR})

openmp/libomptarget/plugins/common/GlobalHandler/GlobalHandler.h

int32_t moveGlobalBetweenDeviceAndHost(int32_t DeviceId,
     if (Err) {
       INFO(OMP_INFOTYPE_DATA_TRANSFER, DeviceId,
            "Failed to read global symbol metadata for '%s' from the device",
            HostGlobal.getName().c_str());
       return Err;
     }
 
     if (Device2Host)
-      Err = DeviceInterfaceTy::memcpyDtoH(DeviceId, HostGlobal.getPtr(),
-                                          DeviceGlobal.getPtr(),
-                                          HostGlobal.getSize());
+      Err = memcpyDtoH(DeviceId, HostGlobal.getPtr(), DeviceGlobal.getPtr(),
     else
-      Err = DeviceInterfaceTy::memcpyHtoD(DeviceId, DeviceGlobal.getPtr(),
-                                          HostGlobal.getPtr(),
-                                          HostGlobal.getSize());
+      Err = memcpyHtoD(DeviceId, DeviceGlobal.getPtr(), HostGlobal.getPtr(),
 
     if (Err) {
       INFO(OMP_INFOTYPE_DATA_TRANSFER, DeviceId,
            "Failed to %s %u bytes associated with global symbol '%s' %s "
            "the device",
            Device2Host ? "read" : "write", HostGlobal.getSize(),
            HostGlobal.getName().c_str(), Device2Host ? "from" : "to");
       return Err;

openmp/libomptarget/plugins/cuda/src/rtl.cpp

bool initializeDeviceStreamPool(const int DeviceId) {
     for (CUstream &S : StreamPool[DeviceId])
       if (!S)
         return false;
 
     return true;
   }
 };
 
-class DeviceRTLTy {
-  /// The debug/configuration kind we read from LIBOMPTARGET_DEVICE_RTL_DEBUG
-  uint32_t DebugKind;
+class DeviceRTLTy final : public DeviceInterfaceTy {
 
   int NumberOfDevices;
   // OpenMP environment properties
   int EnvNumTeams;
   int EnvTeamLimit;
   int EnvTeamThreadLimit;
   // OpenMP requires flags
   int64_t RequiresFlags;
-  // Amount of dynamic shared memory to use at launch.
-  uint64_t DynamicMemorySize;
 
   static constexpr const int HardTeamLimit = 1U << 16U; // 64k
   static constexpr const int HardThreadLimit = 1024;
   static constexpr const int DefaultNumTeams = 128;
   static constexpr const int DefaultNumThreads = 128;
 
   std::unique_ptr<StreamManagerTy> StreamManager;
   std::vector<DeviceDataTy> DeviceData;
class DeviceRTLTy final : public DeviceInterfaceTy {
   /// A vector of memory managers. Since the memory manager is non-copyable and
   // non-removable, we wrap them into std::unique_ptr.
   std::vector<std::unique_ptr<MemoryManagerTy>> MemoryManagers;
 
   /// Whether use memory manager
   bool UseMemoryManager = true;
 
   // Record entry point associated with device
-  void addOffloadEntry(const int DeviceId, const __tgt_offload_entry entry) {
+  void addOffloadEntry(int32_t DeviceId,
+                       const __tgt_offload_entry entry) override {
     FuncOrGblEntryTy &E = DeviceData[DeviceId].FuncGblEntries.back();
     E.Entries.push_back(entry);
   }
 
+  GlobalHandlerTy *getGlobalHandler() override { return &GlobalHandler; }
+
+  int64_t getRequiresFlags() override { return RequiresFlags; }
+
+  void *constructKernelEntry(int32_t DeviceId, const __tgt_device_image *Image,
+                             const __tgt_offload_entry *KernelEntry,
+                             void *Payload) override {
+    CUmodule Module = static_cast<CUmodule>(Payload);
+    CUfunction Func;
+    CUresult Err = cuModuleGetFunction(&Func, Module, KernelEntry->name);
+    // We keep this style here because we need the name
+    if (Err != CUDA_SUCCESS) {
+      REPORT("Loading '%s' Failed\n", KernelEntry->name);
+      CUDA_ERR_STRING(Err);
+      return nullptr;
+    }
+
+    DP("Entry point " DPxMOD " maps to %s (" DPxMOD ")\n", DPxPTR(KernelEntry),
+       KernelEntry->name, DPxPTR(Func));
+
+    StaticGlobalTy<llvm::omp::OMPTgtExecModeFlags> ExecModeGlobal(
+        KernelEntry->name, "_exec_mode");
+    // TODO: We should be able to read it from the image instead.
+    if (!GlobalHandler.readGlobalFromDevice(DeviceId, ExecModeGlobal, Module)) {
+      INFO(OMP_INFOTYPE_DATA_TRANSFER, DeviceId,
+           "Failed to read execution mode for %s, defaulting to SPMD.",
+           KernelEntry->name);
+      ExecModeGlobal.setValue(llvm::omp::OMP_TGT_EXEC_MODE_SPMD);
+    }
+
+    auto &KernelList = DeviceData[DeviceId].KernelsList;
+    KernelList.emplace_back(Func, ExecModeGlobal.getValue());
+    return &KernelList.back();
+  }
+
   // Return a pointer to the entry associated with the pointer
   const __tgt_offload_entry *getOffloadEntry(const int DeviceId,
                                              const void *Addr) const {
     for (const __tgt_offload_entry &Itr :
          DeviceData[DeviceId].FuncGblEntries.back().Entries)
       if (Itr.addr == Addr)
         return &Itr;
 

 
 public:
   // This class should not be copied
   DeviceRTLTy(const DeviceRTLTy &) = delete;
   DeviceRTLTy(DeviceRTLTy &&) = delete;
 
   DeviceRTLTy()
       : NumberOfDevices(0), EnvNumTeams(-1), EnvTeamLimit(-1),
-        EnvTeamThreadLimit(-1), RequiresFlags(OMP_REQ_UNDEFINED),
-        DynamicMemorySize(0) {
+        EnvTeamThreadLimit(-1), RequiresFlags(OMP_REQ_UNDEFINED) {
 
     DP("Start initializing CUDA\n");
 
     CUresult Err = cuInit(0);
     if (Err == CUDA_ERROR_INVALID_HANDLE) {
       // Can't call cuGetErrorString if dlsym failed
       DP("Failed to load CUDA shared library\n");
       return;
     }
     if (!checkResult(Err, "Error returned from cuInit\n")) {
       return;
     }
 
     Err = cuDeviceGetCount(&NumberOfDevices);
     if (!checkResult(Err, "Error returned from cuDeviceGetCount\n"))
       return;
 
     if (NumberOfDevices == 0) {
       DP("There are no devices supporting CUDA.\n");
       return;
     }
 
+    initDeviceInterface(NumberOfDevices);
     DeviceData.resize(NumberOfDevices);
 
     // Get environment variables regarding teams
     if (const char *EnvStr = getenv("OMP_TEAM_LIMIT")) {
       // OMP_TEAM_LIMIT has been set
       EnvTeamLimit = std::stoi(EnvStr);
       DP("Parsed OMP_TEAM_LIMIT=%d\n", EnvTeamLimit);
     }
     if (const char *EnvStr = getenv("OMP_TEAMS_THREAD_LIMIT")) {
       // OMP_TEAMS_THREAD_LIMIT has been set
       EnvTeamThreadLimit = std::stoi(EnvStr);
       DP("Parsed OMP_TEAMS_THREAD_LIMIT=%d\n", EnvTeamThreadLimit);
     }
     if (const char *EnvStr = getenv("OMP_NUM_TEAMS")) {
       // OMP_NUM_TEAMS has been set
       EnvNumTeams = std::stoi(EnvStr);
       DP("Parsed OMP_NUM_TEAMS=%d\n", EnvNumTeams);
     }
-    if (const char *EnvStr = getenv("LIBOMPTARGET_SHARED_MEMORY_SIZE")) {
-      // LIBOMPTARGET_SHARED_MEMORY_SIZE has been set
-      DynamicMemorySize = std::stoi(EnvStr);
-      DP("Parsed LIBOMPTARGET_SHARED_MEMORY_SIZE = %" PRIu64 "\n",
-         DynamicMemorySize);
-    }
 
     StreamManager =
         std::make_unique<StreamManagerTy>(NumberOfDevices, DeviceData);
 
     for (int I = 0; I < NumberOfDevices; ++I)
       DeviceAllocators.emplace_back(I, DeviceData);
 
     // Get the size threshold from environment variable
__tgt_target_table *loadBinary(const int DeviceId,
     Err = cuModuleLoadDataEx(&Module, Image->ImageStart, 0, nullptr, nullptr);
     if (!checkResult(Err, "Error returned from cuModuleLoadDataEx\n"))
       return nullptr;
 
     DP("CUDA module successfully loaded!\n");
 
     Modules.push_back(Module);
 
-    // Find the symbols in the module by name.
-    const __tgt_offload_entry *HostBegin = Image->EntriesBegin;
-    const __tgt_offload_entry *HostEnd = Image->EntriesEnd;
-
-    // TODO: This is basically the same in the AMDGPU and CUDA plugin,
-    //       refactor.
-    std::list<KernelTy> &KernelsList = DeviceData[DeviceId].KernelsList;
-    for (const __tgt_offload_entry *E = HostBegin; E != HostEnd; ++E) {
-      if (!E->addr) {
-        // We return nullptr when something like this happens, the host should
-        // have always something in the address to uniquely identify the target
-        // region.
-        DP("Invalid binary: host entry '<null>' (size = %zd)...\n", E->size);
-        return nullptr;
-      }
-
-      if (E->size) {
-        __tgt_offload_entry Entry = *E;
-        StaticGlobalTy<void *> Global(E->name);
-        if (!GlobalHandler.getGlobalMetadataFromDevice(DeviceId, Global,
-                                                       Module))
-          return nullptr;
-        Entry.addr = Global.getPtr();
+    if (!setupDeviceEnvironment(DeviceId, Module))
 
-        // Note: In the current implementation declare target variables
-        // can either be link or to. This means that once unified
-        // memory is activated via the requires directive, the variable
-        // can be used directly from the host in both cases.
-        // TODO: when variables types other than to or link are added,
-        // the below condition should be changed to explicitly
-        // check for to and link variables types:
-        // (RequiresFlags & OMP_REQ_UNIFIED_SHARED_MEMORY && (e->flags &
-        // OMP_DECLARE_TARGET_LINK || e->flags == OMP_DECLARE_TARGET_TO))
-        if (RequiresFlags & OMP_REQ_UNIFIED_SHARED_MEMORY) {
-          // If unified memory is present any target link or to variables
-          // can access host addresses directly. There is no longer a
-          // need for device copies.
-          Global.setValue(E->addr);
-          if (!GlobalHandler.writeGlobalToDevice(DeviceId, Global, Module))
-            return nullptr;
-        }
-
-        addOffloadEntry(DeviceId, Entry);
-
-        continue;
-      }
-
-      CUfunction Func;
-      Err = cuModuleGetFunction(&Func, Module, E->name);
-      // We keep this style here because we need the name
-      if (Err != CUDA_SUCCESS) {
-        REPORT("Loading '%s' Failed\n", E->name);
-        CUDA_ERR_STRING(Err);
-        return nullptr;
-      }
-
-      DP("Entry point " DPxMOD " maps to %s (" DPxMOD ")\n",
-         DPxPTR(E - HostBegin), E->name, DPxPTR(Func));
-
-      StaticGlobalTy<llvm::omp::OMPTgtExecModeFlags> ExecModeGlobal(
-          E->name, "_exec_mode");
-      // TODO: We should be able to read it from the image instead.
-      if (!GlobalHandler.readGlobalFromDevice(DeviceId, ExecModeGlobal,
-                                              Module)) {
-        INFO(OMP_INFOTYPE_DATA_TRANSFER, DeviceId,
-             "Failed to read execution mode for %s, defaulting to SPMD.",
-             E->name);
-        ExecModeGlobal.setValue(llvm::omp::OMP_TGT_EXEC_MODE_SPMD);
-      }
-
-      KernelsList.emplace_back(Func, ExecModeGlobal.getValue());
-
-      __tgt_offload_entry Entry = *E;
-      Entry.addr = &KernelsList.back();
-      addOffloadEntry(DeviceId, Entry);
-    }
-
-    // send device environment data to the device
-    {
-      if (const char *EnvStr = getenv("LIBOMPTARGET_DEVICE_RTL_DEBUG"))
-        DebugKind = std::stoi(EnvStr);
-
-      // TODO: The device ID used here is not the real device ID used by OpenMP.
-      StaticGlobalTy<DeviceEnvironmentTy> DeviceEnvGlobal(
-          "omptarget_device_environment", DebugKind,
-          static_cast<uint32_t>(NumberOfDevices),
-          static_cast<uint32_t>(DeviceId),
-          static_cast<uint32_t>(DynamicMemorySize));
-      if (!GlobalHandler.writeGlobalToDevice(DeviceId, DeviceEnvGlobal,
-                                             Module)) {
-        // TODO: I think this should be fatal.
-        INFO(OMP_INFOTYPE_DATA_TRANSFER, DeviceId,
-             "Failed to write device environment, continue without.");
-      }
-    }
+    if (!registerOffloadEntries(DeviceId, Image, Module))
 
     return getOffloadEntriesTable(DeviceId);
   }
 
   void *dataAlloc(const int DeviceId, const int64_t Size,
                   const TargetAllocTy Kind) {
     switch (Kind) {
     case TARGET_ALLOC_DEFAULT:
INFO(OMP_INFOTYPE_PLUGIN_KERNEL, DeviceId,
              : "(null)",
          CudaBlocksPerGrid, CudaThreadsPerBlock,
          (!IsSPMDMode ? (IsGenericMode ? "Generic" : "SPMD-Generic") : "SPMD"));
 
     CUstream Stream = getStream(DeviceId, AsyncInfo);
     Err = cuLaunchKernel(KernelInfo->Func, CudaBlocksPerGrid, /* gridDimY */ 1,
                          /* gridDimZ */ 1, CudaThreadsPerBlock,
                          /* blockDimY */ 1, /* blockDimZ */ 1,
-                         DynamicMemorySize, Stream, &Args[0], nullptr);
+                         DeviceEnvironment.DynamicMemSize, Stream, &Args[0],
+                         nullptr);
     if (!checkResult(Err, "Error returned from cuLaunchKernel\n"))
       return OFFLOAD_FAIL;
 
     DP("Launch of entry point at " DPxMOD " successful!\n",
        DPxPTR(TgtEntryPtr));
 
     return OFFLOAD_SUCCESS;
   }
int32_t llvm::omp::plugin::GlobalHandlerTy::getGlobalMetadataFromDevice(
   if (CUSize != DeviceGlobal.getSize()) {
     DP("Loading global '%s' - size mismatch (%zd != %zd)\n", Name, CUSize,
        DeviceGlobal.getSize());
     return OFFLOAD_FAIL;
   }
 
   DP("Entry point " DPxMOD " maps to global %s (" DPxMOD ")\n",
      DPxPTR(E - HostBegin), Name, DPxPTR(CUPtr));
-  DeviceGlobal.setPtr(reinterpret_cast<void*>(CUPtr));
+  DeviceGlobal.setPtr(reinterpret_cast<void *>(CUPtr));
   return OFFLOAD_SUCCESS;
 }
 
 // Exposed library API function
 #ifdef __cplusplus
 extern "C" {
 #endif