tinygrad 0.7.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

tinygrad/runtime/ops_gpu.py

@@ -1,106 +1,103 @@
 from __future__ import annotations
-import pathlib, functools
-import numpy as np
-import pyopencl as cl  # type: ignore
-from typing import Optional, List
-from tinygrad.helpers import DEBUG, getenv, prod, ImageDType, OSX, fromimport
-from tinygrad.ops import Compiled
-from tinygrad.runtime.lib import RawBufferCopyInOut, LRUAllocator, RawBufferTransfer
-from tinygrad.codegen.linearizer import LinearizerOptions
-from tinygrad.renderer.cstyle import uops_to_cstyle, CStyleLanguage
+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.device import BufferOptions, LRUAllocator, Compiled, Compiler, CompileError
 
-OSX_TIMING_RATIO = (125/3) if OSX else 1.0   # see test/external_osx_profiling.py to determine this ratio. it's in like GPU clocks or something
+# 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
 
-# TODO: if you fork and exit the child process after creating anything with cl on AMD, it hangs on e.wait()
-ROCM_LLVM_PATH = pathlib.Path("/opt/rocm/llvm/bin")
-#ROCM_LLVM_PATH = pathlib.Path(__file__).parent.parent.parent.parent / "extra/rocm/build/llvm-project/bin"
-if DEBUG >= 5:
-  early_exec = fromimport("extra.helpers", "enable_early_exec")()
+def check(status):
+  if status != 0: raise RuntimeError(f"OpenCL Error {status}")
+def checked(ret, status): return (check(status.value), ret)[1]
 
-class CLAllocator(LRUAllocator):
-  def _do_alloc(self, size, dtype, device, **kwargs):
-    if isinstance(dtype, ImageDType):
-      # NOTE: the memory is a bit off here due to padding, it's buf.row_pitch * buf.height * 4 * dtype.itemsize
-      assert size == prod(dtype.shape), f"image size mismatch {size} != {dtype.shape}"
-      fmt = cl.ImageFormat(cl.channel_order.RGBA, {2: cl.channel_type.HALF_FLOAT, 4: cl.channel_type.FLOAT}[dtype.itemsize])
-      buf = cl.Image(CL.cl_ctxs[int(device)], cl.mem_flags.READ_WRITE, fmt, shape=(dtype.shape[1], dtype.shape[0]))
-    else:
-      buf = cl.Buffer(CL.cl_ctxs[int(device)], cl.mem_flags.READ_WRITE, size * dtype.itemsize)
-    setattr(buf, 'device', int(device)) # device is tracked on the underlying buffer
-    return buf
-
-class _CL:
-  def __init__(self):
-    cl_platforms = cl.get_platforms()
-    platform_devices: List[List[cl.Device]] = [y for y in ([x.get_devices(device_type=cl.device_type.GPU) for x in cl_platforms] + [x.get_devices(device_type=cl.device_type.CPU) for x in cl_platforms]) if y]
-    self.devices = [device for device in platform_devices[getenv('CL_PLATFORM', 0)] if device.name not in getenv('CL_EXCLUDE', "").split(",")]
-    self.cl_platform = self.devices[0].platform
-  def post_init(self, device=None):
-    self.cl_ctxs: List[cl.Context] = [cl.Context(devices=[x]) for x in self.devices] if device is None else [cl.Context(devices=[self.devices[device]])]
-    if DEBUG >= 1: print(f"using devices: {[ctx.devices[0].hashable_model_and_version_identifier for ctx in self.cl_ctxs]}")
-    self.cl_queue: List[cl.CommandQueue] = [cl.CommandQueue(ctx, device=ctx.devices[0], properties=cl.command_queue_properties.PROFILING_ENABLE) for ctx in self.cl_ctxs]
-    self.cl_allocator = CLAllocator(CL.cl_ctxs[0].devices[0].get_info(cl.device_info.GLOBAL_MEM_SIZE))
-  def synchronize(self):
-    for q in self.cl_queue: q.finish()
-CL = _CL()
-if not getenv("DELAYED_RUNTIME_INIT", False): CL.post_init()
-
-class CLBuffer(RawBufferCopyInOut, RawBufferTransfer):
-  def __init__(self, size, dtype, device='0'): super().__init__(size, dtype, allocator=CL.cl_allocator, **{'device': device})
-  def _copyin(self, x:np.ndarray):
-    assert not self.dtype.name.startswith("image"), f"can't copyin images {self.dtype}"
-    self.event = cl.enqueue_copy(CL.cl_queue[self._buf.device], self._buf, np.require(x, requirements='C'), is_blocking=False)
-  def _copyout(self, x:np.ndarray):
-    assert not self.dtype.name.startswith("image"), f"can't copyout images {self.dtype}"
-    buf = cl.Buffer(CL.cl_ctxs[self._buf.device], cl.mem_flags.WRITE_ONLY | cl.mem_flags.USE_HOST_PTR, 0, hostbuf=x.data)
-    mapped, event = cl.enqueue_map_buffer(CL.cl_queue[self._buf.device], buf, cl.map_flags.WRITE, 0, self.size, dtype=self.dtype.np, is_blocking=False)
-    with mapped.base: cl.enqueue_copy(CL.cl_queue[self._buf.device], mapped, self._buf, is_blocking=True, wait_for=[event] + ([self.event] if hasattr(self, "event") else []))
-  def _transfer(self, x):
-    if "gfx" in CL.cl_ctxs[x._buf.device].devices[0].name:
-      cl.enqueue_copy_buffer_p2p_amd(CL.cl_platform, CL.cl_queue[x._buf.device], x._buf, self._buf, x.size * x.dtype.itemsize).wait()
-    else: raise NotImplementedError("p2p transfer between devices not implemented on non-amd")
+class CLCompiler(Compiler):
+  def __init__(self, device:CLDevice, compile_key:str):
+    self.device = device
+    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)
+    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
+      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
+    check(cl.clReleaseProgram(program))
+    return bytes(binary)
 
 class CLProgram:
