Bind acceleration structures and enable the InlineRT tests (#1245)

Closes #1158 🥳

## Summary
Wire up acceleration-structure descriptor binding end-to-end across all
three backends so shaders can actually consume the TLAS that
`buildPipelineAccelerationStructures()` produced — completing the stack
and promoting the three InlineRT tests from XFAIL to passing.

Per-resource AS handling lands in a new per-backend `createAS()` (paired
with `createSRV()` / `createUAV()` / `createCBV()`): a pure
single-create that queries TLAS sizes via `Dev.getTLASBuildSizes()` and
allocates the handle via `Dev.createTLAS()`, returning the `unique_ptr`
to the caller. No `InvocationState` or `Pipeline` access — the
multi-create (`createBuffers()` / `createResources()`) records the
handle in `InvocationState::TLASes` (a `StringMap` keyed by
`TLASDesc::Name`) and wires a non-owning AS pointer into the
per-resource bundle the binding loop reads. The shared AS-build helper
picks up that map and walks `P.AccelStructs.TLAS` to pair each YAML
descriptor with its pre-allocated handle by name (TLASes without a map
entry are skipped, i.e. declared but unbound). BLAS handles are still
allocated by the helper itself since BLASes aren't user-bindable.

`executeProgram()` in each backend now runs as:

- `createBuffers` / `createResources` (`createAS()` allocates TLAS
handles)
- open encoder → `buildPipelineAccelerationStructures()` → end

- **Vulkan**: `createDescriptorPool()` counts AS descriptors in a
separate scalar (the KHR enum value `1000150000` doesn't fit in the
indexed array used for the core types) and emits one
`VkDescriptorPoolSize` for them. `createDescriptorSets()` reads the
resolved `VulkanAccelerationStructure` handle from `ResourceRef.AS`
(populated by `createResources()`) and writes it through a
`VkWriteDescriptorSetAccelerationStructureKHR` chained on the descriptor
write's `pNext`. The dispatch's pre-barrier dst access now includes
`VK_ACCESS_ACCELERATION_STRUCTURE_READ_BIT_KHR` so the prior AS-build's
writes are visible to the shader's RayQuery reads. Device creation
enables `VK_KHR_ray_query` using the same chain-pre-query +
error-on-flag-mismatch pattern that #1232 set up for the AS / BDA
extensions — without `VK_KHR_ray_query` enabled the shader's
`OpRayQueryProceedKHR` instructions silently no-op and `Output` reads
back zero. `copyResourceDataToDevice()` short-circuits AS bundles via a
new `ResourceBundle::isAccelerationStructure()` predicate (no host
buffer to barrier).
- **DX12**: writes a
`D3D12_SRV_DIMENSION_RAYTRACING_ACCELERATION_STRUCTURE` SRV with the AS
GPU virtual address as `Location` into the heap slot that
`createBuffers()` reserved (`CreateShaderResourceView()` with a null
resource — the AS data lives in the buffer pointed to by `Location`).
- **Metal**: the Metal shader converter doesn't bind the AS directly;
the shader reads a buffer containing an
`IRRaytracingAccelerationStructureGPUHeader` that holds the AS's
`gpuResourceID` plus a pointer to an instance-contributions array.
`createBuffers()` allocates and fills both buffers per AS-descriptor
entry, then points the descriptor at the header buffer's GPU address.
The TLAS itself is built with the `UserID` instance-descriptor variant
so HLSL `CommittedInstanceID()` returns the YAML-specified per-instance
ID instead of the array index.

The three InlineRT tests now actually exercise the AS end-to-end:
`TraceRayInline()` issues a RayQuery against `Scene` and writes a
hit-dependent value into `Output` (the instance ID for `multi-instance`,
1/0 otherwise). The catch-all `XFAIL: *` is dropped; `XFAIL: Clang`
remains. The test shaders also gain explicit `[[vk::binding]]`
annotations because dxc's default HLSL→SPIR-V binding mapping collides
`Scene`'s `t0` with `Output`'s `u0` at binding 0, which VVL flags as a
descriptor type mismatch.

## Test plan

Local on an NVIDIA RTX 3060:
- [x] Linux Vulkan (native `offloader`)
- [ ] Linux D3D12 (Wine + vkd3d-proton + cross-compiled `offloader.exe`)
- [ ] Windows Vulkan (native `offloader.exe`)
- [ ] Windows D3D12 (native `offloader.exe`)

CI (RT-capable runners):
- [ ] windows-nvidia D3D12 (`RaytracingTier 1.2`)
- [ ] windows-intel VK (`VK_KHR_ray_tracing_pipeline`)
- [ ] macOS Metal (`supportsRaytracing`)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: MarijnS95

include/API/Device.h

@@ -25,6 +25,7 @@
 #include "Support/Pipeline.h"
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/Support/Error.h"
@@ -323,23 +324,21 @@ createBufferWithData(Device &Dev, std::string Name,
                      size_t SizeInBytes, ComputeEncoder *Encoder,
                      std::unique_ptr<offloadtest::Buffer> *OutUploadBuffer);
 
-// Builds all BLAS / TLAS objects defined in `P.AccelStructs` using the
-// supplied compute encoder. Uploads each BLAS's vertex/index data, queries
-// sizes via `Dev.getBLASBuildSizes` / `Dev.getTLASBuildSizes`, allocates
-// the handles via `Dev.createBLAS` / `Dev.createTLAS`, and records the GPU
-// builds via two `Enc.batchBuildAS` calls (BLAS batch then TLAS batch — so
-// the AS-build-write barrier between BLAS and TLAS is automatic).
-//
-// Built AS objects are pushed to `OutAS` (in declaration order: BLASes first,
-// then TLASes). Vertex/index buffers used as build inputs are pushed to
-// `OutInputBuffers`; both must outlive command-buffer submission.
+// TLAS handles come in pre-allocated because the caller's binding loop
+// stamps the AS pointer into descriptor bundles before this helper runs;
+// BLAS handles are allocated inline since BLASes aren't user-bindable.
+// BLAS and TLAS builds get separate `Enc.batchBuildAS()` calls so the
+// implicit BLAS-write → TLAS-read barrier sits between them. Outputs
+// (`OutBLAS`, `OutInputBuffers`) must outlive command-buffer submission.
 //
 // TODO: `Pipeline` belongs to the test framework, not the rendering backend
 // API. This helper lives here only because `executeProgram` is still on
 // `Device` — once that moves out, this helper should follow.
 llvm::Error buildPipelineAccelerationStructures(
     Device &Dev, ComputeEncoder &Enc, Pipeline &P,
-    llvm::SmallVectorImpl<std::unique_ptr<AccelerationStructure>> &OutAS,
+    llvm::SmallVectorImpl<std::unique_ptr<AccelerationStructure>> &OutBLAS,
+    const llvm::StringMap<std::unique_ptr<AccelerationStructure>>
+        &PreallocatedTLASes,
     llvm::SmallVectorImpl<std::unique_ptr<Buffer>> &OutInputBuffers);
 
 } // namespace offloadtest

