PIX: Fix instrumentation of libraries; instruction numbering reports (#4817)

* PIX: A few fixes for instrumentation of libraries; instruction numbering reporting

* Fix for tests: don't debug-instrument the patch constant fn

lib/DxilPIXPasses/DxilAnnotateWithVirtualRegister.cpp

@@ -14,6 +14,7 @@
 #include <memory>
 
 #include "dxc/DXIL/DxilModule.h"
+#include "dxc/DXIL/DxilOperations.h"
 #include "dxc/DxilPIXPasses/DxilPIXPasses.h"
 #include "dxc/DxilPIXPasses/DxilPIXVirtualRegisters.h"
 #include "dxc/Support/Global.h"
@@ -70,6 +71,7 @@ public:
   DxilAnnotateWithVirtualRegister() : llvm::ModulePass(ID) {}
 
   bool runOnModule(llvm::Module &M) override;
+  void applyOptions(llvm::PassOptions O) override;
 
 private:
   void AnnotateValues(llvm::Instruction *pI);
@@ -86,6 +88,8 @@ private:
   hlsl::DxilModule* m_DM;
   std::uint32_t m_uVReg;
   std::unique_ptr<llvm::ModuleSlotTracker> m_MST;
+  int m_StartInstruction = 0;
+
   void Init(llvm::Module &M) {
     m_DM = &M.GetOrCreateDxilModule();
     m_uVReg = 0;
@@ -97,6 +101,10 @@ private:
   }
 };
 
+void DxilAnnotateWithVirtualRegister::applyOptions(llvm::PassOptions O) {
+  GetPassOptionInt(O, "startInstruction", &m_StartInstruction, 0);
+}
+
 char DxilAnnotateWithVirtualRegister::ID = 0;
 
 bool DxilAnnotateWithVirtualRegister::runOnModule(llvm::Module &M) {
@@ -111,37 +119,55 @@ bool DxilAnnotateWithVirtualRegister::runOnModule(llvm::Module &M) {
     m_DM->SetValidatorVersion(1, 4);
   }
 
-  std::uint32_t InstNum = 0;
-  auto blocks = PIXPassHelpers::GetAllBlocks(*m_DM);
-  for(auto * block : blocks) {
-    for (llvm::Instruction& I : block->getInstList()) {
-      if (!llvm::isa<llvm::DbgDeclareInst>(&I)) {
-        pix_dxil::PixDxilInstNum::AddMD(M.getContext(), &I, InstNum++);
+  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;
+    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;
+        }
       }
     }
   }
 
   if (OSOverride != nullptr) {
+    // Print a set of strings of the exemplary form "InstructionCount: <n> <fnName>"
     *OSOverride << "\nInstructionCount:" << InstNum << "\n";
-  }
-
-  if (OSOverride != nullptr) {
-    *OSOverride << "\nEnd - instruction ID to line\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";
+    }
 
-  if (OSOverride != nullptr) {
     *OSOverride << "\nBegin - dxil values to virtual register mapping\n";
   }
 
