Towards wiring up images to k4

This patch passes a dynamically sized array of textures to the fine
rasterizer.

A bunch of the low level Vulkan stuff is done, but only enough of the
shaders and encoders to do minimal testing. We'll want to switch from
storage images to sampled images, track the actual array of textures
during encoding, use that to build the descriptor set (which will need
to be more dynamic), and of course run image elements through the
pipeline.

Progress towards #38

piet-gpu-hal/src/hub.rs

@@ -324,6 +324,12 @@ impl PipelineBuilder {
         self
     }
 
+    /// Add a binding with a variable-size array of textures.
+    pub fn add_textures(mut self, max_textures: u32) -> Self {
+        self.0.add_textures(max_textures);
+        self
+    }
+
     pub unsafe fn create_compute_pipeline(
         self,
         session: &Session,
@@ -349,6 +355,12 @@ impl DescriptorSetBuilder {
         self
     }
 
+    pub fn add_textures<'a>(mut self, images: impl IntoRefs<'a, Image>) -> Self {
+        let vk_images = images.into_refs().map(|i| i.vk_image()).collect::<Vec<_>>();
+        self.0.add_textures(&vk_images);
+        self
+    }
+
     pub unsafe fn build(
         self,
         session: &Session,

piet-gpu-hal/src/lib.rs

@@ -212,11 +212,18 @@ pub trait PipelineBuilder<D: Device> {
     fn add_buffers(&mut self, n_buffers: u32);
     /// Add storage images to the pipeline. Each has its own binding.
     fn add_images(&mut self, n_images: u32);
+    /// Add a binding with a variable-size array of textures.
+    fn add_textures(&mut self, max_textures: u32);
     unsafe fn create_compute_pipeline(self, device: &D, code: &[u8]) -> Result<D::Pipeline, Error>;
 }
 
+/// A builder for descriptor sets with more complex layouts.
+///
+/// Note: the order needs to match the pipeline building, and it also needs to
+/// be buffers, then images, then textures.
 pub trait DescriptorSetBuilder<D: Device> {
     fn add_buffers(&mut self, buffers: &[&D::Buffer]);
     fn add_images(&mut self, images: &[&D::Image]);
+    fn add_textures(&mut self, images: &[&D::Image]);
     unsafe fn build(self, device: &D, pipeline: &D::Pipeline) -> Result<D::DescriptorSet, Error>;
 }

piet-gpu-hal/src/vulkan.rs

@@ -71,6 +71,7 @@ pub struct Pipeline {
     pipeline: vk::Pipeline,
     descriptor_set_layout: vk::DescriptorSetLayout,
     pipeline_layout: vk::PipelineLayout,
+    max_textures: u32,
 }
 
 pub struct DescriptorSet {
@@ -92,11 +93,14 @@ pub struct MemFlags(vk::MemoryPropertyFlags);
 
 pub struct PipelineBuilder {
     bindings: Vec<vk::DescriptorSetLayoutBinding>,
+    binding_flags: Vec<vk::DescriptorBindingFlags>,
+    max_textures: u32,
 }
 
 pub struct DescriptorSetBuilder {
     buffers: Vec<vk::Buffer>,
     images: Vec<vk::ImageView>,
+    textures: Vec<vk::ImageView>,
 }
 
 unsafe extern "system" fn vulkan_debug_callback(
@@ -141,6 +145,8 @@ static EXTS: Lazy<Vec<&'static CStr>> = Lazy::new(|| {
     if cfg!(debug_assertions) {
         exts.push(DebugUtils::name());
     }
+    // We'll need this to do runtime query of descriptor indexing.
+    //exts.push(vk::KhrGetPhysicalDeviceProperties2Fn::name());
     exts
 });
 
@@ -270,13 +276,22 @@ impl VkInstance {
             .queue_family_index(qfi)
             .queue_priorities(&queue_priorities)
             .build()];
+
+        // support for descriptor indexing (maybe should be optional for compatibility)
+        let descriptor_indexing = vk::PhysicalDeviceDescriptorIndexingFeatures::builder()
+            .descriptor_binding_variable_descriptor_count(true)
+            .runtime_descriptor_array(true);
+
         let extensions = match surface {
             Some(_) => vec![khr::Swapchain::name().as_ptr()],
             None => vec![],
         };
+        //extensions.push(vk::KhrMaintenance3Fn::name().as_ptr());
+        //extensions.push(vk::ExtDescriptorIndexingFn::name().as_ptr());
         let create_info = vk::DeviceCreateInfo::builder()
             .queue_create_infos(&queue_create_infos)
             .enabled_extension_names(&extensions)
+            .push_next(&mut descriptor_indexing.build())
             .build();
         let device = self.instance.create_device(pdevice, &create_info, None)?;
 
@@ -541,6 +556,8 @@ impl crate::Device for VkDevice {
     unsafe fn pipeline_builder(&self) -> PipelineBuilder {
         PipelineBuilder {
             bindings: Vec::new(),
+            binding_flags: Vec::new(),
+            max_textures: 0,
         }
     }
 
@@ -548,6 +565,7 @@ impl crate::Device for VkDevice {
         DescriptorSetBuilder {
             buffers: Vec::new(),
             images: Vec::new(),
+            textures: Vec::new(),
         }
     }
 
@@ -922,6 +940,8 @@ impl crate::PipelineBuilder<VkDevice> for PipelineBuilder {
                     .stage_flags(vk::ShaderStageFlags::COMPUTE)
                     .build(),
             );
+            self.binding_flags
+                .push(vk::DescriptorBindingFlags::default());
         }
     }
 
