refactor: fixes in ults

Related-To: NEO-13789

Signed-off-by: Mateusz Hoppe <mateusz.hoppe@intel.com>

opencl/test/unit_test/device/device_tests.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2024 Intel Corporation
+ * Copyright (C) 2018-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -257,6 +257,8 @@ TEST(DeviceCleanup, givenDeviceWhenItIsDestroyedThenFlushBatchedSubmissionsIsCal
 }
 
 TEST(DeviceCreation, GiveNonExistingFclWhenCreatingDeviceThenCompilerInterfaceIsNotCreated) {
+    DebugManagerStateRestore restore{};
+    debugManager.flags.ForcePreemptionMode.set(PreemptionMode::Disabled);
     VariableBackup<const char *> frontEndDllName(&Os::frontEndDllName);
     Os::frontEndDllName = "_fake_fcl1_so";
 

shared/source/built_ins/sip.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2024 Intel Corporation
+ * Copyright (C) 2018-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -249,15 +249,6 @@ bool SipKernel::initHexadecimalArraySipKernel(SipKernelType type, Device &device
     return true;
 }
 
-void SipKernel::freeSipKernels(RootDeviceEnvironment *rootDeviceEnvironment, MemoryManager *memoryManager) {
-    for (auto &sipKernel : rootDeviceEnvironment->sipKernels) {
-        if (sipKernel.get()) {
-            memoryManager->freeGraphicsMemory(sipKernel->getSipAllocation());
-            sipKernel.reset();
-        }
-    }
-}
-
 void SipKernel::selectSipClassType(std::string &fileName, Device &device) {
     const GfxCoreHelper &gfxCoreHelper = device.getGfxCoreHelper();
     const std::string unknown("unk");

shared/source/built_ins/sip_init.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2021-2024 Intel Corporation
+ * Copyright (C) 2021-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -8,6 +8,8 @@
 #include "shared/source/built_ins/sip.h"
 #include "shared/source/device/device.h"
 #include "shared/source/execution_environment/execution_environment.h"
+#include "shared/source/execution_environment/root_device_environment.h"
+#include "shared/source/memory_manager/memory_manager.h"
 
 namespace NEO {
 
@@ -15,6 +17,15 @@ bool SipKernel::initSipKernel(SipKernelType type, Device &device) {
     return SipKernel::initSipKernelImpl(type, device, nullptr);
 }
 
+void SipKernel::freeSipKernels(RootDeviceEnvironment *rootDeviceEnvironment, MemoryManager *memoryManager) {
+    for (auto &sipKernel : rootDeviceEnvironment->sipKernels) {
+        if (sipKernel.get()) {
+            memoryManager->freeGraphicsMemory(sipKernel->getSipAllocation());
+            sipKernel.reset();
+        }
+    }
+}
+
 const SipKernel &SipKernel::getSipKernel(Device &device, OsContext *context) {
     if (context && device.getExecutionEnvironment()->getDebuggingMode() == NEO::DebuggingMode::offline) {
         return SipKernel::getDebugSipKernel(device, context);

shared/test/common/helpers/sip_init.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2021-2024 Intel Corporation
+ * Copyright (C) 2021-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -243,6 +243,15 @@ bool SipKernel::initSipKernel(SipKernelType type, Device &device) {
     }
 }
 
+void SipKernel::freeSipKernels(RootDeviceEnvironment *rootDeviceEnvironment, MemoryManager *memoryManager) {
+    for (auto &sipKernel : rootDeviceEnvironment->sipKernels) {
+        if (sipKernel.get()) {
+            memoryManager->freeGraphicsMemory(sipKernel->getSipAllocation());
+            sipKernel.reset();
+        }
+    }
+}
+
 const SipKernel &SipKernel::getSipKernel(Device &device, OsContext *context) {
     if (MockSipData::useMockSip) {
         if (!MockSipData::called) {

shared/test/common/mocks/mock_device.cpp

@@ -60,6 +60,7 @@ MockDevice::MockDevice(ExecutionEnvironment *executionEnvironment, uint32_t root
     UnitTestSetter::setRcsExposure(*executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]);
     UnitTestSetter::setCcsExposure(*executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]);
     executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->initGmm();
+    executionEnvironment->calculateMaxOsContextCount();
 
     if (!executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface) {
         executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = std::make_unique<MockMemoryOperations>();
@@ -96,6 +97,17 @@ void MockDevice::resetCommandStreamReceiver(CommandStreamReceiver *newCsr, uint3
 
     registeredEngine.commandStreamReceiver = newCsr;
     allEngines[engineIndex].commandStreamReceiver = newCsr;
+
+    if (osContext->isPartOfContextGroup()) {
+        auto &secondaryEnginesForType = secondaryEngines[osContext->getEngineType()];
+        for (size_t i = 0; i < secondaryEnginesForType.engines.size(); i++) {
+            if (secondaryEnginesForType.engines[i].commandStreamReceiver == commandStreamReceivers[engineIndex].get()) {
+                secondaryEnginesForType.engines[i].commandStreamReceiver = newCsr;
+                break;
+            }
+        }
+    }
+
     const_cast<EngineControlContainer &>(memoryManager->getRegisteredEngines(rootDeviceIndex)).emplace_back(registeredEngine);
     osContext->incRefInternal();
     newCsr->setupContext(*osContext);

shared/test/common/mocks/mock_memory_manager.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2024 Intel Corporation
+ * Copyright (C) 2018-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -300,6 +300,23 @@ OsContext *MockMemoryManagerOsAgnosticContext::createAndRegisterOsContext(Comman
     return osContext;
 }
 
+OsContext *MockMemoryManagerOsAgnosticContext::createAndRegisterSecondaryOsContext(const OsContext *primaryContext, CommandStreamReceiver *commandStreamReceiver,
+                                                                                   const EngineDescriptor &engineDescriptor) {
+    auto rootDeviceIndex = commandStreamReceiver->getRootDeviceIndex();
+
+    auto osContext = new OsContext(rootDeviceIndex, 0, engineDescriptor);
+    osContext->incRefInternal();
+
+    osContext->setPrimaryContext(primaryContext);
+
+    UNRECOVERABLE_IF(rootDeviceIndex != osContext->getRootDeviceIndex());
+
+    secondaryEngines[rootDeviceIndex].emplace_back(commandStreamReceiver, osContext);
+    allRegisteredEngines[rootDeviceIndex].emplace_back(commandStreamReceiver, osContext);
+
+    return osContext;
+}
+
 OsContext *MockMemoryManagerWithDebuggableOsContext::createAndRegisterOsContext(CommandStreamReceiver *commandStreamReceiver,
                                                                                 const EngineDescriptor &engineDescriptor) {
     auto osContext = new MockOsContext(0, engineDescriptor);

shared/test/common/mocks/mock_memory_manager.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2024 Intel Corporation
+ * Copyright (C) 2018-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -473,6 +473,8 @@ class MockMemoryManagerOsAgnosticContext : public MockMemoryManager {
     MockMemoryManagerOsAgnosticContext(NEO::ExecutionEnvironment &executionEnvironment) : MockMemoryManager(executionEnvironment) {}
     OsContext *createAndRegisterOsContext(CommandStreamReceiver *commandStreamReceiver,
                                           const EngineDescriptor &engineDescriptor) override;
+    OsContext *createAndRegisterSecondaryOsContext(const OsContext *primaryContext, CommandStreamReceiver *commandStreamReceiver,
+                                                   const EngineDescriptor &engineDescriptor) override;
 };
 
 class MockMemoryManagerWithDebuggableOsContext : public MockMemoryManager {

shared/test/unit_test/command_stream/aub_command_stream_receiver_1_tests.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2024 Intel Corporation
+ * Copyright (C) 2018-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -74,7 +74,7 @@ HWTEST_F(AubCommandStreamReceiverTests, givenDebugFlagSetWhenCreatingContextThen
     aubCsr->aubManager = &mockAubManager;
 
     aubCsr->setupContext(*pDevice->getDefaultEngine().osContext);
-    EXPECT_EQ(0x123u, mockAubManager.contextFlags);
+    EXPECT_EQ(0x123u, mockAubManager.contextFlags & 0x123);
 }
 
 TEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverWhenItIsCreatedWithWrongGfxCoreFamilyThenNullPointerShouldBeReturned) {
@@ -350,6 +350,7 @@ HWTEST_F(AubCommandStreamReceiverTests, givenAubCsrAndResidentAllocationWhenProc
 
     auto commandStreamReceiver = std::make_unique<MockAubCsr<FamilyType>>("", true, *pDevice->getExecutionEnvironment(), 0, pDevice->getDeviceBitfield());
 
+    pDevice->getExecutionEnvironment()->memoryManager->reInitLatestContextId();
     auto osContext = pDevice->getExecutionEnvironment()->memoryManager->createAndRegisterOsContext(commandStreamReceiver.get(),
                                                                                                    EngineDescriptorHelper::getDefaultDescriptor({getChosenEngineType(*defaultHwInfo), EngineUsage::regular},
                                                                                                                                                 PreemptionHelper::getDefaultPreemptionMode(*defaultHwInfo)));

shared/test/unit_test/command_stream/command_stream_receiver_with_aub_dump_tests.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2024 Intel Corporation
+ * Copyright (C) 2018-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -155,6 +155,7 @@ struct CommandStreamReceiverWithAubDumpTest : public ::testing::TestWithParam<bo
         auto engineDescriptor = EngineDescriptorHelper::getDefaultDescriptor({getChosenEngineType(DEFAULT_TEST_PLATFORM::hwInfo), EngineUsage::regular},
                                                                              PreemptionHelper::getDefaultPreemptionMode(DEFAULT_TEST_PLATFORM::hwInfo));
 
+        pDevice->getExecutionEnvironment()->memoryManager->reInitLatestContextId();
         auto osContext = executionEnvironment->memoryManager->createAndRegisterOsContext(csrWithAubDump, engineDescriptor);
         csrWithAubDump->setupContext(*osContext);
     }

shared/test/unit_test/command_stream/tbx_command_stream_tests.cpp

@@ -366,6 +366,7 @@ HWTEST_F(TbxCommandSteamSimpleTest, givenTbxCsrAndResidentAllocationWhenProcessR
 
     auto commandStreamReceiver = std::make_unique<MockTbxCsr<FamilyType>>(*pDevice->getExecutionEnvironment(), pDevice->getDeviceBitfield());
 
+    pDevice->getExecutionEnvironment()->memoryManager->reInitLatestContextId();
     auto osContext = pDevice->getExecutionEnvironment()->memoryManager->createAndRegisterOsContext(commandStreamReceiver.get(),
                                                                                                    EngineDescriptorHelper::getDefaultDescriptor({getChosenEngineType(*defaultHwInfo), EngineUsage::regular},
                                                                                                                                                 PreemptionHelper::getDefaultPreemptionMode(*defaultHwInfo)));

shared/test/unit_test/execution_environment/execution_environment_tests.cpp

@@ -504,6 +504,7 @@ TEST(ExecutionEnvironment, givenUnproperSetCsrFlagValueWhenInitializingMemoryMan
 
 TEST(ExecutionEnvironment, whenCalculateMaxOsContexCountThenGlobalVariableHasProperValue) {
     DebugManagerStateRestore restore;
+    debugManager.flags.ContextGroupSize.set(0);
     VariableBackup<uint32_t> osContextCountBackup(&MemoryManager::maxOsContextCount, 0);
     uint32_t numRootDevices = 17u;
     uint32_t expectedOsContextCount = 0u;