[PIX] Fix DXR invocation log bounds

lib/DxilPIXPasses/DxilPIXDXRInvocationsLog.cpp

@@ -18,6 +18,7 @@
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Support/FormattedStream.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Local.h"
 
 #include "PixPassHelpers.h"
@@ -66,6 +67,11 @@ bool DxilPIXDXRInvocationsLog::runOnModule(Module &M) {
   LLVMContext &Ctx = M.getContext();
   OP *HlslOP = DM.GetOP();
 
+  // A zero-entry log has no space for records.
+  if (m_MaxNumEntriesInLog == 0) {
+    return false;
+  }
+
   bool Modified = false;
 
   for (auto entryFunction : DM.GetExportedFunctions()) {
@@ -85,7 +91,9 @@ bool DxilPIXDXRInvocationsLog::runOnModule(Module &M) {
 
     Modified = true;
 
-    IRBuilder<> Builder(dxilutil::FirstNonAllocaInsertionPt(entryFunction));
+    Instruction *InsertionPoint =
+        dxilutil::FirstNonAllocaInsertionPt(entryFunction);
+    IRBuilder<> Builder(InsertionPoint);
 
     // Add the UAVs that we're going to write to
     CallInst *HandleForCountUAV = PIXPassHelpers::CreateUAVOnceForModule(
@@ -163,10 +171,6 @@ bool DxilPIXDXRInvocationsLog::runOnModule(Module &M) {
     Constant *AtomicAdd =
         HlslOP->GetU32Const((unsigned)DXIL::AtomicBinOpCode::Add);
 
-    Function *UMinOpFunc =
-        HlslOP->GetOpFunc(OP::OpCode::UMin, Type::getInt32Ty(Ctx));
-    Constant *UMinOpCode = HlslOP->GetU32Const((unsigned)OP::OpCode::UMin);
-
     Function *StoreFuncFloat =
         HlslOP->GetOpFunc(OP::OpCode::BufferStore, Type::getFloatTy(Ctx));
     Function *StoreFuncInt =
@@ -177,8 +181,8 @@ bool DxilPIXDXRInvocationsLog::runOnModule(Module &M) {
     Constant *WriteMask_XYZW = HlslOP->GetI8Const(15);
     Constant *WriteMask_X = HlslOP->GetI8Const(1);
     Constant *ShaderKindAsConstant = HlslOP->GetU32Const((uint32_t)ShaderKind);
-    Constant *MaxEntryIndexAsConstant =
-        HlslOP->GetU32Const((uint32_t)m_MaxNumEntriesInLog - 1u);
+    Constant *MaxEntryCountAsConstant =
+        HlslOP->GetU32Const((uint32_t)m_MaxNumEntriesInLog);
     Constant *Zero32Arg = HlslOP->GetU32Const(0);
     Constant *One32Arg = HlslOP->GetU32Const(1);
     UndefValue *UndefArg = UndefValue::get(Type::getInt32Ty(Ctx));
@@ -198,19 +202,21 @@ bool DxilPIXDXRInvocationsLog::runOnModule(Module &M) {
         },
         "EntryIndexResult");
 
-    // Clamp the index so that we don't write off the end of the UAV. If we
-    // clamp, then it's up to PIX to replay the work again with a larger log
-    // buffer.
-    auto *EntryIndexClamped = Builder.CreateCall(
-        UMinOpFunc, {UMinOpCode, EntryIndex, MaxEntryIndexAsConstant});
+    // The counter keeps counting past the log capacity. Skip the stores once
+    // the claimed slot is out of range, so the recorded entries stay intact.
+    auto *EntryIndexIsInRange = Builder.CreateICmpULT(
+        EntryIndex, MaxEntryCountAsConstant, "EntryIndexIsInRange");
+    TerminatorInst *StoreEntryBlockTerminator =
+        SplitBlockAndInsertIfThen(EntryIndexIsInRange, InsertionPoint,
+                                  /*Unreachable*/ false);
+    Builder.SetInsertPoint(StoreEntryBlockTerminator);
 
     const auto numBytesPerEntry =
         4 + (3 * 4) + (3 * 4) + (3 * 4) + 4 + 4 +
         4; // See number of bytes we store per shader invocation below
 
-    auto EntryOffset =
-        Builder.CreateMul(EntryIndexClamped,
-                          HlslOP->GetU32Const(numBytesPerEntry), "EntryOffset");
+    auto EntryOffset = Builder.CreateMul(
+        EntryIndex, HlslOP->GetU32Const(numBytesPerEntry), "EntryOffset");
     auto EntryOffsetPlus16 = Builder.CreateAdd(
         EntryOffset, HlslOP->GetU32Const(16), "EntryOffsetPlus16");
     auto EntryOffsetPlus32 = Builder.CreateAdd(

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

@@ -14,7 +14,7 @@
 // Now check that at least three functions were modified (the hit group shaders):
 
 // -------- one ----------
-// Check for out-of-bounds clamp:
+// Check for the per-entry offset calculation:
 // CHECK: mul i32
 // CHECK: 52
 
@@ -28,7 +28,7 @@
 // CHECK: i32 1
 
 // -------- two ----------
-// Check for out-of-bounds clamp:
+// Check for the per-entry offset calculation:
 // CHECK: mul i32
 // CHECK: 52
 
@@ -43,7 +43,7 @@
 // CHECK: dx.op.atomicBinOp.i32
 // CHECK: i32 1
 
-// Check for out-of-bounds clamp:
+// Check for the per-entry offset calculation:
 // CHECK: mul i32
 // CHECK: 52
 

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

@@ -0,0 +1,57 @@
+// RUN: %dxc -Tlib_6_6 %s | %opt -S -hlsl-dxil-pix-dxr-invocations-log,maxNumEntriesInLog=100 | %FileCheck %s
+
+// Each invocation claims one log slot.
+// The counter keeps counting past the log capacity.
+// The stores execute only when the claimed slot is in range.
+
+// CHECK: [[ENTRYINDEX:%EntryIndexResult[0-9]*]] = call i32 @dx.op.atomicBinOp.i32(i32 78,
+// CHECK: [[INRANGE:%EntryIndexIsInRange[0-9]*]] = icmp ult i32 [[ENTRYINDEX]], 100
+// CHECK: br i1 [[INRANGE]]
+// CHECK: mul i32 [[ENTRYINDEX]], 52
+// CHECK: call void @dx.op.bufferStore.i32
+// CHECK: call void @dx.op.bufferStore.f32
+// CHECK: call void @dx.op.bufferStore.f32
+// CHECK: call void @dx.op.bufferStore.i32
+
+// UMin is not part of this shader.
+// CHECK-NOT: @dx.op.binary.i32(i32 40
+// CHECK-NOT: declare i32 @dx.op.binary.i32
+
+struct Payload
+{
+    float4 color;
+};
+
+struct Attribs
+{
+    float2 barycentrics;
+};
+
+RaytracingAccelerationStructure scene : register(t0);
+RWTexture2D<float4> output : register(u0);
+
+[shader("raygeneration")]
+void RayGen()
+{
+    RayDesc ray;
+    ray.Origin = float3(0, 0, 0);
+    ray.Direction = float3(0, 0, 1);
+    ray.TMin = 0.001f;
+    ray.TMax = 1000.f;
+    Payload payload;
+    payload.color = float4(0, 0, 0, 0);
+    TraceRay(scene, RAY_FLAG_NONE, ~0, 0, 1, 0, ray, payload);
+    output[DispatchRaysIndex().xy] = payload.color;
+}
+
+[shader("closesthit")]
+void ClosestHit(inout Payload payload, in Attribs attribs)
+{
+    payload.color = float4(attribs.barycentrics, 0, 1);
+}
+
+[shader("miss")]
+void Miss(inout Payload payload)
+{
+    payload.color = float4(1, 0, 0, 1);
+}

tools/clang/unittests/HLSL/PixTest.cpp

@@ -160,6 +160,10 @@ class PixTest : public ::testing::Test {
 
   TEST_METHOD(DxilPIXDXRInvocationsLog_SanityTest)
   TEST_METHOD(DxilPIXDXRInvocationsLog_EmbeddedRootSigs)
+  TEST_METHOD(DxilPIXDXRInvocationsLog_ZeroCapacityEmitsNothing)
+  TEST_METHOD(DxilPIXDXRInvocationsLog_OneEntryUsesEntryCountBound)
+  TEST_METHOD(DxilPIXDXRInvocationsLog_ExactCapacityUsesEntryCountBound)
+  TEST_METHOD(DxilPIXDXRInvocationsLog_OverflowGuardValidates)
 
   TEST_METHOD(DebugInstrumentation_TextOutput)
   TEST_METHOD(DebugInstrumentation_BlockReport)
@@ -636,7 +640,8 @@ class PixTest : public ::testing::Test {
   CComPtr<IDxcBlob>
   RunDxilPIXAddTidToAmplificationShaderPayloadPass(IDxcBlob *blob);
   CComPtr<IDxcBlob> RunDxilPIXMeshShaderOutputPass(IDxcBlob *blob);
-  CComPtr<IDxcBlob> RunDxilPIXDXRInvocationsLog(IDxcBlob *blob);
+  CComPtr<IDxcBlob>
+  RunDxilPIXDXRInvocationsLog(IDxcBlob *blob, unsigned maxNumEntriesInLog = 24);
   PassOutput
   RunDxilNonUniformResourceIndexInstrumentation(IDxcBlob *blob,
                                                 std::string &outputText);
@@ -676,6 +681,19 @@ static int CountToolsUAVRecords(std::vector<std::string> const &lines) {
   return count;
 }
 
+static bool
+HasDxrInvocationLogEntryCountCheck(std::vector<std::string> const &lines,
+                                   unsigned expectedEntryCount) {
+  const std::string expectedSuffix = ", " + std::to_string(expectedEntryCount);
+  for (auto const &line : lines) {
+    if (line.find("icmp ult i32 %EntryIndexResult") != std::string::npos &&
+        line.find(expectedSuffix) != std::string::npos) {
+      return true;
+    }
+  }
+  return false;
+}
+
 static bool
 RootSignatureHasToolsUAV(const DxilVersionedRootSignatureDesc *rootSignature,
                          uint32_t shaderRegister) {
@@ -962,15 +980,19 @@ CComPtr<IDxcBlob> PixTest::RunDxilPIXMeshShaderOutputPass(IDxcBlob *blob) {
   return pOptimizedModule;
 }
 
-CComPtr<IDxcBlob> PixTest::RunDxilPIXDXRInvocationsLog(IDxcBlob *blob) {
+CComPtr<IDxcBlob>
+PixTest::RunDxilPIXDXRInvocationsLog(IDxcBlob *blob,
+                                     unsigned maxNumEntriesInLog) {
 
   CComPtr<IDxcBlob> dxil = FindModule(DFCC_ShaderDebugInfoDXIL, blob);
   CComPtr<IDxcOptimizer> pOptimizer;
   VERIFY_SUCCEEDED(
       m_dllSupport.CreateInstance(CLSID_DxcOptimizer, &pOptimizer));
+  std::wstring logArg = L"-hlsl-dxil-pix-dxr-invocations-log,"
+                        L"maxNumEntriesInLog=" +
+                        std::to_wstring(maxNumEntriesInLog);
   std::vector<LPCWSTR> Options;
-  Options.push_back(
-      L"-hlsl-dxil-pix-dxr-invocations-log,maxNumEntriesInLog=24");
+  Options.push_back(logArg.c_str());
 
   CComPtr<IDxcBlob> pOptimizedModule;
   CComPtr<IDxcBlobEncoding> pText;
@@ -985,6 +1007,20 @@ CComPtr<IDxcBlob> PixTest::RunDxilPIXDXRInvocationsLog(IDxcBlob *blob) {
   return pOptimizedModule;
 }
 
+static const char *kSingleMissInvocationLogShader = R"x(
+struct [raypayload] MyPayload
+{
+    float2 barycentrics : read(caller) : write(caller,anyhit);
+    uint primitiveIndex : read(caller) : write(caller,anyhit);
+};
+
+[shader("miss")]
+void MissOne(inout MyPayload payload)
+{
+    payload.primitiveIndex = 1;
+}
+)x";
+
 PassOutput PixTest::RunDxilNonUniformResourceIndexInstrumentation(
     IDxcBlob *blob, std::string &outputText) {
 
@@ -3615,6 +3651,47 @@ void MyMiss(inout MyPayload payload)
   RunDxilPIXDXRInvocationsLog(compiledLib);
 }
 
+TEST_F(PixTest, DxilPIXDXRInvocationsLog_ZeroCapacityEmitsNothing) {
+  auto compiledLib =
+      Compile(m_dllSupport, kSingleMissInvocationLogShader, L"lib_6_6", {});
+
+  auto oneEntryOutput = RunDxilPIXDXRInvocationsLog(compiledLib, 1);
+  auto oneEntryLines = Tokenize(Disassemble(oneEntryOutput), "\n");
+  VERIFY_ARE_EQUAL(2, CountToolsUAVRecords(oneEntryLines));
+
+  auto zeroEntryOutput = RunDxilPIXDXRInvocationsLog(compiledLib, 0);
+  auto zeroEntryLines = Tokenize(Disassemble(zeroEntryOutput), "\n");
+  VERIFY_ARE_EQUAL(0, CountToolsUAVRecords(zeroEntryLines));
+}
+
+TEST_F(PixTest, DxilPIXDXRInvocationsLog_OneEntryUsesEntryCountBound) {
+  auto compiledLib =
+      Compile(m_dllSupport, kSingleMissInvocationLogShader, L"lib_6_6", {});
+  auto output = RunDxilPIXDXRInvocationsLog(compiledLib, 1);
+  auto lines = Tokenize(Disassemble(output), "\n");
+
+  VERIFY_IS_TRUE(HasDxrInvocationLogEntryCountCheck(lines, 1));
+}
+
+TEST_F(PixTest, DxilPIXDXRInvocationsLog_ExactCapacityUsesEntryCountBound) {
+  auto compiledLib =
+      Compile(m_dllSupport, kSingleMissInvocationLogShader, L"lib_6_6", {});
+  auto output = RunDxilPIXDXRInvocationsLog(compiledLib, 24);
+  auto lines = Tokenize(Disassemble(output), "\n");
+
+  VERIFY_IS_TRUE(HasDxrInvocationLogEntryCountCheck(lines, 24));
+}
+
+TEST_F(PixTest, DxilPIXDXRInvocationsLog_OverflowGuardValidates) {
+  auto compiledLib =
+      Compile(m_dllSupport, kSingleMissInvocationLogShader, L"lib_6_6", {});
+  auto output = RunDxilPIXDXRInvocationsLog(compiledLib, 1);
+  std::string disassembly = Disassemble(output);
+
+  VERIFY_IS_TRUE(disassembly.find("@dx.op.binary.i32") == std::string::npos);
+  VerifyInstrumentedModuleIsValid(output, "DXR invocations log overflow guard");
+}
+
 uint32_t NuriGetWaveInstructionCount(const std::vector<std::string> &lines) {
   // This is the instruction we'll insert into the shader if we detect dynamic
   // resource indexing