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

tinygrad/runtime/ops_dsp.py

@@ -1,14 +1,64 @@
 from __future__ import annotations
-from typing import Tuple, Any
-import ctypes, os, mmap, tempfile, pathlib, array, functools, threading, contextlib, sys
+import ctypes, os, mmap, tempfile, pathlib, array, functools, threading, contextlib, sys, subprocess, struct
 assert sys.platform != 'win32'
-from tinygrad.device import BufferOptions, Compiled, Allocator
-from tinygrad.helpers import from_mv, getenv, round_up, mv_address, to_mv
-from tinygrad.runtime.ops_clang import ClangCompiler
-from tinygrad.renderer.cstyle import DSPRenderer
+from tinygrad.device import BufferSpec, Compiled, Allocator, Compiler, MallocAllocator
+from tinygrad.dtype import dtypes, DType, PtrDType
+from tinygrad.ops import Ops, UOp
+from tinygrad.helpers import from_mv, getenv, round_up, mv_address, to_mv, cpu_objdump, DEBUG
+from tinygrad.renderer.cstyle import ClangRenderer
 from tinygrad.runtime.autogen import libc, qcom_dsp
 if getenv("IOCTL"): import extra.dsp.run # noqa: F401 # pylint: disable=unused-import
 
+from tinygrad.ops import PatternMatcher, UPat
+
+dsp_pm = PatternMatcher([
+  (((UPat.var('x').maximum(0) ^ -1).maximum(-256) ^ -1).cast(dtypes.uchar.vec(128)),
+   lambda x: UOp(Ops.CUSTOM, dtypes.uchar.vec(128), src=tuple(x.gep(tuple(range(i, i+32))) for i in range(0, 128, 32)),
+     arg="__builtin_HEXAGON_V6_vpackhub_sat_128B(__builtin_HEXAGON_V6_vpackwh_sat_128B({3}, {2}), __builtin_HEXAGON_V6_vpackwh_sat_128B({1}, {0}))")),
+  (UPat(Ops.GEP, name="x"), lambda x: UOp(Ops.CUSTOM, x.dtype, x.src+x.src,
+                      "__builtin_shufflevector({0}, {1}, "+','.join([str(y) for y in x.arg])+")") if len(x.arg) > 1 else None),
+])
+
+dsp_pm_late = PatternMatcher([
+  (UPat.var("x")+UPat(Ops.VECTORIZE, src=UPat.var("y")), lambda x,y: x+UOp(Ops.CUSTOM, x.dtype, (y,), arg="{0}")),
+  (UPat.var("x")*UPat(Ops.VECTORIZE, src=UPat.var("y")), lambda x,y: x*UOp(Ops.CUSTOM, x.dtype, (y,), arg="{0}")),
+  (UPat.var("x")//UPat(Ops.VECTORIZE, src=UPat.var("y")), lambda x,y: x//UOp(Ops.CUSTOM, x.dtype, (y,), arg="{0}")),
+  (UPat(Ops.DEFINE_ACC, src=(UPat(Ops.VECTORIZE, src=UPat(Ops.CONST, arg=0)),), dtype=dtypes.uchar.vec(128), name="d", allow_any_len=True),
+   lambda d: d.replace(src=(UOp(Ops.CUSTOM, d.dtype, arg="__builtin_HEXAGON_V6_vd0_128B()"),)+d.src[1:])),
+])
+
+class DSPRenderer(ClangRenderer):
+  device = "DSP"
+  supports_float4 = True
+  buffer_suffix = " restrict __attribute__((align_value(128)))"
+  kernel_prefix = "__attribute__((noinline)) "
+  pre_matcher = dsp_pm
+  extra_matcher = dsp_pm_late+ClangRenderer.extra_matcher
+  type_map = { **ClangRenderer.type_map, dtypes.uint64: "unsigned long long", dtypes.int64: "long long" }
+  code_for_op = {**ClangRenderer.code_for_op, Ops.SIN: lambda x,dtype: f"__builtin_sin({x})",
+                 Ops.LOG2: lambda x,dtype: f"__builtin_log2l({x})" if dtype == dtypes.float64 else f"__builtin_log2f({x})",
+                 Ops.EXP2: lambda x,dtype: f"__builtin_exp2l({x})" if dtype == dtypes.float64 else f"__builtin_exp2f({x})"}
+
+  def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str,tuple[DType,bool]]], uops:list[UOp], prefix=None) -> str:
+    ret = super().render_kernel(function_name, kernel, bufs, uops, prefix)
+    msrc = ['''/* DSP boilerplate */ struct dcvs_v2_req { int type; int _pad; _Bool dcvs_enable; char dcvs_option; _Bool set_latency; int latency;
+      _Bool set_dcvs_params; short _pad2; char target_corner; char min_corner; char max_corner; int _pad3[3];};''','int HAP_power_set(void*, void*);',
+      'typedef union { struct { void *pv; unsigned int len; } buf; struct { int fd; unsigned int offset; } dma; } remote_arg;',
+      'void* HAP_mmap(void *addr, int len, int prot, int flags, int fd, long offset);', 'int HAP_munmap(void *addr, int len);',
+      'unsigned long long HAP_perf_get_time_us(void);', 'int entry(unsigned long long handle, unsigned int sc, remote_arg* pra) {',
+      'struct dcvs_v2_req req = {.type=7, .dcvs_enable=0, .set_latency=1, .latency=100, .set_dcvs_params=1, .target_corner = 6 /* TURBO */};',
+      'HAP_power_set((void*)handle, (void*)&req);']
+    msrc += ['if ((sc>>24) != 2) return 0;']
+    msrc += [f'int sz_or_val_{i} = ((int*)pra[0].buf.pv)[{i}];' for i,b in enumerate(bufs)]
+    msrc += [f'int off{i} = ((int*)pra[1].buf.pv)[{i}];' for i,b in enumerate(bufs) if isinstance(b[1][0], PtrDType)]
+    msrc += [f'void *buf_{i} = HAP_mmap(0,sz_or_val_{i},3,0,pra[{i+3}].dma.fd,0)+off{i};' for i,b in enumerate(bufs) if isinstance(b[1][0], PtrDType)]
+    msrc += ["unsigned long long start = HAP_perf_get_time_us();"]
+    msrc += [f"{function_name}({', '.join([(f'buf_{i}' if isinstance(b[1][0], PtrDType) else f'sz_or_val_{i}') for i,b in enumerate(bufs)])});"]
+    msrc += ["*(unsigned long long *)(pra[2].buf.pv) = HAP_perf_get_time_us() - start;"]
+    msrc += [f'HAP_munmap(buf_{i}, sz_or_val_{i});' for i,b in enumerate(bufs) if isinstance(b[1][0], PtrDType)]
+    msrc += ["return 0; }"]
+    return ret + '\n' + '\n'.join(msrc)
+
 def rpc_sc(method=0, ins=0, outs=0, fds=0): return (method << 24) | (ins << 16) | (outs << 8) | fds
 def rpc_prep_args(ins=None, outs=None, in_fds=None):
   ins, outs, in_fds = ins or list(), outs or list(), in_fds or list()
