[PIX] Attribute library accesses to reaching entry points

lib/DxilPIXPasses/DxilShaderAccessTracking.cpp

@@ -28,6 +28,8 @@
 #include "llvm/Transforms/Utils/Local.h"
 
 #include <deque>
+#include <set>
+#include <vector>
 
 #include "PixPassHelpers.h"
 
@@ -546,11 +548,16 @@ bool DxilShaderAccessTracking::EmitResourceAccess(DxilModule &DM,
           Builder.CreateMul(ZeroIfOutOfBounds, EncodedFlags);
       uint32_t InstructionNumber = 0;
       (void)pix_dxil::PixDxilInstNum::FromInst(instruction, &InstructionNumber);
-      auto const *shaderModel = DM.GetShaderModel();
-      auto shaderKind = shaderModel->GetKind();
-      uint32_t EncodedInstructionNumber = InstructionNumber |
-                                          InstructionOrdinalndicator |
-                                          EncodeShaderModel(shaderKind);
+      auto *EncodedShaderKindConstant = cast<ConstantInt>(
+          m_FunctionToEncodedAccess.at(Builder.GetInsertBlock()->getParent())
+              .at(ResourceAccessStyle::None));
+      uint32_t EncodedShaderKind = EncodedShaderKindConstant->getLimitedValue();
+      // The ordinal occupies the low 24 bits. Mask it so it does not overlap
+      // the indicator or the shader kind.
+      constexpr uint32_t InstructionOrdinalMask = 0x00FF'FFFF;
+      uint32_t EncodedInstructionNumber =
+          (InstructionNumber & InstructionOrdinalMask) |
+          InstructionOrdinalndicator | EncodedShaderKind;
       auto *MultipliedOutOfBoundsValue = Builder.CreateMul(
           OneIfOutOfBounds, HlslOP->GetU32Const(EncodedInstructionNumber));
       auto *CombinedFlagOrInstructionValue =
@@ -810,6 +817,69 @@ DxilShaderAccessTracking::GetResourceFromHandle(Value *resHandle,
   return ret;
 }
 
+// Map each function to the shader kind of the entry point that reaches it.
+// A library helper has no DxilFunctionProps, so the module kind is Library,
+// which PIX cannot attribute to a pipeline stage. If more than one entry
+// kind reaches the same helper, keep the module kind. Entry points keep
+// their own kind.
+static std::map<llvm::Function *, DXIL::ShaderKind>
+ResolveShaderKindByReachingEntryPoint(DxilModule &DM) {
+  std::map<llvm::Function *, DXIL::ShaderKind> functionToShaderKind;
+
+  const DXIL::ShaderKind ambiguousShaderKind = DM.GetShaderModel()->GetKind();
+  auto entryPoints = DM.GetExportedFunctions();
+
+  for (llvm::Function *entryPoint : entryPoints) {
+    if (entryPoint == nullptr || entryPoint->isDeclaration()) {
+      continue;
+    }
+
+    const DXIL::ShaderKind entryPointShaderKind =
+        PIXPassHelpers::GetFunctionShaderKind(DM, entryPoint);
+
+    std::vector<llvm::Function *> pending{entryPoint};
+    std::set<llvm::Function *> visited;
+    while (!pending.empty()) {
+      llvm::Function *reached = pending.back();
+      pending.pop_back();
+      if (!visited.insert(reached).second) {
+        continue;
+      }
+
+      auto emplaced =
+          functionToShaderKind.emplace(reached, entryPointShaderKind);
+      if (!emplaced.second && emplaced.first->second != entryPointShaderKind) {
+        emplaced.first->second = ambiguousShaderKind;
+      }
+
+      for (llvm::BasicBlock &block : reached->getBasicBlockList()) {
+        for (llvm::Instruction &instruction : block.getInstList()) {
+          auto *call = llvm::dyn_cast<llvm::CallInst>(&instruction);
+          if (call == nullptr) {
+            continue;
+          }
+          llvm::Function *callee = call->getCalledFunction();
+          if (callee == nullptr || callee->isDeclaration() ||
+              callee->isIntrinsic() || hlsl::OP::IsDxilOpFunc(callee)) {
+            continue;
+          }
+          pending.push_back(callee);
+        }
+      }
+    }
+  }
+
+  for (llvm::Function *entryPoint : entryPoints) {
+    if (entryPoint == nullptr || entryPoint->isDeclaration()) {
+      continue;
+    }
+    functionToShaderKind[entryPoint] =
+        PIXPassHelpers::GetFunctionShaderKind(DM, entryPoint);
+  }
+
+  return functionToShaderKind;
+}
+
 bool DxilShaderAccessTracking::runOnModule(Module &M) {
   // This pass adds instrumentation for shader access to resources
 
@@ -840,6 +910,8 @@ bool DxilShaderAccessTracking::runOnModule(Module &M) {
     auto instrumentableFunctions =
         PIXPassHelpers::GetAllInstrumentableFunctions(DM);
 
+    auto functionToShaderKind = ResolveShaderKindByReachingEntryPoint(DM);
+
     if (DM.m_ShaderFlags.GetForceEarlyDepthStencil()) {
       if (OSOverride != nullptr) {
         formatted_raw_ostream FOS(*OSOverride);
@@ -851,17 +923,11 @@ bool DxilShaderAccessTracking::runOnModule(Module &M) {
         PIXPassHelpers::CreateGlobalUAVResource(DM, 0u, "PIX_ShaderAccessUAV");
 
     for (auto *F : instrumentableFunctions) {
-      DXIL::ShaderKind shaderKind = DXIL::ShaderKind::Invalid;
-      if (!DM.HasDxilFunctionProps(F)) {
-        auto ShaderModel = DM.GetShaderModel();
-        shaderKind = ShaderModel->GetKind();
-        if (shaderKind == DXIL::ShaderKind::Library) {
-          continue;
-        }
-      } else {
-        hlsl::DxilFunctionProps const &props = DM.GetDxilFunctionProps(F);
-        shaderKind = props.shaderKind;
-      }
+      auto reachedFrom = functionToShaderKind.find(F);
+      DXIL::ShaderKind shaderKind =
+          reachedFrom != functionToShaderKind.end()
+              ? reachedFrom->second
+              : PIXPassHelpers::GetFunctionShaderKind(DM, F);
 
       IRBuilder<> Builder(F->getEntryBlock().getFirstInsertionPt());
 
@@ -959,9 +1025,13 @@ bool DxilShaderAccessTracking::runOnModule(Module &M) {
           }
 
           for (unsigned iParam : handleParams) {
+            auto uavHandle =
+                m_FunctionToUAVHandle.find(CallerParent->getParent());
+            if (uavHandle == m_FunctionToUAVHandle.end())
+              continue;
+
             // Don't instrument the accesses to the UAV that we just added
-            if (Call->getArgOperand(iParam) ==
-                m_FunctionToUAVHandle[CallerParent->getParent()])
+            if (Call->getArgOperand(iParam) == uavHandle->second)
               continue;
             auto res = GetResourceFromHandle(Call->getArgOperand(iParam), DM);
             if (res.accessStyle == AccessStyle::None) {

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

@@ -0,0 +1,31 @@
+// RUN: %dxc -T lib_6_6 -Od %s | %opt -S -hlsl-dxil-pix-shader-access-instrumentation,config=.256;512;1024. | %FileCheck %s
+
+// A descriptor-heap access record carries the shader kind of the entry
+// point that reached the access. HeapHelper is [noinline] so the access
+// stays in the helper. The descriptor index is a parameter so the helper
+// is not folded away.
+//
+// The kind occupies bits 31:28. An out-of-bounds record sets the
+// instruction-ordinal indicator (bit 27). RayGeneration is 7 and UAVWrite
+// is 3, so the in-bounds flags are 0x73000000 == 1929379840 and the
+// out-of-bounds value is 0x78000000 == 2013265920. Under the module kind,
+// Library (6), those values would be 0x63000000 and 0x68000000.
+
+// CHECK: define void {{.*}}HeapHelper
+// CHECK-NOT: 1660944384
+// CHECK: mul i32 {{.*}}, 1929379840
+// CHECK-NOT: 1744830464
+// CHECK: mul i32 {{.*}}, 2013265920
+
+[noinline]
+export void HeapHelper(uint descriptorIndex)
+{
+    RWByteAddressBuffer heapBuffer = ResourceDescriptorHeap[descriptorIndex];
+    heapBuffer.Store(0, 1);
+}
+
+[shader("raygeneration")]
+void RayGen()
+{
+    HeapHelper(1);
+}

tools/clang/unittests/HLSL/PixTest.cpp

@@ -125,6 +125,8 @@ class PixTest : public ::testing::Test {
   TEST_METHOD(AccessTracking_DynamicRangeRegisterIndex_SM66)
   TEST_METHOD(AccessTracking_ConstantIndexAtRangeLimit)
   TEST_METHOD(AccessTracking_SamplerAccessInLibrary)
+  TEST_METHOD(AccessTracking_OobBindlessUsesFunctionShaderKind)
+  TEST_METHOD(AccessTracking_LibraryNonEntryFunction)
 
   TEST_METHOD(PixStructAnnotation_Lib_DualRaygen)
 
@@ -1318,6 +1320,28 @@ static bool HasBufferStoreWithByteOffset(std::vector<std::string> const &lines,
   return false;
 }
 
+static bool
+HasBufferStoreValueMatchingMask(std::vector<std::string> const &lines,
+                                uint32_t mask, uint32_t maskedValue) {
+  for (auto const &line : lines) {
+    if (line.find("dx.op.bufferStore") == std::string::npos) {
+      continue;
+    }
+
+    size_t position = 0;
+    while ((position = line.find("i32 ", position)) != std::string::npos) {
+      position += 4;
+      char *end = nullptr;
+      uint32_t value =
+          static_cast<uint32_t>(strtoul(line.c_str() + position, &end, 10));
+      if (end != line.c_str() + position && (value & mask) == maskedValue) {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
 TEST_F(PixTest, AccessTracking_MultipleDynamicRangesSameTypeAndSpace) {
   const char *hlsl = R"(
 ByteAddressBuffer g_indices : register(t0);
@@ -1416,6 +1440,64 @@ void RayGen()
       output.blob, "shader access tracking of a library sampler access");
 }
 
+TEST_F(PixTest, AccessTracking_OobBindlessUsesFunctionShaderKind) {
+  if (m_ver.SkipDxilVersion(1, 6)) {
+    return;
+  }
+
+  const char *hlsl = R"(
+[shader("raygeneration")]
+void RayGen()
+{
+    RWByteAddressBuffer output = ResourceDescriptorHeap[1];
+    output.Store(0, 1);
+}
+)";
+
+  auto compiled = Compile(m_dllSupport, hlsl, L"lib_6_6", {L"-Od"});
+  auto output = RunShaderAccessTrackingPass(compiled, L".0;0;0.");
+  auto lines = Split(Disassemble(output.blob), '\n');
+  VERIFY_IS_TRUE(
+      HasBufferStoreValueMatchingMask(lines, 0xF8000000, 0x78000000));
+  VERIFY_IS_TRUE(
+      !HasBufferStoreValueMatchingMask(lines, 0xF8000000, 0x68000000));
+  VerifyInstrumentedModuleIsValid(
+      output.blob,
+      "shader access tracking of an out-of-bounds bindless access");
+}
+
+TEST_F(PixTest, AccessTracking_LibraryNonEntryFunction) {
+  if (m_ver.SkipDxilVersion(1, 6)) {
+    return;
+  }
+
+  const char *hlsl = R"(
+Texture2D<float4> g_texture : register(t0);
+RWByteAddressBuffer g_output : register(u0);
+
+export float4 Helper(uint index)
+{
+    float4 value = g_texture.Load(int3(index, 0, 0));
+    g_output.Store(0, asuint(value.x));
+    return value;
+}
+
+[shader("raygeneration")]
+void RayGen()
+{
+    Helper(0);
+}
+)";
+
+  auto compiled = Compile(m_dllSupport, hlsl, L"lib_6_6", {L"-Od"});
+  auto output =
+      RunShaderAccessTrackingPass(compiled, L"S0:0:4i0;U0:4:4i0;.0;0;0.");
+  auto text = JoinLines(output.lines);
+  VERIFY_IS_TRUE(text.find("NotModified") == std::string::npos);
+  VerifyInstrumentedModuleIsValid(
+      output.blob, "shader access tracking of a library helper function");
+}
+
 TEST_F(PixTest, AddToASGroupSharedPayload) {
 
   const char *hlsl = R"(