-  def __init__(self, name:str, prg:str, binary=False, argdtypes=None, options=None):
-    self.name, self.clprograms = name, [cl.Program(ctx, ctx.devices, [prg]*len(ctx.devices)) if binary else cl.Program(ctx, prg) for ctx in CL.cl_ctxs]  # type: ignore
-    try:
-      self._clprgs = [clprogram.build(options=options) for clprogram in self.clprograms]
-    except cl.RuntimeError as e:
-      if DEBUG >= 3: print("FAILED TO BUILD", prg)
-      raise e
-    self.clprgs = [clprg.__getattr__(name) for clprg in self._clprgs]
-    if DEBUG >= 5 and not OSX:
-      if 'Adreno' in CL.cl_ctxs[0].devices[0].name:
-        fromimport('disassemblers.adreno', 'disasm')(self.binary())
-      elif CL.cl_ctxs[0].devices[0].name.startswith('gfx'):
-        asm = early_exec(([ROCM_LLVM_PATH / "llvm-objdump", '-d', '-'], self.binary()))
-        print('\n'.join([x for x in asm.decode('utf-8').split("\n") if 's_code_end' not in x]))
-      else:
-        # print the PTX for NVIDIA. TODO: probably broken for everything else
-        print(self.binary().decode('utf-8'))
-    if argdtypes is not None: self.set_argdtypes(argdtypes)
-
-  def binary(self): return self.clprograms[0].get_info(cl.program_info.BINARIES)[0]
-  def set_argdtypes(self, argdtypes): self.argdtypes, _ = argdtypes, [clprg.set_scalar_arg_dtypes(argdtypes) for clprg in self.clprgs]
+  def __init__(self, device:CLDevice, name:str, lib:bytes):
+    self.device, 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)
+    check(binary_status.value)
+    check(cl.clBuildProgram(self.program, 1, device.device_id, None, cl.clBuildProgram.argtypes[4](), None)) # NOTE: OSX requires this
+    self.kernel = checked(cl.clCreateKernel(self.program, name.encode(), status := ctypes.c_int32()), status)
 
-  @staticmethod
-  def max_work_group_size(): return CL.cl_ctxs[0].devices[0].max_work_group_size
+  def __del__(self):
+    if hasattr(self, 'kernel'): check(cl.clReleaseKernel(self.kernel))
+    if hasattr(self, 'program'): check(cl.clReleaseProgram(self.program))
 
