use log::{error, info, warn};
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use wgpu;
use super::builder::ComputeConfiguration;
use super::multipass::MultiPassManager;
use super::resource::ResourceLayout;
use crate::{Core, FontSystem, ShaderHotReload, TextureManager, UniformBinding};
crate::uniform_params! {
pub struct ComputeTimeUniform {
pub time: f32,
pub delta: f32,
pub frame: u32,
pub _padding: u32,
}
}
pub struct ComputeShader {
pub pipelines: Vec<wgpu::ComputePipeline>,
pub output_texture: TextureManager,
pub time_uniform: UniformBinding<ComputeTimeUniform>,
pub workgroup_size: [u32; 3],
pub dispatch_once: bool,
pub current_frame: u32,
pub bind_group_layouts: HashMap<u32, wgpu::BindGroupLayout>,
pub pipeline_layout: wgpu::PipelineLayout,
pub group0_bind_group: wgpu::BindGroup, pub group1_bind_group: wgpu::BindGroup, pub group2_bind_group: Option<wgpu::BindGroup>, pub group3_bind_group: Option<wgpu::BindGroup>,
pub custom_uniform: Option<wgpu::Buffer>,
pub custom_uniform_size: Option<u64>,
pub placeholder_input_texture: Option<TextureManager>,
pub multipass_manager: Option<MultiPassManager>,
pub pass_dependencies: Option<HashMap<String, Vec<String>>>,
pub pass_descriptions: Option<Vec<crate::compute::PassDescription>>,
pub font_system: Option<FontSystem>,
pub atomic_buffer_raw: Option<wgpu::Buffer>,
pub atomic_buffer_channels: u32,
pub audio_buffer: Option<wgpu::Buffer>,
pub audio_staging_buffer: Option<wgpu::Buffer>,
pub audio_spectrum_buffer: Option<wgpu::Buffer>,
pub mouse_uniform: Option<UniformBinding<crate::MouseUniform>>,
pub channel_textures: HashMap<u32, Option<(wgpu::TextureView, wgpu::Sampler)>>,
pub num_channels: u32,
pub storage_buffers: Vec<wgpu::Buffer>,
pub empty_bind_groups: std::collections::HashMap<u32, wgpu::BindGroup>,
pub multipass_sampler: wgpu::Sampler,
cached_intermediate_group1: HashMap<String, [wgpu::BindGroup; 2]>,
cached_input_group3: HashMap<String, Vec<wgpu::BindGroup>>,
max_input_deps: usize,
pub entry_points: Vec<String>,
pub hot_reload: Option<ShaderHotReload>,
pub label: String,
pub has_input_texture: bool,
pub texture_format: wgpu::TextureFormat,
}
impl ComputeShader {
pub fn from_builder(core: &Core, shader_source: &str, mut config: ComputeConfiguration) -> Self {
let mut resource_layout = ResourceLayout::new();
resource_layout.add_time_uniform();
resource_layout.add_output_texture(config.texture_format);
if let Some(uniform_size) = config.custom_uniform_size {
resource_layout.add_custom_uniform("params", uniform_size);
}
if config.has_input_texture {
resource_layout.add_input_texture();
}
if config.has_mouse {
resource_layout.add_mouse_uniform();
}
if config.has_fonts {
resource_layout.add_font_resources();
}
if config.has_audio {
resource_layout.add_audio_buffer(config.audio_buffer_size);
}
if config.has_audio_spectrum {
resource_layout.add_audio_spectrum_buffer(config.audio_spectrum_size);
}
if config.has_atomic_buffer {
let atomic_size = (core.size.width * core.size.height * config.atomic_buffer_channels * 4) as u64;
resource_layout.add_atomic_buffer(atomic_size);
}
if let Some(num_channels) = config.num_channels {
resource_layout.add_channel_textures(num_channels);
}
if !config.storage_buffers.is_empty() {
for buffer_spec in &config.storage_buffers {
resource_layout.add_storage_buffer(&buffer_spec.name, buffer_spec.size_bytes);
}
} else if config.passes.is_some() {
resource_layout.add_multipass_input_textures(config.max_input_deps);
}
let bind_group_layouts = resource_layout.create_bind_group_layouts(&core.device);
let mut layouts_vec: Vec<wgpu::BindGroupLayout> = Vec::new();
for i in 0..4 {
if let Some(layout) = bind_group_layouts.get(&i) {
layouts_vec.push(layout.clone()); } else {
let empty_layout =
core.device
.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some(&format!("Empty Group {i} Layout")),
entries: &[],
});
layouts_vec.push(empty_layout);
}
}
let layout_refs: Vec<Option<&wgpu::BindGroupLayout>> = layouts_vec.iter().map(|l| Some(l)).collect();
let pipeline_layout = core
.device
.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some(&format!("{} Pipeline Layout", config.label)),
bind_group_layouts: &layout_refs,
immediate_size: 0,
});
let time_bind_group_layout = bind_group_layouts.get(&0).unwrap();
let time_uniform = UniformBinding::new(
&core.device,
&format!("{} Time Uniform", config.label),
ComputeTimeUniform {
time: 0.0,
delta: 0.0,
frame: 0,
_padding: 0,
},
time_bind_group_layout,
0,
);
let group0_bind_group = time_uniform.bind_group.clone();
let output_texture = Self::create_output_texture(
&core.device,
core.size.width,
core.size.height,
config.texture_format,
&format!("{} Output Texture", config.label),
);
let custom_uniform = if let Some(uniform_size) = config.custom_uniform_size {
Some(core.device.create_buffer(&wgpu::BufferDescriptor {
label: Some(&format!("{} Custom Uniform Buffer", config.label)),
size: uniform_size,
usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
}))
} else {
None
};
let placeholder_input_texture = if config.has_input_texture {
Some(Self::create_placeholder_input_texture(
&core.device,
&format!("{} Placeholder Input", config.label),
))
} else {
None
};
let group1_bind_group = Self::create_group1_bind_group(
&core.device,
bind_group_layouts.get(&1).unwrap(),
&output_texture,
config.custom_uniform_size,
config.has_input_texture,
custom_uniform.as_ref(),
placeholder_input_texture.as_ref().map(|t| &t.view),
placeholder_input_texture.as_ref().map(|t| &t.sampler),
);
let (
font_system,
atomic_buffer_raw,
audio_buffer,
audio_staging_buffer,
audio_spectrum_buffer,
mouse_uniform,
group2_bind_group,
) = Self::create_engine_resources(core, &bind_group_layouts, &config);
let (storage_buffers, group3_bind_group) = if !config.storage_buffers.is_empty() {
Self::create_user_storage_buffers(core, &bind_group_layouts, &config)
} else if config.passes.is_some() {
(Vec::new(), None)
} else {
(Vec::new(), None)
};
let mut empty_bind_groups: std::collections::HashMap<u32, wgpu::BindGroup> =
std::collections::HashMap::new();
for i in 0..4 {
if !bind_group_layouts.contains_key(&i) {
let empty_bind_group = core.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some(&format!("Empty Group {i} Bind Group")),
layout: &layouts_vec[i as usize],
entries: &[],
});
empty_bind_groups.insert(i, empty_bind_group);
}
}
let (multipass_manager, pass_dependencies) = if let Some(passes) = &config.passes {
if config.storage_buffers.is_empty() {
let buffer_names: Vec<String> = passes.iter().map(|p| p.name.clone()).collect();
let dependencies: HashMap<String, Vec<String>> = passes
.iter()
.map(|p| (p.name.clone(), p.inputs.clone()))
.collect();
let manager = MultiPassManager::new(
core,
&buffer_names,
config.texture_format,
bind_group_layouts.get(&3).unwrap().clone(),
config.max_input_deps,
passes,
);
(Some(manager), Some(dependencies))
} else {
let dependencies: HashMap<String, Vec<String>> = passes
.iter()
.map(|p| (p.name.clone(), p.inputs.clone()))
.collect();
(None, Some(dependencies))
}
} else {
(None, None)
};
let shader_module = core
.device
.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some(&format!("{} Module", config.label)),
source: wgpu::ShaderSource::Wgsl(shader_source.into()),
});
let mut pipelines = Vec::new();
for entry_point in &config.entry_points {
let pipeline = core
.device
.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some(&format!("{} Pipeline - {}", config.label, entry_point)),
layout: Some(&pipeline_layout),
module: &shader_module,
entry_point: Some(entry_point),
compilation_options: wgpu::PipelineCompilationOptions::default(),
cache: None,
});
pipelines.push(pipeline);
}
let multipass_sampler = core.device.create_sampler(&wgpu::SamplerDescriptor {
address_mode_u: wgpu::AddressMode::Repeat,
address_mode_v: wgpu::AddressMode::Repeat,
address_mode_w: wgpu::AddressMode::Repeat,
mag_filter: wgpu::FilterMode::Linear,
min_filter: wgpu::FilterMode::Linear,
..Default::default()
});
let hot_reload_path = config.hot_reload_path.take();
let mut shader = Self {
pipelines,
output_texture,
time_uniform,
workgroup_size: config.workgroup_size,
dispatch_once: config.dispatch_once,
current_frame: 0,
bind_group_layouts,
pipeline_layout,
group0_bind_group,
group1_bind_group,
group2_bind_group,
group3_bind_group,
multipass_manager,
pass_dependencies,
pass_descriptions: config.passes.clone(),
font_system,
atomic_buffer_raw,
atomic_buffer_channels: config.atomic_buffer_channels,
audio_buffer,
audio_staging_buffer,
audio_spectrum_buffer,
mouse_uniform,
storage_buffers,
empty_bind_groups,
custom_uniform,
custom_uniform_size: config.custom_uniform_size,
placeholder_input_texture,
channel_textures: Self::initialize_channel_textures(config.num_channels.unwrap_or(0)),
num_channels: config.num_channels.unwrap_or(0),
multipass_sampler,
cached_intermediate_group1: HashMap::new(),
cached_input_group3: HashMap::new(),
max_input_deps: config.max_input_deps,
entry_points: config.entry_points,
hot_reload: None,
label: config.label,
has_input_texture: config.has_input_texture,
texture_format: config.texture_format,
};
shader.rebuild_multipass_caches(&core.device);
if let Some(path) = hot_reload_path {
let reload_module =
core.device
.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some("Hot Reload Module"),
source: wgpu::ShaderSource::Wgsl(shader_source.into()),
});
if let Err(e) = shader.enable_hot_reload(core.device.clone(), path, reload_module) {
warn!("Failed to enable hot reload: {e}");
}
}
shader
}
fn create_output_texture(
device: &wgpu::Device,
width: u32,
height: u32,
format: wgpu::TextureFormat,
label: &str,
) -> TextureManager {
let texture = device.create_texture(&wgpu::TextureDescriptor {
label: Some(label),
size: wgpu::Extent3d {
width,
height,
depth_or_array_layers: 1,
},
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format,
usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::STORAGE_BINDING,
view_formats: &[],
});
let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
address_mode_u: wgpu::AddressMode::ClampToEdge,
address_mode_v: wgpu::AddressMode::ClampToEdge,
address_mode_w: wgpu::AddressMode::ClampToEdge,
mag_filter: wgpu::FilterMode::Linear,
min_filter: wgpu::FilterMode::Linear,
..Default::default()
});
let display_layout = TextureManager::create_display_layout(device);
let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
layout: &display_layout,
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: wgpu::BindingResource::TextureView(&view),
},
wgpu::BindGroupEntry {
binding: 1,
resource: wgpu::BindingResource::Sampler(&sampler),
},
],
label: Some(&format!("{label} Display Bind Group")),
});
TextureManager {
texture,
view,
sampler,
bind_group,
}
}
fn create_placeholder_input_texture(device: &wgpu::Device, label: &str) -> TextureManager {
let texture = device.create_texture(&wgpu::TextureDescriptor {
label: Some(label),
size: wgpu::Extent3d {
width: 1,
height: 1,
depth_or_array_layers: 1,
},
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format: wgpu::TextureFormat::Rgba8UnormSrgb, usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
view_formats: &[],
});
let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
let sampler = device.create_sampler(&wgpu::SamplerDescriptor::default());
let display_layout = TextureManager::create_display_layout(device);
let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
layout: &display_layout,
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: wgpu::BindingResource::TextureView(&view),
},
wgpu::BindGroupEntry {
binding: 1,
resource: wgpu::BindingResource::Sampler(&sampler),
},
],
label: Some(&format!("{label} Placeholder Bind Group")),
});
TextureManager {
texture,
view,
sampler,
bind_group,
}
}
fn create_group1_bind_group(
device: &wgpu::Device,
layout: &wgpu::BindGroupLayout,
output_texture: &TextureManager,
custom_uniform_size: Option<u64>,
has_input_texture: bool,
custom_uniform_buffer: Option<&wgpu::Buffer>,
input_texture_view: Option<&wgpu::TextureView>,
input_sampler: Option<&wgpu::Sampler>,
) -> wgpu::BindGroup {
let storage_view = output_texture
.texture
.create_view(&wgpu::TextureViewDescriptor::default());
let mut entries = vec![wgpu::BindGroupEntry {
binding: 0,
resource: wgpu::BindingResource::TextureView(&storage_view),
}];
if let (Some(buffer), Some(_size)) = (custom_uniform_buffer, custom_uniform_size) {
entries.push(wgpu::BindGroupEntry {
binding: 1, resource: buffer.as_entire_binding(),
});
}
if has_input_texture {
if let (Some(view), Some(sampler)) = (input_texture_view, input_sampler) {
entries.push(wgpu::BindGroupEntry {
binding: 2, resource: wgpu::BindingResource::TextureView(view),
});
entries.push(wgpu::BindGroupEntry {
binding: 3, resource: wgpu::BindingResource::Sampler(sampler),
});
} else {
log::error!("Input texture required but not provided during bind group creation");
}
}
device.create_bind_group(&wgpu::BindGroupDescriptor {
layout,
entries: &entries,
label: Some("Group 1 Bind Group"),
})
}
fn create_engine_resources(
core: &Core,
layouts: &HashMap<u32, wgpu::BindGroupLayout>,
config: &ComputeConfiguration,
) -> (
Option<FontSystem>,
Option<wgpu::Buffer>,
Option<wgpu::Buffer>,
Option<wgpu::Buffer>,
Option<wgpu::Buffer>,
Option<UniformBinding<crate::MouseUniform>>,
Option<wgpu::BindGroup>,
) {
let layout = layouts.get(&2);
if layout.is_none() {
return (None, None, None, None, None, None, None);
}
let layout = layout.unwrap();
let font_system = if config.has_fonts {
Some(FontSystem::new(core))
} else {
None
};
let atomic_buffer_raw = if config.has_atomic_buffer {
let buffer_size = (core.size.width * core.size.height * config.atomic_buffer_channels * 4) as u64;
Some(core.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Atomic Storage Buffer"),
size: buffer_size,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
}))
} else {
None
};
let (audio_buffer, audio_staging_buffer) = if config.has_audio {
let buffer_size = config.audio_buffer_size * std::mem::size_of::<f32>();
let audio_buffer = core.device.create_buffer(&wgpu::BufferDescriptor {
label: Some(&format!("{} Audio Buffer", config.label)),
size: buffer_size as u64,
usage: wgpu::BufferUsages::STORAGE
| wgpu::BufferUsages::COPY_DST
| wgpu::BufferUsages::COPY_SRC,
mapped_at_creation: false,
});
let staging_buffer = core.device.create_buffer(&wgpu::BufferDescriptor {
label: Some(&format!("{} Audio Staging Buffer", config.label)),
size: buffer_size as u64,
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
(Some(audio_buffer), Some(staging_buffer))
} else {
(None, None)
};
let audio_spectrum_buffer = if config.has_audio_spectrum {
let buffer_size = config.audio_spectrum_size * std::mem::size_of::<f32>();
let buffer = core.device.create_buffer(&wgpu::BufferDescriptor {
label: Some(&format!("{} Audio Spectrum Buffer", config.label)),
size: buffer_size as u64,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
Some(buffer)
} else {
None
};
let mouse_uniform = if config.has_mouse {
let temp_layout =
core.device
.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
entries: &[wgpu::BindGroupLayoutEntry {
binding: 0,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Uniform,
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
}],
label: Some("Temp Mouse Layout"),
});
Some(UniformBinding::new(
&core.device,
"Mouse Uniform",
crate::MouseUniform::default(),
&temp_layout,
0,
))
} else {
None
};
let empty_channels = std::collections::HashMap::new();
let num_channels = config.num_channels.unwrap_or(0);
let bind_group = Self::create_group2_bind_group(
&core.device,
&core.queue,
layout,
&font_system,
&atomic_buffer_raw,
&audio_buffer,
&audio_spectrum_buffer,
&mouse_uniform,
&empty_channels,
num_channels,
);
(
font_system,
atomic_buffer_raw,
audio_buffer,
audio_staging_buffer,
audio_spectrum_buffer,
mouse_uniform,
bind_group,
)
}
fn create_group2_bind_group(
device: &wgpu::Device,
queue: &wgpu::Queue,
layout: &wgpu::BindGroupLayout,
font_system: &Option<FontSystem>,
atomic_buffer_raw: &Option<wgpu::Buffer>,
audio_buffer: &Option<wgpu::Buffer>,
audio_spectrum_buffer: &Option<wgpu::Buffer>,
mouse_uniform: &Option<UniformBinding<crate::MouseUniform>>,
channel_textures: &HashMap<u32, Option<(wgpu::TextureView, wgpu::Sampler)>>,
num_channels: u32,
) -> Option<wgpu::BindGroup> {
let default_texture = device.create_texture(&wgpu::TextureDescriptor {
label: Some("Default Channel Texture"),
size: wgpu::Extent3d {
width: 1,
height: 1,
depth_or_array_layers: 1,
},
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format: wgpu::TextureFormat::Rgba8Unorm,
usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
view_formats: &[],
});
let magenta_data: [u8; 4] = [255, 0, 255, 255];
queue.write_texture(
wgpu::TexelCopyTextureInfo {
texture: &default_texture,
mip_level: 0,
origin: wgpu::Origin3d::ZERO,
aspect: wgpu::TextureAspect::All,
},
&magenta_data,
wgpu::TexelCopyBufferLayout {
offset: 0,
bytes_per_row: Some(4),
rows_per_image: Some(1),
},
wgpu::Extent3d {
width: 1,
height: 1,
depth_or_array_layers: 1,
},
);
let default_texture_view =
default_texture.create_view(&wgpu::TextureViewDescriptor::default());
let default_sampler = device.create_sampler(&wgpu::SamplerDescriptor::default());
let mut entries = Vec::new();
let mut binding_counter = 0;
if let Some(mouse) = mouse_uniform {
entries.push(wgpu::BindGroupEntry {
binding: binding_counter,
resource: mouse.buffer.as_entire_binding(),
});
binding_counter += 1;
}
if let Some(font_tex) = font_system {
entries.extend_from_slice(&[
wgpu::BindGroupEntry {
binding: binding_counter,
resource: font_tex.font_uniforms.buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: binding_counter + 1,
resource: wgpu::BindingResource::TextureView(&font_tex.atlas_texture.view),
},
]);
binding_counter += 2;
}
if let Some(audio) = audio_buffer {
entries.push(wgpu::BindGroupEntry {
binding: binding_counter,
resource: audio.as_entire_binding(),
});
binding_counter += 1;
}
if let Some(audio_spectrum) = audio_spectrum_buffer {
entries.push(wgpu::BindGroupEntry {
binding: binding_counter,
resource: audio_spectrum.as_entire_binding(),
});
binding_counter += 1;
}
if let Some(atomic_buf) = atomic_buffer_raw {
entries.push(wgpu::BindGroupEntry {
binding: binding_counter,
resource: atomic_buf.as_entire_binding(),
});
binding_counter += 1;
}
for i in 0..num_channels {
let (texture_view, sampler) = if let Some(Some((view, samp))) = channel_textures.get(&i)
{
(view, samp)
} else {
(&default_texture_view, &default_sampler)
};
entries.push(wgpu::BindGroupEntry {
binding: binding_counter,
resource: wgpu::BindingResource::TextureView(texture_view),
});
binding_counter += 1;
entries.push(wgpu::BindGroupEntry {
binding: binding_counter,
resource: wgpu::BindingResource::Sampler(sampler),
});
binding_counter += 1;
}
if entries.is_empty() {
return None;
}
Some(device.create_bind_group(&wgpu::BindGroupDescriptor {
layout,
entries: &entries,
label: Some("Group 2 Bind Group"),
}))
}
fn create_user_storage_buffers(
core: &Core,
layouts: &HashMap<u32, wgpu::BindGroupLayout>,
config: &ComputeConfiguration,
) -> (Vec<wgpu::Buffer>, Option<wgpu::BindGroup>) {
if config.storage_buffers.is_empty() {
return (Vec::new(), None);
}
let layout = layouts.get(&3);
if layout.is_none() {
return (Vec::new(), None);
}
let layout = layout.unwrap();
let mut storage_buffers = Vec::new();
let mut entries = Vec::new();
for buffer_spec in config.storage_buffers.iter() {
let buffer = core.device.create_buffer(&wgpu::BufferDescriptor {
label: Some(&buffer_spec.name),
size: buffer_spec.size_bytes,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::COPY_SRC,
mapped_at_creation: false,
});
storage_buffers.push(buffer);
}
for (i, buffer) in storage_buffers.iter().enumerate() {
entries.push(wgpu::BindGroupEntry {
binding: i as u32,
resource: buffer.as_entire_binding(),
});
}
let bind_group = core.device.create_bind_group(&wgpu::BindGroupDescriptor {
layout,
entries: &entries,
label: Some("Group 3 Bind Group"),
});
(storage_buffers, Some(bind_group))
}
pub fn dispatch_stage_with_workgroups(
&mut self,
encoder: &mut wgpu::CommandEncoder,
stage_index: usize,
workgroup_count: [u32; 3],
) {
if stage_index >= self.pipelines.len() {
log::error!(
"Stage index {} out of bounds (max: {})",
stage_index,
self.pipelines.len() - 1
);
return;
}
if self.dispatch_once && self.current_frame > 0 {
return;
}
let pipeline = &self.pipelines[stage_index];
let entry_point = &self.entry_points[stage_index];
let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: Some(&format!(
"{} Stage {} - {}",
self.label, stage_index, entry_point
)),
timestamp_writes: None,
});
compute_pass.set_pipeline(pipeline);
compute_pass.set_bind_group(0, &self.group0_bind_group, &[]); compute_pass.set_bind_group(1, &self.group1_bind_group, &[]);
if let Some(ref group2) = self.group2_bind_group {
compute_pass.set_bind_group(2, group2, &[]);
} else if let Some(empty_group2) = self.empty_bind_groups.get(&2) {
compute_pass.set_bind_group(2, empty_group2, &[]);
}
if let Some(ref group3) = self.group3_bind_group {
compute_pass.set_bind_group(3, group3, &[]);
} else if let Some(empty_group3) = self.empty_bind_groups.get(&3) {
compute_pass.set_bind_group(3, empty_group3, &[]);
}
compute_pass.dispatch_workgroups(
workgroup_count[0],
workgroup_count[1],
workgroup_count[2],
);
}
fn workgroup_count_for(&self, width: u32, height: u32) -> [u32; 3] {
[
width.div_ceil(self.workgroup_size[0]),
height.div_ceil(self.workgroup_size[1]),
1,
]
}
pub fn dispatch_stage(
&mut self,
encoder: &mut wgpu::CommandEncoder,
core: &Core,
stage_index: usize,
) {
self.check_hot_reload(&core.device);
let width = self.output_texture.texture.width();
let height = self.output_texture.texture.height();
let workgroup_count = self.workgroup_count_for(width, height);
self.dispatch_stage_with_workgroups(encoder, stage_index, workgroup_count);
}
pub fn dispatch(&mut self, encoder: &mut wgpu::CommandEncoder, core: &Core) {
self.check_hot_reload(&core.device);
if self.dispatch_once && self.current_frame > 0 {
return;
}
let width = self.output_texture.texture.width();
let height = self.output_texture.texture.height();
let workgroup_count = self.workgroup_count_for(width, height);
if self.multipass_manager.is_some() {
self.dispatch_multipass(encoder, workgroup_count);
} else {
self.dispatch_single_pass(encoder, core, workgroup_count);
}
self.current_frame += 1;
}
pub fn dispatch_at_resolution(
&mut self,
encoder: &mut wgpu::CommandEncoder,
core: &Core,
width: u32,
height: u32,
) {
self.check_hot_reload(&core.device);
if self.dispatch_once && self.current_frame > 0 {
return;
}
let workgroup_count = self.workgroup_count_for(width, height);
if self.multipass_manager.is_some() {
self.dispatch_multipass(encoder, workgroup_count);
} else {
self.dispatch_single_pass(encoder, core, workgroup_count);
}
self.current_frame += 1;
}
pub fn flip_buffers(&mut self) {
if let Some(ref mut multipass) = self.multipass_manager {
multipass.flip_buffers();
}
}
pub fn set_custom_params<T: bytemuck::Pod>(&self, params: T, queue: &wgpu::Queue) {
if let Some(ref buffer) = self.custom_uniform {
queue.write_buffer(buffer, 0, bytemuck::bytes_of(¶ms));
} else {
log::warn!("Attempted to set custom params but no custom uniform buffer exists");
}
}
pub fn get_custom_uniform_size(&self) -> Option<u64> {
self.custom_uniform_size
}
pub fn update_input_texture(
&mut self,
texture_view: &wgpu::TextureView,
sampler: &wgpu::Sampler,
device: &wgpu::Device,
) {
if !self.has_input_texture {
log::warn!("Attempted to update input texture but shader was not configured with input texture support");
return;
}
if let Some(ref mut _placeholder) = self.placeholder_input_texture {
}
let group1_layout = self.bind_group_layouts.get(&1).unwrap();
let storage_view = self
.output_texture
.texture
.create_view(&wgpu::TextureViewDescriptor::default());
let mut entries = vec![wgpu::BindGroupEntry {
binding: 0,
resource: wgpu::BindingResource::TextureView(&storage_view),
}];
if let Some(ref buffer) = self.custom_uniform {
entries.push(wgpu::BindGroupEntry {
binding: 1,
resource: buffer.as_entire_binding(),
});
}
entries.push(wgpu::BindGroupEntry {
binding: 2,
resource: wgpu::BindingResource::TextureView(texture_view),
});
entries.push(wgpu::BindGroupEntry {
binding: 3,
resource: wgpu::BindingResource::Sampler(sampler),
});
self.group1_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
layout: group1_layout,
entries: &entries,
label: Some("Updated Group 1 Bind Group with Input Texture"),
});
}
pub fn update_channel_texture(
&mut self,
channel_index: u32,
texture_view: &wgpu::TextureView,
sampler: &wgpu::Sampler,
device: &wgpu::Device,
queue: &wgpu::Queue,
) {
if channel_index >= self.num_channels {
log::warn!(
"Attempted to update channel {} but only {} channels are configured",
channel_index,
self.num_channels
);
return;
}
self.channel_textures
.insert(channel_index, Some((texture_view.clone(), sampler.clone())));
self.recreate_group2_bind_group(device, queue);
}
fn initialize_channel_textures(
num_channels: u32,
) -> HashMap<u32, Option<(wgpu::TextureView, wgpu::Sampler)>> {
let mut channel_textures = HashMap::new();
for i in 0..num_channels {
channel_textures.insert(i, None);
}
channel_textures
}
fn recreate_group2_bind_group(&mut self, device: &wgpu::Device, queue: &wgpu::Queue) {
if let Some(layout) = self.bind_group_layouts.get(&2) {
self.group2_bind_group = Self::create_group2_bind_group(
device,
queue,
layout,
&self.font_system,
&self.atomic_buffer_raw,
&self.audio_buffer,
&self.audio_spectrum_buffer,
&self.mouse_uniform,
&self.channel_textures,
self.num_channels,
);
}
}
fn dispatch_single_pass(
&self,
encoder: &mut wgpu::CommandEncoder,
_core: &Core,
workgroup_count: [u32; 3],
) {
for (i, pipeline) in self.pipelines.iter().enumerate() {
let pass_workgroup_count = if let Some(ref pass_descriptions) = self.pass_descriptions {
if let Some(pass_desc) = pass_descriptions.get(i) {
if let Some(custom_size) = pass_desc.workgroup_size {
custom_size } else {
workgroup_count }
} else {
workgroup_count }
} else {
workgroup_count };
let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: Some(&format!("{} Compute Pass {}", self.label, i)),
timestamp_writes: None,
});
compute_pass.set_pipeline(pipeline);
compute_pass.set_bind_group(0, &self.group0_bind_group, &[]); compute_pass.set_bind_group(1, &self.group1_bind_group, &[]);
if let Some(ref group2) = self.group2_bind_group {
compute_pass.set_bind_group(2, group2, &[]); } else if let Some(empty_group2) = self.empty_bind_groups.get(&2) {
compute_pass.set_bind_group(2, empty_group2, &[]);
} else {
log::error!("No Group 2 bind group available - this shouldn't happen with contiguous layout");
}
if let Some(ref group3) = self.group3_bind_group {
compute_pass.set_bind_group(3, group3, &[]); } else if let Some(empty_group3) = self.empty_bind_groups.get(&3) {
compute_pass.set_bind_group(3, empty_group3, &[]);
} else {
log::error!("No Group 3 bind group available - this shouldn't happen with contiguous layout");
}
compute_pass.dispatch_workgroups(
pass_workgroup_count[0],
pass_workgroup_count[1],
pass_workgroup_count[2],
);
}
}
fn dispatch_multipass(
&mut self,
encoder: &mut wgpu::CommandEncoder,
workgroup_count: [u32; 3],
) {
let num_passes = self.pipelines.len();
for pass_idx in 0..num_passes {
let pipeline = &self.pipelines[pass_idx];
let entry_point = &self.entry_points[pass_idx];
let pass_workgroup_count = if let Some(ref pass_descriptions) = self.pass_descriptions {
if let Some(pass_desc) = pass_descriptions.get(pass_idx) {
if let Some(custom_size) = pass_desc.workgroup_size {
custom_size } else if pass_desc.resolution.is_some() || pass_desc.resolution_scale.is_some() {
if let Some(ref multipass) = self.multipass_manager {
let (bw, bh) = multipass.get_buffer_dimensions(entry_point);
self.workgroup_count_for(bw, bh)
} else {
workgroup_count
}
} else {
workgroup_count }
} else {
workgroup_count
}
} else {
workgroup_count
};
let group3_key = if let (Some(multipass), Some(dependencies)) =
(&self.multipass_manager, &self.pass_dependencies)
{
let empty_deps = Vec::new();
let deps = dependencies.get(entry_point).unwrap_or(&empty_deps);
let first_buf = multipass
.first_buffer_name()
.cloned()
.unwrap_or_else(|| "main".to_string());
let mut key = 0usize;
for i in 0..self.max_input_deps {
let buf_name = if deps.is_empty() {
&first_buf
} else {
deps.get(i).unwrap_or(&deps[0])
};
if multipass.get_write_side(buf_name) {
key |= 1 << i;
}
}
key
} else {
log::warn!("Skipping pass '{entry_point}': multipass manager or dependencies missing");
continue;
};
let input_bind_group = match self.cached_input_group3.get(entry_point) {
Some(cached) => &cached[group3_key],
None => {
log::warn!("No cached input bind group for pass '{entry_point}'");
continue;
}
};
let write_side = self
.multipass_manager
.as_ref()
.map(|m| m.get_write_side(entry_point))
.unwrap_or(false);
let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: Some(&format!("{} Multi-Pass - {}", self.label, entry_point)),
timestamp_writes: None,
});
compute_pass.set_pipeline(pipeline);
compute_pass.set_bind_group(0, &self.group0_bind_group, &[]);
if entry_point == "main_image" {
compute_pass.set_bind_group(1, &self.group1_bind_group, &[]);
} else if let Some(cached) = self.cached_intermediate_group1.get(entry_point) {
compute_pass.set_bind_group(1, &cached[write_side as usize], &[]);
} else {
compute_pass.set_bind_group(1, &self.group1_bind_group, &[]);
log::warn!("No cached Group1 for intermediate pass {entry_point}");
}
if let Some(ref group2) = self.group2_bind_group {
compute_pass.set_bind_group(2, group2, &[]);
} else if let Some(empty_group2) = self.empty_bind_groups.get(&2) {
log::warn!(
"Using empty Group 2 bind group for pass {entry_point} - channels won't work!"
);
compute_pass.set_bind_group(2, empty_group2, &[]);
} else {
log::error!("No Group 2 bind group available - this shouldn't happen with contiguous layout");
}
compute_pass.set_bind_group(3, input_bind_group, &[]);
compute_pass.dispatch_workgroups(
pass_workgroup_count[0],
pass_workgroup_count[1],
pass_workgroup_count[2],
);
if pass_idx < num_passes - 1 {
if let Some(ref mut multipass_mut) = self.multipass_manager {
multipass_mut.mark_written(entry_point);
}
}
}
}
pub fn enable_hot_reload(
&mut self,
device: Arc<wgpu::Device>,
shader_path: PathBuf,
shader_module: wgpu::ShaderModule,
) -> Result<(), notify::Error> {
let entry_point = self
.entry_points
.first()
.cloned()
.unwrap_or_else(|| "main".to_string());
let hot_reload =
ShaderHotReload::new_compute(device, shader_path, shader_module, &entry_point)?;
self.hot_reload = Some(hot_reload);
Ok(())
}
pub fn check_hot_reload(&mut self, device: &wgpu::Device) -> bool {
if let Some(hot_reload) = &mut self.hot_reload {
if let Some(new_module) = hot_reload.reload_compute_shader() {
let mut new_pipelines = Vec::new();
for entry_point in &self.entry_points {
let new_pipeline =
device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some(&format!(
"Updated {} Pipeline - {}",
self.label, entry_point
)),
layout: Some(&self.pipeline_layout),
module: new_module,
entry_point: Some(entry_point),
compilation_options: wgpu::PipelineCompilationOptions::default(),
cache: None,
});
new_pipelines.push(new_pipeline);
}
self.pipelines = new_pipelines;
info!(
"{} shader hot-reloaded at frame: {}",
self.label, self.current_frame
);
return true;
}
}
false
}
pub fn set_time(&mut self, elapsed: f32, delta: f32, queue: &wgpu::Queue) {
self.time_uniform.data.time = elapsed;
self.time_uniform.data.delta = delta;
self.time_uniform.data.frame = self.current_frame;
self.time_uniform.update(queue);
}
pub fn update_audio_spectrum(
&mut self,
resolution_uniform: &crate::ResolutionUniform,
queue: &wgpu::Queue,
) {
if let Some(ref buffer) = self.audio_spectrum_buffer {
let mut spectrum_data = vec![0.0f32; 69];
for i in 0..64 {
let vec_idx = i / 4;
let comp_idx = i % 4;
if vec_idx < 32 {
spectrum_data[i] = resolution_uniform.audio_data[vec_idx][comp_idx];
}
}
spectrum_data[64] = resolution_uniform.bpm;
let total_energy: f32 = spectrum_data[..69].iter().sum();
if total_energy > 0.01 {
log::info!(
"Audio spectrum: energy={:.3}, BPM={:.1}",
total_energy,
spectrum_data[64]
);
}
spectrum_data[65] = resolution_uniform.bass_energy;
spectrum_data[66] = resolution_uniform.mid_energy;
spectrum_data[67] = resolution_uniform.high_energy;
spectrum_data[68] = resolution_uniform.total_energy;
queue.write_buffer(buffer, 0, bytemuck::cast_slice(&spectrum_data));
}
}
pub fn get_output_texture(&self) -> &TextureManager {
&self.output_texture
}
fn rebuild_multipass_caches(&mut self, device: &wgpu::Device) {
self.cached_intermediate_group1.clear();
self.cached_input_group3.clear();
let multipass = match &self.multipass_manager {
Some(m) => m,
None => return,
};
let dependencies = match &self.pass_dependencies {
Some(d) => d,
None => return,
};
let group1_layout = self.bind_group_layouts.get(&1).unwrap();
let first_buf = multipass
.first_buffer_name()
.cloned()
.unwrap_or_else(|| "main".to_string());
for entry_point in &self.entry_points {
if entry_point != "main_image" {
if let Some(textures) = multipass.get_buffer_pair(entry_point) {
let make_bg = |texture: &wgpu::Texture, idx: usize| {
let view =
texture.create_view(&wgpu::TextureViewDescriptor::default());
let mut entries = vec![wgpu::BindGroupEntry {
binding: 0,
resource: wgpu::BindingResource::TextureView(&view),
}];
if let Some(ref uniform_buffer) = self.custom_uniform {
entries.push(wgpu::BindGroupEntry {
binding: 1,
resource: uniform_buffer.as_entire_binding(),
});
}
device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some(&format!(
"{entry_point} Cached Group1 (side={idx})"
)),
layout: group1_layout,
entries: &entries,
})
};
self.cached_intermediate_group1.insert(
entry_point.clone(),
[make_bg(&textures.0, 0), make_bg(&textures.1, 1)],
);
}
}
let empty_deps = Vec::new();
let deps = dependencies.get(entry_point).unwrap_or(&empty_deps);
let n = self.max_input_deps;
let slot_buffers: Vec<&str> = (0..n)
.map(|i| {
if deps.is_empty() {
first_buf.as_str()
} else {
deps.get(i).unwrap_or(&deps[0]).as_str()
}
})
.collect();
let input_layout = multipass.get_input_layout();
let num_combinations = 1usize << n;
let mut cached = Vec::with_capacity(num_combinations);
for key in 0..num_combinations {
let views: Vec<wgpu::TextureView> = (0..n)
.map(|i| {
let write_side_val = (key >> i) & 1 == 1;
let textures = multipass.get_buffer_pair(slot_buffers[i]).unwrap();
let texture = if write_side_val {
&textures.0
} else {
&textures.1
};
texture.create_view(&wgpu::TextureViewDescriptor::default())
})
.collect();
let mut entries = Vec::with_capacity(n * 2);
for i in 0..n {
entries.push(wgpu::BindGroupEntry {
binding: (i * 2) as u32,
resource: wgpu::BindingResource::TextureView(&views[i]),
});
entries.push(wgpu::BindGroupEntry {
binding: (i * 2 + 1) as u32,
resource: wgpu::BindingResource::Sampler(&self.multipass_sampler),
});
}
cached.push(device.create_bind_group(&wgpu::BindGroupDescriptor {
layout: input_layout,
entries: &entries,
label: Some(&format!("{entry_point} Cached Input (key={key})")),
}));
}
self.cached_input_group3
.insert(entry_point.clone(), cached);
}
}
pub fn resize(&mut self, core: &Core, width: u32, height: u32) {
self.output_texture = Self::create_output_texture(
&core.device,
width,
height,
self.texture_format,
&format!("{} Output Texture", self.label),
);
let group1_layout = self.bind_group_layouts.get(&1).unwrap();
self.group1_bind_group = Self::create_group1_bind_group(
&core.device,
group1_layout,
&self.output_texture,
self.custom_uniform_size,
self.has_input_texture,
self.custom_uniform.as_ref(),
self.placeholder_input_texture.as_ref().map(|t| &t.view),
self.placeholder_input_texture.as_ref().map(|t| &t.sampler),
);
if let Some(multipass) = &mut self.multipass_manager {
multipass.resize(core, width, height);
}
self.rebuild_multipass_caches(&core.device);
if let Some(atomic_buffer) = &mut self.atomic_buffer_raw {
let buffer_size = (width * height * self.atomic_buffer_channels * 4) as u64;
*atomic_buffer = core.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Atomic Storage Buffer (resized)"),
size: buffer_size,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
if let Some(layout) = self.bind_group_layouts.get(&2) {
self.group2_bind_group = Self::create_group2_bind_group(
&core.device,
&core.queue,
layout,
&self.font_system,
&self.atomic_buffer_raw,
&self.audio_buffer,
&self.audio_spectrum_buffer,
&self.mouse_uniform,
&self.channel_textures,
self.num_channels,
);
}
}
self.current_frame = 0;
}
pub fn clear_all_buffers(&mut self, core: &Core) {
if let Some(multipass) = &mut self.multipass_manager {
multipass.clear_all(core);
}
self.rebuild_multipass_caches(&core.device);
self.clear_atomic_buffer(core);
self.current_frame = 0;
}
pub fn clear_atomic_buffer(&mut self, core: &Core) {
if self.atomic_buffer_raw.is_some() {
let buffer_size = (core.size.width * core.size.height * self.atomic_buffer_channels * 4) as u64;
self.atomic_buffer_raw = Some(core.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Atomic Storage Buffer (cleared)"),
size: buffer_size,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
}));
if let Some(layout) = self.bind_group_layouts.get(&2) {
self.group2_bind_group = Self::create_group2_bind_group(
&core.device,
&core.queue,
layout,
&self.font_system,
&self.atomic_buffer_raw,
&self.audio_buffer,
&self.audio_spectrum_buffer,
&self.mouse_uniform,
&self.channel_textures,
self.num_channels,
);
}
}
}
pub fn update_mouse_uniform(
&mut self,
mouse_uniform_data: &crate::MouseUniform,
queue: &wgpu::Queue,
) {
if let Some(mouse_uniform) = &mut self.mouse_uniform {
mouse_uniform.data = *mouse_uniform_data;
mouse_uniform.update(queue);
}
}
pub fn get_audio_buffer(&self) -> Option<&wgpu::Buffer> {
self.audio_buffer.as_ref()
}
pub async fn read_audio_buffer(
&self,
device: &wgpu::Device,
queue: &wgpu::Queue,
) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
if let (Some(audio_buffer), Some(staging_buffer)) =
(&self.audio_buffer, &self.audio_staging_buffer)
{
let buffer_size = audio_buffer.size();
let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
label: Some("Audio Buffer Copy"),
});
encoder.copy_buffer_to_buffer(audio_buffer, 0, staging_buffer, 0, buffer_size);
queue.submit(std::iter::once(encoder.finish()));
let buffer_slice = staging_buffer.slice(..);
let (tx, rx) = std::sync::mpsc::channel();
buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
let _ = tx.send(result);
});
let _ = device.poll(wgpu::PollType::wait_indefinitely());
match rx.recv() {
Ok(Ok(())) => {}
Ok(Err(e)) => return Err(e.into()),
Err(_) => return Err("Buffer mapping failed".into()),
}
let samples = {
let data = buffer_slice.get_mapped_range();
let samples: Vec<f32> = bytemuck::cast_slice(&data).to_vec();
samples
};
staging_buffer.unmap();
Ok(samples)
} else {
Ok(Vec::new())
}
}
pub fn handle_export(&mut self, core: &Core, render_kit: &mut crate::RenderKit) {
if let Some((frame, time)) = render_kit.export_manager.try_get_next_frame() {
let settings = render_kit.export_manager.settings();
let export_w = settings.width;
let export_h = settings.height;
if frame == 0 {
self.current_frame = 0;
let current_w = self.output_texture.texture.width();
let current_h = self.output_texture.texture.height();
if current_w != export_w || current_h != export_h {
info!("Export: resizing compute from {}x{} to {}x{}", current_w, current_h, export_w, export_h);
self.resize(core, export_w, export_h);
}
}
match self.capture_export_frame(
core,
time,
render_kit,
None::<fn(&mut Self, &mut wgpu::CommandEncoder, &Core)>,
) {
Ok(data) => {
let settings = render_kit.export_manager.settings();
if let Err(e) = crate::save_frame(data, frame, settings) {
error!("Error saving frame: {e:?}");
}
}
Err(e) => {
error!("Error capturing export frame: {e:?}");
}
}
} else {
let current_w = self.output_texture.texture.width();
let current_h = self.output_texture.texture.height();
if current_w != core.size.width || current_h != core.size.height {
info!("Export complete: resizing compute back to {}x{}", core.size.width, core.size.height);
self.resize(core, core.size.width, core.size.height);
}
render_kit.export_manager.complete_export();
}
}
pub fn handle_export_dispatch(
&mut self,
core: &Core,
render_kit: &mut crate::RenderKit,
custom_dispatch: impl FnOnce(&mut Self, &mut wgpu::CommandEncoder, &Core),
) {
if let Some((frame, time)) = render_kit.export_manager.try_get_next_frame() {
let settings = render_kit.export_manager.settings();
let export_w = settings.width;
let export_h = settings.height;
if frame == 0 {
self.current_frame = 0;
let current_w = self.output_texture.texture.width();
let current_h = self.output_texture.texture.height();
if current_w != export_w || current_h != export_h {
info!("Export: resizing compute from {}x{} to {}x{}", current_w, current_h, export_w, export_h);
self.resize(core, export_w, export_h);
}
}
match self.capture_export_frame(core, time, render_kit, Some(custom_dispatch)) {
Ok(data) => {
let settings = render_kit.export_manager.settings();
if let Err(e) = crate::save_frame(data, frame, settings) {
error!("Error saving frame: {e:?}");
}
}
Err(e) => {
error!("Error capturing export frame: {e:?}");
}
}
} else {
let current_w = self.output_texture.texture.width();
let current_h = self.output_texture.texture.height();
if current_w != core.size.width || current_h != core.size.height {
info!("Export complete: resizing compute back to {}x{}", core.size.width, core.size.height);
self.resize(core, core.size.width, core.size.height);
}
render_kit.export_manager.complete_export();
}
}
pub fn capture_export_frame<F>(
&mut self,
core: &Core,
time: f32,
render_kit: &crate::RenderKit,
custom_dispatch: Option<F>,
) -> Result<Vec<u8>, crate::SurfaceError>
where
F: FnOnce(&mut Self, &mut wgpu::CommandEncoder, &Core),
{
let settings = render_kit.export_manager.settings();
let (capture_texture, output_buffer) =
render_kit.create_capture_texture(&core.device, settings.width, settings.height);
let capture_view = capture_texture.create_view(&wgpu::TextureViewDescriptor::default());
let mut encoder = core
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor {
label: Some("Export Encoder"),
});
let delta = 1.0 / settings.fps as f32;
self.set_time(time, delta, &core.queue);
if let Some(custom_dispatch) = custom_dispatch {
custom_dispatch(self, &mut encoder, core);
} else {
self.dispatch_at_resolution(&mut encoder, core, settings.width, settings.height);
}
{
let mut render_pass = crate::Renderer::begin_render_pass(
&mut encoder,
&capture_view,
wgpu::LoadOp::Clear(wgpu::Color::BLACK),
Some("Export Capture Pass"),
);
render_pass.set_pipeline(&render_kit.renderer.render_pipeline);
render_pass.set_vertex_buffer(0, render_kit.renderer.vertex_buffer.slice(..));
render_pass.set_bind_group(0, &self.output_texture.bind_group, &[]);
render_pass.draw(0..4, 0..1);
}
let align = 256;
let unpadded_bytes_per_row = settings.width * 4;
let padding = (align - unpadded_bytes_per_row % align) % align;
let padded_bytes_per_row = unpadded_bytes_per_row + padding;
encoder.copy_texture_to_buffer(
wgpu::TexelCopyTextureInfo {
texture: &capture_texture,
mip_level: 0,
origin: wgpu::Origin3d::ZERO,
aspect: wgpu::TextureAspect::All,
},
wgpu::TexelCopyBufferInfo {
buffer: &output_buffer,
layout: wgpu::TexelCopyBufferLayout {
offset: 0,
bytes_per_row: Some(padded_bytes_per_row),
rows_per_image: Some(settings.height),
},
},
wgpu::Extent3d {
width: settings.width,
height: settings.height,
depth_or_array_layers: 1,
},
);
core.queue.submit(Some(encoder.finish()));
let buffer_slice = output_buffer.slice(..);
let (tx, rx) = std::sync::mpsc::channel();
buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
tx.send(result).unwrap();
});
let _ = core
.device
.poll(wgpu::PollType::wait_indefinitely())
.unwrap();
rx.recv().unwrap().unwrap();
let padded_data = buffer_slice.get_mapped_range().to_vec();
let mut unpadded_data = Vec::with_capacity((settings.width * settings.height * 4) as usize);
for chunk in padded_data.chunks(padded_bytes_per_row as usize) {
unpadded_data.extend_from_slice(&chunk[..unpadded_bytes_per_row as usize]);
}
Ok(unpadded_data)
}
}