performance: Move preemption allocation init to the first submission

Related-To: NEO-12323 Signed-off-by: Krzysztof Sprzaczkowski <krzysztof.sprzaczkowski@intel.com>
2025-12-19 06:24:51 +08:00 · 2025-01-15 13:44:10 +00:00
parent 92ea7ee7a3
commit a17745532c
23 changed files with 117 additions and 48 deletions
--- a/level_zero/core/source/cmdlist/cmdlist_imp.cpp
+++ b/level_zero/core/source/cmdlist/cmdlist_imp.cpp
@@ -221,7 +221,7 @@ CommandList *CommandList::createImmediate(uint32_t productFamily, Device *device
            bool enabledCmdListSharing = !NEO::EngineHelper::isCopyOnlyEngineType(engineGroupType) && commandList->isFlushTaskSubmissionEnabled;
            commandList->immediateCmdListHeapSharing = L0GfxCoreHelper::enableImmediateCmdListHeapSharing(rootDeviceEnvironment, enabledCmdListSharing);
        }
-        csr->initializeResources(false);
+        csr->initializeResources(false, device->getDevicePreemptionMode());
        csr->initDirectSubmission();

        auto commandQueue = CommandQueue::create(productFamily, device, csr, &cmdQdesc, NEO::EngineHelper::isCopyOnlyEngineType(engineGroupType), internalUsage, true, returnValue);
--- a/level_zero/core/source/cmdqueue/cmdqueue.cpp
+++ b/level_zero/core/source/cmdqueue/cmdqueue.cpp
@@ -250,7 +250,7 @@ CommandQueue *CommandQueue::create(uint32_t productFamily, Device *device, NEO::
        osContext.reInitializeContext();
    }

