src/stg_pass.rs
与本次数学页面一起发布的源码快照。SHA-256: 34355b0b181f6afff032b67568e079409100488c68161b3c4ff523c360201907
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;
8
9#[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}
17
18pub 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}
29
30impl 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 always
41 /// disabled. Choose LessEqual for conventional Z or GreaterEqual for reverse Z.
42 /// Colors retain the baseline's learned RGB convention; the host must decide
43 /// 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 }
52
53 /// 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_code
148 .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::TriangleStrip
196 } else {
197 wgpu::PrimitiveTopology::TriangleList
198 },
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 }
224
225 /// 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_columns
240 .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.preparation
255 .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 0
275 } else {
276 self.preparation.indices().len() * 4
277 },
278 })
279 }
280
281 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 }
292
293 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 false
303 }
304 }
305}
306