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

tinygrad/runtime/driver/hip_comgr.py

@@ -0,0 +1,47 @@
+import ctypes
+import tinygrad.runtime.autogen.comgr as comgr
+
+def check(status):
+  if status != 0:
+    comgr.amd_comgr_status_string(status, ctypes.byref(status_str := ctypes.POINTER(ctypes.c_char)()))
+    raise RuntimeError(f"comgr fail {status}, {ctypes.string_at(status_str).decode()}")
+
+def _get_comgr_data(data_set, data_type):
+  check(comgr.amd_comgr_action_data_get_data(data_set, data_type, 0, ctypes.byref(data_exec := comgr.amd_comgr_data_t())))
+  check(comgr.amd_comgr_get_data(data_exec, ctypes.byref(sz := ctypes.c_uint64()), None))
+  check(comgr.amd_comgr_get_data(data_exec, ctypes.byref(sz), (dat := ctypes.create_string_buffer(sz.value))))
+  check(comgr.amd_comgr_release_data(data_exec))
+  return bytes(dat)
+
+# AMD_COMGR_SAVE_TEMPS=1 AMD_COMGR_REDIRECT_LOGS=stdout AMD_COMGR_EMIT_VERBOSE_LOGS=1
+def compile_hip(prg:str, arch="gfx1100") -> bytes:
+  check(comgr.amd_comgr_create_action_info(ctypes.byref(action_info := comgr.amd_comgr_action_info_t())))
+  check(comgr.amd_comgr_action_info_set_language(action_info, comgr.AMD_COMGR_LANGUAGE_HIP))
+  check(comgr.amd_comgr_action_info_set_isa_name(action_info, b"amdgcn-amd-amdhsa--" + arch.encode()))
+  check(comgr.amd_comgr_action_info_set_logging(action_info, True))
+
+  check(comgr.amd_comgr_create_data_set(ctypes.byref(data_set_src := comgr.amd_comgr_data_set_t())))
+  check(comgr.amd_comgr_create_data_set(ctypes.byref(data_set_bc := comgr.amd_comgr_data_set_t())))
+  check(comgr.amd_comgr_create_data_set(ctypes.byref(data_set_reloc := comgr.amd_comgr_data_set_t())))
+  check(comgr.amd_comgr_create_data_set(ctypes.byref(data_set_exec := comgr.amd_comgr_data_set_t())))
+
+  check(comgr.amd_comgr_create_data(comgr.AMD_COMGR_DATA_KIND_SOURCE, ctypes.byref(data_src := comgr.amd_comgr_data_t())))
+  check(comgr.amd_comgr_set_data(data_src, len(rprg := prg.encode()), rprg))
+  check(comgr.amd_comgr_set_data_name(data_src, b"<null>"))
+
+  check(comgr.amd_comgr_data_set_add(data_set_src, data_src))
+  # -include hiprtc_runtime.h was removed
+  check(comgr.amd_comgr_action_info_set_options(action_info, f"-O3 -mcumode --hip-version=6.0.32830 -DHIP_VERSION_MAJOR=6 -DHIP_VERSION_MINOR=0 -DHIP_VERSION_PATCH=32830 -D__HIPCC_RTC__ -std=c++14 -nogpuinc -Wno-gnu-line-marker -Wno-missing-prototypes --offload-arch={arch} -I/opt/rocm/include -Xclang -disable-llvm-passes".encode())) # noqa: E501
+  status = comgr.amd_comgr_do_action(comgr.AMD_COMGR_ACTION_COMPILE_SOURCE_WITH_DEVICE_LIBS_TO_BC, action_info, data_set_src, data_set_bc)
+  if status != 0:
+    print(_get_comgr_data(data_set_bc, comgr.AMD_COMGR_DATA_KIND_LOG).decode())
+    raise RuntimeError("compile failed")
+  check(comgr.amd_comgr_action_info_set_options(action_info, b"-O3 -mllvm -amdgpu-internalize-symbols"))
+  check(comgr.amd_comgr_do_action(comgr.AMD_COMGR_ACTION_CODEGEN_BC_TO_RELOCATABLE, action_info, data_set_bc, data_set_reloc))
+  check(comgr.amd_comgr_action_info_set_options(action_info, b""))
+  check(comgr.amd_comgr_do_action(comgr.AMD_COMGR_ACTION_LINK_RELOCATABLE_TO_EXECUTABLE, action_info, data_set_reloc, data_set_exec))
+  ret = _get_comgr_data(data_set_exec, comgr.AMD_COMGR_DATA_KIND_EXECUTABLE)
+  check(comgr.amd_comgr_release_data(data_src))
+  for x in [data_set_src, data_set_bc, data_set_reloc, data_set_exec]: check(comgr.amd_comgr_destroy_data_set(x))
+  check(comgr.amd_comgr_destroy_action_info(action_info))
+  return ret

tinygrad/runtime/driver/hsa.py

@@ -0,0 +1,143 @@
+import ctypes, collections
+import tinygrad.runtime.autogen.hsa as hsa
+from tinygrad.helpers import init_c_var
+
+def check(status):
+  if status != 0:
+    hsa.hsa_status_string(status, ctypes.byref(status_str := ctypes.POINTER(ctypes.c_char)()))
+    raise RuntimeError(f"HSA Error {status}: {ctypes.string_at(status_str).decode()}")
+
+# Precalulated AQL info
+AQL_PACKET_SIZE = ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t)
+EMPTY_SIGNAL = hsa.hsa_signal_t()
+
+DISPATCH_KERNEL_SETUP = 3 << hsa.HSA_KERNEL_DISPATCH_PACKET_SETUP_DIMENSIONS
+DISPATCH_KERNEL_HEADER  = 1 << hsa.HSA_PACKET_HEADER_BARRIER
+DISPATCH_KERNEL_HEADER |= hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCACQUIRE_FENCE_SCOPE
+DISPATCH_KERNEL_HEADER |= hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCRELEASE_FENCE_SCOPE
+DISPATCH_KERNEL_HEADER |= hsa.HSA_PACKET_TYPE_KERNEL_DISPATCH << hsa.HSA_PACKET_HEADER_TYPE
+
+BARRIER_HEADER  = 1 << hsa.HSA_PACKET_HEADER_BARRIER
+BARRIER_HEADER |= hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCACQUIRE_FENCE_SCOPE
+BARRIER_HEADER |= hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCRELEASE_FENCE_SCOPE
+BARRIER_HEADER |= hsa.HSA_PACKET_TYPE_BARRIER_AND << hsa.HSA_PACKET_HEADER_TYPE
+
+class AQLQueue:
+  def __init__(self, device, sz=-1):
+    self.device = device
+
+    check(hsa.hsa_agent_get_info(self.device.agent, hsa.HSA_AGENT_INFO_QUEUE_MAX_SIZE, ctypes.byref(max_queue_size := ctypes.c_uint32())))
+    queue_size = min(max_queue_size.value, sz) if sz != -1 else max_queue_size.value
+
+    null_func = ctypes.CFUNCTYPE(None, hsa.hsa_status_t, ctypes.POINTER(hsa.struct_hsa_queue_s), ctypes.c_void_p)()
+    self.hw_queue = init_c_var(ctypes.POINTER(hsa.hsa_queue_t)(), lambda x: check(
+      hsa.hsa_queue_create(self.device.agent, queue_size, hsa.HSA_QUEUE_TYPE_SINGLE, null_func, None, (1<<32)-1, (1<<32)-1, ctypes.byref(x))))
+
+    self.next_doorbell_index = 0
+    self.queue_base = self.hw_queue.contents.base_address
+    self.queue_size = self.hw_queue.contents.size * AQL_PACKET_SIZE # in bytes
+    self.write_addr = self.queue_base
+    self.write_addr_end = self.queue_base + self.queue_size - 1 # precalc saves some time
+    self.available_packet_slots = self.hw_queue.contents.size
+
+    check(hsa.hsa_amd_queue_set_priority(self.hw_queue, hsa.HSA_AMD_QUEUE_PRIORITY_HIGH))
+    check(hsa.hsa_amd_profiling_set_profiler_enabled(self.hw_queue, 1))
+
+  def __del__(self):
+    if hasattr(self, 'hw_queue'): check(hsa.hsa_queue_destroy(self.hw_queue))
+
+  def submit_kernel(self, prg, global_size, local_size, kernargs, completion_signal=None):
+    if self.available_packet_slots == 0: self._wait_queue()
+
+    packet = hsa.hsa_kernel_dispatch_packet_t.from_address(self.write_addr)
+    packet.workgroup_size_x = local_size[0]
+    packet.workgroup_size_y = local_size[1]
+    packet.workgroup_size_z = local_size[2]
+    packet.reserved0 = 0
+    packet.grid_size_x = global_size[0] * local_size[0]
+    packet.grid_size_y = global_size[1] * local_size[1]
+    packet.grid_size_z = global_size[2] * local_size[2]
+    packet.private_segment_size = prg.private_segment_size
+    packet.group_segment_size = prg.group_segment_size
+    packet.kernel_object = prg.handle
+    packet.kernarg_address = kernargs
+    packet.reserved2 = 0
+    packet.completion_signal = completion_signal if completion_signal else EMPTY_SIGNAL
+    packet.setup = DISPATCH_KERNEL_SETUP
+    packet.header = DISPATCH_KERNEL_HEADER
+    self._submit_packet()
+
+  def submit_barrier(self, wait_signals=None, completion_signal=None):
+    assert wait_signals is None or len(wait_signals) <= 5
+    if self.available_packet_slots == 0: self._wait_queue()
+
+    packet = hsa.hsa_barrier_and_packet_t.from_address(self.write_addr)
+    packet.reserved0 = 0
+    packet.reserved1 = 0
+    for i in range(5):
+      packet.dep_signal[i] = wait_signals[i] if wait_signals and len(wait_signals) > i else EMPTY_SIGNAL
+    packet.reserved2 = 0
+    packet.completion_signal = completion_signal if completion_signal else EMPTY_SIGNAL
+    packet.header = BARRIER_HEADER
+    self._submit_packet()
+
+  def blit_packets(self, packet_addr, packet_cnt):
+    if self.available_packet_slots < packet_cnt: self._wait_queue(packet_cnt)
+
+    tail_blit_packets = min((self.queue_base + self.queue_size - self.write_addr) // AQL_PACKET_SIZE, packet_cnt)
+    rem_packet_cnt = packet_cnt - tail_blit_packets
+    ctypes.memmove(self.write_addr, packet_addr, AQL_PACKET_SIZE * tail_blit_packets)
+    if rem_packet_cnt > 0: ctypes.memmove(self.queue_base, packet_addr + AQL_PACKET_SIZE * tail_blit_packets, AQL_PACKET_SIZE * rem_packet_cnt)
+
+    self._submit_packet(packet_cnt)
+
+  def wait(self):
+    self.submit_barrier([], finish_signal := self.device.alloc_signal(reusable=True))
+    hsa.hsa_signal_wait_scacquire(finish_signal, hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
+    self.available_packet_slots = self.queue_size // AQL_PACKET_SIZE
+
+  def _wait_queue(self, need_packets=1):
+    while self.available_packet_slots < need_packets:
+      rindex = hsa.hsa_queue_load_read_index_relaxed(self.hw_queue)
+      self.available_packet_slots = self.queue_size // AQL_PACKET_SIZE - (self.next_doorbell_index - rindex)
+
+  def _submit_packet(self, cnt=1):
+    self.available_packet_slots -= cnt
+    self.next_doorbell_index += cnt
+    hsa.hsa_queue_store_write_index_relaxed(self.hw_queue, self.next_doorbell_index)
+    hsa.hsa_signal_store_screlease(self.hw_queue.contents.doorbell_signal, self.next_doorbell_index-1)
+
+    self.write_addr += AQL_PACKET_SIZE * cnt
+    if self.write_addr > self.write_addr_end:
+      self.write_addr = self.queue_base + (self.write_addr - self.queue_base) % self.queue_size
+
+def scan_agents():
+  agents = collections.defaultdict(list)
+
+  @ctypes.CFUNCTYPE(hsa.hsa_status_t, hsa.hsa_agent_t, ctypes.c_void_p)
+  def __scan_agents(agent, data):
+    status = hsa.hsa_agent_get_info(agent, hsa.HSA_AGENT_INFO_DEVICE, ctypes.byref(device_type := hsa.hsa_device_type_t()))
+    if status == 0: agents[device_type.value].append(agent)
+    return hsa.HSA_STATUS_SUCCESS
+
+  hsa.hsa_iterate_agents(__scan_agents, None)
+  return agents
+
+def find_memory_pool(agent, segtyp=-1, location=-1):
+  @ctypes.CFUNCTYPE(hsa.hsa_status_t, hsa.hsa_amd_memory_pool_t, ctypes.c_void_p)
+  def __filter_amd_memory_pools(mem_pool, data):
+    check(hsa.hsa_amd_memory_pool_get_info(mem_pool, hsa.HSA_AMD_MEMORY_POOL_INFO_SEGMENT, ctypes.byref(segment := hsa.hsa_amd_segment_t())))
+    if segtyp >= 0 and segment.value != segtyp: return hsa.HSA_STATUS_SUCCESS
+
+    check(hsa.hsa_amd_memory_pool_get_info(mem_pool, hsa.HSA_AMD_MEMORY_POOL_INFO_LOCATION, ctypes.byref(loc:=hsa.hsa_amd_memory_pool_location_t())))
+    if location >= 0 and loc.value != location: return hsa.HSA_STATUS_SUCCESS
+
+    check(hsa.hsa_amd_memory_pool_get_info(mem_pool, hsa.HSA_AMD_MEMORY_POOL_INFO_SIZE, ctypes.byref(sz := ctypes.c_size_t())))
+    if sz.value == 0: return hsa.HSA_STATUS_SUCCESS
+
+    ret = ctypes.cast(data, ctypes.POINTER(hsa.hsa_amd_memory_pool_t))
+    ret[0] = mem_pool
+    return hsa.HSA_STATUS_INFO_BREAK
+
+  hsa.hsa_amd_agent_iterate_memory_pools(agent, __filter_amd_memory_pools, ctypes.byref(region := hsa.hsa_amd_memory_pool_t()))
+  return region

tinygrad/runtime/graph/clang.py

@@ -0,0 +1,38 @@
+from typing import List, Dict, cast
+import ctypes
+from tinygrad.helpers import dedup, cpu_time_execution, GraphException, DEBUG
+from tinygrad.engine.jit import GraphRunner
+from tinygrad.device import Buffer, Device
+from tinygrad.engine.realize import ExecItem, CompiledRunner
+from tinygrad.shape.symbolic import Variable
+from tinygrad.runtime.ops_clang import ClangProgram
+from tinygrad.renderer.cstyle import ClangRenderer
+render_dtype = ClangRenderer().render_dtype
+
+class ClangGraph(GraphRunner):
+  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
+    super().__init__(jit_cache, input_rawbuffers, var_vals)
+    if not all(isinstance(ji.prg, CompiledRunner) for ji in jit_cache): raise GraphException
+
+    prgs = '\n'.join(dedup([cast(CompiledRunner, ji.prg).p.src for ji in jit_cache]))
+    args = [f"{render_dtype(x.dtype)}* arg{i}" for i,x in enumerate(input_rawbuffers)]
+    args += [f"int {v.expr}" for v in var_vals]
+    code = ["void batched("+','.join(args)+") {"]
+    for ji in jit_cache:
+      args = []
+      for buf in ji.bufs:
+        assert buf is not None
+        if buf in input_rawbuffers:
+          args.append(f"arg{input_rawbuffers.index(buf)}")
+        else:
+          args.append(f"({render_dtype(buf.dtype)}*)0x{ctypes.addressof(buf._buf):X}")
+      args += [x.expr for x in cast(CompiledRunner, ji.prg).p.vars]
+      code.append(f"  {cast(CompiledRunner, ji.prg).p.function_name}({','.join(args)});")
+    code.append("}")
+    if DEBUG >= 4: print("\n".join(code))
+    compiler = Device["CLANG"].compiler
+    assert compiler is not None
+    self.clprg = ClangProgram("batched", compiler.compile(prgs+"\n"+"\n".join(code))) # no point in caching the pointers
+
+  def __call__(self, rawbufs: List[Buffer], var_vals: Dict[Variable, int], wait=False):
+    return cpu_time_execution(lambda: self.clprg(*[x._buf for x in rawbufs], *[x for x in var_vals.values()]), enable=wait)

tinygrad/runtime/graph/cuda.py

@@ -0,0 +1,81 @@
+import ctypes
+from typing import Any, Optional, Tuple, Dict, List, cast
+import tinygrad.runtime.autogen.cuda as cuda
+from tinygrad.helpers import init_c_var, GraphException
+from tinygrad.device import Buffer, Device
+from tinygrad.runtime.ops_cuda import CUDADevice, check, encode_args, cu_time_execution
+from tinygrad.shape.symbolic import Variable
+from tinygrad.engine.realize import ExecItem, BufferXfer, CompiledRunner
+from tinygrad.engine.jit import MultiGraphRunner
+
+class CUDAGraph(MultiGraphRunner):
+  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
+    super().__init__(jit_cache, input_rawbuffers, var_vals)
+
+    # Check all jit items are compatible.
+    if not all(isinstance(ji.prg, (CompiledRunner, BufferXfer)) for ji in jit_cache): raise GraphException
+
+    self.jc_idx_with_updatable_rawbufs = list(set([x[0] for x in self.input_replace.keys()]))
+    self.updatable_nodes: Dict[int, Tuple[Any, Any, Any, bool]] = {} # Dict[jc index] = tuple(graph node, node params, input kernel params, is memcpy)
+
+    self.graph = init_c_var(cuda.CUgraph(), lambda x: check(cuda.cuGraphCreate(ctypes.byref(x), 0)))
+
+    for j,ji in enumerate(self.jit_cache):
+      if isinstance(ji.prg, CompiledRunner):
+        global_size, local_size = ji.prg.p.launch_dims(var_vals)
+
+        new_node = cuda.CUgraphNode()
+        deps = self._access_resources([x.base for x in ji.bufs[ji.prg.p.outcount:] if x is not None],
+                                      [x.base for x in ji.bufs[:ji.prg.p.outcount] if x is not None], new_dependency=new_node)
+        c_deps = (cuda.CUgraphNode*len(deps))(*deps) if deps else None
+
+        c_args, vargs = encode_args([cast(Buffer, x)._buf for x in ji.bufs], [var_vals[x] for x in ji.prg.p.vars])
+        kern_params = cuda.CUDA_KERNEL_NODE_PARAMS(ji.prg.clprg.prg, *global_size, *local_size, 0, None, vargs)
+        check(cuda.cuGraphAddKernelNode(ctypes.byref(new_node), self.graph, c_deps, len(deps), ctypes.byref(kern_params)))
+
+        if j in self.jc_idx_with_updatable_launch_dims or j in self.jc_idx_with_updatable_var_vals or j in self.jc_idx_with_updatable_rawbufs:
+          self.updatable_nodes[j] = (new_node, kern_params, c_args, False)
+      elif isinstance(ji.prg, BufferXfer):
+        dest, src = [cast(Buffer, x) for x in ji.bufs[0:2]]
+        src_dev = cast(CUDADevice, Device[src.device])
+        node_from = cuda.CUgraphNode()
+        deps = self._access_resources(read=[src.base], write=[dest.base], new_dependency=node_from)
+        c_deps = (cuda.CUgraphNode*len(deps))(*deps) if deps else None
+        cp_params = cuda.CUDA_MEMCPY3D_v2(srcMemoryType=cuda.CU_MEMORYTYPE_DEVICE, srcDevice=src._buf, srcPitch=src.nbytes, srcHeight=1,
+                                          dstMemoryType=cuda.CU_MEMORYTYPE_DEVICE, dstDevice=dest._buf, dstPitch=dest.nbytes, dstHeight=1,
+                                          WidthInBytes=dest.nbytes, Height=1, Depth=1)
+        check(cuda.cuGraphAddMemcpyNode(ctypes.byref(node_from), self.graph, c_deps, len(deps), ctypes.byref(cp_params), src_dev.context))
+        if j in self.jc_idx_with_updatable_rawbufs: self.updatable_nodes[j] = (node_from, cp_params, src_dev.context, True)
+
+    self.instance = init_c_var(cuda.CUgraphExec(), lambda x: check(cuda.cuGraphInstantiate_v2(ctypes.byref(x), self.graph, None, None, 0)))
+
+  def __call__(self, input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int], wait=False) -> Optional[float]:
+    # Update rawbuffers in the c_args struct.
+    for (j,i),input_idx in self.input_replace.items():
+      if not self.updatable_nodes[j][3]: setattr(self.updatable_nodes[j][2], f'f{i}', input_rawbuffers[input_idx]._buf)
+      else:
+        if i == 0: self.updatable_nodes[j][1].destDevice = input_rawbuffers[input_idx]._buf
+        elif i == 1: self.updatable_nodes[j][1].srcDevice = input_rawbuffers[input_idx]._buf
+
+    # Update var_vals in the c_args struct.
+    for j in self.jc_idx_with_updatable_var_vals:
+      for i,v in enumerate(cast(CompiledRunner, self.jit_cache[j].prg).p.vars):
+        setattr(self.updatable_nodes[j][2], f'v{i}', var_vals[v])
+
+    # Update launch dims in the kern_params struct.
+    for j in self.jc_idx_with_updatable_launch_dims:
+      self.set_kernel_node_launch_dims(self.updatable_nodes[j][1], *cast(CompiledRunner, self.jit_cache[j].prg).p.launch_dims(var_vals))
+
+    # Update graph nodes with the updated structs.
+    for node, c_node_params, c_args, is_copy in self.updatable_nodes.values():
+      if not is_copy: check(cuda.cuGraphExecKernelNodeSetParams(self.instance, node, ctypes.byref(c_node_params)))
+      else: check(cuda.cuGraphExecMemcpyNodeSetParams(self.instance, node, ctypes.byref(c_node_params), c_args))
+
+    return cu_time_execution(lambda: check(cuda.cuGraphLaunch(self.instance, None)), enable=wait)
+
+  def __del__(self):
+    if hasattr(self, 'graph'): check(cuda.cuGraphDestroy(self.graph))
+    if hasattr(self, 'instance'): check(cuda.cuGraphExecDestroy(self.instance))
+
+  def set_kernel_node_launch_dims(self, node, global_size: Tuple[int, int, int], local_size: Tuple[int, int, int]):
+    node.blockDimX, node.blockDimY, node.blockDimZ, node.gridDimX, node.gridDimY, node.gridDimZ = *local_size, *global_size

tinygrad/runtime/graph/hcq.py

@@ -0,0 +1,143 @@
+import ctypes, collections, array, time
+from typing import List, Any, Dict, cast, Optional, Tuple, Set
+from tinygrad.helpers import GraphException, round_up, to_mv, init_c_struct_t
+from tinygrad.device import Buffer, BufferOptions, Compiled, Device
+from tinygrad.shape.symbolic import Variable
+from tinygrad.engine.realize import ExecItem, BufferXfer, CompiledRunner
+from tinygrad.engine.jit import MultiGraphRunner
+
+class HCQGraph(MultiGraphRunner):
+  def __init__(self, device_t, comp_hcq_t, copy_hcq_t, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
+    super().__init__(jit_cache, input_rawbuffers, var_vals)
+    self.device_t, self.comp_hcq_t, self.copy_hcq_t = device_t, comp_hcq_t, copy_hcq_t
+
+    # Check all jit items are compatible.
+    self.devices = list(set(cast(self.device_t, d) for ji in jit_cache for d in [Device[cast(Buffer, x).device] for x in ji.bufs])) #type: ignore
+    if any(not isinstance(d, self.device_t) for d in self.devices): raise GraphException
+
+    # Allocate kernel args.
+    kernargs_size: Dict[Compiled, int] = collections.defaultdict(int)
+    for ji in self.jit_cache:
+      if not isinstance(ji.prg, CompiledRunner): continue
+      kernargs_size[ji.prg.device] += round_up(ji.prg.clprg.kernargs_segment_size, 16)
+    kernargs_ptrs: Dict[Compiled, int] = {dev:dev.allocator._alloc(sz, BufferOptions(cpu_access=True)).va_addr for dev,sz in kernargs_size.items()}
+
+    # Fill initial arguments.
+    self.kargs_addrs: Dict[int, int] = {}
+    self.ji_kargs_structs: Dict[int, ctypes.Structure] = {}
+    for j,ji in enumerate(self.jit_cache):
+      if not isinstance(ji.prg, CompiledRunner): continue
+      self.kargs_addrs[j] = kernargs_ptrs[ji.prg.device]
+      kernargs_ptrs[ji.prg.device] += round_up(ji.prg.clprg.kernargs_segment_size, 16)
+
+      args_t = init_c_struct_t(tuple([(f'f{i}', ctypes.c_void_p) for i in range(len(ji.bufs))] +
+                                     [(f'v{i}', ctypes.c_int) for i in range(len(ji.prg.p.vars))]))
+      self.ji_kargs_structs[j] = args_t.from_address(self.kargs_addrs[j] + ji.prg.clprg.kernargs_offset)
+      for i in range(len(ji.bufs)): self.ji_kargs_structs[j].__setattr__(f'f{i}', cast(Buffer, ji.bufs[i])._buf.va_addr)
+      for i in range(len(ji.prg.p.vars)): self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[ji.prg.p.vars[i]])
+
+      # NV needs constbuffer to be set
+      if ji.prg.device.dname.startswith("NV"): to_mv(self.kargs_addrs[j], 0x160).cast('I')[:] = array.array('I', ji.prg.clprg.constbuffer_0)
+
+    # Build queues.
+    self.comp_queues: Dict[Compiled, Any] = collections.defaultdict(self.comp_hcq_t)
+    self.comp_signal = {dev: dev._get_signal(value=0) for dev in self.devices}
+    self.comp_signal_val = {dev: 0 for dev in self.devices}
+
+    self.copy_queues: Dict[Compiled, Any] = collections.defaultdict(self.copy_hcq_t)
+    self.copy_signal = {dev: dev._get_signal(value=0) for dev in self.devices}
+    self.copy_signal_val = {dev: 0 for dev in self.devices}
+
+    self.kickoff_signal = self.devices[0]._get_signal(value=0)
+    self.kickoff_value = 0
+    self.graph_timeline = {dev: 0 for dev in self.devices}
+
+    self.exec_ptrs: Dict[int, Tuple[Any, int]] = {}
+    self.copy_to_devs: Dict[Compiled, Set[Compiled]] = {dev: set() for dev in self.devices}
+
+    for j,ji in enumerate(self.jit_cache):
+      if isinstance(ji.prg, CompiledRunner):
+        exec_params = {}
+        deps = self.access_resources(ji.bufs[(outs:=ji.prg.p.outcount):], ji.bufs[:outs], (self.comp_signal[ji.prg.device], sig_val:=j+1))
+        deps = [x for x in deps if id(x[0]) != id(self.comp_signal[ji.prg.device])]
+
+        # On NV, to synchronize kernel execution, we must either issue a wait or chain executions to schedule them in order.
+        # Chaining executions is preferred when possible, as it is faster.
+        if ji.prg.device.dname.startswith("NV"):
+          if len(deps) == 0 and self.comp_signal_val[ji.prg.device] > 0:
+            exec_params['chain_exec_ptr'] = self.exec_ptrs[self.comp_signal_val[ji.prg.device] - 1][1]
+          else: deps.append((self.comp_signal[ji.prg.device], self.comp_signal_val[ji.prg.device]))
+
+        for sig, val in deps: self.comp_queues[ji.prg.device].wait(sig, val)
+
+        self.exec_ptrs[j] = (self.comp_queues[ji.prg.device], self.comp_queues[ji.prg.device].ptr())
+        self.comp_queues[ji.prg.device].exec(ji.prg.clprg, self.kargs_addrs[j], *ji.prg.p.launch_dims(var_vals),
+                                             signal=self.comp_signal[ji.prg.device], signal_value=sig_val, **exec_params)
+        self.comp_signal_val[ji.prg.device] = sig_val
+      elif isinstance(ji.prg, BufferXfer):
+        dest, src = [cast(Buffer, x) for x in ji.bufs[0:2]]
+        Device[src.device]._gpu_map(dest._buf) #type: ignore
+
+        deps = self.access_resources([src], [dest], (self.copy_signal[Device[src.device]], sig_val:=j+1))
+        deps.append((self.copy_signal[Device[src.device]], self.copy_signal_val[Device[src.device]]))
+        self.copy_signal_val[Device[src.device]] = sig_val
+
+        for sig,val in deps: self.copy_queues[Device[src.device]].wait(sig, val)
+        self.copy_queues[Device[src.device]].copy(dest._buf.va_addr, src._buf.va_addr, dest.nbytes) \
+                                            .signal(self.copy_signal[Device[src.device]], sig_val)
+        self.copy_to_devs[Device[dest.device]].add(Device[src.device])
+
+    for dev in self.devices:
+      if self.copy_signal_val[dev] > 0: self.comp_queues[dev].wait(self.copy_signal[dev], self.copy_signal_val[dev])
+      for dep_dev in self.copy_to_devs[dev]: self.comp_queues[dev].wait(self.copy_signal[dep_dev], self.copy_signal_val[dep_dev])
+
+      if hasattr(self.comp_queues[dev], 'bind'): self.comp_queues[dev].bind(dev)
+      if hasattr(self.copy_queues[dev], 'bind') and self.copy_signal_val[dev] > 0: self.copy_queues[dev].bind(dev)
+
+  def __call__(self, input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int], wait=False) -> Optional[float]:
+    # Wait and restore signals
+    self.kickoff_value += 1
+    for dev in self.devices: dev._wait_signal(dev.timeline_signal, self.graph_timeline[dev])
+    for dev in self.devices:
+      dev._set_signal(self.comp_signal[dev], 0)
+      dev._set_signal(self.copy_signal[dev], 0)
+    dev._set_signal(self.kickoff_signal, self.kickoff_value)
+
+    # Update rawbuffers
+    for (j,i),input_idx in self.input_replace.items():
+      self.ji_kargs_structs[j].__setattr__(f'f{i}', input_rawbuffers[input_idx]._buf.va_addr)
+
+    # Update var_vals
+    for j in self.jc_idx_with_updatable_var_vals:
+      for i,v in enumerate(cast(CompiledRunner, self.jit_cache[j].prg).p.vars):
+        self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[v])
+
+    for j in self.jc_idx_with_updatable_launch_dims:
+      queue, cmd_ptr = self.exec_ptrs[j]
+      queue.update_exec(cmd_ptr, *cast(CompiledRunner, self.jit_cache[j].prg).p.launch_dims(var_vals))
+
+    for dev in self.devices:
+      # Submit sync with world and queues.
+      self.comp_hcq_t().wait(dev.timeline_signal, dev.timeline_value - 1) \
+                       .wait(self.kickoff_signal, self.kickoff_value).submit(dev)
+      self.comp_queues[dev].submit(dev)
+
+      if self.copy_signal_val[dev] > 0:
+        self.copy_hcq_t().wait(dev.timeline_signal, dev.timeline_value - 1) \
+                         .wait(self.kickoff_signal, self.kickoff_value).submit(dev)
+        self.copy_queues[dev].submit(dev)
+
+      # Signal the final value
+      self.comp_hcq_t().signal(dev.timeline_signal, dev.timeline_value).submit(dev)
+      self.graph_timeline[dev] = dev.timeline_value
+      dev.timeline_value += 1
+
+    if wait:
+      st = time.perf_counter()
+      for dev in self.devices: dev._wait_signal(dev.timeline_signal, self.graph_timeline[dev])
+      return time.perf_counter() - st
+    return None
+
+  def access_resources(self, read, write, new_dependency):
+    deps = self._access_resources(read, write, new_dependency)
+    return [(k, max(v for x, v in deps if id(x) == idk)) for idk, k in {id(x[0]): x[0] for x in deps}.items()]

