fix: remove heap resource leak after command list reset

Signed-off-by: Zbigniew Zdanowicz <zbigniew.zdanowicz@intel.com>

level_zero/core/source/cmdlist/cmdlist.cpp

@@ -130,12 +130,9 @@ void CommandList::removeDeallocationContainerData() {
         if (allocData) {
             device->getDriverHandle()->getSvmAllocsManager()->removeSVMAlloc(*allocData);
         }
-        if (!((deallocation->getAllocationType() == NEO::AllocationType::internalHeap) ||
-              (deallocation->getAllocationType() == NEO::AllocationType::linearStream))) {
         memoryManager->freeGraphicsMemory(deallocation);
         eraseDeallocationContainerEntry(deallocation);
     }
-    }
 }
 
 void CommandList::eraseDeallocationContainerEntry(NEO::GraphicsAllocation *allocation) {

level_zero/core/test/unit_tests/sources/cmdlist/test_cmdlist_7.cpp

@@ -1923,23 +1923,6 @@ HWTEST2_F(CommandListCreate, givenImmediateCommandListWhenThereIsNoEnoughSpaceFo
     EXPECT_EQ(latestFlushedTaskCount + 1, whiteBoxCmdList->csr->peekLatestFlushedTaskCount());
 }
 
-HWTEST_F(CommandListCreate, givenCommandListWhenRemoveDeallocationContainerDataThenHeapNotErased) {
-    ze_result_t returnValue;
-    std::unique_ptr<L0::CommandList> commandList(CommandList::create(productFamily,
-                                                                     device,
-                                                                     NEO::EngineGroupType::compute,
-                                                                     0u,
-                                                                     returnValue, false));
-    auto &cmdContainer = commandList->getCmdContainer();
-    auto heapAlloc = cmdContainer.getIndirectHeapAllocation(HeapType::indirectObject);
-    cmdContainer.getDeallocationContainer().push_back(heapAlloc);
-    EXPECT_EQ(cmdContainer.getDeallocationContainer().size(), 1u);
-    commandList->removeDeallocationContainerData();
-    EXPECT_EQ(cmdContainer.getDeallocationContainer().size(), 1u);
-
-    cmdContainer.getDeallocationContainer().clear();
-}
-
 TEST(CommandList, givenContextGroupEnabledWhenCreatingImmediateCommandListThenEachCmdListHasDifferentCsr) {
 
     HardwareInfo hwInfo = *defaultHwInfo;

shared/source/command_container/cmdcontainer.cpp

@@ -35,14 +35,12 @@ CommandContainer::~CommandContainer() {
 
     this->handleCmdBufferAllocations(0u);
 
-    for (auto allocationIndirectHeap : allocationIndirectHeaps) {
     if (heapHelper) {
+        for (auto allocationIndirectHeap : allocationIndirectHeaps) {
             heapHelper->storeHeapAllocation(allocationIndirectHeap);
         }
-    }
+        for (auto heapAllocation : storedHeapsContainer) {
-    for (auto deallocation : deallocationContainer) {
+            heapHelper->storeHeapAllocation(heapAllocation);
-        if (((deallocation->getAllocationType() == AllocationType::internalHeap) || (deallocation->getAllocationType() == AllocationType::linearStream))) {
-            getHeapHelper()->storeHeapAllocation(deallocation);
         }
     }
 }
@@ -185,6 +183,12 @@ void CommandContainer::reset() {
     setDirtyStateForAllHeaps(true);
     slmSize = std::numeric_limits<uint32_t>::max();
     getResidencyContainer().clear();
+    if (getHeapHelper()) {
+        for (auto heapAllocation : storedHeapsContainer) {
+            getHeapHelper()->storeHeapAllocation(heapAllocation);
+        }
+    }
+    storedHeapsContainer.clear();
     getDeallocationContainer().clear();
     sshAllocations.clear();
 
@@ -282,7 +286,7 @@ void CommandContainer::createAndAssignNewHeap(HeapType heapType, size_t size) {
     if (this->immediateCmdListCsr) {
         this->storeAllocationAndFlushTagUpdate(oldAlloc);
     } else {
-        getDeallocationContainer().push_back(oldAlloc);
+        storedHeapsContainer.push_back(oldAlloc);
     }
     setIndirectHeapAllocation(heapType, newAlloc);
     if (oldBase != newBase) {

shared/source/command_container/cmdcontainer.h

@@ -220,6 +220,7 @@ class CommandContainer : public NonCopyableOrMovableClass {
     CmdBufferContainer cmdBufferAllocations;
     ResidencyContainer residencyContainer;
     std::vector<GraphicsAllocation *> deallocationContainer;
+    std::vector<GraphicsAllocation *> storedHeapsContainer;
     HeapContainer sshAllocations;
 
     HeapReserveData dynamicStateHeapReserveData;

shared/test/unit_test/command_container/command_container_tests.cpp

@@ -38,6 +38,7 @@ class MyMockCommandContainer : public CommandContainer {
     using CommandContainer::getAlignedCmdBufferSize;
     using CommandContainer::immediateReusableAllocationList;
     using CommandContainer::secondaryCommandStreamForImmediateCmdList;
+    using CommandContainer::storedHeapsContainer;
 
     GraphicsAllocation *allocateCommandBuffer(bool forceHostMemory) override {
         allocateCommandBufferCalled[!!forceHostMemory]++;
@@ -271,7 +272,7 @@ TEST_F(CommandContainerTest, givenCreateSecondaryCmdBufferInHostMemWhenAllocateS
 TEST_F(CommandContainerTest, givenCmdContainerWithAllocsListWhenAllocateAndResetThenCmdBufferAllocIsReused) {
     AllocationsList allocList;
     auto cmdContainer = std::make_unique<CommandContainer>();
-    cmdContainer->initialize(pDevice, &allocList, true, HeapSize::defaultHeapSize, false);
+    cmdContainer->initialize(pDevice, &allocList, HeapSize::defaultHeapSize, true, false);
     auto &cmdBufferAllocs = cmdContainer->getCmdBufferAllocations();
     auto memoryManager = static_cast<MockMemoryManager *>(pDevice->getMemoryManager());
     EXPECT_EQ(memoryManager->handleFenceCompletionCalled, 0u);
@@ -373,10 +374,10 @@ TEST_F(CommandContainerTest, givenCommandContainerWhenSettingIndirectHeapAllocat
 
 TEST_F(CommandContainerTest, givenHeapAllocationsWhenDestroyCommandContainerThenHeapAllocationsAreReused) {
     std::unique_ptr<CommandContainer> cmdContainer(new CommandContainer);
-    cmdContainer->initialize(pDevice, nullptr, true, HeapSize::defaultHeapSize, false);
+    cmdContainer->initialize(pDevice, nullptr, HeapSize::defaultHeapSize, true, false);
     auto heapAllocationsAddress = cmdContainer->getIndirectHeapAllocation(HeapType::surfaceState)->getUnderlyingBuffer();
     cmdContainer.reset(new CommandContainer);
-    cmdContainer->initialize(pDevice, nullptr, true, HeapSize::defaultHeapSize, false);
+    cmdContainer->initialize(pDevice, nullptr, HeapSize::defaultHeapSize, true, false);
     bool status = true;
     for (uint32_t i = 0; i < HeapType::numTypes && !status; i++) {
         auto heapType = static_cast<HeapType>(i);
@@ -1899,3 +1900,43 @@ TEST_F(CommandContainerTest, givenCmdContainerWhenImmediateCmdListCsrIsSetThenCo
     cmdContainer.initialize(pDevice, nullptr, HeapSize::defaultHeapSize, false, false);
     EXPECT_EQ(cmdContainer.getCommandStream()->getCmdContainer(), nullptr);
 }
+
+TEST_F(CommandContainerTest, givenCmdContainerWhenOldHeapIsStoredAndResetContainerThenUseStorageForReuseForStoredHeap) {
+    MyMockCommandContainer cmdContainer;
+
+    auto status = cmdContainer.initialize(pDevice, nullptr, HeapSize::defaultHeapSize, true, false);
+    EXPECT_EQ(CommandContainer::ErrorCode::success, status);
+    EXPECT_EQ(0u, cmdContainer.storedHeapsContainer.size());
+    auto ioh = cmdContainer.getIndirectHeap(NEO::HeapType::indirectObject);
+    auto iohOldAllocation = ioh->getGraphicsAllocation();
+
+    ioh->getSpace(ioh->getAvailableSpace());
+    cmdContainer.getHeapWithRequiredSizeAndAlignment(NEO::HeapType::indirectObject, 64, 1);
+    EXPECT_EQ(1u, cmdContainer.storedHeapsContainer.size());
+    auto iohAllocIt = std::find(cmdContainer.storedHeapsContainer.begin(),
+                                cmdContainer.storedHeapsContainer.end(),
+                                iohOldAllocation);
+    EXPECT_NE(cmdContainer.storedHeapsContainer.end(), iohAllocIt);
+
+    cmdContainer.reset();
+    EXPECT_EQ(0u, cmdContainer.storedHeapsContainer.size());
+
+    auto internalStorage = pDevice->getDefaultEngine().commandStreamReceiver->getInternalAllocationStorage();
+
+    auto iohReusedAllocation = internalStorage->obtainReusableAllocation(iohOldAllocation->getUnderlyingBufferSize(),
+                                                                         iohOldAllocation->getAllocationType())
+                                   .release();
+    EXPECT_EQ(iohOldAllocation, iohReusedAllocation);
+    pDevice->getMemoryManager()->freeGraphicsMemory(iohReusedAllocation);
+}
+
+TEST_F(CommandContainerTest, givenHeaplessCmdContainerWhenResetContainerThenNoHeapInStorageForReuse) {
+    MyMockCommandContainer cmdContainer;
+
+    auto status = cmdContainer.initialize(pDevice, nullptr, HeapSize::defaultHeapSize, false, false);
+    EXPECT_EQ(CommandContainer::ErrorCode::success, status);
+    EXPECT_EQ(0u, cmdContainer.storedHeapsContainer.size());
+
+    cmdContainer.reset();
+    EXPECT_EQ(0u, cmdContainer.storedHeapsContainer.size());
+}