Add urUSMPoolFree function to spec

include/ur.py

@@ -196,6 +196,7 @@ class ur_function_v(IntEnum):
     ADAPTER_RETAIN = 179                            ## Enumerator for ::urAdapterRetain
     ADAPTER_GET_LAST_ERROR = 180                    ## Enumerator for ::urAdapterGetLastError
     ADAPTER_GET_INFO = 181                          ## Enumerator for ::urAdapterGetInfo
+    USM_POOL_FREE = 182                             ## Enumerator for ::urUSMPoolFree
 
 class ur_function_t(c_int):
     def __str__(self):
@@ -1548,7 +1549,7 @@ class ur_usm_pool_limits_desc_t(Structure):
 ## @brief Get USM memory pool information
 class ur_usm_pool_info_v(IntEnum):
     REFERENCE_COUNT = 0                             ## [uint32_t] Reference count of the pool object.
-                                                    ## The reference count returned should be considered immediately stale. 
+                                                    ## 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.
     CONTEXT = 1                                     ## [::ur_context_handle_t] USM memory pool context info
@@ -3359,6 +3360,13 @@ class ur_bindless_images_exp_dditable_t(Structure):
 else:
     _urUSMFree_t = CFUNCTYPE( ur_result_t, ur_context_handle_t, c_void_p )
 
+###############################################################################
+## @brief Function-pointer for urUSMPoolFree
+if __use_win_types:
+    _urUSMPoolFree_t = WINFUNCTYPE( ur_result_t, ur_usm_pool_handle_t, c_void_p )
+else:
+    _urUSMPoolFree_t = CFUNCTYPE( ur_result_t, ur_usm_pool_handle_t, c_void_p )
+
 ###############################################################################
 ## @brief Function-pointer for urUSMGetMemAllocInfo
 if __use_win_types:
@@ -3403,6 +3411,7 @@ class ur_usm_dditable_t(Structure):
         ("pfnDeviceAlloc", c_void_p),                                   ## _urUSMDeviceAlloc_t
         ("pfnSharedAlloc", c_void_p),                                   ## _urUSMSharedAlloc_t
         ("pfnFree", c_void_p),                                          ## _urUSMFree_t
+        ("pfnPoolFree", c_void_p),                                      ## _urUSMPoolFree_t
         ("pfnGetMemAllocInfo", c_void_p),                               ## _urUSMGetMemAllocInfo_t
         ("pfnPoolCreate", c_void_p),                                    ## _urUSMPoolCreate_t
         ("pfnPoolRetain", c_void_p),                                    ## _urUSMPoolRetain_t
@@ -4057,6 +4066,7 @@ def __init__(self, version : ur_api_version_t):
         self.urUSMDeviceAlloc = _urUSMDeviceAlloc_t(self.__dditable.USM.pfnDeviceAlloc)
         self.urUSMSharedAlloc = _urUSMSharedAlloc_t(self.__dditable.USM.pfnSharedAlloc)
         self.urUSMFree = _urUSMFree_t(self.__dditable.USM.pfnFree)
+        self.urUSMPoolFree = _urUSMPoolFree_t(self.__dditable.USM.pfnPoolFree)
         self.urUSMGetMemAllocInfo = _urUSMGetMemAllocInfo_t(self.__dditable.USM.pfnGetMemAllocInfo)
         self.urUSMPoolCreate = _urUSMPoolCreate_t(self.__dditable.USM.pfnPoolCreate)
         self.urUSMPoolRetain = _urUSMPoolRetain_t(self.__dditable.USM.pfnPoolRetain)

include/ur_api.h