-  def __call__(self, global_size, local_size, *bufs, wait=False) -> Optional[float]:
-    if not hasattr(self, 'argdtypes'): self.set_argdtypes(tuple(None if isinstance(x, CLBuffer) else np.int32 for x in bufs))
-    cl_bufs = [x._buf if isinstance(x, CLBuffer) else x for x in bufs]
-    e = self.clprgs[cl_bufs[0].device](CL.cl_queue[cl_bufs[0].device], [g*l for g,l in zip(global_size, local_size)] if local_size is not None else global_size, local_size, *cl_bufs, wait_for=[x.event for x in bufs if isinstance(x, CLBuffer) and hasattr(x, "event")])
+  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))
+    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
+    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
     if wait:
-      e.wait()
-      try:
-        return ((e.profile.end - e.profile.start) * OSX_TIMING_RATIO) * 1e-9
-      except cl.RuntimeError:   # no profiling info available
-        return None
+      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
 
-renderer = functools.partial(uops_to_cstyle, CStyleLanguage(
-  kernel_prefix = "__kernel", buffer_prefix = "__global ", smem_prefix = "__local ", arg_int_prefix = "const int",
-  half_prekernel = "#pragma OPENCL EXTENSION cl_khr_fp16 : enable",
-  barrier = "barrier(CLK_LOCAL_MEM_FENCE);", float4 = "(float4)",
-  gid = [f'get_group_id({i})' for i in range(3)], lid = [f'get_local_id({i})' for i in range(3)], uses_vload=True))
-GPUBuffer = Compiled(CLBuffer, LinearizerOptions(), renderer, CLProgram, CL.synchronize)
+class CLAllocator(LRUAllocator):
+  def __init__(self, device:CLDevice):
+    self.device = device
+    super().__init__()
+  def _alloc(self, size:int, options:BufferOptions) -> ctypes._CData:
+    if options.image is not None:
+      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)
+    else: return checked(cl.clCreateBuffer(self.device.context, cl.CL_MEM_READ_WRITE, size, None, status := ctypes.c_int32()), status)
+  def _free(self, buf:ctypes._CData, options:BufferOptions): check(cl.clReleaseMemObject(buf))
+  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))
+    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))
+    self.device.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 = init_c_var((cl.cl_device_id * num_devices.value)(), lambda x: check(cl.clGetDeviceIDs(platform_ids[0], device_type, num_devices, x, None)))  # noqa: E501
+
+    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
+    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] = []
+
+    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))
+  def synchronize(self):
+    check(cl.clFinish(self.queue))
+    self.pending_copyin.clear()
+
+GPUDevice = CLDevice # for legacy reasons

tinygrad/runtime/ops_hsa.py

