[PIX] Select callable and node shader debug invocations

Instrument callable shaders, selecting invocations by dispatchRaysIndex, and select node shader invocations according to launch type, reporting whether the selection is exact or approximate.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 40dc9de3-617e-4caf-ab0d-fba0a033ed93

Author: damyanp

lib/DxilPIXPasses/DxilDebugInstrumentation.cpp

@@ -336,7 +336,10 @@ class DxilDebugInstrumentation : public ModulePass {
                                     DXIL::ShaderKind shaderKind);
   Value *addPixelShaderProlog(BuilderContext &BC, SystemValueIndices SVIndices);
   Value *addGeometryShaderProlog(BuilderContext &BC);
+  Value *addThreadIdComparisonProlog(BuilderContext &BC,
+                                     DXIL::OpCode threadIdOpCode);
   Value *addDispatchedShaderProlog(BuilderContext &BC);
+  Value *addNodeShaderProlog(BuilderContext &BC);
   Value *addRaygenShaderProlog(BuilderContext &BC);
   Value *addVertexShaderProlog(BuilderContext &BC,
                                SystemValueIndices SVIndices);
@@ -392,6 +395,7 @@ static bool IsInstrumentableShaderKind(DXIL::ShaderKind shaderKind) {
   case DXIL::ShaderKind::AnyHit:
   case DXIL::ShaderKind::ClosestHit:
   case DXIL::ShaderKind::Miss:
+  case DXIL::ShaderKind::Callable:
   case DXIL::ShaderKind::Node:
     return true;
   default:
@@ -512,6 +516,7 @@ DxilDebugInstrumentation::addRequiredSystemValues(BuilderContext &BC,
   case DXIL::ShaderKind::AnyHit:
   case DXIL::ShaderKind::ClosestHit:
   case DXIL::ShaderKind::Miss:
+  case DXIL::ShaderKind::Callable:
   case DXIL::ShaderKind::Node:
     // Dispatch* thread Id is not in the input signature
     break;
@@ -563,14 +568,15 @@ DxilDebugInstrumentation::addRequiredSystemValues(BuilderContext &BC,
   return SVIndices;
 }
 
-Value *DxilDebugInstrumentation::addDispatchedShaderProlog(BuilderContext &BC) {
+Value *DxilDebugInstrumentation::addThreadIdComparisonProlog(
+    BuilderContext &BC, DXIL::OpCode threadIdOpCode) {
   Constant *Zero32Arg = BC.HlslOP->GetU32Const(0);
   Constant *One32Arg = BC.HlslOP->GetU32Const(1);
   Constant *Two32Arg = BC.HlslOP->GetU32Const(2);
 
   auto ThreadIdFunc =
-      BC.HlslOP->GetOpFunc(DXIL::OpCode::ThreadId, Type::getInt32Ty(BC.Ctx));
-  Constant *Opcode = BC.HlslOP->GetU32Const((unsigned)DXIL::OpCode::ThreadId);
+      BC.HlslOP->GetOpFunc(threadIdOpCode, Type::getInt32Ty(BC.Ctx));
+  Constant *Opcode = BC.HlslOP->GetU32Const((unsigned)threadIdOpCode);
   auto ThreadIdX =
       BC.Builder.CreateCall(ThreadIdFunc, {Opcode, Zero32Arg}, "ThreadIdX");
   auto ThreadIdY =
@@ -598,6 +604,62 @@ Value *DxilDebugInstrumentation::addDispatchedShaderProlog(BuilderContext &BC) {
   return CompareAll;
 }
 
+Value *DxilDebugInstrumentation::addDispatchedShaderProlog(BuilderContext &BC) {
+  return addThreadIdComparisonProlog(BC, DXIL::OpCode::ThreadId);
+}
+
+// The identity of a node shader invocation depends on how the node launches,
+// and the three launch types offer different system values. The validator
+// allows SV_DispatchThreadID and SV_GroupID only to a broadcasting node, and
+// SV_GroupThreadID and SV_GroupIndex only to a broadcasting or a coalescing
+// node. A thread launch node gets none of them.
+//
+//  - A broadcasting node runs over a grid, so SV_DispatchThreadID names an
+//    invocation exactly as it does for a compute shader.
+//  - A coalescing node has only the position of a thread within its group.
+//    That narrows the selection to the group size instead of to a single
+//    invocation.
+//  - A thread launch node has no thread identity at all, so every invocation
+//    stays of interest.
+//
+// The pass report names which of the three applies, so the caller presents the
+// selection as exact or approximate.
+Value *DxilDebugInstrumentation::addNodeShaderProlog(BuilderContext &BC) {
+  llvm::Function *function = BC.Builder.GetInsertBlock()->getParent();
+
+  if (!BC.DM.HasDxilFunctionProps(function)) {
+    *OSOverride << "NodeInvocationSelection:None\n";
+    return BC.HlslOP->GetI1Const(1);
+  }
+
+  hlsl::DxilFunctionProps const &props = BC.DM.GetDxilFunctionProps(function);
+
+  switch (props.Node.LaunchType) {
+  case DXIL::NodeLaunchType::Broadcasting:
+    *OSOverride << "NodeInvocationSelection:DispatchThreadId\n";
+    return addThreadIdComparisonProlog(BC, DXIL::OpCode::ThreadId);
+
+  case DXIL::NodeLaunchType::Coalescing:
+    // The requested coordinates only mean something as a position within the
+    // group. A coordinate outside the declared group size matches no
+    // invocation, and a criterion that nothing satisfies leaves the debugger
+    // with no trace at all. Select every invocation in that case instead.
+    if (m_Parameters.ComputeShader.ThreadIdX < props.numThreads[0] &&
+        m_Parameters.ComputeShader.ThreadIdY < props.numThreads[1] &&
+        m_Parameters.ComputeShader.ThreadIdZ < props.numThreads[2]) {
+      *OSOverride << "NodeInvocationSelection:GroupThreadId\n";
+      return addThreadIdComparisonProlog(BC, DXIL::OpCode::ThreadIdInGroup);
+    }
+    *OSOverride << "NodeInvocationSelection:None\n";
+    return BC.HlslOP->GetI1Const(1);
+
+  case DXIL::NodeLaunchType::Thread:
+  default:
+    *OSOverride << "NodeInvocationSelection:None\n";
+    return BC.HlslOP->GetI1Const(1);
+  }
+}
+
 Value *DxilDebugInstrumentation::addRaygenShaderProlog(BuilderContext &BC) {
   auto DispatchRaysIndexOpFunc = BC.HlslOP->GetOpFunc(
       DXIL::OpCode::DispatchRaysIndex, Type::getInt32Ty(BC.Ctx));
@@ -819,10 +881,16 @@ void DxilDebugInstrumentation::addInvocationSelectionProlog(
   case DXIL::ShaderKind::Intersection:
   case DXIL::ShaderKind::AnyHit:
   case DXIL::ShaderKind::Miss:
+  case DXIL::ShaderKind::Callable:
+    // A callable shader has neither a ray nor a thread of its own, but it runs
+    // as part of the dispatch that calls it, so DispatchRaysIndex names the
+    // invocation that leads to it. PIX selects a raygen invocation on that same
+    // value, so a callable that the selected raygen thread invokes is selected
+    // here too.
     ParameterTestResult = addRaygenShaderProlog(BC);
     break;
   case DXIL::ShaderKind::Node:
-    ParameterTestResult = BC.HlslOP->GetI1Const(1);
+    ParameterTestResult = addNodeShaderProlog(BC);
     break;
   case DXIL::ShaderKind::Compute:
   case DXIL::ShaderKind::Amplification:

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

@@ -0,0 +1,38 @@
+// RUN: %dxc -T lib_6_3 %s | %opt -S -dxil-annotate-with-virtual-regs -hlsl-dxil-debug-instrumentation,parameter0=1,parameter1=2,parameter2=3 -hlsl-dxilemit | %FileCheck %s
+
+// A callable shader is a valid CallShader() target, so PIX must be able to step
+// into it. The pass numbers its instructions and advertises them to PIX, so it
+// must instrument it as well.
+//
+// A callable shader has no ray and no thread of its own, but
+// dx.op.dispatchRaysIndex is legal in it, and reports the index of the ray
+// generation invocation that is responsible for the call. PIX selects a raygen,
+// any-hit, closest-hit or miss invocation on that same identity, so a callable
+// shader uses it too.
+
+// The callable shader is the only shader in the module, so a Block# line means
+// that it is instrumented.
+// CHECK: Block#
+
+// CHECK: %RayX = call i32 @dx.op.dispatchRaysIndex.i32(i32 145, i8 0)
+// CHECK: %RayY = call i32 @dx.op.dispatchRaysIndex.i32(i32 145, i8 1)
+// CHECK: %RayZ = call i32 @dx.op.dispatchRaysIndex.i32(i32 145, i8 2)
+// CHECK: %CompareToThreadIdX = icmp eq i32 %RayX, 1
+// CHECK: %CompareToThreadIdY = icmp eq i32 %RayY, 2
+// CHECK: %CompareToThreadIdZ = icmp eq i32 %RayZ, 3
+// CHECK: %CompareAll = and i1 %CompareXAndY, %CompareToThreadIdZ
+// CHECK: br i1 %CompareAll, label %PIXInterestingBlock, label %PIXNonInterestingBlock
+
+RWStructuredBuffer<float> Output : register(u0);
+
+struct CallableParameters
+{
+  float value;
+};
+
+[shader("callable")]
+void MyCallable(inout CallableParameters parameters)
+{
+  parameters.value = parameters.value * 2.f;
+  Output[0] = parameters.value;
+}

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

@@ -0,0 +1,35 @@
+// RUN: %dxc -T lib_6_8 %s | %opt -S -dxil-annotate-with-virtual-regs -hlsl-dxil-debug-instrumentation,parameter0=10,parameter1=20,parameter2=30 -hlsl-dxilemit | %FileCheck %s
+
+// A broadcasting node runs over a grid, so SV_DispatchThreadID names an
+// invocation exactly as it does for a compute shader, and the debugger selects
+// the one that the user asks for. A constant-true criterion instead makes every
+// invocation believe that it is the selected one, and the debugger shows
+// whichever one reaches the UAV first.
+
+// CHECK: NodeInvocationSelection:DispatchThreadId
+
+// CHECK: %ThreadIdX = call i32 @dx.op.threadId.i32(i32 93, i32 0)
+// CHECK: %ThreadIdY = call i32 @dx.op.threadId.i32(i32 93, i32 1)
+// CHECK: %ThreadIdZ = call i32 @dx.op.threadId.i32(i32 93, i32 2)
+// CHECK: %CompareToThreadIdX = icmp eq i32 %ThreadIdX, 10
+// CHECK: %CompareToThreadIdY = icmp eq i32 %ThreadIdY, 20
+// CHECK: %CompareToThreadIdZ = icmp eq i32 %ThreadIdZ, 30
+// CHECK: %CompareXAndY = and i1 %CompareToThreadIdX, %CompareToThreadIdY
+// CHECK: %CompareAll = and i1 %CompareXAndY, %CompareToThreadIdZ
+// CHECK: br i1 %CompareAll, label %PIXInterestingBlock, label %PIXNonInterestingBlock
+
+RWStructuredBuffer<uint> Output : register(u0);
+
+struct Record
+{
+  uint value;
+};
+
+[Shader("node")]
+[NodeLaunch("broadcasting")]
+[NodeDispatchGrid(2, 1, 1)]
+[NumThreads(4, 2, 1)]
+void BroadcastingNode(DispatchNodeInputRecord<Record> input)
+{
+  Output[0] = input.Get().value;
+}

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

@@ -0,0 +1,36 @@
+// RUN: %dxc -T lib_6_8 %s | %opt -S -dxil-annotate-with-virtual-regs -hlsl-dxil-debug-instrumentation,parameter0=3,parameter1=1,parameter2=0 -hlsl-dxilemit | %FileCheck %s
+
+// A coalescing node has no dispatch grid, so dx.op.threadId is not legal for it.
+// ValidateDxilOperationCallInProfile in DxilValidation.cpp permits ThreadId and
+// GroupId for a broadcasting launch only. The thread group ID is legal, so the
+// debugger discriminates invocations within a group by SV_GroupThreadID.
+
+// CHECK: NodeInvocationSelection:GroupThreadId
+
+// CHECK: %ThreadIdX = call i32 @dx.op.threadIdInGroup.i32(i32 95, i32 0)
+// CHECK: %ThreadIdY = call i32 @dx.op.threadIdInGroup.i32(i32 95, i32 1)
+// CHECK: %ThreadIdZ = call i32 @dx.op.threadIdInGroup.i32(i32 95, i32 2)
+// CHECK: %CompareToThreadIdX = icmp eq i32 %ThreadIdX, 3
+// CHECK: %CompareToThreadIdY = icmp eq i32 %ThreadIdY, 1
+// CHECK: %CompareToThreadIdZ = icmp eq i32 %ThreadIdZ, 0
+// CHECK: %CompareAll = and i1 %CompareXAndY, %CompareToThreadIdZ
+// CHECK: br i1 %CompareAll, label %PIXInterestingBlock, label %PIXNonInterestingBlock
+
+// The requested thread must lie inside the declared thread group, or the pass
+// discriminates nothing. The parameters above lie inside [NumThreads(4, 2, 1)].
+// CHECK-NOT: NodeInvocationSelection:None
+
+RWStructuredBuffer<uint> Output : register(u0);
+
+struct Record
+{
+  uint value;
+};
+
+[Shader("node")]
+[NodeLaunch("coalescing")]
+[NumThreads(4, 2, 1)]
+void CoalescingNode(GroupNodeInputRecords<Record> input)
+{
+  Output[0] = input.Get(0).value;
+}

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

@@ -0,0 +1,30 @@
+// RUN: %dxc -T lib_6_8 %s | %opt -S -dxil-annotate-with-virtual-regs -hlsl-dxil-debug-instrumentation,parameter0=1,parameter1=2,parameter2=3 -hlsl-dxilemit | %FileCheck %s
+
+// A thread launch node has neither a dispatch grid nor a thread group, so none
+// of the thread-identity intrinsics are legal for it.
+// ValidateDxilOperationCallInProfile in DxilValidation.cpp rejects ThreadId,
+// GroupId, ThreadIdInGroup and FlattenedThreadIdInGroup for a thread launch
+// node. Nothing in the shader tells one invocation from another, so the pass
+// keeps the select-everything criterion. This test pins that decision, so that
+// a change cannot emit an intrinsic that the validator rejects.
+
+// CHECK: NodeInvocationSelection:None
+
+// CHECK-NOT: @dx.op.threadId.i32
+// CHECK-NOT: @dx.op.threadIdInGroup.i32
+// CHECK-NOT: @dx.op.flattenedThreadIdInGroup.i32
+// CHECK: br i1 true, label %PIXInterestingBlock, label %PIXNonInterestingBlock
+
+RWStructuredBuffer<uint> Output : register(u0);
+
+struct Record
+{
+  uint value;
+};
+
+[Shader("node")]
+[NodeLaunch("thread")]
+void ThreadLaunchNode(ThreadNodeInputRecord<Record> input)
+{
+  Output[0] = input.Get().value;
+}