Update CUDA allocations

lbushi25 · May 28, 2024 · 1530cf9 · 1530cf9
1 parent e3f526c
commit 1530cf9
Showing 1 changed file with 27 additions and 12 deletions.
diff --git a/source/adapters/cuda/usm.cpp b/source/adapters/cuda/usm.cpp
@@ -130,7 +130,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMFree(ur_context_handle_t hContext,
   return USMFreeImpl(hContext, pMem);
 }
 
-ur_result_t USMDeviceAllocImpl(void **ResultPtr, ur_context_handle_t,
+ur_result_t USMDeviceAllocImpl(void **ResultPtr, ur_context_handle_t hContext,
                                ur_device_handle_t Device,
                                ur_usm_device_mem_flags_t, size_t Size,
                                uint32_t Alignment) {
@@ -141,16 +141,21 @@ ur_result_t USMDeviceAllocImpl(void **ResultPtr, ur_context_handle_t,
     return Err;
   }
 
+  const bool validAlignment =
+      (Alignment == 0 ||
+       reinterpret_cast<std::uintptr_t>(*ResultPtr) % Alignment == 0);
 #ifdef NDEBUG
-  std::ignore = Alignment;
+  if (!validAlignment) {
+    urUSMFree(hContext, *ResultPtr);
+    return UR_RESULT_ERROR_UNSUPPORTED_ALIGNMENT;
+  }
 #else
-  assert((Alignment == 0 ||
-          reinterpret_cast<std::uintptr_t>(*ResultPtr) % Alignment == 0));
+  assert(validAlignment);
 #endif
   return UR_RESULT_SUCCESS;
 }
 
-ur_result_t USMSharedAllocImpl(void **ResultPtr, ur_context_handle_t,
+ur_result_t USMSharedAllocImpl(void **ResultPtr, ur_context_handle_t hContext,
                                ur_device_handle_t Device,
                                ur_usm_host_mem_flags_t,
                                ur_usm_device_mem_flags_t, size_t Size,
@@ -163,16 +168,21 @@ ur_result_t USMSharedAllocImpl(void **ResultPtr, ur_context_handle_t,
     return Err;
   }
 
+  const bool validAlignment =
+      (Alignment == 0 ||
+       reinterpret_cast<std::uintptr_t>(*ResultPtr) % Alignment == 0);
 #ifdef NDEBUG
-  std::ignore = Alignment;
+  if (!validAlignment) {
+    urUSMFree(hContext, *ResultPtr);
+    return UR_RESULT_ERROR_UNSUPPORTED_ALIGNMENT;
+  }
 #else
-  assert((Alignment == 0 ||
-          reinterpret_cast<std::uintptr_t>(*ResultPtr) % Alignment == 0));
+  assert(validAlignment);
 #endif
   return UR_RESULT_SUCCESS;
 }
 
-ur_result_t USMHostAllocImpl(void **ResultPtr, ur_context_handle_t,
+ur_result_t USMHostAllocImpl(void **ResultPtr, ur_context_handle_t hContext,
                              ur_usm_host_mem_flags_t, size_t Size,
                              uint32_t Alignment) {
   try {
@@ -181,11 +191,16 @@ ur_result_t USMHostAllocImpl(void **ResultPtr, ur_context_handle_t,
     return Err;
   }
 
+  const bool validAlignment =
+      (Alignment == 0 ||
+       reinterpret_cast<std::uintptr_t>(*ResultPtr) % Alignment == 0);
 #ifdef NDEBUG
-  std::ignore = Alignment;
+  if (!validAlignment) {
+    urUSMFree(hContext, *ResultPtr);
+    return UR_RESULT_ERROR_UNSUPPORTED_ALIGNMENT;
+  }
 #else
-  assert((Alignment == 0 ||
-          reinterpret_cast<std::uintptr_t>(*ResultPtr) % Alignment == 0));
+  assert(validAlignment);
 #endif
   return UR_RESULT_SUCCESS;
 }