tinygrad 0.8.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

@@ -1,76 +1,81 @@
 import ctypes
 from typing import Any, Optional, Tuple, Dict, List, cast
-import gpuctypes.cuda as cuda
-from tinygrad.helpers import init_c_var, encode_args_cuda_style
-from tinygrad.device import CompiledASTRunner, update_stats, Buffer
-from tinygrad.runtime.ops_cuda import check, cu_time_execution
+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.jit import JitItem, get_input_replace, get_jit_stats, get_jc_idxs_with_updatable_launch_dims, get_jc_idxs_with_updatable_var_vals, GraphException  # noqa: E501
+from tinygrad.engine.realize import ExecItem, BufferXfer, CompiledRunner
+from tinygrad.engine.jit import MultiGraphRunner
 
-class CUDAGraph:
-  def __init__(self, jit_cache: List[JitItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
-    if not all(isinstance(ji.prg, CompiledASTRunner) for ji in jit_cache): raise GraphException
+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)
 
-    self.jit_cache = jit_cache
-    self.input_replace = get_input_replace(jit_cache, input_rawbuffers)
-    self.op_estimate, self.mem_estimate = get_jit_stats(jit_cache)
-    self.jc_idxs_with_updatable_launch_dims = get_jc_idxs_with_updatable_launch_dims(jit_cache)
-    self.jc_idxs_with_updatable_var_vals = get_jc_idxs_with_updatable_var_vals(jit_cache)
-    self.jc_idxs_with_updatable_rawbufs = list(set([x[0] for x in self.input_replace.keys()]))
-    self.updatable_nodes: Dict[int, Tuple[Any, Any, Any]] = {} # Dict[jc index] = tuple(graph node, node params, input kernel params)
+    # Check all jit items are compatible.
+    if not all(isinstance(ji.prg, (CompiledRunner, BufferXfer)) for ji in jit_cache): raise GraphException
 
-    self.graph = self.graph_create()
-    graph_node: Optional[ctypes._CData] = None
+    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,i),input_name in self.input_replace.items(): self.jit_cache[j].rawbufs[i] = input_rawbuffers[input_name]
     for j,ji in enumerate(self.jit_cache):
-      prg: CompiledASTRunner = cast(CompiledASTRunner, ji.prg)
+      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_deps = (type(graph_node)*1)(*(graph_node,)) if graph_node is not None else None
-      c_kernel_input_config, c_input_params = encode_args_cuda_style([cast(Buffer, x)._buf for x in ji.rawbufs], [var_vals[x] for x in prg.vars], *self.encode_args_info())  # noqa: E501
-      c_node_params = self.build_kernel_node_params(prg, *cast(Tuple[List[int], List[int]], prg.launch_dims(var_vals)), c_kernel_input_config)
-      graph_node = self.graph_add_kernel_node(self.graph, c_deps, c_node_params)
+        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_idxs_with_updatable_launch_dims or j in self.jc_idxs_with_updatable_var_vals or j in self.jc_idxs_with_updatable_rawbufs:
-        self.updatable_nodes[j] = (graph_node, c_node_params, c_input_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 = self.graph_instantiate(self.graph)
+    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, jit=False) -> Optional[float]:
-    # Update rawbuffers in the c_input_params struct.
+  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():
-      setattr(self.updatable_nodes[j][2], f'f{i}', input_rawbuffers[input_idx]._buf)
+      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_input_params struct.
-    for j in self.jc_idxs_with_updatable_var_vals:
-      for i,v in enumerate(cast(CompiledASTRunner, self.jit_cache[j].prg).vars):
-        setattr(self.updatable_nodes[j][2], f'f{len(self.jit_cache[j].rawbufs) + i}', var_vals[v])
+    # 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 c_node_params struct.
-    for j in self.jc_idxs_with_updatable_launch_dims:
-      self.set_kernel_node_launch_dims(self.updatable_nodes[j][1], *cast(CompiledASTRunner, self.jit_cache[j].prg).launch_dims(var_vals))
+    # 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, _ in self.updatable_nodes.values():
-      self.graph_exec_kernel_node_set_params(self.instance, node, ctypes.byref(c_node_params))
+    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))
 
-    et = self.graph_launch(self.instance, None, wait=wait)
-    update_stats(f"<batched {len(self.jit_cache)}>", self.op_estimate, self.mem_estimate, var_vals, et, buf_count=len(input_rawbuffers), jit=jit, num_kernels=len(self.jit_cache))  # noqa: E501
-    return et
+    return cu_time_execution(lambda: check(cuda.cuGraphLaunch(self.instance, None)), enable=wait)
 
   def __del__(self):
-    check(cuda.cuGraphDestroy(self.graph))
-    check(cuda.cuGraphExecDestroy(self.instance))
-
-  def encode_args_info(self): return (cuda.CUdeviceptr_v2, (1,2,0))
-  def graph_create(self): return init_c_var(cuda.CUgraph(), lambda x: check(cuda.cuGraphCreate(ctypes.byref(x), 0)))
-  def graph_instantiate(self, graph):
-    return init_c_var(cuda.CUgraphExec(), lambda x: check(cuda.cuGraphInstantiate_v2(ctypes.byref(x), graph, None, None, 0)))
-  def graph_add_kernel_node(self, graph, c_deps, c_node_params):
-    return init_c_var(cuda.CUgraphNode(), lambda x: check(cuda.cuGraphAddKernelNode(ctypes.byref(x), graph, c_deps, ctypes.sizeof(c_deps)//8 if c_deps else 0, ctypes.byref(c_node_params))))  # noqa: E501
-  def graph_launch(self, *args, wait=False): return cu_time_execution(lambda: check(cuda.cuGraphLaunch(*args)), enable=wait)
-  def graph_exec_kernel_node_set_params(self, *args): return check(cuda.cuGraphExecKernelNodeSetParams(*args))
-  def build_kernel_node_params(self, prg, global_size, local_size, c_kernel_config):
-    return cuda.CUDA_KERNEL_NODE_PARAMS(prg.clprg.prg, *global_size, *local_size, 0, None, c_kernel_config)
+    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()]