@@ -21,65 +71,81 @@ def rpc_prep_args(ins=None, outs=None, in_fds=None):
   return pra, fds, attrs, (ins, outs)
 
 class DSPProgram:
-  def __init__(self, device:DSPDevice, name:str, lib:bytes):
-    self.device, self.lib = device, lib
+  def __init__(self, dev:DSPDevice, name:str, lib:bytes):
+    self.dev, self.lib = dev, lib
 
-  def __call__(self, *bufs, vals:Tuple[int, ...]=(), wait=False):
+  def __call__(self, *bufs, vals:tuple[int, ...]=(), wait=False):
     if len(bufs) >= 16: raise RuntimeError(f"Too many buffers to execute: {len(bufs)}")
 
     pra, fds, attrs, _ = rpc_prep_args(ins=[var_vals_mv:=memoryview(bytearray((len(bufs)+len(vals))*4)), off_mv:=memoryview(bytearray(len(bufs)*4))],
                                        outs=[timer:=memoryview(bytearray(8)).cast('Q')], in_fds=[b.share_info.fd for b in bufs])
     var_vals_mv.cast('i')[:] = array.array('i', tuple(b.size for b in bufs) + vals)
     off_mv.cast('I')[:] = array.array('I', tuple(b.offset for b in bufs))
-    self.device.exec_lib(self.lib, rpc_sc(method=2, ins=2, outs=1, fds=len(bufs)), pra, fds, attrs)
+    self.dev.exec_lib(self.lib, rpc_sc(method=2, ins=2, outs=1, fds=len(bufs)), pra, fds, attrs)
     return timer[0] / 1e6
 
 class DSPBuffer:
-  def __init__(self, va_addr:int, size:int, share_info:Any, offset:int=0):
+  def __init__(self, va_addr:int, size:int, share_info, offset:int=0):
     self.va_addr, self.size, self.share_info, self.offset = va_addr, size, share_info, offset
 
 class DSPAllocator(Allocator):
-  def __init__(self, device:DSPDevice):
-    self.device = device
+  def __init__(self, dev:DSPDevice):
+    self.dev = dev
     super().__init__()
 
-  def _alloc(self, size:int, options:BufferOptions):
-    b = qcom_dsp.ION_IOC_ALLOC(self.device.ion_fd, len=size, align=0x200, heap_id_mask=1<<qcom_dsp.ION_SYSTEM_HEAP_ID, flags=qcom_dsp.ION_FLAG_CACHED)
-    share_info = qcom_dsp.ION_IOC_SHARE(self.device.ion_fd, handle=b.handle)
+  def _alloc(self, size:int, options:BufferSpec):
+    b = qcom_dsp.ION_IOC_ALLOC(self.dev.ion_fd, len=size, align=0x200, heap_id_mask=1<<qcom_dsp.ION_SYSTEM_HEAP_ID, flags=qcom_dsp.ION_FLAG_CACHED)
+    share_info = qcom_dsp.ION_IOC_SHARE(self.dev.ion_fd, handle=b.handle)
     va_addr = libc.mmap(0, size, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED, share_info.fd, 0)
     return DSPBuffer(va_addr, size, share_info, offset=0)
 
-  def _free(self, opaque:DSPBuffer, options:BufferOptions):
-    libc.munmap(opaque.va_addr, opaque.size)
-    os.close(opaque.share_info.fd)
-    qcom_dsp.ION_IOC_FREE(self.device.ion_fd, handle=opaque.share_info.handle)
+  def _free(self, opaque:DSPBuffer, options:BufferSpec):
+    if libc is not None and qcom_dsp is not None:
+      libc.munmap(opaque.va_addr, opaque.size)
+      os.close(opaque.share_info.fd)
+      qcom_dsp.ION_IOC_FREE(self.dev.ion_fd, handle=opaque.share_info.handle)
 