-    csr->initializeResources(false);
+    csr->initializeResources(false, device->getDevicePreemptionMode());
    csr->initDirectSubmission();
    if (commandQueue->cmdListHeapAddressModel == NEO::HeapAddressModel::globalStateless) {
        csr->createGlobalStatelessHeap();
--- a/level_zero/core/test/aub_tests/fixtures/aub_fixture.cpp
+++ b/level_zero/core/test/aub_tests/fixtures/aub_fixture.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2021-2024 Intel Corporation
+ * Copyright (C) 2021-2025 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
@@ -75,6 +75,9 @@ void AUBFixtureL0::setUp(const NEO::HardwareInfo *hardwareInfo, bool debuggingEn
    neoDevice = NEO::MockDevice::createWithExecutionEnvironment<NEO::MockDevice>(&hwInfo, executionEnvironment, 0u);

    this->csr = neoDevice->getDefaultEngine().commandStreamReceiver;
+    if (!this->csr->getPreemptionAllocation()) {
+        ASSERT_TRUE(this->csr->createPreemptionAllocation());
+    }

    NEO::DeviceVector devices;
    devices.push_back(std::unique_ptr<NEO::Device>(neoDevice));
--- a/level_zero/core/test/unit_tests/fixtures/device_fixture.cpp
+++ b/level_zero/core/test/unit_tests/fixtures/device_fixture.cpp
@@ -53,6 +53,14 @@ void DeviceFixture::setupWithExecutionEnvironment(NEO::ExecutionEnvironment &exe
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    context = static_cast<ContextImp *>(Context::fromHandle(hContext));
    executionEnvironment.incRefInternal();
+    if (neoDevice->getPreemptionMode() == NEO::PreemptionMode::MidThread) {
+        for (auto &engine : neoDevice->getAllEngines()) {
+            NEO::CommandStreamReceiver *csr = engine.commandStreamReceiver;
+            if (!csr->getPreemptionAllocation()) {
+                csr->createPreemptionAllocation();
+            }
+        }
+    }
 }

 void DeviceFixture::tearDown() {
--- a/level_zero/core/test/unit_tests/fixtures/host_pointer_manager_fixture.cpp
+++ b/level_zero/core/test/unit_tests/fixtures/host_pointer_manager_fixture.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2022-2024 Intel Corporation
+ * Copyright (C) 2022-2025 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
@@ -36,6 +36,15 @@ void HostPointerManagerFixure::setUp() {
    hostDriverHandle = std::make_unique<L0::ult::DriverHandle>();
    hostDriverHandle->initialize(std::move(devices));
    device = hostDriverHandle->devices[0];
+    if (neoDevice->getPreemptionMode() == NEO::PreemptionMode::MidThread) {
+        for (auto &engine : neoDevice->getAllEngines()) {
+            NEO::CommandStreamReceiver *csr = engine.commandStreamReceiver;
+            if (!csr->getPreemptionAllocation()) {
+                csr->createPreemptionAllocation();
+            }
+        }
+    }
+
    openHostPointerManager = static_cast<L0::ult::HostPointerManager *>(hostDriverHandle->hostPointerManager.get());

    heapPointer = hostDriverHandle->getMemoryManager()->allocateSystemMemory(heapSize, MemoryConstants::pageSize);
--- a/level_zero/core/test/unit_tests/sources/cmdlist/test_cmdlist_6.cpp
+++ b/level_zero/core/test/unit_tests/sources/cmdlist/test_cmdlist_6.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2022-2024 Intel Corporation
+ * Copyright (C) 2022-2025 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
@@ -2913,10 +2913,9 @@ HWTEST2_F(CommandListStateBaseAddressGlobalStatelessTest,
    auto otherCsr = std::unique_ptr<UltCommandStreamReceiver<FamilyType>>(static_cast<UltCommandStreamReceiver<FamilyType> *>(createCommandStream(*device->getNEODevice()->getExecutionEnvironment(), 0, 1)));

    otherCsr->setupContext(*neoDevice->getDefaultEngine().osContext);
-    otherCsr->initializeResources(false);
+    otherCsr->initializeResources(false, neoDevice->getPreemptionMode());
    otherCsr->initializeTagAllocation();
    otherCsr->createGlobalFenceAllocation();
-    otherCsr->createPreemptionAllocation();
    otherCsr->createGlobalStatelessHeap();

    ze_command_queue_desc_t desc = {};
--- a/level_zero/core/test/unit_tests/sources/device/test_l0_device.cpp
+++ b/level_zero/core/test/unit_tests/sources/device/test_l0_device.cpp
@@ -753,6 +753,14 @@ struct DeviceTest : public ::testing::Test {
        driverHandle = std::make_unique<Mock<L0::DriverHandleImp>>();
        driverHandle->initialize(std::move(devices));
        device = driverHandle->devices[0];
+        if (neoDevice->getPreemptionMode() == NEO::PreemptionMode::MidThread) {
+            for (auto &engine : neoDevice->getAllEngines()) {
+                NEO::CommandStreamReceiver *csr = engine.commandStreamReceiver;
+                if (!csr->getPreemptionAllocation()) {
+                    csr->createPreemptionAllocation();
+                }
+            }
+        }
    }

    void TearDown() override {
@@ -1799,6 +1807,14 @@ struct GlobalTimestampTest : public ::testing::Test {
    void SetUp() override {
        debugManager.flags.CreateMultipleRootDevices.set(numRootDevices);
        neoDevice = NEO::MockDevice::createWithNewExecutionEnvironment<NEO::MockDevice>(NEO::defaultHwInfo.get(), rootDeviceIndex);
+        if (neoDevice->getPreemptionMode() == NEO::PreemptionMode::MidThread) {
+            for (auto &engine : neoDevice->getAllEngines()) {
+                NEO::CommandStreamReceiver *csr = engine.commandStreamReceiver;
+                if (!csr->getPreemptionAllocation()) {
+                    csr->createPreemptionAllocation();
+                }
+            }
+        }
    }

    DebugManagerStateRestore restorer;
@@ -2467,6 +2483,12 @@ struct MultipleDevicesFixture : public ::testing::Test {
            auto neoDevice = device->getNEODevice();
            context->rootDeviceIndices.pushUnique(neoDevice->getRootDeviceIndex());
            context->deviceBitfields.insert({neoDevice->getRootDeviceIndex(), neoDevice->getDeviceBitfield()});
+            if (neoDevice->getPreemptionMode() == NEO::PreemptionMode::MidThread) {
+                auto &csr = neoDevice->getInternalEngine().commandStreamReceiver;
+                if (!csr->getPreemptionAllocation()) {
+                    csr->createPreemptionAllocation();
+                }
+            }
        }
    }

--- a/level_zero/core/test/unit_tests/sources/event/test_event.cpp
+++ b/level_zero/core/test/unit_tests/sources/event/test_event.cpp
@@ -3465,7 +3465,7 @@ HWTEST_F(EventTests, GivenEventWhenHostSynchronizeCalledThenExpectDownloadEventA
    };

    auto ultCsr = static_cast<UltCommandStreamReceiver<FamilyType> *>(event->csrs[0]);
-    ultCsr->initializeResources(false);
+    ultCsr->initializeResources(false, device->getDevicePreemptionMode());
    VariableBackup<std::function<void(GraphicsAllocation & gfxAllocation)>> backupCsrDownloadImpl(&ultCsr->downloadAllocationImpl);
    ultCsr->downloadAllocationImpl = [&downloadAllocationTrack](GraphicsAllocation &gfxAllocation) {
        downloadAllocationTrack[&gfxAllocation]++;
@@ -3548,7 +3548,7 @@ HWTEST_F(EventContextGroupTests, givenSecondaryCsrWhenDownloadingAllocationThenU
    OsContext osContext(0, static_cast<uint32_t>(neoDevice->getAllEngines().size()), EngineDescriptorHelper::getDefaultDescriptor());

    ultCsr->setupContext(osContext);
-    ultCsr->initializeResources(false);
+    ultCsr->initializeResources(false, device->getDevicePreemptionMode());

    uint32_t downloadCounter = 0;
    ultCsr->downloadAllocationImpl = [&downloadCounter](GraphicsAllocation &gfxAllocation) {
@@ -3588,7 +3588,7 @@ HWTEST_F(EventTests, GivenEventUsedOnNonDefaultCsrWhenHostSynchronizeCalledThenA
    EXPECT_LT(1u, neoDevice->getAllEngines().size());

    auto ultCsr = static_cast<UltCommandStreamReceiver<FamilyType> *>(neoDevice->getAllEngines()[1].commandStreamReceiver);
-    ultCsr->initializeResources(false);
+    ultCsr->initializeResources(false, device->getDevicePreemptionMode());
    EXPECT_NE(event->csrs[0], ultCsr);

    VariableBackup<std::function<void(GraphicsAllocation & gfxAllocation)>> backupCsrDownloadImpl(&ultCsr->downloadAllocationImpl);
@@ -3636,7 +3636,7 @@ HWTEST_F(EventTests, givenInOrderEventWhenHostSynchronizeIsCalledThenAllocationI
    *eventAddress = Event::STATE_SIGNALED;

    auto ultCsr = static_cast<UltCommandStreamReceiver<FamilyType> *>(event->csrs[0]);
-    ultCsr->initializeResources(false);
+    ultCsr->initializeResources(false, device->getDevicePreemptionMode());

    VariableBackup<std::function<void(GraphicsAllocation & gfxAllocation)>> backupCsrDownloadImpl(&ultCsr->downloadAllocationImpl);
    ultCsr->downloadAllocationImpl = [&downloadAllocationTrack](GraphicsAllocation &gfxAllocation) {
@@ -3738,7 +3738,7 @@ HWTEST_F(EventTests, givenInOrderEventWithHostAllocWhenHostSynchronizeIsCalledTh
    *eventAddress = Event::STATE_SIGNALED;

    auto ultCsr = static_cast<UltCommandStreamReceiver<FamilyType> *>(event->csrs[0]);
-    ultCsr->initializeResources(false);
+    ultCsr->initializeResources(false, device->getDevicePreemptionMode());
    VariableBackup<std::function<void(GraphicsAllocation & gfxAllocation)>> backupCsrDownloadImpl(&ultCsr->downloadAllocationImpl);
    ultCsr->downloadAllocationImpl = [&downloadAllocationTrack](GraphicsAllocation &gfxAllocation) {
        downloadAllocationTrack[&gfxAllocation]++;
@@ -4511,9 +4511,9 @@ HWTEST2_F(EventMultiTileDynamicPacketUseTest, givenEventUsedCreatedOnSubDeviceBu
    auto ultCsr1 = static_cast<UltCommandStreamReceiver<FamilyType> *>(subDevice1->getNEODevice()->getDefaultEngine().commandStreamReceiver);
    auto ultCsr2 = static_cast<UltCommandStreamReceiver<FamilyType> *>(subDevice1->getNEODevice()->getInternalEngine().commandStreamReceiver);

-    rootCsr->initializeResources(false);
-    ultCsr0->initializeResources(false);
-    ultCsr1->initializeResources(false);
+    rootCsr->initializeResources(false, device->getDevicePreemptionMode());
+    ultCsr0->initializeResources(false, device->getDevicePreemptionMode());
+    ultCsr1->initializeResources(false, device->getDevicePreemptionMode());

    rootCsr->commandStreamReceiverType = CommandStreamReceiverType::tbx;
    ultCsr0->commandStreamReceiverType = CommandStreamReceiverType::tbx;
@@ -4599,8 +4599,8 @@ HWTEST2_F(EventMultiTileDynamicPacketUseTest, givenEventCounterBasedUsedCreatedO
    auto ultCsr0 = static_cast<UltCommandStreamReceiver<FamilyType> *>(subDevice0->getNEODevice()->getDefaultEngine().commandStreamReceiver);
    auto ultCsr1 = static_cast<UltCommandStreamReceiver<FamilyType> *>(subDevice1->getNEODevice()->getDefaultEngine().commandStreamReceiver);

-    ultCsr0->initializeResources(false);
-    ultCsr1->initializeResources(false);
+    ultCsr0->initializeResources(false, device->getDevicePreemptionMode());
+    ultCsr1->initializeResources(false, device->getDevicePreemptionMode());

    ultCsr0->commandStreamReceiverType = CommandStreamReceiverType::tbx;
    ultCsr1->commandStreamReceiverType = CommandStreamReceiverType::tbx;
--- a/opencl/source/command_queue/command_queue.cpp
+++ b/opencl/source/command_queue/command_queue.cpp
@@ -238,7 +238,7 @@ void CommandQueue::initializeGpgpuInternals() const {
        }
    }

-    gpgpuEngine->commandStreamReceiver->initializeResources(false);
+    gpgpuEngine->commandStreamReceiver->initializeResources(false, device->getPreemptionMode());
    gpgpuEngine->commandStreamReceiver->requestPreallocation();
    gpgpuEngine->commandStreamReceiver->initDirectSubmission();

@@ -421,7 +421,7 @@ void CommandQueue::constructBcsEnginesForSplit() {

            if (bcsEngines[i]) {
                bcsQueueEngineType = engineType;
-                bcsEngines[i]->commandStreamReceiver->initializeResources(false);
+                bcsEngines[i]->commandStreamReceiver->initializeResources(false, device->getPreemptionMode());
                bcsEngines[i]->commandStreamReceiver->initDirectSubmission();
            }
        }
--- a/opencl/test/unit_test/fixtures/cl_device_fixture.cpp
+++ b/opencl/test/unit_test/fixtures/cl_device_fixture.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020-2024 Intel Corporation
+ * Copyright (C) 2020-2025 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
@@ -33,6 +33,10 @@ void ClDeviceFixture::setUpImpl(const NEO::HardwareInfo *hardwareInfo) {
    pTagMemory = commandStreamReceiver.getTagAddress();
    ASSERT_NE(nullptr, const_cast<TagAddressType *>(pTagMemory));
    this->osContext = pDevice->getDefaultEngine().osContext;
+    if (pDevice->getPreemptionMode() == NEO::PreemptionMode::MidThread &&
+        !commandStreamReceiver.getPreemptionAllocation()) {
+        commandStreamReceiver.createPreemptionAllocation();
+    }
 }

 void ClDeviceFixture::tearDown() {
--- a/shared/source/command_stream/command_stream_receiver.cpp
+++ b/shared/source/command_stream/command_stream_receiver.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2024 Intel Corporation
+ * Copyright (C) 2018-2025 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
@@ -331,13 +331,17 @@ uint8_t CommandStreamReceiver::getUmdPowerHintValue() const {
    return this->osContext ? this->osContext->getUmdPowerHintValue() : 0u;
 }

-bool CommandStreamReceiver::initializeResources(bool allocateInterrupt) {
+bool CommandStreamReceiver::initializeResources(bool allocateInterrupt, const PreemptionMode preemptionMode) {
    if (!resourcesInitialized) {
        auto lock = obtainUniqueOwnership();
        if (!resourcesInitialized) {
            if (!osContext->ensureContextInitialized(allocateInterrupt)) {
                return false;
            }
+            if (preemptionMode == NEO::PreemptionMode::MidThread &&
+                !this->getPreemptionAllocation()) {
+                this->createPreemptionAllocation();
+            }
            this->fillReusableAllocationsList();
            this->resourcesInitialized = true;
        }
--- a/shared/source/command_stream/command_stream_receiver.h
+++ b/shared/source/command_stream/command_stream_receiver.h
@@ -284,7 +284,7 @@ class CommandStreamReceiver {
    virtual void setupContext(OsContext &osContext) { this->osContext = &osContext; }
    OsContext &getOsContext() const { return *osContext; }
    uint8_t getUmdPowerHintValue() const;
-    bool initializeResources(bool allocateInterrupt);
+    bool initializeResources(bool allocateInterrupt, const PreemptionMode preemptionMode);
    TagAllocatorBase *getEventTsAllocator();
    TagAllocatorBase *getEventPerfCountAllocator(const uint32_t tagSize);
    virtual TagAllocatorBase *getTimestampPacketAllocator() = 0;
--- a/shared/source/command_stream/command_stream_receiver_hw_base.inl
+++ b/shared/source/command_stream/command_stream_receiver_hw_base.inl
@@ -997,7 +997,7 @@ TaskCountType CommandStreamReceiverHw<GfxFamily>::flushBcsTask(const BlitPropert
    auto newTaskCount = taskCount + 1;
    latestSentTaskCount = newTaskCount;

-    this->initializeResources(false);
+    this->initializeResources(false, device.getPreemptionMode());
    this->initDirectSubmission();

    if (PauseOnGpuProperties::pauseModeAllowed(debugManager.flags.PauseOnBlitCopy.get(), taskCount, PauseOnGpuProperties::PauseMode::BeforeWorkload)) {
--- a/shared/source/command_stream/command_stream_receiver_hw_heap_addressing.inl
+++ b/shared/source/command_stream/command_stream_receiver_hw_heap_addressing.inl
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2024 Intel Corporation
+ * Copyright (C) 2024-2025 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
@@ -14,6 +14,7 @@ SubmissionStatus CommandStreamReceiverHw<GfxFamily>::initializeDeviceWithFirstSu
    if (this->latestFlushedTaskCount > 0) {
        return SubmissionStatus::success;
    }
+
    auto status = flushTagUpdate();

    if (isTbxMode() && (status == SubmissionStatus::success)) {
--- a/shared/source/device/device.cpp
+++ b/shared/source/device/device.cpp
@@ -477,7 +477,7 @@ bool Device::createEngine(EngineTypeUsage engineTypeUsage) {
    commandStreamReceiver->setupContext(*osContext);

    if (osContext->isImmediateContextInitializationEnabled(isDefaultEngine)) {
-        if (!commandStreamReceiver->initializeResources(false)) {
+        if (!commandStreamReceiver->initializeResources(false, this->getPreemptionMode())) {
            return false;
        }
    }
@@ -490,10 +490,6 @@ bool Device::createEngine(EngineTypeUsage engineTypeUsage) {
        return false;
    }

-    if (preemptionMode == PreemptionMode::MidThread && !commandStreamReceiver->createPreemptionAllocation()) {
-        return false;
-    }
-
    EngineControl engine{commandStreamReceiver.get(), osContext};
    allEngines.push_back(engine);
    if (engineUsage == EngineUsage::regular) {
@@ -534,7 +530,7 @@ bool Device::initializeEngines() {
        bool initializeDevice = (engine.osContext->isPartOfContextGroup() || isHeaplessStateInit) && !firstSubmissionDone;

        if (initializeDevice) {
-            engine.commandStreamReceiver->initializeResources(false);
+            engine.commandStreamReceiver->initializeResources(false, this->getPreemptionMode());

            if (debugManager.flags.DeferStateInitSubmissionToFirstRegularUsage.get() != 1) {
                engine.commandStreamReceiver->initializeDeviceWithFirstSubmission(*this);
@@ -604,7 +600,7 @@ EngineControl *Device::getSecondaryEngineCsr(EngineTypeUsage engineTypeUsage, bo

            EngineDescriptor engineDescriptor(engineTypeUsage, getDeviceBitfield(), preemptionMode, false);

-            if (!commandStreamReceiver->initializeResources(allocateInterrupt)) {
+            if (!commandStreamReceiver->initializeResources(allocateInterrupt, this->getPreemptionMode())) {
                return nullptr;
            }

--- a/shared/source/device/root_device.cpp
+++ b/shared/source/device/root_device.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2019-2024 Intel Corporation
+ * Copyright (C) 2019-2025 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
@@ -91,14 +91,11 @@ void RootDevice::initializeRootCommandStreamReceiver() {
    osContext->setIsPrimaryEngine(isPrimaryEngine);

    rootCommandStreamReceiver->setupContext(*osContext);
-    rootCommandStreamReceiver->initializeResources(false);
+    rootCommandStreamReceiver->initializeResources(false, preemptionMode);
    rootCsrCreated = true;
    rootCommandStreamReceiver->initializeTagAllocation();
    rootCommandStreamReceiver->createGlobalFenceAllocation();
    rootCommandStreamReceiver->createWorkPartitionAllocation(*this);
-    if (preemptionMode == PreemptionMode::MidThread) {
-        rootCommandStreamReceiver->createPreemptionAllocation();
-    }
    commandStreamReceivers.push_back(std::move(rootCommandStreamReceiver));

    EngineControl engine{commandStreamReceivers.back().get(), osContext};
--- a/shared/source/helpers/blit_helper.cpp
+++ b/shared/source/helpers/blit_helper.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023-2024 Intel Corporation
+ * Copyright (C) 2023-2025 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
@@ -55,7 +55,7 @@ BlitOperationResult BlitHelper::blitMemoryToAllocationBanks(const Device &device
            return BlitOperationResult::unsupported;
        }

-        bcsEngine->commandStreamReceiver->initializeResources(false);
+        bcsEngine->commandStreamReceiver->initializeResources(false, device.getPreemptionMode());
        bcsEngine->commandStreamReceiver->initDirectSubmission();
        BlitPropertiesContainer blitPropertiesContainer;
        blitPropertiesContainer.push_back(
--- a/shared/source/os_interface/linux/drm_memory_operations_handler_bind.cpp
+++ b/shared/source/os_interface/linux/drm_memory_operations_handler_bind.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020-2024 Intel Corporation
+ * Copyright (C) 2020-2025 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
@@ -30,7 +30,7 @@ MemoryOperationsStatus DrmMemoryOperationsHandlerBind::makeResident(Device *devi
    auto &engines = device->getAllEngines();
    MemoryOperationsStatus result = MemoryOperationsStatus::success;
    for (const auto &engine : engines) {
-        engine.commandStreamReceiver->initializeResources(false);
+        engine.commandStreamReceiver->initializeResources(false, device->getPreemptionMode());
        result = this->makeResidentWithinOsContext(engine.osContext, gfxAllocations, false);
        if (result != MemoryOperationsStatus::success) {
            break;
--- a/shared/test/common/fixtures/command_stream_receiver_fixture.cpp
+++ b/shared/test/common/fixtures/command_stream_receiver_fixture.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2022-2024 Intel Corporation
+ * Copyright (C) 2022-2025 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
@@ -37,6 +37,13 @@ void CommandStreamReceiverFixture::setUp() {
    requiredStreamProperties.initSupport(pDevice->getRootDeviceEnvironment());
    immediateFlushTaskFlags.requiredState = &requiredStreamProperties;
    immediateFlushTaskFlags.sshCpuBase = sshBuffer;
+
+    if (pDevice->getPreemptionMode() == NEO::PreemptionMode::MidThread) {
+        auto &commandStreamReceiver = pDevice->getGpgpuCommandStreamReceiver();
+        if (!commandStreamReceiver.getPreemptionAllocation()) {
+            ASSERT_TRUE(commandStreamReceiver.createPreemptionAllocation());
+        }
+    }
 }

 void CommandStreamReceiverFixture::tearDown() {
--- a/shared/test/common/mocks/ult_device_factory.cpp
+++ b/shared/test/common/mocks/ult_device_factory.cpp
@@ -47,7 +47,16 @@ UltDeviceFactory::UltDeviceFactory(uint32_t rootDevicesCount, uint32_t subDevice
        pCreatedDevice->incRefInternal();
        if (pCreatedDevice->getNumSubDevices() > 1) {
            for (uint32_t i = 0; i < pCreatedDevice->getNumSubDevices(); i++) {
-                this->subDevices.push_back(static_cast<SubDevice *>(pCreatedDevice->getSubDevice(i)));
+                auto *pDevice = static_cast<SubDevice *>(pCreatedDevice->getSubDevice(i));
+                this->subDevices.push_back(pDevice);
+            }
+        }
+        if (pCreatedDevice->getPreemptionMode() == NEO::PreemptionMode::MidThread) {
+            for (auto &engine : pCreatedDevice->getAllEngines()) {
+                NEO::CommandStreamReceiver *csr = engine.commandStreamReceiver;
+                if (!csr->getPreemptionAllocation()) {
+                    csr->createPreemptionAllocation();
+                }
            }
        }
        this->rootDevices.push_back(static_cast<MockDevice *>(pCreatedDevice.release()));
--- a/shared/test/unit_test/command_stream/command_stream_receiver_tests.cpp
+++ b/shared/test/unit_test/command_stream/command_stream_receiver_tests.cpp
@@ -137,9 +137,9 @@ HWTEST_F(CommandStreamReceiverTest, WhenInitializeResourcesThenCallFillReusableA
    ultCsr.fillReusableAllocationsListCalled = 0u;
    ultCsr.resourcesInitialized = false;

-    commandStreamReceiver->initializeResources(false);
+    commandStreamReceiver->initializeResources(false, pDevice->getPreemptionMode());
    EXPECT_EQ(1u, pDevice->getUltCommandStreamReceiver<FamilyType>().fillReusableAllocationsListCalled);
-    commandStreamReceiver->initializeResources(false);
+    commandStreamReceiver->initializeResources(false, pDevice->getPreemptionMode());
    EXPECT_EQ(1u, pDevice->getUltCommandStreamReceiver<FamilyType>().fillReusableAllocationsListCalled);
 }

@@ -161,7 +161,7 @@ HWTEST_F(CommandStreamReceiverTest, whenContextCreateReturnsFalseThenExpectCSRIn
    auto &ultCsr = pDevice->getUltCommandStreamReceiver<FamilyType>();
    ultCsr.resourcesInitialized = false;
    ultCsr.setupContext(osContext);
-    bool ret = ultCsr.initializeResources(false);
+    bool ret = ultCsr.initializeResources(false, pDevice->getPreemptionMode());
    EXPECT_FALSE(ret);
 }

--- a/shared/test/unit_test/preemption/preemption_tests.cpp
+++ b/shared/test/unit_test/preemption/preemption_tests.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2024 Intel Corporation
+ * Copyright (C) 2018-2025 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
@@ -442,6 +442,10 @@ HWTEST_F(MidThreadPreemptionTests, GivenNoWaWhenCreatingCsrSurfaceThenSurfaceIsC
    ASSERT_NE(nullptr, mockDevice.get());

    auto &csr = mockDevice->getUltCommandStreamReceiver<FamilyType>();
+    if (mockDevice->getPreemptionMode() == NEO::PreemptionMode::MidThread &&
+        !csr.getPreemptionAllocation()) {
+        ASSERT_TRUE(csr.createPreemptionAllocation());
+    }
    MemoryAllocation *csrSurface = static_cast<MemoryAllocation *>(csr.getPreemptionAllocation());
    ASSERT_NE(nullptr, csrSurface);
    EXPECT_FALSE(csrSurface->uncacheable);
--- a/shared/test/unit_test/preemption/test_preemption_xe2_and_later.cpp
+++ b/shared/test/unit_test/preemption/test_preemption_xe2_and_later.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2024 Intel Corporation
+ * Copyright (C) 2024-2025 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
@@ -41,6 +41,12 @@ struct Xe2MidThreadPreemptionTests : public Xe2PreemptionTests {
        preemptionMode = PreemptionMode::MidThread;

        Xe2PreemptionTests::SetUp();
+
+        auto &csr = device->getGpgpuCommandStreamReceiver();
+        if (device->getPreemptionMode() == NEO::PreemptionMode::MidThread &&
+            !csr.getPreemptionAllocation()) {
+            csr.createPreemptionAllocation();
+        }
    }
    void TearDown() override {
        Xe2PreemptionTests::TearDown();