From eba306c0991d132b1429df5156cbf84ba5a952a9 Mon Sep 17 00:00:00 2001
From: Igor Venevtsev <igor.venevtsev@intel.com>
Date: Wed, 5 Jul 2023 16:01:51 +0000
Subject: [PATCH] fix: properly set systemMemoryForced flag for secondary
 command buffers

Due to this flag was not properly handled on Windows, command buffer
allocations were never reused in immediate command lists in case of
host secondary buffers. This lead to huge host memory consumption
and performance degradation

Related-To: NEO-8072

Signed-off-by: Igor Venevtsev <igor.venevtsev@intel.com>
---
 .../source/command_container/cmdcontainer.cpp |  7 +++-
 .../source/command_container/cmdcontainer.h   |  2 +-
 .../command_container_tests.cpp               | 40 +++++++++++++++++++
 3 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/shared/source/command_container/cmdcontainer.cpp b/shared/source/command_container/cmdcontainer.cpp
index c87ba06f11..a3c80148ae 100644
--- a/shared/source/command_container/cmdcontainer.cpp
+++ b/shared/source/command_container/cmdcontainer.cpp
@@ -482,7 +482,12 @@ GraphicsAllocation *CommandContainer::allocateCommandBuffer(bool forceHostMemory
                                     device->getDeviceBitfield()};
     properties.flags.forceSystemMemory = forceHostMemory && this->useSecondaryCommandStream;
 
-    return device->getMemoryManager()->allocateGraphicsMemoryWithProperties(properties);
+    auto commandBufferAllocation = device->getMemoryManager()->allocateGraphicsMemoryWithProperties(properties);
+    if (commandBufferAllocation) {
+        commandBufferAllocation->storageInfo.systemMemoryForced = properties.flags.forceSystemMemory;
+    }
+
+    return commandBufferAllocation;
 }
 
 void CommandContainer::fillReusableAllocationLists() {
diff --git a/shared/source/command_container/cmdcontainer.h b/shared/source/command_container/cmdcontainer.h
index 5a14a4d0cc..8a5e4dda1f 100644
--- a/shared/source/command_container/cmdcontainer.h
+++ b/shared/source/command_container/cmdcontainer.h
@@ -150,7 +150,7 @@ class CommandContainer : public NonCopyableOrMovableClass {
     GraphicsAllocation *reuseExistingCmdBuffer();
     GraphicsAllocation *reuseExistingCmdBuffer(bool forceHostMemory);
     GraphicsAllocation *allocateCommandBuffer();
-    GraphicsAllocation *allocateCommandBuffer(bool forceHostMemory);
+    MOCKABLE_VIRTUAL GraphicsAllocation *allocateCommandBuffer(bool forceHostMemory);
     void setCmdBuffer(GraphicsAllocation *cmdBuffer);
     void addCurrentCommandBufferToReusableAllocationList();
 
diff --git a/shared/test/unit_test/command_container/command_container_tests.cpp b/shared/test/unit_test/command_container/command_container_tests.cpp
index 06d8dc1d5f..41b4b6849a 100644
--- a/shared/test/unit_test/command_container/command_container_tests.cpp
+++ b/shared/test/unit_test/command_container/command_container_tests.cpp
@@ -37,6 +37,13 @@ class MyMockCommandContainer : public CommandContainer {
     using CommandContainer::getAlignedCmdBufferSize;
     using CommandContainer::immediateReusableAllocationList;
     using CommandContainer::secondaryCommandStreamForImmediateCmdList;
+
+    GraphicsAllocation *allocateCommandBuffer(bool forceHostMemory) override {
+        allocateCommandBufferCalled[!!forceHostMemory]++;
+        return CommandContainer::allocateCommandBuffer(forceHostMemory);
+    }
+
+    uint32_t allocateCommandBufferCalled[2] = {0, 0};
 };
 
 struct CommandContainerHeapStateTests : public ::testing::Test {
@@ -1432,6 +1439,39 @@ TEST_F(CommandContainerTest, givenCreateSecondaryCmdBufferInHostMemWhenFillReusa
     allocList.freeAllGraphicsAllocations(pDevice);
 }
 
+TEST_F(CommandContainerTest, givenSecondCmdContainerCreatedAfterFirstCmdContainerDestroyedAndReusableAllocationsListUsedThenCommandBuffersAllocationsAreReused) {
+    auto cmdContainer = std::make_unique<MyMockCommandContainer>();
+
+    AllocationsList allocList;
+    cmdContainer->initialize(pDevice, &allocList, HeapSize::defaultHeapSize, true, true);
+
+    EXPECT_EQ(1u, cmdContainer->allocateCommandBufferCalled[0]); // forceHostMemory = 0
+    EXPECT_EQ(1u, cmdContainer->allocateCommandBufferCalled[1]); // forceHostMemory = 1
+    EXPECT_TRUE(allocList.peekIsEmpty());
+
+    cmdContainer.reset();
+    EXPECT_FALSE(allocList.peekIsEmpty());
+
+    cmdContainer = std::make_unique<MyMockCommandContainer>();
+    cmdContainer->initialize(pDevice, &allocList, HeapSize::defaultHeapSize, true, true);
+
+    EXPECT_EQ(0u, cmdContainer->allocateCommandBufferCalled[0]); // forceHostMemory = 0
+    EXPECT_EQ(0u, cmdContainer->allocateCommandBufferCalled[1]); // forceHostMemory = 1
+
+    cmdContainer.reset();
+    allocList.freeAllGraphicsAllocations(pDevice);
+}
+
+TEST_F(CommandContainerTest, givenAllocateCommandBufferInHostMemoryCalledThenForceSystemMemoryFlagSetInAllocationStorageInfo) {
+    auto cmdContainer = std::make_unique<MyMockCommandContainer>();
+    cmdContainer->initialize(pDevice, nullptr, HeapSize::defaultHeapSize, true, true);
+
+    auto commandBufferAllocation = cmdContainer->allocateCommandBuffer(true /*forceHostMemory*/);
+    EXPECT_TRUE(commandBufferAllocation->storageInfo.systemMemoryForced);
+    pDevice->getMemoryManager()->freeGraphicsMemory(commandBufferAllocation);
+    cmdContainer.reset();
+}
+
 TEST_F(CommandContainerTest, givenCmdContainerWhenFillReusableAllocationListsWithSharedHeapsEnabledThenOnlyOneHeapFilled) {
     DebugManagerStateRestore dbgRestore;
     DebugManager.flags.SetAmountOfReusableAllocations.set(1);