ensure safety of indirect dispatch

by injecting a compute shader that validates the content of the indirect buffer

tests/tests/dispatch_workgroups_indirect.rs

@@ -0,0 +1,116 @@
+use wgpu_test::{gpu_test, GpuTestConfiguration, TestParameters};
+
+const SHADER_SRC: &str = "
+@group(0) @binding(0)
+var<storage, read_write> out: u32;
+
+@compute @workgroup_size(1)
+fn main() {
+    out = 1u;
+}
+";
+
+#[gpu_test]
+static CHECK_NUM_WORKGROUPS: GpuTestConfiguration = GpuTestConfiguration::new()
+    .parameters(TestParameters::default().test_features_limits())
+    .run_async(|ctx| async move {
+        let module = ctx
+            .device
+            .create_shader_module(wgpu::ShaderModuleDescriptor {
+                label: None,
+                source: wgpu::ShaderSource::Wgsl(SHADER_SRC.into()),
+            });
+
+        let pipeline = ctx
+            .device
+            .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
+                label: None,
+                layout: None,
+                module: &module,
+                entry_point: "main",
+                compilation_options: Default::default(),
+                cache: None,
+            });
+
+        let out_buffer = ctx.device.create_buffer(&wgpu::BufferDescriptor {
+            label: None,
+            size: 4,
+            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
+            mapped_at_creation: false,
+        });
+
+        let readback_buffer = ctx.device.create_buffer(&wgpu::BufferDescriptor {
+            label: None,
+            size: 4,
+            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
+            mapped_at_creation: false,
+        });
+
+        let indirect_buffer = ctx.device.create_buffer(&wgpu::BufferDescriptor {
+            label: None,
+            size: 12,
+            usage: wgpu::BufferUsages::INDIRECT
+                | wgpu::BufferUsages::COPY_DST
+                | wgpu::BufferUsages::UNIFORM,
+            mapped_at_creation: false,
+        });
+
+        let max = ctx.adapter.limits().max_compute_workgroups_per_dimension;
+        ctx.queue
+            .write_buffer(&indirect_buffer, 0, bytemuck::bytes_of(&[max + 1, 1, 1]));
+
+        let bind_group = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
+            label: None,
+            layout: &pipeline.get_bind_group_layout(0),
+            entries: &[wgpu::BindGroupEntry {
+                binding: 0,
+                resource: out_buffer.as_entire_binding(),
+            }],
+        });
+
+        let mut encoder = ctx
+            .device
+            .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
+
+        {
+            let mut compute_pass =
+                encoder.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
+            compute_pass.set_pipeline(&pipeline);
+            compute_pass.set_bind_group(0, &bind_group, &[]);
+            compute_pass.dispatch_workgroups_indirect(&indirect_buffer, 0);
+        }
+
+        encoder.copy_buffer_to_buffer(&out_buffer, 0, &readback_buffer, 0, 4);
+
+        ctx.queue.submit(Some(encoder.finish()));
+
+        readback_buffer
+            .slice(..)
+            .map_async(wgpu::MapMode::Read, |_| {});
+
+        ctx.async_poll(wgpu::Maintain::wait())
+            .await
+            .panic_on_timeout();
+
+        let view = readback_buffer.slice(..).get_mapped_range();
+        // Make sure the dispatch was discarded
+        assert!(view.iter().all(|v| *v == 0));
+
+        // Test that unsetting the bind group works properly
+        {
+            ctx.device.push_error_scope(wgpu::ErrorFilter::Validation);
+            let mut encoder = ctx
+                .device
+                .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
+            {
+                let mut compute_pass =
+                    encoder.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
+                compute_pass.set_pipeline(&pipeline);
+                compute_pass.dispatch_workgroups_indirect(&indirect_buffer, 0);
+            }
+            let _ = encoder.finish();
+            let error = pollster::block_on(ctx.device.pop_error_scope());
+            assert!(error.map_or(false, |error| format!("{error}")
+                .contains("Expected bind group is missing")));
+        }
+    });

tests/tests/root.rs

@@ -14,6 +14,7 @@ mod clear_texture;
 mod compute_pass_resource_ownership;
 mod create_surface_error;
 mod device;
+mod dispatch_workgroups_indirect;
 mod encoder;
 mod external_texture;
 mod float32_filterable;

wgpu-core/src/command/bind.rs

