[UR][L0] Fix potential memory leaks

unified-runtime/source/adapters/level_zero/memory.cpp

@@ -1664,12 +1664,29 @@ ur_result_t urMemBufferCreate(
         maybeImportUSM(Context->getPlatform()->ZeDriverHandleExpTranslated,
                        Context->ZeContext, Host, Size);
 
+  auto releaseImport = [Context](void *Ptr) {
+    ZeUSMImport.doZeUSMRelease(
+        Context->getPlatform()->ZeDriverHandleExpTranslated, Ptr);
+  };
+  std::unique_ptr<void, decltype(releaseImport)> importGuard(
+      HostPtrImported ? Host : nullptr, releaseImport);
+
   ur_buffer *Buffer = nullptr;
   auto HostPtrOrNull = (Flags & UR_MEM_FLAG_USE_HOST_POINTER)
                            ? reinterpret_cast<char *>(Host)
                            : nullptr;
+
+  std::unique_ptr<ur_buffer, void (*)(ur_buffer *)> bufferGuard = {
+      nullptr, [](ur_buffer *handle) {
+        if (handle)
+          handle->free();
+        delete handle;
+      }};
+
   try {
     Buffer = new ur_buffer(Context, Size, HostPtrOrNull, HostPtrImported);
+    bufferGuard.reset(Buffer);
+    importGuard.release();
   } catch (const std::bad_alloc &) {
     return UR_RESULT_ERROR_OUT_OF_HOST_MEMORY;
   } catch (...) {
@@ -1708,7 +1725,7 @@ ur_result_t urMemBufferCreate(
       die("urMemBufferCreate: not implemented");
   }
 
-  *RetBuffer = v1_cast(static_cast<ur_mem_handle_t_ *>(Buffer));
+  *RetBuffer = v1_cast(static_cast<ur_mem_handle_t_ *>(bufferGuard.release()));
   return UR_RESULT_SUCCESS;
 }
 
@@ -1779,17 +1796,19 @@ ur_result_t urMemBufferPartition(
         "no read-only or write-only yet.");
   }
 
+  std::unique_ptr<ur_buffer> partitionedBuffer;
   try {
-    auto partitionedBuffer =
-        new ur_buffer(static_cast<ur_buffer *>(Buffer),
-                      BufferCreateInfo->origin, BufferCreateInfo->size);
-    *RetMem = v1_cast(static_cast<ur_mem_handle_t_ *>(partitionedBuffer));
+    partitionedBuffer = std::make_unique<ur_buffer>(
+        static_cast<ur_buffer *>(Buffer), BufferCreateInfo->origin,
+        BufferCreateInfo->size);
   } catch (const std::bad_alloc &) {
     return UR_RESULT_ERROR_OUT_OF_HOST_MEMORY;
   } catch (...) {
     return UR_RESULT_ERROR_UNKNOWN;
   }
 
+  *RetMem =
+      v1_cast(static_cast<ur_mem_handle_t_ *>(partitionedBuffer.release()));
   return UR_RESULT_SUCCESS;
 }
 
