Merge pull request #46 from linebender/image_work

Progress towards images

piet-gpu-hal/examples/collatz.rs

@@ -14,9 +14,9 @@ fn main() {
             .unwrap();
         buffer.write(&src).unwrap();
         let code = include_bytes!("./shader/collatz.spv");
-        let pipeline = session.create_simple_compute_pipeline(code, 1, 0).unwrap();
+        let pipeline = session.create_simple_compute_pipeline(code, 1).unwrap();
         let descriptor_set = session
-            .create_descriptor_set(&pipeline, &[buffer.vk_buffer()], &[])
+            .create_simple_descriptor_set(&pipeline, &[&buffer])
             .unwrap();
         let query_pool = session.create_query_pool(2).unwrap();
         let mut cmd_buf = session.cmd_buf().unwrap();

piet-gpu-hal/src/hub.rs

@@ -8,13 +8,16 @@ use std::any::Any;
 use std::sync::{Arc, Mutex, Weak};
 
 use crate::vulkan;
-use crate::{Device, Error};
+use crate::DescriptorSetBuilder as DescriptorSetBuilderTrait;
+use crate::PipelineBuilder as PipelineBuilderTrait;
+use crate::{Device, Error, SamplerParams};
 
 pub type MemFlags = <vulkan::VkDevice as Device>::MemFlags;
 pub type Semaphore = <vulkan::VkDevice as Device>::Semaphore;
 pub type Pipeline = <vulkan::VkDevice as Device>::Pipeline;
 pub type DescriptorSet = <vulkan::VkDevice as Device>::DescriptorSet;
 pub type QueryPool = <vulkan::VkDevice as Device>::QueryPool;
+pub type Sampler = <vulkan::VkDevice as Device>::Sampler;
 
 type Fence = <vulkan::VkDevice as Device>::Fence;
 
@@ -63,6 +66,10 @@ struct BufferInner {
     session: Weak<SessionInner>,
 }
 
+pub struct PipelineBuilder(vulkan::PipelineBuilder);
+
+pub struct DescriptorSetBuilder(vulkan::DescriptorSetBuilder);
+
 impl Session {
     pub fn new(device: vulkan::VkDevice) -> Session {
         Session(Arc::new(SessionInner {
@@ -158,31 +165,28 @@ impl Session {
         self.0.device.create_semaphore()
     }
 
-    /// This creates a pipeline that runs over the buffer.
+    /// This creates a pipeline that operates on some buffers and images.
     ///
     /// The descriptor set layout is just some number of storage buffers and storage images (this might change).
     pub unsafe fn create_simple_compute_pipeline(
         &self,
         code: &[u8],
         n_buffers: u32,
-        n_images: u32,
     ) -> Result<Pipeline, Error> {
-        self.0
+        self.pipeline_builder()
-            .device
+            .add_buffers(n_buffers)
-            .create_simple_compute_pipeline(code, n_buffers, n_images)
+            .create_compute_pipeline(self, code)
     }
 
-    /// Create a descriptor set for a simple pipeline that just references buffers and images.
+    /// Create a descriptor set for a simple pipeline that just references buffers.
-    ///
+    pub unsafe fn create_simple_descriptor_set<'a>(
-    /// Note: when we do portability, the signature will change to not reference the Vulkan types
-    /// directly.
-    pub unsafe fn create_descriptor_set(
         &self,
         pipeline: &Pipeline,
-        bufs: &[&vulkan::Buffer],
+        buffers: impl IntoRefs<'a, Buffer>,
-        images: &[&vulkan::Image],
     ) -> Result<DescriptorSet, Error> {
-        self.0.device.create_descriptor_set(pipeline, bufs, images)
+        self.descriptor_set_builder()
+            .add_buffers(buffers)
+            .build(self, pipeline)
     }
 
     /// Create a query pool for timestamp queries.
@@ -193,6 +197,18 @@ impl Session {
     pub unsafe fn fetch_query_pool(&self, pool: &QueryPool) -> Result<Vec<f64>, Error> {
         self.0.device.fetch_query_pool(pool)
     }
+
+    pub unsafe fn pipeline_builder(&self) -> PipelineBuilder {
+        PipelineBuilder(self.0.device.pipeline_builder())
+    }
+
+    pub unsafe fn descriptor_set_builder(&self) -> DescriptorSetBuilder {
+        DescriptorSetBuilder(self.0.device.descriptor_set_builder())
+    }
+
+    pub unsafe fn create_sampler(&self, params: SamplerParams) -> Result<Sampler, Error> {
+        self.0.device.create_sampler(params)
+    }
 }
 
 impl CmdBuf {
@@ -299,3 +315,134 @@ impl Buffer {
         Ok(())
     }
 }
+
+impl PipelineBuilder {
+    /// Add buffers to the pipeline. Each has its own binding.
+    pub fn add_buffers(mut self, n_buffers: u32) -> Self {
+        self.0.add_buffers(n_buffers);
+        self
+    }
+
+    /// Add storage images to the pipeline. Each has its own binding.
+    pub fn add_images(mut self, n_images: u32) -> Self {
+        self.0.add_images(n_images);
+        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,
+        code: &[u8],
+    ) -> Result<Pipeline, Error> {
+        self.0.create_compute_pipeline(&session.0.device, code)
+    }
+}
+
+impl DescriptorSetBuilder {
+    pub fn add_buffers<'a>(mut self, buffers: impl IntoRefs<'a, Buffer>) -> Self {
+        let vk_buffers = buffers
+            .into_refs()
+            .map(|b| b.vk_buffer())
+            .collect::<Vec<_>>();
+        self.0.add_buffers(&vk_buffers);
+        self
+    }
+
+    pub fn add_images<'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_images(&vk_images);
+        self
+    }
+
+    pub fn add_textures<'a>(mut self, images: impl IntoRefs<'a, Image>, sampler: &Sampler) -> Self {
+        let vk_images = images.into_refs().map(|i| i.vk_image()).collect::<Vec<_>>();
+        self.0.add_textures(&vk_images, sampler);
+        self
+    }
+
+    pub unsafe fn build(
+        self,
+        session: &Session,
+        pipeline: &Pipeline,
+    ) -> Result<DescriptorSet, Error> {
+        self.0.build(&session.0.device, pipeline)
+    }
+}
+
+// This lets us use either a slice or a vector. The type is clunky but it
+// seems fine enough to use.
+pub trait IntoRefs<'a, T: 'a> {
+    type Iterator: Iterator<Item = &'a T>;
+
+    fn into_refs(self) -> Self::Iterator;
+}
+
+impl<'a, T> IntoRefs<'a, T> for &'a [T] {
+    type Iterator = std::slice::Iter<'a, T>;
+    fn into_refs(self) -> Self::Iterator {
+        self.into_iter()
+    }
+}
+
+impl<'a, T> IntoRefs<'a, T> for &'a [&'a T] {
+    type Iterator = std::iter::Copied<std::slice::Iter<'a, &'a T>>;
+    fn into_refs(self) -> Self::Iterator {
+        self.into_iter().copied()
+    }
+}
+
+// TODO: this will benefit from const generics!
+impl<'a, T> IntoRefs<'a, T> for &'a [&'a T; 1] {
+    type Iterator = std::iter::Copied<std::slice::Iter<'a, &'a T>>;
+    fn into_refs(self) -> Self::Iterator {
+        self.into_iter().copied()
+    }
+}
+
+impl<'a, T> IntoRefs<'a, T> for &'a [&'a T; 2] {
+    type Iterator = std::iter::Copied<std::slice::Iter<'a, &'a T>>;
+    fn into_refs(self) -> Self::Iterator {
+        self.into_iter().copied()
+    }
+}
+
+impl<'a, T> IntoRefs<'a, T> for &'a [&'a T; 3] {
+    type Iterator = std::iter::Copied<std::slice::Iter<'a, &'a T>>;
+    fn into_refs(self) -> Self::Iterator {
+        self.into_iter().copied()
+    }
+}
+
+impl<'a, T> IntoRefs<'a, T> for &'a [&'a T; 4] {
+    type Iterator = std::iter::Copied<std::slice::Iter<'a, &'a T>>;
+    fn into_refs(self) -> Self::Iterator {
+        self.into_iter().copied()
+    }
+}
+
+impl<'a, T> IntoRefs<'a, T> for &'a [&'a T; 5] {
+    type Iterator = std::iter::Copied<std::slice::Iter<'a, &'a T>>;
+    fn into_refs(self) -> Self::Iterator {
+        self.into_iter().copied()
+    }
+}
+
+impl<'a, T> IntoRefs<'a, T> for &'a [&'a T; 6] {
+    type Iterator = std::iter::Copied<std::slice::Iter<'a, &'a T>>;
+    fn into_refs(self) -> Self::Iterator {
+        self.into_iter().copied()
+    }
+}
+
+impl<'a, T> IntoRefs<'a, T> for Vec<&'a T> {
+    type Iterator = std::vec::IntoIter<&'a T>;
+    fn into_refs(self) -> Self::Iterator {
+        self.into_iter()
+    }
+}