@@ -131,7 +131,7 @@ mod compat {
                         diff.push(format!("Expected {expected_bgl_type} bind group layout, got {assigned_bgl_type}"))
                     }
                 } else {
-                    diff.push("Assigned bind group layout not found (internal error)".to_owned());
+                    diff.push("Expected bind group is missing".to_owned());
                 }
             } else {
                 diff.push("Expected bind group layout not found (internal error)".to_owned());
@@ -191,6 +191,10 @@ mod compat {
             self.make_range(index)
         }
 
+        pub fn unassign(&mut self, index: usize) {
+            self.entries[index].assigned = None;
+        }
+
         pub fn list_active(&self) -> impl Iterator<Item = usize> + '_ {
             self.entries
                 .iter()
@@ -358,6 +362,14 @@ impl<A: HalApi> Binder<A> {
         &self.payloads[bind_range]
     }
 
+    pub(super) fn unassign_group(&mut self, index: usize) {
+        log::trace!("\tBinding [{}] = null", index);
+
+        self.payloads[index].reset();
+
+        self.manager.unassign(index);
+    }
+
     pub(super) fn list_active<'a>(&'a self) -> impl Iterator<Item = &'a Arc<BindGroup<A>>> + '_ {
         let payloads = &self.payloads;
         self.manager

wgpu-core/src/command/compute.rs

@@ -146,6 +146,20 @@ pub enum ComputePassErrorInner {
     MissingFeatures(#[from] MissingFeatures),
     #[error(transparent)]
     MissingDownlevelFlags(#[from] MissingDownlevelFlags),
+    #[error(transparent)]
+    IndirectValidation(#[from] ComputePassIndirectValidationError),
+}
+
+#[derive(Clone, Debug, Error)]
+pub enum ComputePassIndirectValidationError {
+    #[error(transparent)]
+    ValidationPipeline(
+        #[from] crate::pipeline::CreateDispatchWorkgroupsIndirectValidationPipelineError,
+    ),
+    #[error(transparent)]
+    Buffer(#[from] crate::resource::CreateBufferError),
+    #[error(transparent)]
+    BindGroup(#[from] crate::binding_model::CreateBindGroupError),
 }
 
 impl PrettyError for ComputePassErrorInner {
@@ -283,9 +297,26 @@ impl Global {
         &self,
         pass: &ComputePass<A>,
     ) -> Result<(), ComputePassError> {
+        let mut base = pass.base.as_ref();
+
+        let new_commands = base
+            .commands
+            .iter()
+            .any(|cmd| matches!(cmd, ArcComputeCommand::DispatchIndirect { .. }))
+            .then(|| {
+                self.command_encoder_inject_dispatch_workgroups_indirect_validation(
+                    pass.parent_id,
+                    base.commands,
+                )
+            })
+            .transpose()?;
+        if let Some(new_commands) = new_commands.as_ref() {
+            base.commands = new_commands;
+        }
+
         self.command_encoder_run_compute_pass_impl(
             pass.parent_id,
-            pass.base.as_ref(),
+            base,
             pass.timestamp_writes.as_ref(),
         )
     }
@@ -515,6 +546,20 @@ impl Global {
                         }
                     }
                 }
+                ArcComputeCommand::UnsetBindGroup { index } => {
+                    let scope = PassErrorScope::UnsetBindGroup(*index);
+
+                    let max_bind_groups = cmd_buf.limits.max_bind_groups;
+                    if *index >= max_bind_groups {
+                        return Err(ComputePassErrorInner::BindGroupIndexOutOfRange {
+                            index: *index,
+                            max: max_bind_groups,
+                        })
+                        .map_pass_err(scope);
+                    }
+
+                    state.binder.unassign_group(*index as usize);
+                }
                 ArcComputeCommand::SetPipeline(pipeline) => {
                     let pipeline_id = pipeline.as_info().id();
                     let scope = PassErrorScope::SetPipelineCompute(pipeline_id);
@@ -823,6 +868,166 @@ impl Global {
 
         Ok(())
     }
+
+    fn command_encoder_inject_dispatch_workgroups_indirect_validation<A: HalApi>(
+        &self,
+        encoder_id: id::CommandEncoderId,
+        commands: &[ArcComputeCommand<A>],
+    ) -> Result<Vec<ArcComputeCommand<A>>, ComputePassError> {
+        profiling::scope!("CommandEncoder::inject_dispatch_workgroups_indirect_validation");
+        let scope = PassErrorScope::Pass(encoder_id);
+
+        let hub = A::hub(self);
+
+        let cmd_buf: Arc<CommandBuffer<A>> =
+            CommandBuffer::get_encoder(hub, encoder_id).map_pass_err(scope)?;
+        let device = &cmd_buf.device;
+        if !device.is_valid() {
+            return Err(ComputePassErrorInner::InvalidDevice(
+                cmd_buf.device.as_info().id(),
+            ))
+            .map_pass_err(scope);
+        }
+        let device_id = device.as_info().id();
+
+        let mut new_commands = Vec::with_capacity(commands.len());
+        let mut current_pipeline = None;
+        let mut current_bind_group_0 = None;
+
+        for command in commands {
+            match command {
+                ArcComputeCommand::SetBindGroup {
+                    index: 0,
+                    num_dynamic_offsets,
+                    bind_group,
+                } => {
+                    current_bind_group_0 = Some((*num_dynamic_offsets, bind_group.clone()));
+                    new_commands.push(command.clone());
+                }
+                ArcComputeCommand::SetPipeline(pipeline) => {
+                    current_pipeline = Some(pipeline.clone());
+                    new_commands.push(command.clone());
+                }
+                ArcComputeCommand::DispatchIndirect { buffer, offset } => {
+                    // if there is no pipeline set, don't inject the validation commands as we will error anyway
+                    if let Some(original_pipeline) = current_pipeline.clone() {
+                        let (validation_pipeline_id, validation_bgl_id) = self
+                        .device_get_or_create_dispatch_workgroups_indirect_validation_pipeline::<A>(
+                            device_id,
+                        )
+                        .map_err(ComputePassIndirectValidationError::ValidationPipeline)
+                        .map_pass_err(scope)?;
+
+                        let (dst_buffer_id, error) = self.device_create_buffer::<A>(
+                            device_id,
+                            &crate::resource::BufferDescriptor {
+                                label: None,
+                                size: 4 * 3,
+                                usage: wgt::BufferUsages::INDIRECT | wgt::BufferUsages::STORAGE,
+                                mapped_at_creation: false,
+                            },
+                            None,
+                        );
+                        if let Some(error) = error {
+                            return Err(ComputePassIndirectValidationError::Buffer(error))
+                                .map_pass_err(scope)?;
+                        }
+
+                        let (bind_group_id, error) = self.device_create_bind_group::<A>(
+                            device_id,
+                            &crate::binding_model::BindGroupDescriptor {
+                                label: None,
+                                layout: validation_bgl_id,
+                                entries: std::borrow::Cow::Borrowed(&[
+                                    crate::binding_model::BindGroupEntry {
+                                        binding: 0,
+                                        resource: crate::binding_model::BindingResource::Buffer(
+                                            crate::binding_model::BufferBinding {
+                                                buffer_id: buffer.as_info().id(),
+                                                offset: *offset,
+                                                size: Some(
+                                                    std::num::NonZeroU64::new(4 * 3).unwrap(),
+                                                ),
+                                            },
+                                        ),
+                                    },
+                                    crate::binding_model::BindGroupEntry {
+                                        binding: 1,
+                                        resource: crate::binding_model::BindingResource::Buffer(
+                                            crate::binding_model::BufferBinding {
+                                                buffer_id: dst_buffer_id,
+                                                offset: 0,
+                                                size: Some(
+                                                    std::num::NonZeroU64::new(4 * 3).unwrap(),
+                                                ),
+                                            },
+                                        ),
+                                    },
+                                ]),
+                            },
+                            None,
+                        );
+                        if let Some(error) = error {
+                            return Err(ComputePassIndirectValidationError::BindGroup(error))
+                                .map_pass_err(scope)?;
+                        }
+
+                        let validation_pipeline = hub
+                            .compute_pipelines
+                            .read()
+                            .get_owned(validation_pipeline_id)
+                            .map_err(|_| {
+                                ComputePassErrorInner::InvalidPipeline(validation_pipeline_id)
+                            })
+                            .map_pass_err(scope)?;
+
+                        let bind_group = hub
+                            .bind_groups
+                            .read()
+                            .get_owned(bind_group_id)
+                            .map_err(|_| ComputePassErrorInner::InvalidBindGroup(0))
+                            .map_pass_err(scope)?;
+
+                        let dst_buffer = hub
+                            .buffers
+                            .read()
+                            .get_owned(dst_buffer_id)
+                            .map_err(|_| ComputePassErrorInner::InvalidBuffer(dst_buffer_id))
+                            .map_pass_err(scope)?;
+
+                        new_commands.push(ArcComputeCommand::SetPipeline(validation_pipeline));
+                        new_commands.push(ArcComputeCommand::SetBindGroup {
+                            index: 0,
+                            num_dynamic_offsets: 0,
+                            bind_group,
+                        });
+                        new_commands.push(ArcComputeCommand::Dispatch([1, 1, 1]));
+
+                        new_commands.push(ArcComputeCommand::SetPipeline(original_pipeline));
+                        if let Some((num_dynamic_offsets, bind_group)) =
+                            current_bind_group_0.clone()
+                        {
+                            new_commands.push(ArcComputeCommand::SetBindGroup {
+                                index: 0,
+                                num_dynamic_offsets,
+                                bind_group,
+                            });
+                        } else {
+                            new_commands.push(ArcComputeCommand::UnsetBindGroup { index: 0 });
+                        }
+                        new_commands.push(ArcComputeCommand::DispatchIndirect {
+                            buffer: dst_buffer,
+                            offset: 0,
+                        });
+                    } else {
+                        new_commands.push(command.clone())
+                    }
+                }
+                command => new_commands.push(command.clone()),
+            }
+        }
+        Ok(new_commands)
+    }
 }
 
 // Recording a compute pass.