[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

275

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

 // HostPtr is a address of the host data the allocated target data
 // will be associated with (HostPtr may be NULL if it is not known at
 // allocation time, like for example it would be for target data that
 // 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);
+// case an error occurred on the target device. Kind dictates what allocator
+// to use (e.g. shared, host, device).
+void *__tgt_rtl_data_alloc(int32_t ID, int64_t Size, void *HostPtr,
+                           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,

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

 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
 #include <dlfcn.h>
 #include <elf.h>
 #include <ffi.h>
 #include <fstream>
 #include <iostream>
 #include <libelf.h>
 #include <list>
 #include <memory>
 #include <mutex>
 #include <shared_mutex>
 #include <thread>
 #include <unordered_map>
 #include <vector>
 
for (__tgt_offload_entry *e = HostBegin; e != HostEnd; ++e) {
     entry.addr = (void *)&KernelsList.back();
     DeviceInfo.addOffloadEntry(device_id, entry);
     DP("Entry point %ld maps to %s\n", e - HostBegin, e->name);
   }
 
   return DeviceInfo.getOffloadEntriesTable(device_id);
 }
 
-void *__tgt_rtl_data_alloc(int device_id, int64_t size, void *) {
+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");
+
+  if (kind != TARGET_ALLOC_DEFAULT) {
+    REPORT("Invalid target data allocation kind or requested allocator not "
+           "implemented yet\n");
+    return NULL;
+  }
+
   atmi_status_t err = atmi_malloc(&ptr, size, get_gpu_mem_place(device_id));
   DP("Tgt alloc data %ld bytes, (tgt:%016llx).\n", size,
      (long long unsigned)(Elf64_Addr)ptr);
   ptr = (err == ATMI_STATUS_SUCCESS) ? ptr : NULL;
   return ptr;
 }
 
 int32_t __tgt_rtl_data_submit(int device_id, void *tgt_ptr, void *hst_ptr,

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

 //===----RTLs/cuda/src/rtl.cpp - Target RTLs Implementation ------- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // RTL for CUDA machine
 //
 //===----------------------------------------------------------------------===//
 
 #include <cassert>
 #include <cstddef>
 #include <cuda.h>
 #include <list>
 #include <memory>
 #include <mutex>
 #include <string>
 #include <vector>
 
 #include "Debug.h"
 #include "omptargetplugin.h"

 
 __tgt_target_table *__tgt_rtl_load_binary(int32_t device_id,
                                           __tgt_device_image *image) {
   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
 
   return DeviceRTL.loadBinary(device_id, image);
 }
 
-void *__tgt_rtl_data_alloc(int32_t device_id, int64_t size, void *) {
+void *__tgt_rtl_data_alloc(int32_t device_id, int64_t size, void *,
+                           int32_t kind) {
   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
 
+  if (kind != TARGET_ALLOC_DEFAULT) {
+    REPORT("Invalid target data allocation kind or requested allocator not "
+           "implemented yet\n");
+    return NULL;
+  }
+
   return DeviceRTL.dataAlloc(device_id, size);
 }
 
 int32_t __tgt_rtl_data_submit(int32_t device_id, void *tgt_ptr, void *hst_ptr,
                               int64_t size) {
   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
 
   __tgt_async_info async_info;

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"
 
DP("Entries table range is (" DPxMOD ")->(" DPxMOD ")\n",
      DPxPTR(entries_begin), DPxPTR(entries_end));
   DeviceInfo.createOffloadTable(device_id, entries_begin, entries_end);
 
   elf_end(e);
 
   return DeviceInfo.getOffloadEntriesTable(device_id);
 }
 
-void *__tgt_rtl_data_alloc(int32_t device_id, int64_t size, void *hst_ptr) {
-  void *ptr = malloc(size);
+// Sample implementation of explicit memory allocator. For this plugin all kinds
+// are equivalent to each other.
+void *__tgt_rtl_data_alloc(int32_t device_id, int64_t size, void *hst_ptr,
+                           int32_t kind) {
+  void *ptr = NULL;
+
+  switch (kind) {
+  case TARGET_ALLOC_DEVICE:
+  case TARGET_ALLOC_HOST:
+  case TARGET_ALLOC_SHARED:
+  case TARGET_ALLOC_DEFAULT:
+    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) {
   memcpy(hst_ptr, tgt_ptr, size);
   return OFFLOAD_SUCCESS;
 }

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

int32_t __tgt_rtl_synchronize(int32_t DeviceId,
                               __tgt_async_info *AsyncInfoPtr) {
   return Manager->synchronize(DeviceId, AsyncInfoPtr);
 }
 
 int32_t __tgt_rtl_is_data_exchangable(int32_t SrcDevId, int32_t DstDevId) {
   return Manager->isDataExchangeable(SrcDevId, DstDevId);
 }
 
-void *__tgt_rtl_data_alloc(int32_t DeviceId, int64_t Size, void *HstPtr) {
+void *__tgt_rtl_data_alloc(int32_t DeviceId, int64_t Size, void *HstPtr,
+                           int32_t kind) {
+  if (kind != TARGET_ALLOC_DEFAULT) {
+    REPORT("Invalid target data allocation kind or requested allocator not "
+           "implemented yet\n");
+    return NULL;
+  }
+
   return Manager->dataAlloc(DeviceId, Size, HstPtr);
 }
 
 int32_t __tgt_rtl_data_submit(int32_t DeviceId, void *TgtPtr, void *HstPtr,
                               int64_t Size) {
   return Manager->dataSubmitAsync(DeviceId, TgtPtr, HstPtr, Size, nullptr);
 }
 

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

 #include <algorithm>
 #include <cassert>
 #include <cerrno>
 #include <cstring>
 #include <list>
 #include <stdlib.h>
 #include <string>
 #include <sys/stat.h>
 #include <ve_offload.h>
 #include <vector>
 #include <veosinfo/veosinfo.h>
 
 #include "Debug.h"
 #include "omptargetplugin.h"
 
 #ifndef TARGET_NAME
 #define TARGET_NAME VE

 // will be associated with (HostPtr may be NULL if it is not known at
 // allocation time, like for example it would be for target data that
 // 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) {
+void *__tgt_rtl_data_alloc(int32_t ID, int64_t Size, void *HostPtr,
+                           int32_t kind) {
   int ret;
   uint64_t addr;
 
+  if (kind != TARGET_ALLOC_DEFAULT) {
+    REPORT("Invalid target data allocation kind or requested allocator not "
+           "implemented yet\n");
+    return NULL;
+  }
+
   if (DeviceInfo.ProcHandles[ID] == NULL) {
     struct veo_proc_handle *proc_handle;
     proc_handle = veo_proc_create(DeviceInfo.NodeIds[ID]);
     if (!proc_handle) {
       DP("veo_proc_create() failed for device %d\n", ID);
       return NULL;
     }
     DeviceInfo.ProcHandles[ID] = proc_handle;

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) {
+static void *targetAllocExplicit(size_t size, int device_num, int kind,
+                                 const char *name) {
   TIMESCOPE();
-  DP("Call to omp_target_alloc for device %d requesting %zu bytes\n",
-     device_num, size);
+  DP("Call to %s for device %d requesting %zu bytes\n", name, device_num, size);
 
   if (size <= 0) {
-    DP("Call to omp_target_alloc with non-positive length\n");
+    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("omp_target_alloc returns host ptr " DPxMOD "\n", DPxPTR(rc));
+    DP("%s returns host ptr " DPxMOD "\n", name, DPxPTR(rc));
     return rc;
   }
 
   if (!device_is_ready(device_num)) {
-    DP("omp_target_alloc returns NULL ptr\n");
+    DP("%s returns NULL ptr\n", name);
     return NULL;
   }
 
-  rc = PM->Devices[device_num].allocData(size);
-  DP("omp_target_alloc returns device ptr " DPxMOD "\n", DPxPTR(rc));
+  DeviceTy &Device = PM->Devices[device_num];
+  rc = Device.allocData(size, nullptr, kind);
+  DP("%s returns device ptr " DPxMOD "\n", name, DPxPTR(rc));
   return rc;
 }
 
+EXTERN void *omp_target_alloc(size_t size, int device_num) {
+  return targetAllocExplicit(size, device_num, TARGET_ALLOC_DEFAULT, __func__);
+}
+
+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__);
+}
+
 EXTERN void omp_target_free(void *device_ptr, int device_num) {
   TIMESCOPE();
   DP("Call to omp_target_free for device %d and address " DPxMOD "\n",
      device_num, DPxPTR(device_ptr));
 
   if (!device_ptr) {
     DP("Call to omp_target_free with NULL ptr\n");
     return;

openmp/libomptarget/src/device.h

struct DeviceTy {
   int associatePtr(void *HstPtrBegin, void *TgtPtrBegin, int64_t Size);
   int disassociatePtr(void *HstPtrBegin);
 
   // calls to RTL
   int32_t initOnce();
   __tgt_target_table *load_binary(void *Img);
 
   // device memory allocation/deallocation routines
-  /// Allocates \p Size bytes on the device and returns the address/nullptr when
+  /// Allocates \p Size bytes on the device, host or shared memory space
+  /// (depending on \p Kind) 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);
+  /// implementations ignore \p HstPtr. \p Kind dictates what allocator should
+  /// be used (host, shared, device).
+  void *allocData(int64_t Size, void *HstPtr = nullptr,
+                  int32_t Kind = TARGET_ALLOC_DEFAULT);
   /// 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

 // Load binary to device.
 __tgt_target_table *DeviceTy::load_binary(void *Img) {
   RTL->Mtx.lock();
   __tgt_target_table *rc = RTL->load_binary(RTLDeviceID, Img);
   RTL->Mtx.unlock();
   return rc;
 }
 
-void *DeviceTy::allocData(int64_t Size, void *HstPtr) {
-  return RTL->data_alloc(RTLDeviceID, Size, HstPtr);
+void *DeviceTy::allocData(int64_t Size, void *HstPtr, int32_t Kind) {
+  return RTL->data_alloc(RTLDeviceID, Size, HstPtr, Kind);
 }
 
 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,

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 __tgt_bin_desc;
 
 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_ty)(int32_t, int64_t, void *, 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 *,