[L0] Refcnt Parent Buffer on Sub Buffer Create and die on use of buffer after free

- Refcnt the Parent Buffer on Sub Buffer Creation to avoid free until
  after all sub buffers are released.
- Avoid use after free of buffers by tracking if a buffer has already
  been called with free() and verify that the parent of a sub buffer has
not already been released before attempting to get ze handles from the
parent.
- In either use after free case, die aborts before the memory is accessed.

Signed-off-by: Neil R. Spruit <[email protected]>

Author: nrspruit

source/adapters/level_zero/memory.cpp

@@ -2045,19 +2045,30 @@ ur_result_t _ur_buffer::getZeHandle(char *&ZeHandle, access_mode_t AccessMode,
 
   auto &Allocation = Allocations[Device];
 
+  if (this->isFreed) {
+    die("getZeHandle() buffer already released, no valid handles.");
+  }
+
   // Sub-buffers don't maintain own allocations but rely on parent buffer.
   if (SubBuffer) {
-    UR_CALL(SubBuffer->Parent->getZeHandle(ZeHandle, AccessMode, Device,
-                                           phWaitEvents, numWaitEvents));
-    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,
+                                             phWaitEvents, numWaitEvents));
+      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 an error if the parent buffer is already gone.
+      die("getZeHandle() SubBuffer's parent already released, no valid "
+          "handles.");
+    }
   }
 
   // First handle case where the buffer is represented by only
@@ -2320,6 +2331,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;
 }

source/adapters/level_zero/memory.hpp

@@ -108,7 +108,10 @@ struct _ur_buffer final : ur_mem_handle_t_ {
   // Sub-buffer constructor
   _ur_buffer(_ur_buffer *Parent, size_t Origin, size_t Size)
       : ur_mem_handle_t_(Parent->UrContext),
-        Size(Size), SubBuffer{{Parent, Origin}} {}
+        Size(Size), SubBuffer{{Parent, Origin}} {
+    // Retain the Parent Buffer due to the Creation of the SubBuffer.
+    Parent->RefCount.increment();
+  }
 
   // Interop-buffer constructor
   _ur_buffer(ur_context_handle_t Context, size_t Size,
@@ -136,6 +139,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.