Fix allocating graphics memory for patterns in appendMemoryFill

Signed-off-by: Kacper Nowak <kacper.nowak@intel.com>

level_zero/core/source/cmdlist/cmdlist.h

@@ -257,6 +257,7 @@ struct CommandList : _ze_command_list_handle_t {
 
     std::map<const void *, NEO::GraphicsAllocation *> hostPtrMap;
     std::vector<NEO::GraphicsAllocation *> ownedPrivateAllocations;
+    std::vector<NEO::GraphicsAllocation *> patternAllocations;
 
     NEO::StreamProperties requiredStreamState{};
     NEO::StreamProperties finalStreamState{};

level_zero/core/source/cmdlist/cmdlist_hw.inl

@@ -64,6 +64,10 @@ CommandListCoreFamily<gfxCoreFamily>::~CommandListCoreFamily() {
         device->getNEODevice()->getMemoryManager()->freeGraphicsMemory(alloc);
     }
     this->ownedPrivateAllocations.clear();
+    for (auto &patternAlloc : this->patternAllocations) {
+        device->storeReusableAllocation(*patternAlloc);
+    }
+    this->patternAllocations.clear();
 }
 
 template <GFXCORE_FAMILY gfxCoreFamily>
@@ -1468,16 +1472,15 @@ ze_result_t CommandListCoreFamily<gfxCoreFamily>::appendMemoryFill(void *ptr,
         size_t patternAllocationSize = alignUp(patternSize, MemoryConstants::cacheLineSize);
         uint32_t patternSizeInEls = static_cast<uint32_t>(patternAllocationSize / middleElSize);
 
-        auto patternGfxAlloc = getAllocationFromHostPtrMap(pattern, patternAllocationSize);
+        auto patternGfxAlloc = device->obtainReusableAllocation(patternAllocationSize, NEO::GraphicsAllocation::AllocationType::FILL_PATTERN);
         if (patternGfxAlloc == nullptr) {
             patternGfxAlloc = device->getDriverHandle()->getMemoryManager()->allocateGraphicsMemoryWithProperties({device->getNEODevice()->getRootDeviceIndex(),
                                                                                                                    patternAllocationSize,
                                                                                                                    NEO::GraphicsAllocation::AllocationType::FILL_PATTERN,
                                                                                                                    device->getNEODevice()->getDeviceBitfield()});
-            hostPtrMap.insert(std::make_pair(pattern, patternGfxAlloc));
         }
         void *patternGfxAllocPtr = patternGfxAlloc->getUnderlyingBuffer();
-
+        patternAllocations.push_back(patternGfxAlloc);
         uint64_t patternAllocPtr = reinterpret_cast<uintptr_t>(patternGfxAllocPtr);
         uint64_t patternAllocOffset = 0;
         uint64_t patternSizeToCopy = patternSize;

level_zero/core/source/device/device.h

@@ -133,6 +133,8 @@ struct Device : _ze_device_handle_t {
     virtual ze_result_t getCsrForOrdinalAndIndex(NEO::CommandStreamReceiver **csr, uint32_t ordinal, uint32_t index) = 0;
     virtual ze_result_t getCsrForLowPriority(NEO::CommandStreamReceiver **csr) = 0;
     virtual ze_result_t mapOrdinalForAvailableEngineGroup(uint32_t *ordinal) = 0;
+    virtual NEO::GraphicsAllocation *obtainReusableAllocation(size_t requiredSize, NEO::GraphicsAllocation::AllocationType type) = 0;
+    virtual void storeReusableAllocation(NEO::GraphicsAllocation &alloc) = 0;
 };
 
 } // namespace L0

level_zero/core/source/device/device_imp.cpp

