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

tinygrad/runtime/ops_dsp.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+from typing import Tuple, Any
+import ctypes, os, mmap, tempfile, pathlib, array, functools, threading, contextlib, sys
+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.runtime.autogen import libc, qcom_dsp
+if getenv("IOCTL"): import extra.dsp.run # noqa: F401 # pylint: disable=unused-import
+
+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()
+
+  pra = (qcom_dsp.union_remote_arg * (len(ins) + len(outs) + len(in_fds)))()
+  fds = (ctypes.c_int32 * (len(ins) + len(outs) + len(in_fds)))(*([-1] * (len(ins) + len(outs))), *in_fds)
+  attrs = (ctypes.c_uint32 * (len(ins) + len(outs) + len(in_fds)))(*([0] * (len(ins) + len(outs))), *([1] * (len(in_fds))))
+
+  for i, mv in enumerate(ins + outs): pra[i].buf.pv, pra[i].buf.len = mv_address(mv) if mv.nbytes > 0 else 0, mv.nbytes
+  return pra, fds, attrs, (ins, outs)
+
+class DSPProgram:
+  def __init__(self, device:DSPDevice, name:str, lib:bytes):
+    self.device, self.lib = device, lib
+
+  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)
+    return timer[0] / 1e6
+
+class DSPBuffer:
+  def __init__(self, va_addr:int, size:int, share_info:Any, 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
+    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)
+    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 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)
+
+    # 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()
+
+    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))
+
+    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)
+
+    self.init_dsp()
+    RPCListner(self).start()
+
+  def open_lib(self, lib):
+    self.binded_lib, self.binded_lib_off = lib, 0
+    fp = "file:///tinylib?entry&_modver=1.0&_dom=cdsp\0"
+    pra, _, _, _ = rpc_prep_args(ins=[memoryview(array.array('I', [len(fp), 0xff])), memoryview(bytearray(fp.encode()))],
+                                 outs=[o1:=memoryview(bytearray(0x8)), o2:=memoryview(bytearray(0xff))])
+    qcom_dsp.FASTRPC_IOCTL_INVOKE(self.rpc_fd, handle=0, sc=rpc_sc(method=0, ins=2, outs=2), pra=pra)
+    if o1.cast('i')[1] < 0: raise RuntimeError(f"Cannot open lib: {o2.tobytes().decode()}")
+    return o1.cast('I')[0]
+
+  def close_lib(self, handle):
+    pra, _, _, _ = rpc_prep_args(ins=[memoryview(array.array('I', [handle, 0xff]))], outs=[memoryview(bytearray(0x8)), memoryview(bytearray(0xff))])
+    qcom_dsp.FASTRPC_IOCTL_INVOKE(self.rpc_fd, handle=0, sc=rpc_sc(method=1, ins=1, outs=2), pra=pra)
+
+  def exec_lib(self, lib, sc, args, fds, attrs):
+    def _exec_lib():
+      handle = self.open_lib(lib)
+      qcom_dsp.FASTRPC_IOCTL_INVOKE_ATTRS(self.rpc_fd, fds=fds, attrs=attrs, inv=qcom_dsp.struct_fastrpc_ioctl_invoke(handle=handle, sc=sc, pra=args))
+      self.close_lib(handle)
+    try: _exec_lib()
+    except (OSError, PermissionError):
+      # DSP might ask for a connection reset or just fail with operation not permitted, try to reset connection.
+      self.init_dsp()
+      _exec_lib()
+
+  def init_dsp(self):
+    if hasattr(self, 'rpc_fd'):
+      with contextlib.suppress(OSError):
+        qcom_dsp.FASTRPC_IOCTL_INVOKE(self.rpc_fd, handle=4, sc=rpc_sc(method=2, ins=0, outs=0)) # pylint: disable=access-member-before-definition
+      os.close(self.rpc_fd) # pylint: disable=access-member-before-definition
+
+    self.rpc_fd: int = os.open('/dev/adsprpc-smd', os.O_RDONLY | os.O_NONBLOCK)
+    qcom_dsp.FASTRPC_IOCTL_GETINFO(self.rpc_fd, 3)
+    qcom_dsp.FASTRPC_IOCTL_CONTROL(self.rpc_fd, req=0x3)
+    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):
+  def __init__(self, device:DSPDevice):
+    super().__init__()
+    self.device, self.daemon = device, True
+
+  def run(self):
+    # Setup initial request arguments.
+    context, status, TINYFD = 0, 0xffffffff, 0xffff
+    req_args, _, _, _ = rpc_prep_args(ins=[msg_send:=memoryview(bytearray(0x10)).cast('I'), out_buf:=memoryview(bytearray(0x10000)).cast('I')],
+                                      outs=[msg_recv:=memoryview(bytearray(0x10)).cast('I'), in_buf:=memoryview(bytearray(0x10000)).cast('I')])
+    req_args[1].buf.len = 0
+
+    while True:
+      # Update message request and send it.
+      msg_send[:] = array.array('I', [context, status, req_args[1].buf.len, in_buf.nbytes])
+
+      try: qcom_dsp.FASTRPC_IOCTL_INVOKE(self.device.rpc_fd, handle=0x3, sc=0x04020200, pra=req_args)
+      except OSError: continue # retry
+
+      context, inbufs, outbufs = msg_recv[0], ((sc:=msg_recv[2]) >> 16) & 0xff, (msg_recv[2] >> 8) & 0xff
+
+      in_ptr, out_ptr, objs = mv_address(in_buf), mv_address(out_buf), []
+      for i in range(inbufs + outbufs):
+        obj_ptr = round_up(in_ptr + 4, 8) if i < inbufs else round_up(out_ptr + 4, 8)
+        objs.append(to_mv(obj_ptr, obj_size:=to_mv(in_ptr, 4).cast('I')[0]))
+        if i < inbufs: in_ptr = obj_ptr + obj_size
+        else:
+          to_mv(out_ptr, 4).cast('I')[0] = obj_size
+          out_ptr = obj_ptr + obj_size
+          in_ptr += 4
+
+      in_args, out_args = objs[:inbufs], objs[inbufs:]
+      req_args[1].buf.len = out_ptr - mv_address(out_buf)
+
+      status = 0 # reset status, will set if error
+      if sc == 0x20200: pass # greating
+      elif sc == 0x13050100: # open
+        try: out_args[0].cast('I')[0] = TINYFD if (name:=in_args[3].tobytes()[:-1].decode()) == "tinylib" else os.open(name, os.O_RDONLY)
+        except OSError: status = 1
+      elif sc == 0x3010000:
+        if (fd:=in_args[0].cast('I')[0]) != TINYFD: os.close(fd)
+      elif sc == 0x9010000: # seek
+        if (fd:=in_args[0].cast('I')[0]) == TINYFD:
+          assert in_args[0].cast('I')[2] == qcom_dsp.APPS_STD_SEEK_SET, "Supported only SEEK_SET"
+          res, self.device.binded_lib_off = 0, in_args[0].cast('I')[1]
+        else: res = os.lseek(fd, in_args[0].cast('I')[1], in_args[0].cast('I')[2])
+        status = 0 if res >= 0 else res
+      elif sc == 0x4010200: # read
+        if (fd:=in_args[0].cast('I')[0]) == TINYFD:
+          buf = self.device.binded_lib[self.device.binded_lib_off:self.device.binded_lib_off+in_args[0].cast('I')[1]]
+          self.device.binded_lib_off += len(buf)
+        else: buf = os.read(fd, in_args[0].cast('I')[1])
+        out_args[1][:len(buf)] = buf
+        out_args[0].cast('I')[0:2] = array.array('I', [len(buf), int(len(buf) == 0)])
+      elif sc == 0x1f020100: # stat
+        stat = os.stat(in_args[1].tobytes()[:-1].decode())
+        out_stat = qcom_dsp.struct_apps_std_STAT.from_address(mv_address(out_args[0]))
+        for f in out_stat._fields_: out_stat.__setattr__(f[0], int(getattr(stat, f"st_{f[0]}", 0)))
+      elif sc == 0x2010100: # mmap
+        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}")

