[OpenCL] Add usm_alloc_location_desc struct and handle it in the adap…

…ter. This is the equivalent of buffer_alloc_location_properties_t, it turns out the CL extension flag is applicable to both buffer creation and USM allocations.
oneapi-src · Nov 27, 2023 · 8d4dbcf · 8d4dbcf
1 parent 41228d3
commit 8d4dbcf
Show file tree

Hide file tree

Showing 8 changed files with 193 additions and 70 deletions.
diff --git a/include/ur.py b/include/ur.py
@@ -242,6 +242,7 @@ class ur_structure_type_v(IntEnum):
     KERNEL_EXEC_INFO_PROPERTIES = 31                ## ::ur_kernel_exec_info_properties_t
     KERNEL_ARG_VALUE_PROPERTIES = 32                ## ::ur_kernel_arg_value_properties_t
     KERNEL_ARG_LOCAL_PROPERTIES = 33                ## ::ur_kernel_arg_local_properties_t
+    USM_ALLOC_LOCATION_DESC = 35                    ## ::ur_usm_alloc_location_desc_t
     EXP_COMMAND_BUFFER_DESC = 0x1000                ## ::ur_exp_command_buffer_desc_t
     EXP_SAMPLER_MIP_PROPERTIES = 0x2000             ## ::ur_exp_sampler_mip_properties_t
     EXP_INTEROP_MEM_DESC = 0x2001                   ## ::ur_exp_interop_mem_desc_t
@@ -1531,6 +1532,25 @@ class ur_usm_device_desc_t(Structure):
         ("flags", ur_usm_device_mem_flags_t)                            ## [in] device memory allocation flags.
     ]
 
+###############################################################################
+## @brief USM allocation location desc
+## 
+## @details
+##     - Specify these properties in ::urUSMHostAlloc, ::urUSMDeviceAlloc and
+##       ::urUSMSharedAlloc via ::ur_usm_desc_t as part of a `pNext` chain.
+## 
+## @remarks
+##   _Analogues_
+##     - cl_intel_mem_alloc_buffer_location
+class ur_usm_alloc_location_desc_t(Structure):
+    _fields_ = [
+        ("stype", ur_structure_type_t),                                 ## [in] type of this structure, must be
+                                                                        ## ::UR_STRUCTURE_TYPE_USM_ALLOC_LOCATION_DESC
+        ("pNext", c_void_p),                                            ## [in][optional] pointer to extension-specific structure
+        ("location", c_ulong)                                           ## [in] Identifies the ID of global memory partition to which the memory
+                                                                        ## should be allocated.
+    ]
+
 ###############################################################################
 ## @brief USM pool descriptor type
 class ur_usm_pool_desc_t(Structure):

