PIX: Cp of "Prevent ray query handle writes to PIX debug UAV" (#6362)

Original PR:
#6309

Original commit:

5a31785

AllocateRayQuery returns an int, but it's not really an int: it's a
handle to the query. Thus, it's not sensible for the PIX debugging
instrumentation to attempt to write it out to the debug-data-UAV. The
previous code did an explicit check against the type of the value to be
written, but if that value is actually a phi itself, then that check
would fail. So now we recursively run through the phi values looking to
see if any of its antecedents are a phi, and if so, refuse to send its
value to the UAV.

(cherry picked from commit 5a31785)

Author: jeffnn

lib/DxilPIXPasses/DxilDbgValueToDbgDeclare.cpp

@@ -17,6 +17,7 @@
 
 #include "dxc/DXIL/DxilConstants.h"
 #include "dxc/DXIL/DxilModule.h"
+#include "dxc/DXIL/DxilOperations.h"
 #include "dxc/DXIL/DxilResourceBase.h"
 #include "dxc/DxilPIXPasses/DxilPIXPasses.h"
 #include "llvm/ADT/STLExtras.h"
@@ -597,23 +598,27 @@ bool DxilDbgValueToDbgDeclare::runOnModule(llvm::Module &M) {
 
   auto &Functions = M.getFunctionList();
   for (auto &fn : Functions) {
-    // #DSLTodo: We probably need to merge the list of variables for each export
-    // into one set so that WinPIX shader debugging can follow a thread through
-    // any function within a given module. (Unless PIX chooses to launch a new
-    // debugging session whenever control passes from one function to another.)
-    // For now, it's sufficient to treat each exported function as having
-    // completely separate variables by clearing this member:
+    llvm::SmallPtrSet<Value *, 16> RayQueryHandles;
+    PIXPassHelpers::FindRayQueryHandlesForFunction(&fn, RayQueryHandles);
+    // #DSLTodo: We probably need to merge the list of variables for each
+    // export into one set so that WinPIX shader debugging can follow a
+    // thread through any function within a given module. (Unless PIX
+    // chooses to launch a new debugging session whenever control passes
+    // from one function to another.) For now, it's sufficient to treat each
+    // exported function as having completely separate variables by clearing
+    // this member:
     m_Registers.clear();
-    // Note: they key problem here is variables in common functions called by
-    // multiple exported functions. The DILocalVariables in the common function
-    // will be exactly the same objects no matter which export called the common
-    // function, so the instrumentation here gets a bit confused that the same
-    // variable is present in two functions and ends up pointing one function
-    // to allocas in another function. (This is easy to repro: comment out the
-    // above clear(), and run PixTest::PixStructAnnotation_Lib_DualRaygen.)
-    // Not sure what the right path forward is: might be that we have to tag
-    // m_Registers with the exported function, and maybe write out a function
-    // identifier during debug instrumentation...
+    // Note: they key problem here is variables in common functions called
+    // by multiple exported functions. The DILocalVariables in the common
+    // function will be exactly the same objects no matter which export
+    // called the common function, so the instrumentation here gets a bit
+    // confused that the same variable is present in two functions and ends
+    // up pointing one function to allocas in another function. (This is
+    // easy to repro: comment out the above clear(), and run
+    // PixTest::PixStructAnnotation_Lib_DualRaygen.) Not sure what the right
+    // path forward is: might be that we have to tag m_Registers with the
+    // exported function, and maybe write out a function identifier during
+    // debug instrumentation...
     auto &blocks = fn.getBasicBlockList();
     if (!blocks.empty()) {
       for (auto &block : blocks) {
@@ -640,9 +645,8 @@ bool DxilDbgValueToDbgDeclare::runOnModule(llvm::Module &M) {
           if (auto *DbgValue =
                   llvm::dyn_cast<llvm::DbgValueInst>(instruction)) {
             llvm::Value *V = DbgValue->getValue();
-            if (PIXPassHelpers::IsAllocateRayQueryInstruction(V)) {
+            if (RayQueryHandles.count(V) != 0)
               continue;
-            }
             Changed = true;
             handleDbgValue(M, DbgValue);
             DbgValue->eraseFromParent();
@@ -809,8 +813,8 @@ static llvm::DIType *FindStructMemberTypeAtOffset(llvm::DICompositeType *Ty,
     }
   }
 
-  // Structure resources are expected to fail this (they have no real meaning in
-  // storage)
+  // Structure resources are expected to fail this (they have no real
+  // meaning in storage)
   if (SortedMembers.size() == 1) {
     switch (SortedMembers.begin()->second->getTag()) {
     case llvm::dwarf::DW_TAG_structure_type:
@@ -907,8 +911,9 @@ void DxilDbgValueToDbgDeclare::handleDbgValue(llvm::Module &M,
   }
 
   // Members' "base type" is actually the containing aggregate's type.
-  // To find the actual type of the variable, we must descend the container's
-  // type hierarchy to find the type at the expected offset/size.
+  // To find the actual type of the variable, we must descend the
+  // container's type hierarchy to find the type at the expected
+  // offset/size.
   if (auto *DerivedTy = llvm::dyn_cast<llvm::DIDerivedType>(Ty)) {
     const llvm::DITypeIdentifierMap EmptyMap;
     switch (DerivedTy->getTag()) {
@@ -961,10 +966,10 @@ void DxilDbgValueToDbgDeclare::handleDbgValue(llvm::Module &M,
 
       auto *Zero = B.getInt32(0);
 
-      // Now traverse a list of pairs {Scalar Value, InitialOffset + Offset}.
-      // InitialOffset is the offset from DbgValue's expression (i.e., the
-      // offset from the Variable's start), and Offset is the Scalar Value's
-      // packed offset from DbgValue's value.
+      // Now traverse a list of pairs {Scalar Value, InitialOffset +
+      // Offset}. InitialOffset is the offset from DbgValue's expression
+      // (i.e., the offset from the Variable's start), and Offset is the
+      // Scalar Value's packed offset from DbgValue's value.
       for (const ValueAndOffset &VO : SplitValue(V, InitialOffset, B)) {
 
         OffsetInBits AlignedOffset;
@@ -1022,11 +1027,11 @@ bool DxilDbgValueToDbgDeclare::handleStoreIfDestIsGlobal(
               llvm::dyn_cast<llvm::GetElementPtrInst>(asInstr.Get())) {
         // We are only interested in the case of basic types within an array
         // because the PIX debug instrumentation operates at that level.
-        // Aggregate members will have been descended through to produce their
-        // own entries in the GlobalStorageMap.
-        // Consequently, we're only interested in the GEP's index into the
-        // array. Any deeper indexing in the GEP will be for embedded
-        // aggregates. The three operands in such a GEP mean:
+        // Aggregate members will have been descended through to produce
+        // their own entries in the GlobalStorageMap. Consequently, we're
+        // only interested in the GEP's index into the array. Any deeper
+        // indexing in the GEP will be for embedded aggregates. The three
+        // operands in such a GEP mean:
         //    0 = the pointer
         //    1 = dereference the pointer (expected to be constant int zero)
         //    2 = the index into the array
@@ -1035,7 +1040,8 @@ bool DxilDbgValueToDbgDeclare::handleStoreIfDestIsGlobal(
             llvm::dyn_cast<ConstantInt>(asGEP->getOperand(1))
                     ->getLimitedValue() == 0) {
           // TODO: The case where this index is not a constant int
-          // (Needs changes to the allocas generated elsewhere in this pass.)
+          // (Needs changes to the allocas generated elsewhere in this
+          // pass.)
           if (auto *arrayIndexAsConstInt =
                   llvm::dyn_cast<ConstantInt>(asGEP->getOperand(2))) {
             int MemberIndex = arrayIndexAsConstInt->getLimitedValue();
@@ -1120,9 +1126,9 @@ VariableRegisters::VariableRegisters(
   PopulateAllocaMap(Ty);
   m_Offsets.AlignTo(Ty); // For padding.
 
-  // (min16* types can occupy 16 or 32 bits depending on whether or not they are
-  // natively supported. If non-native, the alignment will be 32, but the
-  // claimed size will still be 16, hence the "max" here)
+  // (min16* types can occupy 16 or 32 bits depending on whether or not they
+  // are natively supported. If non-native, the alignment will be 32, but
+  // the claimed size will still be 16, hence the "max" here)
   assert(m_Offsets.GetCurrentAlignedOffset() ==
          std::max<uint64_t>(DITypePeelTypeAlias(Ty)->getSizeInBits(),
                             DITypePeelTypeAlias(Ty)->getAlignInBits()));

lib/DxilPIXPasses/DxilDebugInstrumentation.cpp

@@ -48,10 +48,10 @@ using namespace hlsl;
 // UAV if the instance is not of interest.
 //
 // In addition, each half of the UAV is further subdivided: the first quarter is
-// the are in which blocks are permitted to start writing their sequence, and
+// the area in which blocks are permitted to start writing their sequence, and
 // that sequence is constrained to be no longer than the size of the second
 // quarter. This allows us to limit writes to the appropriate half of the UAV
-// via a single AND at the beginning of the basic block. An additoinal OR
+// via a single AND at the beginning of the basic block. An additional OR
 // provides the offset, either 0 for threads-of-interest, or UAVSize/2 for
 // not-of-interest.
 //
@@ -341,8 +341,9 @@ class DxilDebugInstrumentation : public ModulePass {
   void reserveDebugEntrySpace(BuilderContext &BC, uint32_t SpaceInDwords);
   std::optional<InstructionAndType> addStoreStepDebugEntry(BuilderContext *BC,
                                                            StoreInst *Inst);
-  std::optional<InstructionAndType> addStepDebugEntry(BuilderContext *BC,
-                                                      Instruction *Inst);
+  std::optional<InstructionAndType>
+  addStepDebugEntry(BuilderContext *BC, Instruction *Inst,
+                    llvm::SmallPtrSetImpl<Value *> const &RayQueryHandles);
   std::optional<DebugShaderModifierRecordType>
   addStepDebugEntryValue(BuilderContext *BC, std::uint32_t InstNum, Value *V,
                          std::uint32_t ValueOrdinal, Value *ValueOrdinalIndex);
@@ -362,8 +363,9 @@ class DxilDebugInstrumentation : public ModulePass {
     uint32_t FirstInstructionOrdinalInBlock;
     std::vector<InstructionToInstrument> Instructions;
   };
-  BlockInstrumentationData FindInstrumentableInstructionsInBlock(BasicBlock &BB,
-                                                                 OP *HlslOP);
+  BlockInstrumentationData FindInstrumentableInstructionsInBlock(
+      BasicBlock &BB, OP *HlslOP,
+      llvm::SmallPtrSetImpl<Value *> const &RayQueryHandles);
   uint32_t
   CountBlockPayloadBytes(std::vector<InstructionToInstrument> const &IsAndTs);
 };
@@ -1003,10 +1005,6 @@ DxilDebugInstrumentation::addStoreStepDebugEntry(BuilderContext *BC,
     return std::nullopt;
   }
 
-  if (PIXPassHelpers::IsAllocateRayQueryInstruction(Inst->getValueOperand())) {
-    return std::nullopt;
-  }
-
   auto Type = addStepDebugEntryValue(BC, InstNum, Inst->getValueOperand(),
                                      ValueOrdinalBase, ValueOrdinalIndex);
   if (Type) {
@@ -1054,20 +1052,25 @@ DxilDebugInstrumentation::addStoreStepDebugEntry(BuilderContext *BC,
   return std::nullopt;
 }
 
-std::optional<InstructionAndType>
-DxilDebugInstrumentation::addStepDebugEntry(BuilderContext *BC,
-                                            Instruction *Inst) {
-  if (PIXPassHelpers::IsAllocateRayQueryInstruction(Inst)) {
+std::optional<InstructionAndType> DxilDebugInstrumentation::addStepDebugEntry(
+    BuilderContext *BC, Instruction *Inst,
+    llvm::SmallPtrSetImpl<Value *> const &RayQueryHandles) {
+
+  std::uint32_t InstNum;
+  if (!pix_dxil::PixDxilInstNum::FromInst(Inst, &InstNum)) {
     return std::nullopt;
   }
 
-  if (auto *St = llvm::dyn_cast<llvm::StoreInst>(Inst)) {
-    return addStoreStepDebugEntry(BC, St);
+  if (RayQueryHandles.count(Inst) != 0) {
+    InstructionAndType ret{};
+    ret.Inst = Inst;
+    ret.InstructionOrdinal = InstNum;
+    ret.Type = DebugShaderModifierRecordTypeDXILStepVoid;
+    return ret;
   }
 
-  std::uint32_t InstNum;
-  if (!pix_dxil::PixDxilInstNum::FromInst(Inst, &InstNum)) {
-    return std::nullopt;
+  if (auto *St = llvm::dyn_cast<llvm::StoreInst>(Inst)) {
+    return addStoreStepDebugEntry(BC, St);
   }
 
   if (auto *Ld = llvm::dyn_cast<llvm::LoadInst>(Inst)) {
@@ -1300,8 +1303,9 @@ Instruction *FindFirstNonPhiInstruction(Instruction *I) {
 // If indicator is "d", a single integer denoting the base for the alloca
 // store.
 DxilDebugInstrumentation::BlockInstrumentationData
-DxilDebugInstrumentation::FindInstrumentableInstructionsInBlock(BasicBlock &BB,
-                                                                OP *HlslOP) {
+DxilDebugInstrumentation::FindInstrumentableInstructionsInBlock(
+    BasicBlock &BB, OP *HlslOP,
+    llvm::SmallPtrSetImpl<Value *> const &RayQueryHandles) {
   BlockInstrumentationData ret{};
   auto &Is = BB.getInstList();
   *OSOverride << "Block#";
@@ -1315,7 +1319,7 @@ DxilDebugInstrumentation::FindInstrumentableInstructionsInBlock(BasicBlock &BB,
         FoundFirstInstruction = true;
       }
     }
-    auto IandT = addStepDebugEntry(nullptr, &Inst);
+    auto IandT = addStepDebugEntry(nullptr, &Inst, RayQueryHandles);
     if (IandT) {
       InstructionToInstrument DebugOutputForThisInstruction{};
       DebugOutputForThisInstruction.ValueType = IandT->Type;
@@ -1394,6 +1398,8 @@ bool DxilDebugInstrumentation::RunOnFunction(Module &M, DxilModule &DM,
   default:
     return false;
   }
+  llvm::SmallPtrSet<Value *, 16> RayQueryHandles;
+  PIXPassHelpers::FindRayQueryHandlesForFunction(function, RayQueryHandles);
 
   Instruction *firstInsertionPt = dxilutil::FirstNonAllocaInsertionPt(function);
   IRBuilder<> Builder(firstInsertionPt);
@@ -1436,7 +1442,7 @@ bool DxilDebugInstrumentation::RunOnFunction(Module &M, DxilModule &DM,
                   values.AddedBlocksToIgnoreForInstrumentation.end(),
                   &BB) == values.AddedBlocksToIgnoreForInstrumentation.end()) {
       auto BlockInstrumentation =
-          FindInstrumentableInstructionsInBlock(BB, BC.HlslOP);
+          FindInstrumentableInstructionsInBlock(BB, BC.HlslOP, RayQueryHandles);
       if (BlockInstrumentation.FirstInstructionOrdinalInBlock <
               m_FirstInstruction ||
           BlockInstrumentation.FirstInstructionOrdinalInBlock >=

lib/DxilPIXPasses/PixPassHelpers.cpp

@@ -37,15 +37,33 @@ using namespace llvm;
 using namespace hlsl;
 
 namespace PIXPassHelpers {
-bool IsAllocateRayQueryInstruction(llvm::Value const *Val) {
-  if (Val != nullptr) {
-    if (llvm::Instruction const *Inst =
-            llvm::dyn_cast<llvm::Instruction>(Val)) {
-      return hlsl::OP::IsDxilOpFuncCallInst(Inst,
-                                            hlsl::OP::OpCode::AllocateRayQuery);
+static void FindRayQueryHandlesFromUse(Value *U,
+                                       SmallPtrSetImpl<Value *> &Handles) {
+  if (Handles.insert(U).second) {
+    auto RayQueryHandleUses = U->uses();
+    for (Use &Use : RayQueryHandleUses) {
+      iterator_range<Value::user_iterator> Users = Use->users();
+      for (User *User : Users) {
+        if (isa<PHINode>(User) || isa<SelectInst>(User))
+          FindRayQueryHandlesFromUse(User, Handles);
+      }
+    }
+  }
+}
+
+void FindRayQueryHandlesForFunction(llvm::Function *F,
+                                    SmallPtrSetImpl<Value *> &RayQueryHandles) {
+  auto &blocks = F->getBasicBlockList();
+  if (!blocks.empty()) {
+    for (auto &block : blocks) {
+      for (auto &instruction : block) {
+        if (hlsl::OP::IsDxilOpFuncCallInst(
+                &instruction, hlsl::OP::OpCode::AllocateRayQuery)) {
+          FindRayQueryHandlesFromUse(&instruction, RayQueryHandles);
+        }
+      }
     }
   }
-  return false;
 }
 
 static unsigned int

lib/DxilPIXPasses/PixPassHelpers.h

@@ -9,6 +9,8 @@
 
 #pragma once
 
+#include <vector>
+
 #include "dxc/DXIL/DxilModule.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/IRBuilder.h"
@@ -30,7 +32,8 @@ class ScopedInstruction {
   llvm::Instruction *Get() const { return m_Instruction; }
 };
 
-bool IsAllocateRayQueryInstruction(llvm::Value const *Val);
+void FindRayQueryHandlesForFunction(
+    llvm::Function *F, llvm::SmallPtrSetImpl<llvm::Value *> &RayQueryHandles);
 llvm::CallInst *CreateUAV(hlsl::DxilModule &DM, llvm::IRBuilder<> &Builder,
                           unsigned int registerId, const char *name);
 llvm::CallInst *CreateHandleForResource(hlsl::DxilModule &DM,

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

@@ -4,14 +4,21 @@
 // CHECK: call void @dx.op.bufferStore.i32(i32 69, %dx.types.Handle
 
 RaytracingAccelerationStructure RTAS;
+RWStructuredBuffer<float> UAV : register(u0);
 
 void DoTrace(RayQuery<RAY_FLAG_FORCE_OPAQUE|RAY_FLAG_SKIP_PROCEDURAL_PRIMITIVES> rayQuery, RayDesc rayDesc) {
   rayQuery.TraceRayInline(RTAS, 0, 1, rayDesc);
 }
 
 float main(RayDesc rayDesc : RAYDESC) : OUT {
-  RayQuery<RAY_FLAG_FORCE_OPAQUE|RAY_FLAG_SKIP_PROCEDURAL_PRIMITIVES> rayQuery;
+  RayQuery<RAY_FLAG_FORCE_OPAQUE | RAY_FLAG_SKIP_PROCEDURAL_PRIMITIVES> rayQuery;
   DoTrace(rayQuery, rayDesc);
   rayQuery.TraceRayInline(RTAS, 1, 2, rayDesc);
+  while (rayQuery.Proceed()) {
+    UAV[0] = rayQuery.CandidateGeometryIndex();
+    UAV[1] = rayQuery.CandidatePrimitiveIndex();
+    UAV[2] = rayQuery.CandidateTriangleBarycentrics().x;
+    rayQuery.CommitNonOpaqueTriangleHit();
+  }
   return 0;
 }

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

@@ -0,0 +1,30 @@
+// RUN: %dxc -T vs_6_5 -Od -E main %s | %opt -S -dxil-annotate-with-virtual-regs -hlsl-dxil-debug-instrumentation | FileCheck %s
+ 
+// CHECK: [[RAYQUERY0:%.*]] = call i32 @dx.op.allocateRayQuery
+// CHECK: [[RAYQUERY1:%.*]] = call i32 @dx.op.allocateRayQuery
+
+// CHECK: [[PHIQUERY:%.*]] = phi i32 [ [[RAYQUERY0:%.*]], {{.*}} ], [ [[RAYQUERY0:%.*]], {{.*}} ]
+
+// The debug instrumentation should NOT try to store this phi value: it's not an i32! (It's an opaque handle).
+// CHECK-NOT: @dx.op.bufferStore{{.*}}[[PHIQUERY]]
+
+RaytracingAccelerationStructure RTAS;
+RWStructuredBuffer<int> UAV : register(u0);
+
+float main(RayDesc rayDesc
+           : RAYDESC) : OUT {
+  RayQuery<RAY_FLAG_FORCE_OPAQUE | RAY_FLAG_SKIP_PROCEDURAL_PRIMITIVES> rayQuery0;
+  RayQuery<RAY_FLAG_FORCE_OPAQUE | RAY_FLAG_SKIP_PROCEDURAL_PRIMITIVES> rayQuery1;
+  rayQuery0.TraceRayInline(RTAS, 1, 2, rayDesc);
+  rayQuery1.TraceRayInline(RTAS, 1, 2, rayDesc);
+  RayQuery<RAY_FLAG_FORCE_OPAQUE | RAY_FLAG_SKIP_PROCEDURAL_PRIMITIVES> usedQuery = rayQuery0;
+  if (UAV[0] == 1)
+    usedQuery = rayQuery1;
+  UAV[1] = usedQuery.CandidatePrimitiveIndex();
+  rayQuery0.CommitNonOpaqueTriangleHit();
+  if (rayQuery0.CommittedStatus() == COMMITTED_TRIANGLE_HIT) {
+    UAV[4] = 42;
+  }
+
+  return 0;
+}