PIX: split vector alloca writes to enable debug instrumentation (microsoft#6990)

jeffnn · jeffnn · commit bbeb9f78ce98 · 2024-12-04T12:35:42.000-08:00
Bit of background: PIX's debug passes add new allocas, stores to which tie debug info to DXIL Values. In the case of a preexisting alloca, PIX will still add its debug writes, but the codegen may have emitted vector-valued stores. Concretely, this happens for the ray payload and RayDesc structs in DXR. Those vector-values stores were being treated incorrectly, resulting in missing values in the PIX shader debugger. This change splits vector-valued stores to such allocas into extractelement/scalar-stores, which enables the rest of the PIX instrumentation to function as expected. (Worth noting that this change only applies to the PIX shader debugging pass.) (cherry picked from commit a03a77f)
diff --git a/lib/DxilPIXPasses/DxilAnnotateWithVirtualRegister.cpp b/lib/DxilPIXPasses/DxilAnnotateWithVirtualRegister.cpp
@@ -77,6 +77,7 @@ class DxilAnnotateWithVirtualRegister : public llvm::ModulePass {
 private:
   void AnnotateValues(llvm::Instruction *pI);
   void AnnotateStore(llvm::Instruction *pI);
+  void SplitVectorStores(hlsl::OP *HlslOP, llvm::Instruction *pI);
   bool IsAllocaRegisterWrite(llvm::Value *V, llvm::AllocaInst **pAI,
                              llvm::Value **pIdx);
   void AnnotateAlloca(llvm::AllocaInst *pAlloca);
@@ -133,6 +134,15 @@ bool DxilAnnotateWithVirtualRegister::runOnModule(llvm::Module &M) {
   auto instrumentableFunctions =
       PIXPassHelpers::GetAllInstrumentableFunctions(*m_DM);
 
+  for (auto *F : instrumentableFunctions) {
+    for (auto &block : F->getBasicBlockList()) {
+      for (auto it = block.begin(); it != block.end();) {
+        llvm::Instruction *I = &*(it++);
+        SplitVectorStores(m_DM->GetOP(), I);
+      }
+    }
+  }
+
   for (auto *F : instrumentableFunctions) {
     for (auto &block : F->getBasicBlockList()) {
       for (llvm::Instruction &I : block.getInstList()) {
@@ -297,15 +307,37 @@ bool DxilAnnotateWithVirtualRegister::IsAllocaRegisterWrite(
           return false;
         }
         // And of course the member we're after might not be at the beginning of
-        // the struct:
-        auto *pStructType = llvm::dyn_cast<llvm::StructType>(
-            pPointerGEP->getPointerOperandType()->getPointerElementType());
-        auto *pStructMember =
-            llvm::dyn_cast<llvm::ConstantInt>(pPointerGEP->getOperand(2));
-        uint64_t memberIndex = pStructMember->getLimitedValue();
-        for (uint64_t i = 0; i < memberIndex; ++i) {
-          precedingMemberCount +=
-              CountStructMembers(pStructType->getStructElementType(i));
+        // any containing struct:
+        if (auto *pStructType = llvm::dyn_cast<llvm::StructType>(
+                pPointerGEP->getPointerOperandType()
+                    ->getPointerElementType())) {
+          auto *pStructMember =
+              llvm::dyn_cast<llvm::ConstantInt>(pPointerGEP->getOperand(2));
+          uint64_t memberIndex = pStructMember->getLimitedValue();
+          for (uint64_t i = 0; i < memberIndex; ++i) {
+            precedingMemberCount +=
+                CountStructMembers(pStructType->getStructElementType(i));
+          }
+        }
+
+        // And the source pointer may be a vector (floatn) type,
+        // and if so, that's another offset to consider.
+        llvm::Type *DestType = pGEP->getPointerOperand()->getType();
+        // We expect this to be a pointer type (it's a GEP after all):
+        if (DestType->isPointerTy()) {
+          llvm::Type *PointedType = DestType->getPointerElementType();
+          // Being careful to check num operands too in order to avoid false
+          // positives:
+          if (PointedType->isVectorTy() && pGEP->getNumOperands() == 3) {
+            // Fetch the second deref (in operand 2).
+            // (the first derefs the pointer to the "floatn",
+            // and the second denotes the index into the floatn.)
+            llvm::Value *vectorIndex = pGEP->getOperand(2);
+            if (auto *constIntIIndex =
+                    llvm::cast<llvm::ConstantInt>(vectorIndex)) {
+              precedingMemberCount += constIntIIndex->getLimitedValue();
+            }
+          }
         }
       } else {
         return false;
@@ -365,6 +397,8 @@ void DxilAnnotateWithVirtualRegister::AnnotateAlloca(
     AssignNewAllocaRegister(pAlloca, 1);
   } else if (auto *AT = llvm::dyn_cast<llvm::ArrayType>(pAllocaTy)) {
     AssignNewAllocaRegister(pAlloca, AT->getNumElements());
+  } else if (auto *VT = llvm::dyn_cast<llvm::VectorType>(pAllocaTy)) {
+    AssignNewAllocaRegister(pAlloca, VT->getNumElements());
   } else if (auto *ST = llvm::dyn_cast<llvm::StructType>(pAllocaTy)) {
     AssignNewAllocaRegister(pAlloca, CountStructMembers(ST));
   } else {
@@ -433,6 +467,36 @@ void DxilAnnotateWithVirtualRegister::AssignNewAllocaRegister(
   m_uVReg += C;
 }
 
+void DxilAnnotateWithVirtualRegister::SplitVectorStores(hlsl::OP *HlslOP,
+                                                        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::Type *SourceType = pSt->getValueOperand()->getType();
+  if (SourceType->isVectorTy()) {
+    if (auto *constIntIIndex = llvm::cast<llvm::ConstantInt>(Index)) {
+      // break vector alloca stores up into individual stores
+      llvm::IRBuilder<> B(pSt);
+      for (uint64_t el = 0; el < SourceType->getVectorNumElements(); ++el) {
+        llvm::Value *destPointer = B.CreateGEP(pSt->getPointerOperand(),
+                                               {B.getInt32(0), B.getInt32(el)});
+        llvm::Value *source =
+            B.CreateExtractElement(pSt->getValueOperand(), el);
+        B.CreateStore(source, destPointer);
+      }
+      pI->eraseFromParent();
+    }
+  }
+}
+
 } // namespace
 
 using namespace llvm;
diff --git a/tools/clang/unittests/HLSL/PixTest.cpp b/tools/clang/unittests/HLSL/PixTest.cpp
@@ -14,6 +14,7 @@
 #endif
 
 #include <algorithm>
+#include <array>
 #include <cassert>
 #include <cfloat>
 #include <map>
@@ -144,6 +145,8 @@ class PixTest : public ::testing::Test {
   TEST_METHOD(DebugInstrumentation_TextOutput)
   TEST_METHOD(DebugInstrumentation_BlockReport)
 
+  TEST_METHOD(DebugInstrumentation_VectorAllocaWrite_Structs)
+
   dxc::DxcDllSupport m_dllSupport;
   VersionSupportInfo m_ver;
 
@@ -435,6 +438,7 @@ class PixTest : public ::testing::Test {
   CComPtr<IDxcBlob> RunDxilPIXDXRInvocationsLog(IDxcBlob *blob);
   void TestPixUAVCase(char const *hlsl, wchar_t const *model,
                       wchar_t const *entry);
+  std::string Disassemble(IDxcBlob *pProgram);
 };
 
 bool PixTest::InitSupport() {
@@ -1083,6 +1087,16 @@ static bool FindStructMemberFromStore(llvm::StoreInst *S,
   return false;
 }
 
+std::string PixTest::Disassemble(IDxcBlob *pProgram) {
+  CComPtr<IDxcCompiler> pCompiler;
+  CComPtr<IDxcOperationResult> pResult;
+  CComPtr<IDxcBlobEncoding> pSource;
+  VERIFY_SUCCEEDED(CreateCompiler(m_dllSupport, &pCompiler));
+  CComPtr<IDxcBlobEncoding> pDisassembly;
+  VERIFY_SUCCEEDED(pCompiler->Disassemble(pProgram, &pDisassembly));
+  return BlobToUtf8(pDisassembly);
+}
+
 // This function lives in lib\DxilPIXPasses\DxilAnnotateWithVirtualRegister.cpp
 // Declared here so we can test it.
 uint32_t CountStructMembers(llvm::Type const *pType);
@@ -2955,3 +2969,172 @@ float4 main() : SV_Target {
   VERIFY_IS_TRUE(foundDoubleAssignment);
   VERIFY_IS_TRUE(found32BitAllocaStore);
 }
+
+std::string ExtractBracedSubstring(std::string const &line) {
+  auto open = line.find('{');
+  auto close = line.find('}');
+  if (open != std::string::npos && close != std::string::npos &&
+      open + 1 < close) {
+    return line.substr(open + 1, close - open - 1);
+  }
+  return {};
+}
+
+int ExtractMetaInt32Value(std::string const &token) {
+  auto findi32 = token.find("i32 ");
+  if (findi32 != std::string_view::npos) {
+    return atoi(
+        std::string(token.data() + findi32 + 4, token.length() - (findi32 + 4))
+            .c_str());
+  }
+  return -1;
+}
+
+std::vector<std::string> Split(std::string str, char delimeter) {
+  std::vector<std::string> lines;
+
+  auto const *p = str.data();
+  auto const *justPastPreviousDelimiter = p;
+  while (p < str.data() + str.length()) {
+    if (*p == delimeter) {
+      lines.emplace_back(std::string(justPastPreviousDelimiter,
+                                     p - justPastPreviousDelimiter));
+      justPastPreviousDelimiter = p + 1;
+      p = justPastPreviousDelimiter;
+    } else {
+      p++;
+    }
+  }
+
+  lines.emplace_back(
+      std::string(justPastPreviousDelimiter, p - justPastPreviousDelimiter));
+
+  return lines;
+}
+
+struct MetadataAllocaDefinition {
+  int base;
+  int count;
+};
+using AllocaDefinitions = std::map<int, MetadataAllocaDefinition>;
+struct MetadataAllocaWrite {
+  int allocaDefMetadataKey;
+  int offset;
+  int size;
+};
+using AllocaWrites = std::map<int, MetadataAllocaWrite>;
+
+struct AllocaMetadata {
+  AllocaDefinitions allocaDefinitions;
+  AllocaWrites allocaWrites;
+  std::vector<int> allocaWritesMetaKeys;
+};
+
+AllocaMetadata
+FindAllocaRelatedMetadata(std::vector<std::string> const &lines) {
+
+  const char *allocaMetaDataAssignment = "= !{i32 1, ";
+  const char *allocaRegWRiteAssignment = "= !{i32 2, !";
+  const char *allocaRegWriteTag = "!pix-alloca-reg-write !";
+
+  AllocaMetadata ret;
+  for (auto const &line : lines) {
+    if (line[0] == '!') {
+      auto key = atoi(std::string(line.data() + 1, line.length() - 1).c_str());
+      if (key != -1) {
+        if (line.find(allocaMetaDataAssignment) != std::string::npos) {
+          std::string bitInBraces = ExtractBracedSubstring(line);
+          if (bitInBraces != "") {
+            auto tokens = Split(bitInBraces, ',');
+            if (tokens.size() == 3) {
+              auto value0 = ExtractMetaInt32Value(tokens[1]);
+              auto value1 = ExtractMetaInt32Value(tokens[2]);
+              if (value0 != -1 && value1 != -1) {
+                MetadataAllocaDefinition def;
+                def.base = value0;
+                def.count = value1;
+                ret.allocaDefinitions[key] = def;
+              }
+            }
+          }
+        } else if (line.find(allocaRegWRiteAssignment) != std::string::npos) {
+          std::string bitInBraces = ExtractBracedSubstring(line);
+          if (bitInBraces != "") {
+            auto tokens = Split(bitInBraces, ',');
+            if (tokens.size() == 4 && tokens[1][1] == '!') {
+              auto allocaKey = atoi(tokens[1].c_str() + 2);
+              auto value0 = ExtractMetaInt32Value(tokens[2]);
+              auto value1 = ExtractMetaInt32Value(tokens[3]);
+              if (value0 != -1 && value1 != -1) {
+                MetadataAllocaWrite aw;
+                aw.allocaDefMetadataKey = allocaKey;
+                aw.size = value0;
+                aw.offset = value1;
+                ret.allocaWrites[key] = aw;
+              }
+            }
+          }
+        }
+      }
+    } else {
+      auto findAw = line.find(allocaRegWriteTag);
+      if (findAw != std::string::npos) {
+        ret.allocaWritesMetaKeys.push_back(
+            atoi(line.c_str() + findAw + strlen(allocaRegWriteTag)));
+      }
+    }
+  }
+  return ret;
+}
+
+TEST_F(PixTest, DebugInstrumentation_VectorAllocaWrite_Structs) {
+  const char *source = R"x(
+RaytracingAccelerationStructure Scene : register(t0, space0);
+struct RayPayload
+{
+    float4 color;
+};
+RWStructuredBuffer<float> UAV: register(u0);
+[shader("raygeneration")]
+void RaygenInternalName()
+{
+    RayDesc ray;
+    ray.Origin = float3(UAV[0], UAV[1],UAV[3]);
+    ray.Direction = float3(4.4,5.5,6.6);
+    ray.TMin = 0.001;
+    ray.TMax = 10000.0;
+    RayPayload payload = { float4(0, 1, 0, 1) };
+    TraceRay(Scene, RAY_FLAG_CULL_BACK_FACING_TRIANGLES, ~0, 0, 1, 0, ray, payload);
+})x";
+
+  auto compiled = Compile(m_dllSupport, source, L"lib_6_6", {L"-Od"});
+  auto output = RunDebugPass(compiled);
+  auto disassembly = Disassemble(output.blob);
+  auto lines = Split(disassembly, '\n');
+  auto metaDataKeyToValue = FindAllocaRelatedMetadata(lines);
+  // To validate that the RayDesc and RayPayload instances were fully covered,
+  // check that there are alloca writes that cover all of them. RayPayload
+  // has four elements, and RayDesc has eight.
+  std::array<bool, 4> RayPayloadElementCoverage;
+  std::array<bool, 8> RayDescElementCoverage;
+
+  for (auto const &write : metaDataKeyToValue.allocaWrites) {
+    // the whole point of the changes with this test is to separate vector
+    // writes into individual elements:
+    VERIFY_ARE_EQUAL(1, write.second.size);
+    auto findAlloca = metaDataKeyToValue.allocaDefinitions.find(
+        write.second.allocaDefMetadataKey);
+    if (findAlloca != metaDataKeyToValue.allocaDefinitions.end()) {
+      if (findAlloca->second.count == 4) {
+        RayPayloadElementCoverage[write.second.offset] = true;
+      } else if (findAlloca->second.count == 8) {
+        RayDescElementCoverage[write.second.offset] = true;
+      }
+    }
+  }
+  // Check that coverage for every element was emitted:
+  for (auto const &b : RayPayloadElementCoverage)
+    VERIFY_IS_TRUE(b);
+  for (auto const &b : RayDescElementCoverage)
+    VERIFY_IS_TRUE(b);
+}
diff --git a/tools/clang/unittests/HLSL/PixTestUtils.cpp b/tools/clang/unittests/HLSL/PixTestUtils.cpp
@@ -54,7 +54,7 @@ int ExtractMetaInt32Value(std::string const &token) {
   return -1;
 }
 std::map<int, std::pair<int, int>>
-MetaDataKeyToRegisterNumber(std::vector<std::string> const &lines) {
+FindAllocaRelatedMetadata(std::vector<std::string> const &lines) {
   // Find lines of the exemplary form
   // "!249 = !{i32 0, i32 20}"  (in the case of a DXIL value)
   // "!196 = !{i32 1, i32 5, i32 1}" (in the case of a DXIL alloca reg)
@@ -148,7 +148,7 @@ DxilRegisterToNameMap BuildDxilRegisterToNameMap(char const *disassembly) {
 
   auto lines = Tokenize(disassembly, "\n");
 
-  auto metaDataKeyToValue = MetaDataKeyToRegisterNumber(lines);
+  auto metaDataKeyToValue = FindAllocaRelatedMetadata(lines);
 
   for (auto const &line : lines) {
     if (line[0] == '!') {
