[libomptarget] Add allocator support for target memory

[grokos]

Summary

This patch adds the infrastructure for allocator support for target memory. Three allocators are introduced for device, host and shared memory. The corresponding API functions have the llvm_ prefix temporarily, until they become part of the OpenMP standard:

llvm_omp_target_alloc_device(size_t size, int device_num);

• Returns memory owned by device.
• Memory is only accessible by device.

llvm_omp_target_alloc_host(size_t size, int device_num);

• Returns non-migratable memory owned by host.
• Memory is accessible by host and device(s).

llvm_omp_target_alloc_shared(size_t size, int device_num);

• Returns migratable memory owned by host and device.
• Memory is accessible by host and device.

The corresponding host OpenMP runtime patch is D96669.

Ultimately, libomptarget needs support from each plugin to allocate memory using the desired method. __tgt_rtl_data_alloc has been extended with an extra kind argument which dictates the allocator to be used. The patch provides a trivial sample implementation using the generic-elf-64 plugin. For all other plugins, calling __tgt_rtl_data_alloc with a specific allocator (other than TARGET_ALLOC_DEFAULT) results in NULL being returned until said allocator has been implemented in the plugin.

[jdoerfert]

LG from my end, someone else?

[tianshilei1992]

openmp/libomptarget/plugins/generic-elf-64bit/src/rtl.cpp

260

I think we can extend __tgt_rtl_data_alloc here by adding a default argument kind such that it has the same semantic as __tgt_rtl_data_alloc. __tgt_rtl_data_alloc_explicit sounds weird...After all, it is internal interface so it will not break anything.

[grokos]

Removed __tgt_rtl_data_alloc_explicit; instead we pass an extra kind argument to __tgt_rtl_data_alloc.

[tianshilei1992]

openmp/libomptarget/src/api.cpp

40

maybe it's better to move this function to omptarget.cpp?

[jdoerfert]

LGTM, let's move functions around separatly. @tianshilei1992 any other comment?

[tianshilei1992]

In D97883#2617913, @jdoerfert wrote:

LGTM, let's move functions around separatly. @tianshilei1992 any other comment?

Sure, but why? Others LGTM.

[grokos]

Rabsed and moved targetAllocExplicit() under omptarget.cpp.

[tianshilei1992]

LGTM.

[JonChesterfield]

Apologies for missing this. Looks reasonable to me. I'll wire up the amdgcn plugin when I find some time.

openmp/libomptarget/include/omptarget.h

203

Is alloc_host expected to ignore the device_num argument?

[grokos]

openmp/libomptarget/include/omptarget.h

203

That's a good question. On the surface of it I would say yes, it makes sense to ignore the device ID when allocating on the host. However, bear in mind that these functions are not standardized yet, so the exact behavior may change to something that makes use of the device ID.

[RaviNarayanaswamy]

Reasonable to drop the device number for alloc_host as all devices should be able to access host memory.

openmp/libomptarget/include/omptarget.h

enum OpenMPOffloadingRequiresDirFlags {
   /// unified_address clause.
   OMP_REQ_UNIFIED_ADDRESS         = 0x004,
   /// unified_shared_memory clause.
   OMP_REQ_UNIFIED_SHARED_MEMORY   = 0x008,
   /// dynamic_allocators clause.
   OMP_REQ_DYNAMIC_ALLOCATORS      = 0x010
 };
 
