Fix QuerySet ownership of ComputePass (#5671)

* add new tests for checking on query set lifetime * Fix ownership management of query sets on compute passes for write_timestamp, timestamp_writes (on desc) and pipeline statistic queries * changelog entry
gfx-rs · Jun 4, 2024 · 9a27ba5 · 9a27ba5
1 parent d258d6c
commit 9a27ba5
Show file tree

Hide file tree

Showing 6 changed files with 338 additions and 206 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -58,7 +58,7 @@ fn independent_cpass<'enc>(encoder: &'enc mut wgpu::CommandEncoder) -> wgpu::Com
 This is very useful for library authors, but opens up an easy way for incorrect use, so use with care.
 `forget_lifetime` is zero overhead and has no side effects on pass recording.
 
-By @wumpf in [#5569](https://github.com/gfx-rs/wgpu/pull/5569), [#5575](https://github.com/gfx-rs/wgpu/pull/5575), [#5620](https://github.com/gfx-rs/wgpu/pull/5620), [#5768](https://github.com/gfx-rs/wgpu/pull/5768) (together with @kpreid).
+By @wumpf in [#5569](https://github.com/gfx-rs/wgpu/pull/5569), [#5575](https://github.com/gfx-rs/wgpu/pull/5575), [#5620](https://github.com/gfx-rs/wgpu/pull/5620), [#5768](https://github.com/gfx-rs/wgpu/pull/5768) (together with @kpreid), [#5671](https://github.com/gfx-rs/wgpu/pull/5671).
 
 #### Querying shader compilation errors
 
@@ -116,6 +116,7 @@ By @stefnotch in [#5410](https://github.com/gfx-rs/wgpu/pull/5410)
 #### General
 
 - Ensure render pipelines have at least 1 target. By @ErichDonGubler in [#5715](https://github.com/gfx-rs/wgpu/pull/5715)
+- `wgpu::ComputePass` now internally takes ownership of `QuerySet` for both `wgpu::ComputePassTimestampWrites` as well as timestamp writes and statistics query, fixing crashes when destroying `QuerySet` before ending the pass. By @wumpf in [#5671](https://github.com/gfx-rs/wgpu/pull/5671)
 
 #### Metal
 

diff --git a/tests/tests/compute_pass_ownership.rs b/tests/tests/compute_pass_ownership.rs
@@ -1,9 +1,5 @@
 //! Tests that compute passes take ownership of resources that are associated with.
 //! I.e. once a resource is passed in to a compute pass, it can be dropped.
-//!
-//! TODO: Also should test resource ownership for:
-//!       * write_timestamp
-//!       * begin_pipeline_statistics_query
 
 use std::num::NonZeroU64;
 
@@ -36,15 +32,10 @@ async fn compute_pass_resource_ownership(ctx: TestingContext) {
 
     let mut encoder = ctx
         .device
-        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
-            label: Some("encoder"),
-        });
+        .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
 
     {
-        let mut cpass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
-            label: Some("compute_pass"),
-            timestamp_writes: None, // TODO: See description above, we should test this as well once we lift the lifetime bound.
-        });
+        let mut cpass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
         cpass.set_pipeline(&pipeline);
         cpass.set_bind_group(0, &bind_group, &[]);
         cpass.dispatch_workgroups_indirect(&indirect_buffer, 0);
@@ -58,18 +49,115 @@ async fn compute_pass_resource_ownership(ctx: TestingContext) {
             .panic_on_timeout();
     }
 
-    // Ensure that the compute pass still executed normally.
-    encoder.copy_buffer_to_buffer(&gpu_buffer, 0, &cpu_buffer, 0, buffer_size);
-    ctx.queue.submit([encoder.finish()]);
-    cpu_buffer.slice(..).map_async(wgpu::MapMode::Read, |_| ());
-    ctx.async_poll(wgpu::Maintain::wait())
-        .await
-        .panic_on_timeout();
+    assert_compute_pass_executed_normally(encoder, gpu_buffer, cpu_buffer, buffer_size, ctx).await;
+}
 
-    let data = cpu_buffer.slice(..).get_mapped_range();
+#[gpu_test]
+static COMPUTE_PASS_QUERY_SET_OWNERSHIP_PIPELINE_STATISTICS: GpuTestConfiguration =
+    GpuTestConfiguration::new()
+        .parameters(
+            TestParameters::default()
+                .test_features_limits()
+                .features(wgpu::Features::PIPELINE_STATISTICS_QUERY),
+        )
+        .run_async(compute_pass_query_set_ownership_pipeline_statistics);
+
+async fn compute_pass_query_set_ownership_pipeline_statistics(ctx: TestingContext) {
+    let ResourceSetup {
+        gpu_buffer,
+        cpu_buffer,
+        buffer_size,
+        indirect_buffer: _,
+        bind_group,
+        pipeline,
+    } = resource_setup(&ctx);
 
-    let floats: &[f32] = bytemuck::cast_slice(&data);
-    assert_eq!(floats, [2.0, 4.0, 6.0, 8.0]);
+    let query_set = ctx.device.create_query_set(&wgpu::QuerySetDescriptor {
+        label: Some("query_set"),
+        ty: wgpu::QueryType::PipelineStatistics(
+            wgpu::PipelineStatisticsTypes::COMPUTE_SHADER_INVOCATIONS,
+        ),
+        count: 1,
+    });
+
+    let mut encoder = ctx
+        .device
+        .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
+
+    {
+        let mut cpass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
+        cpass.set_pipeline(&pipeline);
+        cpass.set_bind_group(0, &bind_group, &[]);
+        cpass.begin_pipeline_statistics_query(&query_set, 0);
+        cpass.dispatch_workgroups(1, 1, 1);
+        cpass.end_pipeline_statistics_query();
+
+        // Drop the query set. Then do a device poll to make sure it's not dropped too early, no matter what.
+        drop(query_set);
+        ctx.async_poll(wgpu::Maintain::wait())
+            .await
+            .panic_on_timeout();
+    }
+
+    assert_compute_pass_executed_normally(encoder, gpu_buffer, cpu_buffer, buffer_size, ctx).await;
+}
+
+#[gpu_test]
+static COMPUTE_PASS_QUERY_TIMESTAMPS: GpuTestConfiguration =
+    GpuTestConfiguration::new()
+        .parameters(TestParameters::default().test_features_limits().features(
+            wgpu::Features::TIMESTAMP_QUERY | wgpu::Features::TIMESTAMP_QUERY_INSIDE_PASSES,
+        ))
+        .run_async(compute_pass_query_timestamps);
+
+async fn compute_pass_query_timestamps(ctx: TestingContext) {
+    let ResourceSetup {
+        gpu_buffer,
+        cpu_buffer,
+        buffer_size,
+        indirect_buffer: _,
+        bind_group,
+        pipeline,
+    } = resource_setup(&ctx);
+
+    let query_set_timestamp_writes = ctx.device.create_query_set(&wgpu::QuerySetDescriptor {
+        label: Some("query_set_timestamp_writes"),
+        ty: wgpu::QueryType::Timestamp,
+        count: 2,
+    });
+    let query_set_write_timestamp = ctx.device.create_query_set(&wgpu::QuerySetDescriptor {
+        label: Some("query_set_write_timestamp"),
+        ty: wgpu::QueryType::Timestamp,
+        count: 1,
+    });
+
+    let mut encoder = ctx
+        .device
+        .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
+
+    {
+        let mut cpass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
+            label: Some("compute_pass"),
+            timestamp_writes: Some(wgpu::ComputePassTimestampWrites {
+                query_set: &query_set_timestamp_writes,
+                beginning_of_pass_write_index: Some(0),
+                end_of_pass_write_index: Some(1),
+            }),
+        });
+        cpass.set_pipeline(&pipeline);
+        cpass.set_bind_group(0, &bind_group, &[]);
+        cpass.write_timestamp(&query_set_write_timestamp, 0);
+        cpass.dispatch_workgroups(1, 1, 1);
+
+        // Drop the query sets. Then do a device poll to make sure they're not dropped too early, no matter what.
+        drop(query_set_timestamp_writes);
+        drop(query_set_write_timestamp);
+        ctx.async_poll(wgpu::Maintain::wait())
+            .await
+            .panic_on_timeout();
+    }
+
+    assert_compute_pass_executed_normally(encoder, gpu_buffer, cpu_buffer, buffer_size, ctx).await;
 }
 
 #[gpu_test]
