validate multi draw indirect calls

Author: teoxoy

wgpu-core/src/command/bundle.rs

@@ -1051,15 +1051,19 @@ impl RenderBundle {
                     instance_limit,
                 } => {
                     let (buffer, offset) = if self.device.indirect_draw_validation.is_some() {
-                        indirect_draw_validation_batcher.add(
+                        let (dst_resource_index, offset) = indirect_draw_validation_batcher.add(
                             indirect_draw_validation_resources,
                             &self.device,
-                            buffer.clone(),
+                            buffer,
                             *offset,
                             *indexed,
                             *vertex_or_index_limit,
                             *instance_limit,
-                        )?
+                        )?;
+
+                        let dst_buffer =
+                            indirect_draw_validation_resources.get_dst_buffer(dst_resource_index);
+                        (dst_buffer, offset)
                     } else {
                         (buffer.try_raw(snatch_guard)?, *offset)
                     };

wgpu-core/src/command/render.rs

@@ -2546,46 +2546,131 @@ fn multi_draw_indirect(
         ),
     );
 
-    let (indirect_raw, offset) = if count == 1 && state.device.indirect_draw_validation.is_some() {
+    fn draw(
+        raw_encoder: &mut dyn hal::DynCommandEncoder,
+        indexed: bool,
+        indirect_buffer: &dyn hal::DynBuffer,
+        offset: u64,
+        count: u32,
+    ) {
+        match indexed {
+            false => unsafe {
+                raw_encoder.draw_indirect(indirect_buffer, offset, count);
+            },
+            true => unsafe {
+                raw_encoder.draw_indexed_indirect(indirect_buffer, offset, count);
+            },
+        }
+    }
+
+    if state.device.indirect_draw_validation.is_some() {
         state
             .info
             .usage_scope
             .buffers
             .merge_single(&indirect_buffer, wgt::BufferUses::STORAGE_READ_ONLY)?;
 
-        indirect_draw_validation_batcher.add(
+        struct DrawData {
+            buffer_index: usize,
+            offset: u64,
+            count: u32,
+        }
+
+        struct DrawContext<'a> {
+            raw_encoder: &'a mut dyn hal::DynCommandEncoder,
+            device: &'a Device,
+
+            indirect_draw_validation_resources:
+                &'a mut crate::indirect_draw_validation::IndirectDrawValidationResources,
+            indirect_draw_validation_batcher:
+                &'a mut crate::indirect_draw_validation::IndirectDrawValidationBatcher,
+
+            indirect_buffer: Arc<crate::resource::Buffer>,
+            indexed: bool,
+            vertex_or_index_limit: u64,
+            instance_limit: u64,
+        }
+
+        impl<'a> DrawContext<'a> {
+            fn add(&mut self, offset: u64) -> Result<DrawData, DeviceError> {
+                let (dst_resource_index, dst_offset) = self.indirect_draw_validation_batcher.add(
+                    self.indirect_draw_validation_resources,
+                    self.device,
+                    &self.indirect_buffer,
+                    offset,
+                    self.indexed,
+                    self.vertex_or_index_limit,
+                    self.instance_limit,
+                )?;
+                Ok(DrawData {
+                    buffer_index: dst_resource_index,
+                    offset: dst_offset,
+                    count: 1,
+                })
+            }
+            fn draw(&mut self, draw_data: DrawData) {
+                let dst_buffer = self
+                    .indirect_draw_validation_resources
+                    .get_dst_buffer(draw_data.buffer_index);
+                draw(
+                    self.raw_encoder,
+                    self.indexed,
+                    dst_buffer,
+                    draw_data.offset,
+                    draw_data.count,
+                );
+            }
+        }
+
+        let mut draw_ctx = DrawContext {
+            raw_encoder: state.raw_encoder,
+            device: state.device,
             indirect_draw_validation_resources,
-            state.device,
+            indirect_draw_validation_batcher,
             indirect_buffer,
-            offset,
             indexed,
-            if indexed {
+            vertex_or_index_limit: if indexed {
                 state.index.limit
             } else {
                 state.vertex.limits.vertex_limit
             },
-            state.vertex.limits.instance_limit,
-        )?
+            instance_limit: state.vertex.limits.instance_limit,
+        };
+
+        let mut current_draw_data = draw_ctx.add(offset)?;
+
+        for i in 1..count {
+            let draw_data = draw_ctx.add(offset + stride * i as u64)?;
+
+            if draw_data.buffer_index == current_draw_data.buffer_index {
+                debug_assert_eq!(
+                    draw_data.offset,
+                    current_draw_data.offset + stride * current_draw_data.count as u64
+                );
+                current_draw_data.count += 1;
+            } else {
+                draw_ctx.draw(current_draw_data);
+                current_draw_data = draw_data;
+            }
+        }
+
+        draw_ctx.draw(current_draw_data);
     } else {
         state
             .info
             .usage_scope
             .buffers
             .merge_single(&indirect_buffer, wgt::BufferUses::INDIRECT)?;
 
-        (indirect_buffer.try_raw(state.snatch_guard)?, offset)
+        draw(
+            state.raw_encoder,
+            indexed,
+            indirect_buffer.try_raw(state.snatch_guard)?,
+            offset,
+            count,
+        );
     };
 
-    match indexed {
-        false => unsafe {
-            state.raw_encoder.draw_indirect(indirect_raw, offset, count);
-        },
-        true => unsafe {
-            state
-                .raw_encoder
-                .draw_indexed_indirect(indirect_raw, offset, count);
-        },
-    }
     Ok(())
 }
 

