WIP improvements to memory cts

include/ur_api.h

@@ -2510,8 +2510,12 @@ typedef enum ur_mem_type_t {
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Memory Information type
 typedef enum ur_mem_info_t {
-    UR_MEM_INFO_SIZE = 0,    ///< [size_t] actual size of of memory object in bytes
-    UR_MEM_INFO_CONTEXT = 1, ///< [::ur_context_handle_t] context in which the memory object was created
+    UR_MEM_INFO_SIZE = 0,            ///< [size_t] actual size of of memory object in bytes
+    UR_MEM_INFO_CONTEXT = 1,         ///< [::ur_context_handle_t] context in which the memory object was created
+    UR_MEM_INFO_REFERENCE_COUNT = 2, ///< [uint32_t] Reference count of the memory object.
+                                     ///< The reference count returned should be considered immediately stale.
+                                     ///< It is unsuitable for general use in applications. This feature is
+                                     ///< provided for identifying memory leaks.
     /// @cond
     UR_MEM_INFO_FORCE_UINT32 = 0x7fffffff
     /// @endcond
@@ -2645,6 +2649,7 @@ typedef struct ur_image_desc_t {
 ///     - ::UR_RESULT_ERROR_INVALID_CONTEXT
 ///     - ::UR_RESULT_ERROR_INVALID_VALUE
 ///     - ::UR_RESULT_ERROR_INVALID_IMAGE_FORMAT_DESCRIPTOR
+///         + `pImageDesc && UR_STRUCTURE_TYPE_IMAGE_DESC != pImageDesc->stype`
 ///         + `pImageDesc && UR_MEM_TYPE_IMAGE1D_ARRAY < pImageDesc->type`
 ///         + `pImageDesc && pImageDesc->numMipLevel != 0`
 ///         + `pImageDesc && pImageDesc->numSamples != 0`
@@ -2985,7 +2990,7 @@ urMemImageCreateWithNativeHandle(
 ///     - ::UR_RESULT_ERROR_INVALID_NULL_HANDLE
 ///         + `NULL == hMemory`
 ///     - ::UR_RESULT_ERROR_INVALID_ENUMERATION
-///         + `::UR_MEM_INFO_CONTEXT < propName`
+///         + `::UR_MEM_INFO_REFERENCE_COUNT < propName`
 ///     - ::UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION
 ///         + If `propName` is not supported by the adapter.
 ///     - ::UR_RESULT_ERROR_INVALID_SIZE

include/ur_print.hpp

@@ -5616,6 +5616,9 @@ inline std::ostream &operator<<(std::ostream &os, enum ur_mem_info_t value) {
     case UR_MEM_INFO_CONTEXT:
         os << "UR_MEM_INFO_CONTEXT";
         break;
+    case UR_MEM_INFO_REFERENCE_COUNT:
+        os << "UR_MEM_INFO_REFERENCE_COUNT";
+        break;
     default:
         os << "unknown enumerator";
         break;
@@ -5657,6 +5660,18 @@ inline ur_result_t printTagged(std::ostream &os, const void *ptr, ur_mem_info_t
 
         os << ")";
     } break;
+    case UR_MEM_INFO_REFERENCE_COUNT: {
+        const uint32_t *tptr = (const uint32_t *)ptr;
+        if (sizeof(uint32_t) > size) {
+            os << "invalid size (is: " << size << ", expected: >=" << sizeof(uint32_t) << ")";
+            return UR_RESULT_ERROR_INVALID_SIZE;
+        }
+        os << (const void *)(tptr) << " (";
+
+        os << *tptr;
+
+        os << ")";
+    } break;
     default:
         os << "unknown enumerator";
         return UR_RESULT_ERROR_INVALID_ENUMERATION;

scripts/core/memory.yml

@@ -62,6 +62,11 @@ etors:
       desc: "[size_t] actual size of of memory object in bytes"
     - name: CONTEXT
       desc: "[$x_context_handle_t] context in which the memory object was created"
+    - name: REFERENCE_COUNT
+      desc: |
+            [uint32_t] Reference count of the memory object.
+            The reference count returned should be considered immediately stale.
+            It is unsuitable for general use in applications. This feature is provided for identifying memory leaks.
 --- #--------------------------------------------------------------------------
 type: enum
 desc: "Image channel order info: number of channels and the channel layout"
@@ -241,6 +246,7 @@ returns:
     - $X_RESULT_ERROR_INVALID_CONTEXT
     - $X_RESULT_ERROR_INVALID_VALUE
     - $X_RESULT_ERROR_INVALID_IMAGE_FORMAT_DESCRIPTOR:
+      - "`pImageDesc && UR_STRUCTURE_TYPE_IMAGE_DESC != pImageDesc->stype`"
       - "`pImageDesc && UR_MEM_TYPE_IMAGE1D_ARRAY < pImageDesc->type`"
       - "`pImageDesc && pImageDesc->numMipLevel != 0`"
       - "`pImageDesc && pImageDesc->numSamples != 0`"

source/adapters/cuda/memory.cpp

@@ -171,6 +171,9 @@ UR_APIEXPORT ur_result_t UR_APICALL urMemGetInfo(ur_mem_handle_t hMemory,
   case UR_MEM_INFO_CONTEXT: {
     return ReturnValue(hMemory->getContext());
   }
+  case UR_MEM_INFO_REFERENCE_COUNT: {
+    return ReturnValue(hMemory->getReferenceCount());
+  }
 
   default:
     return UR_RESULT_ERROR_INVALID_ENUMERATION;

source/adapters/hip/memory.cpp

@@ -231,10 +231,6 @@ UR_APIEXPORT ur_result_t UR_APICALL urMemGetInfo(ur_mem_handle_t hMemory,
                                                  size_t propSize,
                                                  void *pMemInfo,
                                                  size_t *pPropSizeRet) {
-
-  UR_ASSERT(MemInfoType <= UR_MEM_INFO_CONTEXT,
-            UR_RESULT_ERROR_INVALID_ENUMERATION);
-
   // FIXME: Only getting info for the first device in the context. This
   // should be fine in general
   auto Device = hMemory->getContext()->getDevices()[0];
@@ -286,6 +282,9 @@ UR_APIEXPORT ur_result_t UR_APICALL urMemGetInfo(ur_mem_handle_t hMemory,
   case UR_MEM_INFO_CONTEXT: {
     return ReturnValue(hMemory->getContext());
   }
+  case UR_MEM_INFO_REFERENCE_COUNT: {
+    return ReturnValue(hMemory->getReferenceCount());
+  }
 
   default:
     return UR_RESULT_ERROR_INVALID_ENUMERATION;

source/adapters/level_zero/memory.cpp

@@ -1871,9 +1871,6 @@ ur_result_t urMemGetInfo(
     size_t *PropSizeRet ///< [out][optional] pointer to the actual size in
                         ///< bytes of data queried by pMemInfo.
 ) {
-  UR_ASSERT(MemInfoType == UR_MEM_INFO_CONTEXT || !Memory->isImage(),
-            UR_RESULT_ERROR_INVALID_VALUE);
-
   auto Buffer = reinterpret_cast<_ur_buffer *>(Memory);
   std::shared_lock<ur_shared_mutex> Lock(Buffer->Mutex);
   UrReturnHelper ReturnValue(PropSize, MemInfo, PropSizeRet);
@@ -1886,8 +1883,11 @@ ur_result_t urMemGetInfo(
     // Get size of the allocation
     return ReturnValue(size_t{Buffer->Size});
   }
+  case UR_MEM_INFO_REFERENCE_COUNT: {
+    return ReturnValue(Buffer->RefCount.load());
+  }
   default: {
-    die("urMemGetInfo: Parameter is not implemented");
+    return UR_RESULT_ERROR_INVALID_ENUMERATION;
   }
   }
 

source/adapters/level_zero/v2/memory.cpp

@@ -205,6 +205,9 @@ ur_result_t urMemGetInfo(ur_mem_handle_t hMemory, ur_mem_info_t propName,
     // Get size of the allocation
     return returnValue(size_t{hMemory->getSize()});
   }
+  case UR_MEM_INFO_REFERENCE_COUNT: {
+    return returnValue(hMemory->RefCount.load());
+  }
   default: {
     return UR_RESULT_ERROR_INVALID_ENUMERATION;
   }

source/adapters/opencl/memory.cpp

@@ -186,6 +186,8 @@ cl_int mapURMemInfoToCL(ur_mem_info_t URPropName) {
     return CL_MEM_SIZE;
   case UR_MEM_INFO_CONTEXT:
     return CL_MEM_CONTEXT;
+  case UR_MEM_INFO_REFERENCE_COUNT:
+    return CL_MEM_REFERENCE_COUNT;
   default:
     return -1;
   }

source/loader/layers/validation/ur_valddi.cpp

@@ -1104,6 +1104,10 @@ __urdlllocal ur_result_t UR_APICALL urMemImageCreate(
             return UR_RESULT_ERROR_INVALID_ENUMERATION;
         }
 
+        if (pImageDesc && UR_STRUCTURE_TYPE_IMAGE_DESC != pImageDesc->stype) {
+            return UR_RESULT_ERROR_INVALID_IMAGE_FORMAT_DESCRIPTOR;
+        }
+
         if (pImageDesc && UR_MEM_TYPE_IMAGE1D_ARRAY < pImageDesc->type) {
             return UR_RESULT_ERROR_INVALID_IMAGE_FORMAT_DESCRIPTOR;
         }
@@ -1503,7 +1507,7 @@ __urdlllocal ur_result_t UR_APICALL urMemGetInfo(
             return UR_RESULT_ERROR_INVALID_NULL_POINTER;
         }
 
-        if (UR_MEM_INFO_CONTEXT < propName) {
+        if (UR_MEM_INFO_REFERENCE_COUNT < propName) {
             return UR_RESULT_ERROR_INVALID_ENUMERATION;
         }
 

source/loader/ur_libapi.cpp

@@ -1539,6 +1539,7 @@ ur_result_t UR_APICALL urContextSetExtendedDeleter(
 ///     - ::UR_RESULT_ERROR_INVALID_CONTEXT
 ///     - ::UR_RESULT_ERROR_INVALID_VALUE
 ///     - ::UR_RESULT_ERROR_INVALID_IMAGE_FORMAT_DESCRIPTOR
+///         + `pImageDesc && UR_STRUCTURE_TYPE_IMAGE_DESC != pImageDesc->stype`
 ///         + `pImageDesc && UR_MEM_TYPE_IMAGE1D_ARRAY < pImageDesc->type`
 ///         + `pImageDesc && pImageDesc->numMipLevel != 0`
 ///         + `pImageDesc && pImageDesc->numSamples != 0`
@@ -1888,7 +1889,7 @@ ur_result_t UR_APICALL urMemImageCreateWithNativeHandle(
 ///     - ::UR_RESULT_ERROR_INVALID_NULL_HANDLE
 ///         + `NULL == hMemory`
 ///     - ::UR_RESULT_ERROR_INVALID_ENUMERATION
-///         + `::UR_MEM_INFO_CONTEXT < propName`
+///         + `::UR_MEM_INFO_REFERENCE_COUNT < propName`
 ///     - ::UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION
 ///         + If `propName` is not supported by the adapter.
 ///     - ::UR_RESULT_ERROR_INVALID_SIZE

source/ur_api.cpp

@@ -1340,6 +1340,7 @@ ur_result_t UR_APICALL urContextSetExtendedDeleter(
 ///     - ::UR_RESULT_ERROR_INVALID_CONTEXT
 ///     - ::UR_RESULT_ERROR_INVALID_VALUE
 ///     - ::UR_RESULT_ERROR_INVALID_IMAGE_FORMAT_DESCRIPTOR
+///         + `pImageDesc && UR_STRUCTURE_TYPE_IMAGE_DESC != pImageDesc->stype`
 ///         + `pImageDesc && UR_MEM_TYPE_IMAGE1D_ARRAY < pImageDesc->type`
 ///         + `pImageDesc && pImageDesc->numMipLevel != 0`
 ///         + `pImageDesc && pImageDesc->numSamples != 0`
@@ -1634,7 +1635,7 @@ ur_result_t UR_APICALL urMemImageCreateWithNativeHandle(
 ///     - ::UR_RESULT_ERROR_INVALID_NULL_HANDLE
 ///         + `NULL == hMemory`
 ///     - ::UR_RESULT_ERROR_INVALID_ENUMERATION
-///         + `::UR_MEM_INFO_CONTEXT < propName`
+///         + `::UR_MEM_INFO_REFERENCE_COUNT < propName`
 ///     - ::UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION
 ///         + If `propName` is not supported by the adapter.
 ///     - ::UR_RESULT_ERROR_INVALID_SIZE

test/conformance/memory/memory_adapter_level_zero.match

@@ -1,4 +1,6 @@
 {{NONDETERMINISTIC}}
+urMemBufferPartitionWithFlagsTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}__UR_MEM_FLAG_WRITE_ONLY
+urMemBufferPartitionWithFlagsTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}__UR_MEM_FLAG_READ_ONLY
 urMemBufferPartitionTest.InvalidValueCreateType/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 urMemBufferPartitionTest.InvalidValueBufferCreateInfoOutOfBounds/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 {{OPT}}urMemGetInfoImageTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_MEM_INFO_SIZE

test/conformance/memory/memory_adapter_level_zero_v2.match

@@ -1,9 +1,12 @@
 {{NONDETERMINISTIC}}
-urMemBufferPartitionTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
+urMemBufferPartitionWithFlagsTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_MEM_FLAG_WRITE_ONLY
+urMemBufferPartitionWithFlagsTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_MEM_FLAG_READ_ONLY
+urMemBufferPartitionWithFlagsTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_MEM_FLAG_READ_WRITE
 urMemBufferPartitionTest.InvalidValueCreateType/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 urMemBufferPartitionTest.InvalidValueBufferCreateInfoOutOfBounds/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 {{OPT}}urMemGetInfoImageTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_MEM_INFO_SIZE
 {{OPT}}urMemGetInfoImageTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_MEM_INFO_CONTEXT
+{{OPT}}urMemGetInfoImageTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_MEM_INFO_REFERENCE_COUNT
 {{OPT}}urMemImageCreateTestWithImageFormatParam.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_IMAGE_CHANNEL_ORDER_A__UR_IMAGE_CHANNEL_TYPE_SNORM_INT8
 {{OPT}}urMemImageCreateTestWithImageFormatParam.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_IMAGE_CHANNEL_ORDER_A__UR_IMAGE_CHANNEL_TYPE_SNORM_INT16
 {{OPT}}urMemImageCreateTestWithImageFormatParam.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_IMAGE_CHANNEL_ORDER_A__UR_IMAGE_CHANNEL_TYPE_UNORM_INT8

test/conformance/memory/memory_adapter_native_cpu.match

@@ -1,10 +1,12 @@
 {{NONDETERMINISTIC}}
+urMemBufferPartitionWithFlagsTest.Success/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}__UR_MEM_FLAG_WRITE_ONLY
+urMemBufferPartitionWithFlagsTest.Success/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}__UR_MEM_FLAG_READ_ONLY
 urMemBufferPartitionTest.InvalidValueCreateType/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}
 urMemBufferPartitionTest.InvalidValueBufferCreateInfoOutOfBounds/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}
-urMemGetInfoTest.Success/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}__UR_MEM_INFO_SIZE
-urMemGetInfoTest.Success/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}__UR_MEM_INFO_CONTEXT
-urMemGetInfoTest.InvalidSizeSmall/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}__UR_MEM_INFO_SIZE
-urMemGetInfoTest.InvalidSizeSmall/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}__UR_MEM_INFO_CONTEXT
+urMemGetInfoTestWithParam.Success/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}__UR_MEM_INFO_SIZE
+urMemGetInfoTestWithParam.Success/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}__UR_MEM_INFO_CONTEXT
+urMemGetInfoTestWithParam.Success/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}__UR_MEM_INFO_REFERENCE_COUNT
+urMemGetInfoTest.InvalidSizeSmall/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}
 urMemImageCreateTestWithImageFormatParam.Success/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}__UR_IMAGE_CHANNEL_ORDER_A__UR_IMAGE_CHANNEL_TYPE_SNORM_INT8
 urMemImageCreateTestWithImageFormatParam.Success/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}__UR_IMAGE_CHANNEL_ORDER_A__UR_IMAGE_CHANNEL_TYPE_SNORM_INT16
 urMemImageCreateTestWithImageFormatParam.Success/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}__UR_IMAGE_CHANNEL_ORDER_A__UR_IMAGE_CHANNEL_TYPE_UNORM_INT8
