diff --git a/unified-runtime/source/adapters/level_zero/memory.cpp b/unified-runtime/source/adapters/level_zero/memory.cpp index dd12327bf3f61..346b67c69ad80 100644 --- a/unified-runtime/source/adapters/level_zero/memory.cpp +++ b/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 importGuard( + HostPtrImported ? Host : nullptr, releaseImport); + ur_buffer *Buffer = nullptr; auto HostPtrOrNull = (Flags & UR_MEM_FLAG_USE_HOST_POINTER) ? reinterpret_cast(Host) : nullptr; + + std::unique_ptr 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(Buffer)); + *RetBuffer = v1_cast(static_cast(bufferGuard.release())); return UR_RESULT_SUCCESS; } @@ -1779,17 +1796,19 @@ ur_result_t urMemBufferPartition( "no read-only or write-only yet."); } + std::unique_ptr partitionedBuffer; try { - auto partitionedBuffer = - new ur_buffer(static_cast(Buffer), - BufferCreateInfo->origin, BufferCreateInfo->size); - *RetMem = v1_cast(static_cast(partitionedBuffer)); + partitionedBuffer = std::make_unique( + static_cast(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(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 Buffer = { + nullptr, [](ur_buffer *handle) { + if (handle) { + handle->DeviceMappedHostNativePtr = nullptr; + handle->free(); + delete handle; + } + }}; + std::shared_lock 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(NativeMem), - OwnNativeHandle); - *Mem = v1_cast(static_cast(Buffer)); + Buffer.reset(new ur_buffer(Context, Size, Device, + ur_cast(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 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 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(Buffer.release())); return UR_RESULT_SUCCESS; } @@ -2250,6 +2285,10 @@ ur_result_t ur_buffer::getBufferZeHandle(char *&ZeHandle, NeedCopy = false; } + std::unique_ptr + 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;