mirror of
https://jihulab.com/mr-one/openpilot.git
synced 2026-10-27 21:13:42 +08:00
1okko
8d0bf5ee0a
The current recompiled42 model catalog declares tinygrad_ref=fe5d3169b. The old vendored e837 runtime accepted the pkl but left policy outputs at zero, so sync the runtime package to the compiler version used by the downloaded models.
222 lines
13 KiB
Python
222 lines
13 KiB
Python
import functools, struct
|
|
from tinygrad.device import BufferStorage, Compiled, Allocator, BufferSpec, Program, TinyELF
|
|
from tinygrad.renderer.wgsl import WGSLRenderer
|
|
from tinygrad.helpers import round_up, suppress_finalizing, getenv, to_mv
|
|
from tinygrad.runtime.autogen import webgpu
|
|
from tinygrad.runtime.support import c
|
|
from typing import Callable
|
|
import ctypes
|
|
|
|
backend_types = {v: k for k, v in webgpu.enum_WGPUBackendType.items()}
|
|
instance = webgpu.wgpuCreateInstance(webgpu.WGPUInstanceDescriptor(features=webgpu.WGPUInstanceFeatures(timedWaitAnyEnable=True)))
|
|
|
|
def from_wgpu_str(string_view:webgpu.WGPUStringView) -> str: return ctypes.string_at(string_view.data, string_view.length).decode()
|
|
def to_wgpu_str(_str:str) -> webgpu.WGPUStringView: return webgpu.WGPUStringView(data=ctypes.create_string_buffer(_str.encode()), length=len(_str))
|
|
|
|
# gets a memoryview from a buffer, which is assumed to have MAP_READ (see _readable_buffer)
|
|
def buf_to_mv(buf:webgpu.WGPUBuffer) -> memoryview:
|
|
BufferMapAsync(buf, webgpu.WGPUMapMode_Read, 0, size:=webgpu.wgpuBufferGetSize(buf))
|
|
return to_mv(webgpu.wgpuBufferGetConstMappedRange(buf, 0, size), size)
|
|
|
|
# turns a webgpu function returning a future into python-synchronous function
|
|
# the new function handles the status code and optional error message, returning the other callback arguments
|
|
def synchronous(status_enum:dict[int, str], has_emsg:bool=False):
|
|
def wrap(fn:Callable[..., webgpu.WGPUFuture]) -> Callable:
|
|
@functools.wraps(fn)
|
|
def wrapper(*args):
|
|
status, payload, emsg = 0, [], None
|
|
|
|
@next(ty for nm, ty, *_ in fn.argtypes[-1]._real_fields_ if nm == "callback") # type: ignore
|
|
def cb(s:int, *args):
|
|
nonlocal status, payload, emsg
|
|
# the last two arguments are "userdata1" and "userdata2", which we drop
|
|
# we must process wgpu strings in this callback, as they will be freed after we return
|
|
status, (*payload, emsg) = s, [from_wgpu_str(a) if type(a) is webgpu.WGPUStringView else a for a in args[:-2]] + ([] if has_emsg else [None])
|
|
|
|
future = fn(*args, fn.argtypes[-1](mode=webgpu.WGPUCallbackMode_WaitAnyOnly, callback=cb)) # type: ignore
|
|
if (future_status:=webgpu.wgpuInstanceWaitAny(instance, 1, webgpu.WGPUFutureWaitInfo(future), 2**64-1)) != webgpu.WGPUWaitStatus_Success:
|
|
raise RuntimeError(f"error while waiting for future ({fn.__name__}): {webgpu.enum_WGPUWaitStatus.get(future_status)}")
|
|
|
|
if status != 1: raise RuntimeError(f"[{status_enum.get(status)}]{emsg or ''}")
|
|
return payload if len(payload) > 1 else payload[0] if len(payload) == 1 else None
|
|
return wrapper
|
|
return wrap
|
|
|
|
BufferMapAsync = synchronous(webgpu.enum_WGPUBufferMapAsyncStatus, True)(webgpu.wgpuBufferMapAsync2)
|
|
DevicePopErrorScope = synchronous(webgpu.enum_WGPUPopErrorScopeStatus)(webgpu.wgpuDevicePopErrorScope2)
|
|
DeviceCreateComputePipeline = synchronous(webgpu.enum_WGPUCreatePipelineAsyncStatus, True)(webgpu.wgpuDeviceCreateComputePipelineAsync2)
|
|
InstanceRequestAdapter = synchronous(webgpu.enum_WGPURequestAdapterStatus, True)(webgpu.wgpuInstanceRequestAdapter2)
|
|
AdapterRequestDevice = synchronous(webgpu.enum_WGPURequestDeviceStatus, True)(webgpu.wgpuAdapterRequestDevice2)
|
|
QueueOnSubmittedWorkDone = synchronous(webgpu.enum_WGPUQueueWorkDoneStatus)(webgpu.wgpuQueueOnSubmittedWorkDone2)
|
|
|
|
class WebGPUProgram(Program['WebGpuDevice']):
|
|
def __init__(self, dev:'WebGpuDevice', obj:TinyELF):
|
|
self.dev, self.name = dev, to_wgpu_str(obj.name)
|
|
|
|
# Creating shader module
|
|
shader = webgpu.WGPUShaderModuleWGSLDescriptor(code=to_wgpu_str(obj.lib.decode()),
|
|
chain=webgpu.WGPUChainedStruct(sType=webgpu.WGPUSType_ShaderSourceWGSL))
|
|
module = webgpu.WGPUShaderModuleDescriptor(nextInChain=ctypes.cast(ctypes.pointer(shader), ctypes.POINTER(webgpu.struct_WGPUChainedStruct)))
|
|
|
|
# Check compiler error
|
|
webgpu.wgpuDevicePushErrorScope(self.dev.device_res, webgpu.WGPUErrorFilter_Validation)
|
|
self.prg = webgpu.wgpuDeviceCreateShaderModule(self.dev.device_res, module)
|
|
if err := self.dev.pop_error(): raise RuntimeError(f"Shader compilation failed: {err}")
|
|
|
|
@suppress_finalizing
|
|
def __del__(self): webgpu.wgpuShaderModuleRelease(self.prg)
|
|
|
|
def __call__(self, *bufs:webgpu.WGPUBuffer, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
|
|
vals:tuple[int, ...]=(), wait=False, **kw) -> float|None:
|
|
wait = wait and webgpu.WGPUFeatureName_TimestampQuery in self.dev.features
|
|
|
|
# Creating bind group layout
|
|
def bgl_entry(n:int, ty:str):
|
|
return webgpu.WGPUBindGroupLayoutEntry(binding=n, visibility=webgpu.WGPUShaderStage_Compute,
|
|
buffer=webgpu.WGPUBufferBindingLayout(type=getattr(webgpu, f'WGPUBufferBindingType_{ty}')))
|
|
bind_entries = (webgpu.WGPUBindGroupLayoutEntry * (1+len(bufs)+len(vals)))(
|
|
bgl_entry(0, 'Uniform'), *(bgl_entry(i+1, 'Uniform' if i >= len(bufs) else 'Storage') for i in range(len(bufs)+len(vals))))
|
|
|
|
webgpu.wgpuDevicePushErrorScope(self.dev.device_res, webgpu.WGPUErrorFilter_Validation)
|
|
bind_layout = webgpu.wgpuDeviceCreateBindGroupLayout(self.dev.device_res,
|
|
webgpu.WGPUBindGroupLayoutDescriptor(entryCount=len(bind_entries), entries=bind_entries))
|
|
|
|
if err := self.dev.pop_error(): raise RuntimeError(f"Error creating bind group layout: {err}")
|
|
|
|
# Creating pipeline layout
|
|
pipeline_layout_desc = webgpu.WGPUPipelineLayoutDescriptor(bindGroupLayoutCount=1, bindGroupLayouts=(webgpu.WGPUBindGroupLayout*1)(bind_layout))
|
|
|
|
webgpu.wgpuDevicePushErrorScope(self.dev.device_res, webgpu.WGPUErrorFilter_Validation)
|
|
pipeline_layout = webgpu.wgpuDeviceCreatePipelineLayout(self.dev.device_res, pipeline_layout_desc)
|
|
if err := self.dev.pop_error(): raise RuntimeError(f"Error creating pipeline layout: {err}")
|
|
|
|
# Creating bind group
|
|
def bg_entry(n:int, x:webgpu.WGPUBuffer|int|float):
|
|
buf = x if isinstance(x, webgpu.WGPUBuffer) else self.dev.create_uniform(x)
|
|
return webgpu.WGPUBindGroupEntry(binding=n, buffer=buf, offset=0, size=webgpu.wgpuBufferGetSize(buf))
|
|
bindings = (webgpu.WGPUBindGroupEntry * (1+len(bufs)+len(vals)))(bg_entry(0, float('inf')), *(bg_entry(i+1, x) for i,x in enumerate(bufs+vals)))
|
|
|
|
bind_group_desc = webgpu.WGPUBindGroupDescriptor(layout=bind_layout, entryCount=len(bindings), entries=bindings)
|
|
webgpu.wgpuDevicePushErrorScope(self.dev.device_res, webgpu.WGPUErrorFilter_Validation)
|
|
bind_group = webgpu.wgpuDeviceCreateBindGroup(self.dev.device_res, bind_group_desc)
|
|
if err := self.dev.pop_error(): raise RuntimeError(f"Error creating bind group: {err}")
|
|
|
|
# Creating compute pipeline
|
|
compute_desc = webgpu.WGPUComputePipelineDescriptor(layout=pipeline_layout,
|
|
compute=webgpu.WGPUComputeState(module=self.prg, entryPoint=self.name))
|
|
pipeline_result = DeviceCreateComputePipeline(self.dev.device_res, compute_desc)
|
|
|
|
command_encoder = webgpu.wgpuDeviceCreateCommandEncoder(self.dev.device_res, webgpu.WGPUCommandEncoderDescriptor())
|
|
comp_pass_desc = webgpu.WGPUComputePassDescriptor()
|
|
|
|
if wait:
|
|
query_set = webgpu.wgpuDeviceCreateQuerySet(self.dev.device_res, webgpu.WGPUQuerySetDescriptor(type=webgpu.WGPUQueryType_Timestamp, count=2))
|
|
query_buf = webgpu.wgpuDeviceCreateBuffer(
|
|
self.dev.device_res, webgpu.WGPUBufferDescriptor(size=16, usage=webgpu.WGPUBufferUsage_QueryResolve | webgpu.WGPUBufferUsage_CopySrc))
|
|
comp_pass_desc.timestampWrites = c.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(self.dev.queue, 1, (webgpu.WGPUCommandBuffer*1)(cmd_buf))
|
|
|
|
# release created objects
|
|
webgpu.wgpuBindGroupLayoutRelease(bind_layout)
|
|
webgpu.wgpuPipelineLayoutRelease(pipeline_layout)
|
|
webgpu.wgpuBindGroupRelease(bind_group)
|
|
webgpu.wgpuComputePipelineRelease(pipeline_result)
|
|
webgpu.wgpuCommandEncoderRelease(command_encoder)
|
|
webgpu.wgpuComputePassEncoderRelease(compute_pass)
|
|
webgpu.wgpuCommandBufferRelease(cmd_buf)
|
|
|
|
if wait:
|
|
time = ((timestamps:=buf_to_mv(tmp_buf:=self.dev._readable_buffer(query_buf)).cast("Q").tolist())[1] - timestamps[0]) / 1e9
|
|
self.dev.free(query_buf)
|
|
self.dev.free(tmp_buf)
|
|
webgpu.wgpuQuerySetDestroy(query_set)
|
|
webgpu.wgpuQuerySetRelease(query_set)
|
|
return time
|
|
return None
|
|
|
|
class WebGpuAllocator(Allocator['WebGpuDevice']):
|
|
def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
|
|
# WebGPU buffers have to be 4-byte aligned
|
|
return BufferStorage(webgpu.wgpuDeviceCreateBuffer(self.dev.device_res, webgpu.WGPUBufferDescriptor(size=round_up(size, 4),
|
|
usage=webgpu.WGPUBufferUsage_Storage | webgpu.WGPUBufferUsage_CopyDst | webgpu.WGPUBufferUsage_CopySrc)))
|
|
def _copyin(self, dest:webgpu.WGPUBuffer, src:memoryview):
|
|
if src.nbytes % 4:
|
|
padded_src = bytearray(round_up(src.nbytes, 4))
|
|
padded_src[:src.nbytes] = src
|
|
self.dev.write_buffer(dest, padded_src if src.nbytes % 4 else src)
|
|
def _copyout(self, dest:memoryview, src:webgpu.WGPUBuffer):
|
|
dest[:] = buf_to_mv(tmp_buf:=self.dev._readable_buffer(src))[:dest.nbytes]
|
|
self.dev.free(tmp_buf)
|
|
|
|
def _free(self, storage:BufferStorage, options:BufferSpec): self.dev.free(storage.buf)
|
|
|
|
class WebGpuDevice(Compiled):
|
|
def __init__(self, device:str):
|
|
# Requesting an adapter
|
|
adapter_res = InstanceRequestAdapter(instance, webgpu.WGPURequestAdapterOptions(
|
|
powerPreference=webgpu.WGPUPowerPreference_HighPerformance, backendType=backend_types.get(getenv("WEBGPU_BACKEND", ""), 0)))
|
|
|
|
# Get supported features
|
|
webgpu.wgpuAdapterGetFeatures(adapter_res, supported_features:=webgpu.WGPUSupportedFeatures())
|
|
self.features = [feat for i in range(supported_features.featureCount)
|
|
if (feat:=supported_features.features[i]) in [webgpu.WGPUFeatureName_TimestampQuery, webgpu.WGPUFeatureName_ShaderF16]]
|
|
webgpu.wgpuSupportedFeaturesFreeMembers(supported_features)
|
|
dev_desc = webgpu.WGPUDeviceDescriptor(requiredFeatureCount=len(self.features),
|
|
requiredFeatures=(webgpu.WGPUFeatureName * len(self.features))(*self.features))
|
|
|
|
# Limits
|
|
webgpu.wgpuAdapterGetLimits(adapter_res, supported_limits:=webgpu.WGPUSupportedLimits())
|
|
dev_desc.requiredLimits = c.pointer(webgpu.WGPURequiredLimits(limits=supported_limits.limits))
|
|
|
|
# Requesting a device
|
|
self.device_res = AdapterRequestDevice(adapter_res, dev_desc)
|
|
self.queue = webgpu.wgpuDeviceGetQueue(self.device_res)
|
|
|
|
webgpu.wgpuAdapterRelease(adapter_res)
|
|
|
|
super().__init__(device, WebGpuAllocator(self), [WGSLRenderer], WebGPUProgram,
|
|
arch="shader-f16" * (webgpu.WGPUFeatureName_ShaderF16 in self.features))
|
|
|
|
def synchronize(self, timeout:int|None=None): QueueOnSubmittedWorkDone(self.queue)
|
|
|
|
@suppress_finalizing
|
|
def free(self, buf:webgpu.WGPUBuffer):
|
|
if webgpu.wgpuBufferGetMapState(buf) == webgpu.WGPUBufferMapState_Mapped: webgpu.wgpuBufferUnmap(buf)
|
|
webgpu.wgpuBufferDestroy(buf)
|
|
webgpu.wgpuBufferRelease(buf)
|
|
|
|
def pop_error(self) -> str: return DevicePopErrorScope(self.device_res)[1]
|
|
def create_uniform(self, val:int|float) -> webgpu.WGPUBuffer:
|
|
buf = webgpu.wgpuDeviceCreateBuffer(self.device_res,
|
|
webgpu.WGPUBufferDescriptor(size=4, usage=webgpu.WGPUBufferUsage_Uniform | webgpu.WGPUBufferUsage_CopyDst))
|
|
self.write_buffer(buf, val.to_bytes(4, "little", signed=val < 0) if isinstance(val, int) else struct.pack('<f', val))
|
|
return buf
|
|
def _readable_buffer(self, buf:webgpu.WGPUBuffer) -> webgpu.WGPUBuffer:
|
|
size = webgpu.wgpuBufferGetSize(buf)
|
|
ret = webgpu.wgpuDeviceCreateBuffer(self.device_res,
|
|
webgpu.WGPUBufferDescriptor(size=size, usage=webgpu.WGPUBufferUsage_CopyDst | webgpu.WGPUBufferUsage_MapRead, mappedAtCreation=False))
|
|
|
|
# copy_buffer_to_buffer
|
|
encoder = webgpu.wgpuDeviceCreateCommandEncoder(self.device_res, webgpu.WGPUCommandEncoderDescriptor())
|
|
webgpu.wgpuCommandEncoderCopyBufferToBuffer(encoder, buf, 0, ret, 0, size)
|
|
cmd_buf = webgpu.wgpuCommandEncoderFinish(encoder, webgpu.WGPUCommandBufferDescriptor())
|
|
webgpu.wgpuQueueSubmit(self.queue, 1, (webgpu.WGPUCommandBuffer*1)(cmd_buf))
|
|
webgpu.wgpuCommandBufferRelease(cmd_buf)
|
|
webgpu.wgpuCommandEncoderRelease(encoder)
|
|
|
|
return ret
|
|
def write_buffer(self, buf:webgpu.WGPUBuffer, src:memoryview|bytearray|bytes):
|
|
webgpu.wgpuQueueWriteBuffer(self.queue, buf, 0, (ctypes.c_uint8 * len(src)).from_buffer_copy(src), len(src))
|