Trac3r-rust/src/vkprocessor.rs

use vulkano::buffer::{BufferUsage, CpuAccessibleBuffer, DeviceLocalBuffer, ImmutableBuffer, BufferAccess};
use vulkano::command_buffer::AutoCommandBufferBuilder;
use vulkano::descriptor::descriptor_set::PersistentDescriptorSet;
use vulkano::device::{Device, DeviceExtensions, QueuesIter, Queue};
use vulkano::instance::{Instance, InstanceExtensions, PhysicalDevice, QueueFamily};
use vulkano::pipeline::ComputePipeline;
use vulkano::sync::GpuFuture;
use vulkano::sync;
use std::time::SystemTime;
use std::sync::Arc;
use std::ffi::CStr;
use std::path::PathBuf;
use shade_runner as sr;
use image::DynamicImage;

pub struct VkProcessor<'a> {
    instance: Arc<Instance>,
    physical: PhysicalDevice<'a>,
    queue_family: QueueFamily<'a>,
    device: Arc<Device>,
    queues: QueuesIter,
    queue: Arc<Queue>,
    img: Option<DynamicImage>,
    image_buffer: Vec<u8>,
    buffers: Vec::
}

impl VkProcessor {
    pub fn new() -> VkProcessor {
        let instance = Instance::new(None, &InstanceExtensions::none(), None).unwrap();
        let physical = PhysicalDevice::enumerate(&instance).next().unwrap();
        let queue_family = physical.queue_families().find(|&q| q.supports_compute()).unwrap();
        let (device, mut queues) = Device::new(physical,
                                               physical.supported_features(),
                                               &DeviceExtensions::none(),
                                               [(queue_family, 0.5)].iter().cloned()).unwrap();
        VkProcessor {
            instance: instance,
            physical: physical,
            queue_family: queue_family,
            device: device,
            queues: queues,
            queue: queues.next().unwrap(),
            img: Option::None,
            image_buffer: Vec::new(),
            buffers: Vec::new(),
        }
    }

    pub fn compile_kernel(&mut self) {
        let project_root = std::env::current_dir().expect("failed to get root directory");
        let mut compute_path = project_root.clone();
        compute_path.push(PathBuf::from("resources/shaders/simple-homogenize.compute"));

        let shader = sr::load_compute(compute_path).expect("Failed to compile");
        let vulkano_entry = sr::parse_compute(&shader).expect("failed to parse");

        let x = unsafe {
            vulkano::pipeline::shader::ShaderModule::from_words(self.device.clone(), &shader.compute)
        }.unwrap();

        // Compile the shader and add it to a pipeline
        let pipeline = Arc::new({
            unsafe {
                ComputePipeline::new(self.device.clone(), &x.compute_entry_point(
                    CStr::from_bytes_with_nul_unchecked(b"main\0"),
                    vulkano_entry.compute_layout), &(),
                ).unwrap()
            }
        });
    }

    pub fn load_buffers(&mut self) {

        self.img = Option::Some(image::open("resources/images/funky-bird.jpg").unwrap());

        let xy = self.img.dimensions();
        let data_length = xy.0 * xy.1 * 4;
        let pixel_count = self.img.raw_pixels().len();
        println!("Pixel count {}", pixel_count);

        if pixel_count != data_length as usize {
            println!("Creating apha channel...");
            for i in self.img.raw_pixels().iter() {
                if (self.image_buffer.len() + 1) % 4 == 0 {
                    self.image_buffer.push(255);
                }
                self.image_buffer.push(*i);
            }
            self.image_buffer.push(255);
        } else {
            self.image_buffer = self.img.raw_pixels();
        }

        println!("Buffer length {}", self.image_buffer.len());
        println!("Size {:?}", xy);

        println!("Allocating Buffers...");
        {
            // Pull out the image data and place it in a buffer for the kernel to write to and for us to read from
            let write_buffer = {
                let mut buff = image_buffer.iter();
                let data_iter = (0..data_length).map(|n| *(buff.next().unwrap()));
                CpuAccessibleBuffer::from_iter(device.clone(), BufferUsage::all(), data_iter).unwrap()
            };

            // Pull out the image data and place it in a buffer for the kernel to read from
            let read_buffer = {
                let mut buff = image_buffer.iter();
                let data_iter = (0..data_length).map(|n| *(buff.next().unwrap()));
                CpuAccessibleBuffer::from_iter(device.clone(), BufferUsage::all(), data_iter).unwrap()
            };

            // A buffer to hold many i32 values to use as settings
            let settings_buffer = {
                let vec = vec![xy.0, xy.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()
            };
        }

        println!("Done");

        // Create the data descriptor set for our previously created shader pipeline
        let mut set = PersistentDescriptorSet::start(pipeline.clone(), 0)
            .add_buffer(write_buffer.clone()).unwrap()
            .add_buffer(read_buffer.clone()).unwrap()
            .add_buffer(settings_buffer.clone()).unwrap();

        let mut set = Arc::new(set.build().unwrap());
    }

    pub fn run_kernel(&mut self) {
        println!("Running Kernel...");
        // The command buffer I think pretty much serves to define what runs where for how many times
        let command_buffer = AutoCommandBufferBuilder::primary_one_time_submit(device.clone(), queue.family()).unwrap()
            .dispatch([xy.0, xy.1, 1], pipeline.clone(), set.clone(), ()).unwrap()
            .build().unwrap();

        // Create a future for running the command buffer and then just fence it
        let future = sync::now(device.clone())
            .then_execute(queue.clone(), command_buffer).unwrap()
            .then_signal_fence_and_flush().unwrap();

        // I think this is redundant and returns immediately
        future.wait(None).unwrap();
        println!("Done running kernel");
    }

    pub fn read_image() -> Vec<u8> {

        // The buffer is sync'd so we can just read straight from the handle
        let mut data_buffer_content = write_buffer.read().unwrap();

        println!("Reading output");

        let mut image_buffer = Vec::new();

        for y in 0..xy.1 {
            for x in 0..xy.0 {

                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_buffer.push(r);
                image_buffer.push(g);
                image_buffer.push(b);
                image_buffer.push(a);

                img.put_pixel(x, y, image::Rgba([r, g, b, a]))
            }
        }

        image_buffer
    }

    pub fn save_image(&self) {
        println!("Saving output");
        img.save(format!("output/{}.png", SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_secs()));
    }
}