+enum TargetAllocTy : int32_t {
+  TARGET_ALLOC_DEVICE = 0,
+  TARGET_ALLOC_HOST,
+  TARGET_ALLOC_SHARED,
+  TARGET_ALLOC_DEFAULT
+};
+
 /// This struct is a record of an entry point or global. For a function
 /// entry point the size is expected to be zero
 struct __tgt_offload_entry {
   void *addr;   // Pointer to the offload entry info (function or global)
   char *name;   // Name of the function or global
   size_t size;  // Size of the entry info (0 if it is a function)
   int32_t flags; // Flags associated with the entry, e.g. 'link'.
   int32_t reserved; // Reserved, to be used by the runtime library.
int omp_target_memcpy_rect(void *dst, void *src, size_t element_size,
                            const size_t *dst_offsets, const size_t *src_offsets,
                            const size_t *dst_dimensions,
                            const size_t *src_dimensions, int dst_device,
                            int src_device);
 int omp_target_associate_ptr(void *host_ptr, void *device_ptr, size_t size,
                              size_t device_offset, int device_num);
 int omp_target_disassociate_ptr(void *host_ptr, int device_num);
 
+/// Explicit target memory allocators
+/// Using the llvm_ prefix until they become part of the OpenMP standard.
+void *llvm_omp_target_alloc_device(size_t size, int device_num);
+void *llvm_omp_target_alloc_host(size_t size, int device_num);
+void *llvm_omp_target_alloc_shared(size_t size, int device_num);
+
 /// add the clauses of the requires directives in a given file
 void __tgt_register_requires(int64_t flags);
 
 /// adds a target shared library to the target execution image
 void __tgt_register_lib(__tgt_bin_desc *desc);
 
 /// removes a target shared library from the target execution image
 void __tgt_unregister_lib(__tgt_bin_desc *desc);

openmp/libomptarget/include/omptargetplugin.h

 // is allocated by omp_target_alloc() API). Return address of the
 // allocated data on the target that will be used by libomptarget.so to
 // initialize the target data mapping structures. These addresses are
 // used to generate a table of target variables to pass to
 // __tgt_rtl_run_region(). The __tgt_rtl_data_alloc() returns NULL in
 // case an error occurred on the target device.
 void *__tgt_rtl_data_alloc(int32_t ID, int64_t Size, void *HostPtr);
 
+// Explicit target memory allocator
+void *__tgt_rtl_data_alloc_explicit(int32_t ID, int64_t Size, int32_t Kind);
+
 // Pass the data content to the target device using the target address. In case
 // of success, return zero. Otherwise, return an error code.
 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 *AsyncInfoPtr);

