[AMDGPU][Libomptarget][NFC] Split host and device malloc

[pdhaliwal]

Summary

This patch splits the code path for host and device malloc.

[JonChesterfield]

LGTM, thanks. Cleans up the call sites without changing the calls into HSA. Thanks for dropping the extern C shim too.

openmp/libomptarget/plugins/amdgpu/impl/atmi.cpp

if (got != success) {
     return HSA_STATUS_ERROR;
   }
 
   return err;
 }
 
 struct atmiFreePtrDeletor {
   void operator()(void *p) {
-    atmi_free(p); // ignore failure to free
+    core::Runtime::Memfree(p); // ignore failure to free
   }
 };
 
 hsa_status_t atmi_memcpy_h2d(hsa_signal_t signal, void *deviceDest,
                              const void *hostSrc, size_t size,
                              hsa_agent_t agent) {
   hsa_status_t rc = hsa_memory_copy(deviceDest, hostSrc, size);
 
   // hsa_memory_copy sometimes fails in situations where
   // allocate + copy succeeds. Looks like it might be related to
   // locking part of a read only segment. Fall back for now.
   if (rc == HSA_STATUS_SUCCESS) {
     return HSA_STATUS_SUCCESS;
   }
 
   void *tempHostPtr;
-  hsa_status_t ret =
-      atmi_malloc(&tempHostPtr, size, 0 /* DeviceId */, ATMI_DEVTYPE_CPU);
+  hsa_status_t ret = core::Runtime::HostMalloc(&tempHostPtr, size);
   if (ret != HSA_STATUS_SUCCESS) {
-    DEBUG_PRINT("atmi_malloc: Unable to alloc %d bytes for temp scratch\n",
+    DEBUG_PRINT("HostMalloc: Unable to alloc %d bytes for temp scratch\n",
                 size);
     return ret;
   }
   std::unique_ptr<void, atmiFreePtrDeletor> del(tempHostPtr);
   memcpy(tempHostPtr, hostSrc, size);
 
   if (invoke_hsa_copy(signal, deviceDest, tempHostPtr, size, agent) !=
       HSA_STATUS_SUCCESS) {
hsa_status_t atmi_memcpy_d2h(hsa_signal_t signal, void *dest,
   // allocate + copy succeeds. Looks like it might be related to
   // locking part of a read only segment. Fall back for now.
   if (rc == HSA_STATUS_SUCCESS) {
     return HSA_STATUS_SUCCESS;
   }
 
   void *tempHostPtr;
 
-  hsa_status_t ret =
-      atmi_malloc(&tempHostPtr, size, 0 /* DeviceId */, ATMI_DEVTYPE_CPU);
+  hsa_status_t ret = core::Runtime::HostMalloc(&tempHostPtr, size);
   if (ret != HSA_STATUS_SUCCESS) {
-    DEBUG_PRINT("atmi_malloc: Unable to alloc %d bytes for temp scratch\n",
+    DEBUG_PRINT("HostMalloc: Unable to alloc %d bytes for temp scratch\n",
                 size);
     return ret;
   }
   std::unique_ptr<void, atmiFreePtrDeletor> del(tempHostPtr);
 
   if (invoke_hsa_copy(signal, tempHostPtr, deviceSrc, size, agent) !=
       HSA_STATUS_SUCCESS) {
     return HSA_STATUS_ERROR;
   }
 
   memcpy(dest, tempHostPtr, size);
   return HSA_STATUS_SUCCESS;
 }
-
-hsa_status_t atmi_free(void *ptr) { return core::Runtime::Memfree(ptr); }
-
-hsa_status_t atmi_malloc(void **ptr, size_t size, int DeviceId,
-                         atmi_devtype_t DeviceType) {
-  return core::Runtime::Malloc(ptr, size, DeviceId, DeviceType);
-}

openmp/libomptarget/plugins/amdgpu/impl/atmi_runtime.h

 /*===--------------------------------------------------------------------------
  *              ATMI (Asynchronous Task and Memory Interface)
  *
  * This file is distributed under the MIT License. See LICENSE.txt for details.
  *===------------------------------------------------------------------------*/
 #ifndef INCLUDE_ATMI_RUNTIME_H_
 #define INCLUDE_ATMI_RUNTIME_H_
 
 #include "atmi.h"
 #include "hsa.h"
 #include <inttypes.h>
 #include <stdlib.h>
 #ifndef __cplusplus
 #include <stdbool.h>
 #endif
 
 #ifdef __cplusplus
 extern "C" {

  * this function can be called to retrieve the pointer to this global structure.
  *
  * @return Returns a pointer to a global structure of tyoe @p atmi_machine_t.
  * Returns NULL if ATMI is not initialized.
  */
 atmi_machine_t *atmi_machine_get_info();
 /** @} */
 
-/** \defgroup memory_functions ATMI Data Management
- * @{
- */
-/**
- * @brief Allocate memory from the specified memory place.
- *
- * @detail This function allocates memory from the specified memory place. If
- * the memory
- * place belongs primarily to the CPU, then the memory will be accessible by
- * other GPUs and CPUs in the system. If the memory place belongs primarily to a
- * GPU,
- * then it cannot be accessed by other devices in the system.
- *
- * @param[in] ptr The pointer to the memory that will be allocated.
- *
- * @param[in] size The size of the allocation in bytes.
- *
- * @param[in] place The memory place in the system to perform the allocation.
- *
- * @retval ::HSA_STATUS_SUCCESS The function has executed successfully.
- *
- * @retval ::HSA_STATUS_ERROR The function encountered errors.
- *
- */
-hsa_status_t atmi_malloc(void **ptr, size_t size, int DeviceId,
-                         atmi_devtype_t DeviceType);
-/**
- * @brief Frees memory that was previously allocated.
- *
- * @detail This function frees memory that was previously allocated by calling
- * @p atmi_malloc. It throws an error otherwise. It is illegal to access a
- * pointer after a call to this function.
- *
- * @param[in] ptr The pointer to the memory that has to be freed.
- *
- * @retval ::HSA_STATUS_SUCCESS The function has executed successfully.
- *
- * @retval ::HSA_STATUS_ERROR The function encountered errors.
- *
- */
-hsa_status_t atmi_free(void *ptr);
-
 hsa_status_t atmi_memcpy_h2d(hsa_signal_t signal, void *deviceDest,
                              const void *hostSrc, size_t size,
                              hsa_agent_t agent);
 
 hsa_status_t atmi_memcpy_d2h(hsa_signal_t sig, void *hostDest,
                              const void *deviceSrc, size_t size,
                              hsa_agent_t agent);
 
 /** @} */
 
 #ifdef __cplusplus
 }
 #endif
 
 #endif // INCLUDE_ATMI_RUNTIME_H_

openmp/libomptarget/plugins/amdgpu/impl/data.cpp

 hsa_status_t register_allocation(void *ptr, size_t size,
                                  atmi_devtype_t DeviceType) {
   if (DeviceType == ATMI_DEVTYPE_CPU)
     return allow_access_to_all_gpu_agents(ptr);
   else
     return HSA_STATUS_SUCCESS;
 }
 
+hsa_status_t Runtime::DeviceMalloc(void **ptr, size_t size, int DeviceId) {
+  return Runtime::Malloc(ptr, size, DeviceId, ATMI_DEVTYPE_GPU);
+}
+
+hsa_status_t Runtime::HostMalloc(void **ptr, size_t size) {
+  return Runtime::Malloc(ptr, size, 0, ATMI_DEVTYPE_CPU);
+}
+
 hsa_status_t Runtime::Malloc(void **ptr, size_t size, int DeviceId,
                              atmi_devtype_t DeviceType) {
   hsa_amd_memory_pool_t pool =
       get_memory_pool_by_mem_place(DeviceId, DeviceType);
   hsa_status_t err = hsa_amd_memory_pool_allocate(pool, size, 0, ptr);
   DEBUG_PRINT("Malloced [%s %d] %p\n",
               DeviceType == ATMI_DEVTYPE_CPU ? "CPU" : "GPU", DeviceId, *ptr);
 

openmp/libomptarget/plugins/amdgpu/impl/rt.h

static hsa_status_t RegisterModuleFromMemory(
       void *, size_t, atmi_place_t,
       hsa_status_t (*on_deserialized_data)(void *data, size_t size,
                                            void *cb_state),
       void *cb_state, std::vector<hsa_executable_t> &HSAExecutables);
 
   // data
   static hsa_status_t Memcpy(hsa_signal_t, void *, const void *, size_t);
   static hsa_status_t Memfree(void *);
-  static hsa_status_t Malloc(void **ptr, size_t size, int DeviceId,
-                             atmi_devtype_t DeviceType);
+  static hsa_status_t HostMalloc(void **ptr, size_t size);
+  static hsa_status_t DeviceMalloc(void **ptr, size_t size, int DeviceId);
 
   int getMaxQueueSize() const { return env_.getMaxQueueSize(); }
   int getDebugMode() const { return env_.getDebugMode(); }
 
+private:
+  static hsa_status_t Malloc(void **ptr, size_t size, int DeviceId,
+                             atmi_devtype_t DeviceType);
+
 protected:
   Runtime() = default;
   ~Runtime() = default;
   Runtime(const Runtime &) = delete;
   Runtime &operator=(const Runtime &) = delete;
 
 protected:
   // variable to track environment variables
   Environment env_;
 };
 
 } // namespace core
 
 #endif // SRC_RUNTIME_INCLUDE_RT_H_

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>
 
public:
 
   std::vector<hsa_executable_t> HSAExecutables;
 
   std::vector<std::map<std::string, atl_kernel_info_t>> KernelInfoTable;
   std::vector<std::map<std::string, atl_symbol_info_t>> SymbolInfoTable;
 
   struct atmiFreePtrDeletor {
     void operator()(void *p) {
-      atmi_free(p); // ignore failure to free
+      core::Runtime::Memfree(p); // ignore failure to free
     }
   };
 
   // device_State shared across loaded binaries, error if inconsistent size
   std::vector<std::pair<std::unique_ptr<void, atmiFreePtrDeletor>, uint64_t>>
       deviceStateStore;
 
   static const unsigned HardTeamLimit =
hsa_status_t after_loading() {
     }
     return HSA_STATUS_SUCCESS;
   }
 };
 
 static hsa_status_t atmi_calloc(void **ret_ptr, size_t size, int DeviceId) {
   uint64_t rounded = 4 * ((size + 3) / 4);
   void *ptr;
-  hsa_status_t err = atmi_malloc(&ptr, rounded, DeviceId, ATMI_DEVTYPE_GPU);
+  hsa_status_t err = core::Runtime::DeviceMalloc(&ptr, rounded, DeviceId);
   if (err != HSA_STATUS_SUCCESS) {
     return err;
   }
 
   hsa_status_t rc = hsa_amd_memory_fill(ptr, 0, rounded / 4);
   if (rc != HSA_STATUS_SUCCESS) {
     fprintf(stderr, "zero fill device_state failed with %u\n", rc);
-    atmi_free(ptr);
+    core::Runtime::Memfree(ptr);
     return HSA_STATUS_ERROR;
   }
 
   *ret_ptr = ptr;
   return HSA_STATUS_SUCCESS;
 }
 
 static bool image_contains_symbol(void *data, size_t size, const char *sym) {
void *__tgt_rtl_data_alloc(int device_id, int64_t size, void *, int32_t kind) {
   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;
   }
 
-  hsa_status_t err = atmi_malloc(&ptr, size, device_id, ATMI_DEVTYPE_GPU);
+  hsa_status_t err = core::Runtime::DeviceMalloc(&ptr, size, device_id);
   DP("Tgt alloc data %ld bytes, (tgt:%016llx).\n", size,
      (long long unsigned)(Elf64_Addr)ptr);
   ptr = (err == HSA_STATUS_SUCCESS) ? ptr : NULL;
   return ptr;
 }
 
 int32_t __tgt_rtl_data_submit(int device_id, void *tgt_ptr, void *hst_ptr,
                               int64_t size) {
int32_t __tgt_rtl_data_retrieve_async(int device_id, void *hst_ptr,
   initAsyncInfo(AsyncInfo);
   return dataRetrieve(device_id, hst_ptr, tgt_ptr, size, AsyncInfo);
 }
 
 int32_t __tgt_rtl_data_delete(int device_id, void *tgt_ptr) {
   assert(device_id < DeviceInfo.NumberOfDevices && "Device ID too large");
   hsa_status_t err;
   DP("Tgt free data (tgt:%016llx).\n", (long long unsigned)(Elf64_Addr)tgt_ptr);
-  err = atmi_free(tgt_ptr);
+  err = core::Runtime::Memfree(tgt_ptr);
   if (err != HSA_STATUS_SUCCESS) {
     DP("Error when freeing CUDA memory\n");
     return OFFLOAD_FAIL;
   }
   return OFFLOAD_SUCCESS;
 }
 
 // Determine launch values for threadsPerGroup and num_groups.