Add per-geometry Transform on TriangleGeometryDesc

include/API/AccelerationStructure.h

@@ -18,7 +18,9 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Error.h"
 
+#include <array>
 #include <cstdint>
+#include <optional>
 #include <variant>
 
 namespace offloadtest {
@@ -52,6 +54,10 @@ struct TriangleGeometryDesc {
   uint32_t IndexCount = 0;
   IndexFormat IdxFormat = IndexFormat::Uint32;
   bool Opaque = true;
+  // Optional BLAS-side bake transform, 3x4 row-major. Vertices are
+  // multiplied by this before AS build, so the resulting BLAS reports
+  // transformed positions via Object* shader queries.
+  std::optional<std::array<float, 12>> Transform;
 };
 
 struct AABBGeometryDesc {

include/Support/Pipeline.h

@@ -21,6 +21,7 @@
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/YAMLTraits.h"
+#include <array>
 #include <limits>
 #include <memory>
 #include <optional>
@@ -552,6 +553,7 @@ struct TriangleGeometry {
   IndexFormat IdxFormat = IndexFormat::Uint32;
   uint32_t IndexCount = 0;
   bool Opaque = true;
+  std::optional<std::array<float, 12>> Transform;
 };
 
 struct AABBGeometry {

lib/API/DX/Device.cpp

@@ -2356,6 +2356,13 @@ class DXDevice : public offloadtest::Device {
         Tri.IndexFormat = getDXGIIndexFormat(T.IdxFormat);
       }
 
+      // Scratch sizing depends on whether Transform3x4 will be present at
+      // build time; signal that with any non-NULL sentinel here — the DXR
+      // spec lets the value be NULL or non-NULL for the prebuild query and
+      // does not dereference it.
+      if (T.Transform)
+        Tri.Transform3x4 = 1;
+
       GeomDescs.push_back(GD);
     }
     return queryBLASPrebuildSize(GeomDescs);
@@ -3406,6 +3413,18 @@ llvm::Error DXComputeEncoder::batchBuildAS(llvm::ArrayRef<ASBuildItem> Items) {
             GD.Triangles.IndexCount = T.IndexCount;
             GD.Triangles.IndexFormat = getDXGIIndexFormat(T.IdxFormat);
           }
+          if (T.Transform) {
+            const BufferCreateDesc XformDesc = BufferCreateDesc::uploadBuffer();
+            auto XformOrErr = offloadtest::createBufferWithData(
+                *Dev, "AS-Geom-Transform", XformDesc, T.Transform->data(),
+                T.Transform->size() * sizeof(float), nullptr, nullptr);
+            if (!XformOrErr)
+              return XformOrErr.takeError();
+            auto *XformBuf = llvm::cast<DXBuffer>(XformOrErr->get());
+            GD.Triangles.Transform3x4 =
+                XformBuf->Buffer->GetGPUVirtualAddress();
+            CB.KeepAliveOwned.push_back(std::move(*XformOrErr));
+          }
           GeomDescs.push_back(GD);
         }
       } else {

lib/API/Device.cpp

@@ -182,6 +182,7 @@ llvm::Error offloadtest::buildPipelineAccelerationStructures(
       TGD.VertexStride = T.VertexStride;
       TGD.VertexFormat = T.VertexFormat;
       TGD.Opaque = T.Opaque;
+      TGD.Transform = T.Transform;
 
       OutInputBuffers.push_back(std::move(*VBOrErr));
 

lib/API/MTL/MTLDevice.cpp

@@ -3184,6 +3184,18 @@ llvm::Error MTLComputeEncoder::batchBuildAS(llvm::ArrayRef<ASBuildItem> Items) {
             TD->setIndexType(getMetalIndexType(T.IdxFormat));
           }
           TD->setOpaque(T.Opaque);
+          if (T.Transform) {
+            const BufferCreateDesc XformDesc = BufferCreateDesc::uploadBuffer();
+            auto XformBufOrErr = offloadtest::createBufferWithData(
+                *CB->Dev, "BLAS-Transform", XformDesc, T.Transform->data(),
+                T.Transform->size() * sizeof(float), nullptr, nullptr);
+            if (!XformBufOrErr)
+              return XformBufOrErr.takeError();
+            auto *MTLXform = llvm::cast<MTLBuffer>(XformBufOrErr->get());
+            TD->setTransformationMatrixBuffer(MTLXform->Buf);
+            TD->setTransformationMatrixLayout(MTL::MatrixLayoutRowMajor);
+            CB->KeepAliveOwned.push_back(std::move(*XformBufOrErr));
+          }
           Geoms.push_back(TD);
         }
       } else {

lib/API/VK/Device.cpp

@@ -4848,6 +4848,17 @@ llvm::Error VKComputeEncoder::batchBuildAS(llvm::ArrayRef<ASBuildItem> Items) {
           } else {
             Tri.indexType = VK_INDEX_TYPE_NONE_KHR;
           }
+          if (T.Transform) {
+            const BufferCreateDesc XformDesc = BufferCreateDesc::uploadBuffer();
+            auto XformOrErr = offloadtest::createBufferWithData(
+                *Dev, "AS-Geom-Transform", XformDesc, T.Transform->data(),
+                T.Transform->size() * sizeof(float), nullptr, nullptr);
+            if (!XformOrErr)
+              return XformOrErr.takeError();
+            auto *XformBuf = llvm::cast<VulkanBuffer>(XformOrErr->get());
+            Tri.transformData.deviceAddress = XformBuf->getDeviceAddress();
+            CB.KeepAliveOwned.push_back(std::move(*XformOrErr));
+          }
           Geoms[I].push_back(G);
 
           VkAccelerationStructureBuildRangeInfoKHR R = {};

lib/Support/Pipeline.cpp

@@ -638,6 +638,19 @@ void MappingTraits<offloadtest::TriangleGeometry>::mapping(
   I.mapOptional("IndexFormat", G.IdxFormat, IndexFormat::Uint32);
   I.mapOptional("IndexCount", G.IndexCount, 0u);
   I.mapOptional("Opaque", G.Opaque, true);
+  llvm::SmallVector<float> Transform;
+  I.mapOptional("Transform", Transform);
+  if (!Transform.empty()) {
+    if (Transform.size() != 12) {
+      I.setError(llvm::Twine("TriangleGeometry.Transform must have exactly 12 "
+                             "floats (3x4 row-major), got ") +
+                 llvm::Twine(Transform.size()));
+      return;
+    }
+    std::array<float, 12> T;
+    std::copy(Transform.begin(), Transform.end(), T.begin());
+    G.Transform = T;
+  }
 }
 
 void MappingTraits<offloadtest::AABBGeometry>::mapping(

test/Feature/InlineRT/geometry-transform.test

@@ -0,0 +1,84 @@
+#--- source.hlsl
+
+[[vk::binding(0, 0)]] RaytracingAccelerationStructure Scene : register(t0);
+[[vk::binding(1, 0)]] RWStructuredBuffer<uint> Output : register(u0);
+
+[numthreads(2,1,1)]
+void main(uint3 tid : SV_DispatchThreadID) {
+  // The triangle vertices are centered around x=0, but a per-geometry
+  // BLAS-bake transform translates them to x=5 at AS-build time. With an
+  // identity-transform TLAS instance, only the ray at x=5 hits.
+  RayDesc Ray;
+  Ray.Origin = float3(tid.x == 0 ? 5.0 : 0.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] = (uint)Q.CommittedStatus();
+}
+//--- 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: 8
+  - Name: Expected
+    Format: UInt32
+    Stride: 4
+    # Lane 0: ray hits the baked-translated triangle → COMMITTED_TRIANGLE_HIT (1)
+    # Lane 1: ray misses (triangle no longer at origin) → COMMITTED_NOTHING (0)
+    Data: [ 1, 0 ]
+AccelerationStructures:
+  BLAS:
+    - Name: TriangleBLAS
+      Triangles:
+        - VertexBuffer: Vertices
+          VertexFormat: RGB32Float
+          VertexStride: 12
+          VertexCount: 3
+          # 3x4 row-major affine — translate x by +5.
+          Transform: [1, 0, 0, 5,  0, 1, 0, 0,  0, 0, 1, 0]
+  TLAS:
+    - Name: Scene
+      Instances:
+        - BLAS: TriangleBLAS
+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: GeometryTransform
+    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