openmp/libomptarget/plugins/exports

 VERS1.0 {
   global:
     __tgt_rtl_is_valid_binary;
     __tgt_rtl_is_data_exchangable;
     __tgt_rtl_number_of_devices;
     __tgt_rtl_init_requires;
     __tgt_rtl_init_device;
     __tgt_rtl_load_binary;
     __tgt_rtl_data_alloc;
+    __tgt_rtl_data_alloc_explicit;
     __tgt_rtl_data_submit;
     __tgt_rtl_data_submit_async;
     __tgt_rtl_data_retrieve;
     __tgt_rtl_data_retrieve_async;
     __tgt_rtl_data_exchange;
     __tgt_rtl_data_exchange_async;
     __tgt_rtl_data_delete;
     __tgt_rtl_run_target_team_region;

openmp/libomptarget/plugins/generic-elf-64bit/src/rtl.cpp

 //===----------------------------------------------------------------------===//
 
 #include <cassert>
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
 #include <dlfcn.h>
 #include <ffi.h>
 #include <gelf.h>
 #include <link.h>
 #include <list>
 #include <string>
 #include <vector>
 
 #include "Debug.h"
 #include "omptargetplugin.h"
 
__tgt_target_table *__tgt_rtl_load_binary(int32_t device_id,
   return DeviceInfo.getOffloadEntriesTable(device_id);
 }
 
 void *__tgt_rtl_data_alloc(int32_t device_id, int64_t size, void *hst_ptr) {
   void *ptr = malloc(size);
   return ptr;
 }
 
+// Sample implementation of explicit memory allocator. For this plugin all kinds
+// are equivalent to each other.
+void *__tgt_rtl_data_alloc_explicit(int32_t device_id, int64_t size,
+                                    int32_t kind) {
+  void *ptr = NULL;
+
+  switch (kind) {
+  case TARGET_ALLOC_DEVICE:
+  case TARGET_ALLOC_HOST:
+  case TARGET_ALLOC_SHARED:
+    ptr = malloc(size);
+    break;
+  default:
+    REPORT("Invalid target data allocation kind");
+  }
+
+  return ptr;
+}
+
 int32_t __tgt_rtl_data_submit(int32_t device_id, void *tgt_ptr, void *hst_ptr,
                               int64_t size) {
   memcpy(tgt_ptr, hst_ptr, size);
   return OFFLOAD_SUCCESS;
 }
 
 int32_t __tgt_rtl_data_retrieve(int32_t device_id, void *hst_ptr, void *tgt_ptr,
                                 int64_t size) {

openmp/libomptarget/src/api.cpp

 
 EXTERN int omp_get_initial_device(void) {
   TIMESCOPE();
   int hostDevice = omp_get_num_devices();
   DP("Call to omp_get_initial_device returning %d\n", hostDevice);
   return hostDevice;
 }
 
 EXTERN void *omp_target_alloc(size_t size, int device_num) {
   TIMESCOPE();
   DP("Call to omp_target_alloc for device %d requesting %zu bytes\n",
      device_num, size);
 
   if (size <= 0) {
     DP("Call to omp_target_alloc with non-positive length\n");
     return NULL;
   }
if (!device_is_ready(device_num)) {
     return OFFLOAD_FAIL;
   }
 
   DeviceTy &Device = PM->Devices[device_num];
   int rc = Device.disassociatePtr(host_ptr);
   DP("omp_target_disassociate_ptr returns %d\n", rc);
   return rc;
 }
+
+static void *targetAllocExplicit(size_t size, int device_num, int kind,
+                                 const char *name) {
+  TIMESCOPE();
+  DP("Call to %s for device %d requesting %zu bytes\n", name, device_num, size);
+
+  if (size <= 0) {
+    DP("Call to %s with non-positive length\n", name);
+    return NULL;
+  }
+
+  void *rc = NULL;
+
+  if (device_num == omp_get_initial_device()) {
+    rc = malloc(size);
+    DP("%s returns host ptr " DPxMOD "\n", name, DPxPTR(rc));
+    return rc;
+  }
+
+  if (!device_is_ready(device_num)) {
+    DP("%s returns NULL ptr\n", name);
+    return NULL;
+  }
+
+  DeviceTy &Device = PM->Devices[device_num];
+  rc = Device.allocDataExplicit(size, kind);
+  DP("%s returns device ptr " DPxMOD "\n", name, DPxPTR(rc));
+  return rc;
+}
+
+EXTERN void *llvm_omp_target_alloc_device(size_t size, int device_num) {
+  return targetAllocExplicit(size, device_num, TARGET_ALLOC_DEVICE, __func__);
+}
+
+EXTERN void *llvm_omp_target_alloc_host(size_t size, int device_num) {
+  return targetAllocExplicit(size, device_num, TARGET_ALLOC_HOST, __func__);
+}
+
+EXTERN void *llvm_omp_target_alloc_shared(size_t size, int device_num) {
+  return targetAllocExplicit(size, device_num, TARGET_ALLOC_SHARED, __func__);
+}

openmp/libomptarget/src/device.h

struct DeviceTy {
   // device memory allocation/deallocation routines
   /// Allocates \p Size bytes on the device and returns the address/nullptr when
   /// succeeds/fails. \p HstPtr is an address of the host data which the
   /// allocated target data will be associated with. If it is unknown, the
   /// default value of \p HstPtr is nullptr. Note: this function doesn't do
   /// pointer association. Actually, all the __tgt_rtl_data_alloc
   /// implementations ignore \p HstPtr.
   void *allocData(int64_t Size, void *HstPtr = nullptr);
+  /// Allocates \p Size bytes on the host or device or shared memory space,
+  /// depending on \p Kind.
+  void *allocDataExplicit(int64_t Size, int32_t Kind);
   /// Deallocates memory which \p TgtPtrBegin points at and returns
   /// OFFLOAD_SUCCESS/OFFLOAD_FAIL when succeeds/fails.
   int32_t deleteData(void *TgtPtrBegin);
 
   // Data transfer. When AsyncInfoPtr is nullptr, the transfer will be
   // synchronous.
   // Copy data from host to device
   int32_t submitData(void *TgtPtrBegin, void *HstPtrBegin, int64_t Size,

openmp/libomptarget/src/device.cpp

__tgt_target_table *DeviceTy::load_binary(void *Img) {
   RTL->Mtx.unlock();
   return rc;
 }
 
 void *DeviceTy::allocData(int64_t Size, void *HstPtr) {
   return RTL->data_alloc(RTLDeviceID, Size, HstPtr);
 }
 
+void *DeviceTy::allocDataExplicit(int64_t Size, int32_t Kind) {
+  if (RTL->data_alloc_explicit)
+    return RTL->data_alloc_explicit(RTLDeviceID, Size, Kind);
+  else
+    return nullptr;
+}
+
 int32_t DeviceTy::deleteData(void *TgtPtrBegin) {
   return RTL->data_delete(RTLDeviceID, TgtPtrBegin);
 }
 
 // Submit data to device
 int32_t DeviceTy::submitData(void *TgtPtrBegin, void *HstPtrBegin, int64_t Size,
                              AsyncInfoTy &AsyncInfo) {
   if (!AsyncInfo || !RTL->data_submit_async || !RTL->synchronize)

openmp/libomptarget/src/exports

global:
     omp_get_initial_device;
     omp_target_alloc;
     omp_target_free;
     omp_target_is_present;
     omp_target_memcpy;
     omp_target_memcpy_rect;
     omp_target_associate_ptr;
     omp_target_disassociate_ptr;
+    llvm_omp_target_alloc_host;
+    llvm_omp_target_alloc_shared;
+    llvm_omp_target_alloc_device;
     __kmpc_push_target_tripcount;
   local:
     *;
 };
 

openmp/libomptarget/src/rtl.h

 
 struct RTLInfoTy {
   typedef int32_t(is_valid_binary_ty)(void *);
   typedef int32_t(is_data_exchangable_ty)(int32_t, int32_t);
   typedef int32_t(number_of_devices_ty)();
   typedef int32_t(init_device_ty)(int32_t);
   typedef __tgt_target_table *(load_binary_ty)(int32_t, void *);
   typedef void *(data_alloc_ty)(int32_t, int64_t, void *);
+  typedef void *(data_alloc_explicit_ty)(int32_t, int64_t, int32_t);
   typedef int32_t(data_submit_ty)(int32_t, void *, void *, int64_t);
   typedef int32_t(data_submit_async_ty)(int32_t, void *, void *, int64_t,
                                         __tgt_async_info *);
   typedef int32_t(data_retrieve_ty)(int32_t, void *, void *, int64_t);
   typedef int32_t(data_retrieve_async_ty)(int32_t, void *, void *, int64_t,
                                           __tgt_async_info *);
   typedef int32_t(data_exchange_ty)(int32_t, void *, int32_t, void *, int64_t);
   typedef int32_t(data_exchange_async_ty)(int32_t, void *, int32_t, void *,
#endif
 
   // Functions implemented in the RTL.
   is_valid_binary_ty *is_valid_binary = nullptr;
   is_data_exchangable_ty *is_data_exchangable = nullptr;
   number_of_devices_ty *number_of_devices = nullptr;
   init_device_ty *init_device = nullptr;
   load_binary_ty *load_binary = nullptr;
   data_alloc_ty *data_alloc = nullptr;
+  data_alloc_explicit_ty *data_alloc_explicit = nullptr;
   data_submit_ty *data_submit = nullptr;
   data_submit_async_ty *data_submit_async = nullptr;
   data_retrieve_ty *data_retrieve = nullptr;
   data_retrieve_async_ty *data_retrieve_async = nullptr;
   data_exchange_ty *data_exchange = nullptr;
   data_exchange_async_ty *data_exchange_async = nullptr;
   data_delete_ty *data_delete = nullptr;
   run_region_ty *run_region = nullptr;

openmp/libomptarget/src/rtl.cpp

#endif
     *((void **)&R.is_data_exchangable) =
         dlsym(dynlib_handle, "__tgt_rtl_is_data_exchangable");
     *((void **)&R.register_lib) =
         dlsym(dynlib_handle, "__tgt_rtl_register_lib");
     *((void **)&R.unregister_lib) =
         dlsym(dynlib_handle, "__tgt_rtl_unregister_lib");
     *((void **)&R.supports_empty_images) =
         dlsym(dynlib_handle, "__tgt_rtl_supports_empty_images");
+    *((void **)&R.data_alloc_explicit) =
+        dlsym(dynlib_handle, "__tgt_rtl_data_alloc_explicit");
   }
 
   DP("RTLs loaded!\n");
 
   return;
 }
 
 ////////////////////////////////////////////////////////////////////////////////