Refactor copy_buffer_to_buffer

wpmed92 · Jan 1, 2025 · 0782ddb · 0782ddb
1 parent 58878f2
commit 0782ddb
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 7 deletions.
diff --git a/pydawn/utils.py b/pydawn/utils.py
@@ -127,17 +127,19 @@ def cb(status, msg, u1, u2):
     if result.status != webgpu.WGPUBufferMapAsyncStatus_Success:
         raise RuntimeError(f"Failed to map buffer: [{webgpu.WGPUBufferMapAsyncStatus__enumvalues[result.status]}] {result.msg}")
 
-def copy_buffer_to_buffer(cmd_encoder, src, src_offset, dst, dst_offset, size):
-    webgpu.wgpuCommandEncoderCopyBufferToBuffer(cmd_encoder, src, src_offset, dst, dst_offset, size)
+def copy_buffer_to_buffer(device, src, src_offset, dst, dst_offset, size):
+    encoder = create_command_encoder(device)
+    webgpu.wgpuCommandEncoderCopyBufferToBuffer(encoder, src, src_offset, dst, dst_offset, size)
+    cb = command_encoder_finish(encoder)
+    submit(device, [cb])
+    webgpu.wgpuCommandBufferRelease(cb)
+    webgpu.wgpuCommandEncoderRelease(encoder)
 
 def read_buffer(dev, buf):
     size = webgpu.wgpuBufferGetSize(buf)
     tmp_usage = webgpu.WGPUBufferUsage_CopyDst | webgpu.WGPUBufferUsage_MapRead
     tmp_buffer = create_buffer(dev, size, tmp_usage)
-    encoder = create_command_encoder(dev)
-    copy_buffer_to_buffer(encoder, buf, 0, tmp_buffer, 0, size)
-    cb = command_encoder_finish(encoder)
-    submit(dev, [cb])
+    copy_buffer_to_buffer(dev, buf, 0, tmp_buffer, 0, size)
     sync(dev)
     map_buffer(tmp_buffer, size)
     ptr = webgpu.wgpuBufferGetConstMappedRange(tmp_buffer, 0, size)

diff --git a/test/test_utils.py b/test/test_utils.py
@@ -22,7 +22,16 @@ def test_buffer_read_write(self):
         test_src = [1, 2, 3, 4]
         utils.write_buffer(self.device, buf, 0, test_src)
         out = utils.read_buffer(self.device, buf)
-        self.assertEqual(list(out), test_src, f"buffer mismatch: (actual) {test_src} != (expected) {out}")
+        self.assertEqual(list(out), test_src, f"buffer mismatch: (actual) {test_src} != (expected) {list(out)}")
+
+    def test_copy_buffer_to_buffer(self):
+        src = utils.create_buffer(self.device, 4, webgpu.WGPUBufferUsage_Storage | webgpu.WGPUBufferUsage_CopySrc | webgpu.WGPUBufferUsage_CopyDst)
+        dst = utils.create_buffer(self.device, 4, webgpu.WGPUBufferUsage_Storage | webgpu.WGPUBufferUsage_CopySrc | webgpu.WGPUBufferUsage_CopyDst)
+        test_src = [4, 8, 16, 32]
+        utils.write_buffer(self.device, src, 0, test_src)
+        utils.copy_buffer_to_buffer(self.device, src, 0, dst, 0, 4)
+        dst_out = utils.read_buffer(self.device, dst)
+        self.assertEqual(list(dst_out), test_src, f"buffer mismatch: (actual) {test_src} != (expected) {list(dst_out)}")
 
     def test_create_shader_module(self):
         shader_source = """