Make bind_map persistent

We'll be persisting some buffers across recordings, so make the mapping from id to actual resource scoped to the engine rather than a single `run_recording` call.

Part of the change is being explicit about which buffers to free and when. That will enable more fine-grained reuse of buffers, including within a recording.

src/engine.rs

@@ -16,7 +16,7 @@
 
 use std::{
     borrow::Cow,
-    collections::{hash_map::Entry, HashMap},
+    collections::{hash_map::Entry, HashMap, HashSet},
     num::{NonZeroU32, NonZeroU64},
     sync::atomic::{AtomicU64, Ordering},
 };
@@ -42,6 +42,7 @@ static ID_COUNTER: AtomicU64 = AtomicU64::new(0);
 pub struct Engine {
     shaders: Vec<Shader>,
     pool: ResourcePool,
+    bind_map: BindMap,
 }
 
 struct Shader {
@@ -96,6 +97,8 @@ pub enum Command {
     Dispatch(ShaderId, (u32, u32, u32), Vec<ResourceProxy>),
     Download(BufProxy),
     Clear(BufProxy, u64, Option<NonZeroU64>),
+    FreeBuf(BufProxy),
+    FreeImage(ImageProxy),
 }
 
 #[derive(Default)]
@@ -149,6 +152,7 @@ impl Engine {
         Engine {
             shaders: vec![],
             pool: Default::default(),
+            bind_map: Default::default(),
         }
     }
 
@@ -250,8 +254,9 @@ impl Engine {
         recording: &Recording,
         external_resources: &[ExternalResource],
     ) -> Result<Downloads, Error> {
-        let mut bind_map = BindMap::default();
         let mut downloads = Downloads::default();
+        let mut free_bufs: HashSet<Id> = Default::default();
+        let mut free_images: HashSet<Id> = Default::default();
 
         let mut encoder = device.create_command_encoder(&Default::default());
         for command in &recording.commands {
@@ -263,14 +268,14 @@ impl Engine {
                     // TODO: if buffer is newly created, might be better to make it mapped at creation
                     // and copy. However, we expect reuse will be most common.
                     queue.write_buffer(&buf, 0, bytes);
-                    bind_map.insert_buf(buf_proxy, buf);
+                    self.bind_map.insert_buf(buf_proxy, buf);
                 }
                 Command::UploadUniform(buf_proxy, bytes) => {
                     let usage = BufferUsages::UNIFORM | BufferUsages::COPY_DST;
                     // Same consideration as above
                     let buf = self.pool.get_buf(buf_proxy, usage, device);
                     queue.write_buffer(&buf, 0, bytes);
-                    bind_map.insert_buf(buf_proxy, buf);
+                    self.bind_map.insert_buf(buf_proxy, buf);
                 }
                 Command::UploadImage(image_proxy, bytes) => {
                     let buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
@@ -322,12 +327,13 @@ impl Engine {
                             depth_or_array_layers: 1,
                         },
                     );
-                    bind_map.insert_image(image_proxy.id, texture, texture_view)
+                    self.bind_map
+                        .insert_image(image_proxy.id, texture, texture_view)
                 }
                 Command::Dispatch(shader_id, wg_size, bindings) => {
                     // println!("dispatching {:?} with {} bindings", wg_size, bindings.len());
                     let shader = &self.shaders[shader_id.0];
-                    let bind_group = bind_map.create_bind_group(
+                    let bind_group = self.bind_map.create_bind_group(
                         device,
                         &shader.bind_group_layout,
                         bindings,
@@ -340,7 +346,11 @@ impl Engine {
                     cpass.dispatch_workgroups(wg_size.0, wg_size.1, wg_size.2);
                 }
                 Command::Download(proxy) => {
-                    let src_buf = bind_map.buf_map.get(&proxy.id).ok_or("buffer not in map")?;
+                    let src_buf = self
+                        .bind_map
+                        .buf_map
+                        .get(&proxy.id)
+                        .ok_or("buffer not in map")?;
                     let buf = device.create_buffer(&wgpu::BufferDescriptor {
                         label: Some(proxy.name),
                         size: proxy.size,
@@ -351,13 +361,38 @@ impl Engine {
                     downloads.buf_map.insert(proxy.id, buf);
                 }
                 Command::Clear(proxy, offset, size) => {
-                    let buffer = bind_map.get_or_create(*proxy, device, &mut self.pool)?;
+                    let buffer = self
+                        .bind_map
+                        .get_or_create(*proxy, device, &mut self.pool)?;
                     encoder.clear_buffer(buffer, *offset, *size);
                 }
+                Command::FreeBuf(proxy) => {
+                    free_bufs.insert(proxy.id);
+                }
+                Command::FreeImage(proxy) => {
+                    free_images.insert(proxy.id);
+                }
             }
         }
         queue.submit(Some(encoder.finish()));
-        self.pool.reap_bindmap(bind_map);
+        for id in free_bufs {
+            if let Some(buf) = self.bind_map.buf_map.remove(&id) {
+                let props = BufferProperties {
+                    size: buf.buffer.size(),
+                    usages: buf.buffer.usage(),
+                    #[cfg(feature = "buffer_labels")]
+                    name: buf.label,
+                };
+                self.pool.bufs.entry(props).or_default().push(buf.buffer);
+            }
+        }
+        for id in free_images {
+            if let Some((texture, view)) = self.bind_map.image_map.remove(&id) {
+                // TODO: have a pool to avoid needless re-allocation
+                drop(texture);
+                drop(view);
+            }
+        }
         Ok(downloads)
     }
 }
@@ -413,6 +448,21 @@ impl Recording {
     pub fn clear_all(&mut self, buf: BufProxy) {
         self.push(Command::Clear(buf, 0, None));
     }
+
+    pub fn free_buf(&mut self, buf: BufProxy) {
+        self.push(Command::FreeBuf(buf));
+    }
+
+    pub fn free_image(&mut self, image: ImageProxy) {
+        self.push(Command::FreeImage(image));
+    }
+
+    pub fn free_resource(&mut self, resource: ResourceProxy) {
+        match resource {
+            ResourceProxy::Buf(buf) => self.free_buf(buf),
+            ResourceProxy::Image(image) => self.free_image(image),
+        }
+    }
 }
 
 impl BufProxy {
@@ -692,7 +742,6 @@ impl ResourcePool {
         let props = BufferProperties {
             size: rounded_size,
             usages: usage,
-            #[cfg(feature = "buffer_labels")]
             name: proxy.name,
         };
         if let Some(buf_vec) = self.bufs.get_mut(&props) {
@@ -711,19 +760,6 @@ impl ResourcePool {
         })
     }
 
-    fn reap_bindmap(&mut self, bind_map: BindMap) {
-        for (_id, buf) in bind_map.buf_map {
-            let size = buf.buffer.size();
-            let props = BufferProperties {
-                size,
-                usages: buf.buffer.usage(),
-                #[cfg(feature = "buffer_labels")]
-                name: buf.label,
-            };
-            self.bufs.entry(props).or_default().push(buf.buffer);
-        }
-    }
-
     /// Quantize a size up to the nearest size class.
     fn size_class(x: u64, bits: u32) -> u64 {
         if x > 1 << bits {

src/render.rs

@@ -231,6 +231,7 @@ pub fn render_encoding_full(
         [config_buf, scene_buf, reduced_buf],
     );
     let mut pathtag_parent = reduced_buf;
+    let mut large_pathtag_bufs = None;
     if pathtag_large {
         let reduced2_size = shaders::PATHTAG_REDUCE_WG as usize;
         let reduced2_buf =
@@ -250,6 +251,7 @@ pub fn render_encoding_full(
             [reduced_buf, reduced2_buf, reduced_scan_buf],
         );
         pathtag_parent = reduced_scan_buf;
+        large_pathtag_bufs = Some((reduced2_buf, reduced_scan_buf));
     }
 
     let tagmonoid_buf = ResourceProxy::new_buf(
@@ -266,6 +268,11 @@ pub fn render_encoding_full(
         (pathtag_wgs as u32, 1, 1),
         [config_buf, scene_buf, pathtag_parent, tagmonoid_buf],
     );
+    recording.free_resource(reduced_buf);
+    if let Some((reduced2, reduced_scan)) = large_pathtag_bufs {
+        recording.free_resource(reduced2);
+        recording.free_resource(reduced_scan);
+    }
     let drawobj_wgs = (n_drawobj + shaders::PATH_BBOX_WG - 1) / shaders::PATH_BBOX_WG;
     let path_bbox_buf = ResourceProxy::new_buf(n_paths as u64 * PATH_BBOX_SIZE, "path_bbox_buf");
     recording.dispatch(
@@ -311,6 +318,7 @@ pub fn render_encoding_full(
             clip_inp_buf,
         ],
     );
+    recording.free_resource(draw_reduced_buf);
     let clip_el_buf = ResourceProxy::new_buf(encoding.n_clips as u64 * CLIP_EL_SIZE, "clip_el_buf");
     let clip_bic_buf = ResourceProxy::new_buf(
         (n_clip / shaders::CLIP_REDUCE_WG) as u64 * CLIP_BIC_SIZE,
@@ -347,6 +355,9 @@ pub fn render_encoding_full(
             ],
         );
     }
+    recording.free_resource(clip_inp_buf);
+    recording.free_resource(clip_bic_buf);
+    recording.free_resource(clip_el_buf);
     let draw_bbox_buf = ResourceProxy::new_buf(n_paths as u64 * DRAW_BBOX_SIZE, "draw_bbox_buf");
     let bump_buf = BufProxy::new(BUMP_SIZE, "bump_buf");
     let width_in_bins = (config.width_in_tiles + 15) / 16;
@@ -371,6 +382,9 @@ pub fn render_encoding_full(
             bin_header_buf,
         ],
     );
+    recording.free_resource(draw_monoid_buf);
+    recording.free_resource(path_bbox_buf);
+    recording.free_resource(clip_bbox_buf);
     // Note: this only needs to be rounded up because of the workaround to store the tile_offset
     // in storage rather than workgroup memory.
     let n_path_aligned = align_up(n_paths as usize, 256);
@@ -388,6 +402,7 @@ pub fn render_encoding_full(
             tile_buf,
         ],
     );
+    recording.free_resource(draw_bbox_buf);
     recording.dispatch(
         shaders.path_coarse,
         (path_coarse_wgs, 1, 1),
@@ -402,6 +417,8 @@ pub fn render_encoding_full(
             segments_buf,
         ],
     );
+    recording.free_resource(tagmonoid_buf);
+    recording.free_resource(cubic_buf);
     recording.dispatch(
         shaders.backdrop,
         (path_wgs, 1, 1),
@@ -422,6 +439,12 @@ pub fn render_encoding_full(
             ptcl_buf,
         ],
     );
+    recording.free_resource(scene_buf);
+    recording.free_resource(draw_monoid_buf);
+    recording.free_resource(bin_header_buf);
+    recording.free_resource(path_buf);
+    // TODO: bump_buf is special
+    recording.free_resource(bump_buf);
     let out_image = ImageProxy::new(width, height, ImageFormat::Rgba8);
     recording.dispatch(
         shaders.fine,
@@ -436,6 +459,12 @@ pub fn render_encoding_full(
             info_bin_data_buf,
         ],
     );
+    recording.free_resource(config_buf);
+    recording.free_resource(tile_buf);
+    recording.free_resource(segments_buf);
+    recording.free_resource(ptcl_buf);
+    recording.free_resource(gradient_image);
+    recording.free_resource(info_bin_data_buf);
     (recording, ResourceProxy::Image(out_image))
 }