@@ -800,6 +800,7 @@ void DeviceImp::releaseResources() {
     metricContext.reset();
     builtins.reset();
     cacheReservation.reset();
+    allocationsForReuse.freeAllGraphicsAllocations(neoDevice);
 
     if (getSourceLevelDebugger()) {
         getSourceLevelDebugger()->notifyDeviceDestruction();
@@ -918,6 +919,18 @@ NEO::GraphicsAllocation *DeviceImp::allocateMemoryFromHostPtr(const void *buffer
     return allocation;
 }
 
+NEO::GraphicsAllocation *DeviceImp::obtainReusableAllocation(size_t requiredSize, NEO::GraphicsAllocation::AllocationType type) {
+    auto alloc = allocationsForReuse.detachAllocation(requiredSize, nullptr, nullptr, type);
+    if (alloc == nullptr)
+        return nullptr;
+    else
+        return alloc.release();
+}
+
+void DeviceImp::storeReusableAllocation(NEO::GraphicsAllocation &alloc) {
+    allocationsForReuse.pushFrontOne(alloc);
+}
+
 ze_result_t DeviceImp::getCsrForOrdinalAndIndex(NEO::CommandStreamReceiver **csr, uint32_t ordinal, uint32_t index) {
     if (ordinal >= CommonConstants::engineGroupCount) {
         return ZE_RESULT_ERROR_INVALID_ARGUMENT;

level_zero/core/source/device/device_imp.h

@@ -8,6 +8,7 @@
 #pragma once
 
 #include "shared/source/helpers/topology_map.h"
+#include "shared/source/memory_manager/allocations_list.h"
 #include "shared/source/page_fault_manager/cpu_page_fault_manager.h"
 #include "shared/source/utilities/spinlock.h"
 
@@ -110,6 +111,8 @@ struct DeviceImp : public Device {
     ze_result_t getCsrForOrdinalAndIndex(NEO::CommandStreamReceiver **csr, uint32_t ordinal, uint32_t index) override;
     ze_result_t getCsrForLowPriority(NEO::CommandStreamReceiver **csr) override;
     ze_result_t mapOrdinalForAvailableEngineGroup(uint32_t *ordinal) override;
+    NEO::GraphicsAllocation *obtainReusableAllocation(size_t requiredSize, NEO::GraphicsAllocation::AllocationType type) override;
+    void storeReusableAllocation(NEO::GraphicsAllocation &alloc) override;
     NEO::Device *getActiveDevice() const;
 
     bool toPhysicalSliceId(const NEO::TopologyMap &topologyMap, uint32_t &slice, uint32_t &deviceIndex);
@@ -133,6 +136,7 @@ struct DeviceImp : public Device {
     NEO::SVMAllocsManager::MapBasedAllocationTracker peerAllocations;
     NEO::SpinLock peerAllocationsMutex;
     std::map<NEO::SvmAllocationData *, MemAdviseFlags> memAdviseSharedAllocations;
+    NEO::AllocationsList allocationsForReuse;
 
   protected:
     NEO::GraphicsAllocation *debugSurface = nullptr;

level_zero/core/test/unit_tests/mocks/mock_cmdlist.h

@@ -51,6 +51,7 @@ struct WhiteBox<::L0::CommandListCoreFamily<gfxCoreFamily>>
     using BaseClass::hostPtrMap;
     using BaseClass::indirectAllocationsAllowed;
     using BaseClass::initialize;
+    using BaseClass::patternAllocations;
     using BaseClass::requiredStreamState;
     using BaseClass::unifiedMemoryControls;
     using BaseClass::updateStreamProperties;

level_zero/core/test/unit_tests/mocks/mock_device.h

@@ -273,6 +273,12 @@ struct Mock<Device> : public Device {
     }
 
     void removeDebugSession() override {}
+
+    NEO::GraphicsAllocation *obtainReusableAllocation(size_t requiredSize, NEO::GraphicsAllocation::AllocationType type) override {
+        return nullptr;
+    }
+
+    void storeReusableAllocation(NEO::GraphicsAllocation &alloc) override {}
 };
 
 template <>

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

@@ -98,6 +98,7 @@ HWTEST2_F(CommandListCreate, givenHostAllocInMapWhenGettingAllocInRangeThenAlloc
     size_t allocSize = 0x1000;
     NEO::MockGraphicsAllocation alloc(const_cast<void *>(cpuPtr), gpuAddress, allocSize);
     commandList->hostPtrMap.insert(std::make_pair(cpuPtr, &alloc));
+    EXPECT_EQ(commandList->getHostPtrMap().size(), 1u);
 
     auto newBufferPtr = ptrOffset(cpuPtr, 0x10);
     auto newBufferSize = allocSize - 0x20;
@@ -114,6 +115,7 @@ HWTEST2_F(CommandListCreate, givenHostAllocInMapWhenSizeIsOutOfRangeThenNullPtrR
     size_t allocSize = 0x1000;
     NEO::MockGraphicsAllocation alloc(const_cast<void *>(cpuPtr), gpuAddress, allocSize);
     commandList->hostPtrMap.insert(std::make_pair(cpuPtr, &alloc));
+    EXPECT_EQ(commandList->getHostPtrMap().size(), 1u);
 
     auto newBufferPtr = ptrOffset(cpuPtr, 0x10);
     auto newBufferSize = allocSize + 0x20;
@@ -130,6 +132,7 @@ HWTEST2_F(CommandListCreate, givenHostAllocInMapWhenPtrIsOutOfRangeThenNullPtrRe
     size_t allocSize = 0x1000;
     NEO::MockGraphicsAllocation alloc(const_cast<void *>(cpuPtr), gpuAddress, allocSize);
     commandList->hostPtrMap.insert(std::make_pair(cpuPtr, &alloc));
+    EXPECT_EQ(commandList->getHostPtrMap().size(), 1u);
 
     auto newBufferPtr = reinterpret_cast<const void *>(gpuAddress - 0x100);
     auto newBufferSize = allocSize - 0x200;
@@ -146,6 +149,8 @@ HWTEST2_F(CommandListCreate, givenHostAllocInMapWhenGetHostPtrAllocCalledThenCor
     size_t allocSize = 0x1000;
     NEO::MockGraphicsAllocation alloc(const_cast<void *>(cpuPtr), gpuAddress, allocSize);
     commandList->hostPtrMap.insert(std::make_pair(cpuPtr, &alloc));
+    EXPECT_EQ(commandList->getHostPtrMap().size(), 1u);
+
     size_t expectedOffset = 0x10;
     auto newBufferPtr = ptrOffset(cpuPtr, expectedOffset);
     auto newBufferSize = allocSize - 0x20;
@@ -162,6 +167,7 @@ HWTEST2_F(CommandListCreate, givenHostAllocInMapWhenPtrIsInMapThenAllocationRetu
     size_t allocSize = 0x1000;
     NEO::MockGraphicsAllocation alloc(const_cast<void *>(cpuPtr), gpuAddress, allocSize);
     commandList->hostPtrMap.insert(std::make_pair(cpuPtr, &alloc));
+    EXPECT_EQ(commandList->getHostPtrMap().size(), 1u);
 
     auto newBufferPtr = cpuPtr;
     auto newBufferSize = allocSize - 0x20;
@@ -178,6 +184,7 @@ HWTEST2_F(CommandListCreate, givenHostAllocInMapWhenPtrIsInMapButWithBiggerSizeT
     size_t allocSize = 0x1000;
     NEO::MockGraphicsAllocation alloc(const_cast<void *>(cpuPtr), gpuAddress, allocSize);
     commandList->hostPtrMap.insert(std::make_pair(cpuPtr, &alloc));
+    EXPECT_EQ(commandList->getHostPtrMap().size(), 1u);
 
     auto newBufferPtr = cpuPtr;
     auto newBufferSize = allocSize + 0x20;
@@ -194,6 +201,7 @@ HWTEST2_F(CommandListCreate, givenHostAllocInMapWhenPtrLowerThanAnyInMapThenNull
     size_t allocSize = 0x1000;
     NEO::MockGraphicsAllocation alloc(const_cast<void *>(cpuPtr), gpuAddress, allocSize);
     commandList->hostPtrMap.insert(std::make_pair(cpuPtr, &alloc));
+    EXPECT_EQ(commandList->getHostPtrMap().size(), 1u);
 
     auto newBufferPtr = reinterpret_cast<const void *>(gpuAddress - 0x10);
     auto newBufferSize = allocSize - 0x20;

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

@@ -139,7 +139,7 @@ HWTEST2_F(AppendFillFixture,
 }
 
 HWTEST2_F(AppendFillFixture,
-          givenTwoCallsToAppendMemoryFillWithSamePatternThenAllocationIsAddedtoHostPtrMapOnlyOnce, Platforms) {
+          givenTwoCallsToAppendMemoryFillWithSamePatternThenAllocationIsCreatedForEachCall, Platforms) {
     using GfxFamily = typename NEO::GfxFamilyMapper<gfxCoreFamily>::GfxFamily;
 
     auto commandList = std::make_unique<WhiteBox<MockCommandList<gfxCoreFamily>>>();
@@ -147,21 +147,21 @@ HWTEST2_F(AppendFillFixture,
 
     ze_result_t result = commandList->appendMemoryFill(dstPtr, pattern, 4, allocSize, nullptr, 0, nullptr);
     EXPECT_EQ(ZE_RESULT_SUCCESS, result);
-    size_t hostPtrMapSize = commandList->getHostPtrMap().size();
+    size_t patternAllocationsVectorSize = commandList->patternAllocations.size();
-    EXPECT_EQ(hostPtrMapSize, 1u);
+    EXPECT_EQ(patternAllocationsVectorSize, 1u);
 
     uint8_t *newDstPtr = new uint8_t[allocSize];
     result = commandList->appendMemoryFill(newDstPtr, pattern, patternSize, allocSize, nullptr, 0, nullptr);
     EXPECT_EQ(ZE_RESULT_SUCCESS, result);
-    size_t newHostPtrMapSize = commandList->getHostPtrMap().size();
+    size_t newPatternAllocationsVectorSize = commandList->patternAllocations.size();
 
-    EXPECT_EQ(hostPtrMapSize, newHostPtrMapSize);
+    EXPECT_GT(newPatternAllocationsVectorSize, patternAllocationsVectorSize);
 
     delete[] newDstPtr;
 }
 
 HWTEST2_F(AppendFillFixture,
-          givenTwoCallsToAppendMemoryFillWithDifferentPatternsThenHostPtrSizeIncrementsByOne, Platforms) {
+          givenTwoCallsToAppendMemoryFillWithDifferentPatternsThenAllocationIsCreatedForEachPattern, Platforms) {
     using GfxFamily = typename NEO::GfxFamilyMapper<gfxCoreFamily>::GfxFamily;
 
     auto commandList = std::make_unique<WhiteBox<MockCommandList<gfxCoreFamily>>>();
@@ -169,15 +169,15 @@ HWTEST2_F(AppendFillFixture,
 
     ze_result_t result = commandList->appendMemoryFill(dstPtr, pattern, 4, allocSize, nullptr, 0, nullptr);
     EXPECT_EQ(ZE_RESULT_SUCCESS, result);
-    size_t hostPtrMapSize = commandList->getHostPtrMap().size();
+    size_t patternAllocationsVectorSize = commandList->patternAllocations.size();
-    EXPECT_EQ(hostPtrMapSize, 1u);
+    EXPECT_EQ(patternAllocationsVectorSize, 1u);
 
     uint8_t newPattern[patternSize] = {1, 2, 3, 4};
     result = commandList->appendMemoryFill(dstPtr, newPattern, patternSize, allocSize, nullptr, 0, nullptr);
     EXPECT_EQ(ZE_RESULT_SUCCESS, result);
-    size_t newHostPtrMapSize = commandList->getHostPtrMap().size();
+    size_t newPatternAllocationsVectorSize = commandList->patternAllocations.size();
 
-    EXPECT_EQ(hostPtrMapSize + 1u, newHostPtrMapSize);
+    EXPECT_EQ(patternAllocationsVectorSize + 1u, newPatternAllocationsVectorSize);
 }
 
 HWTEST2_F(AppendFillFixture,

level_zero/core/test/unit_tests/sources/device/test_device.cpp

@@ -424,6 +424,62 @@ TEST_F(DeviceTest, givenEmptySVmAllocStorageWhenAllocateMemoryFromHostPtrThenVal
     neoDevice->getMemoryManager()->freeGraphicsMemory(allocation);
 }
 
+TEST_F(DeviceTest, givenNonEmptyAllocationsListWhenRequestingAllocationSmallerOrEqualInSizeThenAllocationFromListIsReturned) {
+    auto deviceImp = static_cast<DeviceImp *>(device);
+    constexpr auto dataSize = 1024u;
+    auto data = std::make_unique<int[]>(dataSize);
+
+    constexpr auto allocationSize = sizeof(int) * dataSize;
+
+    auto allocation = device->getDriverHandle()->getMemoryManager()->allocateGraphicsMemoryWithProperties({device->getNEODevice()->getRootDeviceIndex(),
+                                                                                                           allocationSize,
+                                                                                                           NEO::GraphicsAllocation::AllocationType::FILL_PATTERN,
+                                                                                                           neoDevice->getDeviceBitfield()});
+    device->storeReusableAllocation(*allocation);
+    EXPECT_FALSE(deviceImp->allocationsForReuse.peekIsEmpty());
+    auto obtaindedAllocation = device->obtainReusableAllocation(dataSize, NEO::GraphicsAllocation::AllocationType::FILL_PATTERN);
+    EXPECT_TRUE(deviceImp->allocationsForReuse.peekIsEmpty());
+    EXPECT_NE(nullptr, obtaindedAllocation);
+    EXPECT_EQ(allocation, obtaindedAllocation);
+    neoDevice->getMemoryManager()->freeGraphicsMemory(allocation);
+}
+
+TEST_F(DeviceTest, givenNonEmptyAllocationsListWhenRequestingAllocationBiggerInSizeThenNullptrIsReturned) {
+    auto deviceImp = static_cast<DeviceImp *>(device);
+    constexpr auto dataSize = 1024u;
+    auto data = std::make_unique<int[]>(dataSize);
+
+    constexpr auto allocationSize = sizeof(int) * dataSize;
+
+    auto allocation = device->getDriverHandle()->getMemoryManager()->allocateGraphicsMemoryWithProperties({device->getNEODevice()->getRootDeviceIndex(),
+                                                                                                           allocationSize,
+                                                                                                           NEO::GraphicsAllocation::AllocationType::FILL_PATTERN,
+                                                                                                           neoDevice->getDeviceBitfield()});
+    device->storeReusableAllocation(*allocation);
+    EXPECT_FALSE(deviceImp->allocationsForReuse.peekIsEmpty());
+    auto obtaindedAllocation = device->obtainReusableAllocation(4 * dataSize + 1u, NEO::GraphicsAllocation::AllocationType::FILL_PATTERN);
+    EXPECT_EQ(nullptr, obtaindedAllocation);
+    EXPECT_FALSE(deviceImp->allocationsForReuse.peekIsEmpty());
+}
+
+TEST_F(DeviceTest, givenNonEmptyAllocationsListAndUnproperAllocationTypeWhenRequestingAllocationThenNullptrIsReturned) {
+    auto deviceImp = static_cast<DeviceImp *>(device);
+    constexpr auto dataSize = 1024u;
+    auto data = std::make_unique<int[]>(dataSize);
+
+    constexpr auto allocationSize = sizeof(int) * dataSize;
+
+    auto allocation = device->getDriverHandle()->getMemoryManager()->allocateGraphicsMemoryWithProperties({device->getNEODevice()->getRootDeviceIndex(),
+                                                                                                           allocationSize,
+                                                                                                           NEO::GraphicsAllocation::AllocationType::BUFFER,
+                                                                                                           neoDevice->getDeviceBitfield()});
+    device->storeReusableAllocation(*allocation);
+    EXPECT_FALSE(deviceImp->allocationsForReuse.peekIsEmpty());
+    auto obtaindedAllocation = device->obtainReusableAllocation(4 * dataSize + 1u, NEO::GraphicsAllocation::AllocationType::FILL_PATTERN);
+    EXPECT_EQ(nullptr, obtaindedAllocation);
+    EXPECT_FALSE(deviceImp->allocationsForReuse.peekIsEmpty());
+}
+
 struct DeviceHostPointerTest : public ::testing::Test {
     void SetUp() override {
         executionEnvironment = new NEO::ExecutionEnvironment();

opencl/test/unit_test/command_stream/command_stream_receiver_flush_task_2_tests.cpp

@@ -914,7 +914,8 @@ HWTEST_F(CommandStreamReceiverFlushTaskTests, GivenForced32BitAllocationsModeSto
         auto newScratchAllocation = commandStreamReceiver->getScratchAllocation();
         EXPECT_NE(scratchAllocation, newScratchAllocation); // Allocation changed
 
-        std::unique_ptr<GraphicsAllocation> allocationTemporary = commandStreamReceiver->getTemporaryAllocations().detachAllocation(0, nullptr, *commandStreamReceiver, GraphicsAllocation::AllocationType::SCRATCH_SURFACE);
+        CommandStreamReceiver *csrPtr = reinterpret_cast<CommandStreamReceiver *>(commandStreamReceiver);
+        std::unique_ptr<GraphicsAllocation> allocationTemporary = commandStreamReceiver->getTemporaryAllocations().detachAllocation(0, nullptr, csrPtr, GraphicsAllocation::AllocationType::SCRATCH_SURFACE);
 
         EXPECT_EQ(scratchAllocation, allocationTemporary.get());
         pDevice->getMemoryManager()->freeGraphicsMemory(allocationTemporary.release());

shared/source/memory_manager/CMakeLists.txt

@@ -9,6 +9,7 @@ set(NEO_CORE_MEMORY_MANAGER
     ${CMAKE_CURRENT_SOURCE_DIR}/address_mapper.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/address_mapper.h
     ${CMAKE_CURRENT_SOURCE_DIR}/allocations_list.h
+    ${CMAKE_CURRENT_SOURCE_DIR}/allocations_list.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/alignment_selector.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/alignment_selector.h
     ${CMAKE_CURRENT_SOURCE_DIR}/allocation_properties.h

shared/source/memory_manager/allocations_list.cpp

@@ -0,0 +1,72 @@
+/*
+ * Copyright (C) 2021 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#include "shared/source/memory_manager/allocations_list.h"
+
+#include "shared/source/command_stream/command_stream_receiver.h"
+
+namespace NEO {
+
+struct ReusableAllocationRequirements {
+    size_t requiredMinimalSize;
+    volatile uint32_t *csrTagAddress;
+    GraphicsAllocation::AllocationType allocationType;
+    uint32_t contextId;
+    const void *requiredPtr;
+};
+
+AllocationsList::AllocationsList(AllocationUsage allocationUsage)
+    : allocationUsage(allocationUsage) {}
+
+AllocationsList::AllocationsList()
+    : allocationUsage(REUSABLE_ALLOCATION) {}
+
+std::unique_ptr<GraphicsAllocation> AllocationsList::detachAllocation(size_t requiredMinimalSize, const void *requiredPtr, CommandStreamReceiver *commandStreamReceiver, GraphicsAllocation::AllocationType allocationType) {
+    ReusableAllocationRequirements req;
+    req.requiredMinimalSize = requiredMinimalSize;
+    commandStreamReceiver == nullptr ? req.csrTagAddress = nullptr : req.csrTagAddress = commandStreamReceiver->getTagAddress();
+    req.allocationType = allocationType;
+    commandStreamReceiver == nullptr ? req.contextId = UINT32_MAX : req.contextId = commandStreamReceiver->getOsContext().getContextId();
+    req.requiredPtr = requiredPtr;
+    GraphicsAllocation *a = nullptr;
+    GraphicsAllocation *retAlloc = processLocked<AllocationsList, &AllocationsList::detachAllocationImpl>(a, static_cast<void *>(&req));
+    return std::unique_ptr<GraphicsAllocation>(retAlloc);
+}
+
+GraphicsAllocation *AllocationsList::detachAllocationImpl(GraphicsAllocation *, void *data) {
+    ReusableAllocationRequirements *req = static_cast<ReusableAllocationRequirements *>(data);
+    auto *curr = head;
+    while (curr != nullptr) {
+        if ((req->allocationType == curr->getAllocationType()) &&
+            (curr->getUnderlyingBufferSize() >= req->requiredMinimalSize)) {
+            if (req->csrTagAddress == nullptr) {
+                return removeOneImpl(curr, nullptr);
+            }
+            if ((this->allocationUsage == TEMPORARY_ALLOCATION || *req->csrTagAddress >= curr->getTaskCount(req->contextId)) &&
+                (req->requiredPtr == nullptr || req->requiredPtr == curr->getUnderlyingBuffer())) {
+                if (this->allocationUsage == TEMPORARY_ALLOCATION) {
+                    // We may not have proper task count yet, so set notReady to avoid releasing in a different thread
+                    curr->updateTaskCount(CompletionStamp::notReady, req->contextId);
+                }
+                return removeOneImpl(curr, nullptr);
+            }
+        }
+        curr = curr->next;
+    }
+    return nullptr;
+}
+
+void AllocationsList::freeAllGraphicsAllocations(Device *neoDevice) {
+    auto *curr = head;
+    while (curr != nullptr) {
+        auto currNext = curr->next;
+        neoDevice->getMemoryManager()->freeGraphicsMemory(curr);
+        curr = currNext;
+    }
+    head = nullptr;
+}
+} // namespace NEO

shared/source/memory_manager/allocations_list.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2020 Intel Corporation
+ * Copyright (C) 2018-2021 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -17,7 +17,9 @@ class CommandStreamReceiver;
 class AllocationsList : public IDList<GraphicsAllocation, true, true> {
   public:
     AllocationsList(AllocationUsage allocationUsage);
-    std::unique_ptr<GraphicsAllocation> detachAllocation(size_t requiredMinimalSize, const void *requiredPtr, CommandStreamReceiver &commandStreamReceiver, GraphicsAllocation::AllocationType allocationType);
+    AllocationsList();
+    std::unique_ptr<GraphicsAllocation> detachAllocation(size_t requiredMinimalSize, const void *requiredPtr, CommandStreamReceiver *commandStreamReceiver, GraphicsAllocation::AllocationType allocationType);
+    void freeAllGraphicsAllocations(Device *neoDevice);
 
   private:
     GraphicsAllocation *detachAllocationImpl(GraphicsAllocation *, void *);

shared/source/memory_manager/internal_allocation_storage.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2020 Intel Corporation
+ * Copyright (C) 2018-2021 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -66,57 +66,15 @@ void InternalAllocationStorage::freeAllocationsList(uint32_t waitTaskCount, Allo
 }
 
 std::unique_ptr<GraphicsAllocation> InternalAllocationStorage::obtainReusableAllocation(size_t requiredSize, GraphicsAllocation::AllocationType allocationType) {
-    auto allocation = allocationsForReuse.detachAllocation(requiredSize, nullptr, commandStreamReceiver, allocationType);
+    auto allocation = allocationsForReuse.detachAllocation(requiredSize, nullptr, &commandStreamReceiver, allocationType);
     return allocation;
 }
 
 std::unique_ptr<GraphicsAllocation> InternalAllocationStorage::obtainTemporaryAllocationWithPtr(size_t requiredSize, const void *requiredPtr, GraphicsAllocation::AllocationType allocationType) {
-    auto allocation = temporaryAllocations.detachAllocation(requiredSize, requiredPtr, commandStreamReceiver, allocationType);
+    auto allocation = temporaryAllocations.detachAllocation(requiredSize, requiredPtr, &commandStreamReceiver, allocationType);
     return allocation;
 }
 
-struct ReusableAllocationRequirements {
-    size_t requiredMinimalSize;
-    volatile uint32_t *csrTagAddress;
-    GraphicsAllocation::AllocationType allocationType;
-    uint32_t contextId;
-    const void *requiredPtr;
-};
-
-AllocationsList::AllocationsList(AllocationUsage allocationUsage)
-    : allocationUsage(allocationUsage) {}
-
-std::unique_ptr<GraphicsAllocation> AllocationsList::detachAllocation(size_t requiredMinimalSize, const void *requiredPtr, CommandStreamReceiver &commandStreamReceiver, GraphicsAllocation::AllocationType allocationType) {
-    ReusableAllocationRequirements req;
-    req.requiredMinimalSize = requiredMinimalSize;
-    req.csrTagAddress = commandStreamReceiver.getTagAddress();
-    req.allocationType = allocationType;
-    req.contextId = commandStreamReceiver.getOsContext().getContextId();
-    req.requiredPtr = requiredPtr;
-    GraphicsAllocation *a = nullptr;
-    GraphicsAllocation *retAlloc = processLocked<AllocationsList, &AllocationsList::detachAllocationImpl>(a, static_cast<void *>(&req));
-    return std::unique_ptr<GraphicsAllocation>(retAlloc);
-}
-
-GraphicsAllocation *AllocationsList::detachAllocationImpl(GraphicsAllocation *, void *data) {
-    ReusableAllocationRequirements *req = static_cast<ReusableAllocationRequirements *>(data);
-    auto *curr = head;
-    while (curr != nullptr) {
-        if ((req->allocationType == curr->getAllocationType()) &&
-            (curr->getUnderlyingBufferSize() >= req->requiredMinimalSize) &&
-            (this->allocationUsage == TEMPORARY_ALLOCATION || *req->csrTagAddress >= curr->getTaskCount(req->contextId)) &&
-            (req->requiredPtr == nullptr || req->requiredPtr == curr->getUnderlyingBuffer())) {
-            if (this->allocationUsage == TEMPORARY_ALLOCATION) {
-                // We may not have proper task count yet, so set notReady to avoid releasing in a different thread
-                curr->updateTaskCount(CompletionStamp::notReady, req->contextId);
-            }
-            return removeOneImpl(curr, nullptr);
-        }
-        curr = curr->next;
-    }
-    return nullptr;
-}
-
 DeviceBitfield InternalAllocationStorage::getDeviceBitfield() const {
     return commandStreamReceiver.getOsContext().getDeviceBitfield();
 }

shared/source/memory_manager/internal_allocation_storage.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2020 Intel Corporation
+ * Copyright (C) 2018-2021 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *