Add ComputeEncoder::batchBuildAS() and shared AS-build orchestration helper

Move acceleration-structure build commands behind the abstract
ComputeEncoder interface so the orchestration (data upload, build-request
creation, AS allocation, build recording) can live in one place rather
than splitting across three backends.

ComputeEncoder gains a single batchBuildAS(ArrayRef<ASBuildItem>) method.
Each item carries an AccelerationStructure plus a BLAS or TLAS build
request via PointerUnion. The caller guarantees no inter-item memory
dependencies inside a batch — backends record the whole batch with one
barrier slot, no per-element barriers.

  - Vulkan: single vkCmdBuildAccelerationStructuresKHR call covering the
    whole batch. TLAS items serialize VkAccelerationStructureInstanceKHR
    into a device-address upload buffer, BLAS items pull addresses from
    each VulkanBuffer (new getDeviceAddress accessor). Storage buffers
    transparently gain SHADER_DEVICE_ADDRESS + ACCEL_BUILD_INPUT_READ_ONLY
    flags when ray tracing is supported, with the matching
    VkMemoryAllocateFlagsInfo chained on every allocation.
  - DX12: loop calling BuildRaytracingAccelerationStructure per item with
    no intermediate barriers; D3D12_RAYTRACING_INSTANCE_DESC is
    bit-identical to the Vulkan instance struct.
  - Metal: lazy transition to MTL::AccelerationStructureCommandEncoder,
    deduplicates BLAS handles into the
    MTL::InstanceAccelerationStructureDescriptor's instancedAccelera-
    tionStructures array (Metal references BLASes by index, not GPU
    address).

Each backend's CommandBuffer now carries a back-pointer to its owning
Device so the encoder can reach device-loaded entry points and helpers,
plus a keep-alive list for AS scratch and instance buffers.

A shared helper buildPipelineAccelerationStructures in lib/API/Device.cpp
walks Pipeline::AccelStructs, uploads vertex/index data via the new
createBufferWithData, builds requests, allocates AS objects, and issues
two batchBuildAS calls (BLAS batch then TLAS batch — VUID-03403 forbids
referencing a sibling dstAccelerationStructure in one command). Each
backend's executeProgram calls this helper to build the pipeline's AS
objects.

Descriptor binding for AS resources is intentionally still missing — the
existing InlineRT tests still fail at the descriptor-creation step (an
unbound AS resource trips createSRV / equivalents) and stay caught by
their XFAIL: * annotation; PR #1245 wires the AS descriptor write and
flips them green.

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

Author: MarijnS95

include/API/Device.h

@@ -281,6 +281,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, creates
+// the BLASBuildRequest + AS object via `Dev`, and records the build via
+// `Enc.batchBuildAS`. Then resolves TLAS instance references and records the
+// TLAS batch with a separate call (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.
+llvm::Error buildPipelineAccelerationStructures(
+    Device &Dev, ComputeEncoder &Enc, Pipeline &P,
+    llvm::SmallVectorImpl<std::unique_ptr<AccelerationStructure>> &OutAS,
+    llvm::SmallVectorImpl<std::unique_ptr<Buffer>> &OutInputBuffers);
+
 } // namespace offloadtest
 
 #endif // OFFLOADTEST_API_DEVICE_H

include/API/Encoder.h

@@ -11,6 +11,8 @@
 
 #include "API/API.h"
 
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
 
