src/stg_pass.rs
Repository source snapshot · Plain text · GitHub
1//! Surface-independent STG-Lite pass. The host owns the device, attachments,2//! camera, playback clock, submission and presentation. Time is normalized;3//! this API never guesses the capture duration from a PLY file.4use crate::stg_prepare::StgPreparation;5use crate::{parse_stg_ply, CameraUniform};6use glam::Mat4;7use wgpu::util::DeviceExt;89#[derive(Clone, Copy, Debug)]10pub struct PreparedFrame {11 pub source_count: usize,12 /// CPU path count. In GPU mode the count lives in indirect arguments;13 /// this field is zero/unknown and must not be used to issue a direct draw.14 pub visible_count: usize,15 pub index_upload_bytes: usize,16}1718pub struct StgPass {19 pipeline: wgpu::RenderPipeline,20 bind_group: wgpu::BindGroup,21 camera_buffer: wgpu::Buffer,22 index_buffer: wgpu::Buffer,23 preparation: StgPreparation,24 cpu_variant: u8,25 vertices: u32,26 pub gpu: Option<crate::gpu_prepare::GpuPreparation>,27 gpu_frame: Option<(f32, Mat4, Mat4)>,28}2930impl StgPass {31 /// Opt-in GPU path. Hosts call encode_prepare before beginning their render pass.32 pub fn new_gpu(33 device: &wgpu::Device,34 ply: &[u8],35 color_format: wgpu::TextureFormat,36 depth: Option<(wgpu::TextureFormat, wgpu::CompareFunction)>,37 ) -> Result<Self, String> {38 Self::with_options(device, ply, color_format, depth, 11, 0)39 }40 /// `depth` must match the host attachment. Gaussian depth writes are always41 /// disabled. Choose LessEqual for conventional Z or GreaterEqual for reverse Z.42 /// Colors retain the baseline's learned RGB convention; the host must decide43 /// how to convert that convention before compositing into an HDR scene.44 pub fn new(45 device: &wgpu::Device,46 ply: &[u8],47 color_format: wgpu::TextureFormat,48 depth: Option<(wgpu::TextureFormat, wgpu::CompareFunction)>,49 ) -> Result<Self, String> {50 Self::with_options(device, ply, color_format, depth, 9, 0)51 }5253 /// Benchmark options retain the unmodified baseline for paired comparisons.54 pub fn with_options(55 device: &wgpu::Device,56 ply: &[u8],57 color_format: wgpu::TextureFormat,58 depth: Option<(wgpu::TextureFormat, wgpu::CompareFunction)>,59 cpu_variant: u8,60 gpu_variant: u8,61 ) -> Result<Self, String> {62 if cpu_variant == 11 && gpu_variant != 0 {63 return Err(64 "GPU preparation requires the unchanged STG source layout (gpu_variant=0)".into(),65 );66 }67 let parsed = parse_stg_ply(ply)?;68 let mut payload = ply[parsed.payload_offset..parsed.payload_end].to_vec();69 if gpu_variant >= 1 {70 for row in payload.chunks_exact_mut(128) {71 for field in [4, 20, 21, 22, 23] {72 let v = f32::from_le_bytes(row[field * 4..field * 4 + 4].try_into().unwrap());73 let activated = if field == 20 {74 crate::sigmoid(v)75 } else if field == 4 {76 v.exp().max(1e-6)77 } else {78 v.exp()79 };80 row[field * 4..field * 4 + 4].copy_from_slice(&activated.to_le_bytes());81 }82 }83 }84 if payload.len() as u64 > u64::from(device.limits().max_storage_buffer_binding_size) {85 return Err("STG source exceeds the device storage binding limit".into());86 }87 let camera_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {88 label: Some("stg-host-camera"),89 contents: bytemuck::bytes_of(&<CameraUniform as bytemuck::Zeroable>::zeroed()),90 usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,91 });92 let source_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {93 label: Some("stg-host-source"),94 contents: &payload,95 usage: wgpu::BufferUsages::STORAGE,96 });97 let count = parsed.source.len();98 let gpu = (cpu_variant == 11)99 .then(|| crate::gpu_prepare::GpuPreparation::new(device, &source_buffer, count as u32));100 let index_buffer = device.create_buffer(&wgpu::BufferDescriptor {101 label: Some("stg-host-order"),102 size: (count * 4) as u64,103 usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,104 mapped_at_creation: false,105 });106 let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {107 label: Some("stg-host-layout"),108 entries: &[109 wgpu::BindGroupLayoutEntry {110 binding: 0,111 visibility: wgpu::ShaderStages::VERTEX,112 ty: wgpu::BindingType::Buffer {113 ty: wgpu::BufferBindingType::Uniform,114 has_dynamic_offset: false,115 min_binding_size: None,116 },117 count: None,118 },119 wgpu::BindGroupLayoutEntry {120 binding: 1,121 visibility: wgpu::ShaderStages::VERTEX,122 ty: wgpu::BindingType::Buffer {123 ty: wgpu::BufferBindingType::Storage { read_only: true },124 has_dynamic_offset: false,125 min_binding_size: None,126 },127 count: None,128 },129 ],130 });131 let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {132 label: Some("stg-host-bindings"),133 layout: &layout,134 entries: &[135 wgpu::BindGroupEntry {136 binding: 0,137 resource: camera_buffer.as_entire_binding(),138 },139 wgpu::BindGroupEntry {140 binding: 1,141 resource: source_buffer.as_entire_binding(),142 },143 ],144 });145 let mut shader_code = include_str!("splat.wgsl").to_string();146 if gpu_variant >= 1 {147 shader_code = shader_code148 .replace("max(exp(splat.r1.x), 0.000001)", "splat.r1.x")149 .replace("(1.0 / (1.0 + exp(-splat.r5.x)))", "splat.r5.x")150 .replace("exp(splat.r5.yzw)", "splat.r5.yzw");151 }152 if gpu_variant >= 2 {153 let begin = shader_code.find("var corners =").unwrap();154 let end = shader_code[begin..].find("return corners[index];").unwrap() + begin;155 shader_code.replace_range(begin..end, "var corners = array<vec2<f32>, 4>(vec2<f32>(-1.0,-1.0),vec2<f32>(1.0,-1.0),vec2<f32>(-1.0,1.0),vec2<f32>(1.0,1.0));\n ");156 }157 let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {158 label: Some("stg-shared-wgsl"),159 source: wgpu::ShaderSource::Wgsl(shader_code.into()),160 });161 let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {162 label: Some("stg-host-pipeline-layout"),163 bind_group_layouts: &[Some(&layout)],164 immediate_size: 0,165 });166 let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {167 label: Some("stg-host-pipeline"),168 layout: Some(&pipeline_layout),169 vertex: wgpu::VertexState {170 module: &shader,171 entry_point: Some("vs_main"),172 compilation_options: Default::default(),173 buffers: &[Some(wgpu::VertexBufferLayout {174 array_stride: 4,175 step_mode: wgpu::VertexStepMode::Instance,176 attributes: &[wgpu::VertexAttribute {177 format: wgpu::VertexFormat::Uint32,178 offset: 0,179 shader_location: 0,180 }],181 })],182 },183 fragment: Some(wgpu::FragmentState {184 module: &shader,185 entry_point: Some("fs_main"),186 compilation_options: Default::default(),187 targets: &[Some(wgpu::ColorTargetState {188 format: color_format,189 blend: Some(wgpu::BlendState::PREMULTIPLIED_ALPHA_BLENDING),190 write_mask: wgpu::ColorWrites::ALL,191 })],192 }),193 primitive: wgpu::PrimitiveState {194 topology: if gpu_variant >= 2 {195 wgpu::PrimitiveTopology::TriangleStrip196 } else {197 wgpu::PrimitiveTopology::TriangleList198 },199 ..Default::default()200 },201 depth_stencil: depth.map(|(format, depth_compare)| wgpu::DepthStencilState {202 format,203 depth_write_enabled: Some(false),204 depth_compare: Some(depth_compare),205 stencil: Default::default(),206 bias: Default::default(),207 }),208 multisample: Default::default(),209 multiview_mask: None,210 cache: None,211 });212 Ok(Self {213 pipeline,214 bind_group,215 camera_buffer,216 index_buffer,217 preparation: StgPreparation::new(parsed.source),218 cpu_variant,219 vertices: if gpu_variant >= 2 { 4 } else { 6 },220 gpu,221 gpu_frame: None,222 })223 }224225 /// Matrices are column-major (WebGPU clip Z in 0..1, RH view forward -Z).226 /// Set the render pass viewport to the same pixel dimensions before draw.227 /// Submit each prepared frame before preparing another using these buffers.228 pub fn prepare(229 &mut self,230 queue: &wgpu::Queue,231 normalized_time: f32,232 view_columns: [[f32; 4]; 4],233 projection_columns: [[f32; 4]; 4],234 viewport: [u32; 2],235 ) -> Result<PreparedFrame, String> {236 if !normalized_time.is_finite()237 || !(0.0..=1.0).contains(&normalized_time)238 || viewport.contains(&0)239 || view_columns240 .iter()241 .chain(&projection_columns)242 .flatten()243 .any(|x| !x.is_finite())244 {245 return Err("Invalid time, viewport or camera matrix".into());246 }247 let view = Mat4::from_cols_array_2d(&view_columns);248 let projection = Mat4::from_cols_array_2d(&projection_columns);249 let view_projection = projection * view;250 let [width, height] = viewport.map(|v| v as f32);251 if self.gpu.is_some() {252 self.gpu_frame = Some((normalized_time, view, projection));253 } else {254 self.preparation255 .prepare(normalized_time, view, projection, self.cpu_variant);256 }257 let camera = CameraUniform {258 view_proj: view_projection.to_cols_array_2d(),259 viewport: [width, height, width.recip(), height.recip()],260 scene: [normalized_time, 0.0, 0.0, 0.0],261 };262 queue.write_buffer(&self.camera_buffer, 0, bytemuck::bytes_of(&camera));263 if !self.preparation.indices().is_empty() && !self.preparation.cache_hit {264 queue.write_buffer(265 &self.index_buffer,266 0,267 bytemuck::cast_slice(self.preparation.indices()),268 );269 }270 Ok(PreparedFrame {271 source_count: self.preparation.len(),272 visible_count: self.preparation.indices().len(),273 index_upload_bytes: if self.preparation.cache_hit {274 0275 } else {276 self.preparation.indices().len() * 4277 },278 })279 }280281 pub fn draw(&self, pass: &mut wgpu::RenderPass<'_>) {282 pass.set_pipeline(&self.pipeline);283 pass.set_bind_group(0, &self.bind_group, &[]);284 if let Some(gpu) = &self.gpu {285 pass.set_vertex_buffer(0, gpu.indices.slice(..));286 pass.draw_indirect(&gpu.indirect, 0);287 } else {288 pass.set_vertex_buffer(0, self.index_buffer.slice(..));289 pass.draw(0..self.vertices, 0..self.preparation.indices().len() as u32);290 }291 }292293 pub fn encode_prepare(294 &mut self,295 queue: &wgpu::Queue,296 encoder: &mut wgpu::CommandEncoder,297 queries: Option<(&wgpu::QuerySet, u32, u32)>,298 ) -> bool {299 if let (Some(gpu), Some((time, view, projection))) = (&mut self.gpu, self.gpu_frame) {300 gpu.encode(queue, encoder, time, view, projection, queries)301 } else {302 false303 }304 }305}306