@@ -1823,6 +1842,15 @@ ur_result_t urMemBufferCreateWithNativeHandle(
   auto Context = v1_cast(ContextOpque);
   bool OwnNativeHandle = Properties ? Properties->isNativeHandleOwned : false;
 
+  std::unique_ptr<ur_buffer, void (*)(ur_buffer *)> Buffer = {
+      nullptr, [](ur_buffer *handle) {
+        if (handle) {
+          handle->DeviceMappedHostNativePtr = nullptr;
+          handle->free();
+          delete handle;
+        }
+      }};
+
   std::shared_lock<ur_shared_mutex> Lock(Context->Mutex);
 
   // Get base of the allocation
@@ -1847,36 +1875,16 @@ ur_result_t urMemBufferCreateWithNativeHandle(
     UR_ASSERT(Context->isValidDevice(Device), UR_RESULT_ERROR_INVALID_CONTEXT);
   }
 
-  ur_buffer *Buffer = nullptr;
   try {
-    Buffer = new ur_buffer(Context, Size, Device, ur_cast<char *>(NativeMem),
-                           OwnNativeHandle);
-    *Mem = v1_cast(static_cast<ur_mem_handle_t_ *>(Buffer));
+    Buffer.reset(new ur_buffer(Context, Size, Device,
+                               ur_cast<char *>(NativeMem),
+                               /*OwnZeMemHandle=*/false));
   } catch (const std::bad_alloc &) {
     return UR_RESULT_ERROR_OUT_OF_HOST_MEMORY;
   } catch (...) {
     return UR_RESULT_ERROR_UNKNOWN;
   }
 
-  ur_platform_handle_t Plt = Context->getPlatform();
-  std::unique_lock<ur_shared_mutex> ContextsLock(Plt->ContextsMutex,
-                                                 std::defer_lock);
-  // If we don't own the native handle then we can't control deallocation of
-  // that memory so there is no point of keeping track of the memory
-  // allocation for deferred memory release in the mode when indirect access
-  // tracking is enabled.
-  if (IndirectAccessTrackingEnabled && OwnNativeHandle) {
-    // We need to keep track of all memory allocations in the context
-    ContextsLock.lock();
-    // Retain context to be sure that it is released after all memory
-    // allocations in this context are released.
-    UR_CALL(ur::level_zero::v1::urContextRetain(ContextOpque));
-
-    Context->MemAllocs.emplace(std::piecewise_construct,
-                               std::forward_as_tuple(Ptr),
-                               std::forward_as_tuple(Context, OwnNativeHandle));
-  }
-
   if (Device) {
     // If this allocation is on a device, then we re-use it for the buffer.
     // Nothing to do.
@@ -1904,6 +1912,33 @@ ur_result_t urMemBufferCreateWithNativeHandle(
                 nullptr));
   }
 
+  ur_platform_handle_t Plt = Context->getPlatform();
+  std::unique_lock<ur_shared_mutex> ContextsLock(Plt->ContextsMutex,
+                                                 std::defer_lock);
+  // If we don't own the native handle then we can't control deallocation of
+  // that memory so there is no point of keeping track of the memory
+  // allocation for deferred memory release in the mode when indirect access
+  // tracking is enabled.
+  if (IndirectAccessTrackingEnabled && OwnNativeHandle) {
+    // We need to keep track of all memory allocations in the context
+    ContextsLock.lock();
+    try {
+      Context->MemAllocs.emplace(
+          std::piecewise_construct, std::forward_as_tuple(Ptr),
+          std::forward_as_tuple(Context, OwnNativeHandle));
+    } catch (const std::bad_alloc &) {
+      return UR_RESULT_ERROR_OUT_OF_HOST_MEMORY;
+    }
+    // Retain context to be sure that it is released after all memory
+    // allocations in this context are released.
+    UR_CALL(ur::level_zero::v1::urContextRetain(ContextOpque));
+  }
+
+  if (OwnNativeHandle)
+    Buffer->Allocations[Device].ReleaseAction =
+        ur_buffer::allocation_t::free_native;
+
+  *Mem = v1_cast(static_cast<ur_mem_handle_t_ *>(Buffer.release()));
   return UR_RESULT_SUCCESS;
 }
 
@@ -2250,6 +2285,10 @@ ur_result_t ur_buffer::getBufferZeHandle(char *&ZeHandle,
         NeedCopy = false;
     }
 
+    std::unique_ptr<ze_event_handle_t, void (*)(ze_event_handle_t *)>
+        eventListGuard = {nullptr,
+                          [](ze_event_handle_t *handle) { delete[] handle; }};
+
     if (NeedCopy) {
       // Wait on all dependency events passed in to ensure that the memory which
       // is being init is updated correctly.
@@ -2266,6 +2305,7 @@ ur_result_t ur_buffer::getBufferZeHandle(char *&ZeHandle,
           // dependencies, if they exist for device only.
           if (waitlist.ZeEventList == nullptr) {
             waitlist.ZeEventList = new ze_event_handle_t[numWaitEvents];
+            eventListGuard.reset(waitlist.ZeEventList);
           }
           waitlist.ZeEventList[EventListIndex] = phWaitEvents[i]->ZeEvent;
           waitlist.Length++;
@@ -2332,9 +2372,6 @@ ur_result_t ur_buffer::getBufferZeHandle(char *&ZeHandle,
                    (UrContext->ZeCommandListInit, ZeHandle, ZeHandleSrc, Size,
                     nullptr, 0u, nullptr));
       }
-      if (waitlist.ZeEventList) {
-        delete[] waitlist.ZeEventList;
-      }
     }
     Allocation.Valid = true;
     LastDeviceWithValidAllocation = Device;