@@ -205,6 +205,7 @@ typedef enum ur_function_t {
     UR_FUNCTION_ADAPTER_RETAIN = 179,                                          ///< Enumerator for ::urAdapterRetain
     UR_FUNCTION_ADAPTER_GET_LAST_ERROR = 180,                                  ///< Enumerator for ::urAdapterGetLastError
     UR_FUNCTION_ADAPTER_GET_INFO = 181,                                        ///< Enumerator for ::urAdapterGetInfo
+    UR_FUNCTION_USM_POOL_FREE = 182,                                           ///< Enumerator for ::urUSMPoolFree
     /// @cond
     UR_FUNCTION_FORCE_UINT32 = 0x7fffffff
     /// @endcond
@@ -3386,6 +3387,30 @@ urUSMFree(
     void *pMem                    ///< [in] pointer to USM memory object
 );
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Free the USM memory object from the pool
+///
+/// @details
+///     - Calling this function is equivalent to urUSMFree but without the pool
+///       look-up overhead.
+///
+/// @returns
+///     - ::UR_RESULT_SUCCESS
+///     - ::UR_RESULT_ERROR_UNINITIALIZED
+///     - ::UR_RESULT_ERROR_DEVICE_LOST
+///     - ::UR_RESULT_ERROR_ADAPTER_SPECIFIC
+///     - ::UR_RESULT_ERROR_INVALID_NULL_HANDLE
+///         + `NULL == pool`
+///     - ::UR_RESULT_ERROR_INVALID_NULL_POINTER
+///         + `NULL == pMem`
+///     - ::UR_RESULT_ERROR_INVALID_MEM_OBJECT
+///     - ::UR_RESULT_ERROR_OUT_OF_HOST_MEMORY
+UR_APIEXPORT ur_result_t UR_APICALL
+urUSMPoolFree(
+    ur_usm_pool_handle_t pool, ///< [in] Pointer to a pool created using urUSMPoolCreate
+    void *pMem                 ///< [in] pointer to USM memory object
+);
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Get USM memory object allocation information
 ///
@@ -9719,6 +9744,15 @@ typedef struct ur_usm_free_params_t {
     void **ppMem;
 } ur_usm_free_params_t;
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Function parameters for urUSMPoolFree
+/// @details Each entry is a pointer to the parameter passed to the function;
+///     allowing the callback the ability to modify the parameter's value
+typedef struct ur_usm_pool_free_params_t {
+    ur_usm_pool_handle_t *ppool;
+    void **ppMem;
+} ur_usm_pool_free_params_t;
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Function parameters for urUSMGetMemAllocInfo
 /// @details Each entry is a pointer to the parameter passed to the function;

include/ur_ddi.h

@@ -1500,6 +1500,12 @@ typedef ur_result_t(UR_APICALL *ur_pfnUSMFree_t)(
     ur_context_handle_t,
     void *);
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Function-pointer for urUSMPoolFree
+typedef ur_result_t(UR_APICALL *ur_pfnUSMPoolFree_t)(
+    ur_usm_pool_handle_t,
+    void *);
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Function-pointer for urUSMGetMemAllocInfo
 typedef ur_result_t(UR_APICALL *ur_pfnUSMGetMemAllocInfo_t)(
@@ -1543,6 +1549,7 @@ typedef struct ur_usm_dditable_t {
     ur_pfnUSMDeviceAlloc_t pfnDeviceAlloc;
     ur_pfnUSMSharedAlloc_t pfnSharedAlloc;
     ur_pfnUSMFree_t pfnFree;
+    ur_pfnUSMPoolFree_t pfnPoolFree;
     ur_pfnUSMGetMemAllocInfo_t pfnGetMemAllocInfo;
     ur_pfnUSMPoolCreate_t pfnPoolCreate;
     ur_pfnUSMPoolRetain_t pfnPoolRetain;

scripts/core/registry.yml

@@ -529,6 +529,9 @@ etors:
 - name: ADAPTER_GET_INFO
   desc: Enumerator for $xAdapterGetInfo
   value: '181'
+- name: USM_POOL_FREE
+  desc: Enumerator for $xUSMPoolFree
+  value: '182'
 ---
 type: enum
 desc: Defines structure types

scripts/core/usm.yml

@@ -148,7 +148,7 @@ members:
             Must be zero or a power of 2.
             Must be equal to or smaller than the size of the largest data type supported by `hDevice`.
 --- #--------------------------------------------------------------------------
-type: struct 
+type: struct
 desc: "USM host allocation descriptor type."
 details:
   - Specify these properties in $xUSMHostAlloc and $xUSMSharedAlloc via $x_usm_desc_t
@@ -228,10 +228,10 @@ params:
     - type: void**
       name: ppMem
       desc: "[out] pointer to USM host memory object"
-returns:      
+returns:
     - $X_RESULT_ERROR_INVALID_CONTEXT
     - $X_RESULT_ERROR_INVALID_OPERATION:
-      - "If $X_DEVICE_INFO_USM_HOST_SUPPORT is false."    
+      - "If $X_DEVICE_INFO_USM_HOST_SUPPORT is false."
     - $X_RESULT_ERROR_INVALID_VALUE:
       - "`pUSMDesc && pUSMDesc->align != 0 && ((pUSMDesc->align & (pUSMDesc->align-1)) != 0)`" # alignment must be power of two
       - "If `align` is greater that the size of the largest data type supported by `hDevice`."
@@ -272,7 +272,7 @@ params:
     - type: void**
       name: ppMem
       desc: "[out] pointer to USM device memory object"
-returns:      
+returns:
     - $X_RESULT_ERROR_INVALID_CONTEXT
     - $X_RESULT_ERROR_INVALID_OPERATION:
       - "If $X_DEVICE_INFO_USM_HOST_SUPPORT is false."
@@ -317,7 +317,7 @@ params:
     - type: void**
       name: ppMem
       desc: "[out] pointer to USM shared memory object"
-returns:      
+returns:
     - $X_RESULT_ERROR_INVALID_CONTEXT
     - $X_RESULT_ERROR_INVALID_VALUE:
       - "`pUSMDesc && pUSMDesc->align != 0 && ((pUSMDesc->align & (pUSMDesc->align-1)) != 0)`" # alignment must be power of two
@@ -347,6 +347,24 @@ returns:
     - $X_RESULT_ERROR_OUT_OF_HOST_MEMORY
 --- #--------------------------------------------------------------------------
 type: function
+desc: "Free the USM memory object from the pool"
+class: $xUSM
+name: PoolFree
+ordinal: "0"
+details:
+  - "Calling this function is equivalent to urUSMFree but without the pool look-up overhead."
+params:
+    - type: $x_usm_pool_handle_t
+      name: pool
+      desc: "[in] Pointer to a pool created using urUSMPoolCreate"
+    - type: void*
+      name: pMem
+      desc: "[in] pointer to USM memory object"
+returns:
+    - $X_RESULT_ERROR_INVALID_MEM_OBJECT
+    - $X_RESULT_ERROR_OUT_OF_HOST_MEMORY
+--- #--------------------------------------------------------------------------
+type: function
 desc: "Get USM memory object allocation information"
 class: $xUSM
 name: GetMemAllocInfo
@@ -436,7 +454,7 @@ etors:
     - name: REFERENCE_COUNT
       desc: |
             [uint32_t] Reference count of the pool object.
-            The reference count returned should be considered immediately stale. 
+            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.
     - name: CONTEXT
       desc: "[$x_context_handle_t] USM memory pool context info"

source/adapters/null/ur_nullddi.cpp

@@ -1317,6 +1317,28 @@ __urdlllocal ur_result_t UR_APICALL urUSMFree(
     return exceptionToResult(std::current_exception());
 }
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Intercept function for urUSMPoolFree
+__urdlllocal ur_result_t UR_APICALL urUSMPoolFree(
+    ur_usm_pool_handle_t
+        pool,  ///< [in] Pointer to a pool created using urUSMPoolCreate
+    void *pMem ///< [in] pointer to USM memory object
+    ) try {
+    ur_result_t result = UR_RESULT_SUCCESS;
+
+    // if the driver has created a custom function, then call it instead of using the generic path
+    auto pfnPoolFree = d_context.urDdiTable.USM.pfnPoolFree;
+    if (nullptr != pfnPoolFree) {
+        result = pfnPoolFree(pool, pMem);
+    } else {
+        // generic implementation
+    }
+
+    return result;
+} catch (...) {
+    return exceptionToResult(std::current_exception());
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Intercept function for urUSMGetMemAllocInfo
 __urdlllocal ur_result_t UR_APICALL urUSMGetMemAllocInfo(
@@ -5731,6 +5753,8 @@ UR_DLLEXPORT ur_result_t UR_APICALL urGetUSMProcAddrTable(
 
     pDdiTable->pfnFree = driver::urUSMFree;
 
+    pDdiTable->pfnPoolFree = driver::urUSMPoolFree;
+
     pDdiTable->pfnGetMemAllocInfo = driver::urUSMGetMemAllocInfo;
 
     pDdiTable->pfnPoolCreate = driver::urUSMPoolCreate;

source/common/ur_params.hpp

@@ -1112,6 +1112,10 @@ inline std::ostream &operator<<(std::ostream &os, enum ur_function_t value) {
     case UR_FUNCTION_ADAPTER_GET_INFO:
         os << "UR_FUNCTION_ADAPTER_GET_INFO";
         break;
+
+    case UR_FUNCTION_USM_POOL_FREE:
+        os << "UR_FUNCTION_USM_POOL_FREE";
+        break;
     default:
         os << "unknown enumerator";
         break;
@@ -14444,6 +14448,21 @@ inline std::ostream &operator<<(std::ostream &os,
     return os;
 }
 
+inline std::ostream &
+operator<<(std::ostream &os, const struct ur_usm_pool_free_params_t *params) {
+
+    os << ".pool = ";
+
+    ur_params::serializePtr(os, *(params->ppool));
+
+    os << ", ";
+    os << ".pMem = ";
+
+    ur_params::serializePtr(os, *(params->ppMem));
+
+    return os;
+}
+
 inline std::ostream &
 operator<<(std::ostream &os,
            const struct ur_usm_get_mem_alloc_info_params_t *params) {
@@ -15628,6 +15647,9 @@ inline int serializeFunctionParams(std::ostream &os, uint32_t function,
     case UR_FUNCTION_USM_FREE: {
         os << (const struct ur_usm_free_params_t *)params;
     } break;
+    case UR_FUNCTION_USM_POOL_FREE: {
+        os << (const struct ur_usm_pool_free_params_t *)params;
+    } break;
     case UR_FUNCTION_USM_GET_MEM_ALLOC_INFO: {
         os << (const struct ur_usm_get_mem_alloc_info_params_t *)params;
     } break;

source/loader/layers/tracing/ur_trcddi.cpp

@@ -1489,6 +1489,31 @@ __urdlllocal ur_result_t UR_APICALL urUSMFree(
     return result;
 }
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Intercept function for urUSMPoolFree
+__urdlllocal ur_result_t UR_APICALL urUSMPoolFree(
+    ur_usm_pool_handle_t
+        pool,  ///< [in] Pointer to a pool created using urUSMPoolCreate
+    void *pMem ///< [in] pointer to USM memory object
+) {
+    auto pfnPoolFree = context.urDdiTable.USM.pfnPoolFree;
+
+    if (nullptr == pfnPoolFree) {
+        return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+    }
+
+    ur_usm_pool_free_params_t params = {&pool, &pMem};
+    uint64_t instance = context.notify_begin(UR_FUNCTION_USM_POOL_FREE,
+                                             "urUSMPoolFree", &params);
+
+    ur_result_t result = pfnPoolFree(pool, pMem);
+
+    context.notify_end(UR_FUNCTION_USM_POOL_FREE, "urUSMPoolFree", &params,
+                       &result, instance);
+
+    return result;
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Intercept function for urUSMGetMemAllocInfo
 __urdlllocal ur_result_t UR_APICALL urUSMGetMemAllocInfo(
@@ -6703,6 +6728,9 @@ __urdlllocal ur_result_t UR_APICALL urGetUSMProcAddrTable(
     dditable.pfnFree = pDdiTable->pfnFree;
     pDdiTable->pfnFree = ur_tracing_layer::urUSMFree;
 
+    dditable.pfnPoolFree = pDdiTable->pfnPoolFree;
+    pDdiTable->pfnPoolFree = ur_tracing_layer::urUSMPoolFree;
+
     dditable.pfnGetMemAllocInfo = pDdiTable->pfnGetMemAllocInfo;
     pDdiTable->pfnGetMemAllocInfo = ur_tracing_layer::urUSMGetMemAllocInfo;
 

source/loader/layers/validation/ur_valddi.cpp

@@ -1797,6 +1797,34 @@ __urdlllocal ur_result_t UR_APICALL urUSMFree(
     return result;
 }
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Intercept function for urUSMPoolFree
+__urdlllocal ur_result_t UR_APICALL urUSMPoolFree(
+    ur_usm_pool_handle_t
+        pool,  ///< [in] Pointer to a pool created using urUSMPoolCreate
+    void *pMem ///< [in] pointer to USM memory object
+) {
+    auto pfnPoolFree = context.urDdiTable.USM.pfnPoolFree;
+
+    if (nullptr == pfnPoolFree) {
+        return UR_RESULT_ERROR_UNINITIALIZED;
+    }
+
+    if (context.enableParameterValidation) {
+        if (NULL == pool) {
+            return UR_RESULT_ERROR_INVALID_NULL_HANDLE;
+        }
+
+        if (NULL == pMem) {
+            return UR_RESULT_ERROR_INVALID_NULL_POINTER;
+        }
+    }
+
+    ur_result_t result = pfnPoolFree(pool, pMem);
+
+    return result;
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Intercept function for urUSMGetMemAllocInfo
 __urdlllocal ur_result_t UR_APICALL urUSMGetMemAllocInfo(
@@ -8081,6 +8109,9 @@ UR_DLLEXPORT ur_result_t UR_APICALL urGetUSMProcAddrTable(
     dditable.pfnFree = pDdiTable->pfnFree;
     pDdiTable->pfnFree = ur_validation_layer::urUSMFree;
 
+    dditable.pfnPoolFree = pDdiTable->pfnPoolFree;
+    pDdiTable->pfnPoolFree = ur_validation_layer::urUSMPoolFree;
+
     dditable.pfnGetMemAllocInfo = pDdiTable->pfnGetMemAllocInfo;
     pDdiTable->pfnGetMemAllocInfo = ur_validation_layer::urUSMGetMemAllocInfo;