gpu pipeline

2 files changed, 218 insertions, 0 deletions

diff --git a/flowy/src/main.rs b/flowy/src/main.rs
new file mode 100644
index 0000000..d225c00
--- /dev/null
+++ b/flowy/src/main.rs
@@ -0,0 +1,207 @@
+use std::io::{stdin, stdout, Read, Write};
+
+use pollster::FutureExt;
+use wgpu::{
+    include_wgsl, Backends, BindGroupDescriptor, BindGroupEntry, BindGroupLayoutDescriptor,
+    BindGroupLayoutEntry, BindingType, CommandEncoder, ComputePipelineDescriptor, DeviceDescriptor,
+    Extent3d, Features, ImageCopyTexture, Instance, InstanceDescriptor, InstanceFlags, Limits,
+    MaintainBase, Origin3d, PipelineLayoutDescriptor, PowerPreference, RequestAdapterOptions,
+    ShaderStages, TextureDescriptor, TextureDimension, TextureFormat, TextureSampleType,
+    TextureUsages, TextureViewDimension,
+};
+
+fn main() {
+    let instance = Instance::new(InstanceDescriptor {
+        backends: Backends::all(),
+        ..Default::default()
+    });
+
+    let adapter = instance
+        .request_adapter(&RequestAdapterOptions {
+            power_preference: PowerPreference::HighPerformance,
+            force_fallback_adapter: false,
+            compatible_surface: None,
+        })
+        .block_on()
+        .unwrap();
+
+    let (device, queue) = adapter
+        .request_device(
+            &DeviceDescriptor {
+                label: None,
+                features: Features::BGRA8UNORM_STORAGE,
+                limits: Limits::default(),
+            },
+            None,
+        )
+        .block_on()
+        .unwrap();
+
+    let (width, height) = (1920, 1080);
+    let extent = Extent3d {
+        width: width as u32,
+        height: height as u32,
+        depth_or_array_layers: 1,
+    };
+
+    let [texa, texb] = [(), ()].map(|_| {
+        device.create_texture(&TextureDescriptor {
+            label: None,
+            size: extent,
+            mip_level_count: 1,
+            sample_count: 1,
+            dimension: TextureDimension::D2,
+            format: TextureFormat::Bgra8Unorm,
+            usage: TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_DST,
+            view_formats: &[],
+        })
+    });
+    let out_tex = device.create_texture(&TextureDescriptor {
+        label: None,
+        size: extent,
+        mip_level_count: 1,
+        sample_count: 1,
+        dimension: TextureDimension::D2,
+        format: TextureFormat::Bgra8Unorm,
+        usage: TextureUsages::STORAGE_BINDING | TextureUsages::COPY_SRC,
+        view_formats: &[],
+    });
+
+    let mut buffer = vec![0u8; width * height * 4];
+
+    let download_buffer = device.create_buffer(&wgpu::BufferDescriptor {
+        label: None,
+        size: buffer.len() as u64,
+        usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
+        mapped_at_creation: false,
+    });
+
+    let bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
+        label: None,
+        entries: &[
+            BindGroupLayoutEntry {
+                binding: 0,
+                count: None,
+                ty: wgpu::BindingType::StorageTexture {
+                    access: wgpu::StorageTextureAccess::WriteOnly,
+                    format: TextureFormat::Bgra8Unorm,
+                    view_dimension: wgpu::TextureViewDimension::D2,
+                },
+                visibility: ShaderStages::COMPUTE,
+            },
+            BindGroupLayoutEntry {
+                binding: 1,
+                count: None,
+                ty: BindingType::Texture {
+                    sample_type: TextureSampleType::Float { filterable: false },
+                    view_dimension: TextureViewDimension::D2,
+                    multisampled: false,
+                },
+                visibility: ShaderStages::COMPUTE,
+            },
+        ],
+    });
+
+    let bind_goup = device.create_bind_group(&BindGroupDescriptor {
+        label: None,
+        layout: &bind_group_layout,
+        entries: &[
+            BindGroupEntry {
+                binding: 0,
+                resource: wgpu::BindingResource::TextureView(
+                    &out_tex.create_view(&Default::default()),
+                ),
+            },
+            BindGroupEntry {
+                binding: 1,
+                resource: wgpu::BindingResource::TextureView(
+                    &texa.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!("shader.wgsl"));
+    let pipeline = device.create_compute_pipeline(&ComputePipelineDescriptor {
+        label: None,
+        layout: Some(&pipeline_layout),
+        module: &module,
+        entry_point: "main",
+    });
+
+    loop {
+        eprintln!("read");
+        stdin().read_exact(&mut buffer).unwrap();
+
+        eprintln!("upload");
+        queue.write_texture(
+            ImageCopyTexture {
+                aspect: wgpu::TextureAspect::All,
+                mip_level: 0,
+                origin: Origin3d::ZERO,
+                texture: &texa,
+            },
+            &buffer,
+            wgpu::ImageDataLayout {
+                offset: 0,
+                bytes_per_row: Some(extent.width * 4),
+                rows_per_image: Some(extent.height),
+            },
+            extent,
+        );
+
+        eprintln!("compute");
+        let mut encoder = device.create_command_encoder(&Default::default());
+
+        {
+            let mut cpass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
+                label: None,
+                timestamp_writes: None,
+            });
+            cpass.set_pipeline(&pipeline);
+            cpass.set_bind_group(0, &bind_goup, &[]);
+            cpass.dispatch_workgroups(width as u32, height as u32, 1);
+        }
+
+        encoder.copy_texture_to_buffer(
+            wgpu::ImageCopyTexture {
+                texture: &out_tex,
+                mip_level: 0,
+                origin: wgpu::Origin3d::ZERO,
+                aspect: wgpu::TextureAspect::All,
+            },
+            wgpu::ImageCopyBuffer {
+                buffer: &download_buffer,
+                layout: wgpu::ImageDataLayout {
+                    offset: 0,
+                    bytes_per_row: Some(extent.width * 4),
+                    rows_per_image: Some(extent.height),
+                },
+            },
+            extent,
+        );
+
+        queue.submit(Some(encoder.finish()));
+        device.poll(MaintainBase::Wait);
+
+        eprintln!("download");
+        let buffer_slice = download_buffer.slice(..);
+        let (sender, receiver) = oneshot::channel();
+        buffer_slice.map_async(wgpu::MapMode::Read, move |r| sender.send(r).unwrap());
+        device.poll(wgpu::Maintain::Wait);
+        receiver.recv().unwrap().unwrap();
+        {
+            let view = buffer_slice.get_mapped_range();
+            buffer.copy_from_slice(&view[..]);
+        }
+        download_buffer.unmap();
+        eprintln!("write");
+        stdout().write_all(&buffer).unwrap();
+    }
+}
diff --git a/flowy/src/shader.wgsl b/flowy/src/shader.wgsl
new file mode 100644
index 0000000..1597ff6
--- /dev/null
+++ b/flowy/src/shader.wgsl
@@ -0,0 +1,11 @@
+
+@group(0) @binding(0) var output: texture_storage_2d<bgra8unorm, write>;
+@group(0) @binding(1) var texa: texture_2d<f32>;
+
+
+@compute @workgroup_size(1) fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+    let p = textureLoad(texa, global_id.xy, 0);
+    textureStore(output, global_id.xy, p);
+}
+
+