PIX: Updates to shader debugging pass to support DXR hit groups/miss shaders (microsoft#5499)

Two main themes:
-Add instrumentation to the DXR hit group and miss shader types
-Return a bit more info to the caller (PIX) via output text

Author: jeffnn

lib/DxilPIXPasses/DxilAnnotateWithVirtualRegister.cpp

@@ -107,6 +107,16 @@ void DxilAnnotateWithVirtualRegister::applyOptions(llvm::PassOptions O) {
 
 char DxilAnnotateWithVirtualRegister::ID = 0;
 
+static llvm::StringRef
+PrintableSubsetOfMangledFunctionName(llvm::StringRef mangled) {
+  llvm::StringRef printableNameSubset = mangled;
+  if (mangled.size() > 2 && mangled[0] == '\1' && mangled[1] == '?') {
+    printableNameSubset =
+        llvm::StringRef(mangled.data() + 2, mangled.size() - 2);
+  }
+  return printableNameSubset;
+}
+
 bool DxilAnnotateWithVirtualRegister::runOnModule(llvm::Module &M) {
   Init(M);
   if (m_DM == nullptr) {
@@ -120,39 +130,37 @@ bool DxilAnnotateWithVirtualRegister::runOnModule(llvm::Module &M) {
   }
 
   std::uint32_t InstNum = m_StartInstruction;
-  std::map<llvm::StringRef, std::pair<int, int>> InstructionRangeByFunctionName;
 
   auto instrumentableFunctions = PIXPassHelpers::GetAllInstrumentableFunctions(*m_DM);
 
   for (auto * F : instrumentableFunctions) {
-    auto &EndInstruction = InstructionRangeByFunctionName[F->getName()];
-    EndInstruction.first = InstNum;
+    int InstructionRangeStart = InstNum;
+    int InstructionRangeEnd = InstNum;
     for (auto &block : F->getBasicBlockList()) {
       for (llvm::Instruction &I : block.getInstList()) {
         if (!llvm::isa<llvm::DbgDeclareInst>(&I)) {
           pix_dxil::PixDxilInstNum::AddMD(M.getContext(), &I, InstNum++);
-          EndInstruction.second = InstNum;
+          InstructionRangeEnd = InstNum;
         }
       }
     }
+    if (OSOverride != nullptr) {
+      auto shaderKind = PIXPassHelpers::GetFunctionShaderKind(*m_DM, F);
+      std::string FunctioNamePlusKind =
+          F->getName().str() + " " + hlsl::ShaderModel::GetKindName(shaderKind);
+      *OSOverride << "InstructionRange: ";
+      llvm::StringRef printableNameSubset =
+          PrintableSubsetOfMangledFunctionName(FunctioNamePlusKind);
+      *OSOverride << InstructionRangeStart << " " << InstructionRangeEnd << " "
+                  << printableNameSubset << "\n";
+    }
   }
 
   if (OSOverride != nullptr) {
     // Print a set of strings of the exemplary form "InstructionCount: <n> <fnName>"
+    if (m_DM->GetShaderModel()->GetKind() == hlsl::ShaderModel::Kind::Library)
+        *OSOverride << "\nIsLibrary\n";
     *OSOverride << "\nInstructionCount:" << InstNum << "\n";
-    for (auto const &fn : InstructionRangeByFunctionName) {
-      *OSOverride << "InstructionRange: ";
-      int skipOverLeadingUnprintableCharacters = 0;
-      if (fn.first.size() > 2 && fn.first[0] == '\1' && fn.first[1] == '?') {
-        skipOverLeadingUnprintableCharacters = 2;
-      }
-      *OSOverride << fn.second.first << " " << fn.second.second << " "
-                  << (fn.first.str().c_str() +
-                      skipOverLeadingUnprintableCharacters)
-                  << "\n";
-    }
-
-    *OSOverride << "\nBegin - dxil values to virtual register mapping\n";
   }
 
   for (auto * F : instrumentableFunctions) {
@@ -171,10 +179,6 @@ bool DxilAnnotateWithVirtualRegister::runOnModule(llvm::Module &M) {
     }
   }
 
-  if (OSOverride != nullptr) {
-    *OSOverride << "\nEnd - dxil values to virtual register mapping\n";
-  }
-
   m_DM = nullptr;
   return m_uVReg > 0;
 }

lib/DxilPIXPasses/DxilDebugInstrumentation.cpp

@@ -216,9 +216,9 @@ class DxilDebugInstrumentation : public ModulePass {
   uint64_t m_UAVSize = 1024 * 1024;
   struct PerFunctionValues
   {
-    CallInst *UAVHandle;
-    Constant *CounterOffset;
-    Value *InvocationId;
+    CallInst *UAVHandle = nullptr;
+    Constant *CounterOffset = nullptr;
+    Value *InvocationId = nullptr;
     // Together these two values allow branchless writing to the UAV. An
     // invocation of the shader is either of interest or not (e.g. it writes to
     // the pixel the user selected for debugging or it doesn't). If not of
@@ -227,11 +227,11 @@ class DxilDebugInstrumentation : public ModulePass {
     // will be written to the UAV at sequentially increasing offsets.
     // This value will either be one or zero (one if the invocation is of
     // interest, zero otherwise)
-    Value *OffsetMultiplicand;
+    Value *OffsetMultiplicand = nullptr;
     // This will either be zero (if the invocation is of interest) or
     // (UAVSize)-(SmallValue) if not.
-    Value *OffsetAddend;
-    Constant *OffsetMask;
+    Value *OffsetAddend = nullptr;
+    Constant *OffsetMask = nullptr;
     Value *SelectionCriterion = nullptr;
     Value *CurrentIndex = nullptr;
   };
@@ -341,10 +341,10 @@ DxilDebugInstrumentation::SystemValueIndices
   case DXIL::ShaderKind::Mesh:
   case DXIL::ShaderKind::Compute:
   case DXIL::ShaderKind::RayGeneration:
-  //case DXIL::ShaderKind::Intersection:
-  //case DXIL::ShaderKind::AnyHit:
-  //case DXIL::ShaderKind::ClosestHit:
-  //case DXIL::ShaderKind::Miss:
+  case DXIL::ShaderKind::Intersection:
+  case DXIL::ShaderKind::AnyHit:
+  case DXIL::ShaderKind::ClosestHit:
+  case DXIL::ShaderKind::Miss:
     // Dispatch* thread Id is not in the input signature
     break;
   case DXIL::ShaderKind::Vertex: {
@@ -448,9 +448,11 @@ Value *DxilDebugInstrumentation::addRaygenShaderProlog(BuilderContext &BC) {
   auto CompareToX = BC.Builder.CreateICmpEQ(
       RayX, BC.HlslOP->GetU32Const(m_Parameters.ComputeShader.ThreadIdX),
       "CompareToThreadIdX");
+
   auto CompareToY = BC.Builder.CreateICmpEQ(
       RayY, BC.HlslOP->GetU32Const(m_Parameters.ComputeShader.ThreadIdY),
       "CompareToThreadIdY");
+
   auto CompareToZ = BC.Builder.CreateICmpEQ(
       RayZ, BC.HlslOP->GetU32Const(m_Parameters.ComputeShader.ThreadIdZ),
       "CompareToThreadIdZ");
@@ -619,15 +621,15 @@ void DxilDebugInstrumentation::addInvocationSelectionProlog(
   Value *ParameterTestResult = nullptr;
   switch (shaderKind) {
   case DXIL::ShaderKind::RayGeneration:
+  case DXIL::ShaderKind::ClosestHit:
+  case DXIL::ShaderKind::Intersection:
+  case DXIL::ShaderKind::AnyHit:
+  case DXIL::ShaderKind::Miss:
     ParameterTestResult = addRaygenShaderProlog(BC);
     break;
   case DXIL::ShaderKind::Compute:
   case DXIL::ShaderKind::Amplification:
   case DXIL::ShaderKind::Mesh:
-  //case DXIL::ShaderKind::Intersection:
-  //case DXIL::ShaderKind::AnyHit:
-  //case DXIL::ShaderKind::ClosestHit:
-  //case DXIL::ShaderKind::Miss:
     ParameterTestResult = addDispatchedShaderProlog(BC);
     break;
   case DXIL::ShaderKind::Geometry:
@@ -806,7 +808,6 @@ void DxilDebugInstrumentation::addInvocationStartMarker(BuilderContext &BC) {
 
   marker.Header.Details.SizeDwords =
       DebugShaderModifierRecordPayloadSizeDwords(sizeof(marker));
-  ;
   marker.Header.Details.Flags = 0;
   marker.Header.Details.Type =
       DebugShaderModifierRecordTypeInvocationStartMarker;
@@ -980,16 +981,9 @@ bool DxilDebugInstrumentation::RunOnFunction(
   DxilModule &DM,
   llvm::Function * entryFunction) 
 {
-  DXIL::ShaderKind shaderKind = DXIL::ShaderKind::Invalid;
-  if (!DM.HasDxilFunctionProps(entryFunction)) {
-    auto ShaderModel = DM.GetShaderModel();
-    shaderKind = ShaderModel->GetKind();
-  } else {
-    hlsl::DxilFunctionProps const &props =
-        DM.GetDxilFunctionProps(entryFunction);
-    shaderKind = props.shaderKind;
-  }
-
+  DXIL::ShaderKind shaderKind =
+      PIXPassHelpers::GetFunctionShaderKind(DM, entryFunction);
+  
   switch (shaderKind) {
   case DXIL::ShaderKind::Amplification:
   case DXIL::ShaderKind::Mesh:
@@ -1000,12 +994,11 @@ bool DxilDebugInstrumentation::RunOnFunction(
   case DXIL::ShaderKind::RayGeneration:
   case DXIL::ShaderKind::Hull:
   case DXIL::ShaderKind::Domain:
-    break;
-    //todo:
   case DXIL::ShaderKind::Intersection:
   case DXIL::ShaderKind::AnyHit:
   case DXIL::ShaderKind::ClosestHit:
   case DXIL::ShaderKind::Miss:
+    break;
   default:
     return false;
   }
@@ -1042,8 +1035,24 @@ bool DxilDebugInstrumentation::RunOnFunction(
 
   auto &values = m_FunctionToValues[BC.Builder.GetInsertBlock()->getParent()];
 
+  // PIX binds two UAVs when running this instrumentation: one for raygen shaders
+  // and another for the hitgroups and miss shaders. Since PIX invokes this pass
+  // at the library level, which may contain examples of both types, PIX can't really
+  // specify which UAV index to use per-shader. This pass therefore just has to know this:
+  constexpr unsigned int RayGenUAVRegister = 0;
+  constexpr unsigned int HitGroupAndMissUAVRegister = 1;
+  unsigned int UAVRegisterId = RayGenUAVRegister;
+    switch (shaderKind) {
+  case DXIL::ShaderKind::ClosestHit:
+  case DXIL::ShaderKind::Intersection:
+  case DXIL::ShaderKind::AnyHit:
+  case DXIL::ShaderKind::Miss:
+    UAVRegisterId = HitGroupAndMissUAVRegister;
+    break;
+  }
+
   values.UAVHandle = PIXPassHelpers::CreateUAV(
-      DM, Builder, 0,
+      DM, Builder, UAVRegisterId,
       "PIX_DebugUAV_Handle");
   values.CounterOffset = BC.HlslOP->GetU32Const(UAVDumpingGroundOffset() + CounterOffsetBeyondUsefulData);
 

lib/DxilPIXPasses/PixPassHelpers.cpp

@@ -337,6 +337,20 @@ GetAllInstrumentableFunctions(hlsl::DxilModule &DM) {
   return ret;
 }
 
+hlsl::DXIL::ShaderKind GetFunctionShaderKind(hlsl::DxilModule &DM,
+                                       llvm::Function *fn) {
+  hlsl::DXIL::ShaderKind shaderKind = hlsl::DXIL::ShaderKind::Invalid;
+  if (!DM.HasDxilFunctionProps(fn)) {
+    auto ShaderModel = DM.GetShaderModel();
+    shaderKind = ShaderModel->GetKind();
+  } else {
+    hlsl::DxilFunctionProps const &props =
+        DM.GetDxilFunctionProps(fn);
+    shaderKind = props.shaderKind;
+  }
+  return shaderKind;
+}
+
 std::vector<llvm::BasicBlock*> GetAllBlocks(hlsl::DxilModule& DM) {
     std::vector<llvm::BasicBlock*> ret;
     auto entryPoints = DM.GetExportedFunctions();

lib/DxilPIXPasses/PixPassHelpers.h

@@ -30,6 +30,8 @@ namespace PIXPassHelpers
     llvm::Function* GetEntryFunction(hlsl::DxilModule& DM);
     std::vector<llvm::BasicBlock*> GetAllBlocks(hlsl::DxilModule& DM);
     std::vector<llvm::Function*> GetAllInstrumentableFunctions(hlsl::DxilModule& DM);
+    hlsl::DXIL::ShaderKind GetFunctionShaderKind(hlsl::DxilModule &DM,
+                                           llvm::Function *fn);
 #ifdef PIX_DEBUG_DUMP_HELPER
     void Log(const char* format, ...);
     void LogPartialLine(const char* format, ...);

tools/clang/test/HLSLFileCheck/pix/DebugHitGroupAndMiss.hlsl

@@ -0,0 +1,110 @@
+// RUN: %dxc -T lib_6_3 %s | %opt -S -hlsl-dxil-debug-instrumentation,parameter0=10,parameter1=20,parameter2=30 | %FileCheck %s
+
+
+// Just check that the selection prolog was added to each shader:
+
+// CHECK: call i32 @dx.op.dispatchRaysIndex.i32(i32 145, i8 0)
+// CHECK: call i32 @dx.op.dispatchRaysIndex.i32(i32 145, i8 1)
+// CHECK: call i32 @dx.op.dispatchRaysIndex.i32(i32 145, i8 2)
+// CHECK: call i32 @dx.op.dispatchRaysIndex.i32(i32 145, i8 0)
+
+// There must be three compares:
+// CHECK: = icmp eq i32
+// CHECK: = icmp eq i32
+// CHECK: = icmp eq i32
+
+// Two ANDs of these bools:
+// CHECK: and i1
+// CHECK: and i1
+
+//-----------------------------------------------------------------------------
+struct cb_push_test_raytracing_t
+{
+    uint m_rt_index;
+    uint m_raytracing_acc_struct_index;
+    uint m_per_view_index;
+};
+
+//-----------------------------------------------------------------------------
+struct cb_view_t
+{
+    float4x4 m_view;
+    float4x4 m_proj;
+    float4x4 m_inv_view_proj;
+    float3   m_camera_pos;
+};
+
+RaytracingAccelerationStructure g_bindless_raytracing_acc_struct[] : register(t0, space200);
+ConstantBuffer<cb_push_test_raytracing_t> g_push_constants : register(b0, space999);
+ConstantBuffer<cb_view_t> g_view[] : register(b0, space100);
+RWTexture2D<float4> g_output[] : register(u0, space100);
+
+//-----------------------------------------------------------------------------
+// Payloads
+//-----------------------------------------------------------------------------
+
+struct ray_payload_t
+{
+    float3 m_color;
+};
+
+//-----------------------------------------------------------------------------
+// Ray generation
+//-----------------------------------------------------------------------------
+
+[shader("raygeneration")]
+void raygen_shader()
+{
+    if( g_push_constants.m_rt_index != 0xFFFFFFFF &&
+        g_push_constants.m_raytracing_acc_struct_index != 0xFFFFFFFF &&
+        g_push_constants.m_per_view_index != 0xFFFFFFFF)
+    {
+        RaytracingAccelerationStructure l_acc_struct = g_bindless_raytracing_acc_struct[g_push_constants.m_raytracing_acc_struct_index];
+        RWTexture2D<float4> l_output = g_output[g_push_constants.m_rt_index];
+        cb_view_t l_per_view = g_view[g_push_constants.m_per_view_index];
+
+        uint3 l_launch_index = DispatchRaysIndex();
+
+        // Ray dsc
+        RayDesc ray;
+        ray.Origin      = l_per_view.m_camera_pos; // THIS LINE CAUSES AN ACCESS VIOLATION. Replacing l_per_view.m_camera_pos with a constant makes it compile fine
+        //ray.Origin      = 0;
+        ray.Direction   = float3(0, 1.0, 0);
+        ray.TMin        = 0;
+        ray.TMax        = 100000;
+
+        ray_payload_t l_payload;
+        TraceRay(l_acc_struct, 0 /*rayFlags*/, 0xFF, 0 /* ray index*/, 0, 0, ray, l_payload);
+        l_output[l_launch_index.xy] = float4(l_payload.m_color, 1.0);
+    }
+}
+
+//-----------------------------------------------------------------------------
+// Primary rays
+//-----------------------------------------------------------------------------
+
+[shader("miss")]
+void miss_shader(inout ray_payload_t l_payload)
+{
+    l_payload.m_color = float3(0.4, 0.0, 0.0);
+}
+
+[shader("closesthit")]
+void closest_hit_shader(inout ray_payload_t l_payload, in BuiltInTriangleIntersectionAttributes l_attribs)
+{
+    l_payload.m_color = 1.0f;
+}
+
+
+struct MyAttributes {
+  float2 bary;
+  uint id;
+};
+
+[shader("intersection")]
+void intersection1()
+{
+  float hitT = RayTCurrent();
+  MyAttributes attr = (MyAttributes)0;
+  bool bReported = ReportHit(hitT, 0, attr);
+}