@@ -232,3 +234,5 @@ urMemImageCreateTestWithImageFormatParam.Success/SYCL_NATIVE_CPU___SYCL_Native_C
 urMemImageCreateTestWithImageFormatParam.Success/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}__UR_IMAGE_CHANNEL_ORDER_SRGBA__UR_IMAGE_CHANNEL_TYPE_FLOAT
 urMemReleaseTest.Success/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}
 urMemRetainTest.Success/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}
+urMemReleaseTest.CheckReferenceCount/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}
+urMemRetainTest.CheckReferenceCount/SYCL_NATIVE_CPU___SYCL_Native_CPU__{{.*}}

test/conformance/memory/urMemBufferPartition.cpp

@@ -6,19 +6,32 @@
 #include "uur/fixtures.h"
 #include "uur/raii.h"
 
-using urMemBufferPartitionTest = uur::urMemBufferTest;
-UUR_INSTANTIATE_DEVICE_TEST_SUITE_P(urMemBufferPartitionTest);
+using urMemBufferPartitionWithFlagsTest =
+    uur::urContextTestWithParam<ur_mem_flag_t>;
+UUR_TEST_SUITE_P(urMemBufferPartitionWithFlagsTest,
+                 ::testing::Values(UR_MEM_FLAG_READ_WRITE,
+                                   UR_MEM_FLAG_WRITE_ONLY,
+                                   UR_MEM_FLAG_READ_ONLY),
+                 uur::deviceTestWithParamPrinter<ur_mem_flag_t>);
 
