monolith

public

Star0

Fork 0

Watch0

88 commits 1 branch 0 tags 0 releases1.5 MiB

Commit `04f0de436841`

tarball

atagen <boss@atagen.co>2026-10-17 04:35 UTC

2 files changed+86-9

render-vk: pool staging buffers instead of allocating per upload

Bare-metal spans showed vkQueueSubmit on the main thread stalling for 1 to 3 ms whenever the import worker staged a large shm copy. Allocating and freeing host-visible memory takes a driver lock on NVIDIA that submission also waits on, and the 14.7 MB staging buffer for a full-output layer came and went every second. Dropped staging buffers now go back to a pool of up to 96 MiB and are reused for uploads of up to twice their size.

Modified file: crates/render-vk/src/device.rs+18 additions@@ -7,10 +7,27 @@ use std::sync::Mutex;
 /// Fields drop in declaration order, so the allocator frees its memory before
 /// the device is destroyed.
 pub(crate) struct Shared {
+    /// Staging buffers kept for reuse, see `Staging`.
+    pub staging_pool: Mutex<Vec<crate::image::Pooled>>,
     pub allocator: Mutex<Allocator>,
     pub device: Device,
 }
 
+impl Drop for Shared {
+    fn drop(&mut self) {
+        let pool = std::mem::take(
+            self.staging_pool
+                .get_mut()
+                .unwrap_or_else(|e| e.into_inner()),
+        );
+        let allocator = self.allocator.get_mut().unwrap_or_else(|e| e.into_inner());
+        for p in pool {
+            unsafe { self.device.device.destroy_buffer(p.buffer, None) };
+            let _ = allocator.free(p.allocation);
+        }
+    }
+}
+
 /// Loaders for the extensions that dmabuf import, export and fence export
 /// need. Absent when the device lacks any of them.
 pub(crate) struct DmabufExt {
@@ -171,6 +188,7 @@ impl Shared {
             spare_fences: Mutex::new(Vec::new()),
         };
         Ok(Self {
+            staging_pool: Mutex::new(Vec::new()),
             allocator: Mutex::new(allocator?),
             device,
         })

Modified file: crates/render-vk/src/image.rs+68 additions−9 deletions@@ -180,13 +180,31 @@ pub(crate) fn color_range() -> vk::ImageSubresourceRange {
         .layer_count(1)
 }
 
-/// A host-visible buffer for uploads and readbacks.
+/// A host-visible buffer for uploads and readbacks. Dropped ones go back to
+/// a pool on `Shared` instead of being freed. Allocating and freeing large
+/// host-visible memory takes a driver lock on NVIDIA that `vkQueueSubmit`
+/// also waits on, so a staging buffer made per upload on the import worker
+/// stalled the main thread's submits for 1 to 3 ms.
 pub(crate) struct Staging {
     pub shared: Arc<Shared>,
     pub buffer: vk::Buffer,
     pub allocation: Option<Allocation>,
+    size: usize,
+    usage: vk::BufferUsageFlags,
+    location: MemoryLocation,
 }
 
+pub(crate) struct Pooled {
+    pub buffer: vk::Buffer,
+    pub allocation: Allocation,
+    size: usize,
+    usage: vk::BufferUsageFlags,
+    location: MemoryLocation,
+}
+
+/// Pooled staging memory kept at most, past which dropped buffers are freed.
+const POOL_BYTES: usize = 96 * 1024 * 1024;
+
 impl Staging {
     pub fn new(
         shared: &Arc<Shared>,
@@ -194,6 +212,26 @@ impl Staging {
         usage: vk::BufferUsageFlags,
         location: MemoryLocation,
     ) -> Result<Self, Error> {
+        let reused = {
+            let mut pool = shared
+                .staging_pool
+                .lock()
+                .unwrap_or_else(|e| e.into_inner());
+            let fits = |p: &Pooled| {
+                p.usage == usage && p.location == location && p.size >= size && p.size <= size * 2
+            };
+            pool.iter().position(fits).map(|i| pool.swap_remove(i))
+        };
+        if let Some(p) = reused {
+            return Ok(Self {
+                shared: shared.clone(),
+                buffer: p.buffer,
+                allocation: Some(p.allocation),
+                size: p.size,
+                usage,
+                location,
+            });
+        }
         let d = &shared.device.device;
         let buffer = unsafe {
             d.create_buffer(
@@ -227,6 +265,9 @@ impl Staging {
             shared: shared.clone(),
             buffer,
             allocation: Some(allocation),
+            size,
+            usage,
+            location,
         })
     }
 
@@ -247,15 +288,33 @@ impl Staging {
 
 impl Drop for Staging {
     fn drop(&mut self) {
-        unsafe { self.shared.device.device.destroy_buffer(self.buffer, None) };
-        if let Some(a) = self.allocation.take() {
-            let _ = self
-                .shared
-                .allocator
-                .lock()
-                .unwrap_or_else(|e| e.into_inner())
-                .free(a);
+        let Some(allocation) = self.allocation.take() else {
+            return;
+        };
+        let mut pool = self
+            .shared
+            .staging_pool
+            .lock()
+            .unwrap_or_else(|e| e.into_inner());
+        let pooled: usize = pool.iter().map(|p| p.size).sum();
+        if pooled + self.size <= POOL_BYTES {
+            pool.push(Pooled {
+                buffer: self.buffer,
+                allocation,
+                size: self.size,
+                usage: self.usage,
+                location: self.location,
+            });
+            return;
         }
+        drop(pool);
+        unsafe { self.shared.device.device.destroy_buffer(self.buffer, None) };
+        let _ = self
+            .shared
+            .allocator
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
+            .free(allocation);
     }
 }

monolith

Commit 04f0de436841

Commit `04f0de436841`