Continuing the Vulkano book. Working compute shader.

src/main.rs

@@ -2,6 +2,11 @@ use std::sync::Arc;
 
 use vulkano::command_buffer::{AutoCommandBufferBuilder, CommandBufferUsage, CopyBufferInfo, PrimaryAutoCommandBuffer};
 use vulkano::command_buffer::allocator::{StandardCommandBufferAllocator, StandardCommandBufferAllocatorCreateInfo};
+use vulkano::descriptor_set::{PersistentDescriptorSet, WriteDescriptorSet};
+use vulkano::descriptor_set::allocator::StandardDescriptorSetAllocator;
+use vulkano::pipeline::compute::ComputePipelineCreateInfo;
+use vulkano::pipeline::layout::PipelineDescriptorSetLayoutCreateInfo;
+use vulkano::pipeline::{ComputePipeline, Pipeline, PipelineBindPoint, PipelineLayout, PipelineShaderStageCreateInfo};
 use vulkano::sync::GpuFuture;
 use vulkano::{VulkanLibrary, buffer, sync};
 use vulkano::buffer::{Buffer, BufferCreateInfo, BufferUsage};
@@ -69,7 +74,7 @@ fn main() {
         data
     ).expect("Failed to create buffer");*/
 
-    let source_content = 0..64;
+    /*let source_content = 0..64;
     let source = Buffer::from_iter(
         memory_allocator.clone(), 
         BufferCreateInfo { 
@@ -117,6 +122,108 @@ fn main() {
 
     assert_eq!(&*src_content, &*dst_content);
 
-    println!("Success!");
+    println!("Success!");*/
 
+    let data_iter = 0..65536_u32;
+    let data_buffer = Buffer::from_iter(
+        memory_allocator.clone(), 
+        BufferCreateInfo {
+            usage: BufferUsage::STORAGE_BUFFER,
+            ..Default::default()
+        }, 
+        AllocationCreateInfo { 
+            memory_type_filter: MemoryTypeFilter::PREFER_DEVICE|MemoryTypeFilter::HOST_SEQUENTIAL_WRITE, 
+            ..Default::default()
+        }, 
+        data_iter
+    ).expect("Failed to create buffer");
+
+    let shader = cs::load(device.clone()).expect("Failed to create shader module");
+
+    let cs = shader.entry_point("main").unwrap();
+    let stage = PipelineShaderStageCreateInfo::new(cs);
+    let layout = PipelineLayout::new(
+        device.clone(), 
+        PipelineDescriptorSetLayoutCreateInfo::from_stages([&stage])
+        .into_pipeline_layout_create_info(device.clone())
+        .unwrap(),
+    ).unwrap();
+
+    let compute_pipeline = ComputePipeline::new(
+        device.clone(), 
+        None, 
+        ComputePipelineCreateInfo::stage_layout(stage, layout)
+    ).expect("failed to create compute pipeline");
+
+    let descriptor_set_allocator =
+        StandardDescriptorSetAllocator::new(device.clone(), Default::default());
+    let pipeline_layout = compute_pipeline.layout();
+    let descriptor_set_layouts = pipeline_layout.set_layouts();
+    let descriptor_set_layout_index = 0;
+    let descriptor_set_layout = descriptor_set_layouts.get(descriptor_set_layout_index).unwrap();
+    let descriptor_set = PersistentDescriptorSet::new(
+        &descriptor_set_allocator, 
+        descriptor_set_layout.clone(), 
+        [WriteDescriptorSet::buffer(0, data_buffer.clone())], 
+        [],
+    ).unwrap();
+
+    let command_buffer_allocator = StandardCommandBufferAllocator::new(
+        device.clone(), 
+        StandardCommandBufferAllocatorCreateInfo::default(),
+    );
+    let mut command_buffer_builder = AutoCommandBufferBuilder::primary(
+        &command_buffer_allocator, 
+        queue.queue_family_index(), 
+        CommandBufferUsage::OneTimeSubmit
+    ).unwrap();
+
+    let work_group_counts = [1024,1,1];
+
+    command_buffer_builder
+        .bind_pipeline_compute(compute_pipeline.clone())
+        .unwrap()
+        .bind_descriptor_sets(
+            PipelineBindPoint::Compute, 
+            compute_pipeline.layout().clone(), 
+            descriptor_set_layout_index as u32, 
+            descriptor_set
+        ).unwrap()
+        .dispatch(work_group_counts)
+        .unwrap();
+
+    let command_buffer = command_buffer_builder.build().unwrap();
+
+    let future = sync::now(device.clone())
+        .then_execute(queue.clone(), command_buffer).unwrap().then_signal_fence_and_flush().unwrap();
+
+    future.wait(None).unwrap();
+
+    let content = data_buffer.read().unwrap();
+    for (n, val) in content.iter().enumerate() {
+        assert_eq!(*val, n as u32 *12);
+    }
+
+    println!("Compute shader processing succeeded")
 }
+
+
+mod cs {
+    vulkano_shaders::shader!{
+        ty: "compute",
+        src: r"
+            #version 460
+
+            layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
+
+            layout(set = 0, binding = 0) buffer Data{
+                uint data[];
+            } buf;
+
+            void main() {
+                uint idx = gl_GlobalInvocationID.x;
+                buf.data[idx] *= 12;
+            }
+        ",
+    }
+}