[SYCL] Remove OwnZeMemHandle from USMAllocator (#7853)

igchor · web-flow · commit 00ae1e722a19 · 2023-03-24T16:02:46.000-07:00
OwnZeMemHandle is used to track whether an allocation should be freed by
the runtime or not.

There is no point in passing this flag to USMAllocator and then to
USMFreeHelper since the allocation cannot come from the USMAllocator (it
can only be from zeMemAlloc*).

I've thought about moving the memory allocation code to UR first, but
there are a lot of dependencies: we would have to move the
implementation of context, device, event, and queue at least. It seemed
that doing this small change first made more sense.
diff --git a/sycl/plugins/level_zero/pi_level_zero.cpp b/sycl/plugins/level_zero/pi_level_zero.cpp
@@ -3065,8 +3065,7 @@ pi_result piMemRetain(pi_mem Mem) {
 // If indirect access tracking is not enabled then this functions just performs
 // zeMemFree. If indirect access tracking is enabled then reference counting is
 // performed.
-static pi_result ZeMemFreeHelper(pi_context Context, void *Ptr,
-                                 bool OwnZeMemHandle = true) {
+static pi_result ZeMemFreeHelper(pi_context Context, void *Ptr) {
   pi_platform Plt = Context->getPlatform();
   std::unique_lock<pi_shared_mutex> ContextsLock(Plt->ContextsMutex,
                                                  std::defer_lock);
@@ -3086,8 +3085,7 @@ static pi_result ZeMemFreeHelper(pi_context Context, void *Ptr,
     Context->MemAllocs.erase(It);
   }
 
-  if (OwnZeMemHandle)
-    ZE_CALL(zeMemFree, (Context->ZeContext, Ptr));
+  ZE_CALL(zeMemFree, (Context->ZeContext, Ptr));
 
   if (IndirectAccessTrackingEnabled)
     PI_CALL(ContextReleaseHelper(Context));
@@ -3096,7 +3094,7 @@ static pi_result ZeMemFreeHelper(pi_context Context, void *Ptr,
 }
 
 static pi_result USMFreeHelper(pi_context Context, void *Ptr,
-                               bool OwnZeMemHandle);
+                               bool OwnZeMemHandle = true);
 
 pi_result piMemRelease(pi_mem Mem) {
   PI_ASSERT(Mem, PI_ERROR_INVALID_MEM_OBJECT);
@@ -7050,10 +7048,8 @@ static pi_result USMHostAllocImpl(void **ResultPtr, pi_context Context,
   return PI_SUCCESS;
 }
 
-static pi_result USMFreeImpl(pi_context Context, void *Ptr,
-                             bool OwnZeMemHandle) {
-  if (OwnZeMemHandle)
-    ZE_CALL(zeMemFree, (Context->ZeContext, Ptr));
+static pi_result USMFreeImpl(pi_context Context, void *Ptr) {
+  ZE_CALL(zeMemFree, (Context->ZeContext, Ptr));
   return PI_SUCCESS;
 }
 
@@ -7112,8 +7108,8 @@ void *USMMemoryAllocBase::allocate(size_t Size, size_t Alignment) {
   return Ptr;
 }
 
-void USMMemoryAllocBase::deallocate(void *Ptr, bool OwnZeMemHandle) {
-  auto Res = USMFreeImpl(Context, Ptr, OwnZeMemHandle);
+void USMMemoryAllocBase::deallocate(void *Ptr) {
+  auto Res = USMFreeImpl(Context, Ptr);
   if (Res != PI_SUCCESS) {
     throw UsmAllocationException(Res);
   }
@@ -7361,8 +7357,13 @@ static pi_result USMFreeHelper(pi_context Context, void *Ptr,
     Context->MemAllocs.erase(It);
   }
 
+  if (!OwnZeMemHandle) {
+    // Memory should not be freed
+    return PI_SUCCESS;
+  }
+
   if (!UseUSMAllocator) {
-    pi_result Res = USMFreeImpl(Context, Ptr, OwnZeMemHandle);
+    pi_result Res = USMFreeImpl(Context, Ptr);
     if (IndirectAccessTrackingEnabled)
       PI_CALL(ContextReleaseHelper(Context));
     return Res;
@@ -7381,7 +7382,7 @@ static pi_result USMFreeHelper(pi_context Context, void *Ptr,
   // If memory type is host release from host pool
   if (ZeMemoryAllocationProperties.type == ZE_MEMORY_TYPE_HOST) {
     try {
-      Context->HostMemAllocContext->deallocate(Ptr, OwnZeMemHandle);
+      Context->HostMemAllocContext->deallocate(Ptr);
     } catch (const UsmAllocationException &Ex) {
       return Ex.getError();
     } catch (...) {
@@ -7409,16 +7410,16 @@ static pi_result USMFreeHelper(pi_context Context, void *Ptr,
     PI_ASSERT(Device, PI_ERROR_INVALID_DEVICE);
 
     auto DeallocationHelper =
-        [Context, Device, Ptr,
-         OwnZeMemHandle](std::unordered_map<ze_device_handle_t, USMAllocContext>
-                             &AllocContextMap) {
+        [Context, Device,
+         Ptr](std::unordered_map<ze_device_handle_t, USMAllocContext>
+                  &AllocContextMap) {
           try {
             auto It = AllocContextMap.find(Device->ZeDevice);
             if (It == AllocContextMap.end())
               return PI_ERROR_INVALID_VALUE;
 
             // The right context is found, deallocate the pointer
-            It->second.deallocate(Ptr, OwnZeMemHandle);
+            It->second.deallocate(Ptr);
           } catch (const UsmAllocationException &Ex) {
             return Ex.getError();
           }
@@ -7444,7 +7445,7 @@ static pi_result USMFreeHelper(pi_context Context, void *Ptr,
     }
   }
 
-  pi_result Res = USMFreeImpl(Context, Ptr, OwnZeMemHandle);
+  pi_result Res = USMFreeImpl(Context, Ptr);
   if (SharedReadOnlyAllocsIterator != Context->SharedReadOnlyAllocs.end()) {
     Context->SharedReadOnlyAllocs.erase(SharedReadOnlyAllocsIterator);
   }
@@ -7459,7 +7460,7 @@ pi_result piextUSMFree(pi_context Context, void *Ptr) {
   std::scoped_lock<pi_shared_mutex> Lock(
       IndirectAccessTrackingEnabled ? Plt->ContextsMutex : Context->Mutex);
 
-  return USMFreeHelper(Context, Ptr, true /* OwnZeMemHandle */);
+  return USMFreeHelper(Context, Ptr);
 }
 
 pi_result piextKernelSetArgPointer(pi_kernel Kernel, pi_uint32 ArgIndex,
@@ -8379,11 +8380,11 @@ pi_result _pi_buffer::free() {
       std::scoped_lock<pi_shared_mutex> Lock(
           IndirectAccessTrackingEnabled ? Plt->ContextsMutex : Context->Mutex);
 
-      PI_CALL(USMFreeHelper(Context, ZeHandle, true));
+      PI_CALL(USMFreeHelper(Context, ZeHandle));
       break;
     }
     case allocation_t::free_native:
-      PI_CALL(ZeMemFreeHelper(Context, ZeHandle, true));
+      PI_CALL(ZeMemFreeHelper(Context, ZeHandle));
       break;
     case allocation_t::unimport:
       ZeUSMImport.doZeUSMRelease(Context->getPlatform()->ZeDriver, ZeHandle);
diff --git a/sycl/plugins/level_zero/pi_level_zero.hpp b/sycl/plugins/level_zero/pi_level_zero.hpp
@@ -112,7 +112,7 @@ class USMMemoryAllocBase : public SystemMemory {
       : Context{Ctx}, Device{Dev} {}
   void *allocate(size_t Size) override final;
   void *allocate(size_t Size, size_t Alignment) override final;
-  void deallocate(void *Ptr, bool OwnZeMemHandle) override final;
+  void deallocate(void *Ptr) override final;
 };
 
 // Allocation routines for shared memory type
diff --git a/sycl/plugins/unified_runtime/ur/usm_allocator.cpp b/sycl/plugins/unified_runtime/ur/usm_allocator.cpp
@@ -290,7 +290,7 @@ class USMAllocContext::USMAllocImpl {
 
   void *allocate(size_t Size, size_t Alignment, bool &FromPool);
   void *allocate(size_t Size, bool &FromPool);
-  void deallocate(void *Ptr, bool &ToPool, bool OwnZeMemHandle);
+  void deallocate(void *Ptr, bool &ToPool);
 
   SystemMemory &getMemHandle() { return *MemHandle; }
 
@@ -332,7 +332,7 @@ Slab::Slab(Bucket &Bkt)
 
 Slab::~Slab() {
   unregSlab(*this);
-  bucket.getMemHandle().deallocate(MemPtr, true /* OwnZeMemHandle */);
+  bucket.getMemHandle().deallocate(MemPtr);
 }
 
 // Return the index of the first available chunk, -1 otherwize
@@ -737,8 +737,7 @@ Bucket &USMAllocContext::USMAllocImpl::findBucket(size_t Size) {
   return *(*It);
 }
 
-void USMAllocContext::USMAllocImpl::deallocate(void *Ptr, bool &ToPool,
-                                               bool OwnZeMemHandle) {
+void USMAllocContext::USMAllocImpl::deallocate(void *Ptr, bool &ToPool) {
   auto *SlabPtr = AlignPtrDown(Ptr, SlabMinSize());
 
   // Lock the map on read
@@ -748,7 +747,7 @@ void USMAllocContext::USMAllocImpl::deallocate(void *Ptr, bool &ToPool,
   auto Slabs = getKnownSlabs().equal_range(SlabPtr);
   if (Slabs.first == Slabs.second) {
     Lk.unlock();
-    getMemHandle().deallocate(Ptr, OwnZeMemHandle);
+    getMemHandle().deallocate(Ptr);
     return;
   }
 
@@ -779,7 +778,7 @@ void USMAllocContext::USMAllocImpl::deallocate(void *Ptr, bool &ToPool,
   // There is a rare case when we have a pointer from system allocation next
   // to some slab with an entry in the map. So we find a slab
   // but the range checks fail.
-  getMemHandle().deallocate(Ptr, OwnZeMemHandle);
+  getMemHandle().deallocate(Ptr);
 }
 
 USMAllocContext::USMAllocContext(std::unique_ptr<SystemMemory> MemHandle,
@@ -813,9 +812,9 @@ void *USMAllocContext::allocate(size_t size, size_t alignment) {
   return Ptr;
 }
 
-void USMAllocContext::deallocate(void *ptr, bool OwnZeMemHandle) {
+void USMAllocContext::deallocate(void *ptr) {
   bool ToPool;
-  pImpl->deallocate(ptr, ToPool, OwnZeMemHandle);
+  pImpl->deallocate(ptr, ToPool);
 
   if (pImpl->getParams().PoolTrace > 2) {
     auto MT = pImpl->getParams().memoryTypeName;
diff --git a/sycl/plugins/unified_runtime/ur/usm_allocator.hpp b/sycl/plugins/unified_runtime/ur/usm_allocator.hpp
@@ -17,7 +17,7 @@ class SystemMemory {
 public:
   virtual void *allocate(size_t size) = 0;
   virtual void *allocate(size_t size, size_t aligned) = 0;
-  virtual void deallocate(void *ptr, bool OwnZeMemHandle) = 0;
+  virtual void deallocate(void *ptr) = 0;
   virtual ~SystemMemory() = default;
 };
 
@@ -68,7 +68,7 @@ class USMAllocContext {
 
   void *allocate(size_t size);
   void *allocate(size_t size, size_t alignment);
-  void deallocate(void *ptr, bool OwnZeMemHandle);
+  void deallocate(void *ptr);
 
 private:
   std::unique_ptr<USMAllocImpl> pImpl;
