[UR][CUDA] Cleanup USM allocations code (#18052)

The `Impl` functions were only being called once and doing something
very similar to the regular code so we can fold them in the UR call.

Also we can use the `isPowerOf2` utility function for alignment
validation.

Author: npmiller

unified-runtime/source/adapters/cuda/usm.cpp

@@ -28,18 +28,17 @@ UR_APIEXPORT ur_result_t UR_APICALL
 urUSMHostAlloc(ur_context_handle_t hContext, const ur_usm_desc_t *pUSMDesc,
                ur_usm_pool_handle_t hPool, size_t size, void **ppMem) {
   auto alignment = pUSMDesc ? pUSMDesc->align : 0u;
-  UR_ASSERT(!pUSMDesc ||
-                (alignment == 0 || ((alignment & (alignment - 1)) == 0)),
-            UR_RESULT_ERROR_INVALID_VALUE);
 
-  if (!hPool) {
-    return USMHostAllocImpl(ppMem, hContext, /* flags */ 0, size, alignment);
+  auto pool = hPool ? hPool->HostMemPool.get() : hContext->MemoryPoolHost;
+  if (alignment) {
+    UR_ASSERT(isPowerOf2(alignment), UR_RESULT_ERROR_INVALID_VALUE);
+    *ppMem = umfPoolAlignedMalloc(pool, size, alignment);
+  } else {
+    *ppMem = umfPoolMalloc(pool, size);
   }
 
-  auto UMFPool = hPool->HostMemPool.get();
-  *ppMem = umfPoolAlignedMalloc(UMFPool, size, alignment);
   if (*ppMem == nullptr) {
-    auto umfErr = umfPoolGetLastAllocationError(UMFPool);
+    auto umfErr = umfPoolGetLastAllocationError(pool);
     return umf::umf2urResult(umfErr);
   }
   return UR_RESULT_SUCCESS;
@@ -48,23 +47,22 @@ urUSMHostAlloc(ur_context_handle_t hContext, const ur_usm_desc_t *pUSMDesc,
 /// USM: Implements USM device allocations using a normal CUDA device pointer
 ///
 UR_APIEXPORT ur_result_t UR_APICALL
-urUSMDeviceAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
+urUSMDeviceAlloc(ur_context_handle_t, ur_device_handle_t hDevice,
                  const ur_usm_desc_t *pUSMDesc, ur_usm_pool_handle_t hPool,
                  size_t size, void **ppMem) {
   auto alignment = pUSMDesc ? pUSMDesc->align : 0u;
-  UR_ASSERT(!pUSMDesc ||
-                (alignment == 0 || ((alignment & (alignment - 1)) == 0)),
-            UR_RESULT_ERROR_INVALID_VALUE);
 
-  if (!hPool) {
-    return USMDeviceAllocImpl(ppMem, hContext, hDevice, /* flags */ 0, size,
-                              alignment);
+  ScopedContext SC(hDevice);
+  auto pool = hPool ? hPool->DeviceMemPool.get() : hDevice->MemoryPoolDevice;
+  if (alignment) {
+    UR_ASSERT(isPowerOf2(alignment), UR_RESULT_ERROR_INVALID_VALUE);
+    *ppMem = umfPoolAlignedMalloc(pool, size, alignment);
+  } else {
+    *ppMem = umfPoolMalloc(pool, size);
   }
 
-  auto UMFPool = hPool->DeviceMemPool.get();
-  *ppMem = umfPoolAlignedMalloc(UMFPool, size, alignment);
   if (*ppMem == nullptr) {
-    auto umfErr = umfPoolGetLastAllocationError(UMFPool);
+    auto umfErr = umfPoolGetLastAllocationError(pool);
     return umf::umf2urResult(umfErr);
   }
   return UR_RESULT_SUCCESS;
@@ -73,23 +71,22 @@ urUSMDeviceAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
 /// USM: Implements USM Shared allocations using CUDA Managed Memory
 ///
 UR_APIEXPORT ur_result_t UR_APICALL
-urUSMSharedAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
+urUSMSharedAlloc(ur_context_handle_t, ur_device_handle_t hDevice,
                  const ur_usm_desc_t *pUSMDesc, ur_usm_pool_handle_t hPool,
                  size_t size, void **ppMem) {
   auto alignment = pUSMDesc ? pUSMDesc->align : 0u;
-  UR_ASSERT(!pUSMDesc ||
-                (alignment == 0 || ((alignment & (alignment - 1)) == 0)),
-            UR_RESULT_ERROR_INVALID_VALUE);
 
-  if (!hPool) {
-    return USMSharedAllocImpl(ppMem, hContext, hDevice, /*host flags*/ 0,
-                              /*device flags*/ 0, size, alignment);
+  ScopedContext SC(hDevice);
+  auto pool = hPool ? hPool->SharedMemPool.get() : hDevice->MemoryPoolShared;
+  if (alignment) {
+    UR_ASSERT(isPowerOf2(alignment), UR_RESULT_ERROR_INVALID_VALUE);
+    *ppMem = umfPoolAlignedMalloc(pool, size, alignment);
+  } else {
+    *ppMem = umfPoolMalloc(pool, size);
   }
 
-  auto UMFPool = hPool->SharedMemPool.get();
-  *ppMem = umfPoolAlignedMalloc(UMFPool, size, alignment);
   if (*ppMem == nullptr) {
-    auto umfErr = umfPoolGetLastAllocationError(UMFPool);
+    auto umfErr = umfPoolGetLastAllocationError(pool);
     return umf::umf2urResult(umfErr);
   }
   return UR_RESULT_SUCCESS;
@@ -103,56 +100,6 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMFree(ur_context_handle_t hContext,
   return umf::umf2urResult(umfFree(pMem));
 }
 
-ur_result_t USMDeviceAllocImpl(void **ResultPtr, ur_context_handle_t,
-                               ur_device_handle_t Device,
-                               ur_usm_device_mem_flags_t, size_t Size,
-                               [[maybe_unused]] uint32_t Alignment) {
-  try {
-    ScopedContext Active(Device);
-    *ResultPtr = umfPoolMalloc(Device->MemoryPoolDevice, Size);
-    UMF_CHECK_PTR(*ResultPtr);
-  } catch (ur_result_t Err) {
-    return Err;
-  }
-
-  assert((Alignment == 0 ||
-          reinterpret_cast<std::uintptr_t>(*ResultPtr) % Alignment == 0));
-  return UR_RESULT_SUCCESS;
-}
-
-ur_result_t USMSharedAllocImpl(void **ResultPtr, ur_context_handle_t,
-                               ur_device_handle_t Device,
-                               ur_usm_host_mem_flags_t,
-                               ur_usm_device_mem_flags_t, size_t Size,
-                               [[maybe_unused]] uint32_t Alignment) {
-  try {
-    ScopedContext Active(Device);
-    *ResultPtr = umfPoolMalloc(Device->MemoryPoolShared, Size);
-    UMF_CHECK_PTR(*ResultPtr);
-  } catch (ur_result_t Err) {
-    return Err;
-  }
-
-  assert((Alignment == 0 ||
-          reinterpret_cast<std::uintptr_t>(*ResultPtr) % Alignment == 0));
-  return UR_RESULT_SUCCESS;
-}
-
-ur_result_t USMHostAllocImpl(void **ResultPtr, ur_context_handle_t hContext,
-                             ur_usm_host_mem_flags_t, size_t Size,
-                             [[maybe_unused]] uint32_t Alignment) {
-  try {
-    *ResultPtr = umfPoolMalloc(hContext->MemoryPoolHost, Size);
-    UMF_CHECK_PTR(*ResultPtr);
-  } catch (ur_result_t Err) {
-    return Err;
-  }
-
-  assert((Alignment == 0 ||
-          reinterpret_cast<std::uintptr_t>(*ResultPtr) % Alignment == 0));
-  return UR_RESULT_SUCCESS;
-}
-
 UR_APIEXPORT ur_result_t UR_APICALL
 urUSMGetMemAllocInfo(ur_context_handle_t hContext, const void *pMem,
                      ur_usm_alloc_info_t propName, size_t propValueSize,