-  for (auto* block : blocks) {
-    for (llvm::Instruction& I : block->getInstList()) {
-      AnnotateValues(&I);
+  for (auto * F : instrumentableFunctions) {
+    for (auto &block : F->getBasicBlockList()) {
+      for (llvm::Instruction &I : block.getInstList()) {
+        AnnotateValues(&I);
+      }
     }
   }
 
-  for (auto* block : blocks) {
-    for (llvm::Instruction& I : block->getInstList()) {
-      AnnotateStore(&I);
+  for (auto * F : instrumentableFunctions) {
+    for (auto &block : F->getBasicBlockList()) {
+      for (llvm::Instruction &I : block.getInstList()) {
+        AnnotateStore(&I);
+      }
     }
   }
 

lib/DxilPIXPasses/DxilDebugInstrumentation.cpp

@@ -955,14 +955,16 @@ bool DxilDebugInstrumentation::runOnModule(Module &M) {
 
   bool modified = false;
   if (shaderKind == DXIL::ShaderKind::Library) {
-    for (llvm::Function& F : M.functions()) {
-      modified = modified | RunOnFunction(M, DM, &F); 
-      return modified;
+    auto instrumentableFunctions =
+        PIXPassHelpers::GetAllInstrumentableFunctions(DM);
+    for (auto *F : instrumentableFunctions) {
+      if (RunOnFunction(M, DM, F)) {
+        modified = true;
+      }
     }
-  }
-  else {
+  } else {
     llvm::Function *entryFunction = PIXPassHelpers::GetEntryFunction(DM);
-    modified = RunOnFunction(M, DM, entryFunction);
+    modified = RunOnFunction(M, DM, entryFunction);  
   }
   return modified;
 }

lib/DxilPIXPasses/DxilShaderAccessTracking.cpp

@@ -894,6 +894,10 @@ bool DxilShaderAccessTracking::runOnModule(Module &M) {
       }
     }
   } else {
+
+    auto instrumentableFunctions =
+        PIXPassHelpers::GetAllInstrumentableFunctions(DM);
+
     if (DM.m_ShaderFlags.GetForceEarlyDepthStencil()) {
       if (OSOverride != nullptr) {
         formatted_raw_ostream FOS(*OSOverride);
@@ -901,44 +905,38 @@ bool DxilShaderAccessTracking::runOnModule(Module &M) {
       }
     }
     int uavRegId = 0;
-    for (llvm::Function &F : M.functions()) {
-      if (!F.getBasicBlockList().empty()) {
-
-        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;
+    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;
+      }
 
-        IRBuilder<> Builder(F.getEntryBlock().getFirstInsertionPt());
-
-        m_FunctionToUAVHandle[&F] = PIXPassHelpers::CreateUAV(
-            DM, Builder, uavRegId++, "PIX_CountUAV_Handle");
-        OP *HlslOP = DM.GetOP();
-        for (int accessStyle = static_cast<int>(ResourceAccessStyle::None);
-             accessStyle < static_cast<int>(ResourceAccessStyle::EndOfEnum);
-             ++accessStyle) {
-          ResourceAccessStyle style =
-              static_cast<ResourceAccessStyle>(accessStyle);
-          m_FunctionToEncodedAccess[&F][style] = HlslOP->GetU32Const(
-              EncodeShaderModel(shaderKind) | EncodeAccess(style));
-        }
+      IRBuilder<> Builder(F->getEntryBlock().getFirstInsertionPt());
+
+      m_FunctionToUAVHandle[F] = PIXPassHelpers::CreateUAV(
+          DM, Builder, uavRegId++, "PIX_CountUAV_Handle");
+      OP *HlslOP = DM.GetOP();
+      for (int accessStyle = static_cast<int>(ResourceAccessStyle::None);
+           accessStyle < static_cast<int>(ResourceAccessStyle::EndOfEnum);
+           ++accessStyle) {
+        ResourceAccessStyle style =
+            static_cast<ResourceAccessStyle>(accessStyle);
+        m_FunctionToEncodedAccess[F][style] = HlslOP->GetU32Const(
+            EncodeShaderModel(shaderKind) | EncodeAccess(style));
       }
     }
     DM.ReEmitDxilResources();
 
     for (llvm::Function &F : M.functions()) {
-      // Only used DXIL intrinsics:
-      if (!F.isDeclaration() || F.isIntrinsic() || F.use_empty() ||
-          !OP::IsDxilOpFunc(&F))
+      if (!F.isDeclaration() || F.isIntrinsic() || !OP::IsDxilOpFunc(&F))
         continue;
-
       // Gather handle parameter indices, if any
       FunctionType *fnTy =
           cast<FunctionType>(F.getType()->getPointerElementType());

lib/DxilPIXPasses/PixPassHelpers.cpp

@@ -224,6 +224,22 @@ llvm::Function* GetEntryFunction(hlsl::DxilModule& DM) {
     return DM.GetPatchConstantFunction();
 }
 
+std::vector<llvm::Function *>
+GetAllInstrumentableFunctions(hlsl::DxilModule &DM) {
+
+  std::vector<llvm::Function *> ret;
+
+  for (llvm::Function &F : DM.GetModule()->functions()) {
+    if (F.isDeclaration() || F.isIntrinsic() || hlsl::OP::IsDxilOpFunc(&F))
+      continue;
+    if (F.getBasicBlockList().empty())
+      continue;
+    ret.push_back(&F);
+  }
+
+  return ret;
+}
+
 std::vector<llvm::BasicBlock*> GetAllBlocks(hlsl::DxilModule& DM) {
     std::vector<llvm::BasicBlock*> ret;
     auto entryPoints = DM.GetExportedFunctions();

lib/DxilPIXPasses/PixPassHelpers.h

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

tools/clang/unittests/HLSL/PixTest.cpp

@@ -221,6 +221,8 @@ public:
   TEST_METHOD(PixStructAnnotation_ResourceAsMember)
   TEST_METHOD(PixStructAnnotation_WheresMyDbgValue)
 
+  TEST_METHOD(VirtualRegisters_InstructionCounts)
+
   dxc::DxcDllSupport m_dllSupport;
   VersionSupportInfo m_ver;
 
@@ -717,9 +719,10 @@ public:
   {
     CComPtr<IDxcBlob> blob;
     std::vector<ValueLocation> valueLocations;
+    std::vector<std::string> lines;
   };
 
-  PassOutput RunAnnotationPasses(IDxcBlob * dxil)
+  PassOutput RunAnnotationPasses(IDxcBlob * dxil, int startingLineNumber = 0)
   {
     CComPtr<IDxcOptimizer> pOptimizer;
     VERIFY_SUCCEEDED(
@@ -727,7 +730,10 @@ public:
     std::vector<LPCWSTR> Options;
     Options.push_back(L"-opt-mod-passes");
     Options.push_back(L"-dxil-dbg-value-to-dbg-declare");
-    Options.push_back(L"-dxil-annotate-with-virtual-regs");
+    std::wstring annotationCommandLine =
+        L"-dxil-annotate-with-virtual-regs,startInstruction=" +
+        std::to_wstring(startingLineNumber);
+    Options.push_back(annotationCommandLine.c_str());
 
     CComPtr<IDxcBlob> pOptimizedModule;
     CComPtr<IDxcBlobEncoding> pText;
@@ -772,7 +778,7 @@ public:
       }
     }
 
-    return { std::move(pOptimizedModule), std::move(valueLocations) };
+    return { std::move(pOptimizedModule), std::move(valueLocations), std::move(lines) };
   }
 
   std::wstring Disassemble(IDxcBlob * pProgram)
@@ -3241,5 +3247,183 @@ void main()
     }
 }
 
+TEST_F(PixTest, VirtualRegisters_InstructionCounts) {
+  if (m_ver.SkipDxilVersion(1, 5))
+    return;
+
+  for (auto choice : OptimizationChoices) {
+    auto optimization = choice.Flag;
+    const char *hlsl = R"(
+
+RaytracingAccelerationStructure Scene : register(t0, space0);
+RWTexture2D<float4> RenderTarget : register(u0);
+
+struct SceneConstantBuffer
+{
+    float4x4 projectionToWorld;
+    float4 cameraPosition;
+    float4 lightPosition;
+    float4 lightAmbientColor;
+    float4 lightDiffuseColor;
+};
+
+ConstantBuffer<SceneConstantBuffer> g_sceneCB : register(b0);
+
+struct RayPayload
+{
+    float4 color;
+};
+
+inline void GenerateCameraRay(uint2 index, out float3 origin, out float3 direction)
+{
+    float2 xy = index + 0.5f; // center in the middle of the pixel.
+    float2 screenPos = xy;// / DispatchRaysDimensions().xy * 2.0 - 1.0;
+
+    // Invert Y for DirectX-style coordinates.
+    screenPos.y = -screenPos.y;
+
+    // Unproject the pixel coordinate into a ray.
+    float4 world = /*mul(*/float4(screenPos, 0, 1)/*, g_sceneCB.projectionToWorld)*/;
+
+    //world.xyz /= world.w;
+    origin = world.xyz; //g_sceneCB.cameraPosition.xyz;
+    direction = float3(1,0,0);//normalize(world.xyz - origin);
+}
+
+void RaygenCommon()
+{
+    float3 rayDir;
+    float3 origin;
+    
+    // Generate a ray for a camera pixel corresponding to an index from the dispatched 2D grid.
+    GenerateCameraRay(DispatchRaysIndex().xy, origin, rayDir);
+
+    // Trace the ray.
+    // Set the ray's extents.
+    RayDesc ray;
+    ray.Origin = origin;
+    ray.Direction = rayDir;
+    // Set TMin to a non-zero small value to avoid aliasing issues due to floating - point errors.
+    // TMin should be kept small to prevent missing geometry at close contact areas.
+    ray.TMin = 0.001;
+    ray.TMax = 10000.0;
+    RayPayload payload = { float4(0, 0, 0, 0) };
+    TraceRay(Scene, RAY_FLAG_CULL_BACK_FACING_TRIANGLES, ~0, 0, 1, 0, ray, payload);
+
+    // Write the raytraced color to the output texture.
+   // RenderTarget[DispatchRaysIndex().xy] = payload.color;
+}
+
+[shader("raygeneration")]
+void Raygen0()
+{
+    RaygenCommon();
+}
+
+[shader("raygeneration")]
+void Raygen1()
+{
+    RaygenCommon();
+}
+
+typedef BuiltInTriangleIntersectionAttributes MyAttributes;
+
+[shader("closesthit")]
+void InnerClosestHitShader(inout RayPayload payload, in MyAttributes attr)
+{
+    payload.color = float4(0,1,0,0);
+}
+
+
+[shader("miss")]
+void MyMissShader(inout RayPayload payload)
+{
+    payload.color = float4(1, 0, 0, 0);
+})";
+
+    CComPtr<IDxcBlob> pBlob = Compile(hlsl, L"lib_6_6", {optimization});
+    CComPtr<IDxcBlob> pDxil = FindModule(DFCC_ShaderDebugInfoDXIL, pBlob);
+    auto outputLines = RunAnnotationPasses(pDxil).lines;
+
+    const char instructionRangeLabel[] = "InstructionRange:";
+
+    // The numbering pass should have counted  instructions for each "interesting" (to PIX)
+    // function and output its start and (end+1) instruction ordinal. End should always
+    // be a reasonable number of instructions (>10) and end should always be higher
+    // than start, and all four functions above should be represented.
+    int countOfInstructionRangeLines = 0;
+    for (auto const &line : outputLines) {
+      auto tokens = Tokenize(line, " ");
+      if (tokens.size() >= 4) {
+        if (tokens[0] == instructionRangeLabel) {
+          countOfInstructionRangeLines++;
+          int instructionStart = atoi(tokens[1].c_str());
+          int instructionEnd = atoi(tokens[2].c_str());
+          VERIFY_IS_TRUE(instructionEnd > 10);
+          VERIFY_IS_TRUE(instructionEnd > instructionStart);
+          auto found1 = tokens[3].find("Raygen0@@YAXXZ") != std::string::npos;
+          auto found2 = tokens[3].find("Raygen1@@YAXXZ") != std::string::npos;
+          auto foundClosest = tokens[3].find("InnerClosestHit") != std::string::npos;
+          auto foundMiss = tokens[3].find("MyMiss") != std::string::npos;
+          VERIFY_IS_TRUE(found1 || found2 || foundClosest || foundMiss);
+        }
+      }
+    }
+    VERIFY_ARE_EQUAL(4, countOfInstructionRangeLines);
+
+    // Non-library target:
+    const char *PixelShader= R"(
+    [RootSignature("")]
+    float main(float pos : A) : SV_Target {
+      float x = abs(pos);
+      float y = sin(pos);
+      float z = x + y;
+      return z;
+    }
+  )";
+    pBlob = Compile(PixelShader, L"ps_6_6", {optimization});
+    pDxil = FindModule(DFCC_ShaderDebugInfoDXIL, pBlob);
+    outputLines = RunAnnotationPasses(pDxil).lines;
+
+    countOfInstructionRangeLines = 0;
+    for (auto const &line : outputLines) {
+      auto tokens = Tokenize(line, " ");
+      if (tokens.size() >= 4) {
+        if (tokens[0] == instructionRangeLabel) {
+          countOfInstructionRangeLines++;
+          int instructionStart = atoi(tokens[1].c_str());
+          int instructionEnd = atoi(tokens[2].c_str());
+          VERIFY_IS_TRUE(instructionStart == 0);
+          VERIFY_IS_TRUE(instructionEnd > 10);
+          VERIFY_IS_TRUE(instructionEnd > instructionStart);
+          auto foundMain = tokens[3].find("main") != std::string::npos;
+          VERIFY_IS_TRUE(foundMain);
+        }
+      }
+    }
+    VERIFY_ARE_EQUAL(1, countOfInstructionRangeLines);
+
+    // Now check that the initial value parameter works:
+    const int startingInstructionOrdinal = 1234;
+    outputLines = RunAnnotationPasses(pDxil, startingInstructionOrdinal).lines;
+
+    countOfInstructionRangeLines = 0;
+    for (auto const &line : outputLines) {
+      auto tokens = Tokenize(line, " ");
+      if (tokens.size() >= 4) {
+        if (tokens[0] == instructionRangeLabel) {
+          countOfInstructionRangeLines++;
+          int instructionStart = atoi(tokens[1].c_str());
+          int instructionEnd = atoi(tokens[2].c_str());
+          VERIFY_IS_TRUE(instructionStart == startingInstructionOrdinal);
+          VERIFY_IS_TRUE(instructionEnd > instructionStart);
+          auto foundMain = tokens[3].find("main") != std::string::npos;
+          VERIFY_IS_TRUE(foundMain);
+        }
+      }
+    }
+    VERIFY_ARE_EQUAL(1, countOfInstructionRangeLines);
+  }
+}
 
 #endif

utils/hct/hctdb.py

@@ -2135,7 +2135,8 @@ class db_dxil(object):
             {'n':'parameter0','t':'int','c':1},
             {'n':'parameter1','t':'int','c':1},
             {'n':'parameter2','t':'int','c':1}])
-        add_pass('dxil-annotate-with-virtual-regs', 'DxilAnnotateWithVirtualRegister', 'Annotates each instruction in the DXIL module with a virtual register number', [])
+        add_pass('dxil-annotate-with-virtual-regs', 'DxilAnnotateWithVirtualRegister', 'Annotates each instruction in the DXIL module with a virtual register number', [
+            {'n':'startInstruction','t':'int','c':1}])
         add_pass('dxil-dbg-value-to-dbg-declare', 'DxilDbgValueToDbgDeclare', 'Converts llvm.dbg.value uses to llvm.dbg.declare.', [])
         add_pass('hlsl-dxil-reduce-msaa-to-single', 'DxilReduceMSAAToSingleSample', 'HLSL DXIL Reduce all MSAA reads to single-sample reads', [])
         add_pass('hlsl-dxil-PIX-add-tid-to-as-payload', 'DxilPIXAddTidToAmplificationShaderPayload', 'HLSL DXIL Add flat thread id to payload from AS to MS', [