Implement urUSMGetMemAllocInfo and aligned alloc

Author: PietroGhg

source/adapters/native_cpu/context.hpp

@@ -10,13 +10,49 @@
 
 #pragma once
 
+#include <unordered_map>
 #include <ur_api.h>
 
 #include "common.hpp"
 #include "device.hpp"
+#include "ur/ur.hpp"
+
+namespace native_cpu {
+struct usm_alloc_info {
+  const ur_usm_type_t type;
+  const void *base_ptr;
+  const size_t size;
+  const ur_device_handle_t device;
+  const ur_usm_pool_handle_t pool;
+  usm_alloc_info(ur_usm_type_t type, const void *base_ptr, size_t size,
+                 ur_device_handle_t device, ur_usm_pool_handle_t pool)
+      : type(type), base_ptr(base_ptr), size(size), device(device), pool(pool) {
+  }
+};
+} // namespace native_cpu
 
 struct ur_context_handle_t_ : RefCounted {
   ur_context_handle_t_(ur_device_handle_t_ *phDevices) : _device{phDevices} {}
 
   ur_device_handle_t _device;
+
+  void add_alloc_info_entry(const void *ptr, ur_usm_type_t type, size_t size,
+                            ur_usm_pool_handle_t pool) {
+    native_cpu::usm_alloc_info info(type, ptr, size, this->_device, pool);
+    alloc_info.insert(std::make_pair(ptr, info));
+  }
+
+  native_cpu::usm_alloc_info get_alloc_info_entry(const void *ptr) const {
+    auto it = alloc_info.find(ptr);
+    if (it == alloc_info.end()) {
+      return native_cpu::usm_alloc_info(UR_USM_TYPE_UNKNOWN, ptr, 0, nullptr,
+                                        nullptr);
+    }
+    return it->second;
+  }
+
+  void remove_alloc_info_entry(void *ptr) { alloc_info.erase(ptr); }
+
+private:
+  std::unordered_map<const void *, native_cpu::usm_alloc_info> alloc_info;
 };

source/adapters/native_cpu/usm.cpp

@@ -8,22 +8,42 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "ur/ur.hpp"
 #include "ur_api.h"
 
 #include "common.hpp"
+#include "context.hpp"
+#include <cstdlib>
+
+namespace native_cpu {
+
+static void *malloc_impl(uint32_t alignment, size_t size) {
+  void *ptr = nullptr;
+  if (alignment) {
+    ptr = std::aligned_alloc(alignment, size);
+  } else {
+    ptr = malloc(size);
+  }
+  return ptr;
+}
+
+} // namespace native_cpu
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urUSMHostAlloc(ur_context_handle_t hContext, const ur_usm_desc_t *pUSMDesc,
                ur_usm_pool_handle_t pool, size_t size, void **ppMem) {
-  std::ignore = hContext;
-  std::ignore = pUSMDesc;
   std::ignore = pool;
 
+  auto alignment = pUSMDesc ? pUSMDesc->align : 0u;
+  UR_ASSERT((alignment & (alignment - 1)) == 0, UR_RESULT_ERROR_INVALID_VALUE);
   UR_ASSERT(ppMem, UR_RESULT_ERROR_INVALID_NULL_POINTER);
   // TODO: Check Max size when UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE is implemented
   UR_ASSERT(size > 0, UR_RESULT_ERROR_INVALID_USM_SIZE);
 
-  *ppMem = malloc(size);
+  void *ptr = native_cpu::malloc_impl(alignment, size);
+  UR_ASSERT(ptr != nullptr, UR_RESULT_ERROR_OUT_OF_HOST_MEMORY);
+  hContext->add_alloc_info_entry(ptr, UR_USM_TYPE_HOST, size, nullptr);
+  *ppMem = ptr;
 
   return UR_RESULT_SUCCESS;
 }
