Skip to content

Commit

Permalink
ensure safety of indirect dispatch
Browse files Browse the repository at this point in the history
by injecting a compute shader that validates the content of the indirect buffer
  • Loading branch information
teoxoy committed May 21, 2024
1 parent 4902e47 commit a5bebb0
Show file tree
Hide file tree
Showing 10 changed files with 517 additions and 2 deletions.
125 changes: 125 additions & 0 deletions tests/tests/dispatch_workgroups_indirect.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,125 @@
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()
.downlevel_flags(
wgpu::DownlevelFlags::COMPUTE_SHADERS | wgpu::DownlevelFlags::INDIRECT_EXECUTION,
)
.limits(wgpu::Limits {
max_compute_workgroups_per_dimension: 10,
..wgpu::Limits::downlevel_defaults()
}),
)
.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.device.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")));
}
});
1 change: 1 addition & 0 deletions tests/tests/root.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
14 changes: 13 additions & 1 deletion wgpu-core/src/command/bind.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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
Expand Down
Loading

0 comments on commit a5bebb0

Please sign in to comment.