Adds a PIX pass for assigning register IDs to dxil values. (#1827)

* Adds a PIX pass for assigning register IDs to dxil values.

include/dxc/DxilPIXPasses/DxilPIXPasses.h

@@ -16,6 +16,7 @@ class ModulePass;
 class PassRegistry;
 
 ModulePass *createDxilAddPixelHitInstrumentationPass();
+ModulePass *createDxilAnnotateWithVirtualRegisterPass();
 ModulePass *createDxilOutputColorBecomesConstantPass();
 ModulePass *createDxilRemoveDiscardsPass();
 ModulePass *createDxilReduceMSAAToSingleSamplePass();
@@ -24,6 +25,7 @@ ModulePass *createDxilDebugInstrumentationPass();
 ModulePass *createDxilShaderAccessTrackingPass();
 
 void initializeDxilAddPixelHitInstrumentationPass(llvm::PassRegistry&);
+void initializeDxilAnnotateWithVirtualRegisterPass(llvm::PassRegistry&);
 void initializeDxilOutputColorBecomesConstantPass(llvm::PassRegistry&);
 void initializeDxilRemoveDiscardsPass(llvm::PassRegistry&);
 void initializeDxilReduceMSAAToSingleSamplePass(llvm::PassRegistry&);

lib/DxilPIXPasses/CMakeLists.txt

@@ -2,6 +2,7 @@
 # This file is distributed under the University of Illinois Open Source License. See LICENSE.TXT for details.
 add_llvm_library(LLVMDxilPIXPasses
   DxilAddPixelHitInstrumentation.cpp
+  DxilAnnotateWithVirtualRegister.cpp
   DxilDebugInstrumentation.cpp
   DxilForceEarlyZ.cpp
   DxilOutputColorBecomesConstant.cpp
@@ -9,6 +10,7 @@ add_llvm_library(LLVMDxilPIXPasses
   DxilReduceMSAAToSingleSample.cpp
   DxilShaderAccessTracking.cpp
   DxilPIXPasses.cpp
+  DxilPIXVirtualRegisters.cpp
 
 
   ADDITIONAL_HEADER_DIRS

lib/DxilPIXPasses/DxilAnnotateWithVirtualRegister.cpp

@@ -0,0 +1,198 @@
+///////////////////////////////////////////////////////////////////////////////
+//                                                                           //
+// DxilAnnotateWithVirtualRegister.cpp                                       //
+// Copyright (C) Microsoft Corporation. All rights reserved.                 //
+// This file is distributed under the University of Illinois Open Source     //
+// License. See LICENSE.TXT for details.                                     //
+//                                                                           //
+// Annotates the llvm instructions with a virtual register number to be used //
+// during PIX debugging.                                                     //
+//                                                                           //
+///////////////////////////////////////////////////////////////////////////////
+
+#include "dxc/DxilPIXPasses/DxilPIXPasses.h"
+
+#include <memory>
+
+#include "dxc/DXIL/DxilModule.h"
+#include "dxc/Support/Global.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/IR/Constant.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/Type.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include "DxilPIXVirtualRegisters.h"
+
+#define DEBUG_TYPE "dxil-annotate-with-virtual-regs"
+
+namespace {
+using namespace pix_dxil;
+
+class DxilAnnotateWithVirtualRegister : public llvm::ModulePass {
+public:
+  static char ID;
+  DxilAnnotateWithVirtualRegister() : llvm::ModulePass(ID) {}
+
+  bool runOnModule(llvm::Module &M) override;
+
+private:
+  void AnnotateValues(llvm::Instruction *pI);
+  void AnnotateStore(llvm::Instruction *pI);
+  bool IsAllocaRegisterWrite(llvm::Value *V, llvm::AllocaInst **pAI, llvm::Value **pIdx);
+  void AnnotateAlloca(llvm::AllocaInst *pAlloca);
+  void AnnotateGeneric(llvm::Instruction *pI);
+  void AssignNewDxilRegister(llvm::Instruction *pI);
+  void AssignNewAllocaRegister(llvm::AllocaInst *pAlloca, std::uint32_t C);
+
+  hlsl::DxilModule *m_DM;
+  std::uint32_t m_uVReg;
+  void Init(llvm::Module &M) {
+    m_DM = &M.GetOrCreateDxilModule();
+    m_uVReg = 0;
+  }
+};
+
+char DxilAnnotateWithVirtualRegister::ID = 0;
+
+bool DxilAnnotateWithVirtualRegister::runOnModule(llvm::Module &M) {
+  Init(M);
+  if (m_DM == nullptr) {
+    return false;
+  }
+
+  for (llvm::Instruction &I : llvm::inst_range(m_DM->GetEntryFunction())) {
+    AnnotateValues(&I);
+  }
+
+  for (llvm::Instruction &I : llvm::inst_range(m_DM->GetEntryFunction())) {
+    AnnotateStore(&I);
+  }
+
+  m_DM = nullptr;
+  return m_uVReg > 0;
+}
+
+void DxilAnnotateWithVirtualRegister::AnnotateValues(llvm::Instruction *pI) {
+  if (auto *pAlloca = llvm::dyn_cast<llvm::AllocaInst>(pI)) {
+    AnnotateAlloca(pAlloca);
+  } else if (!pI->getType()->isVoidTy()) {
+    AnnotateGeneric(pI);
+  }
+}
+
+void DxilAnnotateWithVirtualRegister::AnnotateStore(llvm::Instruction *pI) {
+  auto *pSt = llvm::dyn_cast<llvm::StoreInst>(pI);
+  if (pSt == nullptr) {
+    return;
+  }
+
+  llvm::AllocaInst *Alloca;
+  llvm::Value *Index;
+  if (!IsAllocaRegisterWrite(pSt->getPointerOperand(), &Alloca, &Index)) {
+    return;
+  }
+
+  llvm::MDNode *AllocaReg = Alloca->getMetadata(PixAllocaReg::MDName);
+  if (AllocaReg == nullptr) {
+    return;
+  }
+
+  PixAllocaRegWrite::AddMD(m_DM->GetCtx(), pSt, AllocaReg, Index);
+}
+
+bool DxilAnnotateWithVirtualRegister::IsAllocaRegisterWrite(llvm::Value *V, llvm::AllocaInst **pAI, llvm::Value **pIdx) {
+  llvm::IRBuilder<> B(m_DM->GetCtx());
+
+  *pAI = nullptr;
+  *pIdx = nullptr;
+
+  if (auto *pGEP = llvm::dyn_cast<llvm::GetElementPtrInst>(V)) {
+    auto *Alloca = llvm::dyn_cast<llvm::AllocaInst>(pGEP->getPointerOperand());
+    if (Alloca == nullptr) {
+      return false;
+    }
+
+    llvm::SmallVector<llvm::Value *, 2> Indices(pGEP->idx_begin(), pGEP->idx_end());
+    if (Indices.size() != 2) {
+      return false;
+    }
+    auto *pIdx0 = llvm::dyn_cast<llvm::ConstantInt>(Indices[0]);
+
+    if (pIdx0 == nullptr || pIdx0->getLimitedValue() != 0) {
+      return false;
+    }
+    
+    *pAI = Alloca;
+    *pIdx = Indices[1];
+    return true;
+  }
+
+  if (auto *pAlloca = llvm::dyn_cast<llvm::AllocaInst>(V)) {
+    llvm::Type *pAllocaTy = pAlloca->getType()->getElementType();
+    if (!pAllocaTy->isFloatTy() && !pAllocaTy->isIntegerTy()) {
+      return false;
+    }
+
+    *pAI = pAlloca;
+    *pIdx = B.getInt32(0);
+    return true;
+  }
+
+  return false;
+}
+
+void DxilAnnotateWithVirtualRegister::AnnotateAlloca(llvm::AllocaInst *pAlloca) {
+  llvm::Type *pAllocaTy = pAlloca->getType()->getElementType();
+  if (pAllocaTy->isFloatTy() || pAllocaTy->isIntegerTy()) {
+    AssignNewAllocaRegister(pAlloca, 1);
+  } else if (auto *AT = llvm::dyn_cast<llvm::ArrayType>(pAllocaTy)) {
+    AssignNewAllocaRegister(pAlloca, AT->getNumElements());
+  } else {
+    DXASSERT_ARGS(false, "Unhandled alloca kind: %d", pAllocaTy->getTypeID());
+  }
+}
+
+void DxilAnnotateWithVirtualRegister::AnnotateGeneric(llvm::Instruction *pI) {
+  if (!pI->getType()->isFloatTy() && !pI->getType()->isIntegerTy()) {
+    return;
+  }
+  AssignNewDxilRegister(pI);
+}
+
+void DxilAnnotateWithVirtualRegister::AssignNewDxilRegister(llvm::Instruction *pI) {
+  PixDxilReg::AddMD(m_DM->GetCtx(), pI, m_uVReg);
+  if (OSOverride != nullptr) {
+    static constexpr bool DontPrintType = false;
+    pI->printAsOperand(*OSOverride, DontPrintType, m_DM->GetModule());
+    *OSOverride << " dxil " << m_uVReg;
+  }
+  m_uVReg++;
+}
+
+void DxilAnnotateWithVirtualRegister::AssignNewAllocaRegister(llvm::AllocaInst *pAlloca, std::uint32_t C) {
+  PixAllocaReg::AddMD(m_DM->GetCtx(), pAlloca, m_uVReg, C);
+  if (OSOverride != nullptr) {
+    static constexpr bool DontPrintType = false;
+    pAlloca->printAsOperand(*OSOverride, DontPrintType, m_DM->GetModule());
+    *OSOverride << " alloca " << m_uVReg << " " << C;
+  }
+  m_uVReg += C;
+}
+}
+
+using namespace llvm;
+
+INITIALIZE_PASS(DxilAnnotateWithVirtualRegister, DEBUG_TYPE, "Annotates each instruction in the DXIL module with a virtual register number", false, false)
+
+ModulePass *llvm::createDxilAnnotateWithVirtualRegisterPass() {
+  return new DxilAnnotateWithVirtualRegister();
+}

lib/DxilPIXPasses/DxilDebugInstrumentation.cpp

@@ -21,6 +21,8 @@
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/ADT/STLExtras.h"
 
+#include "DxilPIXVirtualRegisters.h"
+
 using namespace llvm;
 using namespace hlsl;
 
@@ -195,8 +197,6 @@ private:
 
   std::map<uint32_t, Value *> m_IncrementInstructionBySize;
 
-  unsigned int m_InstructionIndex = 0;
-
   struct BuilderContext {
     Module &M;
     DxilModule &DM;
@@ -633,6 +633,10 @@ void DxilDebugInstrumentation::addInvocationStartMarker(BuilderContext &BC) {
 
 template<typename ReturnType>
 void DxilDebugInstrumentation::addStepEntryForType(DebugShaderModifierRecordType RecordType, BuilderContext &BC, Instruction *Inst) {
+  std::uint32_t RegNum;
+  if (!pix_dxil::PixDxilReg::FromInst(Inst, &RegNum)) {
+    return;
+  }
 
   DebugShaderModifierRecordDXILStep<ReturnType> step = {};
   reserveDebugEntrySpace(BC, sizeof(step));
@@ -641,7 +645,7 @@ void DxilDebugInstrumentation::addStepEntryForType(DebugShaderModifierRecordType
   step.Header.Details.Type = static_cast<uint8_t>(RecordType);
   addDebugEntryValue(BC, BC.HlslOP->GetU32Const(step.Header.u32Header));
   addDebugEntryValue(BC, m_InvocationId);
-  addDebugEntryValue(BC, BC.HlslOP->GetU32Const(m_InstructionIndex++));
+  addDebugEntryValue(BC, BC.HlslOP->GetU32Const(RegNum));
 
   if (RecordType != DebugShaderModifierRecordTypeDXILStepVoid) {
     addDebugEntryValue(BC, Inst);

lib/DxilPIXPasses/DxilPIXPasses.cpp

@@ -31,6 +31,7 @@ HRESULT SetupRegistryPassForPIX() {
     /* <py::lines('INIT-PASSES')>hctdb_instrhelp.get_init_passes(set(["pix"]))</py>*/
     // INIT-PASSES:BEGIN
     initializeDxilAddPixelHitInstrumentationPass(Registry);
+    initializeDxilAnnotateWithVirtualRegisterPass(Registry);
     initializeDxilDebugInstrumentationPass(Registry);
     initializeDxilForceEarlyZPass(Registry);
     initializeDxilOutputColorBecomesConstantPass(Registry);

lib/DxilPIXPasses/DxilPIXVirtualRegisters.cpp

@@ -0,0 +1,83 @@
+///////////////////////////////////////////////////////////////////////////////
+//                                                                           //
+// DxilPIXVirtualRegisters.cpp                                               //
+// Copyright (C) Microsoft Corporation. All rights reserved.                 //
+// This file is distributed under the University of Illinois Open Source     //
+// License. See LICENSE.TXT for details.                                     //
+//                                                                           //
+// Defines functions for dealing with the virtual register annotations in    //
+// DXIL instructions.                                                        //
+//                                                                           //
+///////////////////////////////////////////////////////////////////////////////
+
+#include "DxilPIXVirtualRegisters.h"
+
+#include "dxc/Support/Global.h"
+#include "llvm/IR/Constant.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/IR/Type.h"
+
+static llvm::Metadata *MetadataForValue(llvm::Value *V) {
+  if (auto *C = llvm::dyn_cast<llvm::Constant>(V)) {
+    return llvm::ConstantAsMetadata::get(C);
+  }
+  return llvm::ValueAsMetadata::get(V);
+}
+
+void pix_dxil::PixDxilReg::AddMD(llvm::LLVMContext &Ctx, llvm::Instruction *pI, std::uint32_t RegNum) {
+  llvm::IRBuilder<> B(Ctx);
+  pI->setMetadata(
+      llvm::StringRef(MDName),
+      llvm::MDNode::get(Ctx, { llvm::ConstantAsMetadata::get(B.getInt32(ID)),
+                               llvm::ConstantAsMetadata::get(B.getInt32(RegNum)) }));
+}
+
+bool pix_dxil::PixDxilReg::FromInst(llvm::Instruction *pI, std::uint32_t *pRegNum) {
+  *pRegNum = 0;
+
+  auto *mdNodes = pI->getMetadata(MDName);
+
+  if (mdNodes == nullptr) {
+    return false;
+  }
+
+  if (mdNodes->getNumOperands() != 2) {
+    return false;
+  }
+
+  auto *mdID = llvm::mdconst::dyn_extract<llvm::ConstantInt>(mdNodes->getOperand(0));
+  if (mdID == nullptr || mdID->getLimitedValue() != ID) {
+    return false;
+  }
+
+  auto *mdRegNum = llvm::mdconst::dyn_extract<llvm::ConstantInt>(mdNodes->getOperand(1));
+  if (mdRegNum == nullptr) {
+    return false;
+  }
+
+  *pRegNum = mdRegNum->getLimitedValue();
+  return true;
+}
+
+void pix_dxil::PixAllocaReg::AddMD(llvm::LLVMContext &Ctx, llvm::AllocaInst *pAlloca, std::uint32_t RegNum, std::uint32_t Count) {
+  llvm::IRBuilder<> B(Ctx);
+  pAlloca->setMetadata(
+      llvm::StringRef(MDName),
+      llvm::MDNode::get(Ctx, { llvm::ConstantAsMetadata::get(B.getInt32(ID)),
+                               llvm::ConstantAsMetadata::get(B.getInt32(RegNum)),
+                               llvm::ConstantAsMetadata::get(B.getInt32(Count)) }));
+}
+
+void pix_dxil::PixAllocaRegWrite::AddMD(llvm::LLVMContext &Ctx, llvm::StoreInst *pSt, llvm::MDNode *pAllocaReg, llvm::Value *Index) {
+  llvm::IRBuilder<> B(Ctx);
+  pSt->setMetadata(
+      llvm::StringRef(MDName),
+      llvm::MDNode::get(Ctx, { llvm::ConstantAsMetadata::get(B.getInt32(ID)),
+                               pAllocaReg,
+                               MetadataForValue(Index) }));
+}

lib/DxilPIXPasses/DxilPIXVirtualRegisters.h

@@ -0,0 +1,45 @@
+///////////////////////////////////////////////////////////////////////////////
+//                                                                           //
+// DxilPIXVirtualRegisters.cpp                                               //
+// Copyright (C) Microsoft Corporation. All rights reserved.                 //
+// This file is distributed under the University of Illinois Open Source     //
+// License. See LICENSE.TXT for details.                                     //
+//                                                                           //
+// Declares functions for dealing with the virtual register annotations in   //
+// DXIL instructions.                                                        //
+//                                                                           //
+///////////////////////////////////////////////////////////////////////////////
+
+#include <cstdint>
+
+namespace llvm {
+class AllocaInst;
+class Instruction;
+class LLVMContext;
+class MDNode;
+class StoreInst;
+class Value;
+}  // namespace llvm
+
+namespace pix_dxil {
+namespace PixDxilReg {
+static constexpr char MDName[] = "pix-dxil-reg";
+static constexpr uint32_t ID = 0;
+
+void AddMD(llvm::LLVMContext &Ctx, llvm::Instruction *pI, std::uint32_t RegNum);
+bool FromInst(llvm::Instruction *pI, std::uint32_t *pRegNum);
+}  // namespace PixDxilReg
+
+namespace PixAllocaReg {
+static constexpr char MDName[] = "pix-alloca-reg";
+static constexpr uint32_t ID = 1;
+
+void AddMD(llvm::LLVMContext &Ctx, llvm::AllocaInst *pAlloca, std::uint32_t RegNum, std::uint32_t Count);
+}  // namespace PixAllocaReg
+
+namespace PixAllocaRegWrite {
+static constexpr char MDName[] = "pix-alloca-reg-write";
+static constexpr uint32_t ID = 2;
+void AddMD(llvm::LLVMContext &Ctx, llvm::StoreInst *pSt, llvm::MDNode *pAllocaReg, llvm::Value *Index);
+}  // namespace PixAllocaRegWrite
+}  // namespace pix_dxil

utils/hct/hctdb.py

@@ -1538,6 +1538,7 @@ 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('hlsl-dxil-reduce-msaa-to-single', 'DxilReduceMSAAToSingleSample', 'HLSL DXIL Reduce all MSAA reads to single-sample reads', [])
 
         category_lib="dxil_gen"