-TEST_P(urMemBufferPartitionTest, Success) {
-    ur_buffer_region_t region{UR_STRUCTURE_TYPE_BUFFER_REGION, nullptr, 0,
-                              1024};
+TEST_P(urMemBufferPartitionWithFlagsTest, Success) {
+    uur::raii::Mem buffer = nullptr;
+
+    ASSERT_SUCCESS(
+        urMemBufferCreate(context, getParam(), 1024, nullptr, buffer.ptr()));
+    ASSERT_NE(nullptr, buffer);
+
+    ur_buffer_region_t region{UR_STRUCTURE_TYPE_BUFFER_REGION, nullptr, 0, 512};
     uur::raii::Mem partition = nullptr;
-    ASSERT_SUCCESS(urMemBufferPartition(buffer, UR_MEM_FLAG_READ_WRITE,
+    ASSERT_SUCCESS(urMemBufferPartition(buffer, getParam(),
                                         UR_BUFFER_CREATE_TYPE_REGION, &region,
                                         partition.ptr()));
     ASSERT_NE(partition, nullptr);
 }
 
+using urMemBufferPartitionTest = uur::urMemBufferTest;
+UUR_INSTANTIATE_DEVICE_TEST_SUITE_P(urMemBufferPartitionTest);
+
 TEST_P(urMemBufferPartitionTest, InvalidNullHandleBuffer) {
     ur_buffer_region_t region{UR_STRUCTURE_TYPE_BUFFER_REGION, nullptr, 0,
                               1024};

test/conformance/memory/urMemGetInfo.cpp

@@ -6,19 +6,21 @@
 #include <map>
 #include <uur/fixtures.h>
 
-using urMemGetInfoTest = uur::urMemBufferTestWithParam<ur_mem_info_t>;
+using urMemGetInfoTestWithParam = uur::urMemBufferTestWithParam<ur_mem_info_t>;
 
-static constexpr std::array<ur_mem_info_t, 2> mem_info_values{
-    UR_MEM_INFO_SIZE, UR_MEM_INFO_CONTEXT};
+static constexpr std::array<ur_mem_info_t, 3> mem_info_values{
+    UR_MEM_INFO_SIZE, UR_MEM_INFO_CONTEXT, UR_MEM_INFO_REFERENCE_COUNT};
 static std::unordered_map<ur_mem_info_t, size_t> mem_info_size_map = {
     {UR_MEM_INFO_SIZE, sizeof(size_t)},
     {UR_MEM_INFO_CONTEXT, sizeof(ur_context_handle_t)},
+    {UR_MEM_INFO_REFERENCE_COUNT, sizeof(uint32_t)},
 };
 
-UUR_TEST_SUITE_P(urMemGetInfoTest, ::testing::ValuesIn(mem_info_values),
+UUR_TEST_SUITE_P(urMemGetInfoTestWithParam,
+                 ::testing::ValuesIn(mem_info_values),
                  uur::deviceTestWithParamPrinter<ur_mem_info_t>);
 
-TEST_P(urMemGetInfoTest, Success) {
+TEST_P(urMemGetInfoTestWithParam, Success) {
     ur_mem_info_t info = getParam();
     size_t size;
     ASSERT_SUCCESS(urMemGetInfo(buffer, info, 0, nullptr, &size));
@@ -44,11 +46,20 @@ TEST_P(urMemGetInfoTest, Success) {
         ASSERT_GE(*returned_size, allocation_size);
         break;
     }
+    case UR_MEM_INFO_REFERENCE_COUNT: {
+        const size_t ReferenceCount =
+            *reinterpret_cast<const uint32_t *>(info_data.data());
+        ASSERT_GT(ReferenceCount, 0);
+        break;
+    }
     default:
         break;
     }
 }
 
+using urMemGetInfoTest = uur::urMemBufferTest;
+UUR_INSTANTIATE_DEVICE_TEST_SUITE_P(urMemGetInfoTest);
+
 TEST_P(urMemGetInfoTest, InvalidNullHandleMemory) {
     size_t mem_size = 0;
     ASSERT_EQ_RESULT(UR_RESULT_ERROR_INVALID_NULL_HANDLE,
@@ -108,16 +119,30 @@ TEST_P(urMemGetInfoImageTest, Success) {
     std::vector<uint8_t> info_data(size);
     ASSERT_SUCCESS(urMemGetInfo(image, info, size, info_data.data(), nullptr));
 
-    if (info == UR_MEM_INFO_SIZE) {
+    switch (info) {
+    case UR_MEM_INFO_SIZE: {
         const size_t ExpectedPixelSize = sizeof(float) * 4 /*NumChannels*/;
         const size_t ExpectedImageSize = ExpectedPixelSize * desc.arraySize *
                                          desc.width * desc.height * desc.depth;
         const size_t ImageSizeBytes =
             *reinterpret_cast<const size_t *>(info_data.data());
         ASSERT_EQ(ImageSizeBytes, ExpectedImageSize);
-    } else if (info == UR_MEM_INFO_CONTEXT) {
+        break;
+    }
+    case UR_MEM_INFO_CONTEXT: {
         ur_context_handle_t InfoContext =
             *reinterpret_cast<ur_context_handle_t *>(info_data.data());
         ASSERT_EQ(InfoContext, context);
+        break;
+    }
+    case UR_MEM_INFO_REFERENCE_COUNT: {
+        const size_t ReferenceCount =
+            *reinterpret_cast<const uint32_t *>(info_data.data());
+        ASSERT_GT(ReferenceCount, 0);
+        break;
+    }
+
+    default:
+        break;
     }
 }

test/conformance/memory/urMemRelease.cpp

@@ -16,3 +16,24 @@ TEST_P(urMemReleaseTest, InvalidNullHandleMem) {
     ASSERT_EQ_RESULT(UR_RESULT_ERROR_INVALID_NULL_HANDLE,
                      urMemRelease(nullptr));
 }
+
+TEST_P(urMemReleaseTest, CheckReferenceCount) {
+    uint32_t referenceCount = 0;
+    ASSERT_SUCCESS(urMemGetInfo(buffer, UR_MEM_INFO_REFERENCE_COUNT,
+                                sizeof(referenceCount), &referenceCount,
+                                nullptr));
+    ASSERT_EQ(referenceCount, 1);
+
+    ASSERT_SUCCESS(urMemRetain(buffer));
+    ASSERT_SUCCESS(urMemGetInfo(buffer, UR_MEM_INFO_REFERENCE_COUNT,
+                                sizeof(referenceCount), &referenceCount,
+                                nullptr));
+    ASSERT_EQ(referenceCount, 2);
+
+    ASSERT_SUCCESS(urMemRelease(buffer));
+
+    ASSERT_SUCCESS(urMemGetInfo(buffer, UR_MEM_INFO_REFERENCE_COUNT,
+                                sizeof(referenceCount), &referenceCount,
+                                nullptr));
+    ASSERT_EQ(referenceCount, 1);
+}