Wire AABB/procedural geometry through the Pipeline-level AS builder

lib/API/Device.cpp

@@ -118,50 +118,89 @@ llvm::Error offloadtest::buildPipelineAccelerationStructures(
   llvm::StringMap<AccelerationStructure *> BLASesByName;
 
   for (const auto &BD : P.AccelStructs.BLAS) {
-    llvm::SmallVector<TriangleGeometryDesc> Triangles;
-    Triangles.reserve(BD.Triangles.size());
-    for (const auto &T : BD.Triangles) {
-      assert(T.VertexBufferPtr && "VertexBufferPtr not resolved");
-      auto VBOrErr = createBufferWithData(
-          Dev, "AS-Vertices", UploadDesc, T.VertexBufferPtr->Data[0].get(),
-          T.VertexBufferPtr->size(), nullptr, nullptr);
-      if (!VBOrErr)
-        return VBOrErr.takeError();
-
-      TriangleGeometryDesc TGD;
-      TGD.VertexBuffer = VBOrErr->get();
-      TGD.VertexCount = T.VertexCount;
-      TGD.VertexStride = T.VertexStride;
-      TGD.VertexFormat = T.VertexFormat;
-      TGD.Opaque = T.Opaque;
-
-      OutInputBuffers.push_back(std::move(*VBOrErr));
-
-      if (T.IndexBufferPtr) {
-        auto IBOrErr = createBufferWithData(
-            Dev, "AS-Indices", UploadDesc, T.IndexBufferPtr->Data[0].get(),
-            T.IndexBufferPtr->size(), nullptr, nullptr);
-        if (!IBOrErr)
-          return IBOrErr.takeError();
-        TGD.IndexBuffer = IBOrErr->get();
-        TGD.IndexCount = T.IndexCount;
-        TGD.IdxFormat = T.IdxFormat;
-        OutInputBuffers.push_back(std::move(*IBOrErr));
+    if (!BD.Triangles.empty() && !BD.AABBs.empty())
+      return llvm::createStringError(
+          std::errc::invalid_argument,
+          "BLAS '%s' mixes triangle and AABB geometry; pick one.",
+          BD.Name.c_str());
+    if (BD.Triangles.empty() && BD.AABBs.empty())
+      return llvm::createStringError(std::errc::invalid_argument,
+                                     "BLAS '%s' has no geometry.",
+                                     BD.Name.c_str());
+
+    auto ReqOrErr = [&]() -> llvm::Expected<BLASBuildRequest> {
+      if (!BD.Triangles.empty()) {
+        llvm::SmallVector<TriangleGeometryDesc> Triangles;
+        Triangles.reserve(BD.Triangles.size());
+        for (const auto &T : BD.Triangles) {
+          assert(T.VertexBufferPtr && "VertexBufferPtr not resolved");
+          auto VBOrErr = createBufferWithData(
+              Dev, "AS-Vertices", UploadDesc, T.VertexBufferPtr->Data[0].get(),
+              T.VertexBufferPtr->size(), nullptr, nullptr);
+          if (!VBOrErr)
+            return VBOrErr.takeError();
+
+          TriangleGeometryDesc TGD;
+          TGD.VertexBuffer = VBOrErr->get();
+          TGD.VertexCount = T.VertexCount;
+          TGD.VertexStride = T.VertexStride;
+          TGD.VertexFormat = T.VertexFormat;
+          TGD.Opaque = T.Opaque;
+
+          OutInputBuffers.push_back(std::move(*VBOrErr));
+
+          if (T.IndexBufferPtr) {
+            auto IBOrErr = createBufferWithData(
+                Dev, "AS-Indices", UploadDesc, T.IndexBufferPtr->Data[0].get(),
+                T.IndexBufferPtr->size(), nullptr, nullptr);
+            if (!IBOrErr)
+              return IBOrErr.takeError();
+            TGD.IndexBuffer = IBOrErr->get();
+            TGD.IndexCount = T.IndexCount;
+            TGD.IdxFormat = T.IdxFormat;
+            OutInputBuffers.push_back(std::move(*IBOrErr));
+          }
+          Triangles.push_back(TGD);
+        }
+        BLASBuildRequest Req;
+        Req.Geometry = std::move(Triangles);
+        return Req;
       }
-      Triangles.push_back(TGD);
-    }
-    // TODO: AABB geometry support (would mirror the triangle path).
-
-    auto SizesOrErr = Dev.getBLASBuildSizes(Triangles);
+      llvm::SmallVector<AABBGeometryDesc> AABBs;
+      AABBs.reserve(BD.AABBs.size());
+      for (const auto &A : BD.AABBs) {
+        assert(A.AABBBufferPtr && "AABBBufferPtr not resolved");
+        auto ABOrErr = createBufferWithData(
+            Dev, "AS-AABBs", UploadDesc, A.AABBBufferPtr->Data[0].get(),
+            A.AABBBufferPtr->size(), nullptr, nullptr);
+        if (!ABOrErr)
+          return ABOrErr.takeError();
+        AABBGeometryDesc AGD;
+        AGD.AABBBuffer = ABOrErr->get();
+        AGD.AABBCount = A.AABBCount;
+        AGD.AABBStride = A.AABBStride;
+        AGD.Opaque = A.Opaque;
+        OutInputBuffers.push_back(std::move(*ABOrErr));
+        AABBs.push_back(AGD);
+      }
+      BLASBuildRequest Req;
+      Req.Geometry = std::move(AABBs);
+      return Req;
+    }();
+
+    if (!ReqOrErr)
+      return ReqOrErr.takeError();
+    auto SizesOrErr =
+        std::visit([&Dev](const auto &Geom) { return Dev.getBLASBuildSizes(Geom); },
+                   ReqOrErr->Geometry);
     if (!SizesOrErr)
       return SizesOrErr.takeError();
     auto ASOrErr = Dev.createBLAS(*SizesOrErr);
     if (!ASOrErr)
       return ASOrErr.takeError();
 
-    BLASBuildRequest Req;
+    BLASBuildRequest Req = std::move(*ReqOrErr);
     Req.AS = ASOrErr->get();
-    Req.Geometry = std::move(Triangles);
 
     BLASesByName[BD.Name] = ASOrErr->get();
     OutBLAS.push_back(std::move(*ASOrErr));

test/Feature/InlineRT/aabb-procedural.test

@@ -0,0 +1,86 @@
+#--- source.hlsl
+
+[[vk::binding(0, 0)]] RaytracingAccelerationStructure Scene : register(t0);
+[[vk::binding(1, 0)]] RWStructuredBuffer<uint> Output : register(u0);
+
+[numthreads(1,1,1)]
+void main() {
+  // A single AABB at [-1, 1]^3. The candidate path fires once per AABB
+  // hit; the shader supplies a fixed hit-t of 1.5 (inside the box on the
+  // -z ray below). After the loop CommittedStatus() must report
+  // COMMITTED_PROCEDURAL_PRIMITIVE_HIT (2).
+  RayDesc Ray;
+  Ray.Origin = float3(0, 0, 2);
+  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);
+  while (Q.Proceed()) {
+    if (Q.CandidateType() == CANDIDATE_PROCEDURAL_PRIMITIVE)
+      Q.CommitProceduralPrimitiveHit(1.5);
+  }
+  Output[0] = (uint)Q.CommittedStatus();
+}
+//--- pipeline.yaml
+---
+Shaders:
+  - Stage: Compute
+    Entry: main
+Buffers:
+  - Name: AABBs
+    Format: Float32
+    Stride: 24
+    # D3D12_RAYTRACING_AABB / VkAabbPositionsKHR layout:
+    # { MinX, MinY, MinZ, MaxX, MaxY, MaxZ }, one box per stride.
+    Data: [ -1.0, -1.0, -1.0, 1.0, 1.0, 1.0 ]
+  - Name: Output
+    Format: UInt32
+    Stride: 4
+    FillSize: 4
+  - Name: Expected
+    Format: UInt32
+    Stride: 4
+    # COMMITTED_PROCEDURAL_PRIMITIVE_HIT = 2
+    Data: [ 2 ]
+AccelerationStructures:
+  BLAS:
+    - Name: AABBBLAS
+      AABBs:
+        - AABBBuffer: AABBs
+          AABBCount: 1
+          AABBStride: 24
+  TLAS:
+    - Name: Scene
+      Instances:
+        - BLAS: AABBBLAS
+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: AABBProcedural
+    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