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

Author: jeffnn

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 @@ class DxilAnnotateWithVirtualRegister : public llvm::ModulePass {
   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 @@ class DxilAnnotateWithVirtualRegister : public llvm::ModulePass {
   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 @@ class DxilAnnotateWithVirtualRegister : public llvm::ModulePass {
   }
 };
 
+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,51 +894,49 @@ bool DxilShaderAccessTracking::runOnModule(Module &M) {
       }
     }
   } else {
+
+    auto instrumentableFunctions =
+        PIXPassHelpers::GetAllInstrumentableFunctions(DM);
+
     if (DM.m_ShaderFlags.GetForceEarlyDepthStencil()) {
       if (OSOverride != nullptr) {
         formatted_raw_ostream FOS(*OSOverride);
         FOS << "ShouldAssumeDsvAccess";
       }
     }
     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, ...);