-  def as_buffer(self, src:DSPBuffer) -> memoryview: return to_mv(src.va_addr, src.size)
-  def copyin(self, dest:DSPBuffer, src:memoryview): ctypes.memmove(dest.va_addr, from_mv(src), src.nbytes)
-  def copyout(self, dest:memoryview, src:DSPBuffer): ctypes.memmove(from_mv(dest), src.va_addr, dest.nbytes)
-  def offset(self, buf, size:int, offset:int): return DSPBuffer(buf.va_addr+offset, size, buf.share_info, buf.offset+offset)
+  def _as_buffer(self, src:DSPBuffer) -> memoryview: return to_mv(src.va_addr, src.size)
+  def _copyin(self, dest:DSPBuffer, src:memoryview): ctypes.memmove(dest.va_addr, from_mv(src), src.nbytes)
+  def _copyout(self, dest:memoryview, src:DSPBuffer): ctypes.memmove(from_mv(dest), src.va_addr, dest.nbytes)
+  def _offset(self, buf, size:int, offset:int): return DSPBuffer(buf.va_addr+offset, size, buf.share_info, buf.offset+offset)
 
-class DSPDevice(Compiled):
-  def __init__(self, device:str=""):
-    self.ion_fd = os.open('/dev/ion', os.O_RDONLY)
+class ClangCompiler(Compiler):
+  def __init__(self, cachekey="compile_clang", args:list[str]|None=None, objdump_tool='objdump'):
+    self.args = ['-shared', '-march=native'] if args is None else args
+    self.objdump_tool = objdump_tool
+    super().__init__(cachekey)
 
-    # Generate link script to pass into clang. Aligning all used sections to 4k fixes invoke problem.
-    sections = ['hash', 'text', 'rela.plt', 'got', 'got.plt', 'dynamic', 'dynsym', 'dynstr', 'plt', 'data', 'bss']
-    sections_link = '\n'.join([f'.{n} : ALIGN(4096) {{ *(.{n}) }}' for n in sections])
-    with tempfile.NamedTemporaryFile(delete=False) as self.link_ld:
-      self.link_ld.write(f"SECTIONS {{ . = 0x0; {sections_link}\n /DISCARD/ : {{ *(.note .note.* .gnu.hash .comment) }} }}".encode())
-      self.link_ld.flush()
+  def compile(self, src:str) -> bytes:
+    # TODO: remove file write. sadly clang doesn't like the use of /dev/stdout here
+    with tempfile.NamedTemporaryFile(delete=True) as output_file:
+      subprocess.check_output([getenv("CC", 'clang'), *self.args, '-O2', '-Wall', '-Werror', '-x', 'c', '-fPIC', '-ffreestanding', '-nostdlib',
+                               '-', '-o', str(output_file.name)], input=src.encode('utf-8'))
+      return pathlib.Path(output_file.name).read_bytes()
 
-    compiler_args = ["--target=hexagon", "-mcpu=hexagonv65", "-fuse-ld=lld", "-nostdlib", "-mhvx=v65", "-mhvx-length=128b", f"-T{self.link_ld.name}"]
-    super().__init__(device, DSPAllocator(self), DSPRenderer(),
-                     ClangCompiler("compile_dsp", args=compiler_args, objdump_tool='llvm-objdump'), functools.partial(DSPProgram, self))
+  def disassemble(self, lib:bytes): return cpu_objdump(lib, self.objdump_tool)
 
-    fastrpc_shell = memoryview(bytearray(pathlib.Path('/dsp/cdsp/fastrpc_shell_3').read_bytes()))
-    self.shell_buf = self.allocator.alloc(round_up(fastrpc_shell.nbytes, 0x1000), BufferOptions(nolru=True))
-    ctypes.memmove(self.shell_buf.va_addr, mv_address(fastrpc_shell), fastrpc_shell.nbytes)
+class DSPDevice(Compiled):
+  def __init__(self, device:str=""):
+    compiler_args = ["--target=hexagon", "-mcpu=hexagonv65", "-fuse-ld=lld", "-nostdlib",  "-mhvx=v65", "-mhvx-length=128b"]
+    try:
+      self.ion_fd = os.open('/dev/ion', os.O_RDONLY)
+      # Generate link script to pass into clang. Aligning all used sections to 4k fixes invoke problem.
+      sections = ['hash', 'text', 'rela.plt', 'got', 'got.plt', 'dynamic', 'dynsym', 'dynstr', 'plt', 'data', 'bss']
+      sections_link = '\n'.join([f'.{n} : ALIGN(4096) {{ *(.{n}) }}' for n in sections])
+      with tempfile.NamedTemporaryFile(delete=False) as self.link_ld:
+        self.link_ld.write(f"SECTIONS {{ . = 0x0; {sections_link}\n /DISCARD/ : {{ *(.note .note.* .gnu.hash .comment) }} }}".encode())
+        self.link_ld.flush()
+      super().__init__(device, DSPAllocator(self), DSPRenderer(),
+        ClangCompiler("compile_dsp", ["-shared"] + compiler_args + [f"-T{self.link_ld.name}"], 'llvm-objdump'), functools.partial(DSPProgram, self))
+      fastrpc_shell = memoryview(bytearray(pathlib.Path('/dsp/cdsp/fastrpc_shell_3').read_bytes()))
+      self.shell_buf = self.allocator.alloc(round_up(fastrpc_shell.nbytes, 0x1000), BufferSpec(nolru=True))
+      ctypes.memmove(self.shell_buf.va_addr, mv_address(fastrpc_shell), fastrpc_shell.nbytes)
 
-    self.init_dsp()
-    RPCListner(self).start()
+      self.init_dsp()
+      RPCListener(self).start()
+    except FileNotFoundError:
+      super().__init__(device, MallocAllocator, MockDSPRenderer(), ClangCompiler(None, ["-static"] + compiler_args, 'llvm-objdump'), MockDSPProgram)
 
   def open_lib(self, lib):
     self.binded_lib, self.binded_lib_off = lib, 0
@@ -117,7 +183,7 @@ class DSPDevice(Compiled):
     qcom_dsp.FASTRPC_IOCTL_INIT(self.rpc_fd, flags=0x1, file=self.shell_buf.va_addr, filelen=self.shell_buf.size, filefd=self.shell_buf.share_info.fd)
     qcom_dsp.FASTRPC_IOCTL_INVOKE(self.rpc_fd, handle=3, sc=rpc_sc(method=3, ins=0, outs=0))
 
