Add ur_usm_pool_flags_t entry

Add an entry so the user can specify if the native USM pool should be used.
kswiecicki · Nov 6, 2024 · 69cee9d · 69cee9d
1 parent af9b768
commit 69cee9d
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 3 deletions.
diff --git a/include/ur_api.h b/include/ur_api.h
@@ -3355,15 +3355,16 @@ typedef enum ur_usm_device_mem_flag_t {
 /// @brief USM memory property flags
 typedef uint32_t ur_usm_pool_flags_t;
 typedef enum ur_usm_pool_flag_t {
-    UR_USM_POOL_FLAG_ZERO_INITIALIZE_BLOCK = UR_BIT(0), ///< All coarse-grain allocations (allocations from the driver) will be
-                                                        ///< zero-initialized.
+    UR_USM_POOL_FLAG_ZERO_INITIALIZE_BLOCK = UR_BIT(0),      ///< All coarse-grain allocations (allocations from the driver) will be
+                                                             ///< zero-initialized.
+    UR_USM_POOL_FLAG_USE_NATIVE_MEMORY_POOL_EXP = UR_BIT(1), ///< Use the native memory pool API
     /// @cond
     UR_USM_POOL_FLAG_FORCE_UINT32 = 0x7fffffff
     /// @endcond
 
 } ur_usm_pool_flag_t;
 /// @brief Bit Mask for validating ur_usm_pool_flags_t
-#define UR_USM_POOL_FLAGS_MASK 0xfffffffe
+#define UR_USM_POOL_FLAGS_MASK 0xfffffffc
 
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief USM allocation type

diff --git a/include/ur_print.hpp b/include/ur_print.hpp
@@ -6535,6 +6535,9 @@ inline std::ostream &operator<<(std::ostream &os, enum ur_usm_pool_flag_t value)
     case UR_USM_POOL_FLAG_ZERO_INITIALIZE_BLOCK:
         os << "UR_USM_POOL_FLAG_ZERO_INITIALIZE_BLOCK";
         break;
+    case UR_USM_POOL_FLAG_USE_NATIVE_MEMORY_POOL_EXP:
+        os << "UR_USM_POOL_FLAG_USE_NATIVE_MEMORY_POOL_EXP";
+        break;
     default:
         os << "unknown enumerator";
         break;
@@ -6559,6 +6562,16 @@ inline ur_result_t printFlag<ur_usm_pool_flag_t>(std::ostream &os, uint32_t flag
         }
         os << UR_USM_POOL_FLAG_ZERO_INITIALIZE_BLOCK;
     }
+
+    if ((val & UR_USM_POOL_FLAG_USE_NATIVE_MEMORY_POOL_EXP) == (uint32_t)UR_USM_POOL_FLAG_USE_NATIVE_MEMORY_POOL_EXP) {
+        val ^= (uint32_t)UR_USM_POOL_FLAG_USE_NATIVE_MEMORY_POOL_EXP;
+        if (!first) {
+            os << " | ";
+        } else {
+            first = false;
+        }
+        os << UR_USM_POOL_FLAG_USE_NATIVE_MEMORY_POOL_EXP;
+    }
     if (val != 0) {
         std::bitset<32> bits(val);
         if (!first) {

diff --git a/scripts/core/exp-async-alloc.yml b/scripts/core/exp-async-alloc.yml
@@ -23,6 +23,16 @@ etors:
       value: "0x2021"
       desc: "[$x_bool_t] returns true if the device supports USM allocation enqueueing"
 
+--- #--------------------------------------------------------------------------
+type: enum
+extend: true
+name: $x_usm_pool_flags_t
+desc: "USM memory property flag for creating a native USM pool"
+etors:
+    - name: USE_NATIVE_MEMORY_POOL_EXP
+      value: "$X_BIT(1)"
+      desc: "Use the native memory pool API"
+
 --- #--------------------------------------------------------------------------
 type: enum
 extend: true