From 658052885a31f4fa3c69336723a65eed9cb9018f Mon Sep 17 00:00:00 2001
From: Vecvec <130132884+Vecvec@users.noreply.github.com>
Date: Tue, 7 Jan 2025 08:01:50 +0000
Subject: [PATCH] Use `transform_buffer_offset` when initialising
 transform_buffer. (#6864)

* use `transform_buffer_offset` instead of `index_buffer_offset` in part of the initialisation of transform_buffer.

* Format.

* Changelog.

* Only unwrap transform alignment once.
---
 CHANGELOG.md                         |  1 +
 tests/tests/ray_tracing/as_build.rs  | 87 ++++++++++++++++++++++++++++
 wgpu-core/src/command/ray_tracing.rs | 11 ++--
 3 files changed, 95 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 12aa7ce6a0..c919e78bd6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -213,6 +213,7 @@ By @wumpf in [#6849](https://github.com/gfx-rs/wgpu/pull/6849).
 - Reduced the overhead of command buffer validation. By @nical in [#6721](https://github.com/gfx-rs/wgpu/pull/6721).
 - Set index type to NONE in `get_acceleration_structure_build_sizes`. By @Vecvec in [#6802](https://github.com/gfx-rs/wgpu/pull/6802).
 - Fix `wgpu-info` not showing dx12 adapters. By @wumpf in [#6844](https://github.com/gfx-rs/wgpu/pull/6844).
+- Use `transform_buffer_offset` when initialising `transform_buffer`. By @Vecvec in [#6864](https://github.com/gfx-rs/wgpu/pull/6864).
 
 #### Naga
 
diff --git a/tests/tests/ray_tracing/as_build.rs b/tests/tests/ray_tracing/as_build.rs
index 8928b84c33..a75b280724 100644
--- a/tests/tests/ray_tracing/as_build.rs
+++ b/tests/tests/ray_tracing/as_build.rs
@@ -328,3 +328,90 @@ fn empty_build(ctx: TestingContext) {
     ctx.queue
         .submit([encoder_safe.finish(), encoder_unsafe.finish()]);
 }
+
+#[gpu_test]
+static BUILD_WITH_TRANSFORM: GpuTestConfiguration = GpuTestConfiguration::new()
+    .parameters(
+        TestParameters::default()
+            .test_features_limits()
+            .features(wgpu::Features::EXPERIMENTAL_RAY_TRACING_ACCELERATION_STRUCTURE)
+            // https://github.com/gfx-rs/wgpu/issues/6727
+            .skip(FailureCase::backend_adapter(wgpu::Backends::VULKAN, "AMD")),
+    )
+    .run_sync(build_with_transform);
+
+fn build_with_transform(ctx: TestingContext) {
+    let vertices = ctx.device.create_buffer_init(&BufferInitDescriptor {
+        label: None,
+        contents: &[0; mem::size_of::<[[f32; 3]; 3]>()],
+        usage: BufferUsages::BLAS_INPUT,
+    });
+
+    let transform = ctx
+        .device
+        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
+            label: Some("Vertex Buffer"),
+            contents: bytemuck::cast_slice(&[
+                1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0,
+            ]),
+            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::BLAS_INPUT,
+        });
+
+    let blas_size = BlasTriangleGeometrySizeDescriptor {
+        vertex_format: VertexFormat::Float32x3,
+        vertex_count: 3,
+        index_format: None,
+        index_count: None,
+        flags: AccelerationStructureGeometryFlags::empty(),
+    };
+
+    let blas = ctx.device.create_blas(
+        &CreateBlasDescriptor {
+            label: Some("BLAS"),
+            flags: AccelerationStructureFlags::PREFER_FAST_TRACE,
+            update_mode: AccelerationStructureUpdateMode::Build,
+        },
+        BlasGeometrySizeDescriptors::Triangles {
+            descriptors: vec![blas_size.clone()],
+        },
+    );
+
+    let tlas = ctx.device.create_tlas(&CreateTlasDescriptor {
+        label: Some("TLAS"),
+        max_instances: 1,
+        flags: AccelerationStructureFlags::PREFER_FAST_TRACE,
+        update_mode: AccelerationStructureUpdateMode::Build,
+    });
+
+    let mut tlas_package = TlasPackage::new(tlas);
+    tlas_package[0] = Some(TlasInstance::new(
+        &blas,
+        [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
+        0,
+        0xFF,
+    ));
+
+    let mut encoder_build = ctx
+        .device
+        .create_command_encoder(&CommandEncoderDescriptor {
+            label: Some("BUILD 1"),
+        });
+
+    encoder_build.build_acceleration_structures(
+        [&BlasBuildEntry {
+            blas: &blas,
+            geometry: BlasGeometries::TriangleGeometries(vec![BlasTriangleGeometry {
+                size: &blas_size,
+                vertex_buffer: &vertices,
+                first_vertex: 0,
+                vertex_stride: mem::size_of::<[f32; 3]>() as BufferAddress,
+                index_buffer: None,
+                index_buffer_offset: None,
+                transform_buffer: Some(&transform),
+                transform_buffer_offset: Some(0),
+            }]),
+        }],
+        [&tlas_package],
+    );
+    ctx.queue.submit([encoder_build.finish()]);
+}
diff --git a/wgpu-core/src/command/ray_tracing.rs b/wgpu-core/src/command/ray_tracing.rs
index 65922524f9..9395c20fc1 100644
--- a/wgpu-core/src/command/ray_tracing.rs
+++ b/wgpu-core/src/command/ray_tracing.rs
@@ -1169,24 +1169,27 @@ fn iter_buffers<'a, 'b>(
             {
                 input_barriers.push(barrier);
             }
-            if mesh.transform_buffer_offset.unwrap() % wgt::TRANSFORM_BUFFER_ALIGNMENT != 0 {
+
+            let offset = mesh.transform_buffer_offset.unwrap();
+
+            if offset % wgt::TRANSFORM_BUFFER_ALIGNMENT != 0 {
                 return Err(
                     BuildAccelerationStructureError::UnalignedTransformBufferOffset(
                         transform_buffer.error_ident(),
                     ),
                 );
             }
-            if transform_buffer.size < 48 + mesh.transform_buffer_offset.unwrap() {
+            if transform_buffer.size < 48 + offset {
                 return Err(BuildAccelerationStructureError::InsufficientBufferSize(
                     transform_buffer.error_ident(),
                     transform_buffer.size,
-                    48 + mesh.transform_buffer_offset.unwrap(),
+                    48 + offset,
                 ));
             }
             cmd_buf_data.buffer_memory_init_actions.extend(
                 transform_buffer.initialization_status.read().create_action(
                     transform_buffer,
-                    mesh.transform_buffer_offset.unwrap()..(mesh.index_buffer_offset.unwrap() + 48),
+                    offset..(offset + 48),
                     MemoryInitKind::NeedsInitializedMemory,
                 ),
             );