tinygrad/runtime/ops_gpu.py

@@ -1,16 +1,17 @@
 from __future__ import annotations
 from typing import Tuple, Optional, List, cast
 import ctypes, functools, hashlib
-import tinygrad.runtime.autogen.opencl as cl
-from tinygrad.helpers import init_c_var, to_char_p_p, from_mv, OSX, DEBUG
-from tinygrad.renderer.cstyle import OpenCLRenderer
+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
 
 # 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}")
+  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):
@@ -43,8 +44,8 @@ class CLProgram:
     if hasattr(self, 'kernel'): check(cl.clReleaseKernel(self.kernel))
     if hasattr(self, 'program'): check(cl.clReleaseProgram(self.program))
 
-  def __call__(self, *bufs:ctypes._CData, 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))
+  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
+    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)))
     event = cl.cl_event() if wait else None
@@ -61,18 +62,27 @@ class CLAllocator(LRUAllocator):
   def __init__(self, device:CLDevice):
     self.device = device
     super().__init__()
-  def _alloc(self, size:int, options:BufferOptions) -> ctypes._CData:
+  def _alloc(self, size:int, options:BufferOptions) -> Tuple[ctypes._CData, BufferOptions]:
     if options.image is not None:
-      return checked(cl.clCreateImage2D(self.device.context, cl.CL_MEM_READ_WRITE,
+      return (checked(cl.clCreateImage2D(self.device.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)
-    return checked(cl.clCreateBuffer(self.device.context, cl.CL_MEM_READ_WRITE, size, None, status := ctypes.c_int32()), status)
-  def _free(self, opaque:ctypes._CData, options:BufferOptions): check(cl.clReleaseMemObject(opaque))
-  def copyin(self, dest:ctypes._CData, src:memoryview):
-    check(cl.clEnqueueWriteBuffer(self.device.queue, dest, False, 0, len(src)*src.itemsize, from_mv(src), 0, None, None))
+                                        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):
+    if dest[1].image is not None:
+      check(cl.clEnqueueWriteImage(self.device.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:ctypes._CData):
-    check(cl.clEnqueueReadBuffer(self.device.queue, src, False, 0, len(dest)*dest.itemsize, from_mv(dest), 0, None, None))
+  def copyout(self, dest:memoryview, src:Tuple[ctypes._CData, BufferOptions]):
+    if src[1].image is not None:
+      check(cl.clEnqueueReadImage(self.device.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()
 
 class CLDevice(Compiled):
@@ -90,12 +100,15 @@ class CLDevice(Compiled):
     self.device_id = CLDevice.device_ids[0 if ":" not in device else int(device.split(":")[1])]
     self.device_name = (cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_NAME, 256, buf := ctypes.create_string_buffer(256), None), buf.value.decode())[1]  # noqa: E501
     self.driver_version = (cl.clGetDeviceInfo(self.device_id, cl.CL_DRIVER_VERSION, 256, buf := ctypes.create_string_buffer(256), None), buf.value.decode())[1]  # noqa: E501
+    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.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()
-    super().__init__(device, CLAllocator(self), OpenCLRenderer(), CLCompiler(self, f"compile_cl_{compile_key}"), functools.partial(CLProgram, self))
+    renderer = IntelRenderer() if "cl_intel_subgroup_matrix_multiply_accumulate" in self.device_exts and getenv("INTEL") else OpenCLRenderer()
+    super().__init__(device, CLAllocator(self), renderer, CLCompiler(self, f"compile_cl_{compile_key}"), functools.partial(CLProgram, self))
   def synchronize(self):
     check(cl.clFinish(self.queue))
     self.pending_copyin.clear()

tinygrad/runtime/ops_hip.py

@@ -0,0 +1,68 @@
+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.runtime.autogen import hip
+from tinygrad.runtime.support.compiler_hip import AMDCompiler
+from tinygrad.renderer.cstyle import HIPRenderer
+if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
+
+def check(status):
+  if status != 0: raise RuntimeError(f"HIP Error {status}, {ctypes.string_at(hip.hipGetErrorString(status)).decode()}")
+
+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))
+    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))
+    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)
+      self.vargs = (ctypes.c_void_p * 5)(1, ctypes.cast(ctypes.byref(self.c_args), ctypes.c_void_p), 2,
+                                         ctypes.cast(ctypes.pointer(ctypes.c_size_t(ctypes.sizeof(self.c_args))), ctypes.c_void_p), 3)
+
+    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))
+
+    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))
+      return ret.value * 1e-3
+
+class HIPAllocator(LRUAllocator):
+  def __init__(self, device:HIPDevice):
+    self.device = device
+    super().__init__()
+  def _alloc(self, size:int, options:BufferOptions):
+    check(hip.hipSetDevice(self.device.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))
+    check(hip.hipMemcpy(dest, from_mv(src), len(src), hip.hipMemcpyHostToDevice))
+  def copyout(self, dest:memoryview, src):
+    self.device.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

