Improvements to align CTS and Spec for Memory:

- Add tests for UR_RESULT_ERROR_INVALID_NULL/HOST_POINTER for urMemImageCreate and urMemBufferCreate
- Add missing error condition to spec for urMemImageCreate for checking type of image description struct
- Add tests for UR_RESULT_ERROR_INVALID_NULL_HANDLE/POINTER for urMemBuffer/ImageCreateWithnativeHandle
- Remove skip for urMemImageCreateWithNativeHandle - buffer was invalid as it was never instantiated in fixture struct
- Update image format used in urMemImageTest fixture as it was invalid and likely to cause previously skipped test to fail (which now pass)
- Add missing DDI table entry for urMemImageCreateWithNativeHandle for OpenCL
- Add test for using different ur_mem_flag_t flags with urMemBufferPartition
- Add missing UR_MEM_INFO_REFERENCE_COUNT to spec for urMemGetInfo and added test for this in urMemRetain/Release
- Removed assert in L0 urMemGetInfo which would fail if mem type is not image or the query is not a context query
- Fixed potential bug with HIP urMemGetNativeHandle which would fail if the mem handle was an image because use of std::variant

Author: martygrant

include/ur_api.h

@@ -2515,8 +2515,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
@@ -2650,6 +2654,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`
@@ -2990,7 +2995,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

@@ -5631,6 +5631,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;
@@ -5672,6 +5675,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;
@@ -316,14 +315,18 @@ urMemGetNativeHandle(ur_mem_handle_t hMem, ur_device_handle_t Device,
       return UR_RESULT_ERROR_INVALID_MEM_OBJECT;
     }
   }
-  *phNativeMem = reinterpret_cast<ur_native_handle_t>(
-      std::get<BufferMem>(hMem->Mem).getPtr(Device));
-#elif defined(__HIP_PLATFORM_AMD__)
-  *phNativeMem = reinterpret_cast<ur_native_handle_t>(
-      std::get<BufferMem>(hMem->Mem).getPtr(Device));
-#else
+#elif !defined(__HIP_PLATFORM_AMD__)
 #error("Must define exactly one of __HIP_PLATFORM_AMD__ or __HIP_PLATFORM_NVIDIA__");
 #endif
+  if (std::holds_alternative<BufferMem>(hMem->Mem)) {
+    *phNativeMem = reinterpret_cast<ur_native_handle_t>(
+        std::get<BufferMem>(hMem->Mem).getPtr(Device));
+  } else if (std::holds_alternative<SurfaceMem>(hMem->Mem)) {
+    *phNativeMem = reinterpret_cast<ur_native_handle_t>(
+        std::get<SurfaceMem>(hMem->Mem).getSurface(Device));
+  } else {
+    return UR_RESULT_ERROR_INVALID_MEM_OBJECT;
+  }
   return UR_RESULT_SUCCESS;
 }
 

source/adapters/level_zero/memory.cpp

@@ -1844,9 +1844,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);
@@ -1859,8 +1856,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

@@ -522,6 +522,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/adapters/opencl/ur_interface_loader.cpp

@@ -154,6 +154,7 @@ urGetMemProcAddrTable(ur_api_version_t Version, ur_mem_dditable_t *pDdiTable) {
   pDdiTable->pfnBufferPartition = urMemBufferPartition;
   pDdiTable->pfnBufferCreateWithNativeHandle =
       urMemBufferCreateWithNativeHandle;
+  pDdiTable->pfnImageCreateWithNativeHandle = urMemImageCreateWithNativeHandle;
   pDdiTable->pfnGetInfo = urMemGetInfo;
   pDdiTable->pfnGetNativeHandle = urMemGetNativeHandle;
   pDdiTable->pfnImageCreate = urMemImageCreate;

source/loader/layers/validation/ur_valddi.cpp

@@ -1108,6 +1108,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;
         }
@@ -1507,7 +1511,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;
         }