[L0 v2] implement urMemBufferPartition

Change access_mode_t to device_access_mode_t.
Previously it was use to specify both host and device
access modes which was confusing. Now, it only
specifies device access mode.

Author: igchor

source/adapters/level_zero/v2/kernel.cpp

@@ -394,21 +394,21 @@ ur_result_t urKernelSetArgPointer(
   return hKernel->setArgPointer(argIndex, pProperties, pArgValue);
 }
 
-static ur_mem_handle_t_::access_mode_t memAccessFromKernelProperties(
+static ur_mem_handle_t_::device_access_mode_t memAccessFromKernelProperties(
     const ur_kernel_arg_mem_obj_properties_t *pProperties) {
   if (pProperties) {
     switch (pProperties->memoryAccess) {
     case UR_MEM_FLAG_READ_WRITE:
-      return ur_mem_handle_t_::access_mode_t::read_write;
+      return ur_mem_handle_t_::device_access_mode_t::read_write;
     case UR_MEM_FLAG_WRITE_ONLY:
-      return ur_mem_handle_t_::access_mode_t::write_only;
+      return ur_mem_handle_t_::device_access_mode_t::write_only;
     case UR_MEM_FLAG_READ_ONLY:
-      return ur_mem_handle_t_::access_mode_t::read_only;
+      return ur_mem_handle_t_::device_access_mode_t::read_only;
     default:
-      return ur_mem_handle_t_::access_mode_t::read_write;
+      return ur_mem_handle_t_::device_access_mode_t::read_write;
     }
   }
-  return ur_mem_handle_t_::access_mode_t::read_write;
+  return ur_mem_handle_t_::device_access_mode_t::read_write;
 }
 
 ur_result_t

source/adapters/level_zero/v2/kernel.hpp

@@ -105,7 +105,7 @@ struct ur_kernel_handle_t_ : _ur_object {
 
   struct pending_memory_allocation_t {
     ur_mem_handle_t hMem;
-    ur_mem_handle_t_::access_mode_t mode;
+    ur_mem_handle_t_::device_access_mode_t mode;
     uint32_t argIndex;
   };
 

source/adapters/level_zero/v2/memory.cpp

@@ -13,17 +13,42 @@
 
 #include "../helpers/memory_helpers.hpp"
 
-ur_mem_handle_t_::ur_mem_handle_t_(ur_context_handle_t hContext, size_t size)
-    : hContext(hContext), size(size) {}
+static ur_mem_handle_t_::device_access_mode_t
+getDeviceAccessMode(ur_mem_flags_t memFlag) {
+  if (memFlag & UR_MEM_FLAG_READ_WRITE) {
+    return ur_mem_handle_t_::device_access_mode_t::read_write;
+  } else if (memFlag & UR_MEM_FLAG_READ_ONLY) {
+    return ur_mem_handle_t_::device_access_mode_t::read_only;
+  } else if (memFlag & UR_MEM_FLAG_WRITE_ONLY) {
+    return ur_mem_handle_t_::device_access_mode_t::write_only;
+  } else {
+    return ur_mem_handle_t_::device_access_mode_t::read_write;
+  }
+}
+
+static bool isAccessCompatible(ur_mem_handle_t_::device_access_mode_t requested,
+                               ur_mem_handle_t_::device_access_mode_t actual) {
+  return requested == actual ||
+         actual == ur_mem_handle_t_::device_access_mode_t::read_write;
+}
+
+ur_mem_handle_t_::ur_mem_handle_t_(ur_context_handle_t hContext, size_t size,
+                                   device_access_mode_t accessMode)
+    : accessMode(accessMode), hContext(hContext), size(size) {}
+
+size_t ur_mem_handle_t_::getSize() const { return size; }
+
+ur_shared_mutex &ur_mem_handle_t_::getMutex() { return Mutex; }
 
 ur_usm_handle_t_::ur_usm_handle_t_(ur_context_handle_t hContext, size_t size,
                                    const void *ptr)
-    : ur_mem_handle_t_(hContext, size), ptr(const_cast<void *>(ptr)) {}
+    : ur_mem_handle_t_(hContext, size, device_access_mode_t::read_write),
+      ptr(const_cast<void *>(ptr)) {}
 
 ur_usm_handle_t_::~ur_usm_handle_t_() {}
 
 void *ur_usm_handle_t_::getDevicePtr(
-    ur_device_handle_t hDevice, access_mode_t access, size_t offset,
+    ur_device_handle_t hDevice, device_access_mode_t access, size_t offset,
     size_t size, std::function<void(void *src, void *dst, size_t)> migrate) {
   std::ignore = hDevice;
   std::ignore = access;
@@ -34,9 +59,9 @@ void *ur_usm_handle_t_::getDevicePtr(
 }
 
 void *ur_usm_handle_t_::mapHostPtr(
-    access_mode_t access, size_t offset, size_t size,
+    ur_map_flags_t flags, size_t offset, size_t size,
     std::function<void(void *src, void *dst, size_t)>) {
-  std::ignore = access;
+  std::ignore = flags;
   std::ignore = offset;
   std::ignore = size;
   return ptr;
@@ -50,8 +75,8 @@ void ur_usm_handle_t_::unmapHostPtr(
 
 ur_integrated_mem_handle_t::ur_integrated_mem_handle_t(
     ur_context_handle_t hContext, void *hostPtr, size_t size,
-    host_ptr_action_t hostPtrAction)
-    : ur_mem_handle_t_(hContext, size) {
+    host_ptr_action_t hostPtrAction, device_access_mode_t accessMode)
+    : ur_mem_handle_t_(hContext, size, accessMode) {
   bool hostPtrImported = false;
   if (hostPtrAction == host_ptr_action_t::import) {
     hostPtrImported =
@@ -83,8 +108,9 @@ ur_integrated_mem_handle_t::ur_integrated_mem_handle_t(
 }
 
 ur_integrated_mem_handle_t::ur_integrated_mem_handle_t(
-    ur_context_handle_t hContext, void *hostPtr, size_t size, bool ownHostPtr)
-    : ur_mem_handle_t_(hContext, size) {
+    ur_context_handle_t hContext, void *hostPtr, size_t size,
+    device_access_mode_t accessMode, bool ownHostPtr)
+    : ur_mem_handle_t_(hContext, size, accessMode) {
   this->ptr = usm_unique_ptr_t(hostPtr, [hContext, ownHostPtr](void *ptr) {
     if (!ownHostPtr) {
       return;
@@ -97,7 +123,7 @@ ur_integrated_mem_handle_t::ur_integrated_mem_handle_t(
 }
 
 void *ur_integrated_mem_handle_t::getDevicePtr(
-    ur_device_handle_t hDevice, access_mode_t access, size_t offset,
+    ur_device_handle_t hDevice, device_access_mode_t access, size_t offset,
     size_t size, std::function<void(void *src, void *dst, size_t)> migrate) {
   std::ignore = hDevice;
   std::ignore = access;
@@ -108,9 +134,9 @@ void *ur_integrated_mem_handle_t::getDevicePtr(
 }
 
 void *ur_integrated_mem_handle_t::mapHostPtr(
-    access_mode_t access, size_t offset, size_t size,
+    ur_map_flags_t flags, size_t offset, size_t size,
     std::function<void(void *src, void *dst, size_t)> migrate) {
-  std::ignore = access;
+  std::ignore = flags;
   std::ignore = offset;
   std::ignore = size;
   std::ignore = migrate;
@@ -178,9 +204,10 @@ ur_discrete_mem_handle_t::migrateBufferTo(ur_device_handle_t hDevice, void *src,
   return UR_RESULT_SUCCESS;
 }
 
-ur_discrete_mem_handle_t::ur_discrete_mem_handle_t(ur_context_handle_t hContext,
-                                                   void *hostPtr, size_t size)
-    : ur_mem_handle_t_(hContext, size),
+ur_discrete_mem_handle_t::ur_discrete_mem_handle_t(
+    ur_context_handle_t hContext, void *hostPtr, size_t size,
+    device_access_mode_t accessMode)
+    : ur_mem_handle_t_(hContext, size, accessMode),
       deviceAllocations(hContext->getPlatform()->getNumDevices()),
       activeAllocationDevice(nullptr), hostAllocations() {
   if (hostPtr) {
@@ -189,12 +216,11 @@ ur_discrete_mem_handle_t::ur_discrete_mem_handle_t(ur_context_handle_t hContext,
   }
 }
 
-ur_discrete_mem_handle_t::ur_discrete_mem_handle_t(ur_context_handle_t hContext,
-                                                   ur_device_handle_t hDevice,
-                                                   void *devicePtr, size_t size,
-                                                   void *writeBackMemory,
-                                                   bool ownZePtr)
-    : ur_mem_handle_t_(hContext, size),
+ur_discrete_mem_handle_t::ur_discrete_mem_handle_t(
+    ur_context_handle_t hContext, ur_device_handle_t hDevice, void *devicePtr,
+    size_t size, device_access_mode_t accessMode, void *writeBackMemory,
+    bool ownZePtr)
+    : ur_mem_handle_t_(hContext, size, accessMode),
       deviceAllocations(hContext->getPlatform()->getNumDevices()),
       activeAllocationDevice(hDevice), writeBackPtr(writeBackMemory),
       hostAllocations() {
@@ -227,7 +253,7 @@ ur_discrete_mem_handle_t::~ur_discrete_mem_handle_t() {
 }
 
 void *ur_discrete_mem_handle_t::getDevicePtr(
-    ur_device_handle_t hDevice, access_mode_t access, size_t offset,
+    ur_device_handle_t hDevice, device_access_mode_t access, size_t offset,
     size_t size, std::function<void(void *src, void *dst, size_t)> migrate) {
   TRACK_SCOPE_LATENCY("ur_discrete_mem_handle_t::getDevicePtr");
 
@@ -265,19 +291,18 @@ void *ur_discrete_mem_handle_t::getDevicePtr(
 }
 
 void *ur_discrete_mem_handle_t::mapHostPtr(
-    access_mode_t access, size_t offset, size_t size,
+    ur_map_flags_t flags, size_t offset, size_t size,
     std::function<void(void *src, void *dst, size_t)> migrate) {
   TRACK_SCOPE_LATENCY("ur_discrete_mem_handle_t::mapHostPtr");
-
   // TODO: use async alloc?
 
   void *ptr;
   UR_CALL_THROWS(hContext->getDefaultUSMPool()->allocate(
       hContext, nullptr, nullptr, UR_USM_TYPE_HOST, size, &ptr));
 
-  hostAllocations.emplace_back(ptr, size, offset, access);
+  hostAllocations.emplace_back(ptr, size, offset, flags);
 
-  if (activeAllocationDevice && access != access_mode_t::write_only) {
+  if (activeAllocationDevice && (flags & UR_MAP_FLAG_READ)) {
     auto srcPtr =
         ur_cast<char *>(
             deviceAllocations[activeAllocationDevice->Id.value()].get()) +
@@ -301,10 +326,11 @@ void ur_discrete_mem_handle_t::unmapHostPtr(
             ur_cast<char *>(
                 deviceAllocations[activeAllocationDevice->Id.value()].get()) +
             hostAllocation.offset;
-      } else if (hostAllocation.access != access_mode_t::write_invalidate) {
-        devicePtr = ur_cast<char *>(
-            getDevicePtr(hContext->getDevices()[0], access_mode_t::read_only,
-                         hostAllocation.offset, hostAllocation.size, migrate));
+      } else if (!(hostAllocation.flags &
+                   UR_MAP_FLAG_WRITE_INVALIDATE_REGION)) {
+        devicePtr = ur_cast<char *>(getDevicePtr(
+            hContext->getDevices()[0], device_access_mode_t::read_only,
+            hostAllocation.offset, hostAllocation.size, migrate));
       }
 
       if (devicePtr) {
@@ -332,6 +358,46 @@ static bool useHostBuffer(ur_context_handle_t hContext) {
              ZE_DEVICE_PROPERTY_FLAG_INTEGRATED;
 }
 
+namespace ur::level_zero {
+ur_result_t urMemRetain(ur_mem_handle_t hMem);
+ur_result_t urMemRelease(ur_mem_handle_t hMem);
+} // namespace ur::level_zero
+
+ur_mem_sub_buffer_t::ur_mem_sub_buffer_t(ur_mem_handle_t hParent, size_t offset,
+                                         size_t size,
+                                         device_access_mode_t accessMode)
+    : ur_mem_handle_t_(hParent->getContext(), size, accessMode),
+      hParent(hParent), offset(offset), size(size) {
+  ur::level_zero::urMemRetain(hParent);
+}
+
+ur_mem_sub_buffer_t::~ur_mem_sub_buffer_t() {
+  ur::level_zero::urMemRelease(hParent);
+}
+
+void *ur_mem_sub_buffer_t::getDevicePtr(
+    ur_device_handle_t hDevice, device_access_mode_t access, size_t offset,
+    size_t size, std::function<void(void *src, void *dst, size_t)> migrate) {
+  return hParent->getDevicePtr(hDevice, access, offset + this->offset, size,
+                               migrate);
+}
+
+void *ur_mem_sub_buffer_t::mapHostPtr(
+    ur_map_flags_t flags, size_t offset, size_t size,
+    std::function<void(void *src, void *dst, size_t)> migrate) {
+  return hParent->mapHostPtr(flags, offset + this->offset, size, migrate);
+}
+
+void ur_mem_sub_buffer_t::unmapHostPtr(
+    void *pMappedPtr,
+    std::function<void(void *src, void *dst, size_t)> migrate) {
+  return hParent->unmapHostPtr(pMappedPtr, migrate);
+}
+
+size_t ur_mem_sub_buffer_t::getSize() const { return size; }
+
+ur_shared_mutex &ur_mem_sub_buffer_t::getMutex() { return hParent->getMutex(); }
+
 namespace ur::level_zero {
 ur_result_t urMemBufferCreate(ur_context_handle_t hContext,
                               ur_mem_flags_t flags, size_t size,
@@ -347,6 +413,7 @@ ur_result_t urMemBufferCreate(ur_context_handle_t hContext,
   }
 
   void *hostPtr = pProperties ? pProperties->pHost : nullptr;
+  auto accessMode = getDeviceAccessMode(flags);
 
   if (useHostBuffer(hContext)) {
     // TODO: assert that if hostPtr is set, either UR_MEM_FLAG_USE_HOST_POINTER
@@ -355,10 +422,11 @@ ur_result_t urMemBufferCreate(ur_context_handle_t hContext,
         flags & UR_MEM_FLAG_USE_HOST_POINTER
             ? ur_integrated_mem_handle_t::host_ptr_action_t::import
             : ur_integrated_mem_handle_t::host_ptr_action_t::copy;
-    *phBuffer =
-        new ur_integrated_mem_handle_t(hContext, hostPtr, size, hostPtrAction);
+    *phBuffer = new ur_integrated_mem_handle_t(hContext, hostPtr, size,
+                                               hostPtrAction, accessMode);
   } else {
-    *phBuffer = new ur_discrete_mem_handle_t(hContext, hostPtr, size);
+    *phBuffer =
+        new ur_discrete_mem_handle_t(hContext, hostPtr, size, accessMode);
   }
 
   return UR_RESULT_SUCCESS;
@@ -368,13 +436,21 @@ ur_result_t urMemBufferPartition(ur_mem_handle_t hBuffer, ur_mem_flags_t flags,
                                  ur_buffer_create_type_t bufferCreateType,
                                  const ur_buffer_region_t *pRegion,
                                  ur_mem_handle_t *phMem) {
-  std::ignore = hBuffer;
-  std::ignore = flags;
-  std::ignore = bufferCreateType;
-  std::ignore = pRegion;
-  std::ignore = phMem;
-  logger::error("{} function not implemented!", __FUNCTION__);
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+  UR_ASSERT(bufferCreateType == UR_BUFFER_CREATE_TYPE_REGION,
+            UR_RESULT_ERROR_INVALID_ENUMERATION);
+  UR_ASSERT((pRegion->origin < hBuffer->getSize() &&
+             pRegion->size <= hBuffer->getSize()),
+            UR_RESULT_ERROR_INVALID_BUFFER_SIZE);
+
+  auto accessMode = getDeviceAccessMode(flags);
+
+  UR_ASSERT(isAccessCompatible(accessMode, hBuffer->getDeviceAccessMode()),
+            UR_RESULT_ERROR_INVALID_VALUE);
+
+  *phMem = new ur_mem_sub_buffer_t(hBuffer, pRegion->origin, pRegion->size,
+                                   accessMode);
+
+  return UR_RESULT_SUCCESS;
 }
 
 ur_result_t urMemBufferCreateWithNativeHandle(
@@ -407,21 +483,24 @@ ur_result_t urMemBufferCreateWithNativeHandle(
               UR_RESULT_ERROR_INVALID_CONTEXT);
   }
 
+  // assume read-write
+  auto accessMode = ur_mem_handle_t_::device_access_mode_t::read_write;
+
   if (useHostBuffer(hContext) && memoryAttrs.type == ZE_MEMORY_TYPE_HOST) {
-    *phMem =
-        new ur_integrated_mem_handle_t(hContext, ptr, size, ownNativeHandle);
+    *phMem = new ur_integrated_mem_handle_t(hContext, ptr, size, accessMode,
+                                            ownNativeHandle);
     // if useHostBuffer(hContext) is true but the allocation is on device, we'll
     // treat it as discrete memory
   } else {
     if (memoryAttrs.type == ZE_MEMORY_TYPE_HOST) {
       // For host allocation, we need to copy the data to a device buffer
       // and then copy it back on release
       *phMem = new ur_discrete_mem_handle_t(hContext, hDevice, nullptr, size,
-                                            ptr, ownNativeHandle);
+                                            accessMode, ptr, ownNativeHandle);
     } else {
       // For device/shared allocation, we can use it directly
-      *phMem = new ur_discrete_mem_handle_t(hContext, hDevice, ptr, size,
-                                            nullptr, ownNativeHandle);
+      *phMem = new ur_discrete_mem_handle_t(
+          hContext, hDevice, ptr, size, accessMode, nullptr, ownNativeHandle);
     }
   }
 
@@ -452,12 +531,12 @@ ur_result_t urMemGetInfo(ur_mem_handle_t hMemory, ur_mem_info_t propName,
 }
 
 ur_result_t urMemRetain(ur_mem_handle_t hMem) {
-  hMem->RefCount.increment();
+  hMem->getRefCount().increment();
   return UR_RESULT_SUCCESS;
 }
 
 ur_result_t urMemRelease(ur_mem_handle_t hMem) {
-  if (hMem->RefCount.decrementAndTest()) {
+  if (hMem->getRefCount().decrementAndTest()) {
     delete hMem;
   }
   return UR_RESULT_SUCCESS;
@@ -468,11 +547,11 @@ ur_result_t urMemGetNativeHandle(ur_mem_handle_t hMem,
                                  ur_native_handle_t *phNativeMem) {
   std::ignore = hDevice;
 
-  std::scoped_lock<ur_shared_mutex> lock(hMem->Mutex);
+  std::scoped_lock<ur_shared_mutex> lock(hMem->getMutex());
 
-  auto ptr =
-      hMem->getDevicePtr(nullptr, ur_mem_handle_t_::access_mode_t::read_write,
-                         0, hMem->getSize(), nullptr);
+  auto ptr = hMem->getDevicePtr(
+      nullptr, ur_mem_handle_t_::device_access_mode_t::read_write, 0,
+      hMem->getSize(), nullptr);
   *phNativeMem = reinterpret_cast<ur_native_handle_t>(ptr);
   return UR_RESULT_SUCCESS;
 }