@@ -89,9 +177,7 @@ async fn compute_pass_keep_encoder_alive(ctx: TestingContext) {
 
     let mut encoder = ctx
         .device
-        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
-            label: Some("encoder"),
-        });
+        .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
 
     let cpass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
         label: Some("compute_pass"),
@@ -119,6 +205,26 @@ async fn compute_pass_keep_encoder_alive(ctx: TestingContext) {
     valid(&ctx.device, || drop(cpass));
 }
 
+async fn assert_compute_pass_executed_normally(
+    mut encoder: wgpu::CommandEncoder,
+    gpu_buffer: wgpu::Buffer,
+    cpu_buffer: wgpu::Buffer,
+    buffer_size: u64,
+    ctx: TestingContext,
+) {
+    encoder.copy_buffer_to_buffer(&gpu_buffer, 0, &cpu_buffer, 0, buffer_size);
+    ctx.queue.submit([encoder.finish()]);
+    cpu_buffer.slice(..).map_async(wgpu::MapMode::Read, |_| ());
+    ctx.async_poll(wgpu::Maintain::wait())
+        .await
+        .panic_on_timeout();
+
+    let data = cpu_buffer.slice(..).get_mapped_range();
+
+    let floats: &[f32] = bytemuck::cast_slice(&data);
+    assert_eq!(floats, [2.0, 4.0, 6.0, 8.0]);
+}
+
 // Setup ------------------------------------------------------------
 
 struct ResourceSetup {