From 1d59ffba82b7a53e5ecd911a9a940c366021c106 Mon Sep 17 00:00:00 2001
From: "Dunajski, Bartosz" <bartosz.dunajski@intel.com>
Date: Mon, 4 Mar 2024 09:17:20 +0000
Subject: [PATCH] fix: handle in-order allocator destruction

Related-To: NEO-10572

Signed-off-by: Dunajski, Bartosz <bartosz.dunajski@intel.com>
---
 level_zero/core/source/device/device_imp.cpp  |  1 +
 .../sources/device/test_l0_device.cpp         | 44 ++++++++++---------
 2 files changed, 25 insertions(+), 20 deletions(-)

diff --git a/level_zero/core/source/device/device_imp.cpp b/level_zero/core/source/device/device_imp.cpp
index 42da440c62..819018448c 100644
--- a/level_zero/core/source/device/device_imp.cpp
+++ b/level_zero/core/source/device/device_imp.cpp
@@ -1465,6 +1465,7 @@ void DeviceImp::releaseResources() {
     metricContext.reset();
     builtins.reset();
     cacheReservation.reset();
+    deviceInOrderCounterAllocator.reset();
 
     if (allocationsForReuse.get()) {
         allocationsForReuse->freeAllGraphicsAllocations(neoDevice);
diff --git a/level_zero/core/test/unit_tests/sources/device/test_l0_device.cpp b/level_zero/core/test/unit_tests/sources/device/test_l0_device.cpp
index f947cb0d53..898e0ae96a 100644
--- a/level_zero/core/test/unit_tests/sources/device/test_l0_device.cpp
+++ b/level_zero/core/test/unit_tests/sources/device/test_l0_device.cpp
@@ -30,8 +30,10 @@
 #include "shared/test/common/mocks/mock_memory_manager.h"
 #include "shared/test/common/mocks/mock_os_context.h"
 #include "shared/test/common/mocks/mock_sip.h"
+#include "shared/test/common/mocks/mock_timestamp_container.h"
 #include "shared/test/common/mocks/ult_device_factory.h"
 #include "shared/test/common/test_macros/hw_test.h"
+#include "shared/test/common/utilities/destructor_counted.h"
 
 #include "level_zero/core/source/cache/cache_reservation.h"
 #include "level_zero/core/source/cmdqueue/cmdqueue_imp.h"
@@ -831,32 +833,34 @@ TEST_F(DeviceHostPointerTest, givenHostPointerNotAcceptedByKernelAndHostPointerC
     delete[] buffer;
 }
 
-TEST_F(DeviceTest, whenCreatingDeviceThenCreateInOrderCounterAllocatorOnDemand) {
+TEST_F(DeviceTest, whenCreatingDeviceThenCreateInOrderCounterAllocatorOnDemandAndHandleDestruction) {
+    uint32_t destructorId = 0u;
+
+    class MyMockDevice : public DestructorCounted<NEO::MockDevice, 1> {
+      public:
+        MyMockDevice(NEO::ExecutionEnvironment *executionEnvironment, uint32_t rootDeviceIndex, uint32_t &destructorId) : DestructorCounted(destructorId, executionEnvironment, rootDeviceIndex) {}
+    };
+
+    class MyMockTagAllocator : public DestructorCounted<MockTagAllocator<NEO::DeviceAllocNodeType<true>>, 0> {
+      public:
+        MyMockTagAllocator(uint32_t rootDeviceIndex, MemoryManager *memoryManager, uint32_t &destructorId) : DestructorCounted(destructorId, rootDeviceIndex, memoryManager, 10) {}
+    };
+
     const uint32_t rootDeviceIndex = 0u;
     auto hwInfo = *NEO::defaultHwInfo;
-    auto *neoMockDevice = NEO::MockDevice::createWithNewExecutionEnvironment<NEO::MockDevice>(&hwInfo, rootDeviceIndex);
-    neoMockDevice->incRefInternal();
-    neoMockDevice->deviceBitfield = 0b111;
+
+    auto executionEnvironment = MyMockDevice::prepareExecutionEnvironment(&hwInfo, rootDeviceIndex);
+    executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->setHwInfoAndInitHelpers(&hwInfo);
+    auto *neoMockDevice = new MyMockDevice(executionEnvironment, rootDeviceIndex, destructorId);
+    neoMockDevice->createDeviceImpl();
 
     {
+        auto allocator = new MyMockTagAllocator(0, neoMockDevice->getMemoryManager(), destructorId);
+
         MockDeviceImp deviceImp(neoMockDevice, neoMockDevice->getExecutionEnvironment());
-        EXPECT_EQ(nullptr, deviceImp.deviceInOrderCounterAllocator.get());
-
-        auto allocator = deviceImp.getDeviceInOrderCounterAllocator();
-        EXPECT_NE(nullptr, deviceImp.deviceInOrderCounterAllocator.get());
-
-        auto expectedOffset = alignUp(sizeof(uint64_t) * neoMockDevice->deviceBitfield.count() * 2, MemoryConstants::cacheLineSize);
-
-        auto node1 = allocator->getTag();
-        auto node2 = allocator->getTag();
-
-        EXPECT_EQ(node1->getGpuAddress() + expectedOffset, node2->getGpuAddress());
-        EXPECT_EQ(ptrOffset(node1->getCpuBase(), expectedOffset), node2->getCpuBase());
-
-        node1->returnTag();
-        node2->returnTag();
+        deviceImp.deviceInOrderCounterAllocator.reset(allocator);
+        EXPECT_EQ(allocator, deviceImp.getDeviceInOrderCounterAllocator());
     }
-    neoMockDevice->decRefInternal();
 }
 
 TEST_F(DeviceTest, givenMoreThanOneExtendedPropertiesStructuresWhenKernelPropertiesCalledThenSuccessIsReturnedAndPropertiesAreSet) {