Stack call related changes

Stack call related changes

Author: jaladreips

IGC/AdaptorCommon/RayTracing/RTArgs.cpp

@@ -31,6 +31,7 @@ void ArgQuery::init(CallableShaderTypeMD FuncType, const FunctionMetaData &FMD)
   uint32_t Idx = 0;
   switch (FuncType) {
   case RayGen:
+  case Prologue:
   case CallStackHandler:
     break;
   case Miss:

IGC/Compiler/CISACodeGen/EmitVISAPass.cpp

@@ -11390,6 +11390,21 @@ void EmitPass::emitStackCall(llvm::CallInst *inst) {
 
   bool isIndirectFCall = !F || F->hasFnAttribute("referenced-indirectly");
   bool isInvokeSIMDTarget = F && F->hasFnAttribute("invoke_simd_target");
+
+  if (inst->doesNotReturn()) {
+
+    IGC_ASSERT_MESSAGE(inst->getType()->isVoidTy(), "Call that does not return can't return a value");
+    IGC_ASSERT_MESSAGE(IGCLLVM::getNumArgOperands(inst) == 0, "Arguments for non-returning call are not implemented");
+
+    auto *funcAddr = GetSymbol(IGCLLVM::getCalledValue(inst));
+      funcAddr = TruncatePointer(funcAddr);
+    IGC_ASSERT_MESSAGE(funcAddr->IsUniform(), "Function address must be uniform for non-returning stack call");
+    m_encoder->IndirectStackCall(nullptr, funcAddr, 0, 0);
+    m_encoder->Push();
+
+    return;
+  }
+
   CVariable *ArgBlkVar = m_currShader->GetARGV();
   uint32_t offsetA = 0; // visa argument offset
   uint32_t offsetS = 0; // visa stack offset

IGC/Compiler/CISACodeGen/GenCodeGenModule.cpp

@@ -746,7 +746,7 @@ void GenXFunctionGroupAnalysis::setGroupAttributes() {
           } else if (!calledF || (calledF->isDeclaration() && calledF->hasFnAttribute("referenced-indirectly"))) {
             // This is the true indirect call case, where either the callee's address is taken, or it belongs
             // to an external module. We do not know the callgraph in this case.
-            hasStackCall = true;
+            hasStackCall = !call->doesNotReturn();
             FG->m_hasIndirectCall = true;
             FG->m_hasPartialCallGraph = true;
           } else if (calledF && calledF->hasFnAttribute("referenced-indirectly")) {

IGC/Compiler/CISACodeGen/Platform.hpp

@@ -957,6 +957,8 @@ class CPlatform {
     }
   }
 
+  uint32_t getMinNumGRF() const { return supportsVRT() ? 32 : 128; }
+
   uint32_t getInlineDataSize() const {
     if (!supportInlineData())
       return 0;

IGC/Compiler/CISACodeGen/RayTracingShaderLowering.cpp

@@ -212,9 +212,14 @@ bool RayTracingShaderLowering::runOnModule(Module &M) {
       return EntryPreemptionVal;
     };
     const bool isFunc = !isEntryFunc(CGCtx->getMetaDataUtils(), &F);
+
     if (ForcePreemptionDisable) {
-      RTB.SetInsertPoint(&F.getEntryBlock().front());
-      RTB.CreatePreemptionDisableIntrinsic();
+      auto &rtInfo = CGCtx->getModuleMetaData()->FuncMD[&F].rtInfo;
+      if (rtInfo.callableShaderType != Prologue) {
+        // For raytracing shaders, disable preemption at entry
+        RTB.SetInsertPoint(&F.getEntryBlock().front());
+        RTB.CreatePreemptionDisableIntrinsic();
+      }
     }
 
     SmallVector<Instruction *, 8> DeadInsts;

IGC/Compiler/CodeGenPublic.h

@@ -524,6 +524,8 @@ struct SRayTracingShadersGroup {
   BindlessShaderVec m_CallableShaders;
   // Continuation shaders
   BindlessShaderVec m_Continuations;
+
+  std::optional<SBindlessProgram> prologueKernel;
 };
 
 struct SRayTracingPipelineConfig {

IGC/common/MDFrameWork.cpp

@@ -446,7 +446,10 @@ void IGC::serialize(const IGC::ModuleMetaData &moduleMD, Module *module) {
 }
 
 bool IGC::isBindless(const IGC::FunctionMetaData &funcMD) {
-  return funcMD.rtInfo.callableShaderType != RayGen || funcMD.rtInfo.isContinuation;
+  if (funcMD.rtInfo.isContinuation)
+    return true;
+
+  return funcMD.rtInfo.callableShaderType != RayGen && funcMD.rtInfo.callableShaderType != Prologue;
 }
 
 bool IGC::isContinuation(const IGC::FunctionMetaData &funcMD) { return funcMD.rtInfo.isContinuation; }
@@ -455,6 +458,8 @@ bool IGC::isCallStackHandler(const IGC::FunctionMetaData &funcMD) {
   return funcMD.rtInfo.callableShaderType == IGC::CallStackHandler;
 }
 
+bool IGC::isPrologue(const IGC::FunctionMetaData &funcMD) { return funcMD.rtInfo.callableShaderType == IGC::Prologue; }
+
 
 int IGC::extractAnnotatedNumThreads(const IGC::FunctionMetaData &funcMD) {
   static constexpr const char *searchedString = "num-thread-per-eu";

IGC/common/MDFrameWork.h

@@ -878,6 +878,7 @@ enum class ShaderTypeMD
     bool isBindless(const IGC::FunctionMetaData &funcMD);
     bool isContinuation(const IGC::FunctionMetaData& funcMD);
     bool isCallStackHandler(const IGC::FunctionMetaData &funcMD);
+    bool isPrologue(const IGC::FunctionMetaData &funcMD);
 
 
     // User annotations query functions

IGC/common/RaytracingShaderTypes.h

@@ -21,5 +21,7 @@ enum CallableShaderTypeMD : uint32_t
     // by any RTUnit or BTD machinery.  It is useful, though, to have all
     // shader types available for reference in this enum.
     RayGen,
+    Prologue,
+    NumberOfAllShaderTypes
 };