From 193c965a1c672e5398a0180dde64745ffafad831 Mon Sep 17 00:00:00 2001
From: aarongreig <aaron.greig@codeplay.com>
Date: Fri, 4 Oct 2024 11:35:10 +0100
Subject: [PATCH] Merge pull request #2121 from
 nrspruit/error_after_free_syclos

[L0] Refcnt Parent Buffer on Sub Buffer Create and die on use of buffer after free
---
 source/adapters/level_zero/memory.cpp | 34 ++++++++++++++++++---------
 source/adapters/level_zero/memory.hpp |  8 ++++++-
 2 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/source/adapters/level_zero/memory.cpp b/source/adapters/level_zero/memory.cpp
index b23f0eb960..0925d2027a 100644
--- a/source/adapters/level_zero/memory.cpp
+++ b/source/adapters/level_zero/memory.cpp
@@ -1990,19 +1990,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
@@ -2263,6 +2274,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 d4a0376eae..9a3e860672 100644
--- a/source/adapters/level_zero/memory.hpp
+++ b/source/adapters/level_zero/memory.hpp
@@ -112,7 +112,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,
@@ -140,6 +143,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.