lib/API/DX/Device.cpp

@@ -1078,21 +1078,26 @@ class DXDevice : public offloadtest::Device {
     ComPtr<ID3D12Resource> Buffer;
     std::unique_ptr<offloadtest::Buffer> Readback;
     ComPtr<ID3D12Heap> Heap;
-    ResourceSet(ComPtr<ID3D12Resource> Upload, ComPtr<ID3D12Resource> Buffer,
-                std::unique_ptr<offloadtest::Buffer> Readback,
-                ComPtr<ID3D12Heap> Heap = nullptr)
+    // AS-only; mutually exclusive with the buffer/heap fields above.
+    DXAccelerationStructure *AS = nullptr;
+    explicit ResourceSet(ComPtr<ID3D12Resource> Upload,
+                         ComPtr<ID3D12Resource> Buffer,
+                         std::unique_ptr<offloadtest::Buffer> Readback,
+                         ComPtr<ID3D12Heap> Heap = nullptr)
         : Upload(Upload), Buffer(Buffer), Readback(std::move(Readback)),
           Heap(Heap) {}
+    explicit ResourceSet(DXAccelerationStructure *AS) : AS(AS) {}
     ResourceSet(const ResourceSet &) = delete;
     ResourceSet(ResourceSet &&A)
         : Upload(A.Upload), Buffer(A.Buffer), Readback(std::move(A.Readback)),
-          Heap(A.Heap) {}
+          Heap(A.Heap), AS(A.AS) {}
     ResourceSet &operator=(const ResourceSet &) = delete;
     ResourceSet &operator=(ResourceSet &&A) {
       Upload = A.Upload;
       Buffer = A.Buffer;
       Readback = std::move(A.Readback);
       Heap = A.Heap;
+      AS = A.AS;
       return *this;
     }
   };