piet-gpu-hal/src/lib.rs

@@ -16,10 +16,21 @@ pub enum ImageLayout {
     BlitSrc,
     BlitDst,
     General,
+    ShaderRead,
+}
+
+/// The type of sampling for image lookup.
+///
+/// This could take a lot more params, such as filtering, repeat, behavior
+/// at edges, etc., but for now we'll keep it simple.
+#[derive(Copy, Clone, Debug)]
+pub enum SamplerParams {
+    Nearest,
+    Linear,
 }
 
 pub trait Device: Sized {
-    type Buffer;
+    type Buffer: 'static;
     type Image;
     type MemFlags: MemFlags;
     type Pipeline;
@@ -28,6 +39,9 @@ pub trait Device: Sized {
     type CmdBuf: CmdBuf<Self>;
     type Fence;
     type Semaphore;
+    type PipelineBuilder: PipelineBuilder<Self>;
+    type DescriptorSetBuilder: DescriptorSetBuilder<Self>;
+    type Sampler;
 
     fn create_buffer(&self, size: u64, mem_flags: Self::MemFlags) -> Result<Self::Buffer, Error>;
 
@@ -56,19 +70,48 @@ pub trait Device: Sized {
     /// Maybe doesn't need result return?
     unsafe fn destroy_image(&self, image: &Self::Image) -> Result<(), Error>;
 
+    /// Start building a pipeline.
+    ///
+    /// A pipeline is a bit of shader IR plus a signature for what kinds of resources
+    /// it expects.
+    unsafe fn pipeline_builder(&self) -> Self::PipelineBuilder;
+
+    /// Start building a descriptor set.
+    ///
+    /// A descriptor set is a binding of resources for a given pipeline.
+    unsafe fn descriptor_set_builder(&self) -> Self::DescriptorSetBuilder;
+
+    /// Create a simple compute pipeline that operates on buffers and storage images.
+    ///
+    /// This is provided as a convenience but will probably go away, as the functionality
+    /// is subsumed by the builder.
     unsafe fn create_simple_compute_pipeline(
         &self,
         code: &[u8],
         n_buffers: u32,
         n_images: u32,
-    ) -> Result<Self::Pipeline, Error>;
+    ) -> Result<Self::Pipeline, Error> {
+        let mut builder = self.pipeline_builder();
+        builder.add_buffers(n_buffers);
+        builder.add_images(n_images);
+        builder.create_compute_pipeline(self, code)
+    }
 
+    /// Create a descriptor set for a given pipeline, binding buffers and images.
+    ///
+    /// This is provided as a convenience but will probably go away, as the functionality
+    /// is subsumed by the builder.
     unsafe fn create_descriptor_set(
         &self,
         pipeline: &Self::Pipeline,
         bufs: &[&Self::Buffer],
         images: &[&Self::Image],
-    ) -> Result<Self::DescriptorSet, Error>;
+    ) -> Result<Self::DescriptorSet, Error> {
+        let mut builder = self.descriptor_set_builder();
+        builder.add_buffers(bufs);
+        builder.add_images(images);
+        builder.build(self, pipeline)
+    }
 
     fn create_cmd_buf(&self) -> Result<Self::CmdBuf, Error>;
 
@@ -107,6 +150,8 @@ pub trait Device: Sized {
     unsafe fn create_fence(&self, signaled: bool) -> Result<Self::Fence, Error>;
     unsafe fn wait_and_reset(&self, fences: &[Self::Fence]) -> Result<(), Error>;
     unsafe fn get_fence_status(&self, fence: Self::Fence) -> Result<bool, Error>;
+
+    unsafe fn create_sampler(&self, params: SamplerParams) -> Result<Self::Sampler, Error>;
 }
 
 pub trait CmdBuf<D: Device> {
@@ -174,3 +219,34 @@ pub trait MemFlags: Sized + Clone + Copy {
 
     fn host_coherent() -> Self;
 }
+
+/// A builder for pipelines with more complex layouts.
+pub trait PipelineBuilder<D: Device> {
+    /// Add buffers to the pipeline. Each has its own binding.
+    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]);
+    /// Add an array of storage images.
+    ///
+    /// The images need to be in `ImageLayout::General` layout.
+    fn add_images(&mut self, images: &[&D::Image]);
+    /// Add an array of textures.
+    ///
+    /// The images need to be in `ImageLayout::ShaderRead` layout.
+    ///
+    /// The same sampler is used for all textures, which is not very sophisticated;
+    /// we should have a way to vary the sampler.
+    fn add_textures(&mut self, images: &[&D::Image], sampler: &D::Sampler);
+    unsafe fn build(self, device: &D, pipeline: &D::Pipeline) -> Result<D::DescriptorSet, Error>;
+}