@@ -0,0 +1,278 @@
+from __future__ import annotations
+import ctypes, functools, subprocess, io, atexit, collections, json
+from typing import Tuple, TypeVar, List, Dict, Any
+import tinygrad.runtime.autogen.hsa as hsa
+from tinygrad.helpers import DEBUG, init_c_var, from_mv, round_up, to_mv, init_c_struct_t, getenv
+from tinygrad.device import Compiled, Compiler, CompileError, BufferOptions, LRUAllocator
+from tinygrad.renderer.cstyle import HIPRenderer
+from tinygrad.runtime.driver.hsa import check, scan_agents, find_memory_pool, AQLQueue
+from tinygrad.runtime.driver.hip_comgr import compile_hip
+if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401
+
+PROFILE = getenv("PROFILE", 0)
+
+class HSAProfiler:
+  def __init__(self):
+    self.tracked_signals = collections.defaultdict(list)
+    self.collected_events: List[Tuple[Any, ...]] = []
+    self.copy_timings = hsa.hsa_amd_profiling_async_copy_time_t()
+    self.disp_timings = hsa.hsa_amd_profiling_dispatch_time_t()
+
+  def track(self, signal, device, name, is_copy=False): self.tracked_signals[device].append((signal, name, is_copy))
+  def process(self, device):
+    # Process all tracked signals, should be called before any of tracked signals are reused.
+    for sig,name,is_copy in self.tracked_signals[device]:
+      if is_copy: check(hsa.hsa_amd_profiling_get_async_copy_time(sig, ctypes.byref(timings :=  self.copy_timings)))
+      else: check(hsa.hsa_amd_profiling_get_dispatch_time(device.agent, sig, ctypes.byref(timings := self.disp_timings))) #type:ignore
+      self.collected_events.append((device.device_id, 1 if is_copy else 0, name, timings.start, timings.end))
+    self.tracked_signals.pop(device)
+
+  def save(self, path):
+    mjson = []
+    for i in range(len(HSADevice.devices)):
+      mjson.append({"name": "process_name", "ph": "M", "pid": i, "args": {"name": "HSA"}})
+      mjson.append({"name": "thread_name", "ph": "M", "pid": i, "tid": 0, "args": {"name": "AQL"}})
+      mjson.append({"name": "thread_name", "ph": "M", "pid": i, "tid": 1, "args": {"name": "SDMA"}})
+
+    for dev_id,queue_id,name,st,et in self.collected_events:
+      mjson.append({"name": name, "ph": "B", "pid": dev_id, "tid": queue_id, "ts": st*1e-3})
+      mjson.append({"name": name, "ph": "E", "pid": dev_id, "tid": queue_id, "ts": et*1e-3})
+    with open(path, "w") as f: f.write(json.dumps({"traceEvents": mjson}))
+    print(f"Saved HSA profile to {path}")
+Profiler = HSAProfiler()
+
+class HSACompiler(Compiler):
+  def __init__(self, arch:str):
+    self.arch = arch
+    super().__init__(f"compile_hip_{self.arch}")
+  def compile(self, src:str) -> bytes:
+    try: return compile_hip(src, self.arch)
+    except RuntimeError as e: raise CompileError(e)
+
+class HSAProgram:
+  def __init__(self, device:HSADevice, name:str, lib:bytes):
+    self.device, self.name, self.lib = device, name, lib
+
+    if DEBUG >= 6:
+      asm = subprocess.check_output(["/opt/rocm/llvm/bin/llvm-objdump", '-d', '-'], input=lib)
+      print('\n'.join([x for x in asm.decode('utf-8').split("\n") if 's_code_end' not in x]))
+
+    self.exec = init_c_var(hsa.hsa_executable_t(), lambda x: check(hsa.hsa_executable_create_alt(hsa.HSA_PROFILE_FULL, hsa.HSA_DEFAULT_FLOAT_ROUNDING_MODE_DEFAULT, None, ctypes.byref(x)))) # noqa: E501
+    self.code_reader = init_c_var(hsa.hsa_code_object_reader_t(),
+                                  lambda x: check(hsa.hsa_code_object_reader_create_from_memory(lib, len(lib), ctypes.byref(x))))
+    check(hsa.hsa_executable_load_agent_code_object(self.exec, self.device.agent, self.code_reader, None, None))
+    check(hsa.hsa_executable_freeze(self.exec, None))
+
+    self.kernel = init_c_var(hsa.hsa_executable_symbol_t(), lambda x: check(hsa.hsa_executable_get_symbol_by_name(self.exec, (name+".kd").encode("utf-8"), ctypes.byref(self.device.agent), ctypes.byref(x)))) # noqa: E501
+    self.handle = init_c_var(ctypes.c_uint64(), lambda x: check(hsa.hsa_executable_symbol_get_info(self.kernel, hsa.HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_OBJECT, ctypes.byref(x)))) # noqa: E501
+    self.kernargs_segment_size = init_c_var(ctypes.c_uint32(), lambda x: check(hsa.hsa_executable_symbol_get_info(self.kernel, hsa.HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_KERNARG_SEGMENT_SIZE, ctypes.byref(x)))).value # noqa: E501
+    self.group_segment_size = init_c_var(ctypes.c_uint32(), lambda x: check(hsa.hsa_executable_symbol_get_info(self.kernel, hsa.HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_GROUP_SEGMENT_SIZE, ctypes.byref(x)))).value # noqa: E501
+    self.private_segment_size = init_c_var(ctypes.c_uint32(), lambda x: check(hsa.hsa_executable_symbol_get_info(self.kernel, hsa.HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_PRIVATE_SEGMENT_SIZE, ctypes.byref(x)))).value # noqa: E501
+
+  def __del__(self):
+    self.device.synchronize()
+    if hasattr(self, 'code_reader'): check(hsa.hsa_code_object_reader_destroy(self.code_reader))
+    if hasattr(self, 'exec'): check(hsa.hsa_executable_destroy(self.exec))
+
+  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):
+    if not hasattr(self, "args_struct_t"):
+      self.args_struct_t = init_c_struct_t(tuple([(f'f{i}', ctypes.c_void_p) for i in range(len(args))] +
+                                                 [(f'v{i}', ctypes.c_int) for i in range(len(vals))]))
+      if ctypes.sizeof(self.args_struct_t) != self.kernargs_segment_size:
+        raise RuntimeError(f"HSAProgram.__call__: incorrect args struct size {ctypes.sizeof(self.args_struct_t)} != {self.kernargs_segment_size}")
+
+    kernargs = None
+    if self.kernargs_segment_size > 0:
+      kernargs = self.device.alloc_kernargs(self.kernargs_segment_size)
+      args_st = self.args_struct_t.from_address(kernargs)
+      for i in range(len(args)): args_st.__setattr__(f'f{i}', args[i])
+      for i in range(len(vals)): args_st.__setattr__(f'v{i}', vals[i])
+      self.device.flush_hdp()
+
+    signal = self.device.alloc_signal(reusable=True) if wait or PROFILE else None
+    self.device.hw_queue.submit_kernel(self, global_size, local_size, kernargs, completion_signal=signal)
+    if PROFILE: Profiler.track(signal, self.device, self.name)
+    if wait:
+      hsa.hsa_signal_wait_scacquire(signal, hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
+      check(hsa.hsa_amd_profiling_get_dispatch_time(self.device.agent, signal, ctypes.byref(timings := hsa.hsa_amd_profiling_dispatch_time_t())))
+      return (timings.end - timings.start) * self.device.clocks_to_time
+
+T = TypeVar("T")
+CHUNK_SIZE, PAGE_SIZE = 256*1024*1024, 0x1000
+class HSAAllocator(LRUAllocator):
+  def __init__(self, device:HSADevice):
+    self.device = device
+    super().__init__()
+
+  def _alloc(self, size:int, options:BufferOptions):
+    if options.host:
+      check(hsa.hsa_amd_memory_pool_allocate(HSADevice.cpu_mempool, size, 0, ctypes.byref(mem := ctypes.c_void_p())))
+      check(hsa.hsa_amd_agents_allow_access(2, (hsa.hsa_agent_t*2)(HSADevice.cpu_agent, self.device.agent), None, mem))
+      return mem.value
+    else:
+      c_agents = (hsa.hsa_agent_t * len(HSADevice.agents[hsa.HSA_DEVICE_TYPE_GPU]))(*HSADevice.agents[hsa.HSA_DEVICE_TYPE_GPU])
+      check(hsa.hsa_amd_memory_pool_allocate(self.device.gpu_mempool, size, 0, ctypes.byref(buf := ctypes.c_void_p())))
+      check(hsa.hsa_amd_agents_allow_access(len(HSADevice.agents[hsa.HSA_DEVICE_TYPE_GPU]), c_agents, None, buf))
+      return buf.value
+
+  def _free(self, opaque:T, options:BufferOptions):
+    HSADevice.synchronize_system()
+    check(hsa.hsa_amd_memory_pool_free(opaque))
+
+  def copyin(self, dest:T, src: memoryview):
+    # Async copyin sync model uses barriers on the main hw queue, since barriers are guaranteed to execute in order with all other packets.
+    self.device.hw_queue.submit_barrier([], sync_signal := self.device.alloc_signal(reusable=True))
+    mem = self._alloc(src.nbytes, BufferOptions(host=True))
+    ctypes.memmove(mem, from_mv(src), src.nbytes)
+    check(hsa.hsa_amd_memory_async_copy_on_engine(dest, self.device.agent, mem, HSADevice.cpu_agent, src.nbytes, 1, ctypes.byref(sync_signal),
+                                                  copy_signal := self.device.alloc_signal(reusable=True), hsa.HSA_AMD_SDMA_ENGINE_0, True))
+    self.device.hw_queue.submit_barrier([copy_signal])
+    self.device.delayed_free.append(mem)
+    if PROFILE: Profiler.track(copy_signal, self.device, f"copyin: CPU -> HSA:{self.device.device_id}", is_copy=True)
+
+  def copy_from_fd(self, dest, fd, offset, size):
+    self.device.hw_queue.submit_barrier([], sync_signal := self.device.alloc_signal(reusable=True))
+
+    if not hasattr(self, 'hb'):
+      self.hb = [self._alloc(CHUNK_SIZE, BufferOptions(host=True)) for _ in range(2)]
+      self.hb_signals = [self.device.alloc_signal(reusable=False) for _ in range(2)]
+      self.hb_polarity = 0
+      self.sdma = [hsa.HSA_AMD_SDMA_ENGINE_0, hsa.HSA_AMD_SDMA_ENGINE_1]
+      for sig in self.hb_signals: hsa.hsa_signal_store_relaxed(sig, 0)
+
+    fo = io.FileIO(fd, "a+b", closefd=False)
+    fo.seek(offset - (minor_offset:=offset % PAGE_SIZE))
+
+    copies_called = 0
+    copied_in = 0
+    for local_offset in range(0, size+minor_offset, CHUNK_SIZE):
+      local_size = min(round_up(size+minor_offset, PAGE_SIZE)-local_offset, CHUNK_SIZE)
+      copy_size = min(local_size-minor_offset, size-copied_in)
+      if copy_size == 0: break
+
+      hsa.hsa_signal_wait_scacquire(self.hb_signals[self.hb_polarity], hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
+      self.device.reusable_signals.append(self.hb_signals[self.hb_polarity]) # it's free now and can be reused
+      self.hb_signals[self.hb_polarity] = self.device.alloc_signal(reusable=False)
+
+      fo.readinto(to_mv(self.hb[self.hb_polarity], local_size))
+      check(hsa.hsa_amd_memory_async_copy_on_engine(dest+copied_in, self.device.agent, self.hb[self.hb_polarity]+minor_offset, HSADevice.cpu_agent,
+                                                    copy_size, 1, ctypes.byref(sync_signal), self.hb_signals[self.hb_polarity],
+                                                    self.sdma[self.hb_polarity], True))
+      copied_in += copy_size
+      self.hb_polarity = (self.hb_polarity + 1) % len(self.hb)
+      minor_offset = 0 # only on the first
+      copies_called += 1
+
+    wait_signals = [self.hb_signals[self.hb_polarity - 1]]
+    if copies_called > 1: wait_signals.append(self.hb_signals[self.hb_polarity])
+    self.device.hw_queue.submit_barrier(wait_signals)
+
+  def copyout(self, dest:memoryview, src:T):
+    HSADevice.synchronize_system()
+    copy_signal = self.device.alloc_signal(reusable=True)
+    c_agents = (hsa.hsa_agent_t*2)(self.device.agent, HSADevice.cpu_agent)
+    check(hsa.hsa_amd_memory_lock_to_pool(from_mv(dest), dest.nbytes, c_agents, 2, HSADevice.cpu_mempool, 0, ctypes.byref(addr:=ctypes.c_void_p())))
+    check(hsa.hsa_amd_memory_async_copy(addr, HSADevice.cpu_agent, src, self.device.agent, dest.nbytes, 0, None, copy_signal))
+    hsa.hsa_signal_wait_scacquire(copy_signal, hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
+    check(hsa.hsa_amd_memory_unlock(from_mv(dest)))
+    if PROFILE: Profiler.track(copy_signal, self.device, f"copyout: HSA:{self.device.device_id} -> CPU", is_copy=True)
+
+  def transfer(self, dest:T, src:T, sz:int, src_dev=None, dest_dev=None):
+    src_dev.hw_queue.submit_barrier([], sync_signal_1 := src_dev.alloc_signal(reusable=True))
+    dest_dev.hw_queue.submit_barrier([], sync_signal_2 := dest_dev.alloc_signal(reusable=True))
+    c_wait_signal = (hsa.hsa_signal_t*2)(sync_signal_1, sync_signal_2)
+    check(hsa.hsa_amd_memory_async_copy_on_engine(dest, dest_dev.agent, src, src_dev.agent, sz, 2, c_wait_signal,
+                                                  copy_signal := dest_dev.alloc_signal(reusable=False), hsa.HSA_AMD_SDMA_ENGINE_0, True))
+    src_dev.hw_queue.submit_barrier([copy_signal])
+    dest_dev.hw_queue.submit_barrier([copy_signal])
+    if PROFILE: Profiler.track(copy_signal, src_dev, f"transfer: HSA:{src_dev.device_id} -> HSA:{dest_dev.device_id}", is_copy=True)
+
+class HSADevice(Compiled):
+  devices: List[HSADevice] = []
+  agents: Dict[int, List[hsa.hsa_agent_t]] = {}
+  cpu_agent: hsa.hsa_agent_t
+  cpu_mempool: hsa.hsa_amd_memory_pool_t
+  def __init__(self, device:str=""):
+    if not HSADevice.agents:
+      check(hsa.hsa_init())
+      atexit.register(hsa_terminate)
+      HSADevice.agents = scan_agents()
+      HSADevice.cpu_agent = HSADevice.agents[hsa.HSA_DEVICE_TYPE_CPU][0]
+      HSADevice.cpu_mempool = find_memory_pool(HSADevice.cpu_agent, segtyp=hsa.HSA_AMD_SEGMENT_GLOBAL, location=hsa.HSA_AMD_MEMORY_POOL_LOCATION_CPU)
+      if PROFILE: check(hsa.hsa_amd_profiling_async_copy_enable(1))
+
+    self.device_id = int(device.split(":")[1]) if ":" in device else 0
+    self.agent = HSADevice.agents[hsa.HSA_DEVICE_TYPE_GPU][self.device_id]
+    self.gpu_mempool = find_memory_pool(self.agent, segtyp=hsa.HSA_AMD_SEGMENT_GLOBAL, location=hsa.HSA_AMD_MEMORY_POOL_LOCATION_GPU)
+    self.hw_queue = AQLQueue(self)
+    HSADevice.devices.append(self)
+
+    check(hsa.hsa_agent_get_info(self.agent, hsa.HSA_AGENT_INFO_NAME, ctypes.byref(agent_name_buf := ctypes.create_string_buffer(256))))
+    self.arch = ctypes.string_at(agent_name_buf).decode()
+
+    check(hsa.hsa_system_get_info(hsa.HSA_SYSTEM_INFO_TIMESTAMP_FREQUENCY, ctypes.byref(gpu_freq := ctypes.c_uint64())))
+    self.clocks_to_time: float = 1 / gpu_freq.value
+
+    check(hsa.hsa_agent_get_info(self.agent, hsa.HSA_AMD_AGENT_INFO_HDP_FLUSH, ctypes.byref(hdp_flush := hsa.hsa_amd_hdp_flush_t())))
+    self.hdp_flush = hdp_flush
+
+    self.delayed_free: List[int] = []
+    self.reusable_signals: List[hsa.hsa_signal_t] = []
+
+    from tinygrad.runtime.graph.hsa import HSAGraph
+    super().__init__(device, HSAAllocator(self), HIPRenderer(), HSACompiler(self.arch), functools.partial(HSAProgram, self), HSAGraph)
+
+    # Finish init: preallocate some signals + space for kernargs
+    self.signal_pool = [init_c_var(hsa.hsa_signal_t(), lambda x: check(hsa.hsa_signal_create(1, 0, None, ctypes.byref(x)))) for _ in range(4096)]
+    self._new_kernargs_region(16 << 20) # initial region size is 16mb
+
+  def synchronize(self):
+    self.hw_queue.wait()
+
+    for sig in self.reusable_signals: hsa.hsa_signal_silent_store_relaxed(sig, 1)
+    self.signal_pool.extend(self.reusable_signals)
+    self.reusable_signals.clear()
+
+    for opaque_to_free in self.delayed_free: check(hsa.hsa_amd_memory_pool_free(opaque_to_free))
+    self.delayed_free.clear()
+
+    self.kernarg_next_addr = self.kernarg_start_addr
+    Profiler.process(self)
+
+  @staticmethod
+  def synchronize_system():
+    for d in HSADevice.devices: d.synchronize()
+
+  def alloc_signal(self, reusable=False):
+    if len(self.signal_pool): signal = self.signal_pool.pop()
+    else: check(hsa.hsa_amd_signal_create(1, 0, None, 0, ctypes.byref(signal := hsa.hsa_signal_t())))
+
+    # reusable means a signal could be reused after synchronize for the device it's allocated from is called.
+    if reusable: self.reusable_signals.append(signal)
+    return signal
+
+  def alloc_kernargs(self, sz):
+    if self.kernarg_next_addr + sz >= self.kernarg_start_addr + self.kernarg_pool_sz: self._new_kernargs_region(int(self.kernarg_pool_sz * 2))
+    result = self.kernarg_next_addr
+    self.kernarg_next_addr = round_up(self.kernarg_next_addr + sz, 16)
+    return result
+
+  def _new_kernargs_region(self, sz:int):
+    if hasattr(self, 'kernarg_start_addr'): self.delayed_free.append(self.kernarg_start_addr)
+    self.kernarg_start_addr: int = self.allocator._alloc(sz, BufferOptions())
+    self.kernarg_next_addr = self.kernarg_start_addr
+    self.kernarg_pool_sz: int = sz
+
+  def flush_hdp(self): self.hdp_flush.HDP_MEM_FLUSH_CNTL[0] = 1
+
+def hsa_terminate():
+  # Need to stop/delete aql queue before hsa shut down, this leads to gpu hangs.
+  for dev in HSADevice.devices:
+    Profiler.process(dev)
+    del dev.hw_queue
+
+  # hsa_shut_down cleans up all hsa-related resources.
+  hsa.hsa_shut_down()
+  HSADevice.synchronize = lambda: None #type:ignore
+  HSAProgram.__del__ = lambda _: None #type:ignore
+  if Profiler.collected_events: Profiler.save("/tmp/profile.json")

tinygrad/runtime/ops_llvm.py

@@ -1,67 +1,46 @@
-import time, hashlib, ctypes
-from typing import ClassVar
-from tinygrad.ops import Compiled
-from tinygrad.helpers import getenv, DEBUG
-from ctypes import CFUNCTYPE
-from tinygrad.codegen.linearizer import LinearizerOptions
-from tinygrad.renderer.llvmir import uops_to_llvm_ir
-from tinygrad.runtime.lib import RawMallocBuffer
+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.renderer.llvmir import LLVMRenderer
+import llvmlite.binding as llvm
+
+class LLVMCompiler(Compiler):
+  def __init__(self, device:LLVMDevice):
+    self.device = device
+    super().__init__("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))
+    return self.device.target_machine.emit_object(mod)
 
-import llvmlite.binding as llvm  # type: ignore
-
-class LLVM:
-  target_machine: ClassVar[llvm.targets.TargetMachine] = None
-  engine: ClassVar[llvm.executionengine.ExecutionEngine] = None
-  optimizer: ClassVar[llvm.passmanagers.ModulePassManager] = None
-
-  def __init__(self):
-    if LLVM.engine is not None: return
+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)
+
+  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)
+
+class LLVMDevice(Compiled):
+  def __init__(self, device:str):
     llvm.initialize()
     llvm.initialize_native_target()
     llvm.initialize_native_asmprinter()
     llvm.initialize_native_asmparser()
-    target = llvm.Target.from_triple(llvm.get_process_triple())
-    LLVM.optimizer = llvm.create_module_pass_manager()
-    LLVM.target_machine = target.create_target_machine(opt=2)  # this opt actually can change things. ex: opt=3 means no FMA, opt=2 means FMA
-    LLVM.target_machine.add_analysis_passes(LLVM.optimizer)
-
-    # TODO: this makes compile times so much faster
-    if getenv("LLVMOPT"):
-      llvm.set_option(str(), '-force-vector-interleave=4')  # this makes sum the same speed as torch, it also doubles the (slow) conv speed
-      if DEBUG >= 4: llvm.set_option(str(), '--debug-only=loop-vectorize')
-      #llvm.set_option(str(), '--debug')
-
-      # does this do anything?
-      builder = llvm.create_pass_manager_builder()
-      builder.opt_level = 3
-      builder.size_level = 0
-      builder.loop_vectorize = True
-      builder.slp_vectorize = True
-      builder.populate(LLVM.optimizer)
-
-    LLVM.target_machine.set_asm_verbosity(True)
+    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()
-    LLVM.engine = llvm.create_mcjit_compiler(backing_mod, LLVM.target_machine)
-
-class LLVMProgram:
-  def __init__(self, name:str, prg:str, binary=False):
-    self.mod = llvm.parse_assembly(prg)
-    self.mod.verify()
-    LLVM().optimizer.run(self.mod)
-    self.mod.name = hashlib.sha1(prg.encode('utf-8')).hexdigest()
-    if DEBUG >= 5: print(LLVM.target_machine.emit_assembly(self.mod))
-    LLVM.engine.add_module(self.mod)
-    LLVM.engine.finalize_object()
-    self.fxn = LLVM.engine.get_function_address(name)
-
-  def __del__(self):
-    if hasattr(self, 'mod'): LLVM.engine.remove_module(self.mod)
-
-  def __call__(self, unused_global_size, unused_local_size, *bufs, wait=False):
-    cfunc = CFUNCTYPE(ctypes.c_int, *[ctypes.c_void_p for _ in bufs])(self.fxn)
-    if wait: st = time.monotonic()
-    cfunc(*[x._buf for x in bufs])
-    if wait: return time.monotonic()-st
-
-LLVMBuffer = Compiled(RawMallocBuffer, LinearizerOptions(supports_float4=False, has_local=False), uops_to_llvm_ir, LLVMProgram)
+    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))