diff --git a/source/adapters/level_zero/memory.cpp b/source/adapters/level_zero/memory.cpp index 69edf83a78..de65d062aa 100644 --- a/source/adapters/level_zero/memory.cpp +++ b/source/adapters/level_zero/memory.cpp @@ -1682,6 +1682,7 @@ ur_result_t urMemRelease( Buffer->free(); } delete Mem; + Mem = nullptr; return UR_RESULT_SUCCESS; } @@ -2024,18 +2025,37 @@ ur_result_t _ur_buffer::getZeHandle(char *&ZeHandle, access_mode_t AccessMode, auto &Allocation = Allocations[Device]; + if (this->isFreed) { + logger::debug("getZeHandle(ur_device_handle{}) buffer already released, no " + "valid handles.", + (void *)Device); + ZeHandle = nullptr; + return UR_RESULT_ERROR_INVALID_NULL_HANDLE; + } + // Sub-buffers don't maintain own allocations but rely on parent buffer. if (SubBuffer) { - UR_CALL(SubBuffer->Parent->getZeHandle(ZeHandle, AccessMode, Device)); - ZeHandle += SubBuffer->Origin; - // Still store the allocation info in the PI sub-buffer for - // getZeHandlePtr to work. At least zeKernelSetArgumentValue needs to - // be given a pointer to the allocation handle rather than its value. - // - Allocation.ZeHandle = ZeHandle; - Allocation.ReleaseAction = allocation_t::keep; - LastDeviceWithValidAllocation = Device; - return UR_RESULT_SUCCESS; + // Verify that the Parent Buffer is still valid or if it has been freed. + if (SubBuffer->Parent && !SubBuffer->Parent->isFreed) { + UR_CALL(SubBuffer->Parent->getZeHandle(ZeHandle, AccessMode, Device)); + ZeHandle += SubBuffer->Origin; + // Still store the allocation info in the PI sub-buffer for + // getZeHandlePtr to work. At least zeKernelSetArgumentValue needs to + // be given a pointer to the allocation handle rather than its value. + // + Allocation.ZeHandle = ZeHandle; + Allocation.ReleaseAction = allocation_t::keep; + LastDeviceWithValidAllocation = Device; + return UR_RESULT_SUCCESS; + } else { + // Return a null pointer handle and an error if the parent buffer is + // already gone. + logger::debug("getZeHandle(ur_device_handle{}) SubBuffer's parent " + "already released, no valid handles.", + (void *)Device); + ZeHandle = nullptr; + return UR_RESULT_ERROR_INVALID_NULL_HANDLE; + } } // First handle case where the buffer is represented by only @@ -2266,6 +2286,7 @@ ur_result_t _ur_buffer::free() { die("_ur_buffer::free(): Unhandled release action"); } ZeHandle = nullptr; // don't leave hanging pointers + this->isFreed = true; } return UR_RESULT_SUCCESS; } diff --git a/source/adapters/level_zero/memory.hpp b/source/adapters/level_zero/memory.hpp index 71d102e9dd..b39dd58eb8 100644 --- a/source/adapters/level_zero/memory.hpp +++ b/source/adapters/level_zero/memory.hpp @@ -129,6 +129,9 @@ struct _ur_buffer final : ur_mem_handle_t_ { // Frees all allocations made for the buffer. ur_result_t free(); + // Tracks if this buffer is freed already or should be considered valid. + bool isFreed{false}; + // Information about a single allocation representing this buffer. struct allocation_t { // Level Zero memory handle is really just a naked pointer.