@@ -32,16 +52,19 @@ UR_APIEXPORT ur_result_t UR_APICALL
 urUSMDeviceAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
                  const ur_usm_desc_t *pUSMDesc, ur_usm_pool_handle_t pool,
                  size_t size, void **ppMem) {
-  std::ignore = hContext;
   std::ignore = hDevice;
-  std::ignore = pUSMDesc;
   std::ignore = pool;
 
+  auto alignment = pUSMDesc ? pUSMDesc->align : 0u;
+  UR_ASSERT((alignment & (alignment - 1)) == 0, UR_RESULT_ERROR_INVALID_VALUE);
   UR_ASSERT(ppMem, UR_RESULT_ERROR_INVALID_NULL_POINTER);
   // TODO: Check Max size when UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE is implemented
   UR_ASSERT(size > 0, UR_RESULT_ERROR_INVALID_USM_SIZE);
 
-  *ppMem = malloc(size);
+  void *ptr = native_cpu::malloc_impl(alignment, size);
+  UR_ASSERT(ptr != nullptr, UR_RESULT_ERROR_OUT_OF_RESOURCES);
+  *ppMem = ptr;
+  hContext->add_alloc_info_entry(ptr, UR_USM_TYPE_DEVICE, size, nullptr);
 
   return UR_RESULT_SUCCESS;
 }
@@ -50,26 +73,29 @@ UR_APIEXPORT ur_result_t UR_APICALL
 urUSMSharedAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
                  const ur_usm_desc_t *pUSMDesc, ur_usm_pool_handle_t pool,
                  size_t size, void **ppMem) {
-  std::ignore = hContext;
   std::ignore = hDevice;
-  std::ignore = pUSMDesc;
   std::ignore = pool;
 
+  auto alignment = pUSMDesc ? pUSMDesc->align : 0u;
+  UR_ASSERT((alignment & (alignment - 1)) == 0, UR_RESULT_ERROR_INVALID_VALUE);
   UR_ASSERT(ppMem, UR_RESULT_ERROR_INVALID_NULL_POINTER);
   // TODO: Check Max size when UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE is implemented
   UR_ASSERT(size > 0, UR_RESULT_ERROR_INVALID_USM_SIZE);
 
-  *ppMem = malloc(size);
+  void *ptr = native_cpu::malloc_impl(alignment, size);
+  UR_ASSERT(ptr != nullptr, UR_RESULT_ERROR_OUT_OF_HOST_MEMORY);
+  *ppMem = ptr;
+  hContext->add_alloc_info_entry(ptr, UR_USM_TYPE_SHARED, size, nullptr);
 
   return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urUSMFree(ur_context_handle_t hContext,
                                               void *pMem) {
-  std::ignore = hContext;
 
   UR_ASSERT(pMem, UR_RESULT_ERROR_INVALID_NULL_POINTER);
 
+  hContext->remove_alloc_info_entry(pMem);
   free(pMem);
 
   return UR_RESULT_SUCCESS;
@@ -79,19 +105,25 @@ UR_APIEXPORT ur_result_t UR_APICALL
 urUSMGetMemAllocInfo(ur_context_handle_t hContext, const void *pMem,
                      ur_usm_alloc_info_t propName, size_t propSize,
                      void *pPropValue, size_t *pPropSizeRet) {
-  std::ignore = hContext;
-  std::ignore = pMem;
-  std::ignore = propName;
-  std::ignore = propSize;
-  std::ignore = pPropValue;
-  std::ignore = pPropSizeRet;
 
+  UR_ASSERT(pMem != nullptr, UR_RESULT_ERROR_INVALID_NULL_POINTER);
   UrReturnHelper ReturnValue(propSize, pPropValue, pPropSizeRet);
+  if (propName == UR_USM_ALLOC_INFO_BASE_PTR) {
+    // TODO: logic to compute base ptr given ptr
+    DIE_NO_IMPLEMENTATION;
+  }
 
+  native_cpu::usm_alloc_info alloc_info = hContext->get_alloc_info_entry(pMem);
   switch (propName) {
   case UR_USM_ALLOC_INFO_TYPE:
-    // Todo implement this in context
-    return ReturnValue(UR_USM_TYPE_DEVICE);
+    return ReturnValue(alloc_info.type);
+  case UR_USM_ALLOC_INFO_SIZE:
+    return ReturnValue(alloc_info.size);
+  case UR_USM_ALLOC_INFO_DEVICE:
+    return ReturnValue(alloc_info.device);
+  case UR_USM_ALLOC_INFO_POOL:
+    return ReturnValue(alloc_info.pool);
+    ;
   default:
     DIE_NO_IMPLEMENTATION;
   }