Revision control
Copy as Markdown
Other Tools
use std::mem::size_of;
use std::num::NonZeroU64;
use thiserror::Error;
use crate::{
device::DeviceError,
pipeline::{CreateComputePipelineError, CreateShaderModuleError},
};
#[derive(Clone, Debug, Error)]
#[non_exhaustive]
pub enum CreateDispatchIndirectValidationPipelineError {
#[error(transparent)]
DeviceError(#[from] DeviceError),
#[error(transparent)]
ShaderModule(#[from] CreateShaderModuleError),
#[error(transparent)]
ComputePipeline(#[from] CreateComputePipelineError),
}
/// This machinery requires the following limits:
///
/// - max_bind_groups: 2,
/// - max_dynamic_storage_buffers_per_pipeline_layout: 1,
/// - max_storage_buffers_per_shader_stage: 2,
/// - max_storage_buffer_binding_size: 3 * min_storage_buffer_offset_alignment,
/// - max_push_constant_size: 4,
/// - max_compute_invocations_per_workgroup 1
///
/// These are all indirectly satisfied by `DownlevelFlags::INDIRECT_EXECUTION`, which is also
/// required for this module's functionality to work.
#[derive(Debug)]
pub struct IndirectValidation {
module: Box<dyn hal::DynShaderModule>,
dst_bind_group_layout: Box<dyn hal::DynBindGroupLayout>,
src_bind_group_layout: Box<dyn hal::DynBindGroupLayout>,
pipeline_layout: Box<dyn hal::DynPipelineLayout>,
pipeline: Box<dyn hal::DynComputePipeline>,
dst_buffer: Box<dyn hal::DynBuffer>,
dst_bind_group: Box<dyn hal::DynBindGroup>,
}
pub struct Params<'a> {
pub pipeline_layout: &'a dyn hal::DynPipelineLayout,
pub pipeline: &'a dyn hal::DynComputePipeline,
pub dst_buffer: &'a dyn hal::DynBuffer,
pub dst_bind_group: &'a dyn hal::DynBindGroup,
pub aligned_offset: u64,
pub offset_remainder: u64,
}
impl IndirectValidation {
pub fn new(
device: &dyn hal::DynDevice,
limits: &wgt::Limits,
) -> Result<Self, CreateDispatchIndirectValidationPipelineError> {
let max_compute_workgroups_per_dimension = limits.max_compute_workgroups_per_dimension;
let src = format!(
"
@group(0) @binding(0)
var<storage, read_write> dst: array<u32, 6>;
@group(1) @binding(0)
var<storage, read> src: array<u32>;
struct OffsetPc {{
inner: u32,
}}
var<push_constant> offset: OffsetPc;
@compute @workgroup_size(1)
fn main() {{
let src = vec3(src[offset.inner], src[offset.inner + 1], src[offset.inner + 2]);
let max_compute_workgroups_per_dimension = {max_compute_workgroups_per_dimension}u;
if (
src.x > max_compute_workgroups_per_dimension ||
src.y > max_compute_workgroups_per_dimension ||
src.z > max_compute_workgroups_per_dimension
) {{
dst = array(0u, 0u, 0u, 0u, 0u, 0u);
}} else {{
dst = array(src.x, src.y, src.z, src.x, src.y, src.z);
}}
}}
"
);
// SAFETY: The value we are passing to `new_unchecked` is not zero, so this is safe.
const SRC_BUFFER_SIZE: NonZeroU64 =
unsafe { NonZeroU64::new_unchecked(size_of::<u32>() as u64 * 3) };
// SAFETY: The value we are passing to `new_unchecked` is not zero, so this is safe.
const DST_BUFFER_SIZE: NonZeroU64 = unsafe {
NonZeroU64::new_unchecked(
SRC_BUFFER_SIZE.get() * 2, // From above: `dst: array<u32, 6>`
)
};
let module = naga::front::wgsl::parse_str(&src).map_err(|inner| {
CreateShaderModuleError::Parsing(naga::error::ShaderError {
source: src.clone(),
label: None,
inner: Box::new(inner),
})
})?;
let info = crate::device::create_validator(
wgt::Features::PUSH_CONSTANTS,
wgt::DownlevelFlags::empty(),
naga::valid::ValidationFlags::all(),
)
.validate(&module)
.map_err(|inner| {
CreateShaderModuleError::Validation(naga::error::ShaderError {
source: src,
label: None,
inner: Box::new(inner),
})
})?;
let hal_shader = hal::ShaderInput::Naga(hal::NagaShader {
module: std::borrow::Cow::Owned(module),
info,
debug_source: None,
});
let hal_desc = hal::ShaderModuleDescriptor {
label: None,
runtime_checks: false,
};
let module =
unsafe { device.create_shader_module(&hal_desc, hal_shader) }.map_err(|error| {
match error {
hal::ShaderError::Device(error) => {
CreateShaderModuleError::Device(DeviceError::from_hal(error))
}
hal::ShaderError::Compilation(ref msg) => {
log::error!("Shader error: {}", msg);
CreateShaderModuleError::Generation
}
}
})?;
let dst_bind_group_layout_desc = hal::BindGroupLayoutDescriptor {
label: None,
flags: hal::BindGroupLayoutFlags::empty(),
entries: &[wgt::BindGroupLayoutEntry {
binding: 0,
visibility: wgt::ShaderStages::COMPUTE,
ty: wgt::BindingType::Buffer {
ty: wgt::BufferBindingType::Storage { read_only: false },
has_dynamic_offset: false,
min_binding_size: Some(DST_BUFFER_SIZE),
},
count: None,
}],
};
let dst_bind_group_layout = unsafe {
device
.create_bind_group_layout(&dst_bind_group_layout_desc)
.map_err(DeviceError::from_hal)?
};
let src_bind_group_layout_desc = hal::BindGroupLayoutDescriptor {
label: None,
flags: hal::BindGroupLayoutFlags::empty(),
entries: &[wgt::BindGroupLayoutEntry {
binding: 0,
visibility: wgt::ShaderStages::COMPUTE,
ty: wgt::BindingType::Buffer {
ty: wgt::BufferBindingType::Storage { read_only: true },
has_dynamic_offset: true,
min_binding_size: Some(SRC_BUFFER_SIZE),
},
count: None,
}],
};
let src_bind_group_layout = unsafe {
device
.create_bind_group_layout(&src_bind_group_layout_desc)
.map_err(DeviceError::from_hal)?
};
let pipeline_layout_desc = hal::PipelineLayoutDescriptor {
label: None,
flags: hal::PipelineLayoutFlags::FIRST_VERTEX_INSTANCE,
bind_group_layouts: &[
dst_bind_group_layout.as_ref(),
src_bind_group_layout.as_ref(),
],
push_constant_ranges: &[wgt::PushConstantRange {
stages: wgt::ShaderStages::COMPUTE,
range: 0..4,
}],
};
let pipeline_layout = unsafe {
device
.create_pipeline_layout(&pipeline_layout_desc)
.map_err(DeviceError::from_hal)?
};
let pipeline_desc = hal::ComputePipelineDescriptor {
label: None,
layout: pipeline_layout.as_ref(),
stage: hal::ProgrammableStage {
module: module.as_ref(),
entry_point: "main",
constants: &Default::default(),
zero_initialize_workgroup_memory: false,
},
cache: None,
};
let pipeline =
unsafe { device.create_compute_pipeline(&pipeline_desc) }.map_err(|err| match err {
hal::PipelineError::Device(error) => {
CreateComputePipelineError::Device(DeviceError::from_hal(error))
}
hal::PipelineError::Linkage(_stages, msg) => {
CreateComputePipelineError::Internal(msg)
}
hal::PipelineError::EntryPoint(_stage) => CreateComputePipelineError::Internal(
crate::device::ENTRYPOINT_FAILURE_ERROR.to_string(),
),
hal::PipelineError::PipelineConstants(_, error) => {
CreateComputePipelineError::PipelineConstants(error)
}
})?;
let dst_buffer_desc = hal::BufferDescriptor {
label: None,
size: DST_BUFFER_SIZE.get(),
usage: hal::BufferUses::INDIRECT | hal::BufferUses::STORAGE_READ_WRITE,
memory_flags: hal::MemoryFlags::empty(),
};
let dst_buffer =
unsafe { device.create_buffer(&dst_buffer_desc) }.map_err(DeviceError::from_hal)?;
let dst_bind_group_desc = hal::BindGroupDescriptor {
label: None,
layout: dst_bind_group_layout.as_ref(),
entries: &[hal::BindGroupEntry {
binding: 0,
resource_index: 0,
count: 1,
}],
buffers: &[hal::BufferBinding {
buffer: dst_buffer.as_ref(),
offset: 0,
size: Some(DST_BUFFER_SIZE),
}],
samplers: &[],
textures: &[],
acceleration_structures: &[],
};
let dst_bind_group = unsafe {
device
.create_bind_group(&dst_bind_group_desc)
.map_err(DeviceError::from_hal)
}?;
Ok(Self {
module,
dst_bind_group_layout,
src_bind_group_layout,
pipeline_layout,
pipeline,
dst_buffer,
dst_bind_group,
})
}
/// `Ok(None)` will only be returned if `buffer_size` is `0`.
pub fn create_src_bind_group(
&self,
device: &dyn hal::DynDevice,
limits: &wgt::Limits,
buffer_size: u64,
buffer: &dyn hal::DynBuffer,
) -> Result<Option<Box<dyn hal::DynBindGroup>>, DeviceError> {
let binding_size = calculate_src_buffer_binding_size(buffer_size, limits);
let Some(binding_size) = NonZeroU64::new(binding_size) else {
return Ok(None);
};
let hal_desc = hal::BindGroupDescriptor {
label: None,
layout: self.src_bind_group_layout.as_ref(),
entries: &[hal::BindGroupEntry {
binding: 0,
resource_index: 0,
count: 1,
}],
buffers: &[hal::BufferBinding {
buffer,
offset: 0,
size: Some(binding_size),
}],
samplers: &[],
textures: &[],
acceleration_structures: &[],
};
unsafe {
device
.create_bind_group(&hal_desc)
.map(Some)
.map_err(DeviceError::from_hal)
}
}
pub fn params<'a>(&'a self, limits: &wgt::Limits, offset: u64, buffer_size: u64) -> Params<'a> {
// The offset we receive is only required to be aligned to 4 bytes.
//
// Binding offsets and dynamic offsets are required to be aligned to
// min_storage_buffer_offset_alignment (256 bytes by default).
//
// So, we work around this limitation by calculating an aligned offset
// and pass the remainder through a push constant.
//
// We could bind the whole buffer and only have to pass the offset
// through a push constant but we might run into the
// max_storage_buffer_binding_size limit.
//
// See the inner docs of `calculate_src_buffer_binding_size` to
// see how we get the appropriate `binding_size`.
let alignment = limits.min_storage_buffer_offset_alignment as u64;
let binding_size = calculate_src_buffer_binding_size(buffer_size, limits);
let aligned_offset = offset - offset % alignment;
// This works because `binding_size` is either `buffer_size` or `alignment * 2 + buffer_size % alignment`.
let max_aligned_offset = buffer_size - binding_size;
let aligned_offset = aligned_offset.min(max_aligned_offset);
let offset_remainder = offset - aligned_offset;
Params {
pipeline_layout: self.pipeline_layout.as_ref(),
pipeline: self.pipeline.as_ref(),
dst_buffer: self.dst_buffer.as_ref(),
dst_bind_group: self.dst_bind_group.as_ref(),
aligned_offset,
offset_remainder,
}
}
pub fn dispose(self, device: &dyn hal::DynDevice) {
let IndirectValidation {
module,
dst_bind_group_layout,
src_bind_group_layout,
pipeline_layout,
pipeline,
dst_buffer,
dst_bind_group,
} = self;
unsafe {
device.destroy_bind_group(dst_bind_group);
device.destroy_buffer(dst_buffer);
device.destroy_compute_pipeline(pipeline);
device.destroy_pipeline_layout(pipeline_layout);
device.destroy_bind_group_layout(src_bind_group_layout);
device.destroy_bind_group_layout(dst_bind_group_layout);
device.destroy_shader_module(module);
}
}
}
fn calculate_src_buffer_binding_size(buffer_size: u64, limits: &wgt::Limits) -> u64 {
let alignment = limits.min_storage_buffer_offset_alignment as u64;
// We need to choose a binding size that can address all possible sets of 12 contiguous bytes in the buffer taking
// into account that the dynamic offset needs to be a multiple of `min_storage_buffer_offset_alignment`.
// Given the know variables: `offset`, `buffer_size`, `alignment` and the rule `offset + 12 <= buffer_size`.
// Let `chunks = floor(buffer_size / alignment)`.
// Let `chunk` be the interval `[0, chunks]`.
// Let `offset = alignment * chunk + r` where `r` is the interval [0, alignment - 4].
// Let `binding` be the interval `[offset, offset + 12]`.
// Let `aligned_offset = alignment * chunk`.
// Let `aligned_binding` be the interval `[aligned_offset, aligned_offset + r + 12]`.
// Let `aligned_binding_size = r + 12 = [12, alignment + 8]`.
// Let `min_aligned_binding_size = alignment + 8`.
// `min_aligned_binding_size` is the minimum binding size required to address all 12 contiguous bytes in the buffer
// but the last aligned_offset + min_aligned_binding_size might overflow the buffer. In order to avoid this we must
// pick a larger `binding_size` that satisfies: `last_aligned_offset + binding_size = buffer_size` and
// `binding_size >= min_aligned_binding_size`.
// Let `buffer_size = alignment * chunks + sr` where `sr` is the interval [0, alignment - 4].
// Let `last_aligned_offset = alignment * (chunks - u)` where `u` is the interval [0, chunks].
// => `binding_size = buffer_size - last_aligned_offset`
// => `binding_size = alignment * chunks + sr - alignment * (chunks - u)`
// => `binding_size = alignment * chunks + sr - alignment * chunks + alignment * u`
// => `binding_size = sr + alignment * u`
// => `min_aligned_binding_size <= sr + alignment * u`
// => `alignment + 8 <= sr + alignment * u`
// => `u` must be at least 2
// => `binding_size = sr + alignment * 2`
let binding_size = 2 * alignment + (buffer_size % alignment);
binding_size.min(buffer_size)
}