mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-01-25 06:48:22 -05:00
10618aba98
* Start from andredaprato:webgpu-clean * Fix infs * inf wgsl function is not needed * Emulated ulong for threefry, more tests passing * Randomness tests passing * Update model export to support new changes in webgpu, efficientnet export works again * Simplify shift emulation in wgsl * Delete test file * Fix bigger than u32 u32 literal * Why was skip copies added here? * Python3.12 for webgpu tests * Fix model export syntax error * Get test ops passing with some skips * Fix lint * Much simpler shift * Run more tests * Timestamp queries are not supported in CI, so skip search tests * All fancy indexing passing * r is ctx * Run more dtype tests by using is_dtype_supported * Cleanup ulong shift rendering * UPat -> Pat, UOps -> Ops * Pat -> UPat * Refactor render_ushift if-else * Pattern to avoid ulong mul * Remove vals_dtype * is_nan trick + rewrite, test_isnan passing * Rewrite a * select(1, nan, gate) -> select(a, nan, gate) * No arg, just op * Support char, uchar, short, ushort * Run test_index_mnis now that we have uint8 * Fix pyling * Save 3 lines by using base Compiler * No more long emulation * Remove fixup_binops * No more external_local_bufx wgsl specific cstyle modif, use base extra_pm * Simpler, faster copyin/out * Skip some new tests that use long * Fix typo * copyout touchup * Save lines by using render_cast * WebGL is not supported in core, delete it from is_dtype_supported * More narrow test skips for some unary tests * TernaryOps, UnaryOps -> Ops * TinyGrad supports WebGPU * StableDiffusion demo: f16tof32 gpu is a lib, update UI * Packed load/store, no more scale_size, no core tinygrad changes * Rename copyin, copyout * Device -> dev * Fix lint * Pattern matcher rule for packed load/store * Refactor * Shorter packed load/store * this should fix lint * Fix mypy * SD compile script working * New SD webgpu UI * New default prompt * New SD weights * Fix title when webgpu not available * Run symbolic tests, simplify is_nan, use round_up * Show step time on UI * Bump minimum wgpu version to v0.19 * Fix latent --------- Co-authored-by: George Hotz <72895+geohot@users.noreply.github.com>
66 lines
4.3 KiB
Python
66 lines
4.3 KiB
Python
import functools
|
|
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)
|
|
if isinstance(val, int): wgpu_device.queue.write_buffer(buf, 0, val.to_bytes(4, "little"))
|
|
else: wgpu_device.queue.write_buffer(buf, 0, struct.pack('<f', val))
|
|
return buf
|
|
|
|
class WebGPUProgram:
|
|
def __init__(self, dev, name:str, lib:bytes):
|
|
(self.dev, self.timestamp_supported) = dev
|
|
self.name, self.lib, self.prg = name, lib, self.dev.create_shader_module(code=lib.decode()) # NOTE: this is the compiler
|
|
def __call__(self, *bufs, global_size=(1,1,1), local_size=(1,1,1), vals=(), wait=False):
|
|
wait = wait and self.timestamp_supported
|
|
binding_layouts = [{"binding": 0, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.uniform }}]
|
|
binding_layouts += [{"binding": i+1, "visibility": wgpu.ShaderStage.COMPUTE,
|
|
"buffer": {"type": wgpu.BufferBindingType.uniform if i >= len(bufs) else wgpu.BufferBindingType.storage }} for i in range(len(bufs)+len(vals))] # noqa: E501
|
|
bindings = [{"binding": 0, "resource": {"buffer": create_uniform(self.dev, float('inf')), "offset": 0, "size": 4}}]
|
|
bindings += [{"binding": i+1, "resource": {"buffer": create_uniform(self.dev, x) if i >= len(bufs) else x, "offset": 0,
|
|
"size": 4 if i >= len(bufs) else x.size}} for i,x in enumerate(bufs+vals)] # noqa: E501
|
|
bind_group_layout = self.dev.create_bind_group_layout(entries=binding_layouts)
|
|
pipeline_layout = self.dev.create_pipeline_layout(bind_group_layouts=[bind_group_layout])
|
|
bind_group = self.dev.create_bind_group(layout=bind_group_layout, entries=bindings)
|
|
compute_pipeline = self.dev.create_compute_pipeline(layout=pipeline_layout,compute={"module": self.prg, "entry_point": self.name},)
|
|
command_encoder = self.dev.create_command_encoder()
|
|
if wait:
|
|
query_set = self.dev.create_query_set(type=wgpu.QueryType.timestamp, count=2)
|
|
query_buf = self.dev.create_buffer(size=16, usage=wgpu.BufferUsage.QUERY_RESOLVE | wgpu.BufferUsage.COPY_SRC)
|
|
timestamp_writes = {"query_set": query_set, "beginning_of_pass_write_index": 0, "end_of_pass_write_index": 1}
|
|
compute_pass = command_encoder.begin_compute_pass(timestamp_writes=timestamp_writes if wait else None) # pylint: disable=E0606
|
|
compute_pass.set_pipeline(compute_pipeline)
|
|
compute_pass.set_bind_group(0, bind_group, [], 0, 999999) # last 2 not used
|
|
compute_pass.dispatch_workgroups(*global_size) # x y z
|
|
compute_pass.end()
|
|
if wait:
|
|
command_encoder.resolve_query_set(query_set=query_set, first_query=0, query_count=2, destination=query_buf, destination_offset=0)
|
|
self.dev.queue.submit([command_encoder.finish()])
|
|
return ((timestamps:=self.dev.queue.read_buffer(query_buf).cast("Q").tolist())[1] - timestamps[0]) / 1e9 if wait else None
|
|
|
|
# WebGPU buffers have to be 4-byte aligned
|
|
class WebGpuAllocator(Allocator):
|
|
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)
|
|
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
|
|
|
|
class WebGpuDevice(Compiled):
|
|
def __init__(self, device:str):
|
|
adapter = wgpu.gpu.request_adapter_sync(power_preference="high-performance")
|
|
timestamp_supported = wgpu.FeatureName.timestamp_query in adapter.features
|
|
wgpu_device = adapter.request_device_sync(required_features=[wgpu.FeatureName.timestamp_query] if timestamp_supported else [])
|
|
super().__init__(device, WebGpuAllocator(wgpu_device), WGSLRenderer(), Compiler(),
|
|
functools.partial(WebGPUProgram, (wgpu_device, timestamp_supported)))
|