From 2ed673f220e0cb9afc5ac7d9ae485629d85ef19e Mon Sep 17 00:00:00 2001 From: Chen-Yu Yang Date: Thu, 12 Dec 2024 08:35:34 -0800 Subject: [PATCH] simpler webgpu copyin [pr] can directly construct the buffer instead of init then assign. also put standard lib import together --- tinygrad/runtime/ops_webgpu.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/tinygrad/runtime/ops_webgpu.py b/tinygrad/runtime/ops_webgpu.py index fe67491cc6fe6..19553a295965f 100644 --- a/tinygrad/runtime/ops_webgpu.py +++ b/tinygrad/runtime/ops_webgpu.py @@ -1,9 +1,8 @@ -import functools +import functools, struct from tinygrad.device import Compiled, Allocator, Compiler from tinygrad.renderer.wgsl import WGSLRenderer from tinygrad.helpers import round_up import wgpu -import struct def create_uniform(wgpu_device, val) -> wgpu.GPUBuffer: buf = wgpu_device.create_buffer(size=4, usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST) @@ -47,10 +46,8 @@ def __init__(self, dev): self.dev = dev def _alloc(self, size: int, options): return self.dev.create_buffer(size=round_up(size, 4), usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_DST | wgpu.BufferUsage.COPY_SRC) def _copyin(self, dest, src: memoryview): - if src.nbytes % 4: - padded_src = bytearray(round_up(src.nbytes, 4)) - padded_src[:src.nbytes] = src - self.dev.queue.write_buffer(dest, 0, padded_src if src.nbytes % 4 else src) + if (pad := src.nbytes % 4): src = memoryview(bytearray(src) + bytearray(pad)) + self.dev.queue.write_buffer(dest, 0, src) def _copyout(self, dest: memoryview, src): buffer_data = self.dev.queue.read_buffer(src, 0) dest[:] = buffer_data[:dest.nbytes] if src._nbytes > dest.nbytes else buffer_data