PIX:Descend Type Hierarchy for Alloca'd structs (#5028)

In some cases structs like the RayDesc built-in type (one of the arguments to TraceRays) will be placed in an alloca.
The PIX annotation code wasn't properly counting all the leaf-node values of aggregates within such structs when determining offsets for later members.
The result of this was erroneous attribution of writes to those members during PIX shader debugging.

* PIX:Descend Type Hierarchy for Alloca'd structs

* Add test, vectors are scalar-only

* New file-check for metadata nodes

lib/DxilPIXPasses/DxilAnnotateWithVirtualRegister.cpp

@@ -43,8 +43,10 @@
 
 uint32_t CountStructMembers(llvm::Type const *pType) {
   uint32_t Count = 0;
-
-  if (auto *ST = llvm::dyn_cast<llvm::StructType>(pType)) {
+  if (auto *VT = llvm::dyn_cast<llvm::VectorType>(pType)) {
+    // Vector types can only contain scalars:
+    Count = VT->getVectorNumElements();
+  } else if (auto *ST = llvm::dyn_cast<llvm::StructType>(pType)) {
     for (auto &El : ST->elements()) {
       Count += CountStructMembers(El);
     }
@@ -385,11 +387,26 @@ void DxilAnnotateWithVirtualRegister::AnnotateGeneric(llvm::Instruction *pI) {
             llvm::dyn_cast<llvm::ConstantInt>(GEP->getOperand(2));
         if (OffsetAsInt != nullptr)
         {
-          std::uint32_t Offset = static_cast<std::uint32_t>(
+          std::uint32_t OffsetInElementsFromStructureStart = static_cast<std::uint32_t>(
             OffsetAsInt->getValue().getLimitedValue());
-          DXASSERT(Offset < regSize,
+          DXASSERT(OffsetInElementsFromStructureStart < regSize,
             "Structure member offset out of expected range");
-          PixDxilReg::AddMD(m_DM->GetCtx(), pI, baseStructRegNum + Offset);
+          std::uint32_t OffsetInValuesFromStructureStart =
+              OffsetInElementsFromStructureStart; 
+          if (auto *ST = llvm::dyn_cast<llvm::StructType>(GEP->getPointerOperandType()
+                                                       ->getPointerElementType())) {
+            DXASSERT(OffsetInElementsFromStructureStart < ST->getNumElements(),
+                     "Offset into struct is bigger than struct");
+            OffsetInValuesFromStructureStart = 0;
+            for (std::uint32_t Element = 0;
+                 Element < OffsetInElementsFromStructureStart; ++Element) {
+              OffsetInValuesFromStructureStart +=
+                  CountStructMembers(ST->getElementType(Element));
+            }
+          }
+          PixDxilReg::AddMD(m_DM->GetCtx(), pI,
+                            baseStructRegNum +
+                                OffsetInValuesFromStructureStart);
         }
       }
     }

tools/clang/test/HLSLFileCheck/pix/AnnotateVirtualRegs-Raygen.hlsl

@@ -0,0 +1,40 @@
+// RUN: %dxc -Od -T lib_6_6 %s | %opt -S -dxil-annotate-with-virtual-regs | FileCheck %s
+
+
+/* To run locally run:
+%dxc -Od -T lib_6_6 %s -Fc %t.ll
+%opt %t.ll -S -dxil-annotate-with-virtual-regs | FileCheck %s
+*/
+
+RaytracingAccelerationStructure scene : register(t0);
+
+struct RayPayload
+{
+    int3 color;
+};
+
+[shader("raygeneration")]
+void ENTRY()
+{
+    RayDesc ray = {{0,0,0}, {0,0,1}, 0.05, 1000.0};
+    RayPayload pld;
+    TraceRay(scene, 0 /*rayFlags*/, 0xFF /*rayMask*/, 0 /*sbtRecordOffset*/, 1 /*sbtRecordStride*/, 0 /*missIndex*/, ray, pld);
+}
+
+// CHECK: {{.*}} = alloca %struct.RayDesc, align 4, !pix-dxil-inst-num {{.*}}, !pix-alloca-reg [[RDAlloca:![0-9]+]]
+// CHECK: {{.*}} = alloca %struct.RayPayload, align 4, !pix-dxil-inst-num {{.*}}, !pix-alloca-reg [[RPAlloca:![0-9]+]]
+// CHECK: {{.*}} = getelementptr inbounds %struct.RayDesc, %struct.RayDesc* {{.*}}, i32 0, i32 0, !pix-dxil-inst-num {{.*}}, !pix-dxil-reg [[RDGEP:![0-9]+]]
+// CHECK: {{.*}} = load i32, i32* getelementptr inbounds ([1 x i32], [1 x i32]* @dx.nothing.a, i32 0, i32 0), !pix-dxil-inst-num {{.*}}, !pix-dxil-reg [[NothGEP:![0-9]+]]
+// CHECK: {{.*}} = getelementptr inbounds %struct.RayDesc, %struct.RayDesc* {{.*}}, i32 0, i32 1, !pix-dxil-inst-num {{.*}}, !pix-dxil-reg [[RDGEP2:![0-9]+]]
+// CHECK: {{.*}} = load i32, i32* getelementptr inbounds ([1 x i32], [1 x i32]* @dx.nothing.a, i32 0, i32 0), !pix-dxil-inst-num {{.*}}, !pix-dxil-reg [[NothGEP2:![0-9]+]]
+// CHECK: {{.*}} = getelementptr inbounds %struct.RayDesc, %struct.RayDesc* {{.*}}, i32 0, i32 2, !pix-dxil-inst-num {{.*}}, !pix-dxil-reg [[RDGEP3:![0-9]+]]
+// CHECK: {{.*}} = load i32, i32* getelementptr inbounds ([1 x i32], [1 x i32]* @dx.nothing.a, i32 0, i32 0), !pix-dxil-inst-num {{.*}}, !pix-dxil-reg [[NothGEP3:![0-9]+]]
+
+// CHECK-DAG: [[RDAlloca]] = !{i32 1, i32 0, i32 8}
+// CHECK-DAG: [[RPAlloca]] = !{i32 1, i32 8, i32 3}
+// CHECK-DAG: [[RDGEP]] = !{i32 0, i32 0}
+// CHECK-DAG: [[NothGEP]] = !{i32 0, i32 11}
+// CHECK-DAG: [[RDGEP2]] = !{i32 0, i32 3}
+// CHECK-DAG: [[NothGEP2]] = !{i32 0, i32 12}
+// CHECK-DAG: [[RDGEP3]] = !{i32 0, i32 4}
+// CHECK-DAG: [[NothGEP3]] = !{i32 0, i32 13}

tools/clang/unittests/HLSL/PixTest.cpp

@@ -203,6 +203,7 @@ public:
   TEST_METHOD(AddToASPayload)
 
   TEST_METHOD(PixStructAnnotation_Lib_DualRaygen)
+  TEST_METHOD(PixStructAnnotation_Lib_RaygenAllocaStructAlignment)
 
   TEST_METHOD(PixStructAnnotation_Simple)
   TEST_METHOD(PixStructAnnotation_CopiedStruct)
@@ -2737,6 +2738,99 @@ void Raygen1()
   }
 }
 
+TEST_F(PixTest, PixStructAnnotation_Lib_RaygenAllocaStructAlignment) {
+  if (m_ver.SkipDxilVersion(1, 5)) return;
+
+  const char* hlsl = R"(
+
+RaytracingAccelerationStructure Scene : register(t0, space0);
+RWTexture2D<float4> RenderTarget : register(u0);
+
+struct SceneConstantBuffer
+{
+    float4x4 projectionToWorld;
+    float4 cameraPosition;
+    float4 lightPosition;
+    float4 lightAmbientColor;
+    float4 lightDiffuseColor;
+};
+
+ConstantBuffer<SceneConstantBuffer> g_sceneCB : register(b0);
+
+struct RayPayload
+{
+    float4 color;
+};
+
+inline void GenerateCameraRay(uint2 index, out float3 origin, out float3 direction)
+{
+    float2 xy = index + 0.5f; // center in the middle of the pixel.
+    float2 screenPos = xy;// / DispatchRaysDimensions().xy * 2.0 - 1.0;
+
+    // Invert Y for DirectX-style coordinates.
+    screenPos.y = -screenPos.y;
+
+    // Unproject the pixel coordinate into a ray.
+    float4 world = /*mul(*/float4(screenPos, 0, 1)/*, g_sceneCB.projectionToWorld)*/;
+
+    //world.xyz /= world.w;
+    origin = world.xyz; //g_sceneCB.cameraPosition.xyz;
+    direction = float3(1,0,0);//normalize(world.xyz - origin);
+}
+
+void RaygenCommon()
+{
+    float3 rayDir;
+    float3 origin;
+    
+    // Generate a ray for a camera pixel corresponding to an index from the dispatched 2D grid.
+    GenerateCameraRay(DispatchRaysIndex().xy, origin, rayDir);
+
+    // Trace the ray.
+    // Set the ray's extents.
+    RayDesc ray;
+    ray.Origin = origin;
+    ray.Direction = rayDir;
+    // Set TMin to a non-zero small value to avoid aliasing issues due to floating - point errors.
+    // TMin should be kept small to prevent missing geometry at close contact areas.
+    ray.TMin = 0.001;
+    ray.TMax = 10000.0;
+    RayPayload payload = { float4(0, 0, 0, 0) };
+    TraceRay(Scene, RAY_FLAG_CULL_BACK_FACING_TRIANGLES, ~0, 0, 1, 0, ray, payload);
+
+    // Write the raytraced color to the output texture.
+   // RenderTarget[DispatchRaysIndex().xy] = payload.color;
+}
+
+[shader("raygeneration")]
+void Raygen()
+{
+    RaygenCommon();
+}
+)";
+
+  auto Testables = TestStructAnnotationCase(hlsl, L"-Od", true, L"lib_6_6");
+
+  // Built-in type "RayDesc" has this structure: struct { float3 Origin; float
+  // TMin; float3 Direction; float TMax; } This is 8 floats, with members at
+  // offsets 0,3,4,7 respectively.
+
+  auto FindAtLeastOneOf = [=](char const *name, uint32_t index) {
+    VERIFY_IS_TRUE(std::find_if(Testables.AllocaWrites.begin(),
+                                Testables.AllocaWrites.end(),
+                                [&name, &index](AllocaWrite const &aw) {
+                                  return 0 == strcmp(aw.memberName.c_str(),
+                                                     name) &&
+                                         aw.index == index;
+                                }) != Testables.AllocaWrites.end());
+  };
+
+  FindAtLeastOneOf("Origin.x", 0);
+  FindAtLeastOneOf("TMin", 3);
+  FindAtLeastOneOf("Direction.x", 4);
+  FindAtLeastOneOf("TMax", 7);
+}
+
 TEST_F(PixTest, PixStructAnnotation_Simple) {
   if (m_ver.SkipDxilVersion(1, 5))
     return;