Add InstanceContributionToHitGroupIndex YAML field and shader query

include/API/AccelerationStructure.h

@@ -55,6 +55,8 @@ struct AccelerationStructureInstance {
   float Transform[3][4] = {{1, 0, 0, 0}, {0, 1, 0, 0}, {0, 0, 1, 0}};
   uint32_t InstanceID = 0;
   uint8_t InstanceMask = 0xFF;
+  // 24-bit; high bits are truncated by the backend to match DXR's bitfield.
+  uint32_t InstanceContributionToHitGroupIndex = 0;
   AccelerationStructure *BLAS = nullptr;
 };
 

include/Support/Pipeline.h

@@ -515,6 +515,7 @@ struct InstanceDesc {
   float Transform[12] = {1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0};
   uint32_t InstanceID = 0;
   uint8_t InstanceMask = 0xFF;
+  uint32_t InstanceContributionToHitGroupIndex = 0;
 };
 
 struct TLASDesc {

lib/API/DX/Device.cpp

@@ -3065,7 +3065,8 @@ llvm::Error DXComputeEncoder::batchBuildAS(llvm::ArrayRef<ASBuildItem> Items) {
         // silent narrowing.
         NI.InstanceID = Inst.InstanceID & 0xFFFFFFu;
         NI.InstanceMask = Inst.InstanceMask;
-        NI.InstanceContributionToHitGroupIndex = 0;
+        NI.InstanceContributionToHitGroupIndex =
+            Inst.InstanceContributionToHitGroupIndex & 0xFFFFFFu;
         NI.Flags = D3D12_RAYTRACING_INSTANCE_FLAG_NONE;
         auto *BLASPtr = llvm::cast<DXAccelerationStructure>(Inst.BLAS);
         NI.AccelerationStructure = BLASPtr->getGPUVirtualAddress();

lib/API/Device.cpp

@@ -200,6 +200,8 @@ llvm::Error offloadtest::buildPipelineAccelerationStructures(
       memcpy(Inst.Transform, I.Transform, sizeof(I.Transform));
       Inst.InstanceID = I.InstanceID;
       Inst.InstanceMask = I.InstanceMask;
+      Inst.InstanceContributionToHitGroupIndex =
+          I.InstanceContributionToHitGroupIndex;
       Inst.BLAS = It->second;
       Req.Instances.push_back(Inst);
     }

lib/API/MTL/MTLDevice.cpp

@@ -1381,7 +1381,11 @@ class MTLDevice : public offloadtest::Device {
           // `InstanceContributionToHitGroupIndex`).
           const uint32_t InstCount =
               static_cast<uint32_t>(R.first->TLASPtr->Instances.size());
-          llvm::SmallVector<uint32_t> Contributions(InstCount, 0);
+          llvm::SmallVector<uint32_t> Contributions;
+          Contributions.reserve(InstCount);
+          for (const auto &Inst : R.first->TLASPtr->Instances)
+            Contributions.push_back(Inst.InstanceContributionToHitGroupIndex &
+                                    0xFFFFFFu);
           const BufferCreateDesc Desc{MemoryLocation::GpuToCpu,
                                       BufferUsage::Storage};
           auto ContribBufOrErr = createBufferWithData(

lib/API/VK/Device.cpp

@@ -4397,7 +4397,8 @@ llvm::Error VKComputeEncoder::batchBuildAS(llvm::ArrayRef<ASBuildItem> Items) {
         memcpy(&NI.transform.matrix, Inst.Transform, sizeof(Inst.Transform));
         NI.instanceCustomIndex = Inst.InstanceID & 0xFFFFFFu;
         NI.mask = Inst.InstanceMask;
-        NI.instanceShaderBindingTableRecordOffset = 0;
+        NI.instanceShaderBindingTableRecordOffset =
+            Inst.InstanceContributionToHitGroupIndex & 0xFFFFFFu;
         NI.flags = 0;
         auto *BLASPtr = llvm::cast<VulkanAccelerationStructure>(Inst.BLAS);
         NI.accelerationStructureReference = BLASPtr->getDeviceAddress();

lib/Support/Pipeline.cpp

@@ -668,6 +668,8 @@ void MappingTraits<offloadtest::InstanceDesc>::mapping(
   uint32_t Mask = D.InstanceMask;
   I.mapOptional("InstanceMask", Mask, 255u);
   D.InstanceMask = static_cast<uint8_t>(Mask);
+  I.mapOptional("InstanceContributionToHitGroupIndex",
+                D.InstanceContributionToHitGroupIndex, 0u);
 }
 
 void MappingTraits<offloadtest::TLASDesc>::mapping(IO &I,

test/Feature/InlineRT/instance-contribution.test

@@ -0,0 +1,93 @@
+#--- source.hlsl
+
+[[vk::binding(0, 0)]] RaytracingAccelerationStructure Scene : register(t0);
+[[vk::binding(1, 0)]] RWStructuredBuffer<uint> Output : register(u0);
+
+[numthreads(3,1,1)]
+void main(uint3 tid : SV_DispatchThreadID) {
+  // Three instances of the same triangle BLAS tiled along x at x = -4, 0, +4
+  // with distinct InstanceContributionToHitGroupIndex values. Each lane fires
+  // straight down at its own instance, so CommittedInstanceContributionToHit-
+  // GroupIndex() must equal the per-instance value.
+  RayDesc Ray;
+  Ray.Origin = float3((float(tid.x) - 1.0) * 4.0, 0, 1);
+  Ray.Direction = float3(0, 0, -1);
+  Ray.TMin = 0.0;
+  Ray.TMax = 100.0;
+  RayQuery<RAY_FLAG_NONE> Q;
+  Q.TraceRayInline(Scene, RAY_FLAG_NONE, 0xFF, Ray);
+  Q.Proceed();
+  Output[tid.x] = Q.CommittedStatus() == COMMITTED_TRIANGLE_HIT
+                      ? Q.CommittedInstanceContributionToHitGroupIndex()
+                      : 0xFFFFFFFF;
+}
+//--- pipeline.yaml
+---
+Shaders:
+  - Stage: Compute
+    Entry: main
+Buffers:
+  - Name: Vertices
+    Format: Float32
+    Stride: 12
+    Data: [ 0.0, 1.0, 0.0, -1.0, -1.0, 0.0, 1.0, -1.0, 0.0 ]
+  - Name: Output
+    Format: UInt32
+    Stride: 4
+    FillSize: 12
+  - Name: Expected
+    Format: UInt32
+    Stride: 4
+    # 24-bit field: highest legal value is 0xFFFFFF. Pick three distinct
+    # values across the range, including one that exercises the top bits.
+    Data: [ 7, 99, 16777215 ]
+AccelerationStructures:
+  BLAS:
+    - Name: TriangleBLAS
+      Triangles:
+        - VertexBuffer: Vertices
+          VertexFormat: RGB32Float
+          VertexStride: 12
+          VertexCount: 3
+  TLAS:
+    - Name: Scene
+      Instances:
+        - BLAS: TriangleBLAS
+          Transform: [1, 0, 0, -4,  0, 1, 0, 0,  0, 0, 1, 0]
+          InstanceContributionToHitGroupIndex: 7
+        - BLAS: TriangleBLAS
+          Transform: [1, 0, 0, 0,  0, 1, 0, 0,  0, 0, 1, 0]
+          InstanceContributionToHitGroupIndex: 99
+        - BLAS: TriangleBLAS
+          Transform: [1, 0, 0, 4,  0, 1, 0, 0,  0, 0, 1, 0]
+          InstanceContributionToHitGroupIndex: 16777215
+DescriptorSets:
+  - Resources:
+    - Name: Scene
+      Kind: AccelerationStructure
+      DirectXBinding:
+        Register: 0
+        Space: 0
+      VulkanBinding:
+        Binding: 0
+    - Name: Output
+      Kind: RWStructuredBuffer
+      DirectXBinding:
+        Register: 0
+        Space: 0
+      VulkanBinding:
+        Binding: 1
+Results:
+  - Result: InstanceContribution
+    Rule: BufferExact
+    Actual: Output
+    Expected: Expected
+...
+#--- end
+
+# REQUIRES: acceleration-structure
+# XFAIL: Clang
+
+# RUN: split-file %s %t
+# RUN: %dxc_target -T cs_6_5 -Fo %t.o %t/source.hlsl
+# RUN: %offloader %t/pipeline.yaml %t.o