Unify preemption programming in PreemptionHelper

Change-Id: I70a02ed963c809f1998a56c94169992b1a611450 Signed-off-by: Zbigniew Zdanowicz <zbigniew.zdanowicz@intel.com>
2026-01-04 15:53:45 +08:00 · 2020-02-07 14:46:11 +01:00
parent 9a7dbfe4f1
commit 888d2ddeae
6 changed files with 249 additions and 107 deletions
--- a/core/command_stream/preemption.inl
+++ b/core/command_stream/preemption.inl
@@ -14,50 +14,6 @@

 namespace NEO {

-template <typename GfxFamily>
-size_t PreemptionHelper::getPreemptionWaCsSize(const Device &device) {
-    typedef typename GfxFamily::MI_LOAD_REGISTER_IMM MI_LOAD_REGISTER_IMM;
-    size_t size = 0;
-    PreemptionMode preemptionMode = device.getPreemptionMode();
-    if (preemptionMode == PreemptionMode::ThreadGroup ||
-        preemptionMode == PreemptionMode::MidThread) {
-        if (device.getHardwareInfo().workaroundTable.waModifyVFEStateAfterGPGPUPreemption) {
-            size += 2 * sizeof(MI_LOAD_REGISTER_IMM);
-        }
-    }
-    return size;
-}
-
-template <typename GfxFamily>
-void PreemptionHelper::applyPreemptionWaCmdsBegin(LinearStream *pCommandStream, const Device &device) {
-    typedef typename GfxFamily::MI_LOAD_REGISTER_IMM MI_LOAD_REGISTER_IMM;
-    PreemptionMode preemptionMode = device.getPreemptionMode();
-    if (preemptionMode == PreemptionMode::ThreadGroup ||
-        preemptionMode == PreemptionMode::MidThread) {
-        if (device.getHardwareInfo().workaroundTable.waModifyVFEStateAfterGPGPUPreemption) {
-            auto pCmd = reinterpret_cast<MI_LOAD_REGISTER_IMM *>(pCommandStream->getSpace(sizeof(MI_LOAD_REGISTER_IMM)));
-            *pCmd = GfxFamily::cmdInitLoadRegisterImm;
-            pCmd->setRegisterOffset(CS_GPR_R0);
-            pCmd->setDataDword(GPGPU_WALKER_COOKIE_VALUE_BEFORE_WALKER);
-        }
-    }
-}
-
-template <typename GfxFamily>
-void PreemptionHelper::applyPreemptionWaCmdsEnd(LinearStream *pCommandStream, const Device &device) {
-    typedef typename GfxFamily::MI_LOAD_REGISTER_IMM MI_LOAD_REGISTER_IMM;
-    PreemptionMode preemptionMode = device.getPreemptionMode();
-    if (preemptionMode == PreemptionMode::ThreadGroup ||
-        preemptionMode == PreemptionMode::MidThread) {
-        if (device.getHardwareInfo().workaroundTable.waModifyVFEStateAfterGPGPUPreemption) {
-            auto pCmd = reinterpret_cast<MI_LOAD_REGISTER_IMM *>(pCommandStream->getSpace(sizeof(MI_LOAD_REGISTER_IMM)));
-            *pCmd = GfxFamily::cmdInitLoadRegisterImm;
-            pCmd->setRegisterOffset(CS_GPR_R0);
-            pCmd->setDataDword(GPGPU_WALKER_COOKIE_VALUE_AFTER_WALKER);
-        }
-    }
-}
-
 template <typename GfxFamily>
 void PreemptionHelper::programCsrBaseAddress(LinearStream &preambleCmdStream, Device &device, const GraphicsAllocation *preemptionCsr) {
    using GPGPU_CSR_BASE_ADDRESS = typename GfxFamily::GPGPU_CSR_BASE_ADDRESS;
@@ -131,8 +87,26 @@ size_t PreemptionHelper::getRequiredStateSipCmdSize(const Device &device) {
    return size;
 }

+template <typename GfxFamily>
+size_t PreemptionHelper::getPreemptionWaCsSize(const Device &device) {
+    return 0u;
+}
+template <typename GfxFamily>
+void PreemptionHelper::applyPreemptionWaCmdsBegin(LinearStream *pCommandStream, const Device &device) {
+}
+
+template <typename GfxFamily>
+void PreemptionHelper::applyPreemptionWaCmdsEnd(LinearStream *pCommandStream, const Device &device) {
+}
+
 template <typename GfxFamily>
 void PreemptionHelper::programInterfaceDescriptorDataPreemption(INTERFACE_DESCRIPTOR_DATA<GfxFamily> *idd, PreemptionMode preemptionMode) {
+    using INTERFACE_DESCRIPTOR_DATA = typename GfxFamily::INTERFACE_DESCRIPTOR_DATA;
+    if (preemptionMode == PreemptionMode::MidThread) {
+        idd->setThreadPreemptionDisable(INTERFACE_DESCRIPTOR_DATA::THREAD_PREEMPTION_DISABLE_DISABLE);
+    } else {
+        idd->setThreadPreemptionDisable(INTERFACE_DESCRIPTOR_DATA::THREAD_PREEMPTION_DISABLE_ENABLE);
+    }
 }

 template <typename GfxFamily>
--- a/core/gen11/preemption_gen11.cpp
+++ b/core/gen11/preemption_gen11.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2019 Intel Corporation
+ * Copyright (C) 2018-2020 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
@@ -7,24 +7,10 @@

 #include "core/command_stream/preemption.h"
 #include "core/command_stream/preemption.inl"
-#include "core/memory_manager/graphics_allocation.h"

 namespace NEO {

-typedef ICLFamily GfxFamily;
-
-template <>
-size_t PreemptionHelper::getPreemptionWaCsSize<GfxFamily>(const Device &device) {
-    return 0;
-}
-
-template <>
-void PreemptionHelper::applyPreemptionWaCmdsBegin<GfxFamily>(LinearStream *pCommandStream, const Device &device) {
-}
-
-template <>
-void PreemptionHelper::applyPreemptionWaCmdsEnd<GfxFamily>(LinearStream *pCommandStream, const Device &device) {
-}
+using GfxFamily = ICLFamily;

 template void PreemptionHelper::programCmdStream<GfxFamily>(LinearStream &cmdStream, PreemptionMode newPreemptionMode,
                                                            PreemptionMode oldPreemptionMode, GraphicsAllocation *preemptionCsr);
@@ -33,14 +19,8 @@ template void PreemptionHelper::programCsrBaseAddress<GfxFamily>(LinearStream &p
 template void PreemptionHelper::programStateSip<GfxFamily>(LinearStream &preambleCmdStream, Device &device);
 template size_t PreemptionHelper::getRequiredStateSipCmdSize<GfxFamily>(const Device &device);
 template size_t PreemptionHelper::getRequiredCmdStreamSize<GfxFamily>(PreemptionMode newPreemptionMode, PreemptionMode oldPreemptionMode);
-
-template <>
-void PreemptionHelper::programInterfaceDescriptorDataPreemption<GfxFamily>(INTERFACE_DESCRIPTOR_DATA<GfxFamily> *idd, PreemptionMode preemptionMode) {
-    using INTERFACE_DESCRIPTOR_DATA = typename GfxFamily::INTERFACE_DESCRIPTOR_DATA;
-    if (preemptionMode == PreemptionMode::MidThread) {
-        idd->setThreadPreemptionDisable(INTERFACE_DESCRIPTOR_DATA::THREAD_PREEMPTION_DISABLE_DISABLE);
-    } else {
-        idd->setThreadPreemptionDisable(INTERFACE_DESCRIPTOR_DATA::THREAD_PREEMPTION_DISABLE_ENABLE);
-    }
-}
+template size_t PreemptionHelper::getPreemptionWaCsSize<GfxFamily>(const Device &device);
+template void PreemptionHelper::applyPreemptionWaCmdsBegin<GfxFamily>(LinearStream *pCommandStream, const Device &device);
+template void PreemptionHelper::applyPreemptionWaCmdsEnd<GfxFamily>(LinearStream *pCommandStream, const Device &device);
+template void PreemptionHelper::programInterfaceDescriptorDataPreemption<GfxFamily>(INTERFACE_DESCRIPTOR_DATA<GfxFamily> *idd, PreemptionMode preemptionMode);
 } // namespace NEO
--- a/core/gen12lp/preemption_gen12lp.cpp
+++ b/core/gen12lp/preemption_gen12lp.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2019 Intel Corporation
+ * Copyright (C) 2019-2020 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
@@ -10,20 +10,7 @@

 namespace NEO {

-typedef TGLLPFamily GfxFamily;
-
-template <>
-size_t PreemptionHelper::getPreemptionWaCsSize<GfxFamily>(const Device &device) {
-    return 0;
-}
-
-template <>
-void PreemptionHelper::applyPreemptionWaCmdsBegin<GfxFamily>(LinearStream *pCommandStream, const Device &device) {
-}
-
-template <>
-void PreemptionHelper::applyPreemptionWaCmdsEnd<GfxFamily>(LinearStream *pCommandStream, const Device &device) {
-}
+using GfxFamily = TGLLPFamily;

 template void PreemptionHelper::programCmdStream<GfxFamily>(LinearStream &cmdStream, PreemptionMode newPreemptionMode,
                                                            PreemptionMode oldPreemptionMode, GraphicsAllocation *preemptionCsr);
@@ -32,14 +19,8 @@ template void PreemptionHelper::programCsrBaseAddress<GfxFamily>(LinearStream &p
 template void PreemptionHelper::programStateSip<GfxFamily>(LinearStream &preambleCmdStream, Device &device);
 template size_t PreemptionHelper::getRequiredStateSipCmdSize<GfxFamily>(const Device &device);
 template size_t PreemptionHelper::getRequiredCmdStreamSize<GfxFamily>(PreemptionMode newPreemptionMode, PreemptionMode oldPreemptionMode);
-
-template <>
-void PreemptionHelper::programInterfaceDescriptorDataPreemption<GfxFamily>(INTERFACE_DESCRIPTOR_DATA<GfxFamily> *idd, PreemptionMode preemptionMode) {
-    using INTERFACE_DESCRIPTOR_DATA = typename GfxFamily::INTERFACE_DESCRIPTOR_DATA;
-    if (preemptionMode == PreemptionMode::MidThread) {
-        idd->setThreadPreemptionDisable(INTERFACE_DESCRIPTOR_DATA::THREAD_PREEMPTION_DISABLE_DISABLE);
-    } else {
-        idd->setThreadPreemptionDisable(INTERFACE_DESCRIPTOR_DATA::THREAD_PREEMPTION_DISABLE_ENABLE);
-    }
-}
+template size_t PreemptionHelper::getPreemptionWaCsSize<GfxFamily>(const Device &device);
+template void PreemptionHelper::applyPreemptionWaCmdsBegin<GfxFamily>(LinearStream *pCommandStream, const Device &device);
+template void PreemptionHelper::applyPreemptionWaCmdsEnd<GfxFamily>(LinearStream *pCommandStream, const Device &device);
+template void PreemptionHelper::programInterfaceDescriptorDataPreemption<GfxFamily>(INTERFACE_DESCRIPTOR_DATA<GfxFamily> *idd, PreemptionMode preemptionMode);
 } // namespace NEO
--- a/core/gen8/preemption_gen8.cpp
+++ b/core/gen8/preemption_gen8.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2019 Intel Corporation
+ * Copyright (C) 2018-2020 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
@@ -10,7 +10,7 @@

 namespace NEO {

-typedef BDWFamily GfxFamily;
+using GfxFamily = BDWFamily;

 template <>
 struct PreemptionConfig<GfxFamily> {
@@ -59,10 +59,53 @@ template <>
 void PreemptionHelper::programStateSip<GfxFamily>(LinearStream &preambleCmdStream, Device &device) {
 }

-template size_t PreemptionHelper::getRequiredCmdStreamSize<GfxFamily>(PreemptionMode newPreemptionMode, PreemptionMode oldPreemptionMode);
-template size_t PreemptionHelper::getPreemptionWaCsSize<GfxFamily>(const Device &device);
-template void PreemptionHelper::applyPreemptionWaCmdsBegin<GfxFamily>(LinearStream *pCommandStream, const Device &device);
-template void PreemptionHelper::applyPreemptionWaCmdsEnd<GfxFamily>(LinearStream *pCommandStream, const Device &device);
-template void PreemptionHelper::programInterfaceDescriptorDataPreemption<GfxFamily>(INTERFACE_DESCRIPTOR_DATA<GfxFamily> *idd, PreemptionMode preemptionMode);
+template <>
+size_t PreemptionHelper::getPreemptionWaCsSize<GfxFamily>(const Device &device) {
+    typedef typename GfxFamily::MI_LOAD_REGISTER_IMM MI_LOAD_REGISTER_IMM;
+    size_t size = 0;
+    PreemptionMode preemptionMode = device.getPreemptionMode();
+    if (preemptionMode == PreemptionMode::ThreadGroup ||
+        preemptionMode == PreemptionMode::MidThread) {
+        if (device.getHardwareInfo().workaroundTable.waModifyVFEStateAfterGPGPUPreemption) {
+            size += 2 * sizeof(MI_LOAD_REGISTER_IMM);
+        }
+    }
+    return size;
+}

+template <>
+void PreemptionHelper::applyPreemptionWaCmdsBegin<GfxFamily>(LinearStream *pCommandStream, const Device &device) {
+    typedef typename GfxFamily::MI_LOAD_REGISTER_IMM MI_LOAD_REGISTER_IMM;
+    PreemptionMode preemptionMode = device.getPreemptionMode();
+    if (preemptionMode == PreemptionMode::ThreadGroup ||
+        preemptionMode == PreemptionMode::MidThread) {
+        if (device.getHardwareInfo().workaroundTable.waModifyVFEStateAfterGPGPUPreemption) {
+            auto pCmd = reinterpret_cast<MI_LOAD_REGISTER_IMM *>(pCommandStream->getSpace(sizeof(MI_LOAD_REGISTER_IMM)));
+            *pCmd = GfxFamily::cmdInitLoadRegisterImm;
+            pCmd->setRegisterOffset(CS_GPR_R0);
+            pCmd->setDataDword(GPGPU_WALKER_COOKIE_VALUE_BEFORE_WALKER);
+        }
+    }
+}
+
+template <>
+void PreemptionHelper::applyPreemptionWaCmdsEnd<GfxFamily>(LinearStream *pCommandStream, const Device &device) {
+    typedef typename GfxFamily::MI_LOAD_REGISTER_IMM MI_LOAD_REGISTER_IMM;
+    PreemptionMode preemptionMode = device.getPreemptionMode();
+    if (preemptionMode == PreemptionMode::ThreadGroup ||
+        preemptionMode == PreemptionMode::MidThread) {
+        if (device.getHardwareInfo().workaroundTable.waModifyVFEStateAfterGPGPUPreemption) {
+            auto pCmd = reinterpret_cast<MI_LOAD_REGISTER_IMM *>(pCommandStream->getSpace(sizeof(MI_LOAD_REGISTER_IMM)));
+            *pCmd = GfxFamily::cmdInitLoadRegisterImm;
+            pCmd->setRegisterOffset(CS_GPR_R0);
+            pCmd->setDataDword(GPGPU_WALKER_COOKIE_VALUE_AFTER_WALKER);
+        }
+    }
+}
+
+template <>
+void PreemptionHelper::programInterfaceDescriptorDataPreemption<GfxFamily>(INTERFACE_DESCRIPTOR_DATA<GfxFamily> *idd, PreemptionMode preemptionMode) {
+}
+
+template size_t PreemptionHelper::getRequiredCmdStreamSize<GfxFamily>(PreemptionMode newPreemptionMode, PreemptionMode oldPreemptionMode);
 } // namespace NEO
--- a/core/gen9/preemption_gen9.cpp
+++ b/core/gen9/preemption_gen9.cpp
@@ -14,7 +14,55 @@

 namespace NEO {

-typedef SKLFamily GfxFamily;
+using GfxFamily = SKLFamily;
+
+template <>
+size_t PreemptionHelper::getPreemptionWaCsSize<GfxFamily>(const Device &device) {
+    typedef typename GfxFamily::MI_LOAD_REGISTER_IMM MI_LOAD_REGISTER_IMM;
+    size_t size = 0;
+    PreemptionMode preemptionMode = device.getPreemptionMode();
+    if (preemptionMode == PreemptionMode::ThreadGroup ||
+        preemptionMode == PreemptionMode::MidThread) {
+        if (device.getHardwareInfo().workaroundTable.waModifyVFEStateAfterGPGPUPreemption) {
+            size += 2 * sizeof(MI_LOAD_REGISTER_IMM);
+        }
+    }
+    return size;
+}
+
+template <>
+void PreemptionHelper::applyPreemptionWaCmdsBegin<GfxFamily>(LinearStream *pCommandStream, const Device &device) {
+    typedef typename GfxFamily::MI_LOAD_REGISTER_IMM MI_LOAD_REGISTER_IMM;
+    PreemptionMode preemptionMode = device.getPreemptionMode();
+    if (preemptionMode == PreemptionMode::ThreadGroup ||
+        preemptionMode == PreemptionMode::MidThread) {
+        if (device.getHardwareInfo().workaroundTable.waModifyVFEStateAfterGPGPUPreemption) {
+            auto pCmd = reinterpret_cast<MI_LOAD_REGISTER_IMM *>(pCommandStream->getSpace(sizeof(MI_LOAD_REGISTER_IMM)));
+            *pCmd = GfxFamily::cmdInitLoadRegisterImm;
+            pCmd->setRegisterOffset(CS_GPR_R0);
+            pCmd->setDataDword(GPGPU_WALKER_COOKIE_VALUE_BEFORE_WALKER);
+        }
+    }
+}
+
+template <>
+void PreemptionHelper::applyPreemptionWaCmdsEnd<GfxFamily>(LinearStream *pCommandStream, const Device &device) {
+    typedef typename GfxFamily::MI_LOAD_REGISTER_IMM MI_LOAD_REGISTER_IMM;
+    PreemptionMode preemptionMode = device.getPreemptionMode();
+    if (preemptionMode == PreemptionMode::ThreadGroup ||
+        preemptionMode == PreemptionMode::MidThread) {
+        if (device.getHardwareInfo().workaroundTable.waModifyVFEStateAfterGPGPUPreemption) {
+            auto pCmd = reinterpret_cast<MI_LOAD_REGISTER_IMM *>(pCommandStream->getSpace(sizeof(MI_LOAD_REGISTER_IMM)));
+            *pCmd = GfxFamily::cmdInitLoadRegisterImm;
+            pCmd->setRegisterOffset(CS_GPR_R0);
+            pCmd->setDataDword(GPGPU_WALKER_COOKIE_VALUE_AFTER_WALKER);
+        }
+    }
+}
+
+template <>
+void PreemptionHelper::programInterfaceDescriptorDataPreemption<GfxFamily>(INTERFACE_DESCRIPTOR_DATA<GfxFamily> *idd, PreemptionMode preemptionMode) {
+}

 template void PreemptionHelper::programCmdStream<GfxFamily>(LinearStream &cmdStream, PreemptionMode newPreemptionMode,
                                                            PreemptionMode oldPreemptionMode, GraphicsAllocation *preemptionCsr);
@@ -24,8 +72,5 @@ template void PreemptionHelper::programCsrBaseAddress<GfxFamily>(LinearStream &p
 template void PreemptionHelper::programStateSip<GfxFamily>(LinearStream &preambleCmdStream, Device &device);
 template size_t PreemptionHelper::getRequiredStateSipCmdSize<GfxFamily>(const Device &device);
 template size_t PreemptionHelper::getRequiredCmdStreamSize<GfxFamily>(PreemptionMode newPreemptionMode, PreemptionMode oldPreemptionMode);
-template size_t PreemptionHelper::getPreemptionWaCsSize<GfxFamily>(const Device &device);
-template void PreemptionHelper::applyPreemptionWaCmdsBegin<GfxFamily>(LinearStream *pCommandStream, const Device &device);
-template void PreemptionHelper::applyPreemptionWaCmdsEnd<GfxFamily>(LinearStream *pCommandStream, const Device &device);
-template void PreemptionHelper::programInterfaceDescriptorDataPreemption<GfxFamily>(INTERFACE_DESCRIPTOR_DATA<GfxFamily> *idd, PreemptionMode preemptionMode);
+
 } // namespace NEO
--- a/core/unit_tests/gen8/test_preemption_gen8.cpp
+++ b/core/unit_tests/gen8/test_preemption_gen8.cpp
@@ -5,8 +5,10 @@
 *
 */

+#include "core/command_stream/linear_stream.h"
 #include "core/command_stream/preemption.h"
 #include "core/helpers/hw_helper.h"
+#include "core/memory_manager/memory_constants.h"
 #include "core/unit_tests/fixtures/preemption_fixture.h"
 #include "unit_tests/command_queue/enqueue_fixture.h"
 #include "unit_tests/fixtures/hello_world_fixture.h"
@@ -194,3 +196,120 @@ GEN8TEST_F(Gen8PreemptionTests, givenInterfaceDescriptorDataWhenAnyPreemptionMod
    ret = memcmp(&idd, &iddArg, sizeof(INTERFACE_DESCRIPTOR_DATA));
    EXPECT_EQ(0, ret);
 }
+
+struct Gen8PreemptionTestsLinearStream : public Gen8PreemptionTests {
+    void SetUp() override {
+        Gen8PreemptionTests::SetUp();
+        cmdBufferAllocation = alignedMalloc(MemoryConstants::pageSize, MemoryConstants::pageSize);
+        cmdBuffer.replaceBuffer(cmdBufferAllocation, MemoryConstants::pageSize);
+    }
+
+    void TearDown() override {
+        alignedFree(cmdBufferAllocation);
+        Gen8PreemptionTests::TearDown();
+    }
+
+    LinearStream cmdBuffer;
+    void *cmdBufferAllocation;
+    HardwareParse cmdBufferParser;
+};
+
+GEN8TEST_F(Gen8PreemptionTestsLinearStream, givenMidBatchPreemptionWhenProgrammingWaCmdsBeginThenExpectNoCmds) {
+    device->setPreemptionMode(PreemptionMode::MidBatch);
+    PreemptionHelper::applyPreemptionWaCmdsBegin<FamilyType>(&cmdBuffer, device->getDevice());
+    EXPECT_EQ(0u, cmdBuffer.getUsed());
+}
+
+GEN8TEST_F(Gen8PreemptionTestsLinearStream, givenMidBatchPreemptionWhenProgrammingWaCmdsEndThenExpectNoCmds) {
+    device->setPreemptionMode(PreemptionMode::MidBatch);
+    PreemptionHelper::applyPreemptionWaCmdsEnd<FamilyType>(&cmdBuffer, device->getDevice());
+    EXPECT_EQ(0u, cmdBuffer.getUsed());
+}
+
+GEN8TEST_F(Gen8PreemptionTestsLinearStream, givenThreadGroupPreemptionNoWaSetWhenProgrammingWaCmdsBeginThenExpectNoCmd) {
+    device->setPreemptionMode(PreemptionMode::ThreadGroup);
+    device->getExecutionEnvironment()->getMutableHardwareInfo()->workaroundTable.waModifyVFEStateAfterGPGPUPreemption = false;
+    PreemptionHelper::applyPreemptionWaCmdsBegin<FamilyType>(&cmdBuffer, device->getDevice());
+    EXPECT_EQ(0u, cmdBuffer.getUsed());
+}
+
+GEN8TEST_F(Gen8PreemptionTestsLinearStream, givenThreadGroupPreemptionNoWaSetWhenProgrammingWaCmdsEndThenExpectNoCmd) {
+    device->setPreemptionMode(PreemptionMode::ThreadGroup);
+    device->getExecutionEnvironment()->getMutableHardwareInfo()->workaroundTable.waModifyVFEStateAfterGPGPUPreemption = false;
+    PreemptionHelper::applyPreemptionWaCmdsEnd<FamilyType>(&cmdBuffer, device->getDevice());
+    EXPECT_EQ(0u, cmdBuffer.getUsed());
+}
+
+GEN8TEST_F(Gen8PreemptionTestsLinearStream, givenThreadGroupPreemptionWaSetWhenProgrammingWaCmdsBeginThenExpectMmioCmd) {
+    using MI_LOAD_REGISTER_IMM = typename FamilyType::MI_LOAD_REGISTER_IMM;
+    device->setPreemptionMode(PreemptionMode::ThreadGroup);
+    device->getExecutionEnvironment()->getMutableHardwareInfo()->workaroundTable.waModifyVFEStateAfterGPGPUPreemption = true;
+    PreemptionHelper::applyPreemptionWaCmdsBegin<FamilyType>(&cmdBuffer, device->getDevice());
+
+    cmdBufferParser.parseCommands<FamilyType>(cmdBuffer);
+    cmdBufferParser.findHardwareCommands<FamilyType>();
+    GenCmdList::iterator itMmioCmd = cmdBufferParser.lriList.begin();
+    ASSERT_NE(cmdBufferParser.lriList.end(), itMmioCmd);
+    MI_LOAD_REGISTER_IMM *mmioCmd = genCmdCast<MI_LOAD_REGISTER_IMM *>(*itMmioCmd);
+    EXPECT_EQ(0x2600u, mmioCmd->getRegisterOffset());
+    EXPECT_EQ(0xFFFFFFFFu, mmioCmd->getDataDword());
+}
+
+GEN8TEST_F(Gen8PreemptionTestsLinearStream, givenThreadGroupPreemptionWaSetWhenProgrammingWaCmdsEndThenExpectMmioCmd) {
+    using MI_LOAD_REGISTER_IMM = typename FamilyType::MI_LOAD_REGISTER_IMM;
+    device->setPreemptionMode(PreemptionMode::ThreadGroup);
+    device->getExecutionEnvironment()->getMutableHardwareInfo()->workaroundTable.waModifyVFEStateAfterGPGPUPreemption = true;
+    PreemptionHelper::applyPreemptionWaCmdsEnd<FamilyType>(&cmdBuffer, device->getDevice());
+
+    cmdBufferParser.parseCommands<FamilyType>(cmdBuffer);
+    cmdBufferParser.findHardwareCommands<FamilyType>();
+    GenCmdList::iterator itMmioCmd = cmdBufferParser.lriList.begin();
+    ASSERT_NE(cmdBufferParser.lriList.end(), itMmioCmd);
+    MI_LOAD_REGISTER_IMM *mmioCmd = genCmdCast<MI_LOAD_REGISTER_IMM *>(*itMmioCmd);
+    EXPECT_EQ(0x2600u, mmioCmd->getRegisterOffset());
+    EXPECT_EQ(0x00000000u, mmioCmd->getDataDword());
+}
+
+GEN8TEST_F(Gen8PreemptionTestsLinearStream, givenMidThreadPreemptionNoWaSetWhenProgrammingWaCmdsBeginThenExpectNoCmd) {
+    device->setPreemptionMode(PreemptionMode::MidThread);
+    device->getExecutionEnvironment()->getMutableHardwareInfo()->workaroundTable.waModifyVFEStateAfterGPGPUPreemption = false;
+    PreemptionHelper::applyPreemptionWaCmdsBegin<FamilyType>(&cmdBuffer, device->getDevice());
+    EXPECT_EQ(0u, cmdBuffer.getUsed());
+}
+
+GEN8TEST_F(Gen8PreemptionTestsLinearStream, givenMidThreadPreemptionNoWaSetWhenProgrammingWaCmdsEndThenExpectNoCmd) {
+    device->setPreemptionMode(PreemptionMode::MidThread);
+    device->getExecutionEnvironment()->getMutableHardwareInfo()->workaroundTable.waModifyVFEStateAfterGPGPUPreemption = false;
+    PreemptionHelper::applyPreemptionWaCmdsEnd<FamilyType>(&cmdBuffer, device->getDevice());
+    EXPECT_EQ(0u, cmdBuffer.getUsed());
+}
+
+GEN8TEST_F(Gen8PreemptionTestsLinearStream, givenMidThreadPreemptionWaSetWhenProgrammingWaCmdsBeginThenExpectMmioCmd) {
+    using MI_LOAD_REGISTER_IMM = typename FamilyType::MI_LOAD_REGISTER_IMM;
+    device->setPreemptionMode(PreemptionMode::MidThread);
+    device->getExecutionEnvironment()->getMutableHardwareInfo()->workaroundTable.waModifyVFEStateAfterGPGPUPreemption = true;
+    PreemptionHelper::applyPreemptionWaCmdsBegin<FamilyType>(&cmdBuffer, device->getDevice());
+
+    cmdBufferParser.parseCommands<FamilyType>(cmdBuffer);
+    cmdBufferParser.findHardwareCommands<FamilyType>();
+    GenCmdList::iterator itMmioCmd = cmdBufferParser.lriList.begin();
+    ASSERT_NE(cmdBufferParser.lriList.end(), itMmioCmd);
+    MI_LOAD_REGISTER_IMM *mmioCmd = genCmdCast<MI_LOAD_REGISTER_IMM *>(*itMmioCmd);
+    EXPECT_EQ(0x2600u, mmioCmd->getRegisterOffset());
+    EXPECT_EQ(0xFFFFFFFFu, mmioCmd->getDataDword());
+}
+
+GEN8TEST_F(Gen8PreemptionTestsLinearStream, givenMidThreadPreemptionWaSetWhenProgrammingWaCmdsEndThenExpectMmioCmd) {
+    using MI_LOAD_REGISTER_IMM = typename FamilyType::MI_LOAD_REGISTER_IMM;
+    device->setPreemptionMode(PreemptionMode::MidThread);
+    device->getExecutionEnvironment()->getMutableHardwareInfo()->workaroundTable.waModifyVFEStateAfterGPGPUPreemption = true;
+    PreemptionHelper::applyPreemptionWaCmdsEnd<FamilyType>(&cmdBuffer, device->getDevice());
+
+    cmdBufferParser.parseCommands<FamilyType>(cmdBuffer);
+    cmdBufferParser.findHardwareCommands<FamilyType>();
+    GenCmdList::iterator itMmioCmd = cmdBufferParser.lriList.begin();
+    ASSERT_NE(cmdBufferParser.lriList.end(), itMmioCmd);
+    MI_LOAD_REGISTER_IMM *mmioCmd = genCmdCast<MI_LOAD_REGISTER_IMM *>(*itMmioCmd);
+    EXPECT_EQ(0x2600u, mmioCmd->getRegisterOffset());
+    EXPECT_EQ(0x00000000u, mmioCmd->getDataDword());
+}