piet-gpu-hal/src/vulkan.rs

@@ -9,7 +9,7 @@ use ash::version::{DeviceV1_0, EntryV1_0, InstanceV1_0};
 use ash::{vk, Device, Entry, Instance};
 use once_cell::sync::Lazy;
 
-use crate::{Device as DeviceTrait, Error, ImageLayout};
+use crate::{Device as DeviceTrait, Error, ImageLayout, SamplerParams};
 
 pub struct VkInstance {
     /// Retain the dynamic lib.
@@ -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 {
@@ -90,6 +91,19 @@ pub struct QueryPool {
 #[derive(Clone, Copy)]
 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>,
+    sampler: vk::Sampler,
+}
+
 unsafe extern "system" fn vulkan_debug_callback(
     message_severity: vk::DebugUtilsMessageSeverityFlagsEXT,
     message_type: vk::DebugUtilsMessageTypeFlagsEXT,
@@ -132,6 +146,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
 });
 
@@ -261,13 +277,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)?;
 
@@ -375,6 +400,9 @@ impl crate::Device for VkDevice {
     type MemFlags = MemFlags;
     type Fence = vk::Fence;
     type Semaphore = vk::Semaphore;
+    type PipelineBuilder = PipelineBuilder;
+    type DescriptorSetBuilder = DescriptorSetBuilder;
+    type Sampler = vk::Sampler;
 
     fn create_buffer(&self, size: u64, mem_flags: MemFlags) -> Result<Buffer, Error> {
         unsafe {
@@ -435,7 +463,8 @@ impl crate::Device for VkDevice {
         // want to add sampling for images and so on.
         let usage = vk::ImageUsageFlags::STORAGE
             | vk::ImageUsageFlags::TRANSFER_SRC
-            | vk::ImageUsageFlags::TRANSFER_DST;
+            | vk::ImageUsageFlags::TRANSFER_DST
+            | vk::ImageUsageFlags::SAMPLED;
         let image = device.create_image(
             &vk::ImageCreateInfo::builder()
                 .image_type(vk::ImageType::TYPE_2D)
@@ -527,151 +556,21 @@ impl crate::Device for VkDevice {
         Ok(device.get_fence_status(fence)?)
     }
 
-    /// This creates a pipeline that runs over the buffer.
+    unsafe fn pipeline_builder(&self) -> PipelineBuilder {
-    ///
+        PipelineBuilder {
-    /// The descriptor set layout is just some number of storage buffers and storage images (this might change).
+            bindings: Vec::new(),
-    unsafe fn create_simple_compute_pipeline(
+            binding_flags: Vec::new(),
-        &self,
+            max_textures: 0,
-        code: &[u8],
-        n_buffers: u32,
-        n_images: u32,
-    ) -> Result<Pipeline, Error> {
-        let device = &self.device.device;
-        let mut bindings = Vec::new();
-        for i in 0..n_buffers {
-            bindings.push(
-                vk::DescriptorSetLayoutBinding::builder()
-                    .binding(i)
-                    .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
-                    .descriptor_count(1)
-                    .stage_flags(vk::ShaderStageFlags::COMPUTE)
-                    .build(),
-            );
         }
-        for i in n_buffers..n_buffers + n_images {
-            bindings.push(
-                vk::DescriptorSetLayoutBinding::builder()
-                    .binding(i)
-                    .descriptor_type(vk::DescriptorType::STORAGE_IMAGE)
-                    .descriptor_count(1)
-                    .stage_flags(vk::ShaderStageFlags::COMPUTE)
-                    .build(),
-            );
-        }
-        let descriptor_set_layout = device.create_descriptor_set_layout(
-            &vk::DescriptorSetLayoutCreateInfo::builder().bindings(&bindings),
-            None,
-        )?;
-
-        let descriptor_set_layouts = [descriptor_set_layout];
-
-        // Create compute pipeline.
-        let code_u32 = convert_u32_vec(code);
-        let compute_shader_module = device
-            .create_shader_module(&vk::ShaderModuleCreateInfo::builder().code(&code_u32), None)?;
-        let entry_name = CString::new("main").unwrap();
-        let pipeline_layout = device.create_pipeline_layout(
-            &vk::PipelineLayoutCreateInfo::builder().set_layouts(&descriptor_set_layouts),
-            None,
-        )?;
-
-        let pipeline = device
-            .create_compute_pipelines(
-                vk::PipelineCache::null(),
-                &[vk::ComputePipelineCreateInfo::builder()
-                    .stage(
-                        vk::PipelineShaderStageCreateInfo::builder()
-                            .stage(vk::ShaderStageFlags::COMPUTE)
-                            .module(compute_shader_module)
-                            .name(&entry_name)
-                            .build(),
-                    )
-                    .layout(pipeline_layout)
-                    .build()],
-                None,
-            )
-            .map_err(|(_pipeline, err)| err)?[0];
-        Ok(Pipeline {
-            pipeline,
-            pipeline_layout,
-            descriptor_set_layout,
-        })
     }
 
-    unsafe fn create_descriptor_set(
+    unsafe fn descriptor_set_builder(&self) -> DescriptorSetBuilder {
-        &self,
+        DescriptorSetBuilder {
-        pipeline: &Pipeline,
+            buffers: Vec::new(),
-        bufs: &[&Buffer],
+            images: Vec::new(),
-        images: &[&Image],
+            textures: Vec::new(),
-    ) -> Result<DescriptorSet, Error> {
+            sampler: vk::Sampler::null(),
-        let device = &self.device.device;
-        let mut descriptor_pool_sizes = Vec::new();
-        if !bufs.is_empty() {
-            descriptor_pool_sizes.push(
-                vk::DescriptorPoolSize::builder()
-                    .ty(vk::DescriptorType::STORAGE_BUFFER)
-                    .descriptor_count(bufs.len() as u32)
-                    .build(),
-            );
         }
-        if !images.is_empty() {
-            descriptor_pool_sizes.push(
-                vk::DescriptorPoolSize::builder()
-                    .ty(vk::DescriptorType::STORAGE_IMAGE)
-                    .descriptor_count(images.len() as u32)
-                    .build(),
-            );
-        }
-        let descriptor_pool = device.create_descriptor_pool(
-            &vk::DescriptorPoolCreateInfo::builder()
-                .pool_sizes(&descriptor_pool_sizes)
-                .max_sets(1),
-            None,
-        )?;
-        let descriptor_set_layouts = [pipeline.descriptor_set_layout];
-        let descriptor_sets = device
-            .allocate_descriptor_sets(
-                &vk::DescriptorSetAllocateInfo::builder()
-                    .descriptor_pool(descriptor_pool)
-                    .set_layouts(&descriptor_set_layouts),
-            )
-            .unwrap();
-        for (i, buf) in bufs.iter().enumerate() {
-            let buf_info = vk::DescriptorBufferInfo::builder()
-                .buffer(buf.buffer)
-                .offset(0)
-                .range(vk::WHOLE_SIZE)
-                .build();
-            device.update_descriptor_sets(
-                &[vk::WriteDescriptorSet::builder()
-                    .dst_set(descriptor_sets[0])
-                    .dst_binding(i as u32)
-                    .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
-                    .buffer_info(&[buf_info])
-                    .build()],
-                &[],
-            );
-        }
-        for (i, image) in images.iter().enumerate() {
-            let binding = i + bufs.len();
-            let image_info = vk::DescriptorImageInfo::builder()
-                .sampler(vk::Sampler::null())
-                .image_view(image.image_view)
-                .image_layout(vk::ImageLayout::GENERAL)
-                .build();
-            device.update_descriptor_sets(
-                &[vk::WriteDescriptorSet::builder()
-                    .dst_set(descriptor_sets[0])
-                    .dst_binding(binding as u32)
-                    .descriptor_type(vk::DescriptorType::STORAGE_IMAGE)
-                    .image_info(&[image_info])
-                    .build()],
-                &[],
-            );
-        }
-        Ok(DescriptorSet {
-            descriptor_set: descriptor_sets[0],
-        })
     }
 
     fn create_cmd_buf(&self) -> Result<CmdBuf, Error> {
@@ -797,6 +696,30 @@ impl crate::Device for VkDevice {
         device.unmap_memory(buffer.buffer_memory);
         Ok(())
     }
+
+    unsafe fn create_sampler(&self, params: SamplerParams) -> Result<Self::Sampler, Error> {
+        let device = &self.device.device;
+        let filter = match params {
+            SamplerParams::Linear => vk::Filter::LINEAR,
+            SamplerParams::Nearest => vk::Filter::NEAREST,
+        };
+        let sampler = device.create_sampler(&vk::SamplerCreateInfo::builder()
+            .mag_filter(filter)
+            .min_filter(filter)
+            .mipmap_mode(vk::SamplerMipmapMode::LINEAR)
+            .address_mode_u(vk::SamplerAddressMode::CLAMP_TO_BORDER)
+            .address_mode_v(vk::SamplerAddressMode::CLAMP_TO_BORDER)
+            .address_mode_w(vk::SamplerAddressMode::CLAMP_TO_BORDER)
+            .mip_lod_bias(0.0)
+            .compare_op(vk::CompareOp::NEVER)
+            .min_lod(0.0)
+            .max_lod(0.0)
+            .border_color(vk::BorderColor::FLOAT_TRANSPARENT_BLACK)
+            .max_anisotropy(1.0)
+            .anisotropy_enable(false)
+        , None)?;
+        Ok(sampler)
+    }
 }
 
 impl crate::CmdBuf<VkDevice> for CmdBuf {
@@ -1033,6 +956,227 @@ impl crate::MemFlags for MemFlags {
     }
 }
 
+impl crate::PipelineBuilder<VkDevice> for PipelineBuilder {
+    fn add_buffers(&mut self, n_buffers: u32) {
+        let start = self.bindings.len() as u32;
+        for i in 0..n_buffers {
+            self.bindings.push(
+                vk::DescriptorSetLayoutBinding::builder()
+                    .binding(start + i)
+                    .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
+                    .descriptor_count(1)
+                    .stage_flags(vk::ShaderStageFlags::COMPUTE)
+                    .build(),
+            );
+            self.binding_flags
+                .push(vk::DescriptorBindingFlags::default());
+        }
+    }
+
+    fn add_images(&mut self, n_images: u32) {
+        let start = self.bindings.len() as u32;
+        for i in 0..n_images {
+            self.bindings.push(
+                vk::DescriptorSetLayoutBinding::builder()
+                    .binding(start + i)
+                    .descriptor_type(vk::DescriptorType::STORAGE_IMAGE)
+                    .descriptor_count(1)
+                    .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::COMBINED_IMAGE_SAMPLER)
+                .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,
+        code: &[u8],
+    ) -> Result<Pipeline, Error> {
+        let device = &device.device.device;
+        let descriptor_set_layout = device.create_descriptor_set_layout(
+            &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];
+
+        // Create compute pipeline.
+        let code_u32 = convert_u32_vec(code);
+        let compute_shader_module = device
+            .create_shader_module(&vk::ShaderModuleCreateInfo::builder().code(&code_u32), None)?;
+        let entry_name = CString::new("main").unwrap();
+        let pipeline_layout = device.create_pipeline_layout(
+            &vk::PipelineLayoutCreateInfo::builder().set_layouts(&descriptor_set_layouts),
+            None,
+        )?;
+
+        let pipeline = device
+            .create_compute_pipelines(
+                vk::PipelineCache::null(),
+                &[vk::ComputePipelineCreateInfo::builder()
+                    .stage(
+                        vk::PipelineShaderStageCreateInfo::builder()
+                            .stage(vk::ShaderStageFlags::COMPUTE)
+                            .module(compute_shader_module)
+                            .name(&entry_name)
+                            .build(),
+                    )
+                    .layout(pipeline_layout)
+                    .build()],
+                None,
+            )
+            .map_err(|(_pipeline, err)| err)?[0];
+        Ok(Pipeline {
+            pipeline,
+            pipeline_layout,
+            descriptor_set_layout,
+            max_textures: self.max_textures,
+        })
+    }
+}
+
+impl crate::DescriptorSetBuilder<VkDevice> for DescriptorSetBuilder {
+    fn add_buffers(&mut self, buffers: &[&Buffer]) {
+        self.buffers.extend(buffers.iter().map(|b| b.buffer));
+    }
+
+    fn add_images(&mut self, images: &[&Image]) {
+        self.images.extend(images.iter().map(|i| i.image_view));
+    }
+
+    fn add_textures(&mut self, images: &[&Image], sampler: &vk::Sampler) {
+        self.textures.extend(images.iter().map(|i| i.image_view));
+        self.sampler = *sampler;
+    }
+
+    unsafe fn build(self, device: &VkDevice, pipeline: &Pipeline) -> Result<DescriptorSet, Error> {
+        let device = &device.device.device;
+        let mut descriptor_pool_sizes = Vec::new();
+        if !self.buffers.is_empty() {
+            descriptor_pool_sizes.push(
+                vk::DescriptorPoolSize::builder()
+                    .ty(vk::DescriptorType::STORAGE_BUFFER)
+                    .descriptor_count(self.buffers.len() as u32)
+                    .build(),
+            );
+        }
+        if !self.images.is_empty() {
+            descriptor_pool_sizes.push(
+                vk::DescriptorPoolSize::builder()
+                    .ty(vk::DescriptorType::STORAGE_IMAGE)
+                    .descriptor_count(self.images.len() as u32)
+                    .build(),
+            );
+        }
+        if pipeline.max_textures > 0 {
+            descriptor_pool_sizes.push(
+                vk::DescriptorPoolSize::builder()
+                    .ty(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)
+                    .descriptor_count(pipeline.max_textures)
+                    .build(),
+            );
+        }
+        let descriptor_pool = device.create_descriptor_pool(
+            &vk::DescriptorPoolCreateInfo::builder()
+                .pool_sizes(&descriptor_pool_sizes)
+                .max_sets(1),
+            None,
+        )?;
+        let descriptor_set_layouts = [pipeline.descriptor_set_layout];
+        let descriptor_sets = device
+            .allocate_descriptor_sets(
+                &vk::DescriptorSetAllocateInfo::builder()
+                    .descriptor_pool(descriptor_pool)
+                    .set_layouts(&descriptor_set_layouts),
+            )
+            .unwrap();
+        let mut binding = 0;
+        // Maybe one call to update_descriptor_sets with an array of descriptor_writes?
+        for buf in &self.buffers {
+            device.update_descriptor_sets(
+                &[vk::WriteDescriptorSet::builder()
+                    .dst_set(descriptor_sets[0])
+                    .dst_binding(binding)
+                    .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
+                    .buffer_info(&[vk::DescriptorBufferInfo::builder()
+                        .buffer(*buf)
+                        .offset(0)
+                        .range(vk::WHOLE_SIZE)
+                        .build()])
+                    .build()],
+                &[],
+            );
+            binding += 1;
+        }
+        for image in &self.images {
+            device.update_descriptor_sets(
+                &[vk::WriteDescriptorSet::builder()
+                    .dst_set(descriptor_sets[0])
+                    .dst_binding(binding)
+                    .descriptor_type(vk::DescriptorType::STORAGE_IMAGE)
+                    .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(self.sampler)
+                        .image_view(*texture)
+                        .image_layout(vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL)
+                        .build()
+                })
+                .collect::<Vec<_>>();
+            device.update_descriptor_sets(
+                &[vk::WriteDescriptorSet::builder()
+                    .dst_set(descriptor_sets[0])
+                    .dst_binding(binding)
+                    .descriptor_type(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)
+                    .image_info(&infos)
+                    .build()],
+                &[],
+            );
+            //binding += 1;
+        }
+        Ok(DescriptorSet {
+            descriptor_set: descriptor_sets[0],
+        })
+    }
+}
+
 impl VkSwapchain {
     pub unsafe fn next(&mut self) -> Result<(usize, vk::Semaphore), Error> {
         let acquisition_semaphore = self.acquisition_semaphores[self.acquisition_idx];
@@ -1137,5 +1281,6 @@ fn map_image_layout(layout: ImageLayout) -> vk::ImageLayout {
         ImageLayout::BlitSrc => vk::ImageLayout::TRANSFER_SRC_OPTIMAL,
         ImageLayout::BlitDst => vk::ImageLayout::TRANSFER_DST_OPTIMAL,
         ImageLayout::General => vk::ImageLayout::GENERAL,
+        ImageLayout::ShaderRead => vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL,
     }
 }

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(set = 0, binding = 4) uniform sampler2D 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 < 1024 && xy_uint.y < 1024) {
+            rgb[i] = texture(textures[gl_WorkGroupID.x / 64], vec2(xy_uint.x, xy_uint.y + CHUNK_DY * i) / 1024.0).rgb;
+        }
         mask[i] = 1.0;
     }
 

piet-gpu/src/lib.rs

@@ -12,7 +12,7 @@ use piet::{Color, ImageFormat, RenderContext};
 
 use piet_gpu_types::encoder::Encode;
 
-use piet_gpu_hal::hub;
+use piet_gpu_hal::{SamplerParams, hub};
 use piet_gpu_hal::{CmdBuf, Error, ImageLayout, MemFlags};
 
 use pico_svg::PicoSvg;
@@ -201,6 +201,8 @@ pub struct Renderer {
     n_elements: usize,
     n_paths: usize,
     n_pathseg: usize,
+
+    bg_image: hub::Image,
 }
 
 impl Renderer {
@@ -236,16 +238,10 @@ impl Renderer {
         let image_dev = session.create_image2d(WIDTH as u32, HEIGHT as u32, dev)?;
 
         let el_code = include_bytes!("../shader/elements.spv");
-        let el_pipeline = session.create_simple_compute_pipeline(el_code, 4, 0)?;
+        let el_pipeline = session.create_simple_compute_pipeline(el_code, 4)?;
-        let el_ds = session.create_descriptor_set(
+        let el_ds = session.create_simple_descriptor_set(
             &el_pipeline,
-            &[
+            &[&scene_dev, &state_buf, &anno_buf, &pathseg_buf],
-                scene_dev.vk_buffer(),
-                state_buf.vk_buffer(),
-                anno_buf.vk_buffer(),
-                pathseg_buf.vk_buffer(),
-            ],
-            &[],
         )?;
 
         let mut tile_alloc_buf_host = session.create_buffer(12, host)?;
@@ -256,40 +252,24 @@ impl Renderer {
         let tile_alloc_start = ((n_paths + 31) & !31) * PATH_SIZE;
         tile_alloc_buf_host.write(&[n_paths as u32, n_pathseg as u32, tile_alloc_start as u32])?;
         let tile_alloc_code = include_bytes!("../shader/tile_alloc.spv");
-        let tile_pipeline = session.create_simple_compute_pipeline(tile_alloc_code, 3, 0)?;
+        let tile_pipeline = session.create_simple_compute_pipeline(tile_alloc_code, 3)?;
-        let tile_ds = session.create_descriptor_set(
+        let tile_ds = session.create_simple_descriptor_set(
             &tile_pipeline,
-            &[
+            &[&anno_buf, &tile_alloc_buf_dev, &tile_buf],
-                anno_buf.vk_buffer(),
-                tile_alloc_buf_dev.vk_buffer(),
-                tile_buf.vk_buffer(),
-            ],
-            &[],
         )?;
 
         let path_alloc_code = include_bytes!("../shader/path_coarse.spv");
-        let path_pipeline = session.create_simple_compute_pipeline(path_alloc_code, 3, 0)?;
+        let path_pipeline = session.create_simple_compute_pipeline(path_alloc_code, 3)?;
-        let path_ds = session.create_descriptor_set(
+        let path_ds = session.create_simple_descriptor_set(
             &path_pipeline,
-            &[
+            &[&pathseg_buf, &tile_alloc_buf_dev, &tile_buf],
-                pathseg_buf.vk_buffer(),
-                tile_alloc_buf_dev.vk_buffer(),
-                tile_buf.vk_buffer(),
-            ],
-            &[],
         )?;
 
         let backdrop_alloc_code = include_bytes!("../shader/backdrop.spv");
-        let backdrop_pipeline =
+        let backdrop_pipeline = session.create_simple_compute_pipeline(backdrop_alloc_code, 3)?;
-            session.create_simple_compute_pipeline(backdrop_alloc_code, 3, 0)?;
+        let backdrop_ds = session.create_simple_descriptor_set(
-        let backdrop_ds = session.create_descriptor_set(
             &backdrop_pipeline,
-            &[
+            &[&anno_buf, &tile_alloc_buf_dev, &tile_buf],
-                anno_buf.vk_buffer(),
-                tile_alloc_buf_dev.vk_buffer(),
-                tile_buf.vk_buffer(),
-            ],
-            &[],
         )?;
 
         let mut bin_alloc_buf_host = session.create_buffer(8, host)?;
@@ -299,15 +279,10 @@ impl Renderer {
         let bin_alloc_start = ((n_paths + 255) & !255) * 8;
         bin_alloc_buf_host.write(&[n_paths as u32, bin_alloc_start as u32])?;
         let bin_code = include_bytes!("../shader/binning.spv");
-        let bin_pipeline = session.create_simple_compute_pipeline(bin_code, 3, 0)?;
+        let bin_pipeline = session.create_simple_compute_pipeline(bin_code, 3)?;
-        let bin_ds = session.create_descriptor_set(
+        let bin_ds = session.create_simple_descriptor_set(
             &bin_pipeline,
-            &[
+            &[&anno_buf, &bin_alloc_buf_dev, &bin_buf],
-                anno_buf.vk_buffer(),
-                bin_alloc_buf_dev.vk_buffer(),
-                bin_buf.vk_buffer(),
-            ],
-            &[],
         )?;
 
         let clip_scratch_buf = session.create_buffer(1024 * 1024, dev)?;
@@ -318,30 +293,41 @@ impl Renderer {
         let coarse_alloc_start = WIDTH_IN_TILES * HEIGHT_IN_TILES * PTCL_INITIAL_ALLOC;
         coarse_alloc_buf_host.write(&[n_paths as u32, coarse_alloc_start as u32])?;
         let coarse_code = include_bytes!("../shader/coarse.spv");
-        let coarse_pipeline = session.create_simple_compute_pipeline(coarse_code, 5, 0)?;
+        let coarse_pipeline = session.create_simple_compute_pipeline(coarse_code, 5)?;
-        let coarse_ds = session.create_descriptor_set(
+        let coarse_ds = session.create_simple_descriptor_set(
             &coarse_pipeline,
             &[
-                anno_buf.vk_buffer(),
+                &anno_buf,
-                bin_buf.vk_buffer(),
+                &bin_buf,
-                tile_buf.vk_buffer(),
+                &tile_buf,
-                coarse_alloc_buf_dev.vk_buffer(),
+                &coarse_alloc_buf_dev,
-                ptcl_buf.vk_buffer(),
+                &ptcl_buf,
             ],
-            &[],
         )?;
 
+        let bg_image = Self::make_test_bg_image(&session);
+
         let k4_code = include_bytes!("../shader/kernel4.spv");
-        let k4_pipeline = session.create_simple_compute_pipeline(k4_code, 3, 1)?;
+        // This is an arbitrary limit on the number of textures that can be referenced by
-        let k4_ds = session.create_descriptor_set(
+        // the fine rasterizer. To set it for real, we probably want to pay attention both
-            &k4_pipeline,
+        // to the device limit (maxDescriptorSetSampledImages) but also to the number of
-            &[
+        // images encoded (I believe there's an cost when allocating descriptor pools). If
-                ptcl_buf.vk_buffer(),
+        // it can't be satisfied, then for compatibility we'll probably want to fall back
-                tile_buf.vk_buffer(),
+        // to an atlasing approach.
-                clip_scratch_buf.vk_buffer(),
+        let max_textures = 256;
-            ],
+        let sampler = session.create_sampler(SamplerParams::Linear)?;
-            &[image_dev.vk_image()],
+        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], &sampler)
+            .build(&session, &k4_pipeline)?;
 
         Ok(Renderer {
             scene_buf,
@@ -377,6 +363,7 @@ impl Renderer {
             n_elements,
             n_paths,
             n_pathseg,
+            bg_image,
         })
     }
 
@@ -488,8 +475,7 @@ impl Renderer {
                 ImageLayout::BlitDst,
             );
             cmd_buf.copy_buffer_to_image(buffer.vk_buffer(), image.vk_image());
-            // TODO: instead of General, we might want ShaderReadOnly
+            cmd_buf.image_barrier(image.vk_image(), ImageLayout::BlitDst, ImageLayout::ShaderRead);
-            cmd_buf.image_barrier(image.vk_image(), ImageLayout::BlitDst, ImageLayout::General);
             cmd_buf.finish();
             // Make sure not to drop the buffer and image until the command buffer completes.
             cmd_buf.add_resource(&buffer);
@@ -499,4 +485,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()
+    }
 }