-class RPCListner(threading.Thread):
+class RPCListener(threading.Thread):
   def __init__(self, device:DSPDevice):
     super().__init__()
     self.device, self.daemon = device, True
@@ -179,3 +245,54 @@ class RPCListner(threading.Thread):
         st = qcom_dsp.FASTRPC_IOCTL_MMAP(self.device.rpc_fd, fd=-1, flags=in_args[0].cast('I')[2], vaddrin=0, size=in_args[0].cast('Q')[3])
         out_args[0].cast('Q')[0:2] = array.array('Q', [0, st.vaddrout])
       else: raise RuntimeError(f"Unknown op: {sc=:X}")
+
+# ***** mock DSP *****
+
+mockdsp_boilerplate = '''/* DSP boilerplate */ static long syscall(long r0, long r1, long r2, long r3, long r4, long r5, long r6) {
+long retval; __asm__ volatile("r0 = %1; r1 = %2; r2 = %3; r3 = %4; r4 = %5; r5 = %6; r6 = %7; trap0(#1); %0 = r0" : "=r" (retval)
+  : "r" (r0), "r" (r1), "r" (r2), "r" (r3), "r" (r4), "r" (r5), "r" (r6) : "r0", "r1", "r2", "r3", "r4", "r5", "r6"); return retval; }
+static int read(int fd, void* buf, int len) {{ return syscall(fd, (long)buf, len, 0, 0, 0, 63); }}
+static int write(int fd, void* buf, int len) {{ return syscall(fd, (long)buf, len, 0, 0, 0, 64); }}
+static int exit(int ret) {{ return syscall(ret, 0, 0, 0, 0, 0, 93); }}
+static unsigned int inscount(void) {{ unsigned int ret; __asm__ volatile(".word 0x6a15c000; %0 = R0" : "=r" (ret) : : "r0"); return ret; }}
+static void *mmap2(void *addr, unsigned int length, int prot, int flags, int fd, unsigned long offset) {{
+return (void*)syscall((long)addr, length, prot, flags, fd, offset, 222); }}'''
+
+class MockDSPRenderer(DSPRenderer):
+  def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str,tuple[DType,bool]]], uops:list[UOp], prefix=None) -> str:
+    ret = ClangRenderer.render_kernel(self, function_name, kernel, bufs, uops, prefix)
+    # https://gpages.juszkiewicz.com.pl/syscalls-table/syscalls.html
+    # control register 21 is HEX_REG_QEMU_INSN_CNT, 0x6a15c000 loads it
+    msrc = [mockdsp_boilerplate, 'void _start(void) {']
+    for i,b in enumerate(bufs):
+      if isinstance(b[1][0], PtrDType):
+        sz = b[1][0].size*b[1][0].itemsize
+        # for loop for big reads
+        msrc.append(f"void *buf{i} = mmap2(0, {sz}, 3, 0x21, -1, 0); for(int rd = 0; rd < {sz}; rd += read(0, buf{i}+rd, {sz}-rd));")
+      else:
+        msrc.append(f"unsigned int val{i}; read(0, &val{i}, 4);")
+    msrc.append("unsigned int st = inscount();")
+    msrc.append(f"{function_name}({', '.join([(f'(void*)buf{i}' if isinstance(b[1][0], PtrDType) else f'val{i}') for i,b in enumerate(bufs)])});")
+    msrc.append("unsigned int et = inscount() - st; write(1, &et, sizeof(et));")
+    for i,b in enumerate(bufs):
+      if isinstance(b[1][0], PtrDType): msrc.append(f"write(1, buf{i}, {b[1][0].size*b[1][0].itemsize});")
+    msrc.append('exit(0); }')
+    return ret + '\n' + '\n'.join(msrc)
+
+class MockDSPProgram:
+  def __init__(self, name:str, lib:bytes): self.lib = lib
+  def __call__(self, *bufs, vals:tuple[int, ...]=(), wait=False):
+    with tempfile.NamedTemporaryFile(suffix=".out") as dsp_lib:
+      dsp_lib.write(self.lib)
+      dsp_lib.flush()
+      os.chmod(dsp_lib.name, 0o0777)
+      # NOTE: this timing includes a docker launch
+      proc = subprocess.run(["docker", "run", "--rm", "-i", "-v", f"{os.path.abspath(os.path.dirname(dsp_lib.name))}:/work", "-w", "/work",
+                            "qemu-hexagon", "-c", f"qemu-hexagon {'-strace' if DEBUG >= 5 else ''} /work/"+os.path.basename(dsp_lib.name)],
+                            input=b''.join([bytes(x) for x in bufs] + [struct.pack("I", x) for x in vals]), stdout=subprocess.PIPE, check=True)
+    offset = 4
+    for x in bufs:
+      x[:] = proc.stdout[offset:offset+len(x)]
+      offset += len(x)
+    assert offset == len(proc.stdout)
+    return struct.unpack("I", proc.stdout[0:4])[0] / 1e9  # pretend it's 1 Ghz, but this is an inscount, not a time

tinygrad/runtime/ops_gpu.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
-from typing import Tuple, Optional, List, cast
-import ctypes, functools, hashlib
+from typing import Optional, cast
+import ctypes, functools, hashlib, contextlib
 from tinygrad.runtime.autogen import opencl as cl
 from tinygrad.helpers import init_c_var, to_char_p_p, from_mv, OSX, DEBUG, getenv, mv_address
 from tinygrad.renderer.cstyle import OpenCLRenderer, IntelRenderer
