mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-10-27 16:43:59 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
131 lines
8.8 KiB
Python
131 lines
8.8 KiB
Python
from __future__ import annotations
|
|
from typing import cast
|
|
import ctypes, hashlib
|
|
from tinygrad.runtime.autogen import opencl as cl
|
|
from tinygrad.runtime.support import c
|
|
from tinygrad.helpers import to_char_p_p, from_mv, OSX, DEBUG, suppress_finalizing, unwrap, round_up, is_image_shape
|
|
from tinygrad.renderer.cstyle import OpenCLRenderer
|
|
from tinygrad.device import BufferStorage, BufferSpec, Allocator, Compiled, Compiler, CompileError, TinyELF, Program
|
|
|
|
CC_CB = c.CFUNCTYPE[None, [c.POINTER[ctypes.c_char], c.POINTER[None], cl.size_t, c.POINTER[None]]]
|
|
BP_CB = c.CFUNCTYPE[None, [cl.cl_program, c.POINTER[None]]]
|
|
|
|
# see test/external/external_osx_profiling.py to determine this ratio. it's in like GPU clocks or something
|
|
OSX_TIMING_RATIO = (125/3) if OSX else 1.0
|
|
|
|
cl_errors = {attr: k for k in dir(cl) if k.startswith("CL_") and isinstance(attr:=getattr(cl, k), int) and attr <= 0}
|
|
def check(status):
|
|
if status != 0: raise RuntimeError(f"OpenCL Error {status}: {cl_errors.get(status, 'Unknown error')}")
|
|
def checked(ret, status): return (check(status.value), ret)[1]
|
|
|
|
class CLCompiler(Compiler):
|
|
def __init__(self, dev:CLDevice, compile_key:str):
|
|
self.dev = dev
|
|
super().__init__(f"compile_cl_{compile_key}")
|
|
def compile(self, src:str) -> bytes:
|
|
program = checked(cl.clCreateProgramWithSource(self.dev.context, 1, to_char_p_p([src.encode()]), None, status := ctypes.c_int32()), status)
|
|
build_status: int = cl.clBuildProgram(program, 1, self.dev.cl_dev, None, BP_CB(), None)
|
|
if build_status != 0:
|
|
cl.clGetProgramBuildInfo(program, self.dev.cl_dev, cl.CL_PROGRAM_BUILD_LOG, 0, None, log_size := ctypes.c_size_t())
|
|
cl.clGetProgramBuildInfo(program, self.dev.cl_dev, cl.CL_PROGRAM_BUILD_LOG,
|
|
log_size.value, mstr := ctypes.create_string_buffer(log_size.value), None)
|
|
raise CompileError(f"OpenCL Compile Error\n\n{mstr.value.decode()}")
|
|
check(cl.clGetProgramInfo(program, cl.CL_PROGRAM_BINARY_SIZES, ctypes.sizeof(ctypes.c_size_t), binary_sizes := (ctypes.c_size_t * 1)(), None))
|
|
check(cl.clGetProgramInfo(program, cl.CL_PROGRAM_BINARIES, ctypes.sizeof(ctypes.c_void_p),
|
|
(ctypes.c_void_p * 1)(ctypes.addressof(binary := ctypes.create_string_buffer(binary_sizes[0]))), None))
|
|
check(cl.clReleaseProgram(program))
|
|
return bytes(binary)
|
|
|
|
class CLProgram(Program['CLDevice']):
|
|
def __init__(self, device:CLDevice, obj:TinyELF):
|
|
self.dev, self.lib, self.signature = device, device.cl_compiler.compile_cached(obj.lib.decode()), obj.signature
|
|
self.program = checked(cl.clCreateProgramWithBinary(device.context, 1, device.cl_dev, (ctypes.c_size_t * 1)(len(self.lib)),
|
|
to_char_p_p([self.lib], ctypes.c_ubyte), binary_status := ctypes.c_int32(),
|
|
errcode_ret := ctypes.c_int32()), errcode_ret)
|
|
check(binary_status.value)
|
|
check(cl.clBuildProgram(self.program, 1, device.cl_dev, None, BP_CB(), None)) # NOTE: OSX requires this
|
|
self.kernel = checked(cl.clCreateKernel(self.program, obj.name.encode(), status := ctypes.c_int32()), status)
|
|
|
|
def __del__(self):
|
|
try: check(cl.clReleaseKernel(self.kernel))
|
|
except (TypeError, AttributeError): pass
|
|
try: check(cl.clReleaseProgram(self.program))
|
|
except (TypeError, AttributeError): pass
|
|
|
|
def __call__(self, *bufs:cl.cl_mem, 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:
|
|
for i, (_, slot, dt, shape) in enumerate(self.signature):
|
|
b = bufs[slot] if slot < len(bufs) else getattr(ctypes, f"c_int{dt.bitsize}")(vals[slot-len(bufs)])
|
|
if is_image_shape(shape):
|
|
pitch = (round_up(shape[1], 256) if OSX else shape[1]) * 4 * dt.itemsize
|
|
fmt = cl.cl_image_format(cl.CL_RGBA, {2:cl.CL_HALF_FLOAT, 4:cl.CL_FLOAT}[dt.itemsize])
|
|
desc = cl.cl_image_desc(cl.CL_MEM_OBJECT_IMAGE2D, shape[1], shape[0], image_row_pitch=pitch, buffer=b)
|
|
img = checked(cl.clCreateImage(self.dev.context, cl.CL_MEM_READ_WRITE, fmt, desc, None, status:=ctypes.c_int32()), status)
|
|
check(cl.clSetKernelArg(self.kernel, i, ctypes.sizeof(img), ctypes.byref(img)))
|
|
else: check(cl.clSetKernelArg(self.kernel, i, ctypes.sizeof(b), ctypes.byref(b)))
|
|
global_size = cast(tuple[int,int,int], tuple(int(g*l) for g,l in zip(global_size, local_size)))
|
|
event = cl.cl_event() if wait else None
|
|
check(cl.clEnqueueNDRangeKernel(self.dev.queue, self.kernel, len(global_size), None, (ctypes.c_size_t * len(global_size))(*global_size),
|
|
(ctypes.c_size_t * len(local_size))(*local_size), 0, None, event))
|
|
if wait:
|
|
assert event is not None
|
|
check(cl.clWaitForEvents(1, event))
|
|
check(cl.clGetEventProfilingInfo(event, cl.CL_PROFILING_COMMAND_START, 8, ctypes.byref(start := ctypes.c_uint64()), None))
|
|
check(cl.clGetEventProfilingInfo(event, cl.CL_PROFILING_COMMAND_END, 8, ctypes.byref(end := ctypes.c_uint64()), None))
|
|
return float(end.value-start.value) * OSX_TIMING_RATIO * 1e-9
|
|
return None
|
|
|
|
class CLAllocator(Allocator['CLDevice']):
|
|
def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
|
|
return BufferStorage(checked(cl.clCreateBuffer(self.dev.context, cl.CL_MEM_READ_WRITE, size, None, status := ctypes.c_int32()), status))
|
|
|
|
@suppress_finalizing
|
|
def _free(self, storage:BufferStorage, options:BufferSpec): check(cl.clReleaseMemObject(storage.buf))
|
|
def _copyin(self, dest:cl.cl_mem, src:memoryview):
|
|
self.dev.pending_copyin.append(src:=memoryview(bytearray(src))) # NOTE: these can't be freed until the GPU actually executes this command
|
|
check(cl.clEnqueueWriteBuffer(self.dev.queue, dest, False, 0, len(src)*src.itemsize, from_mv(src), 0, None, None))
|
|
def _copyout(self, dest:memoryview, src:cl.cl_mem):
|
|
check(cl.clEnqueueReadBuffer(self.dev.queue, src, False, 0, len(dest)*dest.itemsize, from_mv(dest), 0, None, None))
|
|
self.dev.synchronize()
|
|
|
|
class CLDevice(Compiled):
|
|
device_ids = None # this is global and only initted once
|
|
def __init__(self, device:str=""):
|
|
if CLDevice.device_ids is None:
|
|
check(cl.clGetPlatformIDs(0, None, num_platforms := ctypes.c_uint32()))
|
|
check(cl.clGetPlatformIDs(num_platforms.value, platform_ids := (cl.cl_platform_id * num_platforms.value)(), None))
|
|
for device_type in [cl.CL_DEVICE_TYPE_GPU, cl.CL_DEVICE_TYPE_DEFAULT]:
|
|
err = cl.clGetDeviceIDs(platform_ids[0], device_type, 0, None, num_devices := ctypes.c_uint32())
|
|
if err == 0 and num_devices.value != 0: break
|
|
if DEBUG >= 1: print(f"CLDevice: got {num_platforms.value} platforms and {num_devices.value} devices")
|
|
CLDevice.device_ids = c.init_c_var((cl.cl_device_id * num_devices.value),
|
|
lambda x: check(cl.clGetDeviceIDs(platform_ids[0], device_type, num_devices, x, None)))
|
|
|
|
self.cl_dev = CLDevice.device_ids[0 if ":" not in device else int(device.split(":")[1])]
|
|
self.device_name = (cl.clGetDeviceInfo(self.cl_dev, cl.CL_DEVICE_NAME, 256,
|
|
buf:=ctypes.create_string_buffer(256), None), buf.value.decode())[1]
|
|
self.driver_version = (cl.clGetDeviceInfo(self.cl_dev, cl.CL_DRIVER_VERSION, 256,
|
|
buf:=ctypes.create_string_buffer(256), None), buf.value.decode())[1]
|
|
if DEBUG >= 1: print(f"CLDevice: opening {self.device_name} with version {self.driver_version}")
|
|
self.context = checked(cl.clCreateContext(None, 1, self.cl_dev, CC_CB(), None, status := ctypes.c_int32()), status)
|
|
self.queue = checked(cl.clCreateCommandQueue(self.context, self.cl_dev, cl.CL_QUEUE_PROFILING_ENABLE, status), status)
|
|
self.pending_copyin: list[memoryview] = []
|
|
check(cl.clGetDeviceInfo(self.cl_dev, cl.CL_DEVICE_EXTENSIONS, 0, None, ctypes.byref(exts_len:=ctypes.c_size_t())))
|
|
self.device_exts = (cl.clGetDeviceInfo(self.cl_dev, cl.CL_DEVICE_EXTENSIONS, exts_len.value,
|
|
ctypes.byref(buf := ctypes.create_string_buffer(exts_len.value)), None),
|
|
ctypes.string_at(buf).decode().split())[1]
|
|
|
|
self.cl_compiler = CLCompiler(self, f"{hashlib.md5(self.device_name.encode() + self.driver_version.encode()).hexdigest()}")
|
|
|
|
arch = ",".join(self.device_exts)
|
|
if "cl_khr_image2d_from_buffer" in self.device_exts:
|
|
check(cl.clGetDeviceInfo(self.cl_dev, cl.CL_DEVICE_IMAGE_PITCH_ALIGNMENT, 4, ctypes.byref(ipa := ctypes.c_uint32()), None))
|
|
arch += f",IMAGE_PITCH_ALIGNMENT={ipa.value}"
|
|
super().__init__(device, CLAllocator(self), [OpenCLRenderer], CLProgram, arch=arch)
|
|
|
|
def count(self) -> int: return len(unwrap(self.device_ids))
|
|
|
|
def synchronize(self, timeout:int|None=None):
|
|
check(cl.clFinish(self.queue))
|
|
self.pending_copyin.clear()
|