Implement Queue::write_buffer_with (#2777)

* implement Queue::write_buffer_with

* address comments

* update doc

* Fix copy span location

Co-authored-by: Connor Fitzgerald <connorwadefitzgerald@gmail.com>

wgpu-core/src/hub.rs

@@ -5,7 +5,7 @@ use crate::{
  id,
  instance::{Adapter, HalSurface, Instance, Surface},
  pipeline::{ComputePipeline, RenderPipeline, ShaderModule},
- resource::{Buffer, QuerySet, Sampler, Texture, TextureClearMode, TextureView},
+ resource::{Buffer, QuerySet, Sampler, StagingBuffer, Texture, TextureClearMode, TextureView},
  Epoch, Index,
 };
 
@@ -351,6 +351,7 @@ impl<A: HalApi> Access<Buffer<A>> for CommandBuffer<A> {}
 impl<A: HalApi> Access<Buffer<A>> for ComputePipeline<A> {}
 impl<A: HalApi> Access<Buffer<A>> for RenderPipeline<A> {}
 impl<A: HalApi> Access<Buffer<A>> for QuerySet<A> {}
+impl<A: HalApi> Access<StagingBuffer<A>> for Device<A> {}
 impl<A: HalApi> Access<Texture<A>> for Root {}
 impl<A: HalApi> Access<Texture<A>> for Device<A> {}
 impl<A: HalApi> Access<Texture<A>> for Buffer<A> {}
@@ -452,6 +453,7 @@ pub trait GlobalIdentityHandlerFactory:
  + IdentityHandlerFactory<id::ComputePipelineId>
  + IdentityHandlerFactory<id::QuerySetId>
  + IdentityHandlerFactory<id::BufferId>
+ + IdentityHandlerFactory<id::StagingBufferId>
  + IdentityHandlerFactory<id::TextureId>
  + IdentityHandlerFactory<id::TextureViewId>
  + IdentityHandlerFactory<id::SamplerId>
@@ -639,6 +641,7 @@ pub struct Hub<A: HalApi, F: GlobalIdentityHandlerFactory> {
  pub compute_pipelines: Registry<ComputePipeline<A>, id::ComputePipelineId, F>,
  pub query_sets: Registry<QuerySet<A>, id::QuerySetId, F>,
  pub buffers: Registry<Buffer<A>, id::BufferId, F>,
+ pub staging_buffers: Registry<StagingBuffer<A>, id::StagingBufferId, F>,
  pub textures: Registry<Texture<A>, id::TextureId, F>,
  pub texture_views: Registry<TextureView<A>, id::TextureViewId, F>,
  pub samplers: Registry<Sampler<A>, id::SamplerId, F>,
@@ -659,6 +662,7 @@ impl<A: HalApi, F: GlobalIdentityHandlerFactory> Hub<A, F> {
  compute_pipelines: Registry::new(A::VARIANT, factory),
  query_sets: Registry::new(A::VARIANT, factory),
  buffers: Registry::new(A::VARIANT, factory),
+ staging_buffers: Registry::new(A::VARIANT, factory),
  textures: Registry::new(A::VARIANT, factory),
  texture_views: Registry::new(A::VARIANT, factory),
  samplers: Registry::new(A::VARIANT, factory),

wgpu-core/src/id.rs

@@ -196,6 +196,7 @@ pub type DeviceId = Id<crate::device::Device<Dummy>>;
 pub type QueueId = DeviceId;
 // Resource
 pub type BufferId = Id<crate::resource::Buffer<Dummy>>;
+pub type StagingBufferId = Id<crate::resource::StagingBuffer<Dummy>>;
 pub type TextureViewId = Id<crate::resource::TextureView<Dummy>>;
 pub type TextureId = Id<crate::resource::Texture<Dummy>>;
 pub type SamplerId = Id<crate::resource::Sampler<Dummy>>;

wgpu-core/src/resource.rs

@@ -186,6 +186,24 @@ impl<A: hal::Api> Resource for Buffer<A> {
  }
 }
 
+pub struct StagingBuffer<A: hal::Api> {
+ pub(crate) raw: A::Buffer,
+ pub(crate) size: wgt::BufferAddress,
+ pub(crate) is_coherent: bool,
+}
+
+impl<A: hal::Api> Resource for StagingBuffer<A> {
+ const TYPE: &'static str = "StagingBuffer";
+
+ fn life_guard(&self) -> &LifeGuard {
+ unreachable!()
+ }
+
+ fn label(&self) -> &str {
+ "<StagingBuffer>"
+ }
+}
+
 pub type TextureDescriptor<'a> = wgt::TextureDescriptor<Label<'a>>;
 
 #[derive(Debug)]

wgpu/src/backend/direct.rs

@@ -2192,6 +2192,58 @@ impl crate::Context for Context {
  }
  }
 
+ fn queue_validate_write_buffer(
+ &self,
+ queue: &Self::QueueId,
+ buffer: &Self::BufferId,
+ offset: wgt::BufferAddress,
+ size: wgt::BufferSize,
+ ) {
+ let global = &self.0;
+ match wgc::gfx_select!(
+ *queue => global.queue_validate_write_buffer(*queue, buffer.id, offset, size.get())
+ ) {
+ Ok(()) => (),
+ Err(err) => self.handle_error_fatal(err, "Queue::write_buffer_with"),
+ }
+ }
+
+ fn queue_create_staging_buffer(
+ &self,
+ queue: &Self::QueueId,
+ size: wgt::BufferSize,
+ ) -> QueueWriteBuffer {
+ let global = &self.0;
+ match wgc::gfx_select!(
+ *queue => global.queue_create_staging_buffer(*queue, size, PhantomData)
+ ) {
+ Ok((buffer_id, ptr)) => QueueWriteBuffer {
+ buffer_id,
+ mapping: BufferMappedRange {
+ ptr,
+ size: size.get() as usize,
+ },
+ },
+ Err(err) => self.handle_error_fatal(err, "Queue::write_buffer_with"),
+ }
+ }
+
+ fn queue_write_staging_buffer(
+ &self,
+ queue: &Self::QueueId,
+ buffer: &Self::BufferId,
+ offset: wgt::BufferAddress,
+ staging_buffer: &QueueWriteBuffer,
+ ) {
+ let global = &self.0;
+ match wgc::gfx_select!(
+ *queue => global.queue_write_staging_buffer(*queue, buffer.id, offset, staging_buffer.buffer_id)
+ ) {
+ Ok(()) => (),
+ Err(err) => self.handle_error_fatal(err, "Queue::write_buffer_with"),
+ }
+ }
+
  fn queue_write_texture(
  &self,
  queue: &Self::QueueId,
@@ -2324,6 +2376,27 @@ fn default_error_handler(err: crate::Error) {
  panic!("wgpu error: {}\n", err);
 }
 
+#[derive(Debug)]
+pub struct QueueWriteBuffer {
+ buffer_id: wgc::id::StagingBufferId,
+ mapping: BufferMappedRange,
+}
+
+impl std::ops::Deref for QueueWriteBuffer {
+ type Target = [u8];
+
+ fn deref(&self) -> &Self::Target {
+ panic!("QueueWriteBuffer is write-only!");
+ }
+}
+
+impl std::ops::DerefMut for QueueWriteBuffer {
+ fn deref_mut(&mut self) -> &mut Self::Target {
+ use crate::BufferMappedRangeSlice;
+ self.mapping.slice_mut()
+ }
+}
+
 #[derive(Debug)]
 pub struct BufferMappedRange {
  ptr: *mut u8,

wgpu/src/backend/mod.rs

@@ -1,9 +1,9 @@
 #[cfg(all(target_arch = "wasm32", not(feature = "webgl")))]
 mod web;
 #[cfg(all(target_arch = "wasm32", not(feature = "webgl")))]
-pub(crate) use web::{BufferMappedRange, Context};
+pub(crate) use web::{BufferMappedRange, Context, QueueWriteBuffer};
 
 #[cfg(any(not(target_arch = "wasm32"), feature = "webgl"))]
 mod direct;
 #[cfg(any(not(target_arch = "wasm32"), feature = "webgl"))]
-pub(crate) use direct::{BufferMappedRange, Context};
+pub(crate) use direct::{BufferMappedRange, Context, QueueWriteBuffer};

wgpu/src/backend/web.rs

@@ -2197,6 +2197,34 @@ impl crate::Context for Context {
  );
  }
 
+ fn queue_validate_write_buffer(
+ &self,
+ _queue: &Self::QueueId,
+ _buffer: &Self::BufferId,
+ _offset: wgt::BufferAddress,
+ _size: wgt::BufferSize,
+ ) {
+ // TODO
+ }
+
+ fn queue_create_staging_buffer(
+ &self,
+ _queue: &Self::QueueId,
+ size: wgt::BufferSize,
+ ) -> QueueWriteBuffer {
+ QueueWriteBuffer(vec![0; size.get() as usize].into_boxed_slice())
+ }
+
+ fn queue_write_staging_buffer(
+ &self,
+ queue: &Self::QueueId,
+ buffer: &Self::BufferId,
+ offset: wgt::BufferAddress,
+ staging_buffer: &QueueWriteBuffer,
+ ) {
+ self.queue_write_buffer(queue, buffer, offset, staging_buffer)
+ }
+
  fn queue_write_texture(
  &self,
  queue: &Self::QueueId,
@@ -2262,6 +2290,23 @@ impl crate::Context for Context {
 
 pub(crate) type SurfaceOutputDetail = ();
 
+#[derive(Debug)]
+pub struct QueueWriteBuffer(Box<[u8]>);
+
+impl std::ops::Deref for QueueWriteBuffer {
+ type Target = [u8];
+
+ fn deref(&self) -> &Self::Target {
+ panic!("QueueWriteBuffer is write-only!");
+ }
+}
+
+impl std::ops::DerefMut for QueueWriteBuffer {
+ fn deref_mut(&mut self) -> &mut Self::Target {
+ &mut self.0
+ }
+}
+
 #[derive(Debug)]
 pub struct BufferMappedRange {
  actual_mapping: js_sys::Uint8Array,

wgpu/src/lib.rs

@@ -43,7 +43,7 @@ pub use wgt::{
  QUERY_SIZE, VERTEX_STRIDE_ALIGNMENT,
 };
 
-use backend::{BufferMappedRange, Context as C};
+use backend::{BufferMappedRange, Context as C, QueueWriteBuffer};
 
 /// Filter for error scopes.
 #[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
@@ -482,6 +482,25 @@ trait Context: Debug + Send + Sized + Sync {
  offset: BufferAddress,
  data: &[u8],
  );
+ fn queue_validate_write_buffer(
+ &self,
+ queue: &Self::QueueId,
+ buffer: &Self::BufferId,
+ offset: wgt::BufferAddress,
+ size: wgt::BufferSize,
+ );
+ fn queue_create_staging_buffer(
+ &self,
+ queue: &Self::QueueId,
+ size: BufferSize,
+ ) -> QueueWriteBuffer;
+ fn queue_write_staging_buffer(
+ &self,
+ queue: &Self::QueueId,
+ buffer: &Self::BufferId,
+ offset: BufferAddress,
+ staging_buffer: &QueueWriteBuffer,
+ );
  fn queue_write_texture(
  &self,
  queue: &Self::QueueId,
@@ -3376,6 +3395,40 @@ impl<'a> RenderBundleEncoder<'a> {
  }
 }
 
+/// A write-only view into a staging buffer
+pub struct QueueWriteBufferView<'a> {
+ queue: &'a Queue,
+ buffer: &'a Buffer,
+ offset: BufferAddress,
+ inner: QueueWriteBuffer,
+}
+
+impl<'a> std::ops::Deref for QueueWriteBufferView<'a> {
+ type Target = [u8];
+
+ fn deref(&self) -> &Self::Target {
+ panic!("QueueWriteBufferView is write-only!");
+ }
+}
+
+impl<'a> std::ops::DerefMut for QueueWriteBufferView<'a> {
+ fn deref_mut(&mut self) -> &mut Self::Target {
+ &mut self.inner
+ }
+}
+
+impl<'a> Drop for QueueWriteBufferView<'a> {
+ fn drop(&mut self) {
+ Context::queue_write_staging_buffer(
+ &*self.queue.context,
+ &self.queue.id,
+ &self.buffer.id,
+ self.offset,
+ &self.inner,
+ );
+ }
+}
+
 impl Queue {
  /// Schedule a data write into `buffer` starting at `offset`.
  ///
@@ -3388,6 +3441,32 @@ impl Queue {
  Context::queue_write_buffer(&*self.context, &self.id, &buffer.id, offset, data)
  }
 
+ /// Schedule a data write into `buffer` starting at `offset` via the returned [QueueWriteBufferView].
+ ///
+ /// The returned value can be dereferenced to a `&mut [u8]`; dereferencing it to a `&[u8]` panics!
+ ///
+ /// This method is intended to have low performance costs.
+ /// As such, the write is not immediately submitted, and instead enqueued
+ /// internally to happen at the start of the next `submit()` call.
+ ///
+ /// This method fails if `size` is greater than the size of `buffer` starting at `offset`.
+ #[must_use]
+ pub fn write_buffer_with<'a>(
+ &'a self,
+ buffer: &'a Buffer,
+ offset: BufferAddress,
+ size: BufferSize,
+ ) -> QueueWriteBufferView<'a> {
+ Context::queue_validate_write_buffer(&*self.context, &self.id, &buffer.id, offset, size);
+ let staging_buffer = Context::queue_create_staging_buffer(&*self.context, &self.id, size);
+ QueueWriteBufferView {
+ queue: self,
+ buffer,
+ offset,
+ inner: staging_buffer,
+ }
+ }
+
  /// Schedule a data write into `texture`.
  ///
  /// This method is intended to have low performance costs.