diff options
Diffstat (limited to 'flowy/src/motion/dec.rs')
-rw-r--r-- | flowy/src/motion/dec.rs | 156 |
1 files changed, 155 insertions, 1 deletions
diff --git a/flowy/src/motion/dec.rs b/flowy/src/motion/dec.rs index 84aeff7..9d15259 100644 --- a/flowy/src/motion/dec.rs +++ b/flowy/src/motion/dec.rs @@ -1,4 +1,158 @@ -pub struct MotionDecoder {} +use super::{CommonBuffers, Params, RoundParams}; +use bytemuck::{Pod, Zeroable}; +use std::mem::size_of; +use wgpu::{ + include_wgsl, BindGroup, BindGroupDescriptor, BindGroupEntry, BindGroupLayoutDescriptor, + BindGroupLayoutEntry, BindingType, Buffer, BufferDescriptor, BufferUsages, CommandEncoder, + ComputePipeline, ComputePipelineDescriptor, Device, PipelineLayoutDescriptor, Queue, + ShaderStages, TextureSampleType, TextureViewDimension, +}; + +pub struct MotionDecoder { + pipeline: ComputePipeline, + bind_groups: [BindGroup; 2], + + uniform_buffer: Buffer, + uniform: DecoderUniform, +} + +#[repr(C)] +#[derive(Debug, Clone, Copy, Pod, Zeroable, Default)] +pub struct DecoderUniform { + block_size: [i32; 2], + output_stride: i32, + _pad: u32, +} + +impl MotionDecoder { + pub fn create(device: &Device, params: &Params, bufs: &CommonBuffers) -> Self { + let uniform_buffer = device.create_buffer(&BufferDescriptor { + label: Some("encoder uniforms"), + size: size_of::<DecoderUniform>() as u64, + usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + let uniform = DecoderUniform { + block_size: [8, 8], + output_stride: (params.width / params.block_width) as i32, + ..Default::default() + }; + + let bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor { + label: None, + entries: &[ + BindGroupLayoutEntry { + binding: 0, + count: None, + ty: BindingType::Buffer { + ty: wgpu::BufferBindingType::Uniform, + has_dynamic_offset: false, + min_binding_size: None, + }, + visibility: ShaderStages::COMPUTE, + }, + BindGroupLayoutEntry { + binding: 1, + count: None, + ty: BindingType::Buffer { + has_dynamic_offset: false, + min_binding_size: None, + ty: wgpu::BufferBindingType::Storage { read_only: true }, + }, + visibility: ShaderStages::COMPUTE, + }, + BindGroupLayoutEntry { + binding: 2, + count: None, + ty: BindingType::StorageTexture { + access: wgpu::StorageTextureAccess::WriteOnly, + format: wgpu::TextureFormat::Bgra8Unorm, + view_dimension: TextureViewDimension::D2, + }, + visibility: ShaderStages::COMPUTE, + }, + BindGroupLayoutEntry { + binding: 3, + count: None, + ty: BindingType::Texture { + sample_type: TextureSampleType::Float { filterable: false }, + view_dimension: TextureViewDimension::D2, + multisampled: false, + }, + visibility: ShaderStages::COMPUTE, + }, + ], + }); + + let bind_groups = [0, 1].map(|i| { + device.create_bind_group(&BindGroupDescriptor { + label: None, + layout: &bind_group_layout, + entries: &[ + BindGroupEntry { + binding: 0, + resource: uniform_buffer.as_entire_binding(), + }, + BindGroupEntry { + binding: 1, + resource: bufs.offsets.as_entire_binding(), + }, + BindGroupEntry { + binding: 2, + resource: wgpu::BindingResource::TextureView( + &bufs.textures[i].create_view(&Default::default()), + ), + }, + BindGroupEntry { + binding: 3, + resource: wgpu::BindingResource::TextureView( + &bufs.textures[1 - i].create_view(&Default::default()), + ), + }, + ], + }) + }); + + let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor { + label: None, + bind_group_layouts: &[&bind_group_layout], + push_constant_ranges: &[], + }); + + let module = device.create_shader_module(include_wgsl!("dec.wgsl")); + let pipeline = device.create_compute_pipeline(&ComputePipelineDescriptor { + label: None, + layout: Some(&pipeline_layout), + module: &module, + entry_point: "main", + }); + + Self { + bind_groups, + uniform, + uniform_buffer, + pipeline, + } + } + + pub fn write_uniforms(&self, queue: &Queue) { + queue.write_buffer( + &self.uniform_buffer, + 0, + bytemuck::cast_slice(&[self.uniform]), + ) + } + + pub fn pass(&self, encoder: &mut CommandEncoder, params: &Params, rp: &RoundParams) { + let mut cpass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: None, + timestamp_writes: None, + }); + cpass.set_pipeline(&self.pipeline); + cpass.set_bind_group(0, &self.bind_groups[rp.swap], &[]); + cpass.dispatch_workgroups(params.blocks_x as u32, params.blocks_y as u32, 1); + } +} /* BindGroupLayoutEntry { |