[fix] invalidate state caches only for heaps used by initialized context

This is number of small tweaks to state cache invalidation:
1. Invalidate if heap was actually created.
2. Check if os context was actually initialized.
3. Heap allocation was actually submitted, as it might attain zero task count
value, when allocation is stored in csr internal storage, as csr wasn't used,
but the csr task count being zero is assigned to heap allocation when stored.

Related-To: NEO-5055

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

level_zero/core/source/cmdlist/cmdlist_imp.cpp

@@ -53,16 +53,21 @@ ze_result_t CommandListImp::destroy() {
         this->stateBaseAddressTracking &&
         this->cmdListHeapAddressModel == NEO::HeapAddressModel::PrivateHeaps) {
 
-        auto memoryManager = device->getNEODevice()->getMemoryManager();
+        auto surfaceStateHeap = this->commandContainer.getIndirectHeap(NEO::HeapType::SURFACE_STATE);
+        if (surfaceStateHeap) {
+            auto heapAllocation = surfaceStateHeap->getGraphicsAllocation();
+            NEO::WaitParams defaultWaitParams{false, false, NEO::TimeoutControls::maxTimeout};
 
-        auto heapAllocation = this->commandContainer.getIndirectHeap(NEO::HeapType::SURFACE_STATE)->getGraphicsAllocation();
-        for (auto &engine : memoryManager->getRegisteredEngines()) {
+            auto &deviceEngines = device->getNEODevice()->getMemoryManager()->getRegisteredEngines();
+            for (auto &engine : deviceEngines) {
                 if (NEO::EngineHelpers::isComputeEngine(engine.getEngineType())) {
                     auto contextId = engine.osContext->getContextId();
-
-                if (heapAllocation->isUsedByOsContext(contextId)) {
-                    engine.commandStreamReceiver->sendRenderStateCacheFlush();
-                    engine.commandStreamReceiver->waitForCompletionWithTimeout(NEO::WaitParams{false, false, NEO::TimeoutControls::maxTimeout}, engine.commandStreamReceiver->peekTaskCount());
+                    if (heapAllocation->isUsedByOsContext(contextId) && engine.osContext->isInitialized() && heapAllocation->getTaskCount(contextId) > 0) {
+                        auto submissionStatus = engine.commandStreamReceiver->sendRenderStateCacheFlush();
+                        if (submissionStatus == NEO::SubmissionStatus::SUCCESS) {
+                            engine.commandStreamReceiver->waitForCompletionWithTimeout(defaultWaitParams, engine.commandStreamReceiver->peekTaskCount());
+                        }
+                    }
                 }
             }
         }

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

@@ -38,8 +38,11 @@ struct MemoryManagerCommandListCreateNegativeTest : public NEO::MockMemoryManage
     bool forceFailureInPrimaryAllocation = false;
 };
 
-struct CommandListCreateNegativeTest : public ::testing::Test {
-    void SetUp() override {
+template <int32_t stateBaseAddressTracking>
+struct CommandListCreateNegativeFixture {
+    void setUp() {
+        DebugManager.flags.EnableStateBaseAddressTracking.set(stateBaseAddressTracking);
+
         executionEnvironment = new NEO::ExecutionEnvironment();
         executionEnvironment->prepareRootDeviceEnvironments(numRootDevices);
         for (uint32_t i = 0; i < numRootDevices; i++) {
@@ -61,9 +64,11 @@ struct CommandListCreateNegativeTest : public ::testing::Test {
 
         device = driverHandle->devices[0];
     }
-    void TearDown() override {
+    void tearDown() {
     }
 
+    DebugManagerStateRestore restorer;
+
     NEO::ExecutionEnvironment *executionEnvironment = nullptr;
     std::unique_ptr<Mock<L0::DriverHandleImp>> driverHandle;
     NEO::MockDevice *neoDevice = nullptr;
@@ -72,6 +77,8 @@ struct CommandListCreateNegativeTest : public ::testing::Test {
     const uint32_t numRootDevices = 1u;
 };
 
+using CommandListCreateNegativeTest = Test<CommandListCreateNegativeFixture<0>>;
+
 TEST_F(CommandListCreateNegativeTest, whenDeviceAllocationFailsDuringCommandListCreateThenAppropriateValueIsReturned) {
     ze_result_t returnValue;
     memoryManager->forceFailureInPrimaryAllocation = true;
@@ -80,6 +87,23 @@ TEST_F(CommandListCreateNegativeTest, whenDeviceAllocationFailsDuringCommandList
     ASSERT_EQ(nullptr, commandList);
 }
 
+using CommandListCreateNegativeStateBaseAddressTest = Test<CommandListCreateNegativeFixture<1>>;
+
+HWTEST2_F(CommandListCreateNegativeStateBaseAddressTest,
+          GivenStateBaseAddressTrackingWhenDeviceAllocationFailsDuringCommandListCreateThenCacheIsNotInvalidatedAndAppropriateValueIsReturned,
+          IsAtLeastSkl) {
+    auto &csr = neoDevice->getUltCommandStreamReceiver<FamilyType>();
+    auto &csrStream = csr.commandStream;
+
+    ze_result_t returnValue;
+    memoryManager->forceFailureInPrimaryAllocation = true;
+    std::unique_ptr<L0::CommandList> commandList(CommandList::create(productFamily, device, NEO::EngineGroupType::RenderCompute, 0u, returnValue));
+    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY, returnValue);
+    ASSERT_EQ(nullptr, commandList);
+
+    EXPECT_EQ(0u, csrStream.getUsed());
+}
+
 TEST_F(CommandListCreateNegativeTest, whenDeviceAllocationFailsDuringCommandListImmediateCreateThenAppropriateValueIsReturned) {
     ze_result_t returnValue;
     const ze_command_queue_desc_t desc = {};

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

@@ -12,6 +12,7 @@
 #include "shared/test/common/cmd_parse/gen_cmd_parse.h"
 #include "shared/test/common/helpers/unit_test_helper.h"
 #include "shared/test/common/libult/ult_command_stream_receiver.h"
+#include "shared/test/common/mocks/mock_os_context.h"
 #include "shared/test/common/test_macros/hw_test.h"
 
 #include "level_zero/core/source/kernel/kernel_imp.h"
@@ -2584,5 +2585,80 @@ HWTEST2_F(CommandListStateBaseAddressPrivateHeapTest,
     EXPECT_TRUE(NEO::UnitTestHelper<FamilyType>::findStateCacheFlushPipeControl(csrStream));
 }
 
+HWTEST2_F(CommandListStateBaseAddressPrivateHeapTest,
+          givenCommandListUsingPrivateSurfaceHeapWhenOsContextNotInitializedAndCommandListDestroyedThenCsrDoNotDispatchesStateCacheFlush,
+          IsAtLeastSkl) {
+    auto &csr = neoDevice->getUltCommandStreamReceiver<FamilyType>();
+    EngineControl &engine = neoDevice->getDefaultEngine();
+    static_cast<NEO::MockOsContext *>(engine.osContext)->contextInitialized = false;
+    auto &csrStream = csr.commandStream;
+
+    ze_result_t returnValue;
+    L0::ult::CommandList *cmdListObject = whiteboxCast(CommandList::create(productFamily, device, engineGroupType, 0u, returnValue));
+
+    ze_group_count_t groupCount{1, 1, 1};
+    CmdListKernelLaunchParams launchParams = {};
+    cmdListObject->appendLaunchKernel(kernel->toHandle(), &groupCount, nullptr, 0, nullptr, launchParams, false);
+
+    returnValue = cmdListObject->close();
+    EXPECT_EQ(ZE_RESULT_SUCCESS, returnValue);
+
+    returnValue = cmdListObject->destroy();
+    EXPECT_EQ(ZE_RESULT_SUCCESS, returnValue);
+
+    EXPECT_EQ(0u, csrStream.getUsed());
+}
+
+HWTEST2_F(CommandListStateBaseAddressPrivateHeapTest,
+          givenCommandListUsingPrivateSurfaceHeapWhenTaskCountZeroAndCommandListDestroyedThenCsrDoNotDispatchesStateCacheFlush,
+          IsAtLeastSkl) {
+    auto &csr = neoDevice->getUltCommandStreamReceiver<FamilyType>();
+    auto &csrStream = csr.commandStream;
+
+    ze_result_t returnValue;
+    L0::ult::CommandList *cmdListObject = whiteboxCast(CommandList::create(productFamily, device, engineGroupType, 0u, returnValue));
+
+    returnValue = cmdListObject->destroy();
+    EXPECT_EQ(ZE_RESULT_SUCCESS, returnValue);
+
+    cmdListObject = whiteboxCast(CommandList::create(productFamily, device, engineGroupType, 0u, returnValue));
+
+    returnValue = cmdListObject->destroy();
+    EXPECT_EQ(ZE_RESULT_SUCCESS, returnValue);
+
+    EXPECT_EQ(0u, csrStream.getUsed());
+}
+
+HWTEST2_F(CommandListStateBaseAddressPrivateHeapTest,
+          givenCommandListUsingPrivateSurfaceHeapWhenCommandListDestroyedAndCsrStateCacheFlushDispatchFailsThenWaitNotCalled,
+          IsAtLeastSkl) {
+    auto &csr = neoDevice->getUltCommandStreamReceiver<FamilyType>();
+
+    ze_result_t returnValue;
+    L0::ult::CommandList *cmdListObject = whiteboxCast(CommandList::create(productFamily, device, engineGroupType, 0u, returnValue));
+
+    ze_group_count_t groupCount{1, 1, 1};
+    CmdListKernelLaunchParams launchParams = {};
+    cmdListObject->appendLaunchKernel(kernel->toHandle(), &groupCount, nullptr, 0, nullptr, launchParams, false);
+
+    returnValue = cmdListObject->close();
+    EXPECT_EQ(ZE_RESULT_SUCCESS, returnValue);
+
+    auto cmdListHandle = cmdListObject->toHandle();
+    returnValue = commandQueue->executeCommandLists(1, &cmdListHandle, nullptr, true);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, returnValue);
+
+    csr.callBaseSendRenderStateCacheFlush = false;
+    csr.flushReturnValue = SubmissionStatus::DEVICE_UNINITIALIZED;
+    csr.waitForCompletionWithTimeoutTaskCountCalled = 0;
+
+    returnValue = cmdListObject->destroy();
+    EXPECT_EQ(ZE_RESULT_SUCCESS, returnValue);
+
+    EXPECT_EQ(0u, csr.waitForCompletionWithTimeoutTaskCountCalled);
+
+    csr.callBaseSendRenderStateCacheFlush = true;
+}
+
 } // namespace ult
 } // namespace L0

level_zero/core/test/unit_tests/sources/cmdqueue/test_cmdqueue_enqueue_cmdlist.cpp

@@ -44,6 +44,9 @@ struct CommandQueueExecuteCommandLists : public Test<DeviceFixture> {
     }
 
     void TearDown() override {
+        auto tagAddress = device->getNEODevice()->getDefaultEngine().commandStreamReceiver->getTagAddress();
+        *tagAddress = std::numeric_limits<TagAddressType>::max();
+
         for (auto i = 0u; i < numCommandLists; i++) {
             auto commandList = CommandList::fromHandle(commandLists[i]);
             commandList->destroy();

shared/test/common/libult/ult_command_stream_receiver.h

@@ -382,6 +382,13 @@ class UltCommandStreamReceiver : public CommandStreamReceiverHw<GfxFamily>, publ
         pollForCompletionCalled++;
     }
 
+    SubmissionStatus sendRenderStateCacheFlush() override {
+        if (callBaseSendRenderStateCacheFlush) {
+            return BaseClass::sendRenderStateCacheFlush();
+        }
+        return *flushReturnValue;
+    }
+
     std::vector<std::string> aubCommentMessages;
 
     BatchBuffer latestFlushedBatchBuffer = {};
@@ -430,6 +437,7 @@ class UltCommandStreamReceiver : public CommandStreamReceiverHw<GfxFamily>, publ
     bool shouldFlushBatchedSubmissionsReturnSuccess = false;
     bool callBaseFillReusableAllocationsList = false;
     bool callBaseFlushBcsTask{true};
+    bool callBaseSendRenderStateCacheFlush = true;
 };
 
 } // namespace NEO

shared/test/common/mocks/mock_os_context.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2019-2022 Intel Corporation
+ * Copyright (C) 2019-2023 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -12,6 +12,7 @@ namespace NEO {
 class MockOsContext : public OsContext {
   public:
     using OsContext::checkDirectSubmissionSupportsEngine;
+    using OsContext::contextInitialized;
     using OsContext::debuggableContext;
     using OsContext::engineType;
     using OsContext::engineUsage;