refactor preemption check into function

intel · Dec 12, 2024 · b46d1c2 · b46d1c2
1 parent d047fc5
commit b46d1c2
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/IGC/Compiler/CISACodeGen/EmitVISAPass.cpp b/IGC/Compiler/CISACodeGen/EmitVISAPass.cpp
@@ -1465,10 +1465,8 @@ bool EmitPass::runOnFunction(llvm::Function& F)
 
     if ((m_currShader->GetShaderType() == ShaderType::COMPUTE_SHADER ||
         m_currShader->GetShaderType() == ShaderType::OPENCL_SHADER) &&
-        m_currShader->m_Platform->supportDisableMidThreadPreemptionSwitch() &&
-        IGC_IS_FLAG_ENABLED(EnableDisableMidThreadPreemptionOpt) &&
-        (m_currShader->GetContext()->m_instrTypes.numLoopInsts == 0) &&
-        (m_currShader->ProgramOutput()->m_InstructionCount < IGC_GET_FLAG_VALUE(MidThreadPreemptionDisableThreshold)))
+        m_currShader->m_State.shouldDisablePreemption(
+            m_currShader->ProgramOutput()->m_InstructionCount))
     {
 
         {

diff --git a/IGC/Compiler/CISACodeGen/GenericShaderState.hpp b/IGC/Compiler/CISACodeGen/GenericShaderState.hpp
@@ -107,6 +107,14 @@ namespace IGC
         CodeGenContext& Ctx;
 
         CodeGenContext& GetContext() const { return Ctx; }
+
+        bool shouldDisablePreemption(unsigned NumInst) const
+        {
+            return (Ctx.platform.supportDisableMidThreadPreemptionSwitch() &&
+                IGC_IS_FLAG_ENABLED(EnableDisableMidThreadPreemptionOpt) &&
+                (Ctx.m_instrTypes.numLoopInsts == 0) &&
+                (NumInst < IGC_GET_FLAG_VALUE(MidThreadPreemptionDisableThreshold)));
+        }
     private:
         bool m_HasSample = false;
         int m_BarrierNumber = 0;