-from tinygrad.device import BufferOptions, LRUAllocator, Compiled, Compiler, CompileError
+from tinygrad.device import BufferSpec, LRUAllocator, Compiled, Compiler, CompileError
 
 # 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
@@ -15,15 +15,15 @@ def check(status):
 def checked(ret, status): return (check(status.value), ret)[1]
 
 class CLCompiler(Compiler):
-  def __init__(self, device:CLDevice, compile_key:str):
-    self.device = device
+  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.device.context, 1, to_char_p_p([src.encode()]), None, status := ctypes.c_int32()), status)
-    build_status: int = cl.clBuildProgram(program, 1, self.device.device_id, None, cl.clBuildProgram.argtypes[4](), None)
+    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.device_id, None, cl.clBuildProgram.argtypes[4](), None)
     if build_status != 0:
-      cl.clGetProgramBuildInfo(program, self.device.device_id, cl.CL_PROGRAM_BUILD_LOG, 0, None, log_size := ctypes.c_size_t())
-      cl.clGetProgramBuildInfo(program, self.device.device_id, cl.CL_PROGRAM_BUILD_LOG, log_size.value, mstr := ctypes.create_string_buffer(log_size.value), None)  # noqa: E501
+      cl.clGetProgramBuildInfo(program, self.dev.device_id, cl.CL_PROGRAM_BUILD_LOG, 0, None, log_size := ctypes.c_size_t())
+      cl.clGetProgramBuildInfo(program, self.dev.device_id, cl.CL_PROGRAM_BUILD_LOG, log_size.value, mstr := ctypes.create_string_buffer(log_size.value), None)  # noqa: E501
       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))  # noqa: E501
@@ -32,7 +32,7 @@ class CLCompiler(Compiler):
 
 class CLProgram:
   def __init__(self, device:CLDevice, name:str, lib:bytes):
-    self.device, self.name, self.lib = device, name, lib
+    self.dev, self.name, self.lib = device, name, lib
     self.program = checked(cl.clCreateProgramWithBinary(device.context, 1, device.device_id, (ctypes.c_size_t * 1)(len(lib)),
                                                         to_char_p_p([lib], ctypes.c_ubyte), binary_status := ctypes.c_int32(),
                                                         errcode_ret := ctypes.c_int32()), errcode_ret)
@@ -41,15 +41,15 @@ class CLProgram:
     self.kernel = checked(cl.clCreateKernel(self.program, name.encode(), status := ctypes.c_int32()), status)
 
   def __del__(self):
-    if hasattr(self, 'kernel'): check(cl.clReleaseKernel(self.kernel))
-    if hasattr(self, 'program'): check(cl.clReleaseProgram(self.program))
+    with contextlib.suppress(TypeError, AttributeError): check(cl.clReleaseKernel(self.kernel))
+    with contextlib.suppress(TypeError, AttributeError): check(cl.clReleaseProgram(self.program))
 
-  def __call__(self, *bufs:Tuple[ctypes._CData, BufferOptions], global_size:Tuple[int,int,int]=(1,1,1), local_size:Optional[Tuple[int,int,int]]=None, vals:Tuple[int, ...]=(), wait=False) -> Optional[float]:  # noqa: E501
+  def __call__(self, *bufs:tuple[ctypes._CData, BufferSpec], global_size:tuple[int,int,int]=(1,1,1), local_size:Optional[tuple[int,int,int]]=None, vals:tuple[int, ...]=(), wait=False) -> Optional[float]:  # noqa: E501
     for i,(b,_) in enumerate(bufs): cl.clSetKernelArg(self.kernel, i, ctypes.sizeof(b), ctypes.byref(b))
     for i,v in enumerate(vals,start=len(bufs)): cl.clSetKernelArg(self.kernel, i, 4, ctypes.byref(ctypes.c_int32(v)))
-    if local_size is not None: global_size = cast(Tuple[int,int,int], tuple(int(g*l) for g,l in zip(global_size, local_size)))
+    if local_size is not None: 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.device.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) if local_size else None, 0, None, event))  # noqa: E501
+    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) if local_size else None, 0, None, event))  # noqa: E501
     if wait:
       assert event is not None
       check(cl.clWaitForEvents(1, event))
@@ -59,31 +59,31 @@ class CLProgram:
     return None
 
 class CLAllocator(LRUAllocator):
-  def __init__(self, device:CLDevice):
-    self.device = device
+  def __init__(self, dev:CLDevice):
+    self.dev = dev
     super().__init__()
-  def _alloc(self, size:int, options:BufferOptions) -> Tuple[ctypes._CData, BufferOptions]:
+  def _alloc(self, size:int, options:BufferSpec) -> tuple[ctypes._CData, BufferSpec]:
     if options.image is not None:
