abstract away usage of scratch space for buffer barriers and unique index iterators

Author: teoxoy

wgpu-core/src/indirect_validation/draw.rs

@@ -1,4 +1,7 @@
-use super::CreateIndirectValidationPipelineError;
+use super::{
+    utils::{BufferBarrierScratch, BufferBarriers, UniqueIndexExt as _, UniqueIndexScratch},
+    CreateIndirectValidationPipelineError,
+};
 use crate::{
     device::{queue::TempResource, Device, DeviceError},
     lock::{rank, Mutex},
@@ -256,47 +259,37 @@ impl Draw {
             batch.metadata_buffer_offset = metadata_buffer_offset;
         }
 
-        let mut buffer_barriers = Vec::new();
-        let mut buffer_index_set = bit_set::BitSet::new();
-
-        for index in batches
-            .values()
-            .map(|batch| batch.staging_buffer_index)
-            .filter(|index| buffer_index_set.insert(*index))
-        {
-            let staging_buffer = &staging_buffers[index];
-
-            buffer_barriers.push(hal::BufferBarrier {
-                buffer: staging_buffer.raw(),
-                usage: hal::StateTransition {
-                    from: wgt::BufferUses::MAP_WRITE,
-                    to: wgt::BufferUses::COPY_SRC,
-                },
-            });
-        }
-        buffer_index_set.clear();
-
-        for index in batches
-            .values()
-            .map(|batch| batch.metadata_resource_index)
-            .filter(|index| buffer_index_set.insert(*index))
-        {
-            let metadata_buffer = resources.get_metadata_buffer(index);
-
-            buffer_barriers.push(hal::BufferBarrier {
-                buffer: metadata_buffer,
-                usage: hal::StateTransition {
-                    from: wgt::BufferUses::STORAGE_READ_ONLY,
-                    to: wgt::BufferUses::COPY_DST,
-                },
-            });
-        }
-        buffer_index_set.clear();
-
-        unsafe {
-            encoder.transition_buffers(&buffer_barriers);
-        }
-        buffer_barriers.clear();
+        let buffer_barrier_scratch = &mut BufferBarrierScratch::new();
+        let unique_index_scratch = &mut UniqueIndexScratch::new();
+
+        BufferBarriers::new(buffer_barrier_scratch)
+            .extend(
+                batches
+                    .values()
+                    .map(|batch| batch.staging_buffer_index)
+                    .unique(unique_index_scratch)
+                    .map(|index| hal::BufferBarrier {
+                        buffer: staging_buffers[index].raw(),
+                        usage: hal::StateTransition {
+                            from: wgt::BufferUses::MAP_WRITE,
+                            to: wgt::BufferUses::COPY_SRC,
+                        },
+                    }),
+            )
+            .extend(
+                batches
+                    .values()
+                    .map(|batch| batch.metadata_resource_index)
+                    .unique(unique_index_scratch)
+                    .map(|index| hal::BufferBarrier {
+                        buffer: resources.get_metadata_buffer(index),
+                        usage: hal::StateTransition {
+                            from: wgt::BufferUses::STORAGE_READ_ONLY,
+                            to: wgt::BufferUses::COPY_DST,
+                        },
+                    }),
+            )
+            .encode(encoder);
 
         for batch in batches.values() {
             let data = batch.metadata();
@@ -319,44 +312,38 @@ impl Draw {
             }
         }
 
-        for index in batches
-            .values()
-            .map(|batch| batch.metadata_resource_index)
-            .filter(|index| buffer_index_set.insert(*index))
-        {
-            let metadata_buffer = resources.get_metadata_buffer(index);
-
-            buffer_barriers.push(hal::BufferBarrier {
-                buffer: metadata_buffer,
-                usage: hal::StateTransition {
-                    from: wgt::BufferUses::COPY_DST,
-                    to: wgt::BufferUses::STORAGE_READ_ONLY,
-                },
-            });
-        }
-        buffer_index_set.clear();
-
-        for index in batches
-            .values()
-            .map(|batch| batch.dst_resource_index)
-            .filter(|index| buffer_index_set.insert(*index))
-        {
-            let dst_buffer = resources.get_dst_buffer(index);
-
-            buffer_barriers.push(hal::BufferBarrier {
-                buffer: dst_buffer,
-                usage: hal::StateTransition {
-                    from: wgt::BufferUses::INDIRECT,
-                    to: wgt::BufferUses::STORAGE_READ_WRITE,
-                },
-            });
+        for staging_buffer in staging_buffers {
+            temp_resources.push(TempResource::StagingBuffer(staging_buffer));
         }
-        buffer_index_set.clear();
 
-        unsafe {
-            encoder.transition_buffers(&buffer_barriers);
-        }
-        buffer_barriers.clear();
+        BufferBarriers::new(buffer_barrier_scratch)
+            .extend(
+                batches
+                    .values()
+                    .map(|batch| batch.metadata_resource_index)
+                    .unique(unique_index_scratch)
+                    .map(|index| hal::BufferBarrier {
+                        buffer: resources.get_metadata_buffer(index),
+                        usage: hal::StateTransition {
+                            from: wgt::BufferUses::COPY_DST,
+                            to: wgt::BufferUses::STORAGE_READ_ONLY,
+                        },
+                    }),
+            )
+            .extend(
+                batches
+                    .values()
+                    .map(|batch| batch.dst_resource_index)
+                    .unique(unique_index_scratch)
+                    .map(|index| hal::BufferBarrier {
+                        buffer: resources.get_dst_buffer(index),
+                        usage: hal::StateTransition {
+                            from: wgt::BufferUses::INDIRECT,
+                            to: wgt::BufferUses::STORAGE_READ_WRITE,
+                        },
+                    }),
+            )
+            .encode(encoder);
 
         let desc = hal::ComputePassDescriptor {
             label: None,
@@ -420,31 +407,21 @@ impl Draw {
             encoder.end_compute_pass();
         }
 
-        for index in batches
-            .values()
-            .map(|batch| batch.dst_resource_index)
-            .filter(|index| buffer_index_set.insert(*index))
-        {
-            let dst_buffer = resources.get_dst_buffer(index);
-
-            buffer_barriers.push(hal::BufferBarrier {
-                buffer: dst_buffer,
-                usage: hal::StateTransition {
-                    from: wgt::BufferUses::STORAGE_READ_WRITE,
-                    to: wgt::BufferUses::INDIRECT,
-                },
-            });
-        }
-        buffer_index_set.clear();
-
-        unsafe {
-            encoder.transition_buffers(&buffer_barriers);
-        }
-        buffer_barriers.clear();
-
-        for staging_buffer in staging_buffers {
-            temp_resources.push(TempResource::StagingBuffer(staging_buffer));
-        }
+        BufferBarriers::new(buffer_barrier_scratch)
+            .extend(
+                batches
+                    .values()
+                    .map(|batch| batch.dst_resource_index)
+                    .unique(unique_index_scratch)
+                    .map(|index| hal::BufferBarrier {
+                        buffer: resources.get_dst_buffer(index),
+                        usage: hal::StateTransition {
+                            from: wgt::BufferUses::STORAGE_READ_WRITE,
+                            to: wgt::BufferUses::INDIRECT,
+                        },
+                    }),
+            )
+            .encode(encoder);
 
         Ok(())
     }

wgpu-core/src/indirect_validation/mod.rs

@@ -7,6 +7,7 @@ use thiserror::Error;
 
 mod dispatch;
 mod draw;
+mod utils;
 
 pub(crate) use dispatch::Dispatch;
 pub(crate) use draw::{Draw, DrawBatcher, DrawResources};

wgpu-core/src/indirect_validation/utils.rs

@@ -0,0 +1,84 @@
+use alloc::vec::Vec;
+
+pub(crate) struct UniqueIndexScratch(bit_set::BitSet);
+
+impl UniqueIndexScratch {
+    pub(crate) fn new() -> Self {
+        Self(bit_set::BitSet::new())
+    }
+}
+
+pub(crate) struct UniqueIndex<'a, I: Iterator<Item = usize>> {
+    inner: I,
+    scratch: &'a mut UniqueIndexScratch,
+}
+
+impl<'a, I: Iterator<Item = usize>> UniqueIndex<'a, I> {
+    fn new(inner: I, scratch: &'a mut UniqueIndexScratch) -> Self {
+        scratch.0.clear();
+        Self { inner, scratch }
+    }
+}
+
+impl<'a, I: Iterator<Item = usize>> Iterator for UniqueIndex<'a, I> {
+    type Item = usize;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.inner.find(|&i| self.scratch.0.insert(i))
+    }
+}
+
+pub(crate) trait UniqueIndexExt: Iterator<Item = usize> {
+    fn unique<'a>(self, scratch: &'a mut UniqueIndexScratch) -> UniqueIndex<'a, Self>
+    where
+        Self: Sized,
+    {
+        UniqueIndex::new(self, scratch)
+    }
+}
+
+impl<T: Iterator<Item = usize>> UniqueIndexExt for T {}
+
+type BufferBarrier<'b> = hal::BufferBarrier<'b, dyn hal::DynBuffer>;
+
+pub(crate) struct BufferBarrierScratch<'b>(Vec<BufferBarrier<'b>>);
+
+impl<'b> BufferBarrierScratch<'b> {
+    pub(crate) fn new() -> Self {
+        Self(Vec::new())
+    }
+}
+
+pub(crate) struct BufferBarriers<'a, 'b> {
+    scratch: &'a mut BufferBarrierScratch<'b>,
+}
+
+impl<'a, 'b> BufferBarriers<'a, 'b> {
+    pub(crate) fn new(scratch: &'a mut BufferBarrierScratch<'_>) -> Self {
+        // change lifetime of buffer reference, this is safe since `scratch` is empty,
+        // it was either just created or it has been cleared on `BufferBarriers::drop`
+        let scratch = unsafe {
+            core::mem::transmute::<&'a mut BufferBarrierScratch<'_>, &'a mut BufferBarrierScratch<'b>>(
+                scratch,
+            )
+        };
+        Self { scratch }
+    }
+
+    pub(crate) fn extend(self, iter: impl Iterator<Item = BufferBarrier<'b>>) -> Self {
+        self.scratch.0.extend(iter);
+        self
+    }
+
+    pub(crate) fn encode(self, encoder: &mut dyn hal::DynCommandEncoder) {
+        unsafe {
+            encoder.transition_buffers(&self.scratch.0);
+        }
+    }
+}
+
+impl<'a, 'b> Drop for BufferBarriers<'a, 'b> {
+    fn drop(&mut self) {
+        self.scratch.0.clear();
+    }
+}