Add size validation for USM device/shared allocs

RossBrunton · RossBrunton · commit ffdbe093bbb5 · 2025-01-14T12:17:18.000Z
The spec states that an allocation size higher than
`DEVICE_INFO_MAX_MEM_ALLOC_SIZE` should report an error. This has been
added to the Cuda, Hip and Native CPU backends now.
diff --git a/source/adapters/cuda/usm.cpp b/source/adapters/cuda/usm.cpp
@@ -63,6 +63,10 @@ urUSMDeviceAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
                 (alignment == 0 || ((alignment & (alignment - 1)) == 0)),
             UR_RESULT_ERROR_INVALID_VALUE);
 
+  if (size > hDevice->getMaxAllocSize()) {
+    return UR_RESULT_ERROR_INVALID_USM_SIZE;
+  }
+
   if (!hPool) {
     return USMDeviceAllocImpl(ppMem, hContext, hDevice, /* flags */ 0, size,
                               alignment);
@@ -88,6 +92,10 @@ urUSMSharedAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
                 (alignment == 0 || ((alignment & (alignment - 1)) == 0)),
             UR_RESULT_ERROR_INVALID_VALUE);
 
+  if (size > hDevice->getMaxAllocSize()) {
+    return UR_RESULT_ERROR_INVALID_USM_SIZE;
+  }
+
   if (!hPool) {
     return USMSharedAllocImpl(ppMem, hContext, hDevice, /*host flags*/ 0,
                               /*device flags*/ 0, size, alignment);
diff --git a/source/adapters/hip/device.cpp b/source/adapters/hip/device.cpp
@@ -13,6 +13,7 @@
 #include "context.hpp"
 #include "event.hpp"
 #include "logger/ur_logger.hpp"
+#include "usm.hpp"
 
 #include <sstream>
 
@@ -206,22 +207,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
     return ReturnValue(Bits);
   }
   case UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE: {
-    // Max size of memory object allocation in bytes.
-    // The minimum value is max(min(1024 × 1024 ×
-    // 1024, 1/4th of CL_DEVICE_GLOBAL_MEM_SIZE),
-    // 32 × 1024 × 1024) for devices that are not of type
-    // CL_DEVICE_TYPE_CUSTOM.
-
-    size_t Global = 0;
-    detail::ur::assertion(hipDeviceTotalMem(&Global, hDevice->get()) ==
-                          hipSuccess);
-
-    auto QuarterGlobal = static_cast<uint32_t>(Global / 4u);
-
-    auto MaxAlloc = std::max(std::min(1024u * 1024u * 1024u, QuarterGlobal),
-                             32u * 1024u * 1024u);
-
-    return ReturnValue(uint64_t{MaxAlloc});
+    return ReturnValue(uint64_t{maxUSMAllocationSize(hDevice)});
   }
   case UR_DEVICE_INFO_IMAGE_SUPPORTED: {
     bool Enabled = false;
diff --git a/source/adapters/hip/usm.cpp b/source/adapters/hip/usm.cpp
@@ -49,6 +49,10 @@ urUSMDeviceAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
   UR_ASSERT(checkUSMAlignment(alignment, pUSMDesc),
             UR_RESULT_ERROR_INVALID_VALUE);
 
+  if (size > maxUSMAllocationSize(hDevice)) {
+    return UR_RESULT_ERROR_INVALID_USM_SIZE;
+  }
+
   if (!hPool) {
     return USMDeviceAllocImpl(ppMem, hContext, hDevice, /* flags */ 0, size,
                               alignment);
@@ -66,6 +70,10 @@ urUSMSharedAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
   UR_ASSERT(checkUSMAlignment(alignment, pUSMDesc),
             UR_RESULT_ERROR_INVALID_VALUE);
 
+  if (size > maxUSMAllocationSize(hDevice)) {
+    return UR_RESULT_ERROR_INVALID_USM_SIZE;
+  }
+
   if (!hPool) {
     return USMSharedAllocImpl(ppMem, hContext, hDevice, /*host flags*/ 0,
                               /*device flags*/ 0, size, alignment);
@@ -454,6 +462,24 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMPoolGetInfo(
   }
 }
 
+uint64_t maxUSMAllocationSize(const ur_device_handle_t &Device) {
+  // Max size of memory object allocation in bytes.
+  // The minimum value is max(min(1024 × 1024 ×
+  // 1024, 1/4th of CL_DEVICE_GLOBAL_MEM_SIZE),
+  // 32 × 1024 × 1024) for devices that are not of type
+  // CL_DEVICE_TYPE_CUSTOM.
+  size_t Global = 0;
+  detail::ur::assertion(hipDeviceTotalMem(&Global, Device->get()) ==
+                        hipSuccess);
+
+  auto QuarterGlobal = static_cast<uint32_t>(Global / 4u);
+
+  auto MaxAlloc = std::max(std::min(1024u * 1024u * 1024u, QuarterGlobal),
+                           32u * 1024u * 1024u);
+
+  return MaxAlloc;
+}
+
 bool checkUSMAlignment(uint32_t &alignment, const ur_usm_desc_t *pUSMDesc) {
   alignment = pUSMDesc ? pUSMDesc->align : 0u;
   return (!pUSMDesc ||
diff --git a/source/adapters/hip/usm.hpp b/source/adapters/hip/usm.hpp
@@ -137,6 +137,8 @@ ur_result_t USMHostAllocImpl(void **ResultPtr, ur_context_handle_t Context,
                              ur_usm_host_mem_flags_t Flags, size_t Size,
                              uint32_t Alignment);
 
+uint64_t maxUSMAllocationSize(const ur_device_handle_t &Device);
+
 bool checkUSMAlignment(uint32_t &alignment, const ur_usm_desc_t *pUSMDesc);
 
 bool checkUSMImplAlignment(uint32_t Alignment, void **ResultPtr);
diff --git a/source/adapters/native_cpu/device.cpp b/source/adapters/native_cpu/device.cpp
@@ -12,6 +12,7 @@
 
 #include "common.hpp"
 #include "platform.hpp"
+#include "usm.hpp"
 
 #if defined(_MSC_VER) || defined(__MINGW32__) || defined(__MINGW64__)
 #ifndef NOMINMAX
@@ -321,14 +322,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
   case UR_DEVICE_INFO_PARTITION_AFFINITY_DOMAIN:
     return ReturnValue(ur_device_affinity_domain_flags_t{0});
   case UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE: {
-    size_t Global = hDevice->mem_size;
-
-    auto QuarterGlobal = static_cast<uint32_t>(Global / 4u);
-
-    auto MaxAlloc = std::max(std::min(1024u * 1024u * 1024u, QuarterGlobal),
-                             32u * 1024u * 1024u);
-
-    return ReturnValue(uint64_t{MaxAlloc});
+    return ReturnValue(uint64_t{native_cpu::detail::maxUSMAllocationSize(hDevice)});
   }
   case UR_DEVICE_INFO_EXECUTION_CAPABILITIES:
     // TODO : CHECK
diff --git a/source/adapters/native_cpu/usm.cpp b/source/adapters/native_cpu/usm.cpp
@@ -13,6 +13,7 @@
 
 #include "common.hpp"
 #include "context.hpp"
+#include "usm.hpp"
 #include <cstdlib>
 
 namespace umf {
@@ -27,7 +28,8 @@ static ur_result_t alloc_helper(ur_context_handle_t hContext,
   auto alignment = (pUSMDesc && pUSMDesc->align) ? pUSMDesc->align : 1u;
   UR_ASSERT(isPowerOf2(alignment), UR_RESULT_ERROR_UNSUPPORTED_ALIGNMENT);
   UR_ASSERT(ppMem, UR_RESULT_ERROR_INVALID_NULL_POINTER);
-  // TODO: Check Max size when UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE is implemented
+  // TODO: Check Max size for host allocations when
+  // UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE is implemented
   UR_ASSERT(size > 0, UR_RESULT_ERROR_INVALID_USM_SIZE);
 
   auto *ptr = hContext->add_alloc(alignment, type, size, nullptr);
@@ -49,8 +51,9 @@ 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 = hDevice;
   std::ignore = pool;
+  UR_ASSERT(size < native_cpu::detail::maxUSMAllocationSize(hDevice),
+            UR_RESULT_ERROR_INVALID_USM_SIZE);
 
   return alloc_helper(hContext, pUSMDesc, size, ppMem, UR_USM_TYPE_DEVICE);
 }
@@ -59,8 +62,9 @@ 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 = hDevice;
   std::ignore = pool;
+  UR_ASSERT(size < native_cpu::detail::maxUSMAllocationSize(hDevice),
+            UR_RESULT_ERROR_INVALID_USM_SIZE);
 
   return alloc_helper(hContext, pUSMDesc, size, ppMem, UR_USM_TYPE_SHARED);
 }
@@ -155,3 +159,12 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMReleaseExp(ur_context_handle_t Context,
   std::ignore = HostPtr;
   DIE_NO_IMPLEMENTATION;
 }
+
+uint64_t maxUSMAllocationSize(const ur_device_handle_t &Device) {
+  size_t Global = Device->mem_size;
+
+  auto QuarterGlobal = static_cast<uint32_t>(Global / 4u);
+
+  return std::max(std::min(1024u * 1024u * 1024u, QuarterGlobal),
+                  32u * 1024u * 1024u);
+}
diff --git a/source/adapters/native_cpu/usm.hpp b/source/adapters/native_cpu/usm.hpp
@@ -0,0 +1,17 @@
+//===--------- usm.hpp - Native CPU Adapter --------------------------===//
+//
+// Copyright (C) 2025 Intel Corporation
+//
+// Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM
+// Exceptions. See LICENSE.TXT
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===-----------------------------------------------------------------===//
+
+#include "common.hpp"
+
+namespace native_cpu {
+namespace detail {
+uint64_t maxUSMAllocationSize(const ur_device_handle_t &Device);
+} // namespace detail
+} // namespace native_cpu
diff --git a/test/conformance/usm/urUSMDeviceAlloc.cpp b/test/conformance/usm/urUSMDeviceAlloc.cpp
@@ -118,13 +118,19 @@ TEST_P(urUSMDeviceAllocTest, InvalidNullPtrResult) {
 }
 
 TEST_P(urUSMDeviceAllocTest, InvalidUSMSize) {
-    UUR_KNOWN_FAILURE_ON(uur::CUDA{}, uur::HIP{}, uur::LevelZero{},
-                         uur::NativeCPU{});
+    UUR_KNOWN_FAILURE_ON(uur::LevelZero{});
+
+    size_t max_size;
+    ASSERT_SUCCESS(urDeviceGetInfo(device, UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE,
+                                   sizeof(max_size), &max_size, nullptr));
+    if (max_size == std::numeric_limits<size_t>::max()) {
+        GTEST_SKIP() << "Device has no max allocation size";
+    }
 
     void *ptr = nullptr;
     ASSERT_EQ_RESULT(
         UR_RESULT_ERROR_INVALID_USM_SIZE,
-        urUSMDeviceAlloc(context, device, nullptr, pool, -1, &ptr));
+        urUSMDeviceAlloc(context, device, nullptr, pool, max_size + 1, &ptr));
 }
 
 TEST_P(urUSMDeviceAllocTest, InvalidValueAlignPowerOfTwo) {
diff --git a/test/conformance/usm/urUSMSharedAlloc.cpp b/test/conformance/usm/urUSMSharedAlloc.cpp
@@ -141,11 +141,16 @@ TEST_P(urUSMSharedAllocTest, InvalidNullPtrMem) {
 }
 
 TEST_P(urUSMSharedAllocTest, InvalidUSMSize) {
-    UUR_KNOWN_FAILURE_ON(uur::CUDA{}, uur::HIP{}, uur::NativeCPU{});
+    size_t max_size;
+    ASSERT_SUCCESS(urDeviceGetInfo(device, UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE,
+                                   sizeof(max_size), &max_size, nullptr));
+    if (max_size == std::numeric_limits<size_t>::max()) {
+        GTEST_SKIP() << "Device has no max allocation size";
+    }
 
     ASSERT_EQ_RESULT(
         UR_RESULT_ERROR_INVALID_USM_SIZE,
-        urUSMSharedAlloc(context, device, nullptr, pool, -1, &ptr));
+        urUSMSharedAlloc(context, device, nullptr, pool, max_size + 1, &ptr));
 }
 
 TEST_P(urUSMSharedAllocTest, InvalidValueAlignPowerOfTwo) {