@@ -936,9 +956,27 @@ impl crate::PipelineBuilder<VkDevice> for PipelineBuilder {
                     .stage_flags(vk::ShaderStageFlags::COMPUTE)
                     .build(),
             );
+            self.binding_flags
+                .push(vk::DescriptorBindingFlags::default());
         }
     }
 
+    fn add_textures(&mut self, max_textures: u32) {
+        let start = self.bindings.len() as u32;
+        self.bindings.push(
+            vk::DescriptorSetLayoutBinding::builder()
+                .binding(start)
+                // TODO: we do want these to be sampled images
+                .descriptor_type(vk::DescriptorType::STORAGE_IMAGE)
+                .descriptor_count(max_textures)
+                .stage_flags(vk::ShaderStageFlags::COMPUTE)
+                .build(),
+        );
+        self.binding_flags
+            .push(vk::DescriptorBindingFlags::VARIABLE_DESCRIPTOR_COUNT);
+        self.max_textures += max_textures;
+    }
+
     unsafe fn create_compute_pipeline(
         self,
         device: &VkDevice,
@@ -946,7 +984,14 @@ impl crate::PipelineBuilder<VkDevice> for PipelineBuilder {
     ) -> Result<Pipeline, Error> {
         let device = &device.device.device;
         let descriptor_set_layout = device.create_descriptor_set_layout(
-            &vk::DescriptorSetLayoutCreateInfo::builder().bindings(&self.bindings),
+            &vk::DescriptorSetLayoutCreateInfo::builder()
+                .bindings(&self.bindings)
+                // It might be a slight optimization not to push this if max_textures = 0
+                .push_next(
+                    &mut vk::DescriptorSetLayoutBindingFlagsCreateInfo::builder()
+                        .binding_flags(&self.binding_flags)
+                        .build(),
+                ),
             None,
         )?;
         let descriptor_set_layouts = [descriptor_set_layout];
@@ -981,6 +1026,7 @@ impl crate::PipelineBuilder<VkDevice> for PipelineBuilder {
             pipeline,
             pipeline_layout,
             descriptor_set_layout,
+            max_textures: self.max_textures,
         })
     }
 }
@@ -994,6 +1040,10 @@ impl crate::DescriptorSetBuilder<VkDevice> for DescriptorSetBuilder {
         self.images.extend(images.iter().map(|i| i.image_view));
     }
 
+    fn add_textures(&mut self, images: &[&Image]) {
+        self.textures.extend(images.iter().map(|i| i.image_view));
+    }
+
     unsafe fn build(self, device: &VkDevice, pipeline: &Pipeline) -> Result<DescriptorSet, Error> {
         let device = &device.device.device;
         let mut descriptor_pool_sizes = Vec::new();
@@ -1005,11 +1055,12 @@ impl crate::DescriptorSetBuilder<VkDevice> for DescriptorSetBuilder {
                     .build(),
             );
         }
