tinygrad 0.10.0__py3-none-any.whl → 0.10.2__py3-none-any.whl

tinygrad/runtime/ops_webgpu.py

@@ -0,0 +1,225 @@
+import functools, struct
+from tinygrad.device import  Compiled, Allocator, Compiler
+from tinygrad.renderer.wgsl import WGSLRenderer
+from tinygrad.helpers import round_up, OSX
+from tinygrad.runtime.autogen import webgpu
+from typing import List, Any
+import ctypes
+import os
+
+backend_types = {v: k for k, v in webgpu.WGPUBackendType__enumvalues.items() }
+
+try:
+  instance = webgpu.wgpuCreateInstance(webgpu.WGPUInstanceDescriptor(features = webgpu.WGPUInstanceFeatures(timedWaitAnyEnable = True)))
+except AttributeError:
+  raise RuntimeError("Cannot find dawn library. Install it with: " + ("brew tap wpmed92/dawn && brew install dawn" if OSX else
+  "sudo curl -L https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.so -o /usr/lib/libwebgpu_dawn.so"))
+
+def to_c_string(_str): return ctypes.create_string_buffer(_str.encode('utf-8'))
+
+def from_wgpu_str(string_view): return ctypes.string_at(string_view.data, string_view.length).decode("utf-8")
+
+def to_wgpu_str(_str):
+  return webgpu.WGPUStringView(data=ctypes.cast(ctypes.pointer(to_c_string(_str)), ctypes.POINTER(ctypes.c_char)), length=len(_str))
+
+def _wait(future):
+  assert webgpu.wgpuInstanceWaitAny(instance, 1, webgpu.WGPUFutureWaitInfo(future=future), 2**64-1) == webgpu.WGPUWaitStatus_Success, "Future failed"
+
+def write_buffer(device, buf, offset, src):
+  src = bytearray(src)
+  webgpu.wgpuQueueWriteBuffer(webgpu.wgpuDeviceGetQueue(device), buf, offset, (ctypes.c_uint8 * len(src)).from_buffer(src), len(src))
+
+def _run(async_fun, cb_info_type, cb_type, status_enum, res_idx, msg_idx, *params):
+  result: List[Any] = []
+
+  def cb(*params):
+    result[:] = params
+    if msg_idx: result[msg_idx] = from_wgpu_str(result[msg_idx])
+
+  cb_info = cb_info_type(nextInChain=None, mode=webgpu.WGPUCallbackMode_WaitAnyOnly, callback=cb_type(cb))
+  _wait(async_fun(*params, cb_info))
+
+  if result[0] != 1: raise RuntimeError(f"[{status_enum[result[0]] if status_enum else 'ERROR'}]{result[msg_idx] if msg_idx else ''}")
+  return result[res_idx] if res_idx else None
+
+def copy_buffer_to_buffer(dev, src, src_offset, dst, dst_offset, size):
+  encoder = webgpu.wgpuDeviceCreateCommandEncoder(dev, webgpu.WGPUCommandEncoderDescriptor())
+  webgpu.wgpuCommandEncoderCopyBufferToBuffer(encoder, src, src_offset, dst, dst_offset, size)
+  cb = webgpu.wgpuCommandEncoderFinish(encoder, webgpu.WGPUCommandBufferDescriptor())
+  webgpu.wgpuQueueSubmit(webgpu.wgpuDeviceGetQueue(dev), 1, (webgpu.WGPUCommandBuffer*1)(cb))
+  webgpu.wgpuCommandBufferRelease(cb)
+  webgpu.wgpuCommandEncoderRelease(encoder)
+
+def read_buffer(dev, buf):
+  size = webgpu.wgpuBufferGetSize(buf)
+  tmp_buffer = webgpu.wgpuDeviceCreateBuffer(dev, webgpu.WGPUBufferDescriptor(size=size,
+    usage=webgpu.WGPUBufferUsage_CopyDst | webgpu.WGPUBufferUsage_MapRead, mappedAtCreation=False))
+  copy_buffer_to_buffer(dev, buf, 0, tmp_buffer, 0, size)
+  _run(webgpu.wgpuBufferMapAsync2, webgpu.WGPUBufferMapCallbackInfo2, webgpu.WGPUBufferMapCallback2, webgpu.WGPUBufferMapAsyncStatus__enumvalues,
+    None, 0, tmp_buffer, webgpu.WGPUMapMode_Read, 0, size)
+  void_ptr = ctypes.cast(webgpu.wgpuBufferGetConstMappedRange(tmp_buffer, 0, size), ctypes.c_void_p)
+  buf_copy = bytearray((ctypes.c_uint8 * size).from_address(void_ptr.value))
+  webgpu.wgpuBufferUnmap(tmp_buffer)
+  webgpu.wgpuBufferDestroy(tmp_buffer)
+  return memoryview(buf_copy).cast("B")
+
+def pop_error(device):
+  return _run(webgpu.wgpuDevicePopErrorScopeF, webgpu.WGPUPopErrorScopeCallbackInfo, webgpu.WGPUPopErrorScopeCallback, None, 2, 2, device)
+
+def create_uniform(wgpu_device, val):
+  buf = webgpu.wgpuDeviceCreateBuffer(wgpu_device,
+    webgpu.WGPUBufferDescriptor(size=4, usage=webgpu.WGPUBufferUsage_Uniform | webgpu.WGPUBufferUsage_CopyDst))
+  write_buffer(wgpu_device, buf, 0, val.to_bytes(4, "little") if isinstance(val, int) else struct.pack('<f', val))
+  return buf
+
+class WebGPUProgram:
+  def __init__(self, dev, name:str, lib:bytes):
+    (self.dev, self.timestamp_supported) = dev
+
+    # Creating shader module
+    shader = webgpu.WGPUShaderModuleWGSLDescriptor(code=to_wgpu_str(lib.decode()),
+      chain=webgpu.WGPUChainedStruct(sType=webgpu.WGPUSType_ShaderSourceWGSL))
+    module = webgpu.WGPUShaderModuleDescriptor()
+    module.nextInChain = ctypes.cast(ctypes.pointer(shader), ctypes.POINTER(webgpu.struct_WGPUChainedStruct))
+
+    # Check compiler error
+    webgpu.wgpuDevicePushErrorScope(self.dev, webgpu.WGPUErrorFilter_Validation)
+    shader_module = webgpu.wgpuDeviceCreateShaderModule(self.dev, module)
+
+    if err := pop_error(self.dev): raise RuntimeError(f"Shader compilation failed: {err}")
+
+    self.name, self.lib, self.prg = name, lib, shader_module
+  def __call__(self, *bufs, global_size=(1,1,1), local_size=(1,1,1), vals=(), wait=False):
+    wait = wait and self.timestamp_supported
+    tmp_bufs = [*bufs]
+    buf_patch = False
+
+    # WebGPU does not allow using the same buffer for input and output
+    for i in range(1, len(bufs)):
+      if bufs[i] == bufs[0]:
+        tmp_bufs[0] = webgpu.wgpuDeviceCreateBuffer(self.dev,
+          webgpu.WGPUBufferDescriptor(size=webgpu.wgpuBufferGetSize(bufs[0]), usage=webgpu.wgpuBufferGetUsage(bufs[0])))
+        buf_patch = True
+
+    # Creating bind group layout
+    binding_layouts = [webgpu.WGPUBindGroupLayoutEntry(binding=0, visibility= webgpu.WGPUShaderStage_Compute,
+      buffer=webgpu.WGPUBufferBindingLayout(type=webgpu.WGPUBufferBindingType_Uniform))]
+    binding_layouts += [webgpu.WGPUBindGroupLayoutEntry(binding=i+1, visibility=webgpu.WGPUShaderStage_Compute,
+      buffer=webgpu.WGPUBufferBindingLayout(type=webgpu.WGPUBufferBindingType_Uniform if i >= len(tmp_bufs)
+      else webgpu.WGPUBufferBindingType_Storage)) for i in range(len(tmp_bufs)+len(vals))]
+
+    bl_arr_type = webgpu.WGPUBindGroupLayoutEntry * len(binding_layouts)
+    webgpu.wgpuDevicePushErrorScope(self.dev, webgpu.WGPUErrorFilter_Validation)
+    bind_group_layouts = [webgpu.wgpuDeviceCreateBindGroupLayout(self.dev, webgpu.WGPUBindGroupLayoutDescriptor(
+      entryCount=len(binding_layouts), entries=ctypes.cast(bl_arr_type(*binding_layouts), ctypes.POINTER(webgpu.WGPUBindGroupLayoutEntry))))]
+
+    if bg_layout_err := pop_error(self.dev): raise RuntimeError(f"Error creating bind group layout: {bg_layout_err}")
+
+    # Creating pipeline layout
+    pipeline_layout_desc = webgpu.WGPUPipelineLayoutDescriptor(bindGroupLayoutCount=len(bind_group_layouts),
+      bindGroupLayouts = (webgpu.WGPUBindGroupLayout * len(bind_group_layouts))(*bind_group_layouts))
+
+    webgpu.wgpuDevicePushErrorScope(self.dev, webgpu.WGPUErrorFilter_Validation)
+    pipeline_layout = webgpu.wgpuDeviceCreatePipelineLayout(self.dev, pipeline_layout_desc)
+
+    if pipe_err := pop_error(self.dev): raise RuntimeError(f"Error creating pipeline layout: {pipe_err}")
+
+    # Creating bind group
+    bindings = [webgpu.WGPUBindGroupEntry(binding=0, buffer=create_uniform(self.dev, float('inf')), offset=0, size=4)]
+    bindings += [webgpu.WGPUBindGroupEntry(binding=i+1, buffer=create_uniform(self.dev, x) if i >= len(tmp_bufs) else x, offset=0,
+      size=4 if i >= len(tmp_bufs) else webgpu.wgpuBufferGetSize(x)) for i,x in enumerate(tuple(tmp_bufs)+vals)]
+
+    bg_arr_type = webgpu.WGPUBindGroupEntry * len(bindings)
+    bind_group_desc = webgpu.WGPUBindGroupDescriptor(layout=bind_group_layouts[0], entryCount=len(bindings), entries=bg_arr_type(*bindings))
+    webgpu.wgpuDevicePushErrorScope(self.dev, webgpu.WGPUErrorFilter_Validation)
+    bind_group = webgpu.wgpuDeviceCreateBindGroup(self.dev, bind_group_desc)
+
+    if bind_err := pop_error(self.dev): raise RuntimeError(f"Error creating bind group: {bind_err}")
+
+    # Creating compute pipeline
+    compute_desc = webgpu.WGPUComputePipelineDescriptor(layout=pipeline_layout,
+      compute=webgpu.WGPUComputeState(module=self.prg, entryPoint=to_wgpu_str(self.name)))
+    pipeline_result = _run(webgpu.wgpuDeviceCreateComputePipelineAsync2, webgpu.WGPUCreateComputePipelineAsyncCallbackInfo2,
+    webgpu.WGPUCreateComputePipelineAsyncCallback2, webgpu.WGPUCreatePipelineAsyncStatus__enumvalues, 1, None, self.dev, compute_desc)
+
+    command_encoder = webgpu.wgpuDeviceCreateCommandEncoder(self.dev, webgpu.WGPUCommandEncoderDescriptor())
+    comp_pass_desc = webgpu.WGPUComputePassDescriptor(nextInChain=None)
+
+    if wait:
+      query_set = webgpu.wgpuDeviceCreateQuerySet(self.dev, webgpu.WGPUQuerySetDescriptor(type=webgpu.WGPUQueryType_Timestamp, count=2))
+      query_buf = webgpu.wgpuDeviceCreateBuffer(self.dev,
+        webgpu.WGPUBufferDescriptor(size=16, usage=webgpu.WGPUBufferUsage_QueryResolve | webgpu.WGPUBufferUsage_CopySrc))
+      comp_pass_desc.timestampWrites = ctypes.pointer(webgpu.WGPUComputePassTimestampWrites(
+        querySet=query_set, beginningOfPassWriteIndex=0, endOfPassWriteIndex=1))
+
+    # Begin compute pass
+    compute_pass = webgpu.wgpuCommandEncoderBeginComputePass(command_encoder, comp_pass_desc)
+    webgpu.wgpuComputePassEncoderSetPipeline(compute_pass, pipeline_result)
+    webgpu.wgpuComputePassEncoderSetBindGroup(compute_pass, 0, bind_group, 0, None)
+    webgpu.wgpuComputePassEncoderDispatchWorkgroups(compute_pass, *global_size)
+    webgpu.wgpuComputePassEncoderEnd(compute_pass)
+
+    if wait: webgpu.wgpuCommandEncoderResolveQuerySet(command_encoder, query_set, 0, 2, query_buf, 0)
+
+    cmd_buf = webgpu.wgpuCommandEncoderFinish(command_encoder, webgpu.WGPUCommandBufferDescriptor())
+    webgpu.wgpuQueueSubmit(webgpu.wgpuDeviceGetQueue(self.dev), 1, (webgpu.WGPUCommandBuffer*1)(cmd_buf))
+
+    if buf_patch:
+      copy_buffer_to_buffer(self.dev, tmp_bufs[0], 0, bufs[0], 0, webgpu.wgpuBufferGetSize(bufs[0]))
+      webgpu.wgpuBufferDestroy(tmp_bufs[0])
+
+    if wait:
+      time = ((timestamps:=read_buffer(self.dev, query_buf).cast("Q").tolist())[1] - timestamps[0]) / 1e9
+      webgpu.wgpuBufferDestroy(query_buf)
+      webgpu.wgpuQuerySetDestroy(query_set)
+      return time
+
+class WebGpuAllocator(Allocator):
+  def __init__(self, dev): self.dev = dev
+  def _alloc(self, size: int, options):
+    # WebGPU buffers have to be 4-byte aligned
+    return webgpu.wgpuDeviceCreateBuffer(self.dev, webgpu.WGPUBufferDescriptor(size=round_up(size, 4),
+      usage=webgpu.WGPUBufferUsage_Storage | webgpu.WGPUBufferUsage_CopyDst | webgpu.WGPUBufferUsage_CopySrc))
+  def _copyin(self, dest, src: memoryview):
+    if src.nbytes % 4:
+      padded_src = bytearray(round_up(src.nbytes, 4))
+      padded_src[:src.nbytes] = src
+    write_buffer(self.dev, dest, 0, padded_src if src.nbytes % 4 else src)
+  def _copyout(self, dest: memoryview, src):
+    buffer_data = read_buffer(self.dev, src)
+    dest[:] = buffer_data[:dest.nbytes] if webgpu.wgpuBufferGetSize(src)  > dest.nbytes else buffer_data
+  def _free(self, opaque, options):
+    webgpu.wgpuBufferDestroy(opaque)
+
+class WebGpuDevice(Compiled):
+  def __init__(self, device:str):
+    # Requesting an adapter
+    adapter_res = _run(webgpu.wgpuInstanceRequestAdapterF, webgpu.WGPURequestAdapterCallbackInfo, webgpu.WGPURequestAdapterCallback,
+    webgpu.WGPURequestAdapterStatus__enumvalues, 1, 2, instance,
+
+    webgpu.WGPURequestAdapterOptions(powerPreference=webgpu.WGPUPowerPreference_HighPerformance,
+      backendType=backend_types.get(os.getenv("WEBGPU_BACKEND", ""), 0)))
+
+    # Get supported features
+    supported_features = webgpu.WGPUSupportedFeatures()
+    webgpu.wgpuAdapterGetFeatures(adapter_res, supported_features)
+    supported = [supported_features.features[i] for i in range(supported_features.featureCount)]
+    features = [feat for feat in [webgpu.WGPUFeatureName_TimestampQuery, webgpu.WGPUFeatureName_ShaderF16] if feat in supported]
+    dev_desc = webgpu.WGPUDeviceDescriptor(requiredFeatureCount=len(features),requiredFeatures=(webgpu.WGPUFeatureName * len(features))(*features))
+
+    # Limits
+    supported_limits = webgpu.WGPUSupportedLimits()
+    webgpu.wgpuAdapterGetLimits(adapter_res, ctypes.cast(ctypes.pointer(supported_limits),ctypes.POINTER(webgpu.struct_WGPUSupportedLimits)))
+    limits = webgpu.WGPURequiredLimits(limits=supported_limits.limits)
+    dev_desc.requiredLimits = ctypes.cast(ctypes.pointer(limits),ctypes.POINTER(webgpu.struct_WGPURequiredLimits))
+
+    # Requesting a device
+    device_res = _run(webgpu.wgpuAdapterRequestDeviceF, webgpu.WGPURequestDeviceCallbackInfo, webgpu.WGPURequestDeviceCallback,
+    webgpu.WGPURequestDeviceStatus__enumvalues, 1, 2, adapter_res, dev_desc)
+
+    super().__init__(device, WebGpuAllocator(device_res), WGSLRenderer(), Compiler(),
+      functools.partial(WebGPUProgram, (device_res, webgpu.WGPUFeatureName_TimestampQuery in supported)))
+
+  def synchronize(self):
+    _run(webgpu.wgpuQueueOnSubmittedWorkDone2, webgpu.WGPUQueueWorkDoneCallbackInfo2, webgpu.WGPUQueueWorkDoneCallback2,
+    webgpu.WGPUQueueWorkDoneStatus__enumvalues, None, None, webgpu.wgpuDeviceGetQueue(self.runtime.args[0][0]))

