Update SBA in hybrid immediate and regular commandlist usages

Fix to check and update heap states.

Related-To: LOCI-3379
Signed-off-by: Aravind Gopalakrishnan <[email protected]>

Author: aravindksg

level_zero/core/source/cmdqueue/cmdqueue_hw.h

@@ -191,6 +191,8 @@ struct CommandQueueHw : public CommandQueueImp {
                                                               const NEO::StreamProperties &cmdListRequired,
                                                               const NEO::StreamProperties &cmdListFinal);
 
+    inline void updateBaseAddressState(CommandList *lastCommandList);
+
     size_t alignedChildStreamPadding{};
 };
 

level_zero/core/source/cmdqueue/cmdqueue_hw.inl

@@ -192,6 +192,7 @@ ze_result_t CommandQueueHw<gfxCoreFamily>::executeCommandListsRegular(
         this->mergeOneCmdListPipelinedState(commandList);
     }
 
+    this->updateBaseAddressState(CommandList::fromHandle(phCommandLists[numCommandLists - 1]));
     this->collectPrintfContentsFromAllCommandsLists(phCommandLists, numCommandLists);
     this->migrateSharedAllocationsIfRequested(ctx.isMigrationRequested, phCommandLists[0]);
     this->prefetchMemoryIfRequested(ctx.performMemoryPrefetch);
@@ -1226,4 +1227,18 @@ void CommandQueueHw<gfxCoreFamily>::programRequiredStateComputeModeForCommandLis
     csrState.stateComputeMode.setProperties(cmdListFinal.stateComputeMode);
 }
 
+template <GFXCORE_FAMILY gfxCoreFamily>
+void CommandQueueHw<gfxCoreFamily>::updateBaseAddressState(CommandList *lastCommandList) {
+    auto csrHw = static_cast<NEO::CommandStreamReceiverHw<GfxFamily> *>(csr);
+    auto dsh = lastCommandList->commandContainer.getIndirectHeap(NEO::HeapType::DYNAMIC_STATE);
+    if (dsh != nullptr) {
+        csrHw->getDshState().updateAndCheck(dsh);
+    }
+
+    auto ssh = lastCommandList->commandContainer.getIndirectHeap(NEO::HeapType::SURFACE_STATE);
+    if (ssh != nullptr) {
+        csrHw->getSshState().updateAndCheck(ssh);
+    }
+}
+
 } // namespace L0

shared/source/command_stream/command_stream_receiver_hw.h

@@ -141,6 +141,9 @@ class CommandStreamReceiverHw : public CommandStreamReceiver {
     }
     void initializeDeviceWithFirstSubmission() override;
 
+    HeapDirtyState &getDshState() {
+        return dshState;
+    }
     HeapDirtyState &getSshState() {
         return sshState;
     }

shared/test/unit_test/command_stream/command_stream_receiver_tests.cpp

@@ -2218,6 +2218,30 @@ HWTEST_F(CommandStreamReceiverTest, givenSshDirtyStateWhenUpdatingStateWithNewHe
     EXPECT_FALSE(check);
 }
 
+HWTEST_F(CommandStreamReceiverTest, givenDshDirtyStateWhenUpdatingStateWithNewHeapThenExpectDirtyStateTrue) {
+    MockGraphicsAllocation allocation{};
+    allocation.gpuAddress = 0xABCD00;
+    allocation.size = 0x1000;
+
+    IndirectHeap dummyHeap(&allocation, false);
+
+    auto dirtyStateCopy = static_cast<CommandStreamReceiverHw<FamilyType> *>(commandStreamReceiver)->getDshState();
+
+    bool check = dirtyStateCopy.updateAndCheck(&dummyHeap);
+    EXPECT_TRUE(check);
+
+    check = dirtyStateCopy.updateAndCheck(&dummyHeap);
+    EXPECT_FALSE(check);
+
+    auto dirtyState = static_cast<CommandStreamReceiverHw<FamilyType> *>(commandStreamReceiver)->getDshState();
+
+    check = dirtyState.updateAndCheck(&dummyHeap);
+    EXPECT_TRUE(check);
+
+    check = dirtyState.updateAndCheck(&dummyHeap);
+    EXPECT_FALSE(check);
+}
+
 using CommandStreamReceiverHwTest = Test<CommandStreamReceiverFixture>;
 
 HWTEST2_F(CommandStreamReceiverHwTest, givenSshHeapNotProvidedWhenFlushTaskPerformedThenSbaProgammedSurfaceBaseAddressToZero, IsAtLeastXeHpCore) {