From 7176e609a18f228bae1713aa74437f4afb1f8607 Mon Sep 17 00:00:00 2001
From: Lukasz Jobczyk <lukasz.jobczyk@intel.com>
Date: Wed, 17 Nov 2021 15:31:39 +0000
Subject: [PATCH] Do not wait for task count when free empty allocation list

Signed-off-by: Lukasz Jobczyk <lukasz.jobczyk@intel.com>
---
 .../command_stream_receiver.cpp               |  5 ++-
 .../command_stream_receiver_tests.cpp         | 33 +++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/shared/source/command_stream/command_stream_receiver.cpp b/shared/source/command_stream/command_stream_receiver.cpp
index a236c5e5d7..4a5a4b8881 100644
--- a/shared/source/command_stream/command_stream_receiver.cpp
+++ b/shared/source/command_stream/command_stream_receiver.cpp
@@ -151,7 +151,10 @@ void CommandStreamReceiver::waitForTaskCount(uint32_t requiredTaskCount) {
 }
 
 void CommandStreamReceiver::waitForTaskCountAndCleanAllocationList(uint32_t requiredTaskCount, uint32_t allocationUsage) {
-    this->CommandStreamReceiver::waitForTaskCount(requiredTaskCount);
+    auto &list = allocationUsage == TEMPORARY_ALLOCATION ? internalAllocationStorage->getTemporaryAllocations() : internalAllocationStorage->getAllocationsForReuse();
+    if (!list.peekIsEmpty()) {
+        this->CommandStreamReceiver::waitForTaskCount(requiredTaskCount);
+    }
     internalAllocationStorage->cleanAllocationList(requiredTaskCount, allocationUsage);
 }
 
diff --git a/shared/test/unit_test/command_stream/command_stream_receiver_tests.cpp b/shared/test/unit_test/command_stream/command_stream_receiver_tests.cpp
index 5fdbc3315f..7566161834 100644
--- a/shared/test/unit_test/command_stream/command_stream_receiver_tests.cpp
+++ b/shared/test/unit_test/command_stream/command_stream_receiver_tests.cpp
@@ -1044,6 +1044,15 @@ TEST(CommandStreamReceiverSimpleTest, givenMultipleActivePartitionsWhenWaitingFo
     executionEnvironment.initializeMemoryManager();
     DeviceBitfield deviceBitfield(0b11);
     MockCommandStreamReceiver csr(executionEnvironment, 0, deviceBitfield);
+    auto osContext = std::unique_ptr<OsContext>(OsContext::create(nullptr, 0,
+                                                                  EngineDescriptorHelper::getDefaultDescriptor({aub_stream::ENGINE_BCS, EngineUsage::Regular})));
+    csr.setupContext(*osContext);
+
+    auto hostPtr = reinterpret_cast<void *>(0x1234);
+    size_t size = 100;
+    auto temporaryAllocation = std::make_unique<MemoryAllocation>(0, GraphicsAllocation::AllocationType::EXTERNAL_HOST_PTR, hostPtr, size, 0, MemoryPool::System4KBPages, MemoryManager::maxOsContextCount);
+    temporaryAllocation->updateTaskCount(0u, 0u);
+    csr.getInternalAllocationStorage()->storeAllocationWithTaskCount(std::move(temporaryAllocation), TEMPORARY_ALLOCATION, 2u);
 
     csr.mockTagAddress[0] = 0u;
     csr.mockTagAddress[2] = 0u;
@@ -1065,6 +1074,30 @@ TEST(CommandStreamReceiverSimpleTest, givenMultipleActivePartitionsWhenWaitingFo
     CpuIntrinsicsTests::pauseAddress = nullptr;
 }
 
+TEST(CommandStreamReceiverSimpleTest, givenEmptyTemporaryAllocationListWhenWaitingForTaskCountForCleaningTemporaryAllocationsThenDoNotWait) {
+    MockExecutionEnvironment executionEnvironment;
+    executionEnvironment.prepareRootDeviceEnvironments(1);
+    executionEnvironment.initializeMemoryManager();
+    DeviceBitfield deviceBitfield(1);
+    MockCommandStreamReceiver csr(executionEnvironment, 0, deviceBitfield);
+
+    csr.mockTagAddress[0] = 0u;
+    csr.taskCount = 3u;
+
+    VariableBackup<volatile uint32_t *> backupPauseAddress(&CpuIntrinsicsTests::pauseAddress);
+    VariableBackup<uint32_t> backupPauseValue(&CpuIntrinsicsTests::pauseValue);
+    VariableBackup<uint32_t> backupPauseOffset(&CpuIntrinsicsTests::pauseOffset);
+
+    CpuIntrinsicsTests::pauseAddress = &csr.mockTagAddress[0];
+    CpuIntrinsicsTests::pauseValue = 3u;
+
+    CpuIntrinsicsTests::pauseCounter = 0;
+    csr.waitForTaskCountAndCleanTemporaryAllocationList(3u);
+    EXPECT_EQ(0u, CpuIntrinsicsTests::pauseCounter);
+
+    CpuIntrinsicsTests::pauseAddress = nullptr;
+}
+
 TEST(CommandStreamReceiverMultiContextTests, givenMultipleCsrsWhenSameResourcesAreUsedThenResidencyIsProperlyHandled) {
     std::unique_ptr<MockDevice> device(MockDevice::createWithNewExecutionEnvironment<MockDevice>(defaultHwInfo.get(), 0u));