For rayquery check/release, enable preemption after all releases are executed

For rayquery check/release, enable preemption after all releases are executed

Author: jaladreips

IGC/AdaptorCommon/RayTracing/RTBuilder.cpp

@@ -185,15 +185,14 @@ RayQueryReleaseIntrinsic* RTBuilder::CreateRayQueryReleaseIntrinsic(Value* predi
     return cast<RayQueryReleaseIntrinsic>(rayQueryRelease);
 }
 
-PreemptionDisableIntrinsic* RTBuilder::CreatePreemptionDisableIntrinsic(Value* Flag)
+PreemptionDisableIntrinsic* RTBuilder::CreatePreemptionDisableIntrinsic()
 {
     Module* M = this->GetInsertBlock()->getModule();
 
     auto* GII = CreateCall(
         GenISAIntrinsic::getDeclaration(
             M,
-            GenISAIntrinsic::GenISA_PreemptionDisable),
-        Flag ? Flag : getTrue());
+            GenISAIntrinsic::GenISA_PreemptionDisable));
 
     return cast<PreemptionDisableIntrinsic>(GII);
 }

IGC/AdaptorCommon/RayTracing/RTBuilder.h

@@ -259,7 +259,7 @@ class RTBuilder : public IGCIRBuilder<>
     RayQueryCheckIntrinsic* CreateRayQueryCheckIntrinsic(Value* predicate = nullptr);
     RayQueryReleaseIntrinsic* CreateRayQueryReleaseIntrinsic(Value* predicate = nullptr);
 
-    PreemptionDisableIntrinsic* CreatePreemptionDisableIntrinsic(Value* Flag = nullptr);
+    PreemptionDisableIntrinsic* CreatePreemptionDisableIntrinsic();
     PreemptionEnableIntrinsic* CreatePreemptionEnableIntrinsic(Value* Flag = nullptr);
 
     SyncStackPointerVal* getSyncStackPointer();

IGC/Compiler/CISACodeGen/EmitVISAPass.cpp

@@ -23612,19 +23612,6 @@ void EmitPass::emitPreemptionDisable(PreemptionDisableIntrinsic* PDI)
         ~getEncoderPreemptionMode(PREEMPTION_ENABLED),
         ISA_TYPE_UD);
 
-    CVariable* Flag = nullptr;
-    if (auto* CI = dyn_cast<ConstantInt>(PDI->getFlag()))
-    {
-        if (CI->isZero())
-            return;
-    }
-    else
-    {
-        Flag = GetSymbol(PDI->getFlag());
-    }
-
-    m_encoder->SetPredicate(Flag);
-
     m_encoder->And(m_currShader->GetCR0(), m_currShader->GetCR0(), Mask);
     m_encoder->Push();
 }

IGC/Compiler/CISACodeGen/RayTracingShaderLowering.cpp

@@ -40,6 +40,8 @@ class RayTracingShaderLowering : public ModulePass
     {
         AU.setPreservesCFG();
         AU.addRequired<CodeGenContextWrapper>();
+        AU.addRequired<DominatorTreeWrapperPass>();
+        AU.addRequired<PostDominatorTreeWrapperPass>();
     }
 
     StringRef getPassName() const override
@@ -192,6 +194,12 @@ bool RayTracingShaderLowering::runOnModule(Module& M)
 
         SmallVector<Instruction*, 8> DeadInsts;
 
+        BasicBlock* disablePreemptionBB = nullptr;
+        BasicBlock* enablePreemptionBB = nullptr;
+
+        DominatorTree& DT = getAnalysis<DominatorTreeWrapperPass>(F).getDomTree();
+        PostDominatorTree& PDT = getAnalysis<PostDominatorTreeWrapperPass>(F).getPostDomTree();
+
         for (auto II = inst_begin(&F), IE = inst_end(&F); II != IE; /* empty */)
         {
             Instruction& I = *II++;
@@ -245,25 +253,44 @@ bool RayTracingShaderLowering::runOnModule(Module& M)
             case GenISAIntrinsic::GenISA_RayQueryCheck:
                 if (!ForcePreemptionDisable)
                 {
-                    RTB.SetInsertPoint(GII);
-                    RTB.CreatePreemptionDisableIntrinsic(I.getOperand(0));
+                    disablePreemptionBB =
+                        disablePreemptionBB ?
+                        DT.findNearestCommonDominator(disablePreemptionBB, GII->getParent()) :
+                        GII->getParent();
                 }
                 break;
             case GenISAIntrinsic::GenISA_RayQueryRelease:
                 if (!ForcePreemptionDisable)
                 {
-                    Value* Flag = isFunc ? checkPreemption(F) : nullptr;
-                    Flag = Flag ? RTB.CreateAnd(Flag, I.getOperand(0)) : I.getOperand(0);
-
-                    RTB.SetInsertPoint(GII->getNextNode());
-                    RTB.CreatePreemptionEnableIntrinsic(Flag);
+                    enablePreemptionBB =
+                        enablePreemptionBB ?
+                        PDT.findNearestCommonDominator(enablePreemptionBB, GII->getParent()) :
+                        GII->getParent();
                 }
                 break;
             default:
                 break;
             }
         }
 
+        if (disablePreemptionBB)
+        {
+            IGC_ASSERT(enablePreemptionBB);
+
+            RTB.SetInsertPoint(disablePreemptionBB->getFirstNonPHI());
+            RTB.CreatePreemptionDisableIntrinsic();
+        }
+
+        if (enablePreemptionBB)
+        {
+            IGC_ASSERT(disablePreemptionBB);
+
+            Value* Flag = isFunc ? checkPreemption(F) : nullptr;
+
+            RTB.SetInsertPoint(enablePreemptionBB->getTerminator());
+            RTB.CreatePreemptionEnableIntrinsic(Flag);
+        }
+
         for (auto* I : DeadInsts)
             I->eraseFromParent();
     }
@@ -279,6 +306,8 @@ namespace IGC {
 #define PASS_ANALYSIS false
 IGC_INITIALIZE_PASS_BEGIN(RayTracingShaderLowering, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
 IGC_INITIALIZE_PASS_DEPENDENCY(CodeGenContextWrapper)
+IGC_INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+IGC_INITIALIZE_PASS_DEPENDENCY(PostDominatorTreeWrapperPass)
 IGC_INITIALIZE_PASS_END(RayTracingShaderLowering, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
 
 ModulePass* CreateRayTracingShaderLowering()

IGC/GenISAIntrinsics/Intrinsic_definitions.py

@@ -3208,7 +3208,7 @@
 ####################################################################################################
 "GenISA_PreemptionDisable":["Disable Preemption",
     [("void",                          ""),
-    [("bool",                          "predication flag")],
+    [],
     "None"]],
 ####################################################################################################
 "GenISA_RayQueryCheck": ["Raytracing: RayQueryCheck",