Repository: denoland/deno
Pull request: fix(ext/webgpu): GPUQuene writeBuffer size given in elements of TA

Fixes https://github.com/denoland/deno/issues/31212


Changed files:
filename: ext/webgpu/queue.rs
patch: @@ -9,6 +9,7 @@ use deno_core::WebIDL;
 use deno_core::cppgc::Ref;
 use deno_core::futures::channel::oneshot;
 use deno_core::op2;
+use deno_core::v8;
 use deno_error::JsErrorBox;
 
 use crate::Instance;
@@ -128,19 +129,62 @@ impl GPUQueue {
 
   #[required(3)]
   #[undefined]
-  fn write_buffer(
+  fn write_buffer<'a>(
     &self,
+    scope: &mut v8::PinScope<'a, '_>,
     #[webidl] buffer: Ref<GPUBuffer>,
     #[webidl(options(enforce_range = true))] buffer_offset: u64,
-    #[anybuffer] buf: &[u8],
+    data_arg: v8::Local<'a, v8::Value>,
     #[webidl(default = 0, options(enforce_range = true))] data_offset: u64,
     #[webidl(options(enforce_range = true))] size: Option<u64>,
-  ) {
+  ) -> Result<(), JsErrorBox> {
+    // Per the WebGPU spec, dataOffset and size are in elements (not bytes)
+    // when data is a TypedArray, and in bytes otherwise.
+    let (buf, bytes_per_element) = if let Ok(typed_array) =
+      v8::Local::<v8::TypedArray>::try_from(data_arg)
+    {
+      let len = typed_array.length();
+      let bpe = if len > 0 {
+        typed_array.byte_length() / len
+      } else {
+        1
+      };
+      let byte_offset = typed_array.byte_offset();
+      let byte_len = typed_array.byte_length();
+      let ab = typed_array.buffer(scope).unwrap();
+      // SAFETY: Pointer is non-null, and V8 guarantees that the
+      // byte_offset is within the buffer backing store.
+      let ptr = unsafe { ab.data().unwrap().as_ptr().add(byte_offset) };
+      let buf =
+          // SAFETY: the slice is within the bounds of the backing store
+          unsafe { std::slice::from_raw_parts(ptr as *const u8, byte_len) };
+      (buf, bpe)
+    } else if let Ok(view) =
+      v8::Local::<v8::ArrayBufferView>::try_from(data_arg)
+    {
+      let byte_offset = view.byte_offset();
+      let byte_len = view.byte_length();
+      let ab = view.buffer(scope).unwrap();
+      // SAFETY: Pointer is non-null, and V8 guarantees that the
+      // byte_offset is within the buffer backing store.
+      let ptr = unsafe { ab.data().unwrap().as_ptr().add(byte_offset) };
+      // SAFETY: the slice is within the bounds of the backing store
+      let buf =
+        unsafe { std::slice::from_raw_parts(ptr as *const u8, byte_len) };
+      (buf, 1)
+    } else {
+      return Err(JsErrorBox::type_error(
+        "data must be an ArrayBuffer or ArrayBufferView",
+      ));
+    };
+
+    let data_offset_bytes = data_offset as usize * bytes_per_element;
     let data = match size {
       Some(size) => {
-        &buf[(data_offset as usize)..((data_offset + size) as usize)]
+        let size_bytes = size as usize * bytes_per_element;
+        &buf[data_offset_bytes..(data_offset_bytes + size_bytes)]
       }
-      None => &buf[(data_offset as usize)..],
+      None => &buf[data_offset_bytes..],
     };
 
     let err = self
@@ -149,6 +193,8 @@ impl GPUQueue {
       .err();
 
     self.error_handler.push_error(err);
+
+    Ok(())
   }
 
   #[required(4)]

filename: tests/unit/webgpu_test.ts
patch: @@ -649,6 +649,47 @@ Deno.test({
   device.destroy();
 });
 
+Deno.test({
+  permissions: { read: true, env: true },
+  ignore: isWsl || isCIWithoutGPU,
+}, async function webgpuWriteBufferTypedArrayElementOffsets() {
+  const adapter = await navigator.gpu.requestAdapter();
+  assert(adapter);
+  const device = await adapter.requestDevice();
+  assert(device);
+
+  const source = new Float32Array([1.0, 2.0, 3.0, 4.0]);
+  const gpuBuffer = device.createBuffer({
+    size: 2 * Float32Array.BYTES_PER_ELEMENT,
+    usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
+  });
+
+  // dataOffset=1 and size=2 should be in elements for a Float32Array
+  device.queue.writeBuffer(gpuBuffer, 0, source, 1, 2);
+
+  const readBuffer = device.createBuffer({
+    size: 2 * Float32Array.BYTES_PER_ELEMENT,
+    usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
+  });
+
+  const encoder = device.createCommandEncoder();
+  encoder.copyBufferToBuffer(
+    gpuBuffer,
+    0,
+    readBuffer,
+    0,
+    2 * Float32Array.BYTES_PER_ELEMENT,
+  );
+  device.queue.submit([encoder.finish()]);
+
+  await readBuffer.mapAsync(GPUMapMode.READ);
+  const result = new Float32Array(readBuffer.getMappedRange());
+  assertEquals(result, new Float32Array([2.0, 3.0]));
+
+  readBuffer.unmap();
+  device.destroy();
+});
+
 async function checkIsWsl() {
   return Deno.build.os === "linux" && await hasMicrosoftProcVersion();