fix: update context base address state after execution of regular command list

Related-To: NEO-7808

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

level_zero/core/source/cmdqueue/cmdqueue_hw.inl

@@ -1470,19 +1470,26 @@ void CommandQueueHw<gfxCoreFamily>::programRequiredStateBaseAddressForCommandLis
 template <GFXCORE_FAMILY gfxCoreFamily>
 void CommandQueueHw<gfxCoreFamily>::updateBaseAddressState(CommandList *lastCommandList) {
     auto csrHw = static_cast<NEO::CommandStreamReceiverHw<GfxFamily> *>(csr);
+    auto &streamProperties = this->csr->getStreamProperties();
+
     auto &commandContainer = lastCommandList->getCmdContainer();
 
     if (lastCommandList->getCmdListHeapAddressModel() == NEO::HeapAddressModel::GlobalStateless) {
-        csrHw->getSshState().updateAndCheck(csr->getGlobalStatelessHeap());
+        auto globalStateless = csr->getGlobalStatelessHeap();
+        csrHw->getSshState().updateAndCheck(globalStateless);
+        streamProperties.stateBaseAddress.setPropertiesSurfaceState(globalStateless->getHeapGpuBase(), globalStateless->getHeapSizeInPages());
     } else {
         auto dsh = commandContainer.getIndirectHeap(NEO::HeapType::DYNAMIC_STATE);
         if (dsh != nullptr) {
             csrHw->getDshState().updateAndCheck(dsh);
+            streamProperties.stateBaseAddress.setPropertiesDynamicState(dsh->getHeapGpuBase(), dsh->getHeapSizeInPages());
         }
 
         auto ssh = commandContainer.getIndirectHeap(NEO::HeapType::SURFACE_STATE);
         if (ssh != nullptr) {
             csrHw->getSshState().updateAndCheck(ssh);
+            streamProperties.stateBaseAddress.setPropertiesBindingTableSurfaceState(ssh->getHeapGpuBase(), ssh->getHeapSizeInPages(),
+                                                                                    ssh->getHeapGpuBase(), ssh->getHeapSizeInPages());
         }
     }
 

level_zero/core/source/gfx_core_helpers/l0_gfx_core_helper_xehp_and_later.inl

@@ -84,7 +84,7 @@ bool L0GfxCoreHelperHw<Family>::platformSupportsPrimaryBatchBufferCmdList() cons
 
 template <typename Family>
 bool L0GfxCoreHelperHw<Family>::platformSupportsImmediateComputeFlushTask() const {
-    return false;
+    return true;
 }
 
 } // namespace L0

level_zero/core/source/xe_hpc_core/l0_gfx_core_helper_xe_hpc_core.cpp

@@ -29,11 +29,6 @@ bool L0GfxCoreHelperHw<Family>::alwaysAllocateEventInLocalMem() const {
     return true;
 }
 
-template <>
-bool L0GfxCoreHelperHw<Family>::platformSupportsImmediateComputeFlushTask() const {
-    return true;
-}
-
 template class L0GfxCoreHelperHw<Family>;
 
 } // namespace L0

level_zero/core/source/xe_hpg_core/l0_gfx_core_helper_xe_hpg_core.cpp

@@ -24,6 +24,11 @@ bool L0GfxCoreHelperHw<Family>::isResumeWARequired() {
     return true;
 }
 
+template <>
+bool L0GfxCoreHelperHw<Family>::platformSupportsImmediateComputeFlushTask() const {
+    return false;
+}
+
 template class L0GfxCoreHelperHw<Family>;
 
 } // namespace L0

level_zero/core/test/unit_tests/fixtures/cmdlist_fixture.cpp