@@ -21,6 +23,19 @@ namespace offloadtest {
 
 class Buffer;
 class PipelineState;
+class AccelerationStructure;
+struct BLASBuildRequest;
+struct TLASBuildRequest;
+
+/// One acceleration-structure build to record in a batch. The caller is
+/// responsible for ensuring no item in a batch has a memory dependency on
+/// another (e.g. a TLAS that reads a BLAS being built in the same batch must
+/// be in a separate batch — that barrier is inserted between batchBuildAS
+/// calls automatically).
+struct ASBuildItem {
+  AccelerationStructure *AS;
+  llvm::PointerUnion<const BLASBuildRequest *, const TLASBuildRequest *> Req;
+};
 
 /// Base class for all command encoders. An encoder records commands into a
 /// command buffer. Call endEncoding() when done recording. Barriers are
@@ -82,6 +97,14 @@ class ComputeEncoder : public CommandEncoder {
   virtual llvm::Error copyBufferToBuffer(Buffer &Src, size_t SrcOffset,
                                          Buffer &Dst, size_t DstOffset,
                                          size_t Size) = 0;
+
+  /// Build a batch of acceleration structures in a single barrier slot. All
+  /// items in `Items` must be independent — no item may depend on another's
+  /// build output. Backends may issue this as one native batch call (Vulkan)
+  /// or as a sequence of single-AS calls without intermediate barriers (DX12,
+  /// Metal). A barrier covering AS-build writes is implicitly emitted before
+  /// any subsequent command that reads from the freshly-built structures.
+  virtual llvm::Error batchBuildAS(llvm::ArrayRef<ASBuildItem> Items) = 0;
 };
 
 struct Viewport {

lib/API/DX/Device.cpp

@@ -571,12 +571,20 @@ class DXQueue : public offloadtest::Queue {
       override;
 };
 
+class DXDevice; // forward decl — defined below in this same anon ns
+
 class DXCommandBuffer : public offloadtest::CommandBuffer {
 public:
   ComPtr<ID3D12CommandAllocator> Allocator;
   ComPtr<ID3D12GraphicsCommandListX> CmdList;
+  /// Back-pointer to the owning device. Used by encoders that need access to
+  /// device-level resources (e.g. allocating AS scratch buffers).
+  DXDevice *Dev = nullptr;
   /// Whether a UAV barrier is pending from a prior compute command.
   bool PendingUAVBarrier = false;
+  /// Buffers that must outlive command-buffer submission (e.g. AS scratch
+  /// and TLAS instance buffers used during builds).
+  llvm::SmallVector<std::unique_ptr<offloadtest::Buffer>> KeepAliveOwned;
 
   static llvm::Expected<std::unique_ptr<DXCommandBuffer>>
   create(ComPtr<ID3D12DeviceX> Device) {
@@ -717,6 +725,10 @@ class DXComputeEncoder : public offloadtest::ComputeEncoder {
     return llvm::Error::success();
   }
 
+  // Defined out-of-line below — needs DXDevice's full type for access to the
+  // ID3D12Device5 entry point and helper allocators.
+  llvm::Error batchBuildAS(llvm::ArrayRef<ASBuildItem> Items) override;
+
   void endEncodingImpl() override { popDebugGroup(); }
 };
 
@@ -947,6 +959,10 @@ DXCommandBuffer::createRenderEncoder(
 }
 
 class DXDevice : public offloadtest::Device {
+  // DXComputeEncoder needs access to Device5 for AS build commands and to the
+  // raw ID3D12Device for scratch buffer allocation.
+  friend class DXComputeEncoder;
+
 private:
   ComPtr<IDXCoreAdapter> Adapter;
   ComPtr<ID3D12DeviceX> Device;
@@ -1002,6 +1018,12 @@ class DXDevice : public offloadtest::Device {
 
     llvm::SmallVector<DescriptorTable> DescTables;
     llvm::SmallVector<ResourcePair> RootResources;
+
+    // Built acceleration structures, kept alive for the pipeline lifetime.
+    llvm::SmallVector<std::unique_ptr<offloadtest::AccelerationStructure>>
+        AccelStructs;
+    // Vertex/index buffers consumed during AS builds; must outlive submission.
+    llvm::SmallVector<std::unique_ptr<offloadtest::Buffer>> ASInputBuffers;
   };
 
 public:
@@ -1588,7 +1610,11 @@ class DXDevice : public offloadtest::Device {
 
   llvm::Expected<std::unique_ptr<offloadtest::CommandBuffer>>
   createCommandBuffer() override {
-    return DXCommandBuffer::create(Device);
+    auto CBOrErr = DXCommandBuffer::create(Device);
+    if (!CBOrErr)
+      return CBOrErr.takeError();
+    (*CBOrErr)->Dev = this;
+    return std::unique_ptr<offloadtest::CommandBuffer>(std::move(*CBOrErr));
   }
 
   llvm::Expected<std::unique_ptr<offloadtest::RenderPass>>
@@ -2630,8 +2656,19 @@ class DXDevice : public offloadtest::Device {
     if (!CBOrErr)
       return CBOrErr.takeError();
     State.CB = std::move(*CBOrErr);
+    State.CB->Dev = this;
     llvm::outs() << "Command buffer 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))
+        return Err;
+      (*EncOrErr)->endEncoding();
+    }
+
     if (auto Err = createBuffers(P, State))
       return Err;
     llvm::outs() << "Buffers created.\n";
@@ -2822,6 +2859,157 @@ class DXDevice : public offloadtest::Device {
     return llvm::Error::success();
   }
 };
+
+llvm::Error DXComputeEncoder::batchBuildAS(llvm::ArrayRef<ASBuildItem> Items) {
+  if (Items.empty())
+    return llvm::Error::success();
+  if (!CB.Dev || !CB.Dev->Device)
+    return llvm::createStringError(
+        std::errc::not_supported,
+        "Ray tracing not supported on this command buffer's device.");
+  DXDevice *Dev = CB.Dev;
+
+  // BuildRaytracingAccelerationStructure lives on ID3D12GraphicsCommandList4.
+  ComPtr<ID3D12GraphicsCommandList4> CmdList4;
+  if (auto Err = HR::toError(CB.CmdList.As(&CmdList4),
+                             "Failed to query ID3D12GraphicsCommandList4."))
+    return Err;
+
+  // Flush a pending barrier before reading, like dispatch(): a TLAS build must
+  // observe BLASes built in the previous batch.
+  CB.flushBarrier();
+
+  // Per the ComputeEncoder::batchBuildAS contract, the caller guarantees no
+  // inter-item memory dependencies within a batch (BLAS and TLAS go in
+  // separate batches, so a TLAS never sees BLASes from the same call). Each
+  // item also gets its own scratch resource, so there's no aliasing between
+  // the builds — no intra-loop UAV barrier is needed.
+  for (const auto &Item : Items) {
+    auto *DXAS = llvm::cast<DXAccelerationStructure>(Item.AS);
+
+    D3D12_BUILD_RAYTRACING_ACCELERATION_STRUCTURE_DESC Desc = {};
+    Desc.DestAccelerationStructureData = DXAS->getGPUVirtualAddress();
+
+    llvm::SmallVector<D3D12_RAYTRACING_GEOMETRY_DESC> GeomDescs;
+    uint64_t ScratchSize = 0;
+
+    if (const auto *BLAS = llvm::dyn_cast<const BLASBuildRequest *>(Item.Req)) {
+      Desc.Inputs.Type =
+          D3D12_RAYTRACING_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL;
+      Desc.Inputs.DescsLayout = D3D12_ELEMENTS_LAYOUT_ARRAY;
+
+      if (const auto *Tris =
+              std::get_if<llvm::SmallVector<TriangleGeometryDesc>>(
+                  &BLAS->Geometry)) {
+        GeomDescs.reserve(Tris->size());
+        for (const auto &T : *Tris) {
+          D3D12_RAYTRACING_GEOMETRY_DESC GD = {};
+          GD.Type = D3D12_RAYTRACING_GEOMETRY_TYPE_TRIANGLES;
+          if (T.Opaque)
+            GD.Flags = D3D12_RAYTRACING_GEOMETRY_FLAG_OPAQUE;
+          auto *VB = llvm::cast<DXBuffer>(T.VertexBuffer);
+          GD.Triangles.VertexBuffer.StartAddress =
+              VB->Buffer->GetGPUVirtualAddress() + T.VertexBufferOffset;
+          GD.Triangles.VertexBuffer.StrideInBytes = T.VertexStride;
+          GD.Triangles.VertexCount = T.VertexCount;
+          GD.Triangles.VertexFormat = getDXGIFormat(T.VertexFormat);
+          if (T.IndexBuffer) {
+            auto *IB = llvm::cast<DXBuffer>(T.IndexBuffer);
+            GD.Triangles.IndexBuffer =
+                IB->Buffer->GetGPUVirtualAddress() + T.IndexBufferOffset;
+            GD.Triangles.IndexCount = T.IndexCount;
+            GD.Triangles.IndexFormat = getDXGIIndexFormat(T.IdxFormat);
+          }
+          GeomDescs.push_back(GD);
+        }
+      } else {
+        const auto &AABBs =
+            std::get<llvm::SmallVector<AABBGeometryDesc>>(BLAS->Geometry);
+        GeomDescs.reserve(AABBs.size());
+        for (const auto &A : AABBs) {
+          D3D12_RAYTRACING_GEOMETRY_DESC GD = {};
+          GD.Type = D3D12_RAYTRACING_GEOMETRY_TYPE_PROCEDURAL_PRIMITIVE_AABBS;
+          if (A.Opaque)
+            GD.Flags = D3D12_RAYTRACING_GEOMETRY_FLAG_OPAQUE;
+          auto *AB = llvm::cast<DXBuffer>(A.AABBBuffer);
+          GD.AABBs.AABBs.StartAddress =
+              AB->Buffer->GetGPUVirtualAddress() + A.AABBBufferOffset;
+          GD.AABBs.AABBs.StrideInBytes = A.AABBStride;
+          GD.AABBs.AABBCount = A.AABBCount;
+          GeomDescs.push_back(GD);
+        }
+      }
+      Desc.Inputs.NumDescs = static_cast<UINT>(GeomDescs.size());
+      Desc.Inputs.pGeometryDescs = GeomDescs.data();
+      ScratchSize = BLAS->Sizes.ScratchDataSizeInBytes;
+    } else {
+      const auto *TLAS = llvm::cast<const TLASBuildRequest *>(Item.Req);
+      Desc.Inputs.Type = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL;
+      Desc.Inputs.DescsLayout = D3D12_ELEMENTS_LAYOUT_ARRAY;
+      Desc.Inputs.NumDescs = static_cast<UINT>(TLAS->Instances.size());
+
+      // D3D12_RAYTRACING_INSTANCE_DESC has the same byte layout as
+      // VkAccelerationStructureInstanceKHR. Serialize and upload via the
+      // shared abstract-API helper using an upload-heap (CpuToGpu) buffer.
+      llvm::SmallVector<D3D12_RAYTRACING_INSTANCE_DESC> Native;
+      Native.reserve(TLAS->Instances.size());
+      for (const auto &Inst : TLAS->Instances) {
+        D3D12_RAYTRACING_INSTANCE_DESC NI = {};
+        static_assert(sizeof(NI.Transform) == sizeof(Inst.Transform),
+                      "Transform layout mismatch");
+        memcpy(&NI.Transform, Inst.Transform, sizeof(Inst.Transform));
+        // D3D12_RAYTRACING_INSTANCE_DESC packs InstanceID into a 24-bit
+        // bitfield; truncate explicitly so the value matches the VK path
+        // (vkInstanceCustomIndex is likewise 24-bit) instead of relying on
+        // silent narrowing.
+        NI.InstanceID = Inst.InstanceID & 0xFFFFFFu;
+        NI.InstanceMask = Inst.InstanceMask;
+        NI.InstanceContributionToHitGroupIndex = 0;
+        NI.Flags = D3D12_RAYTRACING_INSTANCE_FLAG_NONE;
+        auto *BLASPtr = llvm::cast<DXAccelerationStructure>(Inst.BLAS);
+        NI.AccelerationStructure = BLASPtr->getGPUVirtualAddress();
+        Native.push_back(NI);
+      }
+      const size_t Bytes =
+          Native.size() * sizeof(D3D12_RAYTRACING_INSTANCE_DESC);
+
+      const BufferCreateDesc UploadDesc{MemoryLocation::CpuToGpu,
+                                        BufferUsage::Storage};
+      auto InstBufOrErr = offloadtest::createBufferWithData(
+          *Dev, "TLAS-Instances", UploadDesc, Native.data(), Bytes, nullptr,
+          nullptr);
+      if (!InstBufOrErr)
+        return InstBufOrErr.takeError();
+      auto *DXInstBuf = llvm::cast<DXBuffer>(InstBufOrErr->get());
+      Desc.Inputs.InstanceDescs = DXInstBuf->Buffer->GetGPUVirtualAddress();
+
+      CB.KeepAliveOwned.push_back(std::move(*InstBufOrErr));
+      ScratchSize = TLAS->Sizes.ScratchDataSizeInBytes;
+    }
+
+    // Allocate scratch in the default heap; createBuffer applies
+    // ALLOW_UNORDERED_ACCESS for default-heap allocations and creates the
+    // resource in COMMON state, which D3D12 implicitly promotes to
+    // UNORDERED_ACCESS on first GPU access during the AS build.
+    const BufferCreateDesc ScratchDesc{MemoryLocation::GpuOnly,
+                                       BufferUsage::Storage};
+    auto ScratchOrErr =
+        Dev->createBuffer("AS-Scratch", ScratchDesc, ScratchSize);
+    if (!ScratchOrErr)
+      return ScratchOrErr.takeError();
+    auto *DXScratchBuf = llvm::cast<DXBuffer>(ScratchOrErr->get());
+    Desc.ScratchAccelerationStructureData =
+        DXScratchBuf->Buffer->GetGPUVirtualAddress();
+    CB.KeepAliveOwned.push_back(std::move(*ScratchOrErr));
+
+    insertDebugSignpost("BuildRaytracingAccelerationStructure");
+    CmdList4->BuildRaytracingAccelerationStructure(&Desc, 0, nullptr);
+  }
+
+  // Signal that this batch's AS writes need a barrier before the next reader.
+  CB.addPendingUAVBarrier();
+  return llvm::Error::success();
+}
 } // namespace
 
 llvm::Expected<offloadtest::SubmitResult> DXQueue::submit(