diff --git a/include/ur_api.h b/include/ur_api.h
@@ -251,6 +251,7 @@ typedef enum ur_structure_type_t {
     UR_STRUCTURE_TYPE_KERNEL_EXEC_INFO_PROPERTIES = 31,     ///< ::ur_kernel_exec_info_properties_t
     UR_STRUCTURE_TYPE_KERNEL_ARG_VALUE_PROPERTIES = 32,     ///< ::ur_kernel_arg_value_properties_t
     UR_STRUCTURE_TYPE_KERNEL_ARG_LOCAL_PROPERTIES = 33,     ///< ::ur_kernel_arg_local_properties_t
+    UR_STRUCTURE_TYPE_USM_ALLOC_LOCATION_DESC = 35,         ///< ::ur_usm_alloc_location_desc_t
     UR_STRUCTURE_TYPE_EXP_COMMAND_BUFFER_DESC = 0x1000,     ///< ::ur_exp_command_buffer_desc_t
     UR_STRUCTURE_TYPE_EXP_SAMPLER_MIP_PROPERTIES = 0x2000,  ///< ::ur_exp_sampler_mip_properties_t
     UR_STRUCTURE_TYPE_EXP_INTEROP_MEM_DESC = 0x2001,        ///< ::ur_exp_interop_mem_desc_t
@@ -3221,6 +3222,25 @@ typedef struct ur_usm_device_desc_t {
 
 } ur_usm_device_desc_t;
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief USM allocation location desc
+///
+/// @details
+///     - Specify these properties in ::urUSMHostAlloc, ::urUSMDeviceAlloc and
+///       ::urUSMSharedAlloc via ::ur_usm_desc_t as part of a `pNext` chain.
+///
+/// @remarks
+///   _Analogues_
+///     - cl_intel_mem_alloc_buffer_location
+typedef struct ur_usm_alloc_location_desc_t {
+    ur_structure_type_t stype; ///< [in] type of this structure, must be
+                               ///< ::UR_STRUCTURE_TYPE_USM_ALLOC_LOCATION_DESC
+    const void *pNext;         ///< [in][optional] pointer to extension-specific structure
+    uint32_t location;         ///< [in] Identifies the ID of global memory partition to which the memory
+                               ///< should be allocated.
+
+} ur_usm_alloc_location_desc_t;
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief USM pool descriptor type
 typedef struct ur_usm_pool_desc_t {
@@ -3258,6 +3278,7 @@ typedef struct ur_usm_pool_limits_desc_t {
 ///     - Any flags/hints passed through pUSMDesc only affect the single
 ///       allocation.
 ///     - See also ::ur_usm_host_desc_t.
+///     - See also ::ur_usm_alloc_location_desc_t.
 ///
 /// @returns
 ///     - ::UR_RESULT_SUCCESS
@@ -3301,6 +3322,7 @@ urUSMHostAlloc(
 ///     - Any flags/hints passed through pUSMDesc only affect the single
 ///       allocation.
 ///     - See also ::ur_usm_device_desc_t.
+///     - See also ::ur_usm_alloc_location_desc_t.
 ///
 /// @returns
 ///     - ::UR_RESULT_SUCCESS
@@ -3347,6 +3369,7 @@ urUSMDeviceAlloc(
 ///       allocation.
 ///     - See also ::ur_usm_host_desc_t.
 ///     - See also ::ur_usm_device_desc_t.
+///     - See also ::ur_usm_alloc_location_desc_t.
 ///
 /// @returns
 ///     - ::UR_RESULT_SUCCESS

diff --git a/scripts/core/registry.yml b/scripts/core/registry.yml
@@ -645,3 +645,6 @@ etors:
 - name: KERNEL_ARG_LOCAL_PROPERTIES
   desc: $x_kernel_arg_local_properties_t
   value: '33'
+- name: USM_ALLOC_LOCATION_DESC
+  desc: $x_usm_alloc_location_desc_t
+  value: '35'
diff --git a/scripts/core/usm.yml b/scripts/core/usm.yml
@@ -175,6 +175,23 @@ members:
       desc: "[in] device memory allocation flags."
 --- #--------------------------------------------------------------------------
 type: struct
+desc: "USM allocation location desc"
+details:
+    - Specify these properties in $xUSMHostAlloc, $xUSMDeviceAlloc and
+      $xUSMSharedAlloc via $x_usm_desc_t as part of a `pNext` chain.
+analogue:
+    - "cl_intel_mem_alloc_buffer_location"
+class: $xUSM
+name: $x_usm_alloc_location_desc_t
+base: $x_base_desc_t
+members:
+    - type: uint32_t
+      name: location
+      desc: >
+          [in] Identifies the ID of global memory partition to which the memory
+          should be allocated.
+--- #--------------------------------------------------------------------------
+type: struct
 desc: "USM pool descriptor type"
 class: $xUSM
 name: $x_usm_pool_desc_t
@@ -212,6 +229,7 @@ details:
   - "Allocations served from different memory pools must be isolated and must not reside on the same page."
   - "Any flags/hints passed through pUSMDesc only affect the single allocation."
   - "See also $x_usm_host_desc_t."
+  - "See also $x_usm_alloc_location_desc_t."
 params:
     - type: $x_context_handle_t
       name: hContext
@@ -253,6 +271,7 @@ details:
   - "Allocations served from different memory pools must be isolated and must not reside on the same page."
   - "Any flags/hints passed through pUSMDesc only affect the single allocation."
   - "See also $x_usm_device_desc_t."
+  - "See also $x_usm_alloc_location_desc_t."
 params:
     - type: $x_context_handle_t
       name: hContext
@@ -298,6 +317,7 @@ details:
   - "Any flags/hints passed through pUSMDesc only affect the single allocation."
   - "See also $x_usm_host_desc_t."
   - "See also $x_usm_device_desc_t."
+  - "See also $x_usm_alloc_location_desc_t."
 params:
     - type: $x_context_handle_t
       name: hContext

diff --git a/source/adapters/opencl/usm.cpp b/source/adapters/opencl/usm.cpp
@@ -10,30 +10,86 @@
 
 #include "common.hpp"
 
+inline cl_mem_alloc_flags_intel
+hostDescToClFlags(const ur_usm_host_desc_t &desc) {
+  cl_mem_alloc_flags_intel allocFlags = 0;
+  if (desc.flags & UR_USM_HOST_MEM_FLAG_INITIAL_PLACEMENT) {
+    allocFlags |= CL_MEM_ALLOC_INITIAL_PLACEMENT_HOST_INTEL;
+  }
+  return allocFlags;
+}
+
+inline cl_mem_alloc_flags_intel
+deviceDescToClFlags(const ur_usm_device_desc_t &desc) {
+  cl_mem_alloc_flags_intel allocFlags = 0;
+  if (desc.flags & UR_USM_DEVICE_MEM_FLAG_INITIAL_PLACEMENT) {
+    allocFlags |= CL_MEM_ALLOC_INITIAL_PLACEMENT_DEVICE_INTEL;
+  }
+  if (desc.flags & UR_USM_DEVICE_MEM_FLAG_WRITE_COMBINED) {
+    allocFlags |= CL_MEM_ALLOC_WRITE_COMBINED_INTEL;
+  }
+  return allocFlags;
+}
+
+ur_result_t
+usmDescToCLMemProperties(const ur_base_desc_t *Desc,
+                         std::vector<cl_mem_properties_intel> &Properties) {
+  cl_mem_alloc_flags_intel AllocFlags = 0;
+  const auto *Next = Desc;
+  do {
+    switch (Next->stype) {
+    case UR_STRUCTURE_TYPE_USM_HOST_DESC: {
+      auto HostDesc = reinterpret_cast<const ur_usm_host_desc_t *>(Next);
+      if (UR_USM_HOST_MEM_FLAGS_MASK & HostDesc->flags) {
+        return UR_RESULT_ERROR_INVALID_ENUMERATION;
+      }
+      AllocFlags |= hostDescToClFlags(*HostDesc);
+      break;
+    }
+    case UR_STRUCTURE_TYPE_USM_DEVICE_DESC: {
+      auto DeviceDesc = reinterpret_cast<const ur_usm_device_desc_t *>(Next);
+      if (UR_USM_HOST_MEM_FLAGS_MASK & DeviceDesc->flags) {
+        return UR_RESULT_ERROR_INVALID_ENUMERATION;
+      }
+      AllocFlags |= deviceDescToClFlags(*DeviceDesc);
+      break;
+    }
+    case UR_STRUCTURE_TYPE_USM_ALLOC_LOCATION_DESC: {
+      auto LocationDesc =
+          reinterpret_cast<const ur_usm_alloc_location_desc_t *>(Next);
+      Properties.push_back(CL_MEM_ALLOC_BUFFER_LOCATION_INTEL);
+      // CL bitfields are cl_ulong
+      Properties.push_back(static_cast<cl_ulong>(LocationDesc->location));
+      break;
+    }
+    default:
+      return UR_RESULT_ERROR_INVALID_VALUE;
+    }
+
+    Next = Next->pNext ? static_cast<const ur_base_desc_t *>(Next->pNext)
+                       : nullptr;
+  } while (Next);
+
+  if (AllocFlags) {
+    Properties.push_back(CL_MEM_ALLOC_FLAGS_INTEL);
+    Properties.push_back(AllocFlags);
+  }
+  Properties.push_back(0);
+
+  return UR_RESULT_SUCCESS;
+}
+
 UR_APIEXPORT ur_result_t UR_APICALL
 urUSMHostAlloc(ur_context_handle_t hContext, const ur_usm_desc_t *pUSMDesc,
                ur_usm_pool_handle_t, size_t size, void **ppMem) {
 
   void *Ptr = nullptr;
   uint32_t Alignment = pUSMDesc ? pUSMDesc->align : 0;
 
-  cl_mem_alloc_flags_intel Flags = 0;
-  cl_mem_properties_intel Properties[3];
-
-  if (pUSMDesc && pUSMDesc->pNext &&
-      static_cast<const ur_base_desc_t *>(pUSMDesc->pNext)->stype ==
-          UR_STRUCTURE_TYPE_USM_HOST_DESC) {
-    const auto *HostDesc =
-        static_cast<const ur_usm_host_desc_t *>(pUSMDesc->pNext);
-
-    if (HostDesc->flags & UR_USM_HOST_MEM_FLAG_INITIAL_PLACEMENT) {
-      Flags |= CL_MEM_ALLOC_INITIAL_PLACEMENT_HOST_INTEL;
-    }
-    Properties[0] = CL_MEM_ALLOC_FLAGS_INTEL;
-    Properties[1] = Flags;
-    Properties[2] = 0;
-  } else {
-    Properties[0] = 0;
+  std::vector<cl_mem_properties_intel> AllocProperties;
+  if (pUSMDesc && pUSMDesc->pNext) {
+    UR_RETURN_ON_FAILURE(usmDescToCLMemProperties(
+        static_cast<const ur_base_desc_t *>(pUSMDesc->pNext), AllocProperties));
   }
 
   // First we need to look up the function pointer
@@ -47,7 +103,9 @@ urUSMHostAlloc(ur_context_handle_t hContext, const ur_usm_desc_t *pUSMDesc,
 
   if (FuncPtr) {
     cl_int ClResult = CL_SUCCESS;
-    Ptr = FuncPtr(CLContext, Properties, size, Alignment, &ClResult);
+    Ptr = FuncPtr(CLContext,
+                  AllocProperties.empty() ? nullptr : AllocProperties.data(),
+                  size, Alignment, &ClResult);
     if (ClResult == CL_INVALID_BUFFER_SIZE) {
       return UR_RESULT_ERROR_INVALID_USM_SIZE;
     }
@@ -71,25 +129,10 @@ urUSMDeviceAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
   void *Ptr = nullptr;
   uint32_t Alignment = pUSMDesc ? pUSMDesc->align : 0;
 
-  cl_mem_alloc_flags_intel Flags = 0;
-  cl_mem_properties_intel Properties[3];
-  if (pUSMDesc && pUSMDesc->pNext &&
-      static_cast<const ur_base_desc_t *>(pUSMDesc->pNext)->stype ==
-          UR_STRUCTURE_TYPE_USM_DEVICE_DESC) {
-    const auto *HostDesc =
-        static_cast<const ur_usm_device_desc_t *>(pUSMDesc->pNext);
-
-    if (HostDesc->flags & UR_USM_DEVICE_MEM_FLAG_INITIAL_PLACEMENT) {
-      Flags |= CL_MEM_ALLOC_INITIAL_PLACEMENT_DEVICE_INTEL;
-    }
-    if (HostDesc->flags & UR_USM_DEVICE_MEM_FLAG_WRITE_COMBINED) {
-      Flags |= CL_MEM_ALLOC_WRITE_COMBINED_INTEL;
-    }
-    Properties[0] = CL_MEM_ALLOC_FLAGS_INTEL;
-    Properties[1] = Flags;
-    Properties[2] = 0;
-  } else {
-    Properties[0] = 0;
+  std::vector<cl_mem_properties_intel> AllocProperties;
+  if (pUSMDesc && pUSMDesc->pNext) {
+    UR_RETURN_ON_FAILURE(usmDescToCLMemProperties(
+        static_cast<const ur_base_desc_t *>(pUSMDesc->pNext), AllocProperties));
   }
 
   // First we need to look up the function pointer
@@ -104,8 +147,8 @@ urUSMDeviceAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
   if (FuncPtr) {
     cl_int ClResult = CL_SUCCESS;
     Ptr = FuncPtr(CLContext, cl_adapter::cast<cl_device_id>(hDevice),
-                  cl_adapter::cast<cl_mem_properties_intel *>(Properties), size,
-                  Alignment, &ClResult);
+                  AllocProperties.empty() ? nullptr : AllocProperties.data(),
+                  size, Alignment, &ClResult);
     if (ClResult == CL_INVALID_BUFFER_SIZE) {
       return UR_RESULT_ERROR_INVALID_USM_SIZE;
     }
@@ -129,35 +172,10 @@ urUSMSharedAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
   void *Ptr = nullptr;
   uint32_t Alignment = pUSMDesc ? pUSMDesc->align : 0;
 
-  cl_mem_alloc_flags_intel Flags = 0;
-  const auto *NextStruct =
-      (pUSMDesc ? static_cast<const ur_base_desc_t *>(pUSMDesc->pNext)
-                : nullptr);
-  while (NextStruct) {
-    if (NextStruct->stype == UR_STRUCTURE_TYPE_USM_HOST_DESC) {
-      const auto *HostDesc =
-          reinterpret_cast<const ur_usm_host_desc_t *>(NextStruct);
-      if (HostDesc->flags & UR_USM_HOST_MEM_FLAG_INITIAL_PLACEMENT) {
-        Flags |= CL_MEM_ALLOC_INITIAL_PLACEMENT_HOST_INTEL;
-      }
-    } else if (NextStruct->stype == UR_STRUCTURE_TYPE_USM_DEVICE_DESC) {
-      const auto *DevDesc =
-          reinterpret_cast<const ur_usm_device_desc_t *>(NextStruct);
-      if (DevDesc->flags & UR_USM_DEVICE_MEM_FLAG_INITIAL_PLACEMENT) {
-        Flags |= CL_MEM_ALLOC_INITIAL_PLACEMENT_DEVICE_INTEL;
-      }
-      if (DevDesc->flags & UR_USM_DEVICE_MEM_FLAG_WRITE_COMBINED) {
-        Flags |= CL_MEM_ALLOC_WRITE_COMBINED_INTEL;
-      }
-    }
-    NextStruct = static_cast<const ur_base_desc_t *>(NextStruct->pNext);
-  }
-
-  cl_mem_properties_intel Properties[3] = {CL_MEM_ALLOC_FLAGS_INTEL, Flags, 0};
-
-  // Passing a flags value of 0 doesn't work, so truncate the properties
-  if (Flags == 0) {
-    Properties[0] = 0;
+  std::vector<cl_mem_properties_intel> AllocProperties;
+  if (pUSMDesc && pUSMDesc->pNext) {
+    UR_RETURN_ON_FAILURE(usmDescToCLMemProperties(
+        static_cast<const ur_base_desc_t *>(pUSMDesc->pNext), AllocProperties));
   }
 
   // First we need to look up the function pointer
@@ -172,8 +190,8 @@ urUSMSharedAlloc(ur_context_handle_t hContext, ur_device_handle_t hDevice,
   if (FuncPtr) {
     cl_int ClResult = CL_SUCCESS;
     Ptr = FuncPtr(CLContext, cl_adapter::cast<cl_device_id>(hDevice),
-                  cl_adapter::cast<cl_mem_properties_intel *>(Properties), size,
-                  Alignment, cl_adapter::cast<cl_int *>(&ClResult));
+                  AllocProperties.empty() ? nullptr : AllocProperties.data(),
+                  size, Alignment, cl_adapter::cast<cl_int *>(&ClResult));
     if (ClResult == CL_INVALID_BUFFER_SIZE) {
       return UR_RESULT_ERROR_INVALID_USM_SIZE;
     }