@@ -386,8 +386,21 @@ void ImmediateFlushTaskGlobalStatelessCmdListFixture::setUp() {
 void ImmediateFlushTaskCsrSharedHeapCmdListFixture::setUp() {
     DebugManager.flags.EnableImmediateCmdListHeapSharing.set(1);
     DebugManager.flags.SelectCmdListHeapAddressModel.set(static_cast<int32_t>(NEO::HeapAddressModel::PrivateHeaps));
+    DebugManager.flags.ForceDefaultHeapSize.set(64);
 
     ImmediateFlushTaskCmdListFixture::setUp();
+
+    mockKernelImmData->kernelDescriptor->payloadMappings.samplerTable.numSamplers = 1;
+    mockKernelImmData->kernelDescriptor->payloadMappings.samplerTable.tableOffset = 16;
+    mockKernelImmData->kernelDescriptor->payloadMappings.samplerTable.borderColor = 0;
+    kernel->dynamicStateHeapData.reset(new uint8_t[512]);
+    kernel->dynamicStateHeapDataSize = 512;
+
+    mockKernelImmData->mockKernelInfo->heapInfo.surfaceStateHeapSize = 128;
+    mockKernelImmData->kernelDescriptor->payloadMappings.bindingTable.numEntries = 1;
+    mockKernelImmData->kernelDescriptor->payloadMappings.bindingTable.tableOffset = 64;
+    kernel->surfaceStateHeapDataSize = 128;
+    kernel->surfaceStateHeapData.reset(new uint8_t[256]);
 }
 
 void ImmediateFlushTaskPrivateHeapCmdListFixture::setUp() {

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

@@ -2220,6 +2220,93 @@ HWTEST2_F(ImmediateFlushTaskCsrSharedHeapCmdListTest,
     EXPECT_EQ(0u, sbaCmd->getSurfaceStateBaseAddress());
 }
 
+HWTEST2_F(ImmediateFlushTaskCsrSharedHeapCmdListTest,
+          givenImmediateFlushOnCsrSharedHeapsFirstWhenExecuteRegularCommandListSecondAndImmediateThirdThenExpectSharedHeapBaseAddressRestored,
+          IsAtLeastXeHpCore) {
+    using STATE_BASE_ADDRESS = typename FamilyType::STATE_BASE_ADDRESS;
+
+    auto &csrImmediate = neoDevice->getUltCommandStreamReceiver<FamilyType>();
+    auto &csrStream = csrImmediate.commandStream;
+
+    ze_group_count_t groupCount{1, 1, 1};
+    CmdListKernelLaunchParams launchParams = {};
+    size_t csrUsedBefore = csrStream.getUsed();
+    auto result = commandListImmediate->appendLaunchKernel(kernel->toHandle(), &groupCount, nullptr, 0, nullptr, launchParams, false);
+    size_t csrUsedAfter = csrStream.getUsed();
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+
+    auto ssSharedHeap = commandListImmediate->getCmdContainer().getSurfaceStateHeapReserve().indirectHeapReservation;
+    EXPECT_NE(nullptr, ssSharedHeap->getGraphicsAllocation());
+    auto ssShareBaseAddress = ssSharedHeap->getHeapGpuBase();
+
+    auto dsSharedHeap = commandListImmediate->getCmdContainer().getDynamicStateHeapReserve().indirectHeapReservation;
+    uint64_t dsShareBaseAddress = 0;
+    if (this->dshRequired) {
+        EXPECT_NE(nullptr, dsSharedHeap->getGraphicsAllocation());
+        dsShareBaseAddress = dsSharedHeap->getHeapGpuBase();
+    } else {
+        EXPECT_EQ(nullptr, dsSharedHeap->getGraphicsAllocation());
+    }
+
+    GenCmdList cmdList;
+    ASSERT_TRUE(FamilyType::PARSE::parseCommandBuffer(
+        cmdList,
+        ptrOffset(csrStream.getCpuBase(), csrUsedBefore),
+        csrUsedAfter - csrUsedBefore));
+    auto sbaCmds = findAll<STATE_BASE_ADDRESS *>(cmdList.begin(), cmdList.end());
+    ASSERT_EQ(expectedSbaCmds, sbaCmds.size());
+    auto sbaCmd = reinterpret_cast<STATE_BASE_ADDRESS *>(*sbaCmds[0]);
+
+    EXPECT_TRUE(sbaCmd->getSurfaceStateBaseAddressModifyEnable());
+    EXPECT_EQ(ssShareBaseAddress, sbaCmd->getSurfaceStateBaseAddress());
+
+    EXPECT_EQ(this->dshRequired, sbaCmd->getDynamicStateBaseAddressModifyEnable());
+    EXPECT_EQ(dsShareBaseAddress, sbaCmd->getDynamicStateBaseAddress());
+
+    result = commandList->appendLaunchKernel(kernel->toHandle(), &groupCount, nullptr, 0, nullptr, launchParams, false);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+    result = commandList->close();
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+
+    auto &container = commandList->getCmdContainer();
+    auto sshRegularHeap = container.getIndirectHeap(NEO::HeapType::SURFACE_STATE);
+    auto ssRegularBaseAddress = sshRegularHeap->getHeapGpuBase();
+
+    uint64_t dsRegularBaseAddress = static_cast<uint64_t>(-1);
+    if (this->dshRequired) {
+        auto dshRegularHeap = container.getIndirectHeap(NEO::HeapType::DYNAMIC_STATE);
+        dsRegularBaseAddress = dshRegularHeap->getHeapGpuBase();
+    }
+
+    ze_command_list_handle_t cmdListHandle = commandList->toHandle();
+    result = commandQueue->executeCommandLists(1, &cmdListHandle, nullptr, true);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+
+    auto &csrBaseAddressState = csrImmediate.getStreamProperties().stateBaseAddress;
+    EXPECT_EQ(ssRegularBaseAddress, static_cast<uint64_t>(csrBaseAddressState.surfaceStateBaseAddress.value));
+    EXPECT_EQ(dsRegularBaseAddress, static_cast<uint64_t>(csrBaseAddressState.dynamicStateBaseAddress.value));
+
+    csrUsedBefore = csrStream.getUsed();
+    result = commandListImmediate->appendLaunchKernel(kernel->toHandle(), &groupCount, nullptr, 0, nullptr, launchParams, false);
+    csrUsedAfter = csrStream.getUsed();
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+
+    cmdList.clear();
+    ASSERT_TRUE(FamilyType::PARSE::parseCommandBuffer(
+        cmdList,
+        ptrOffset(csrStream.getCpuBase(), csrUsedBefore),
+        csrUsedAfter - csrUsedBefore));
+    sbaCmds = findAll<STATE_BASE_ADDRESS *>(cmdList.begin(), cmdList.end());
+    ASSERT_EQ(expectedSbaCmds, sbaCmds.size());
+    sbaCmd = reinterpret_cast<STATE_BASE_ADDRESS *>(*sbaCmds[0]);
+
+    EXPECT_TRUE(sbaCmd->getSurfaceStateBaseAddressModifyEnable());
+    EXPECT_EQ(ssShareBaseAddress, sbaCmd->getSurfaceStateBaseAddress());
+
+    EXPECT_EQ(this->dshRequired, sbaCmd->getDynamicStateBaseAddressModifyEnable());
+    EXPECT_EQ(dsShareBaseAddress, sbaCmd->getDynamicStateBaseAddress());
+}
+
 using ImmediateFlushTaskPrivateHeapCmdListTest = Test<ImmediateFlushTaskPrivateHeapCmdListFixture>;
 
 HWTEST2_F(ImmediateFlushTaskPrivateHeapCmdListTest,