diff --git a/include/API/AccelerationStructure.h b/include/API/AccelerationStructure.h index df357771a..90dae2eda 100644 --- a/include/API/AccelerationStructure.h +++ b/include/API/AccelerationStructure.h @@ -18,7 +18,9 @@ #include "llvm/ADT/SmallVector.h" #include "llvm/Support/Error.h" +#include #include +#include #include 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> Transform; }; struct AABBGeometryDesc { diff --git a/include/Support/Pipeline.h b/include/Support/Pipeline.h index f8dc1785d..1e58f0048 100644 --- a/include/Support/Pipeline.h +++ b/include/Support/Pipeline.h @@ -21,6 +21,7 @@ #include "llvm/Support/Error.h" #include "llvm/Support/MemoryBuffer.h" #include "llvm/Support/YAMLTraits.h" +#include #include #include #include @@ -552,6 +553,7 @@ struct TriangleGeometry { IndexFormat IdxFormat = IndexFormat::Uint32; uint32_t IndexCount = 0; bool Opaque = true; + std::optional> Transform; }; struct AABBGeometry { diff --git a/lib/API/DX/Device.cpp b/lib/API/DX/Device.cpp index afce98aa6..f0b499b6f 100644 --- a/lib/API/DX/Device.cpp +++ b/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 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(XformOrErr->get()); + GD.Triangles.Transform3x4 = + XformBuf->Buffer->GetGPUVirtualAddress(); + CB.KeepAliveOwned.push_back(std::move(*XformOrErr)); + } GeomDescs.push_back(GD); } } else { diff --git a/lib/API/Device.cpp b/lib/API/Device.cpp index 33260a6cf..20ed77ad7 100644 --- a/lib/API/Device.cpp +++ b/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)); diff --git a/lib/API/MTL/MTLDevice.cpp b/lib/API/MTL/MTLDevice.cpp index bf3554bb2..1820d13a6 100644 --- a/lib/API/MTL/MTLDevice.cpp +++ b/lib/API/MTL/MTLDevice.cpp @@ -3184,6 +3184,18 @@ llvm::Error MTLComputeEncoder::batchBuildAS(llvm::ArrayRef 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(XformBufOrErr->get()); + TD->setTransformationMatrixBuffer(MTLXform->Buf); + TD->setTransformationMatrixLayout(MTL::MatrixLayoutRowMajor); + CB->KeepAliveOwned.push_back(std::move(*XformBufOrErr)); + } Geoms.push_back(TD); } } else { diff --git a/lib/API/VK/Device.cpp b/lib/API/VK/Device.cpp index 5e5f5c9f7..f13dc8eac 100644 --- a/lib/API/VK/Device.cpp +++ b/lib/API/VK/Device.cpp @@ -4848,6 +4848,17 @@ llvm::Error VKComputeEncoder::batchBuildAS(llvm::ArrayRef 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(XformOrErr->get()); + Tri.transformData.deviceAddress = XformBuf->getDeviceAddress(); + CB.KeepAliveOwned.push_back(std::move(*XformOrErr)); + } Geoms[I].push_back(G); VkAccelerationStructureBuildRangeInfoKHR R = {}; diff --git a/lib/Support/Pipeline.cpp b/lib/Support/Pipeline.cpp index 41fa3f7d6..01846e690 100644 --- a/lib/Support/Pipeline.cpp +++ b/lib/Support/Pipeline.cpp @@ -638,6 +638,19 @@ void MappingTraits::mapping( I.mapOptional("IndexFormat", G.IdxFormat, IndexFormat::Uint32); I.mapOptional("IndexCount", G.IndexCount, 0u); I.mapOptional("Opaque", G.Opaque, true); + llvm::SmallVector 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 T; + std::copy(Transform.begin(), Transform.end(), T.begin()); + G.Transform = T; + } } void MappingTraits::mapping( diff --git a/test/Feature/InlineRT/geometry-transform.test b/test/Feature/InlineRT/geometry-transform.test new file mode 100644 index 000000000..1db616755 --- /dev/null +++ b/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 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 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