@@ -1121,9 +1126,11 @@ class DXDevice : public offloadtest::Device {
     llvm::SmallVector<DescriptorTable> DescTables;
     llvm::SmallVector<ResourcePair> RootResources;
 
-    // Built acceleration structures, kept alive for the pipeline lifetime.
+    // Parallel-indexed to `P.AccelStructs.BLAS`.
     llvm::SmallVector<std::unique_ptr<offloadtest::AccelerationStructure>>
-        AccelStructs;
+        BLASes;
+    // Keyed by `TLASDesc::Name`.
+    llvm::StringMap<std::unique_ptr<offloadtest::AccelerationStructure>> TLASes;
     // Vertex/index buffers consumed during AS builds; must outlive submission.
     llvm::SmallVector<std::unique_ptr<offloadtest::Buffer>> ASInputBuffers;
   };
@@ -2007,21 +2014,40 @@ class DXDevice : public offloadtest::Device {
   // returns the next available HeapIdx
   uint32_t bindSRV(Resource &R, InvocationState &IS, uint32_t HeapIdx,
                    const ResourceBundle &ResBundle) {
-    const uint32_t EltSize = R.getElementSize();
-    const uint32_t NumElts = R.size() / EltSize;
-    const D3D12_SHADER_RESOURCE_VIEW_DESC SRVDesc = getSRVDescription(R);
     const uint32_t DescHandleIncSize = Device->GetDescriptorHandleIncrementSize(
         D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV);
     const D3D12_CPU_DESCRIPTOR_HANDLE SRVHandleHeapStart =
         IS.DescHeap->GetCPUDescriptorHandleForHeapStart();
 
-    for (const ResourceSet &RS : ResBundle) {
-      llvm::outs() << "SRV: HeapIdx = " << HeapIdx << " EltSize = " << EltSize
-                   << " NumElts = " << NumElts << "\n";
-      D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle = SRVHandleHeapStart;
-      SRVHandle.ptr += HeapIdx * DescHandleIncSize;
-      Device->CreateShaderResourceView(RS.Buffer.Get(), &SRVDesc, SRVHandle);
-      HeapIdx++;
+    if (R.isAccelerationStructure()) {
+      // AS SRVs are created with a null resource; the AS lives in the
+      // buffer referenced by Location.
+      D3D12_SHADER_RESOURCE_VIEW_DESC SRVDesc = {};
+      SRVDesc.Format = DXGI_FORMAT_UNKNOWN;
+      SRVDesc.ViewDimension =
+          D3D12_SRV_DIMENSION_RAYTRACING_ACCELERATION_STRUCTURE;
+      SRVDesc.Shader4ComponentMapping =
+          D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
+      for (const ResourceSet &RS : ResBundle) {
+        SRVDesc.RaytracingAccelerationStructure.Location =
+            RS.AS->getGPUVirtualAddress();
+        D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle = SRVHandleHeapStart;
+        SRVHandle.ptr += HeapIdx * DescHandleIncSize;
+        Device->CreateShaderResourceView(nullptr, &SRVDesc, SRVHandle);
+        HeapIdx++;
+      }
+    } else {
+      const uint32_t EltSize = R.getElementSize();
+      const uint32_t NumElts = R.size() / EltSize;
+      const D3D12_SHADER_RESOURCE_VIEW_DESC SRVDesc = getSRVDescription(R);
+      for (const ResourceSet &RS : ResBundle) {
+        llvm::outs() << "SRV: HeapIdx = " << HeapIdx << " EltSize = " << EltSize
+                     << " NumElts = " << NumElts << "\n";
+        D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle = SRVHandleHeapStart;
+        SRVHandle.ptr += HeapIdx * DescHandleIncSize;
+        Device->CreateShaderResourceView(RS.Buffer.Get(), &SRVDesc, SRVHandle);
+        HeapIdx++;
+      }
     }
     return HeapIdx;
   }
@@ -2228,11 +2254,35 @@ class DXDevice : public offloadtest::Device {
     return HeapIdx;
   }
 
+  llvm::Expected<std::unique_ptr<AccelerationStructure>> createAS(Resource &R) {
+    assert(R.TLASPtr && "AS resource must be resolved to a TLAS");
+    assert(R.getArraySize() == 1 && "AS arrays not yet supported");
+    auto SizesOrErr =
+        getTLASBuildSizes(static_cast<uint32_t>(R.TLASPtr->Instances.size()));
+    if (!SizesOrErr)
+      return SizesOrErr.takeError();
+    return createTLAS(*SizesOrErr);
+  }
+
   llvm::Error createBuffers(Pipeline &P, InvocationState &IS) {
     auto CreateBuffer =
         [&IS,
          this](Resource &R,
                llvm::SmallVectorImpl<ResourcePair> &Resources) -> llvm::Error {
+      if (R.isAccelerationStructure()) {
+        auto ASOrErr = createAS(R);
+        if (!ASOrErr)
+          return ASOrErr.takeError();
+        ResourceBundle Bundle;
+        Bundle.emplace_back(
+            llvm::cast<DXAccelerationStructure>(ASOrErr->get()));
+        auto Inserted =
+            IS.TLASes.try_emplace(R.TLASPtr->Name, std::move(*ASOrErr));
+        assert(Inserted.second && "TLAS bound to multiple resources NYI");
+        (void)Inserted;
+        Resources.push_back(std::make_pair(&R, std::move(Bundle)));
+        return llvm::Error::success();
+      }
       switch (getDescriptorKind(R.Kind)) {
       case DescriptorKind::SRV: {
         auto ExRes = createSRV(R, IS);
@@ -2723,20 +2773,21 @@ class DXDevice : public offloadtest::Device {
     State.CB->Dev = this;
     llvm::outs() << "Command buffer created.\n";
 
+    if (auto Err = createBuffers(P, State))
+      return Err;
+    llvm::outs() << "Buffers created.\n";
+
     if (!P.AccelStructs.BLAS.empty() || !P.AccelStructs.TLAS.empty()) {
       auto EncOrErr = State.CB->createComputeEncoder();
       if (!EncOrErr)
         return EncOrErr.takeError();
       if (auto Err = offloadtest::buildPipelineAccelerationStructures(
-              *this, **EncOrErr, P, State.AccelStructs, State.ASInputBuffers))
+              *this, **EncOrErr, P, State.BLASes, State.TLASes,
+              State.ASInputBuffers))
         return Err;
       (*EncOrErr)->endEncoding();
     }
 
-    if (auto Err = createBuffers(P, State))
-      return Err;
-    llvm::outs() << "Buffers created.\n";
-
     BindingsDesc BndDesc = {};
     for (auto &S : P.Sets) {
       DescriptorSetLayoutDesc Layout;

lib/API/Device.cpp

@@ -97,7 +97,9 @@ offloadtest::createRenderTargetFromCPUBuffer(Device &Dev,
 
 llvm::Error offloadtest::buildPipelineAccelerationStructures(
     Device &Dev, ComputeEncoder &Enc, Pipeline &P,
-    llvm::SmallVectorImpl<std::unique_ptr<AccelerationStructure>> &OutAS,
+    llvm::SmallVectorImpl<std::unique_ptr<AccelerationStructure>> &OutBLAS,
+    const llvm::StringMap<std::unique_ptr<AccelerationStructure>>
+        &PreallocatedTLASes,
     llvm::SmallVectorImpl<std::unique_ptr<Buffer>> &OutInputBuffers) {
   if (P.AccelStructs.BLAS.empty() && P.AccelStructs.TLAS.empty())
     return llvm::Error::success();
@@ -113,7 +115,7 @@ llvm::Error offloadtest::buildPipelineAccelerationStructures(
   // them through pointers stored in ASBuildItem.
   llvm::SmallVector<BLASBuildRequest> BLASRequests;
   BLASRequests.reserve(P.AccelStructs.BLAS.size());
-  llvm::StringMap<size_t> BLASIndex;
+  llvm::StringMap<AccelerationStructure *> BLASesByName;
 
   for (const auto &BD : P.AccelStructs.BLAS) {
     llvm::SmallVector<TriangleGeometryDesc> Triangles;
@@ -161,8 +163,8 @@ llvm::Error offloadtest::buildPipelineAccelerationStructures(
     Req.AS = ASOrErr->get();
     Req.Geometry = std::move(Triangles);
 
-    BLASIndex[BD.Name] = OutAS.size();
-    OutAS.push_back(std::move(*ASOrErr));
+    BLASesByName[BD.Name] = ASOrErr->get();
+    OutBLAS.push_back(std::move(*ASOrErr));
     BLASRequests.push_back(std::move(Req));
   }
 
@@ -174,16 +176,20 @@ llvm::Error offloadtest::buildPipelineAccelerationStructures(
     if (auto Err = Enc.batchBuildAS(BLASBatch))
       return Err;
 
-  // TLAS pass — references BLASes built in the previous batch.
+  // Separate `batchBuildAS()` from the BLAS batch so the BLAS-write →
+  // TLAS-read barrier between them is implicit.
   llvm::SmallVector<TLASBuildRequest> TLASRequests;
-  TLASRequests.reserve(P.AccelStructs.TLAS.size());
-
-  for (const auto &TD : P.AccelStructs.TLAS) {
+  TLASRequests.reserve(PreallocatedTLASes.size());
+  for (const TLASDesc &TD : P.AccelStructs.TLAS) {
+    auto ASIt = PreallocatedTLASes.find(TD.Name);
+    if (ASIt == PreallocatedTLASes.end())
+      continue; // TLAS declared but not bound to any resource.
     TLASBuildRequest Req;
+    Req.AS = ASIt->second.get();
     Req.Instances.reserve(TD.Instances.size());
     for (const auto &I : TD.Instances) {
-      auto It = BLASIndex.find(I.BLAS);
-      if (It == BLASIndex.end())
+      auto It = BLASesByName.find(I.BLAS);
+      if (It == BLASesByName.end())
         return llvm::createStringError(std::errc::invalid_argument,
                                        "TLAS '%s' references unknown BLAS '%s'",
                                        TD.Name.c_str(), I.BLAS.c_str());
@@ -194,21 +200,11 @@ llvm::Error offloadtest::buildPipelineAccelerationStructures(
       memcpy(Inst.Transform, I.Transform, sizeof(I.Transform));
       Inst.InstanceID = I.InstanceID;
       Inst.InstanceMask = I.InstanceMask;
-      Inst.BLAS = OutAS[It->second].get();
+      Inst.BLAS = It->second;
       Req.Instances.push_back(Inst);
     }
     if (auto Err = validateTLASBuildRequest(Req))
       return Err;
-    auto SizesOrErr =
-        Dev.getTLASBuildSizes(static_cast<uint32_t>(Req.Instances.size()));
-    if (!SizesOrErr)
-      return SizesOrErr.takeError();
-    auto ASOrErr = Dev.createTLAS(*SizesOrErr);
-    if (!ASOrErr)
-      return ASOrErr.takeError();
-
-    Req.AS = ASOrErr->get();
-    OutAS.push_back(std::move(*ASOrErr));
     TLASRequests.push_back(std::move(Req));
   }