From 58795943ad0fbc2ad330c9ca1696704127f4d048 Mon Sep 17 00:00:00 2001
From: Martin Morrison-Grant <martin.morrisongrant@codeplay.com>
Date: Mon, 3 Feb 2025 18:01:57 +0000
Subject: [PATCH] address feedback - use UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE in
 InvalidUSMSize test

---
 test/conformance/usm/helpers.h            | 20 ++++++++++----------
 test/conformance/usm/urUSMDeviceAlloc.cpp | 10 +++++++++-
 test/conformance/usm/urUSMHostAlloc.cpp   | 10 +++++++++-
 test/conformance/usm/urUSMSharedAlloc.cpp | 10 +++++++++-
 4 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/test/conformance/usm/helpers.h b/test/conformance/usm/helpers.h
index 6c9085cc8f..de5d3c1ccb 100644
--- a/test/conformance/usm/helpers.h
+++ b/test/conformance/usm/helpers.h
@@ -11,13 +11,13 @@
 
 namespace uur {
 
-using USMDeviceAllocParams =
+using USMAllocTestParams =
     std::tuple<uur::BoolTestParam, uint32_t, size_t, ur_usm_advice_flags_t>;
 
-struct urUSMAllocTest : uur::urQueueTestWithParam<uur::USMDeviceAllocParams> {
+struct urUSMAllocTest : uur::urQueueTestWithParam<uur::USMAllocTestParams> {
   void SetUp() override {
     UUR_RETURN_ON_FATAL_FAILURE(
-        uur::urQueueTestWithParam<uur::USMDeviceAllocParams>::SetUp());
+        uur::urQueueTestWithParam<uur::USMAllocTestParams>::SetUp());
     if (usePool) {
       ur_bool_t poolSupport = false;
       ASSERT_SUCCESS(uur::GetDeviceUSMPoolSupport(device, poolSupport));
@@ -34,7 +34,7 @@ struct urUSMAllocTest : uur::urQueueTestWithParam<uur::USMDeviceAllocParams> {
       ASSERT_SUCCESS(urUSMPoolRelease(pool));
     }
     UUR_RETURN_ON_FATAL_FAILURE(
-        uur::urQueueTestWithParam<uur::USMDeviceAllocParams>::TearDown());
+        uur::urQueueTestWithParam<uur::USMAllocTestParams>::TearDown());
   }
 
   ur_usm_pool_handle_t pool = nullptr;
@@ -49,21 +49,21 @@ struct urUSMAllocTest : uur::urQueueTestWithParam<uur::USMDeviceAllocParams> {
 template <typename T>
 inline std::string printUSMAllocTestString(
     const testing::TestParamInfo<typename T::ParamType> &info) {
-  // ParamType will be std::tuple<ur_device_handle_t, USMDeviceAllocParams>
+  // ParamType will be std::tuple<ur_device_handle_t, USMAllocTestParams>
   const auto device_handle = std::get<0>(info.param).device;
   const auto platform_device_name =
       uur::GetPlatformAndDeviceName(device_handle);
-  const auto &usmDeviceAllocParams = std::get<1>(info.param);
-  const auto &BoolParam = std::get<0>(usmDeviceAllocParams);
+  const auto &USMAllocTestParams = std::get<1>(info.param);
+  const auto &BoolParam = std::get<0>(USMAllocTestParams);
 
   std::stringstream ss;
   ss << BoolParam.name << (BoolParam.value ? "Enabled" : "Disabled"); // UsePool
   ss << "_";
-  ss << std::get<1>(usmDeviceAllocParams); // alignment
+  ss << std::get<1>(USMAllocTestParams); // alignment
   ss << "_";
-  ss << std::get<2>(usmDeviceAllocParams); // size
+  ss << std::get<2>(USMAllocTestParams); // size
   ss << "_";
-  ss << std::get<3>(usmDeviceAllocParams); // ur_usm_advice_flags_t
+  ss << std::get<3>(USMAllocTestParams); // ur_usm_advice_flags_t
 
   return platform_device_name + "__" + ss.str();
 }
diff --git a/test/conformance/usm/urUSMDeviceAlloc.cpp b/test/conformance/usm/urUSMDeviceAlloc.cpp
index 0e824893fc..83e2664dda 100644
--- a/test/conformance/usm/urUSMDeviceAlloc.cpp
+++ b/test/conformance/usm/urUSMDeviceAlloc.cpp
@@ -101,7 +101,15 @@ TEST_P(urUSMDeviceAllocTest, InvalidUSMSize) {
       UR_RESULT_ERROR_INVALID_USM_SIZE,
       urUSMDeviceAlloc(context, device, nullptr, pool, size, &ptr));
 
-  size = std::numeric_limits<size_t>::max() + 1;
+  size_t max_mem_alloc_info_size;
+  urDeviceGetInfo(device, UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE, 0, nullptr,
+                  &max_mem_alloc_info_size);
+
+  uint64_t max_mem_alloc_size = 0;
+  urDeviceGetInfo(device, UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE,
+                  max_mem_alloc_info_size, &max_mem_alloc_size, nullptr);
+
+  size = max_mem_alloc_size + 1;
   ASSERT_EQ_RESULT(
       UR_RESULT_ERROR_INVALID_USM_SIZE,
       urUSMDeviceAlloc(context, device, nullptr, pool, size, &ptr));
diff --git a/test/conformance/usm/urUSMHostAlloc.cpp b/test/conformance/usm/urUSMHostAlloc.cpp
index ec95736712..1109f2ad80 100644
--- a/test/conformance/usm/urUSMHostAlloc.cpp
+++ b/test/conformance/usm/urUSMHostAlloc.cpp
@@ -109,7 +109,15 @@ TEST_P(urUSMHostAllocTest, InvalidUSMSize) {
   ASSERT_EQ_RESULT(UR_RESULT_ERROR_INVALID_USM_SIZE,
                    urUSMHostAlloc(context, nullptr, pool, size, &ptr));
 
-  size = std::numeric_limits<size_t>::max() + 1;
+  size_t max_mem_alloc_info_size;
+  urDeviceGetInfo(device, UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE, 0, nullptr,
+                  &max_mem_alloc_info_size);
+
+  uint64_t max_mem_alloc_size = 0;
+  urDeviceGetInfo(device, UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE,
+                  max_mem_alloc_info_size, &max_mem_alloc_size, nullptr);
+
+  size = max_mem_alloc_size + 1;
   ASSERT_EQ_RESULT(UR_RESULT_ERROR_INVALID_USM_SIZE,
                    urUSMHostAlloc(context, nullptr, pool, size, &ptr));
 }
diff --git a/test/conformance/usm/urUSMSharedAlloc.cpp b/test/conformance/usm/urUSMSharedAlloc.cpp
index 697012a592..1021cb58a2 100644
--- a/test/conformance/usm/urUSMSharedAlloc.cpp
+++ b/test/conformance/usm/urUSMSharedAlloc.cpp
@@ -133,7 +133,15 @@ TEST_P(urUSMSharedAllocTest, InvalidUSMSize) {
       UR_RESULT_ERROR_INVALID_USM_SIZE,
       urUSMSharedAlloc(context, device, nullptr, pool, size, &ptr));
 
-  size = std::numeric_limits<size_t>::max() + 1;
+  size_t max_mem_alloc_info_size;
+  urDeviceGetInfo(device, UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE, 0, nullptr,
+                  &max_mem_alloc_info_size);
+
+  uint64_t max_mem_alloc_size = 0;
+  urDeviceGetInfo(device, UR_DEVICE_INFO_MAX_MEM_ALLOC_SIZE,
+                  max_mem_alloc_info_size, &max_mem_alloc_size, nullptr);
+
+  size = max_mem_alloc_size + 1;
   ASSERT_EQ_RESULT(
       UR_RESULT_ERROR_INVALID_USM_SIZE,
       urUSMSharedAlloc(context, device, nullptr, pool, size, &ptr));