Instrument alloca register write (#1831)

* instruments alloca register writes

lib/DxilPIXPasses/DxilAnnotateWithVirtualRegister.cpp

@@ -69,6 +69,9 @@ bool DxilAnnotateWithVirtualRegister::runOnModule(llvm::Module &M) {
     return false;
   }
 
+  if (OSOverride != nullptr) {
+    *OSOverride << "\nBegin - dxil values to virtual register mapping\n";
+  }
   for (llvm::Instruction &I : llvm::inst_range(m_DM->GetEntryFunction())) {
     AnnotateValues(&I);
   }
@@ -77,6 +80,10 @@ bool DxilAnnotateWithVirtualRegister::runOnModule(llvm::Module &M) {
     AnnotateStore(&I);
   }
 
+  if (OSOverride != nullptr) {
+    *OSOverride << "\nEnd - dxil values to virtual register mapping\n";
+  }
+
   m_DM = nullptr;
   return m_uVReg > 0;
 }
@@ -173,7 +180,7 @@ void DxilAnnotateWithVirtualRegister::AssignNewDxilRegister(llvm::Instruction *p
   if (OSOverride != nullptr) {
     static constexpr bool DontPrintType = false;
     pI->printAsOperand(*OSOverride, DontPrintType, m_DM->GetModule());
-    *OSOverride << " dxil " << m_uVReg;
+    *OSOverride << " dxil " << m_uVReg << "\n";
   }
   m_uVReg++;
 }
@@ -183,7 +190,7 @@ void DxilAnnotateWithVirtualRegister::AssignNewAllocaRegister(llvm::AllocaInst *
   if (OSOverride != nullptr) {
     static constexpr bool DontPrintType = false;
     pAlloca->printAsOperand(*OSOverride, DontPrintType, m_DM->GetModule());
-    *OSOverride << " alloca " << m_uVReg << " " << C;
+    *OSOverride << " alloca " << m_uVReg << " " << C << "\n";
   }
   m_uVReg += C;
 }

lib/DxilPIXPasses/DxilDebugInstrumentation.cpp

@@ -226,10 +226,12 @@ private:
   void addDebugEntryValue(BuilderContext &BC, Value * TheValue);
   void addInvocationStartMarker(BuilderContext &BC);
   void reserveDebugEntrySpace(BuilderContext &BC, uint32_t SpaceInDwords);
+  void addStoreStepDebugEntry(BuilderContext &BC, StoreInst *Inst);
   void addStepDebugEntry(BuilderContext &BC, Instruction *Inst);
+  void addStepDebugEntryValue(BuilderContext &BC, Value *V, std::uint32_t RegNum);
   uint32_t UAVDumpingGroundOffset();
   template<typename ReturnType>
-  void addStepEntryForType(DebugShaderModifierRecordType RecordType, BuilderContext &BC, Instruction *Inst);
+  void addStepEntryForType(DebugShaderModifierRecordType RecordType, BuilderContext &BC, Value *V, std::uint32_t RegNum);
 
 };
 
