new streaming system

1 files changed, 156 insertions, 0 deletions

diff --git a/flowy/src/motion/enc.rs b/flowy/src/motion/enc.rs
new file mode 100644
index 0000000..aa45d7b
--- /dev/null
+++ b/flowy/src/motion/enc.rs
@@ -0,0 +1,156 @@
+use std::mem::size_of;
+
+use super::{CommonBuffers, Params, RoundParams};
+use bytemuck::{Pod, Zeroable};
+use wgpu::{
+    include_wgsl, BindGroup, BindGroupDescriptor, BindGroupEntry, BindGroupLayoutDescriptor,
+    BindGroupLayoutEntry, BindingType, Buffer, BufferDescriptor, BufferUsages, CommandEncoder,
+    ComputePipeline, ComputePipelineDescriptor, Device, PipelineLayoutDescriptor, Queue,
+    ShaderStages, TextureFormat, TextureSampleType, TextureViewDimension,
+};
+
+pub struct MotionEncoder {
+    pipeline: ComputePipeline,
+    bind_groups: [BindGroup; 2],
+
+    uniform_buffer: Buffer,
+    uniform: EncoderUniform,
+}
+
+#[repr(C)]
+#[derive(Debug, Clone, Copy, Pod, Zeroable, Default)]
+pub struct EncoderUniform {
+    block_size: [i32; 2],
+    output_stride: i32,
+    _pad: u32,
+}
+
+impl MotionEncoder {
+    pub fn create(device: &Device, params: &Params, bufs: &CommonBuffers) -> Self {
+        let uniform_buffer = device.create_buffer(&BufferDescriptor {
+            label: Some("encoder uniforms"),
+            size: size_of::<EncoderUniform>() as u64,
+            usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST,
+            mapped_at_creation: false,
+        });
+        let uniform = EncoderUniform {
+            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: false },
+                    },
+                    visibility: ShaderStages::COMPUTE,
+                },
+                BindGroupLayoutEntry {
+                    binding: 2,
+                    count: None,
+                    ty: BindingType::Texture {
+                        sample_type: TextureSampleType::Float { filterable: false },
+                        view_dimension: TextureViewDimension::D2,
+                        multisampled: false,
+                    },
+                    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[1 - 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!("enc.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.extent.width, params.extent.height, 1);
+    }
+}