fix: inherit work partition allocation from primary root csr

Related-To: NEO-8171

Signed-off-by: Bartosz Dunajski <bartosz.dunajski@intel.com>

shared/source/command_stream/command_stream_receiver.h

@@ -216,7 +216,13 @@ class CommandStreamReceiver {
 
         return globalFenceAllocation;
     }
-    GraphicsAllocation *getWorkPartitionAllocation() const { return workPartitionAllocation; }
+    GraphicsAllocation *getWorkPartitionAllocation() const {
+        if (primaryCsr) {
+            return primaryCsr->getWorkPartitionAllocation();
+        }
+        return workPartitionAllocation;
+    }
+
     GraphicsAllocation *getGlobalStatelessHeapAllocation() const {
         if (primaryCsr) {
             return primaryCsr->getGlobalStatelessHeapAllocation();

shared/source/command_stream/command_stream_receiver_hw_base.inl

@@ -530,12 +530,12 @@ CompletionStamp CommandStreamReceiverHw<GfxFamily>::flushTask(
 
     this->makeResident(*tagAllocation);
 
-    if (globalFenceAllocation) {
+    if (getGlobalFenceAllocation()) {
-        makeResident(*globalFenceAllocation);
+        makeResident(*getGlobalFenceAllocation());
     }
 
-    if (preemptionAllocation) {
+    if (getPreemptionAllocation()) {
-        makeResident(*preemptionAllocation);
+        makeResident(*getPreemptionAllocation());
     }
 
     bool debuggingEnabled = device.getDebugger() != nullptr;
@@ -549,8 +549,8 @@ CompletionStamp CommandStreamReceiverHw<GfxFamily>::flushTask(
         makeResident(*debugSurface);
     }
 
-    if (workPartitionAllocation) {
+    if (getWorkPartitionAllocation()) {
-        makeResident(*workPartitionAllocation);
+        makeResident(*getWorkPartitionAllocation());
     }
 
     auto rtBuffer = device.getRTMemoryBackedBuffer();
@@ -846,7 +846,7 @@ inline WaitStatus CommandStreamReceiverHw<GfxFamily>::waitForTaskCountWithKmdNot
 
 template <typename GfxFamily>
 inline void CommandStreamReceiverHw<GfxFamily>::programPreemption(LinearStream &csr, DispatchFlags &dispatchFlags) {
-    PreemptionHelper::programCmdStream<GfxFamily>(csr, dispatchFlags.preemptionMode, this->lastPreemptionMode, preemptionAllocation);
+    PreemptionHelper::programCmdStream<GfxFamily>(csr, dispatchFlags.preemptionMode, this->lastPreemptionMode, getPreemptionAllocation());
     this->lastPreemptionMode = dispatchFlags.preemptionMode;
 }
 
@@ -866,7 +866,7 @@ inline void CommandStreamReceiverHw<GfxFamily>::programStateSip(LinearStream &cm
 template <typename GfxFamily>
 inline void CommandStreamReceiverHw<GfxFamily>::programPreamble(LinearStream &csr, Device &device, uint32_t &newL3Config) {
     if (!this->isPreambleSent) {
-        PreambleHelper<GfxFamily>::programPreamble(&csr, device, newL3Config, this->preemptionAllocation, EngineHelpers::isBcs(osContext->getEngineType()));
+        PreambleHelper<GfxFamily>::programPreamble(&csr, device, newL3Config, getPreemptionAllocation(), EngineHelpers::isBcs(osContext->getEngineType()));
         this->isPreambleSent = true;
         this->lastSentL3Config = newL3Config;
     }
@@ -2116,12 +2116,12 @@ void CommandStreamReceiverHw<GfxFamily>::handleImmediateFlushAllocationsResidenc
                                                                                   LinearStream &csrStream) {
     this->makeResident(*tagAllocation);
 
-    if (globalFenceAllocation) {
+    if (getGlobalFenceAllocation()) {
-        makeResident(*globalFenceAllocation);
+        makeResident(*getGlobalFenceAllocation());
     }
 
-    if (workPartitionAllocation) {
+    if (getWorkPartitionAllocation()) {
-        makeResident(*workPartitionAllocation);
+        makeResident(*getWorkPartitionAllocation());
     }
 
     if (device.getRTMemoryBackedBuffer()) {
@@ -2132,8 +2132,8 @@ void CommandStreamReceiverHw<GfxFamily>::handleImmediateFlushAllocationsResidenc
         makeResident(*csrStream.getGraphicsAllocation());
     }
 
-    if (preemptionAllocation) {
+    if (getPreemptionAllocation()) {
-        makeResident(*preemptionAllocation);
+        makeResident(*getPreemptionAllocation());
     }
 
     if (device.isStateSipRequired()) {

shared/source/xe2_hpg_core/command_stream_receiver_hw_xe2_hpg_core.cpp

@@ -39,8 +39,8 @@ void populateFactoryTable<CommandStreamReceiverHw<Family>>() {
 template <>
 void CommandStreamReceiverHw<Family>::programEnginePrologue(LinearStream &csr) {
     if (!this->isEnginePrologueSent) {
-        if (globalFenceAllocation) {
+        if (getGlobalFenceAllocation()) {
-            EncodeMemoryFence<Family>::encodeSystemMemoryFence(csr, globalFenceAllocation);
+            EncodeMemoryFence<Family>::encodeSystemMemoryFence(csr, getGlobalFenceAllocation());
         }
         this->isEnginePrologueSent = true;
     }
@@ -49,7 +49,7 @@ void CommandStreamReceiverHw<Family>::programEnginePrologue(LinearStream &csr) {
 template <>
 size_t CommandStreamReceiverHw<Family>::getCmdSizeForPrologue() const {
     if (!this->isEnginePrologueSent) {
-        if (globalFenceAllocation) {
+        if (getGlobalFenceAllocation()) {
             return EncodeMemoryFence<Family>::getSystemMemoryFenceSize();
         }
     }

shared/source/xe_hpc_core/command_stream_receiver_hw_xe_hpc_core.cpp

@@ -39,8 +39,8 @@ void populateFactoryTable<CommandStreamReceiverHw<Family>>() {
 template <>
 void CommandStreamReceiverHw<Family>::programEnginePrologue(LinearStream &csr) {
     if (!this->isEnginePrologueSent) {
-        if (globalFenceAllocation) {
+        if (getGlobalFenceAllocation()) {
-            EncodeMemoryFence<Family>::encodeSystemMemoryFence(csr, globalFenceAllocation);
+            EncodeMemoryFence<Family>::encodeSystemMemoryFence(csr, getGlobalFenceAllocation());
         }
         this->isEnginePrologueSent = true;
     }
@@ -49,7 +49,7 @@ void CommandStreamReceiverHw<Family>::programEnginePrologue(LinearStream &csr) {
 template <>
 size_t CommandStreamReceiverHw<Family>::getCmdSizeForPrologue() const {
     if (!this->isEnginePrologueSent) {
-        if (globalFenceAllocation) {
+        if (getGlobalFenceAllocation()) {
             return EncodeMemoryFence<Family>::getSystemMemoryFenceSize();
         }
     }

shared/test/unit_test/command_stream/command_stream_receiver_tests.cpp

@@ -5699,6 +5699,60 @@ HWTEST_F(CommandStreamReceiverContextGroupTest, givenSecondaryCsrWhenGettingInte
     }
 }
 
+HWTEST_F(CommandStreamReceiverContextGroupTest, givenSecondaryRootCsrWhenGettingInternalAllocationsThenAllocationFromPrimnaryCsrAreReturned) {
+    HardwareInfo hwInfo = *defaultHwInfo;
+    if (hwInfo.capabilityTable.defaultEngineType != aub_stream::EngineType::ENGINE_CCS) {
+        GTEST_SKIP();
+    }
+
+    DebugManagerStateRestore dbgRestorer;
+    debugManager.flags.ContextGroupSize.set(5);
+
+    hwInfo.featureTable.flags.ftrCCSNode = true;
+    hwInfo.capabilityTable.defaultEngineType = aub_stream::ENGINE_CCS;
+    hwInfo.gtSystemInfo.CCSInfo.NumberOfCCSEnabled = 1;
+    hwInfo.capabilityTable.defaultPreemptionMode = PreemptionMode::MidThread;
+
+    UltDeviceFactory deviceFactory{1, 2};
+    auto device = deviceFactory.rootDevices[0];
+    const auto &gfxCoreHelper = device->getRootDeviceEnvironment().getHelper<GfxCoreHelper>();
+
+    const auto ccsIndex = 0;
+    auto &secondaryEngines = device->secondaryEngines[EngineHelpers::mapCcsIndexToEngineType(ccsIndex)];
+    auto secondaryEnginesCount = secondaryEngines.engines.size();
+    ASSERT_EQ(5u, secondaryEnginesCount);
+
+    EXPECT_TRUE(secondaryEngines.engines[0].commandStreamReceiver->isInitialized());
+
+    auto primaryCsr = secondaryEngines.engines[0].commandStreamReceiver;
+    primaryCsr->createGlobalStatelessHeap();
+
+    for (uint32_t secondaryIndex = 1; secondaryIndex < secondaryEnginesCount; secondaryIndex++) {
+        device->getSecondaryEngineCsr({EngineHelpers::mapCcsIndexToEngineType(ccsIndex), EngineUsage::regular}, false);
+    }
+
+    for (uint32_t i = 0; i < secondaryEngines.highPriorityEnginesTotal; i++) {
+        device->getSecondaryEngineCsr({EngineHelpers::mapCcsIndexToEngineType(ccsIndex), EngineUsage::highPriority}, false);
+    }
+
+    for (uint32_t secondaryIndex = 0; secondaryIndex < secondaryEnginesCount; secondaryIndex++) {
+
+        if (secondaryIndex > 0) {
+            EXPECT_NE(primaryCsr->getTagAllocation(), secondaryEngines.engines[secondaryIndex].commandStreamReceiver->getTagAllocation());
+        }
+
+        if (gfxCoreHelper.isFenceAllocationRequired(hwInfo)) {
+            EXPECT_EQ(primaryCsr->getGlobalFenceAllocation(), secondaryEngines.engines[secondaryIndex].commandStreamReceiver->getGlobalFenceAllocation());
+        }
+
+        EXPECT_EQ(primaryCsr->getPreemptionAllocation(), secondaryEngines.engines[secondaryIndex].commandStreamReceiver->getPreemptionAllocation());
+        EXPECT_EQ(primaryCsr->getGlobalStatelessHeapAllocation(), secondaryEngines.engines[secondaryIndex].commandStreamReceiver->getGlobalStatelessHeapAllocation());
+        EXPECT_EQ(primaryCsr->getGlobalStatelessHeap(), secondaryEngines.engines[secondaryIndex].commandStreamReceiver->getGlobalStatelessHeap());
+        EXPECT_EQ(primaryCsr->getPrimaryScratchSpaceController(), secondaryEngines.engines[secondaryIndex].commandStreamReceiver->getPrimaryScratchSpaceController());
+        EXPECT_EQ(primaryCsr->getWorkPartitionAllocation(), secondaryEngines.engines[secondaryIndex].commandStreamReceiver->getWorkPartitionAllocation());
+    }
+}
+
 HWTEST_F(CommandStreamReceiverContextGroupTest, givenContextGroupWhenCreatingEnginesThenSetCorrectMaxOsContextCount) {
 
     HardwareInfo hwInfo = *defaultHwInfo;