-        if !self.images.is_empty() {
+        let n_images_total = self.images.len() + pipeline.max_textures as usize;
+        if n_images_total > 0 {
             descriptor_pool_sizes.push(
                 vk::DescriptorPoolSize::builder()
                     .ty(vk::DescriptorType::STORAGE_IMAGE)
-                    .descriptor_count(self.images.len() as u32)
+                    .descriptor_count(n_images_total as u32)
                     .build(),
             );
         }
@@ -1030,39 +1081,60 @@ impl crate::DescriptorSetBuilder<VkDevice> for DescriptorSetBuilder {
         let mut binding = 0;
         // Maybe one call to update_descriptor_sets with an array of descriptor_writes?
         for buf in &self.buffers {
-            let buf_info = vk::DescriptorBufferInfo::builder()
-                .buffer(*buf)
-                .offset(0)
-                .range(vk::WHOLE_SIZE)
-                .build();
             device.update_descriptor_sets(
                 &[vk::WriteDescriptorSet::builder()
                     .dst_set(descriptor_sets[0])
                     .dst_binding(binding)
                     .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
-                    .buffer_info(&[buf_info])
+                    .buffer_info(&[vk::DescriptorBufferInfo::builder()
+                        .buffer(*buf)
+                        .offset(0)
+                        .range(vk::WHOLE_SIZE)
+                        .build()])
                     .build()],
                 &[],
             );
             binding += 1;
         }
         for image in &self.images {
-            let image_info = vk::DescriptorImageInfo::builder()
-                .sampler(vk::Sampler::null())
-                .image_view(*image)
-                .image_layout(vk::ImageLayout::GENERAL)
-                .build();
             device.update_descriptor_sets(
                 &[vk::WriteDescriptorSet::builder()
                     .dst_set(descriptor_sets[0])
                     .dst_binding(binding)
                     .descriptor_type(vk::DescriptorType::STORAGE_IMAGE)
-                    .image_info(&[image_info])
+                    .image_info(&[vk::DescriptorImageInfo::builder()
+                        .sampler(vk::Sampler::null())
+                        .image_view(*image)
+                        .image_layout(vk::ImageLayout::GENERAL)
+                        .build()])
                     .build()],
                 &[],
             );
             binding += 1;
         }
+        if !self.textures.is_empty() {
+            let infos = self
+                .textures
+                .iter()
+                .map(|texture| {
+                    vk::DescriptorImageInfo::builder()
+                        .sampler(vk::Sampler::null())
+                        .image_view(*texture)
+                        .image_layout(vk::ImageLayout::GENERAL)
+                        .build()
+                })
+                .collect::<Vec<_>>();
+            device.update_descriptor_sets(
+                &[vk::WriteDescriptorSet::builder()
+                    .dst_set(descriptor_sets[0])
+                    .dst_binding(binding)
+                    .descriptor_type(vk::DescriptorType::STORAGE_IMAGE)
+                    .image_info(&infos)
+                    .build()],
+                &[],
+            );
+            //binding += 1;
+        }
         Ok(DescriptorSet {
             descriptor_set: descriptor_sets[0],
         })

piet-gpu/shader/kernel4.comp

@@ -6,6 +6,7 @@
 
 #version 450
 #extension GL_GOOGLE_include_directive : enable
+#extension GL_EXT_nonuniform_qualifier : enable
 
 #include "setup.h"
 
@@ -28,6 +29,8 @@ layout(set = 0, binding = 2) buffer ClipScratchBuf {
 
 layout(rgba8, set = 0, binding = 3) uniform writeonly image2D image;
 
+layout(rgba8, set = 0, binding = 4) uniform readonly image2D textures[];
+
 #include "ptcl.h"
 #include "tile.h"
 
@@ -103,6 +106,9 @@ void main() {
     uint clip_tos = 0;
     for (uint i = 0; i < CHUNK; i++) {
         rgb[i] = vec3(0.5);
+        if (xy_uint.x < 256 && xy_uint.y < 256) {
+            rgb[i] = imageLoad(textures[gl_WorkGroupID.x / 16], ivec2(xy_uint.x, xy_uint.y + CHUNK_DY * i)).rgb;
+        }
         mask[i] = 1.0;
     }
 

piet-gpu/src/lib.rs

@@ -199,6 +199,8 @@ pub struct Renderer {
     n_elements: usize,
     n_paths: usize,
     n_pathseg: usize,
+
+    bg_image: hub::Image,
 }
 
 impl Renderer {
@@ -301,16 +303,27 @@ impl Renderer {
             ],
         )?;
 
+        let bg_image = Self::make_test_bg_image(&session);
+
         let k4_code = include_bytes!("../shader/kernel4.spv");
+        // This is an arbitrary limit on the number of textures that can be referenced by
+        // the fine rasterizer. To set it for real, we probably want to pay attention both
+        // to the device limit (maxDescriptorSetSampledImages) but also to the number of
+        // images encoded (I believe there's an cost when allocating descriptor pools). If
+        // it can't be satisfied, then for compatibility we'll probably want to fall back
+        // to an atlasing approach.
+        let max_textures = 256;
         let k4_pipeline = session
             .pipeline_builder()
             .add_buffers(3)
             .add_images(1)
+            .add_textures(max_textures)
             .create_compute_pipeline(&session, k4_code)?;
         let k4_ds = session
             .descriptor_set_builder()
             .add_buffers(&[&ptcl_buf, &tile_buf, &clip_scratch_buf])
             .add_images(&[&image_dev])
+            .add_textures(&[&bg_image])
             .build(&session, &k4_pipeline)?;
 
         Ok(Renderer {
@@ -347,6 +360,7 @@ impl Renderer {
             n_elements,
             n_paths,
             n_pathseg,
+            bg_image,
         })
     }
 
@@ -469,4 +483,22 @@ impl Renderer {
             Ok(image)
         }
     }
+
+    /// Make a test image.
+    fn make_test_bg_image(session: &hub::Session) -> hub::Image {
+        const WIDTH: usize = 256;
+        const HEIGHT: usize = 256;
+        let mut buf = vec![255u8; WIDTH * HEIGHT * 4];
+        for y in 0..HEIGHT {
+            for x in 0..WIDTH {
+                let r = x as u8;
+                let g = y as u8;
+                let b = r ^ g;
+                buf[(y * WIDTH + x) * 4] = r;
+                buf[(y * WIDTH + x) * 4 + 1] = g;
+                buf[(y * WIDTH + x) * 4 + 2] = b;
+            }
+        }
+        Self::make_image(session, WIDTH, HEIGHT, &buf, ImageFormat::RgbaPremul).unwrap()
+    }
 }