tinygrad/runtime/graph/hsa.py

@@ -0,0 +1,171 @@
+import ctypes, collections, time, itertools
+from typing import List, Any, Dict, cast, Optional, Tuple
+from tinygrad.helpers import GraphException, init_c_var, round_up
+from tinygrad.device import Buffer, BufferOptions
+from tinygrad.device import Compiled, Device
+from tinygrad.shape.symbolic import Variable
+from tinygrad.runtime.ops_hsa import HSADevice, PROFILE, Profiler
+from tinygrad.engine.realize import ExecItem, BufferXfer, CompiledRunner
+from tinygrad.engine.jit import MultiGraphRunner
+import tinygrad.runtime.autogen.hsa as hsa
+from tinygrad.runtime.driver.hsa import check, AQLQueue, AQL_PACKET_SIZE, EMPTY_SIGNAL
+
+def dedup_signals(signals): return [hsa.hsa_signal_t(hndl) for hndl in set([x.handle for x in signals if isinstance(x, hsa.hsa_signal_t)])]
+
+class VirtAQLQueue(AQLQueue):
+  def __init__(self, device, sz):
+    self.device = device
+    self.virt_queue = (hsa.hsa_kernel_dispatch_packet_t * sz)()
+    self.queue_base = self.write_addr = ctypes.addressof(self.virt_queue)
+    self.packets_count = 0
+    self.available_packet_slots = sz
+  def _wait_queue(self, need_packets=1): assert False, f"VirtQueue is too small to handle {self.packets_count+need_packets} packets!"
+  def _submit_packet(self):
+    self.write_addr += AQL_PACKET_SIZE
+    self.packets_count += 1
+    self.available_packet_slots -= 1
+
+class HSAGraph(MultiGraphRunner):
+  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
+    super().__init__(jit_cache, input_rawbuffers, var_vals)
+
+    # Check all jit items are compatible.
+    compiled_devices = set()
+    for ji in self.jit_cache:
+      if isinstance(ji.prg, CompiledRunner): compiled_devices.add(ji.prg.device)
+      elif isinstance(ji.prg, BufferXfer):
+        for x in ji.bufs[0:2]: compiled_devices.add(Device[cast(Buffer, x).device])
+      else: raise GraphException
+    if any(not isinstance(d, HSADevice) for d in compiled_devices): raise GraphException
+
+    self.devices: List[HSADevice] = list(compiled_devices) #type:ignore
+
+    # Allocate kernel args.
+    kernargs_size: Dict[Compiled, int] = collections.defaultdict(int)
+    for ji in self.jit_cache:
+      if isinstance(ji.prg, CompiledRunner): kernargs_size[ji.prg.device] += round_up(ctypes.sizeof(ji.prg.clprg.args_struct_t), 16)
+    kernargs_ptrs: Dict[Compiled, int] = {dev:dev.allocator._alloc(sz, BufferOptions()) for dev,sz in kernargs_size.items()}
+
+    # Fill initial arguments.
+    self.ji_kargs_structs: Dict[int, ctypes.Structure] = {}
+    for j,ji in enumerate(self.jit_cache):
+      if not isinstance(ji.prg, CompiledRunner): continue
+      self.ji_kargs_structs[j] = ji.prg.clprg.args_struct_t.from_address(kernargs_ptrs[ji.prg.device])
+      kernargs_ptrs[ji.prg.device] += round_up(ctypes.sizeof(ji.prg.clprg.args_struct_t), 16)
+      for i in range(len(ji.bufs)): self.ji_kargs_structs[j].__setattr__(f'f{i}', cast(Buffer, ji.bufs[i])._buf)
+      for i in range(len(ji.prg.p.vars)): self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[ji.prg.p.vars[i]])
+
+    # Build queues.
+    self.virt_aql_queues: Dict[Compiled, VirtAQLQueue] = {dev:VirtAQLQueue(dev, 2*len(self.jit_cache)+16) for dev in self.devices}
+    self.packets = {}
+    self.transfers = []
+    self.ji_to_transfer: Dict[int, int] = {} # faster to store transfers as list and update using this mapping table.
+    self.signals_to_reset: List[hsa.hsa_signal_t] = []
+    self.signals_to_devices: Dict[ctypes.c_uint64, List[HSADevice]] = {}
+    self.profile_info: Dict[Compiled, List[Tuple[Any, ...]]] = collections.defaultdict(list)
+
+    # Special packet to wait for the world.
+    self.kickoff_signals: Dict[HSADevice, hsa.hsa_signal_t] = {dev:self.alloc_signal(reset_on_start=True) for dev in self.devices}
+    for dev in self.devices: self.virt_aql_queues[dev].submit_barrier([], self.kickoff_signals[dev])
+
+    for j,ji in enumerate(self.jit_cache):
+      if isinstance(ji.prg, CompiledRunner):
+        wait_signals = self.access_resources(ji.bufs[(outs:=ji.prg.p.outcount):], ji.bufs[:outs], new_dependency=j, sync_with_aql_packets=False)
+        for i in range(0, len(wait_signals), 5):
+          self.virt_aql_queues[ji.prg.device].submit_barrier(wait_signals[i:i+5])
+        self.packets[j] = hsa.hsa_kernel_dispatch_packet_t.from_address(self.virt_aql_queues[ji.prg.device].write_addr)
+
+        sync_signal = self.alloc_signal(reset_on_start=True) if PROFILE else None
+        self.virt_aql_queues[ji.prg.device].submit_kernel(ji.prg.clprg, *ji.prg.p.launch_dims(var_vals), #type:ignore
+                                                          ctypes.addressof(self.ji_kargs_structs[j]), completion_signal=sync_signal)
+        if PROFILE: self.profile_info[ji.prg.device].append((sync_signal, ji.prg.clprg.name, False))
+      elif isinstance(ji.prg, BufferXfer):
+        dest, src = [cast(Buffer, x) for x in ji.bufs[0:2]]
+        dest_dev, src_dev = cast(HSADevice, Device[dest.device]), cast(HSADevice, Device[src.device])
+        sync_signal = self.alloc_signal(reset_on_start=True, wait_on=[dest_dev, src_dev])
+
+        wait_signals = self.access_resources(read=[src], write=[dest], new_dependency=sync_signal, sync_with_aql_packets=True)
+        self.transfers.append([dest._buf, dest_dev.agent, src._buf, src_dev.agent, dest.nbytes, len(wait_signals),
+                              (hsa.hsa_signal_t*len(wait_signals))(*wait_signals), sync_signal, hsa.HSA_AMD_SDMA_ENGINE_0, True])
+        self.ji_to_transfer[j] = len(self.transfers) - 1
+        if PROFILE: self.profile_info[src_dev].append((sync_signal, f"transfer: HSA:{src_dev.device_id} -> HSA:{dest_dev.device_id}", True))
+
+    # Wait for all active signals to finish the graph
+    wait_signals_to_finish: Dict[HSADevice, List[hsa.hsa_signal_t]] = collections.defaultdict(list)
+    for v in dedup_signals(list(self.w_dependency_map.values()) + list(itertools.chain.from_iterable(self.r_dependency_map.values()))):
+      for dev in self.signals_to_devices[v.handle]:
+        wait_signals_to_finish[dev].append(v)
+
+    self.finish_signal = init_c_var(hsa.hsa_signal_t(), lambda x: check(hsa.hsa_amd_signal_create(1, 0, None, 0, ctypes.byref(x))))
+    for dev in self.devices:
+      wait_signals = wait_signals_to_finish[dev]
+      for i in range(0, max(1, len(wait_signals)), 5):
+        self.virt_aql_queues[dev].submit_barrier(wait_signals[i:i+5], completion_signal=self.finish_signal if i+5>=len(wait_signals) else None)
+
+    # Zero signals to allow graph to start and execute.
+    for sig in self.signals_to_reset: hsa.hsa_signal_silent_store_relaxed(sig, 0)
+    hsa.hsa_signal_silent_store_relaxed(self.finish_signal, 0)
+
+  def __call__(self, input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int], wait=False) -> Optional[float]:
+    # Wait and restore signals
+    hsa.hsa_signal_wait_scacquire(self.finish_signal, hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
+    for sig in self.signals_to_reset: hsa.hsa_signal_silent_store_relaxed(sig, 1)
+    hsa.hsa_signal_silent_store_relaxed(self.finish_signal, len(self.devices))
+
+    # Update rawbuffers
+    for (j,i),input_idx in self.input_replace.items():
+      if j in self.ji_kargs_structs:
+        self.ji_kargs_structs[j].__setattr__(f'f{i}', input_rawbuffers[input_idx]._buf)
+      else:
+        if i == 0: self.transfers[self.ji_to_transfer[j]][0] = input_rawbuffers[input_idx]._buf # dest
+        elif i == 1: self.transfers[self.ji_to_transfer[j]][2] = input_rawbuffers[input_idx]._buf # src
+
+    # Update var_vals
+    for j in self.jc_idx_with_updatable_var_vals:
+      for i,v in enumerate(cast(CompiledRunner, self.jit_cache[j].prg).p.vars):
+        self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[v])
+
+    # Update launch dims
+    for j in self.jc_idx_with_updatable_launch_dims:
+      gl, lc = cast(CompiledRunner, self.jit_cache[j].prg).p.launch_dims(var_vals)
+      self.packets[j].workgroup_size_x = lc[0]
+      self.packets[j].workgroup_size_y = lc[1]
+      self.packets[j].workgroup_size_z = lc[2]
+      self.packets[j].grid_size_x = gl[0] * lc[0]
+      self.packets[j].grid_size_y = gl[1] * lc[1]
+      self.packets[j].grid_size_z = gl[2] * lc[2]
+
+    for dev in self.devices:
+      dev.flush_hdp()
+      dev.hw_queue.blit_packets(self.virt_aql_queues[dev].queue_base, self.virt_aql_queues[dev].packets_count)
+
+    for transfer_data in self.transfers:
+      check(hsa.hsa_amd_memory_async_copy_on_engine(*transfer_data))
+
+    et = None
+    if wait:
+      st = time.perf_counter()
+      hsa.hsa_signal_wait_scacquire(self.finish_signal, hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
+      et = time.perf_counter() - st
+
+    for profdev,profdata in self.profile_info.items(): Profiler.tracked_signals[profdev] += profdata
+    return et
+
+  def alloc_signal(self, reset_on_start=False, wait_on=None):
+    sync_signal = init_c_var(hsa.hsa_signal_t(), lambda x: check(hsa.hsa_amd_signal_create(1, 0, None, 0, ctypes.byref(x))))
+    if reset_on_start: self.signals_to_reset.append(sync_signal)
+    if wait_on is not None: self.signals_to_devices[sync_signal.handle] = wait_on
+    return sync_signal
+
+  def dependency_as_signal(self, dep, sync_with_aql_packets) -> Optional[hsa.hsa_signal_t]:
+    if isinstance(dep, hsa.hsa_signal_t): return dep
+    elif sync_with_aql_packets and isinstance(packet := self.packets.get(dep), hsa.hsa_kernel_dispatch_packet_t):
+      if packet.completion_signal.handle == EMPTY_SIGNAL.handle: packet.completion_signal = self.alloc_signal(reset_on_start=True)
+      return packet.completion_signal
+    return None
+
+  def access_resources(self, read, write, new_dependency, sync_with_aql_packets=False):
+    rdeps = self._access_resources(read, write, new_dependency)
+    wait_signals = [self.dependency_as_signal(dep, sync_with_aql_packets=sync_with_aql_packets) for dep in rdeps]
+    if sync_with_aql_packets: wait_signals += [self.kickoff_signals[cast(HSADevice, Device[rawbuf.device])] for rawbuf in read+write]
+    return dedup_signals(wait_signals)