[DxilPIXPasses] Assigns unique ID to each instruction. (#1888)

This PR modifies the DxilAnnotateWithVirtualRegister pass to also annotate
each instruction in the DXIL module with a unique ID. This ID can then be
used by trace debuggers to uniquely identify an instruction in DXIL module.

This PR also modifies the DxilDebugInstrumentation pass to output this
unique ID on each entry in the debug trace. The debug trace is also expanded
to include the Value Ordinal (i.e., the virtual register modified by the
instruction).

Finally, this PR bundles a small bugfix to the dxassembler. The bug prevented
the DxcAssembler to add the debug DXIL to the container it assembles.

lib/DxilPIXPasses/DxilAnnotateWithVirtualRegister.cpp

@@ -72,6 +72,12 @@ bool DxilAnnotateWithVirtualRegister::runOnModule(llvm::Module &M) {
   if (OSOverride != nullptr) {
     *OSOverride << "\nBegin - dxil values to virtual register mapping\n";
   }
+
+  std::uint32_t InstNum = 0;
+  for (llvm::Instruction &I : llvm::inst_range(m_DM->GetEntryFunction())) {
+    pix_dxil::PixDxilInstNum::AddMD(M.getContext(), &I, ++InstNum);
+  }
+
   for (llvm::Instruction &I : llvm::inst_range(m_DM->GetEntryFunction())) {
     AnnotateValues(&I);
   }

lib/DxilPIXPasses/DxilDebugInstrumentation.cpp

@@ -127,6 +127,13 @@ struct DebugShaderModifierRecordDXILStepBase {
 template< typename ReturnType >
 struct DebugShaderModifierRecordDXILStep : public DebugShaderModifierRecordDXILStepBase {
   ReturnType ReturnValue;
+  union {
+    struct {
+      uint32_t ValueOrdinalBase : 16;
+      uint32_t ValueOrdinalIndex : 16;
+    } Details;
+    uint32_t u32ValueOrdinal;
+  } ValueOrdinal;
 };
 
 template< >
@@ -228,10 +235,10 @@ private:
   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);
+  void addStepDebugEntryValue(BuilderContext &BC, std::uint32_t InstNum, Value *V, std::uint32_t ValueOrdinal, Value *ValueOrdinalIndex);
   uint32_t UAVDumpingGroundOffset();
   template<typename ReturnType>
-  void addStepEntryForType(DebugShaderModifierRecordType RecordType, BuilderContext &BC, Value *V, std::uint32_t RegNum);
+  void addStepEntryForType(DebugShaderModifierRecordType RecordType, BuilderContext &BC, std::uint32_t InstNum, Value *V, std::uint32_t ValueOrdinal, Value *ValueOrdinalIndex);
 
 };
 
@@ -634,7 +641,7 @@ void DxilDebugInstrumentation::addInvocationStartMarker(BuilderContext &BC) {
 }
 
 template<typename ReturnType>
-void DxilDebugInstrumentation::addStepEntryForType(DebugShaderModifierRecordType RecordType, BuilderContext &BC, Value *V, std::uint32_t RegNum) {
+void DxilDebugInstrumentation::addStepEntryForType(DebugShaderModifierRecordType RecordType, BuilderContext &BC, std::uint32_t InstNum, Value *V, std::uint32_t ValueOrdinal, Value *ValueOrdinalIndex) {
   DebugShaderModifierRecordDXILStep<ReturnType> step = {};
   reserveDebugEntrySpace(BC, sizeof(step));
 
@@ -642,32 +649,34 @@ 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(RegNum));
+  addDebugEntryValue(BC, BC.HlslOP->GetU32Const(InstNum));
 
   if (RecordType != DebugShaderModifierRecordTypeDXILStepVoid) {
     addDebugEntryValue(BC, V);
+
+    IRBuilder<> &B = BC.Builder;
+
+    Value *VO = BC.HlslOP->GetU32Const(ValueOrdinal << 16);
+    Value *VOI = B.CreateAnd(ValueOrdinalIndex, BC.HlslOP->GetU32Const(0xFFFF), "ValueOrdinalIndex");
+    Value *EncodedValueOrdinalAndIndex = BC.Builder.CreateOr(VO, VOI, "ValueOrdinal");
+    addDebugEntryValue(BC, EncodedValueOrdinalAndIndex);
   }
 }
 
 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)) {
+  std::uint32_t ValueOrdinalBase;
+  std::uint32_t UnusedValueOrdinalSize;
+  llvm::Value *ValueOrdinalIndex;
+  if (!pix_dxil::PixAllocaRegWrite::FromInst(Inst, &ValueOrdinalBase, &UnusedValueOrdinalSize, &ValueOrdinalIndex)) {
     return;
   }
 
-  auto *RegIndexConst = llvm::dyn_cast<llvm::ConstantInt>(RegIndex);
-  if (RegIndexConst == nullptr) {
+  std::uint32_t InstNum;
+  if (!pix_dxil::PixDxilInstNum::FromInst(Inst, &InstNum)) {
     return;
   }
 
-  if (RegIndexConst->getLimitedValue() >= RegSize) {
-    return;
-  }
-
-  const std::uint32_t ActualReg = RegBase + RegIndexConst->getLimitedValue();
-  addStepDebugEntryValue(BC, Inst->getValueOperand(), ActualReg);
+  addStepDebugEntryValue(BC, InstNum, Inst->getValueOperand(), ValueOrdinalBase, ValueOrdinalIndex);
 }
 
 void DxilDebugInstrumentation::addStepDebugEntry(BuilderContext &BC, Instruction *Inst) {
@@ -685,33 +694,38 @@ void DxilDebugInstrumentation::addStepDebugEntry(BuilderContext &BC, Instruction
     return;
   }
 
-  addStepDebugEntryValue(BC, Inst, RegNum);
+  std::uint32_t InstNum;
+  if (!pix_dxil::PixDxilInstNum::FromInst(Inst, &InstNum)) {
+    return;
+  }
+
+  addStepDebugEntryValue(BC, InstNum, Inst, RegNum, BC.Builder.getInt32(0));
 }
 
-void DxilDebugInstrumentation::addStepDebugEntryValue(BuilderContext &BC, Value *V, std::uint32_t RegNum) {
+void DxilDebugInstrumentation::addStepDebugEntryValue(BuilderContext &BC, std::uint32_t InstNum, Value *V, std::uint32_t ValueOrdinal, Value *ValueOrdinalIndex) {
   const Type::TypeID ID = V->getType()->getTypeID();
 
   switch (ID) {
   case Type::TypeID::StructTyID:
   case Type::TypeID::VoidTyID:
-    addStepEntryForType<void>(DebugShaderModifierRecordTypeDXILStepVoid, BC, V, RegNum);
+    addStepEntryForType<void>(DebugShaderModifierRecordTypeDXILStepVoid, BC, InstNum, V, ValueOrdinal, ValueOrdinalIndex);
     break;
   case Type::TypeID::FloatTyID:
-    addStepEntryForType<float>(DebugShaderModifierRecordTypeDXILStepFloat, BC, V, RegNum);
+    addStepEntryForType<float>(DebugShaderModifierRecordTypeDXILStepFloat, BC, InstNum, V, ValueOrdinal, ValueOrdinalIndex);
     break;
   case Type::TypeID::IntegerTyID:
     if (V->getType()->getIntegerBitWidth() == 64) {
-      addStepEntryForType<uint64_t>(DebugShaderModifierRecordTypeDXILStepUint64, BC, V, RegNum);
+      addStepEntryForType<uint64_t>(DebugShaderModifierRecordTypeDXILStepUint64, BC, InstNum, V, ValueOrdinal, ValueOrdinalIndex);
     }
     else {
-      addStepEntryForType<uint32_t>(DebugShaderModifierRecordTypeDXILStepUint32, BC, V, RegNum);
+      addStepEntryForType<uint32_t>(DebugShaderModifierRecordTypeDXILStepUint32, BC, InstNum, V, ValueOrdinal, ValueOrdinalIndex);
     }
     break;
   case Type::TypeID::DoubleTyID:
-    addStepEntryForType<double>(DebugShaderModifierRecordTypeDXILStepDouble, BC, V, RegNum);
+    addStepEntryForType<double>(DebugShaderModifierRecordTypeDXILStepDouble, BC, InstNum, V, ValueOrdinal, ValueOrdinalIndex);
     break;
   case Type::TypeID::HalfTyID:
-    addStepEntryForType<float>(DebugShaderModifierRecordTypeDXILStepFloat, BC, V, RegNum);
+    addStepEntryForType<float>(DebugShaderModifierRecordTypeDXILStepFloat, BC, InstNum, V, ValueOrdinal, ValueOrdinalIndex);
     break;
   case Type::TypeID::PointerTyID:
     // Skip pointer calculation instructions. They aren't particularly meaningful to the user (being a mere

lib/DxilPIXPasses/DxilPIXVirtualRegisters.cpp

@@ -29,6 +29,41 @@ static llvm::Metadata *MetadataForValue(llvm::Value *V) {
   return llvm::ValueAsMetadata::get(V);
 }
 
+void pix_dxil::PixDxilInstNum::AddMD(llvm::LLVMContext &Ctx, llvm::Instruction *pI, std::uint32_t InstNum) {
+  llvm::IRBuilder<> B(Ctx);
+  pI->setMetadata(
+    llvm::StringRef(MDName),
+    llvm::MDNode::get(Ctx, { llvm::ConstantAsMetadata::get(B.getInt32(ID)),
+                             llvm::ConstantAsMetadata::get(B.getInt32(InstNum)) }));
+}
+
+bool pix_dxil::PixDxilInstNum::FromInst(llvm::Instruction *pI, std::uint32_t *pInstNum) {
+  *pInstNum = 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 *mdInstNum = llvm::mdconst::dyn_extract<llvm::ConstantInt>(mdNodes->getOperand(1));
+  if (mdInstNum == nullptr) {
+    return false;
+  }
+
+  *pInstNum = mdInstNum->getLimitedValue();
+  return true;
+}
+
 void pix_dxil::PixDxilReg::AddMD(llvm::LLVMContext &Ctx, llvm::Instruction *pI, std::uint32_t RegNum) {
   llvm::IRBuilder<> B(Ctx);
   pI->setMetadata(

lib/DxilPIXPasses/DxilPIXVirtualRegisters.h

@@ -22,6 +22,14 @@ class Value;
 }  // namespace llvm
 
 namespace pix_dxil {
+namespace PixDxilInstNum {
+static constexpr char MDName[] = "pix-dxil-inst-num";
+static constexpr uint32_t ID = 3;
+
+void AddMD(llvm::LLVMContext &Ctx, llvm::Instruction *pI, std::uint32_t InstNum);
+bool FromInst(llvm::Instruction *pI, std::uint32_t *pInstNum);
+}  // namespace PixDxilInstNum
+
 namespace PixDxilReg {
 static constexpr char MDName[] = "pix-dxil-reg";
 static constexpr uint32_t ID = 0;

tools/clang/test/CodeGenHLSL/pix/DebugFlowControl.hlsl

@@ -8,7 +8,7 @@
 // CHECK:  %MaskedForUAVLimit3 = and i32 %UAVIncResult2, 983039
 // CHECK:  %MultipliedForInterest4 = mul i32 %MaskedForUAVLimit3, %OffsetMultiplicand
 // CHECK:  %AddedForInterest5 = add i32 %MultipliedForInterest4, %OffsetAddend
-// CHECK:  call void @dx.op.bufferStore.i32(i32 69, %dx.types.Handle %PIX_DebugUAV_Handle, i32 %AddedForInterest5, i32 undef, i32 64770, i32 undef, i32 undef, i32 undef, i8 1)
+// CHECK:  call void @dx.op.bufferStore.i32(i32 69, %dx.types.Handle %PIX_DebugUAV_Handle, i32 %AddedForInterest5, i32 undef, i32 64771, i32 undef, i32 undef, i32 undef, i8 1)
 // CHECK:  switch i32
 // CHECK:    i32 0, label 
 // CHECK:    i32 32, label

tools/clang/tools/dxcompiler/dxcassembler.cpp

@@ -130,8 +130,11 @@ HRESULT STDMETHODCALLTYPE DxcAssembler::AssembleToContainer(
     outStream.flush();
 
     CComPtr<IDxcBlob> pResultBlob;
+    static constexpr hlsl::SerializeDxilFlags flags = static_cast<hlsl::SerializeDxilFlags>(
+        static_cast<uint32_t>(SerializeDxilFlags::IncludeDebugNamePart) |
+        static_cast<uint32_t>(SerializeDxilFlags::IncludeDebugInfoPart));
     dxcutil::AssembleToContainer(std::move(M), pResultBlob,
-                                         TM.p, SerializeDxilFlags::IncludeDebugNamePart,
+                                         TM.p, flags,
                                          pOutputStream);
 
     IFT(DxcOperationResult::CreateFromResultErrorStatus(pResultBlob, nullptr, S_OK, ppResult));