tinygrad/runtime/support/allocator.py

@@ -0,0 +1,94 @@
+import collections
+from tinygrad.helpers import round_up
+
+class TLSFAllocator:
+  """
+  The allocator is based on the Two-Level Segregated Fit (TLSF) algorithm. The allocator maintains 2 level of buckets:
+    * 1st level is determined by the most significant bit of the size.
+    * 2nd level splits the covered memory of 1st level into @lv2_cnt entries.
+
+  For each allocation request, the allocator searches for the smallest block that can fit the requested size.
+  For each deallocation request, the allocator merges the block with its neighbors if they are free.
+  """
+
+  def __init__(self, size:int, base:int=0, block_size:int=16, lv2_cnt:int=16):
+    self.size, self.base, self.block_size, self.l2_cnt = size, base, block_size, lv2_cnt.bit_length()
+    self.storage:list = [collections.defaultdict(list) for _ in range(size.bit_length() + 1)]
+    self.lv1_entries:list[int] = [0] * len(self.storage)
+
+    # self.blocks is more like a linked list, where each entry is a contigous block.
+    self.blocks:dict[int, tuple[int, int|None, int|None, bool]] = {0: (size, None, None, True)} # size, next, prev, is_free
+    self._insert_block(0, size)
+
+  def lv1(self, size): return size.bit_length()
+  def lv2(self, size): return (size - (1 << (size.bit_length() - 1))) // (1 << max(0, size.bit_length() - self.l2_cnt))
+
+  def _insert_block(self, start:int, size:int, prev:int|None=None):
+    if prev is None: prev = self.blocks[start][2]
+    self.storage[self.lv1(size)][self.lv2(size)].append(start)
+    self.lv1_entries[self.lv1(size)] += 1
+    self.blocks[start] = (size, start + size, prev, True)
+    return self
+
+  def _remove_block(self, start:int, size:int, prev:int|None=None):
+    if prev is None: prev = self.blocks[start][2]
+    self.storage[self.lv1(size)][self.lv2(size)].remove(start)
+    self.lv1_entries[self.lv1(size)] -= 1
+    self.blocks[start] = (size, start + size, prev, False)
+    return self
+
+  def _split_block(self, start:int, size:int, new_size:int):
+    nxt = self.blocks[start][1]
+    assert self.blocks[start][3], "block must be free"
+    self._remove_block(start, size)._insert_block(start, new_size)._insert_block(start + new_size, size - new_size, prev=start)
+    if nxt in self.blocks: self.blocks[nxt] = (self.blocks[nxt][0], self.blocks[nxt][1], start + new_size, self.blocks[nxt][3])
+    return self
+
+  def _merge_right(self, start:int):
+    size, nxt, _, is_free = self.blocks[start]
+    assert is_free, "block must be free"
+
+    while is_free and nxt in self.blocks:
+      if (blk:=self.blocks[nxt])[3] is False: break
+      self._remove_block(start, size)._remove_block(nxt, blk[0])._insert_block(start, size:=size + blk[0])
+      assert self.blocks[start][1] == blk[1]
+      _, nxt, _, _ = self.blocks.pop(nxt)
+
+    if nxt in self.blocks: self.blocks[nxt] = (self.blocks[nxt][0], self.blocks[nxt][1], start, self.blocks[nxt][3])
+
+  def _merge_block(self, start:int):
+    # Go left while blocks are free. Then merge all them right.
+    while (x:=self.blocks[start][2]) is not None and self.blocks[x][3] is True: start = x
+    self._merge_right(start)
+
+  def alloc(self, req_size:int, align:int=1) -> int:
+    req_size = max(self.block_size, req_size) # at least block size.
+    size = max(self.block_size, req_size + align - 1)
+
+    # Round up the allocation size to the next bucket, so any entry there can fit the requested size.
+    size = round_up(size, (1 << size.bit_length() - self.l2_cnt))
+
+    # Search for the smallest block that can fit the requested size. Start with the it's bucket and go up until any block is found.
+    for l1 in range(self.lv1(size), len(self.storage)):
+      if self.lv1_entries[l1] == 0: continue
+      for l2 in range(self.lv2(size) if l1 == size.bit_length() else 0, (1 << self.l2_cnt)):
+        if len(self.storage[l1][l2]) > 0:
+          nsize = self.blocks[self.storage[l1][l2][0]][0]
+          assert nsize >= size, "block must be larger"
+
+          # Block start address.
+          start = self.storage[l1][l2][0]
+
+          # If request contains alignment, split the block into two parts.
+          if (new_start:=round_up(start, align)) != start:
+            self._split_block(start, nsize, new_start - start)
+            start, nsize = new_start, self.blocks[new_start][0]
+
+          # If the block is larger than the requested size, split it into two parts.
+          if nsize > req_size: self._split_block(start, nsize, req_size)
+          self._remove_block(start, req_size) # Mark the block as allocated.
+          return start + self.base
+    raise MemoryError(f"Can't allocate {req_size} bytes")
+
+  def free(self, start:int):
+    self._insert_block(start - self.base, self.blocks[start - self.base][0])._merge_block(start - self.base)