Files
mandelbrot/src/fractal/renderer.rs
T

681 lines
25 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! wgpu resources for the fractal: the render pipeline, the uniform buffer, the
//! reference-orbit storage buffer, and the egui paint callback that drives them.
//!
//! Rendering strategy: the expensive per-pixel perturbation shader renders into
//! an offscreen **cache texture**, and only when the view/coloring/size actually
//! change (tracked by `rendered`). Every egui frame then just blits that cached
//! texture onto egui's surface with a cheap textured fullscreen triangle — so
//! incidental repaints (mouse-move, hover, the worker-pending poll) cost a blit,
//! not a full fractal recompute. The fragment shader iterates each pixel as an
//! f32 perturbation delta from the reference orbit stored in `ref_buffer`.
use std::sync::Arc;
use eframe::egui_wgpu::{self, wgpu};
/// Maximum reference-orbit length (points) the storage buffer can hold. Also
/// bounds the iteration count. 128k points * 8 bytes = 1 MiB.
pub const MAX_REF_POINTS: usize = 1 << 17;
/// GPU-side view + coloring parameters. Layout must match `Uniforms` in the
/// WGSL shader; total size is a multiple of 16 bytes for uniform-buffer rules.
#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
pub struct Uniforms {
/// Complex-plane span (width, height) covered by the view. Per-pixel `dc`
/// is `centered * span`, where `centered` is in [-0.5, 0.5].
pub span: [f32; 2],
pub max_iter: u32,
pub ref_len: u32,
pub color_offset: f32,
pub color_scale: f32,
pub bailout_sq: f32,
/// 0 = Mandelbrot, 1 = Julia.
pub is_julia: u32,
pub palette_id: u32,
/// Supersampling factor per axis: 1 = off, 2 = 2×2 (4 samples).
pub aa_level: u32,
/// Iteration formula (`FractalKind::shader_id`).
pub kind: u32,
/// Exponent for the Multibrot kind.
pub power: u32,
/// Complex offset of the view center from the reference center, so a stale
/// or reused reference (computed at a slightly different center) still maps
/// correctly. Added to every pixel's per-pixel offset.
pub dc_offset: [f32; 2],
/// 0 = escape-time coloring, 1 = distance-estimation shading.
pub de_coloring: u32,
/// Padding to a 16-byte multiple (uniform buffer requirement).
pub _pad: u32,
}
/// Offscreen texture the fractal is rendered into, plus the bind group used to
/// blit it. Recreated whenever the widget's pixel size changes.
struct CacheTarget {
view: wgpu::TextureView,
blit_bind_group: wgpu::BindGroup,
width: u32,
height: u32,
}
/// State the cache texture was last rendered with. If the next frame's inputs
/// match this, the cache is still valid and the fractal shader is skipped.
struct RenderedState {
uniforms: Uniforms,
generation: u64,
width: u32,
height: u32,
}
pub struct FractalRenderer {
pipeline: wgpu::RenderPipeline,
bind_group_layout: wgpu::BindGroupLayout,
uniform_buffer: wgpu::Buffer,
ref_buffer: wgpu::Buffer,
bind_group: wgpu::BindGroup,
target_format: wgpu::TextureFormat,
/// Generation of the reference orbit currently uploaded to `ref_buffer`.
uploaded_generation: u64,
/// Blit pipeline + resources that copy the cache texture to egui's surface.
blit_pipeline: wgpu::RenderPipeline,
blit_bind_group_layout: wgpu::BindGroupLayout,
blit_sampler: wgpu::Sampler,
/// The offscreen cache; `None` until the first frame sizes it.
cache: Option<CacheTarget>,
/// What the cache currently holds; `None` forces a re-render.
rendered: Option<RenderedState>,
}
impl FractalRenderer {
pub fn new(device: &wgpu::Device, target_format: wgpu::TextureFormat) -> Self {
let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some("mandelbrot"),
source: wgpu::ShaderSource::Wgsl(include_str!("../shaders/mandelbrot.wgsl").into()),
});
let uniform_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("fractal uniforms"),
size: std::mem::size_of::<Uniforms>() as u64,
usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let ref_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("reference orbit"),
size: (MAX_REF_POINTS * std::mem::size_of::<[f32; 2]>()) as u64,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("fractal bind group layout"),
entries: &[
wgpu::BindGroupLayoutEntry {
binding: 0,
visibility: wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Uniform,
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 1,
visibility: wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Storage { read_only: true },
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
],
});
let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("fractal bind group"),
layout: &bind_group_layout,
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: uniform_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: ref_buffer.as_entire_binding(),
},
],
});
let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("fractal pipeline layout"),
bind_group_layouts: &[Some(&bind_group_layout)],
immediate_size: 0,
});
let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
label: Some("fractal pipeline"),
layout: Some(&pipeline_layout),
vertex: wgpu::VertexState {
module: &shader,
entry_point: Some("vs_main"),
buffers: &[],
compilation_options: Default::default(),
},
fragment: Some(wgpu::FragmentState {
module: &shader,
entry_point: Some("fs_main"),
targets: &[Some(wgpu::ColorTargetState {
format: target_format,
blend: None,
write_mask: wgpu::ColorWrites::ALL,
})],
compilation_options: Default::default(),
}),
primitive: wgpu::PrimitiveState::default(),
depth_stencil: None,
multisample: wgpu::MultisampleState::default(),
multiview_mask: None,
cache: None,
});
// Blit pipeline: samples the cache texture onto egui's surface.
let blit_shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some("blit"),
source: wgpu::ShaderSource::Wgsl(include_str!("../shaders/blit.wgsl").into()),
});
let blit_bind_group_layout =
device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("blit bind group layout"),
entries: &[
wgpu::BindGroupLayoutEntry {
binding: 0,
visibility: wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Texture {
sample_type: wgpu::TextureSampleType::Float { filterable: true },
view_dimension: wgpu::TextureViewDimension::D2,
multisampled: false,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 1,
visibility: wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
count: None,
},
],
});
let blit_sampler = device.create_sampler(&wgpu::SamplerDescriptor {
label: Some("blit sampler"),
mag_filter: wgpu::FilterMode::Linear,
min_filter: wgpu::FilterMode::Linear,
..Default::default()
});
let blit_pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("blit pipeline layout"),
bind_group_layouts: &[Some(&blit_bind_group_layout)],
immediate_size: 0,
});
let blit_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
label: Some("blit pipeline"),
layout: Some(&blit_pipeline_layout),
vertex: wgpu::VertexState {
module: &blit_shader,
entry_point: Some("vs_main"),
buffers: &[],
compilation_options: Default::default(),
},
fragment: Some(wgpu::FragmentState {
module: &blit_shader,
entry_point: Some("fs_main"),
targets: &[Some(wgpu::ColorTargetState {
format: target_format,
blend: None,
write_mask: wgpu::ColorWrites::ALL,
})],
compilation_options: Default::default(),
}),
primitive: wgpu::PrimitiveState::default(),
depth_stencil: None,
multisample: wgpu::MultisampleState::default(),
multiview_mask: None,
cache: None,
});
Self {
pipeline,
bind_group_layout,
uniform_buffer,
ref_buffer,
bind_group,
target_format,
uploaded_generation: u64::MAX,
blit_pipeline,
blit_bind_group_layout,
blit_sampler,
cache: None,
rendered: None,
}
}
/// Ensure the cache texture exists at `width`×`height`. Recreates it (and its
/// blit bind group) on a size change, invalidating any previous render.
fn ensure_cache(&mut self, device: &wgpu::Device, width: u32, height: u32) {
if let Some(c) = &self.cache
&& c.width == width
&& c.height == height
{
return;
}
let texture = device.create_texture(&wgpu::TextureDescriptor {
label: Some("fractal cache"),
size: wgpu::Extent3d {
width,
height,
depth_or_array_layers: 1,
},
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format: self.target_format,
usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
view_formats: &[],
});
let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
let blit_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("blit bind group"),
layout: &self.blit_bind_group_layout,
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: wgpu::BindingResource::TextureView(&view),
},
wgpu::BindGroupEntry {
binding: 1,
resource: wgpu::BindingResource::Sampler(&self.blit_sampler),
},
],
});
self.cache = Some(CacheTarget {
view,
blit_bind_group,
width,
height,
});
// New texture → old render is gone.
self.rendered = None;
}
/// Handles needed to build a standalone [`ExportRender`] off the UI thread:
/// the (immutable) pipeline and its bind-group layout, plus the target
/// format. Cloned so the caller can drop the render-state lock before use.
pub fn export_handles(&self) -> (wgpu::RenderPipeline, wgpu::BindGroupLayout, wgpu::TextureFormat) {
(
self.pipeline.clone(),
self.bind_group_layout.clone(),
self.target_format,
)
}
}
/// A self-contained render of one export image. It owns its own uniform and
/// reference buffers (a snapshot of the view at export time), so it is unaffected
/// by panning/zooming on the main thread, and can run on a background thread.
/// The image is rendered in horizontal tiles so progress can be reported as the
/// GPU works through it.
pub struct ExportRender {
pipeline: wgpu::RenderPipeline,
bind_group: wgpu::BindGroup,
texture: wgpu::Texture,
view: wgpu::TextureView,
readback: wgpu::Buffer,
/// Padded bytes-per-row of the readback buffer.
pub padded_bpr: u32,
pub width: u32,
pub height: u32,
/// Number of horizontal tiles the render is split into.
pub tiles: u32,
pub swap_rb: bool,
}
impl ExportRender {
/// Allocate the export's dedicated GPU resources and upload the snapshot.
#[allow(clippy::too_many_arguments)]
pub fn new(
device: &wgpu::Device,
queue: &wgpu::Queue,
pipeline: wgpu::RenderPipeline,
bind_group_layout: &wgpu::BindGroupLayout,
target_format: wgpu::TextureFormat,
width: u32,
height: u32,
uniforms: Uniforms,
reference: &[[f32; 2]],
) -> Self {
let uniform_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("export uniforms"),
size: std::mem::size_of::<Uniforms>() as u64,
usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
queue.write_buffer(&uniform_buffer, 0, bytemuck::bytes_of(&uniforms));
let count = reference.len().min(MAX_REF_POINTS);
let ref_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("export reference orbit"),
size: (count.max(1) * std::mem::size_of::<[f32; 2]>()) as u64,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
if count > 0 {
queue.write_buffer(&ref_buffer, 0, bytemuck::cast_slice(&reference[..count]));
}
let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("export bind group"),
layout: bind_group_layout,
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: uniform_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: ref_buffer.as_entire_binding(),
},
],
});
let texture = device.create_texture(&wgpu::TextureDescriptor {
label: Some("export target"),
size: wgpu::Extent3d {
width,
height,
depth_or_array_layers: 1,
},
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format: target_format,
usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
view_formats: &[],
});
let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
let padded_bpr = (width * 4).div_ceil(align) * align;
let readback = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("export readback"),
size: (padded_bpr * height) as u64,
usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
mapped_at_creation: false,
});
// ~128px bands, kept to a sane range so progress is smooth without too
// many submissions.
let tiles = (height / 128).clamp(8, 64).min(height.max(1));
let swap_rb = matches!(
target_format,
wgpu::TextureFormat::Bgra8Unorm | wgpu::TextureFormat::Bgra8UnormSrgb
);
Self {
pipeline,
bind_group,
texture,
view,
readback,
padded_bpr,
width,
height,
tiles,
swap_rb,
}
}
/// Pixel row range `[y0, y1)` covered by tile `t`.
fn tile_rows(&self, t: u32) -> (u32, u32) {
let band = self.height.div_ceil(self.tiles);
let y0 = (t * band).min(self.height);
let y1 = (y0 + band).min(self.height);
(y0, y1)
}
/// Render one horizontal tile into the export texture and submit it. Tile 0
/// clears the whole attachment; later tiles preserve earlier ones.
pub fn render_tile(&self, device: &wgpu::Device, queue: &wgpu::Queue, t: u32) {
let (y0, y1) = self.tile_rows(t);
if y1 <= y0 {
return;
}
let load = if t == 0 {
wgpu::LoadOp::Clear(wgpu::Color::BLACK)
} else {
wgpu::LoadOp::Load
};
let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
label: Some("export tile"),
});
{
let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
label: Some("export tile pass"),
color_attachments: &[Some(wgpu::RenderPassColorAttachment {
view: &self.view,
depth_slice: None,
resolve_target: None,
ops: wgpu::Operations {
load,
store: wgpu::StoreOp::Store,
},
})],
depth_stencil_attachment: None,
timestamp_writes: None,
occlusion_query_set: None,
multiview_mask: None,
});
// Full-viewport triangle (so pixel→plane mapping matches the whole
// image), scissored to this tile's rows.
pass.set_scissor_rect(0, y0, self.width, y1 - y0);
pass.set_pipeline(&self.pipeline);
pass.set_bind_group(0, &self.bind_group, &[]);
pass.draw(0..3, 0..1);
}
queue.submit(std::iter::once(encoder.finish()));
}
/// Copy the finished texture into the mappable readback buffer and submit.
pub fn copy_to_readback(&self, device: &wgpu::Device, queue: &wgpu::Queue) {
let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
label: Some("export copy"),
});
encoder.copy_texture_to_buffer(
wgpu::TexelCopyTextureInfo {
texture: &self.texture,
mip_level: 0,
origin: wgpu::Origin3d::ZERO,
aspect: wgpu::TextureAspect::All,
},
wgpu::TexelCopyBufferInfo {
buffer: &self.readback,
layout: wgpu::TexelCopyBufferLayout {
offset: 0,
bytes_per_row: Some(self.padded_bpr),
rows_per_image: Some(self.height),
},
},
wgpu::Extent3d {
width: self.width,
height: self.height,
depth_or_array_layers: 1,
},
);
queue.submit(std::iter::once(encoder.finish()));
}
/// The mappable readback buffer (valid after [`copy_to_readback`]).
pub fn readback(&self) -> &wgpu::Buffer {
&self.readback
}
}
/// Convert a padded BGRA/RGBA readback into tightly-packed RGBA8 and encode it
/// as PNG bytes, reporting progress in `[0, 1]` via `on_progress` as rows are
/// streamed to the compressor (encoding is the slow, subdividable phase).
pub fn encode_png_with_progress(
padded: &[u8],
width: u32,
height: u32,
padded_bpr: u32,
swap_rb: bool,
mut on_progress: impl FnMut(f32),
) -> Vec<u8> {
use std::io::Write as _;
let row = (width * 4) as usize;
let mut out = Vec::new();
{
let mut encoder = png::Encoder::new(&mut out, width, height);
encoder.set_color(png::ColorType::Rgba);
encoder.set_depth(png::BitDepth::Eight);
let mut writer = encoder.write_header().expect("png header");
let mut stream = writer.stream_writer().expect("png stream");
let mut line = vec![0u8; row];
for y in 0..height as usize {
let src_off = y * padded_bpr as usize;
let src = &padded[src_off..src_off + row];
if swap_rb {
for x in 0..width as usize {
line[x * 4] = src[x * 4 + 2];
line[x * 4 + 1] = src[x * 4 + 1];
line[x * 4 + 2] = src[x * 4];
line[x * 4 + 3] = src[x * 4 + 3];
}
stream.write_all(&line).expect("png data");
} else {
stream.write_all(src).expect("png data");
}
if y % 64 == 0 {
on_progress(y as f32 / height as f32);
}
}
stream.finish().expect("png finish");
}
on_progress(1.0);
out
}
/// A per-frame paint callback. Carries this frame's uniforms plus a reference to
/// the current reference orbit (cheap `Arc` clone). The orbit is only re-uploaded
/// to the GPU when its `generation` changes, and the fractal is only re-rendered
/// into the cache when the uniforms, generation, or `size_px` change.
pub struct FractalCallback {
pub uniforms: Uniforms,
pub reference: Arc<Vec<[f32; 2]>>,
pub generation: u64,
/// Widget size in physical pixels — the cache texture resolution.
pub size_px: [u32; 2],
}
impl egui_wgpu::CallbackTrait for FractalCallback {
fn prepare(
&self,
device: &wgpu::Device,
queue: &wgpu::Queue,
_screen_descriptor: &egui_wgpu::ScreenDescriptor,
egui_encoder: &mut wgpu::CommandEncoder,
resources: &mut egui_wgpu::CallbackResources,
) -> Vec<wgpu::CommandBuffer> {
let Some(renderer) = resources.get_mut::<FractalRenderer>() else {
return Vec::new();
};
let width = self.size_px[0].max(1);
let height = self.size_px[1].max(1);
renderer.ensure_cache(device, width, height);
if renderer.uploaded_generation != self.generation && !self.reference.is_empty() {
let count = self.reference.len().min(MAX_REF_POINTS);
queue.write_buffer(
&renderer.ref_buffer,
0,
bytemuck::cast_slice(&self.reference[..count]),
);
renderer.uploaded_generation = self.generation;
}
// Re-render the cache only when what it depends on changed.
let dirty = renderer.rendered.as_ref().is_none_or(|r| {
r.generation != self.generation
|| r.width != width
|| r.height != height
|| bytemuck::bytes_of(&r.uniforms) != bytemuck::bytes_of(&self.uniforms)
});
if !dirty {
return Vec::new();
}
queue.write_buffer(
&renderer.uniform_buffer,
0,
bytemuck::bytes_of(&self.uniforms),
);
if let Some(cache) = &renderer.cache {
let mut pass = egui_encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
label: Some("fractal cache pass"),
color_attachments: &[Some(wgpu::RenderPassColorAttachment {
view: &cache.view,
depth_slice: None,
resolve_target: None,
ops: wgpu::Operations {
load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
store: wgpu::StoreOp::Store,
},
})],
depth_stencil_attachment: None,
timestamp_writes: None,
occlusion_query_set: None,
multiview_mask: None,
});
pass.set_pipeline(&renderer.pipeline);
pass.set_bind_group(0, &renderer.bind_group, &[]);
pass.draw(0..3, 0..1);
}
renderer.rendered = Some(RenderedState {
uniforms: self.uniforms,
generation: self.generation,
width,
height,
});
Vec::new()
}
fn paint(
&self,
_info: egui::PaintCallbackInfo,
render_pass: &mut wgpu::RenderPass<'static>,
resources: &egui_wgpu::CallbackResources,
) {
if let Some(renderer) = resources.get::<FractalRenderer>()
&& let Some(cache) = &renderer.cache
{
render_pass.set_pipeline(&renderer.blit_pipeline);
render_pass.set_bind_group(0, &cache.blit_bind_group, &[]);
render_pass.draw(0..3, 0..1);
}
}
}