@@ -632,12 +634,7 @@ 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;
-  }
-
+void DxilDebugInstrumentation::addStepEntryForType(DebugShaderModifierRecordType RecordType, BuilderContext &BC, Value *V, std::uint32_t RegNum) {
   DebugShaderModifierRecordDXILStep<ReturnType> step = {};
   reserveDebugEntrySpace(BC, sizeof(step));
 
@@ -648,8 +645,29 @@ void DxilDebugInstrumentation::addStepEntryForType(DebugShaderModifierRecordType
   addDebugEntryValue(BC, BC.HlslOP->GetU32Const(RegNum));
 
   if (RecordType != DebugShaderModifierRecordTypeDXILStepVoid) {
-    addDebugEntryValue(BC, Inst);
+    addDebugEntryValue(BC, V);
+  }
+}
+
+void DxilDebugInstrumentation::addStoreStepDebugEntry(BuilderContext &BC, StoreInst *Inst) {
+  std::uint32_t RegBase;
+  std::uint32_t RegSize;
+  llvm::Value *RegIndex;
+  if (!pix_dxil::PixAllocaRegWrite::FromInst(Inst, &RegBase, &RegSize, &RegIndex)) {
+    return;
+  }
+
+  auto *RegIndexConst = llvm::dyn_cast<llvm::ConstantInt>(RegIndex);
+  if (RegIndexConst == nullptr) {
+    return;
+  }
+
+  if (RegIndexConst->getLimitedValue() >= RegSize) {
+    return;
   }
+
+  const std::uint32_t ActualReg = RegBase + RegIndexConst->getLimitedValue();
+  addStepDebugEntryValue(BC, Inst->getValueOperand(), ActualReg);
 }
 
 void DxilDebugInstrumentation::addStepDebugEntry(BuilderContext &BC, Instruction *Inst) {
@@ -657,29 +675,43 @@ void DxilDebugInstrumentation::addStepDebugEntry(BuilderContext &BC, Instruction
     return;
   }
 
-  Type::TypeID ID = Inst->getType()->getTypeID();
+  if (auto *St = llvm::dyn_cast<llvm::StoreInst>(Inst)) {
+    addStoreStepDebugEntry(BC, St);
+    return;
+  }
+
+  std::uint32_t RegNum;
+  if (!pix_dxil::PixDxilReg::FromInst(Inst, &RegNum)) {
+    return;
+  }
+
+  addStepDebugEntryValue(BC, Inst, RegNum);
+}
+
+void DxilDebugInstrumentation::addStepDebugEntryValue(BuilderContext &BC, Value *V, std::uint32_t RegNum) {
+  const Type::TypeID ID = V->getType()->getTypeID();
 
   switch (ID) {
   case Type::TypeID::StructTyID:
   case Type::TypeID::VoidTyID:
-    addStepEntryForType<void>(DebugShaderModifierRecordTypeDXILStepVoid, BC, Inst);
+    addStepEntryForType<void>(DebugShaderModifierRecordTypeDXILStepVoid, BC, V, RegNum);
     break;
   case Type::TypeID::FloatTyID:
-    addStepEntryForType<float>(DebugShaderModifierRecordTypeDXILStepFloat, BC, Inst);
+    addStepEntryForType<float>(DebugShaderModifierRecordTypeDXILStepFloat, BC, V, RegNum);
     break;
   case Type::TypeID::IntegerTyID:
-    if (Inst->getType()->getIntegerBitWidth() == 64) {
-      addStepEntryForType<uint64_t>(DebugShaderModifierRecordTypeDXILStepUint64, BC, Inst);
+    if (V->getType()->getIntegerBitWidth() == 64) {
+      addStepEntryForType<uint64_t>(DebugShaderModifierRecordTypeDXILStepUint64, BC, V, RegNum);
     }
     else {
-      addStepEntryForType<uint32_t>(DebugShaderModifierRecordTypeDXILStepUint32, BC, Inst);
+      addStepEntryForType<uint32_t>(DebugShaderModifierRecordTypeDXILStepUint32, BC, V, RegNum);
     }
     break;
   case Type::TypeID::DoubleTyID:
-    addStepEntryForType<double>(DebugShaderModifierRecordTypeDXILStepDouble, BC, Inst);
+    addStepEntryForType<double>(DebugShaderModifierRecordTypeDXILStepDouble, BC, V, RegNum);
     break;
   case Type::TypeID::HalfTyID:
-    addStepEntryForType<float>(DebugShaderModifierRecordTypeDXILStepFloat, BC, Inst);
+    addStepEntryForType<float>(DebugShaderModifierRecordTypeDXILStepFloat, BC, V, RegNum);
     break;
   case Type::TypeID::PointerTyID:
     // Skip pointer calculation instructions. They aren't particularly meaningful to the user (being a mere
@@ -745,10 +777,9 @@ bool DxilDebugInstrumentation::runOnModule(Module &M) {
 
   // Instrument original instructions:
   for (auto & Inst : AllInstructions) {
-    // Instrumentation goes after the instruction if it has a return value.
-    // Otherwise, the instruction might be a terminator so we HAVE to put the instrumentation before
-    if (Inst->getType()->getTypeID() != Type::TypeID::VoidTyID) {
-      // Has a return type, so can't be a terminator, so start inserting before the next instruction
+    // Instrumentation goes after the instruction if it is not a terminator. Otherwise,
+    // Instrumentation goes prior to the instruction.
+    if (!Inst->isTerminator()) {
       IRBuilder<> Builder(Inst->getNextNode());
       BuilderContext BC2{ BC.M, BC.DM, BC.Ctx, BC.HlslOP, Builder };
       addStepDebugEntry(BC2, Inst);

lib/DxilPIXPasses/DxilPIXVirtualRegisters.cpp

@@ -64,6 +64,28 @@ bool pix_dxil::PixDxilReg::FromInst(llvm::Instruction *pI, std::uint32_t *pRegNu
   return true;
 }
 
+static bool ParsePixAllocaReg(llvm::MDNode *MD, std::uint32_t *RegNum, std::uint32_t *Count) {
+  if (MD->getNumOperands() != 3) {
+    return false;
+  }
+
+  auto *mdID = llvm::mdconst::dyn_extract<llvm::ConstantInt>(MD->getOperand(0));
+  if (mdID == nullptr || mdID->getLimitedValue() != pix_dxil::PixAllocaReg::ID) {
+    return false;
+  }
+
+  auto *mdRegNum = llvm::mdconst::dyn_extract<llvm::ConstantInt>(MD->getOperand(1));
+  auto *mdCount = llvm::mdconst::dyn_extract<llvm::ConstantInt>(MD->getOperand(2));
+
+  if (mdRegNum == nullptr || mdCount == nullptr) {
+    return false;
+  }
+
+  *RegNum = mdRegNum->getLimitedValue();
+  *Count = mdCount->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(
@@ -81,3 +103,32 @@ void pix_dxil::PixAllocaRegWrite::AddMD(llvm::LLVMContext &Ctx, llvm::StoreInst
                                pAllocaReg,
                                MetadataForValue(Index) }));
 }
+
+bool pix_dxil::PixAllocaRegWrite::FromInst(llvm::StoreInst *pI, std::uint32_t *pRegBase, std::uint32_t *pRegSize, llvm::Value **pIndex) {
+  *pRegBase = 0;
+  *pRegSize = 0;
+  *pIndex = nullptr;
+
+  auto *mdNodes = pI->getMetadata(MDName);
+  if (mdNodes == nullptr || mdNodes->getNumOperands() != 3) {
+    return false;
+  }
+
+  auto *mdID = llvm::mdconst::dyn_extract<llvm::ConstantInt>(mdNodes->getOperand(0));
+  if (mdID == nullptr || mdID->getLimitedValue() != ID) {
+    return false;
+  }
+
+  auto *mdAllocaReg = llvm::dyn_cast<llvm::MDNode>(mdNodes->getOperand(1));
+  if (mdAllocaReg == nullptr || !ParsePixAllocaReg(mdAllocaReg, pRegBase, pRegSize)) {
+    return false;
+  }
+
+  auto *mdIndex = llvm::dyn_cast<llvm::ValueAsMetadata>(mdNodes->getOperand(2));
+  if (mdIndex == nullptr) {
+    return false;
+  }
+  *pIndex = mdIndex->getValue();
+
+  return true;
+}

lib/DxilPIXPasses/DxilPIXVirtualRegisters.h

@@ -41,5 +41,6 @@ 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);
+bool FromInst(llvm::StoreInst *pI, std::uint32_t *pRegBase, std::uint32_t *pRegSize, llvm::Value **pIndex);
 }  // namespace PixAllocaRegWrite
 }  // namespace pix_dxil