wgpu-core/src/indirect_draw_validation.rs

@@ -751,7 +751,7 @@ impl IndirectDrawValidationResources {
         }
     }
 
-    fn get_dst_buffer(&self, index: usize) -> &dyn hal::DynBuffer {
+    pub(crate) fn get_dst_buffer(&self, index: usize) -> &dyn hal::DynBuffer {
         self.dst_entries.get(index).unwrap().buffer.as_ref()
     }
 
@@ -933,17 +933,18 @@ impl IndirectDrawValidationBatcher {
 
     /// Add an indirect draw to be validated.
     ///
-    /// Returns an indirect buffer and offset to be used for the draw.
+    /// Returns the index of the indirect buffer in `indirect_draw_validation_resources`
+    /// and the offset to be used for the draw.
     pub(crate) fn add<'a>(
         &mut self,
         indirect_draw_validation_resources: &'a mut IndirectDrawValidationResources,
         device: &Device,
-        src_buffer: Arc<crate::resource::Buffer>,
+        src_buffer: &Arc<crate::resource::Buffer>,
         offset: u64,
         indexed: bool,
         vertex_or_index_limit: u64,
         instance_limit: u64,
-    ) -> Result<(&'a dyn hal::DynBuffer, u64), DeviceError> {
+    ) -> Result<(usize, u64), DeviceError> {
         let stride = crate::command::get_stride_of_indirect_args(indexed);
 
         let (dst_resource_index, dst_offset) = indirect_draw_validation_resources
@@ -973,7 +974,7 @@ impl IndirectDrawValidationBatcher {
             }
             hashbrown::hash_map::Entry::Vacant(vacant_entry) => {
                 vacant_entry.insert(DrawIndirectValidationBatch {
-                    src_buffer,
+                    src_buffer: src_buffer.clone(),
                     src_dynamic_offset,
                     dst_resource_index,
                     entries: vec![entry],
@@ -987,8 +988,6 @@ impl IndirectDrawValidationBatcher {
             }
         }
 
-        let dst_buffer = indirect_draw_validation_resources.get_dst_buffer(dst_resource_index);
-
-        Ok((dst_buffer, dst_offset))
+        Ok((dst_resource_index, dst_offset))
     }
 }