@@ -2,27 +2,35 @@ from __future__ import annotations
 import ctypes, functools
 from typing import Tuple
 from tinygrad.device import Compiled, Compiler, MallocAllocator
-from tinygrad.helpers import DEBUG, cpu_time_execution, cpu_objdump
+from tinygrad.helpers import cpu_time_execution, getenv, cpu_objdump
 from tinygrad.renderer.llvmir import LLVMRenderer
 import llvmlite.binding as llvm
 
 class LLVMCompiler(Compiler):
-  def __init__(self, device:LLVMDevice):
+  def __init__(self, device:LLVMDevice, opt:bool=False):
     self.device = device
-    super().__init__("compile_llvm")
+    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 compile(self, src:str) -> bytes:
     mod = llvm.parse_assembly(src)
     mod.verify()
-    self.device.optimizer.run(mod)
-    if DEBUG >= 5: print(self.device.target_machine.emit_assembly(mod))
+    self.optimizer.run(mod)
     return self.device.target_machine.emit_object(mod)
 
+  def disassemble(self, lib:bytes): cpu_objdump(lib)
+
 class LLVMProgram:
   def __init__(self, device:LLVMDevice, name:str, lib:bytes):
-    if DEBUG >= 6: cpu_objdump(lib)
     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"
 
   def __call__(self, *bufs, vals:Tuple[int, ...]=(), wait=False):
     if not hasattr(self, 'cfunc'):
@@ -35,12 +43,9 @@ class LLVMDevice(Compiled):
     llvm.initialize_native_target()
     llvm.initialize_native_asmprinter()
     llvm.initialize_native_asmparser()
-    self.optimizer: llvm.passmanagers.ModulePassManager = llvm.create_module_pass_manager()
     # 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)
-    self.target_machine.add_analysis_passes(self.optimizer)
-    self.target_machine.set_asm_verbosity(True)
     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), functools.partial(LLVMProgram, self))
+    super().__init__(device, MallocAllocator, LLVMRenderer(), LLVMCompiler(self, getenv("LLVMOPT")), functools.partial(LLVMProgram, self))