From a2a154ba4a25404d727921c427ab45bf5af11385 Mon Sep 17 00:00:00 2001
From: Mateusz Hoppe <mateusz.hoppe@intel.com>
Date: Fri, 24 May 2024 13:46:53 +0000
Subject: [PATCH] fix: initialize context with first submission in heapless
 execution

- early return from initializeDeviceWithFirstSubmission if flush was
already called

Related-To: NEO-11219

Signed-off-by: Mateusz Hoppe <mateusz.hoppe@intel.com>
---
 level_zero/core/source/cmdqueue/cmdqueue_hw.inl   |  4 +++-
 ...command_stream_receiver_hw_heap_addressing.inl |  3 +++
 .../command_stream_receiver_tests.cpp             | 15 +++++++++++++++
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/level_zero/core/source/cmdqueue/cmdqueue_hw.inl b/level_zero/core/source/cmdqueue/cmdqueue_hw.inl
index 99fba1c07a..c6e58dd9e7 100644
--- a/level_zero/core/source/cmdqueue/cmdqueue_hw.inl
+++ b/level_zero/core/source/cmdqueue/cmdqueue_hw.inl
@@ -135,10 +135,12 @@ ze_result_t CommandQueueHw<gfxCoreFamily>::executeCommandListsRegularHeapless(
     ze_event_handle_t hSignalEvent, uint32_t numWaitEvents,
     ze_event_handle_t *phWaitEvents) {
 
+    auto neoDevice = this->device->getNEODevice();
+    this->csr->initializeDeviceWithFirstSubmission(*neoDevice);
+
     this->setupCmdListsAndContextParams(ctx, commandListHandles, numCommandLists, hFence);
     ctx.isDirectSubmissionEnabled = this->csr->isDirectSubmissionEnabled();
     bool instructionCacheFlushRequired = this->csr->isInstructionCacheFlushRequired();
-    auto neoDevice = this->device->getNEODevice();
     bool stateCacheFlushRequired = neoDevice->getBindlessHeapsHelper() ? neoDevice->getBindlessHeapsHelper()->getStateDirtyForContext(this->csr->getOsContext().getContextId()) : false;
 
     std::unique_lock<std::mutex> lockForIndirect;
diff --git a/shared/source/command_stream/command_stream_receiver_hw_heap_addressing.inl b/shared/source/command_stream/command_stream_receiver_hw_heap_addressing.inl
index a101711db2..18440cb9a4 100644
--- a/shared/source/command_stream/command_stream_receiver_hw_heap_addressing.inl
+++ b/shared/source/command_stream/command_stream_receiver_hw_heap_addressing.inl
@@ -11,6 +11,9 @@ namespace NEO {
 
 template <typename GfxFamily>
 SubmissionStatus CommandStreamReceiverHw<GfxFamily>::initializeDeviceWithFirstSubmission(Device &device) {
+    if (this->latestFlushedTaskCount > 0) {
+        return SubmissionStatus::success;
+    }
     return flushTagUpdate();
 }
 
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 4d3d42c364..dcc59f1a5b 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
@@ -5236,6 +5236,21 @@ HWTEST_F(CommandStreamReceiverTest, givenBcsCsrWhenInitializeDeviceWithFirstSubm
     EXPECT_EQ(SubmissionStatus::success, commandStreamReceiver.initializeDeviceWithFirstSubmission(*pDevice));
 }
 
+HWTEST_F(CommandStreamReceiverTest, givenCsrWhenInitializeDeviceWithFirstSubmissionIsCalledThenFlushOnlyForTheFirstTime) {
+    MockOsContext mockOsContext(0, EngineDescriptorHelper::getDefaultDescriptor({defaultHwInfo->capabilityTable.defaultEngineType, EngineUsage::regular}));
+    pDevice->setPreemptionMode(PreemptionMode::Disabled);
+
+    MockCsrHw<FamilyType> commandStreamReceiver(*pDevice->executionEnvironment, pDevice->getRootDeviceIndex(), pDevice->getDeviceBitfield());
+    commandStreamReceiver.setupContext(mockOsContext);
+    commandStreamReceiver.initializeTagAllocation();
+    EXPECT_EQ(0u, commandStreamReceiver.taskCount);
+    EXPECT_EQ(SubmissionStatus::success, commandStreamReceiver.initializeDeviceWithFirstSubmission(*pDevice));
+    EXPECT_EQ(1u, commandStreamReceiver.taskCount);
+
+    EXPECT_EQ(SubmissionStatus::success, commandStreamReceiver.initializeDeviceWithFirstSubmission(*pDevice));
+    EXPECT_EQ(1u, commandStreamReceiver.taskCount);
+}
+
 using CommandStreamReceiverHwHeaplessTest = Test<DeviceFixture>;
 
 HWTEST_F(CommandStreamReceiverHwHeaplessTest, whenHeaplessCommandStreamReceiverFunctionsAreCalledThenExceptionIsThrown) {