1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106

use std::sync::Arc;
use vulkano::device::Device;
use vulkano::buffer::{CpuAccessibleBuffer, BufferUsage};
use vulkano::pipeline::ComputePipeline;
use vulkano::descriptor::pipeline_layout::PipelineLayout;
use vulkano::descriptor::descriptor_set::{PersistentDescriptorSet, PersistentDescriptorSetBuf};
use image::ImageBuffer;
use image::Rgba;

#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
pub struct CompuBufferHandle {
    pub handle: u32,
}


#[derive(Clone)]
pub struct CompuBuffers {
    dimensions: (u32, u32),
    device: Arc<Device>,
    handle: Arc<CompuBufferHandle>,

    io_buffers: Vec<Arc<CpuAccessibleBuffer<[u8]>>>,
    settings_buffer: Arc<CpuAccessibleBuffer<[u32]>>,
}

impl CompuBuffers {
    pub fn new(device: Arc<Device>, data: Vec<u8>,
               dimensions: (u32, u32), stride: u32,
               handle: Arc<CompuBufferHandle>) -> CompuBuffers {

        let data_length = dimensions.0 * dimensions.1 * stride;

        let input_buffer = {
            let mut buff = data.iter();
            let data_iter = (0..data_length).map(|n| *(buff.next().unwrap()));
            CpuAccessibleBuffer::from_iter(device.clone(), BufferUsage::all(), data_iter).unwrap()
        };

        let output_buffer = {
            let mut buff = data.iter();
            let data_iter = (0..data_length).map(|n| *(buff.next().unwrap()));
            CpuAccessibleBuffer::from_iter(device.clone(), BufferUsage::all(), data_iter).unwrap()
        };

        // Settings buffer which holds i32's
        // Compile macros into the kernel eventually to index them
        let settings_buffer = {
            let vec = vec![dimensions.0, dimensions.1];
            let mut buff = vec.iter();
            let data_iter =
                (0..2).map(|n| *(buff.next().unwrap()));
            CpuAccessibleBuffer::from_iter(device.clone(),
                                           BufferUsage::all(),
                                           data_iter).unwrap()
        };

        CompuBuffers {
            dimensions: dimensions,
            device: device.clone(),

            handle: handle,
            io_buffers: vec![input_buffer, output_buffer],
            settings_buffer: settings_buffer,
        }
    }

    pub fn get_size(&self) -> (u32, u32) {
        self.dimensions
    }

    pub fn get_descriptor_set(&self, compute_pipeline: std::sync::Arc<ComputePipeline<PipelineLayout<shade_runner::layouts::ComputeLayout>>>)
                              -> Arc<PersistentDescriptorSet<std::sync::Arc<ComputePipeline<PipelineLayout<shade_runner::layouts::ComputeLayout>>>, ((((),
                       PersistentDescriptorSetBuf<std::sync::Arc<vulkano::buffer::cpu_access::CpuAccessibleBuffer<[u8]>>>),
                      PersistentDescriptorSetBuf<std::sync::Arc<vulkano::buffer::cpu_access::CpuAccessibleBuffer<[u8]>>>),
                     PersistentDescriptorSetBuf<std::sync::Arc<vulkano::buffer::cpu_access::CpuAccessibleBuffer<[u32]>>>)>> {

        Arc::new(PersistentDescriptorSet::start(compute_pipeline.clone(), 0)
            .add_buffer(self.io_buffers.get(0).unwrap().clone()).unwrap()
            .add_buffer(self.io_buffers.get(1).unwrap().clone()).unwrap()
            .add_buffer(self.settings_buffer.clone()).unwrap()
            .build().unwrap())
    }

    pub fn read_output_buffer(&self) -> ImageBuffer<Rgba<u8>, Vec<u8>> {
        let xy = self.get_size();

        self.io_buffers.get(1).unwrap().write().unwrap().map(|x| x);
        let data_buffer_content = self.io_buffers.get(1).unwrap().read().unwrap();
        ImageBuffer::from_fn(xy.0, xy.1, |x, y| {
            let r = data_buffer_content[((xy.0 * y + x) * 4 + 0) as usize] as u8;
            let g = data_buffer_content[((xy.0 * y + x) * 4 + 1) as usize] as u8;
            let b = data_buffer_content[((xy.0 * y + x) * 4 + 2) as usize] as u8;
            let a = data_buffer_content[((xy.0 * y + x) * 4 + 3) as usize] as u8;

            image::Rgba([r, g, b, a])
        })
    }

    pub fn get_input_buffer(&self) -> Arc<CpuAccessibleBuffer<[u8]>> {
        self.io_buffers.get(0).unwrap().clone()
    }

    pub fn get_output_buffer(&self) -> Arc<CpuAccessibleBuffer<[u8]>> {
        self.io_buffers.get(1).unwrap().clone()
    }
}