-      return (checked(cl.clCreateImage2D(self.device.context, cl.CL_MEM_READ_WRITE,
+      return (checked(cl.clCreateImage2D(self.dev.context, cl.CL_MEM_READ_WRITE,
                                         cl.cl_image_format(cl.CL_RGBA, {2: cl.CL_HALF_FLOAT, 4: cl.CL_FLOAT}[options.image.itemsize]),
                                         options.image.shape[1], options.image.shape[0], 0, None, status := ctypes.c_int32()), status), options)
-    return (checked(cl.clCreateBuffer(self.device.context, cl.CL_MEM_READ_WRITE, size, None, status := ctypes.c_int32()), status), options)
-  def _free(self, opaque:Tuple[ctypes._CData, BufferOptions], options:BufferOptions): check(cl.clReleaseMemObject(opaque[0]))
-  def copyin(self, dest:Tuple[ctypes._CData, BufferOptions], src:memoryview):
+    return (checked(cl.clCreateBuffer(self.dev.context, cl.CL_MEM_READ_WRITE, size, None, status := ctypes.c_int32()), status), options)
+  def _free(self, opaque:tuple[ctypes._CData, BufferSpec], options:BufferSpec): check(cl.clReleaseMemObject(opaque[0]))
+  def _copyin(self, dest:tuple[ctypes._CData, BufferSpec], src:memoryview):
     if dest[1].image is not None:
-      check(cl.clEnqueueWriteImage(self.device.queue, dest[0], False, (ctypes.c_size_t * 3)(0,0,0),
+      check(cl.clEnqueueWriteImage(self.dev.queue, dest[0], False, (ctypes.c_size_t * 3)(0,0,0),
                                    (ctypes.c_size_t * 3)(dest[1].image.shape[1],dest[1].image.shape[0],1), 0, 0, from_mv(src), 0, None, None))
     else:
       if mv_address(src) % 16: src = memoryview(bytearray(src))
-      check(cl.clEnqueueWriteBuffer(self.device.queue, dest[0], False, 0, len(src)*src.itemsize, from_mv(src), 0, None, None))
-    self.device.pending_copyin.append(src)    # NOTE: these can't be freed until the GPU actually executes this command
-  def copyout(self, dest:memoryview, src:Tuple[ctypes._CData, BufferOptions]):
+      check(cl.clEnqueueWriteBuffer(self.dev.queue, dest[0], False, 0, len(src)*src.itemsize, from_mv(src), 0, None, None))
+    self.dev.pending_copyin.append(src)    # NOTE: these can't be freed until the GPU actually executes this command
+  def _copyout(self, dest:memoryview, src:tuple[ctypes._CData, BufferSpec]):
     if src[1].image is not None:
-      check(cl.clEnqueueReadImage(self.device.queue, src[0], False, (ctypes.c_size_t * 3)(0,0,0),
+      check(cl.clEnqueueReadImage(self.dev.queue, src[0], False, (ctypes.c_size_t * 3)(0,0,0),
                                   (ctypes.c_size_t * 3)(src[1].image.shape[1],src[1].image.shape[0],1), 0, 0, from_mv(dest), 0, None, None))
     else:
-      check(cl.clEnqueueReadBuffer(self.device.queue, src[0], False, 0, len(dest)*dest.itemsize, from_mv(dest), 0, None, None))
-    self.device.synchronize()
+      check(cl.clEnqueueReadBuffer(self.dev.queue, src[0], 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
@@ -103,7 +103,7 @@ class CLDevice(Compiled):
     if DEBUG >= 1: print(f"CLDevice: opening {self.device_name} with version {self.driver_version}")
     self.context = checked(cl.clCreateContext(None, 1, self.device_id, cl.clCreateContext.argtypes[3](), None, status := ctypes.c_int32()), status)
     self.queue = checked(cl.clCreateCommandQueue(self.context, self.device_id, cl.CL_QUEUE_PROFILING_ENABLE, status), status)
-    self.pending_copyin: List[memoryview] = []
+    self.pending_copyin: list[memoryview] = []
     self.device_exts = (cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_EXTENSIONS, 4096, ctypes.byref(buf := ctypes.create_string_buffer(4096)), ctypes.byref(total := ctypes.c_size_t())), ctypes.string_at(buf, size=total.value).decode())[1]  # noqa: E501
 
     compile_key = hashlib.md5(self.device_name.encode() + self.driver_version.encode()).hexdigest()

tinygrad/runtime/ops_hip.py

@@ -1,8 +1,6 @@
-from __future__ import annotations
 import ctypes, functools
-from typing import Tuple
 from tinygrad.helpers import init_c_var, from_mv, init_c_struct_t, getenv
-from tinygrad.device import Compiled, LRUAllocator, BufferOptions
+from tinygrad.device import Compiled, LRUAllocator, BufferSpec
 from tinygrad.runtime.autogen import hip
 from tinygrad.runtime.support.compiler_hip import AMDCompiler
 from tinygrad.renderer.cstyle import HIPRenderer
@@ -11,18 +9,28 @@ if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint
 def check(status):
   if status != 0: raise RuntimeError(f"HIP Error {status}, {ctypes.string_at(hip.hipGetErrorString(status)).decode()}")
 
+class HIPDevice(Compiled):
+  def __init__(self, device:str=""):
+    self.device_id = int(device.split(":")[1]) if ":" in device else 0
+    self.arch = init_c_var(hip.hipDeviceProp_t(), lambda x: check(hip.hipGetDeviceProperties(x, self.device_id))).gcnArchName.decode()
+    self.time_event_st, self.time_event_en = [init_c_var(hip.hipEvent_t(), lambda x: hip.hipEventCreate(ctypes.byref(x), 0)) for _ in range(2)]
+    super().__init__(device, HIPAllocator(self), HIPRenderer(), AMDCompiler(self.arch), functools.partial(HIPProgram, self))
+  def synchronize(self):
+    check(hip.hipSetDevice(self.device_id))
+    check(hip.hipDeviceSynchronize())
+
 class HIPProgram:
-  def __init__(self, device:HIPDevice, name:str, lib:bytes):
-    self.device, self.name, self.lib = device, name, lib
-    check(hip.hipSetDevice(self.device.device_id))
+  def __init__(self, dev:HIPDevice, name:str, lib:bytes):
+    self.dev, self.name, self.lib = dev, name, lib
+    check(hip.hipSetDevice(self.dev.device_id))
     self.module = init_c_var(hip.hipModule_t(), lambda x: check(hip.hipModuleLoadData(ctypes.byref(x), lib)))
     self.prg = init_c_var(hip.hipFunction_t(), lambda x: check(hip.hipModuleGetFunction(ctypes.byref(x), self.module, name.encode("utf-8"))))
 
   def __del__(self):
     if hasattr(self, 'module'): check(hip.hipModuleUnload(self.module))
 
-  def __call__(self, *args, global_size:Tuple[int,int,int]=(1,1,1), local_size:Tuple[int,int,int]=(1,1,1), vals:Tuple[int, ...]=(), wait=False):
-    check(hip.hipSetDevice(self.device.device_id))
+  def __call__(self, *args, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
+    check(hip.hipSetDevice(self.dev.device_id))
     if not hasattr(self, "vargs"):
       self.c_args = init_c_struct_t(tuple([(f'f{i}', hip.hipDeviceptr_t) for i in range(len(args))] +
                                           [(f'v{i}', ctypes.c_int) for i in range(len(vals))]))(*args, *vals)
@@ -32,37 +40,27 @@ class HIPProgram:
     for i in range(len(args)): self.c_args.__setattr__(f'f{i}', args[i])
     for i in range(len(vals)): self.c_args.__setattr__(f'v{i}', vals[i])
 
-    if wait: check(hip.hipEventRecord(self.device.time_event_st, None))
+    if wait: check(hip.hipEventRecord(self.dev.time_event_st, None))
 
     check(hip.hipModuleLaunchKernel(self.prg, *global_size, *local_size, 0, None, None, self.vargs))
 
     if wait:
-      check(hip.hipEventRecord(self.device.time_event_en, None))
-      check(hip.hipEventSynchronize(self.device.time_event_en))
-      check(hip.hipEventElapsedTime(ctypes.byref(ret := ctypes.c_float()), self.device.time_event_st, self.device.time_event_en))
+      check(hip.hipEventRecord(self.dev.time_event_en, None))
+      check(hip.hipEventSynchronize(self.dev.time_event_en))
+      check(hip.hipEventElapsedTime(ctypes.byref(ret := ctypes.c_float()), self.dev.time_event_st, self.dev.time_event_en))
       return ret.value * 1e-3
 
 class HIPAllocator(LRUAllocator):
-  def __init__(self, device:HIPDevice):
-    self.device = device
+  def __init__(self, dev:HIPDevice):
+    self.dev = dev
     super().__init__()
-  def _alloc(self, size:int, options:BufferOptions):
-    check(hip.hipSetDevice(self.device.device_id))
+  def _alloc(self, size:int, options:BufferSpec):
+    check(hip.hipSetDevice(self.dev.device_id))
     return init_c_var(hip.hipDeviceptr_t(), lambda x: check(hip.hipMalloc(ctypes.byref(x), size)))
-  def _free(self, opaque, options:BufferOptions): check(hip.hipFree(opaque))
-  def copyin(self, dest, src: memoryview):
-    check(hip.hipSetDevice(self.device.device_id))
+  def _free(self, opaque, options:BufferSpec): check(hip.hipFree(opaque))
+  def _copyin(self, dest, src: memoryview):
+    check(hip.hipSetDevice(self.dev.device_id))
     check(hip.hipMemcpy(dest, from_mv(src), len(src), hip.hipMemcpyHostToDevice))
-  def copyout(self, dest:memoryview, src):
-    self.device.synchronize()
+  def _copyout(self, dest:memoryview, src):
+    self.dev.synchronize()
     check(hip.hipMemcpy(from_mv(dest), src, len(dest), hip.hipMemcpyDeviceToHost))
-
-class HIPDevice(Compiled):
-  def __init__(self, device:str=""):
-    self.device_id = int(device.split(":")[1]) if ":" in device else 0
-    self.arch = init_c_var(hip.hipDeviceProp_t(), lambda x: check(hip.hipGetDeviceProperties(x, self.device_id))).gcnArchName.decode()
-    self.time_event_st, self.time_event_en = [init_c_var(hip.hipEvent_t(), lambda x: hip.hipEventCreate(ctypes.byref(x), 0)) for _ in range(2)]
-    super().__init__(device, HIPAllocator(self), HIPRenderer(), AMDCompiler(self.arch), functools.partial(HIPProgram, self))
-  def synchronize(self):
-    check(hip.hipSetDevice(self.device_id))
-    check(hip.hipDeviceSynchronize())

tinygrad/runtime/ops_llvm.py

@@ -1,51 +1,58 @@
-from __future__ import annotations
-import ctypes, functools
-from typing import Tuple
-from tinygrad.device import Compiled, Compiler, MallocAllocator
-from tinygrad.helpers import cpu_time_execution, getenv, cpu_objdump
+import ctypes, platform
+from tinygrad.device import Compiled, Compiler, MallocAllocator, CPUProgram
+from tinygrad.helpers import OSX, getenv, capstone_flatdump, DEBUG
 from tinygrad.renderer.llvmir import LLVMRenderer
-import llvmlite.binding as llvm
+import tinygrad.runtime.autogen.llvm as llvm
+from tinygrad.runtime.support.elf import jit_loader
+
+def cerr(): return ctypes.pointer(ctypes.pointer(ctypes.c_char()))
+
+def expect(x, err, ret=None):
+  if x: raise RuntimeError(llvm.string_cast(err.contents) if not isinstance(err, str) else err)
+  return ret
 
 class LLVMCompiler(Compiler):
-  def __init__(self, device:LLVMDevice, opt:bool=False):
-    self.device = device
-    self.optimizer: llvm.passmanagers.ModulePassManager = llvm.create_module_pass_manager()
-    self.device.target_machine.add_analysis_passes(self.optimizer)
-    if opt:
-      with llvm.create_pass_manager_builder() as builder:
-        builder.opt_level = 3; builder.size_level = 0; builder.loop_vectorize = True; builder.slp_vectorize = True  # noqa: E702
-        builder.populate(self.optimizer)
-    super().__init__("compile_llvm_opt" if opt else "compile_llvm")
+  def __init__(self, host_arch:str):
+    for component in ['Target', 'TargetInfo', 'TargetMC', 'AsmParser', 'AsmPrinter']: getattr(llvm, f'LLVMInitialize{host_arch}{component}')()
 
-  def compile(self, src:str) -> bytes:
-    mod = llvm.parse_assembly(src)
-    mod.verify()
-    self.optimizer.run(mod)
-    return self.device.target_machine.emit_object(mod)
+    triple = {'AArch64': b'aarch64', 'X86': b'x86_64'}[host_arch] + b'-none-unknown-elf'
+    target = expect(llvm.LLVMGetTargetFromTriple(triple, ctypes.pointer(tgt:=llvm.LLVMTargetRef()), err:=cerr()), err, tgt)
+    # +reserve-x18 here does the same thing as -ffixed-x18 in ops_cpu.py, see comments there for why it's needed on arm osx
+    cpu, feats = ctypes.string_at(llvm.LLVMGetHostCPUName()), (b'+reserve-x18,' if OSX else b'') + ctypes.string_at(llvm.LLVMGetHostCPUFeatures())
+    if DEBUG >= 2: print(f"LLVM init for {cpu!r} with {feats!r}")
+    self.target_machine = llvm.LLVMCreateTargetMachine(target, triple, cpu, feats,
+                                                       llvm.LLVMCodeGenLevelDefault, llvm.LLVMRelocPIC, llvm.LLVMCodeModelDefault)
 
-  def disassemble(self, lib:bytes): cpu_objdump(lib)
+    self.pbo = llvm.LLVMCreatePassBuilderOptions()
+    if (opt:=bool(getenv("LLVMOPT", "1"))):
+      self.passes = b'default<O2>'
+      llvm.LLVMPassBuilderOptionsSetLoopUnrolling(self.pbo, True)
+      llvm.LLVMPassBuilderOptionsSetLoopVectorization(self.pbo, True)
+      llvm.LLVMPassBuilderOptionsSetSLPVectorization(self.pbo, True)
+      llvm.LLVMPassBuilderOptionsSetVerifyEach(self.pbo, True)
+    else:
+      self.passes = b'default<O0>'
 
-class LLVMProgram:
-  def __init__(self, device:LLVMDevice, name:str, lib:bytes):
-    self.name, self.lib = name, lib
-    device.engine.add_object_file(llvm.object_file.ObjectFileRef.from_data(lib))
-    self.fxn = device.engine.get_function_address(name)
-    assert self.fxn != 0, "LLVM failed to get function address"
+    super().__init__(f"compile_llvm_jit{'_opt' if opt else ''}")
+
+  def __del__(self): llvm.LLVMDisposePassBuilderOptions(self.pbo)
+
+  def compile(self, src:str) -> bytes:
+    src_buf = llvm.LLVMCreateMemoryBufferWithMemoryRangeCopy(ctypes.create_string_buffer(src_bytes:=src.encode()), len(src_bytes), b'src')
+    mod = expect(llvm.LLVMParseIRInContext(llvm.LLVMGetGlobalContext(), src_buf, ctypes.pointer(m:=llvm.LLVMModuleRef()), err:=cerr()), err, m)
+    expect(llvm.LLVMVerifyModule(mod, llvm.LLVMReturnStatusAction, err:=cerr()), err)
+    expect(llvm.LLVMRunPasses(mod, self.passes, self.target_machine, self.pbo), 'failed to run passes')
+    if DEBUG >= 7: print(ctypes.string_at(llvm.LLVMPrintModuleToString(mod)).decode())
+    obj_buf = expect(llvm.LLVMTargetMachineEmitToMemoryBuffer(self.target_machine, mod, llvm.LLVMObjectFile, err:=cerr(),
+                                                              ctypes.pointer(buf:=llvm.LLVMMemoryBufferRef())), err, buf)
+    llvm.LLVMDisposeModule(mod)
+    obj = ctypes.string_at(llvm.LLVMGetBufferStart(obj_buf), llvm.LLVMGetBufferSize(obj_buf))
+    llvm.LLVMDisposeMemoryBuffer(obj_buf)
+    return jit_loader(obj)
 
-  def __call__(self, *bufs, vals:Tuple[int, ...]=(), wait=False):
-    if not hasattr(self, 'cfunc'):
-      self.cfunc = ctypes.CFUNCTYPE(ctypes.c_int, *([ctypes.c_void_p]*len(bufs)), *([ctypes.c_int32]*len(vals)))(self.fxn)
-    return cpu_time_execution(lambda: self.cfunc(*bufs, *vals), enable=wait)
+  def disassemble(self, lib:bytes): capstone_flatdump(lib)
 
 class LLVMDevice(Compiled):
   def __init__(self, device:str):
-    llvm.initialize()
-    llvm.initialize_native_target()
-    llvm.initialize_native_asmprinter()
-    llvm.initialize_native_asmparser()
-    # this opt actually can change things. ex: opt=3 means no FMA, opt=2 means FMA
-    self.target_machine: llvm.targets.TargetMachine = llvm.Target.from_triple(llvm.get_process_triple()).create_target_machine(opt=2)
-    backing_mod = llvm.parse_assembly(str())
-    backing_mod.triple = llvm.get_process_triple()
-    self.engine: llvm.executionengine.ExecutionEngine = llvm.create_mcjit_compiler(backing_mod, self.target_machine)
-    super().__init__(device, MallocAllocator, LLVMRenderer(), LLVMCompiler(self, getenv("LLVMOPT")), functools.partial(LLVMProgram, self))
+    compiler = LLVMCompiler({'arm64': 'AArch64', 'aarch64': 'AArch64', 'x86_64': 'X86', 'AMD64': 'X86'}[platform.machine()])
+    super().__init__(device, MallocAllocator, LLVMRenderer(), compiler, CPUProgram)