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

tinygrad/runtime/driver/hip_comgr.py

@@ -0,0 +1,56 @@
+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", asm=False) -> 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))
+
+  if asm:
+    check(comgr.amd_comgr_set_data_name(data_src, b"<null>.s"))
+    check(comgr.amd_comgr_data_set_add(data_set_src, data_src))
+    status = comgr.amd_comgr_do_action(comgr.AMD_COMGR_ACTION_ASSEMBLE_SOURCE_TO_RELOCATABLE, action_info, data_set_src, data_set_reloc)
+    if status != 0:
+      print(_get_comgr_data(data_set_reloc, comgr.AMD_COMGR_DATA_KIND_LOG).decode())
+      raise RuntimeError("assemble failed")
+  else:
+    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/graph/clang.py

@@ -0,0 +1,39 @@
+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 += sorted([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[1] for x in sorted(var_vals.items(), key=lambda x: x[0].expr)]), 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, dedup
+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 = dedup([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,187 @@
+import collections, array, time
+from typing import List, Any, Dict, cast, Optional, Tuple, Set
+from tinygrad.helpers import round_up, to_mv, PROFILE
+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, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
+    super().__init__(jit_cache, input_rawbuffers, var_vals)
+    self.devices = list(set(cast(Any, d) for ji in jit_cache for d in [Device[cast(Buffer, x).device] for x in ji.bufs]))
+
+    # 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_alloc_size, 16)
+    self.kernargs_bufs: Dict[Compiled, Any] = {dev:dev.allocator._alloc(sz, BufferOptions(cpu_access=True)) for dev,sz in kernargs_size.items()}
+    kernargs_ptrs: Dict[Compiled, int] = {dev:buf.va_addr for dev,buf in self.kernargs_bufs.items()}
+
+    # Fill initial arguments.
+    self.kargs_addrs: Dict[int, int] = {}
+    self.ji_args_bufs: Dict[int, memoryview] = {}
+    self.ji_args_vars: Dict[int, memoryview] = {}
+    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_alloc_size, 16)
+
+      self.ji_args_bufs[j] = to_mv(self.kargs_addrs[j] + ji.prg.clprg.kernargs_offset, len(ji.bufs) * 8).cast('Q')
+      self.ji_args_vars[j] = to_mv(self.kargs_addrs[j] + ji.prg.clprg.kernargs_offset + len(ji.bufs) * 8, len(ji.prg.p.vars) * 4).cast('I')
+      for i in range(len(ji.bufs)): self.ji_args_bufs[j][i] = cast(Buffer, ji.bufs[i])._buf.va_addr
+      for i in range(len(ji.prg.p.vars)): self.ji_args_vars[j][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)
+
+    # Schedule Dependencies.
+    # There are two types of queues on each device: copy and compute. Both must synchronize with all external operations before launching any
+    # graph-related tasks. This synchronization uses a global timeline signal per device. Within the graph, the compute queue coordinates with
+    # global operations and sets a kickoff signal. Any queue accessing a buffer from another device waits for this signal from the device’s
+    # compute queue to ensure exclusive access. The compute queue signals the completion of the graph, synchronizing with the device's copy queue.
+    self.comp_queues: Dict[Compiled, Any] = {dev: dev.hw_compute_queue_t() for dev in self.devices}
+    self.copy_queues: Dict[Compiled, Any] = {dev: dev.hw_copy_queue_t() for dev in self.devices}
+
+    self.signal_sched: Dict[int, Tuple[List, Optional[int], Optional[List]]] = {} # Dict[ji_idx, (deps, sigval, prof_info)]
+    self.signals: Dict[Any, Any] = {q: self.devices[0]._get_signal(value=0) for q in list(self.comp_queues.values())+list(self.copy_queues.values())}
+    self.dev_kickoff_signal = {dev: self.devices[0]._get_signal(value=0) for dev in self.devices + ['CPU']} # Dict[dev, signal]
+    self.kickoff_value = 0
+
+    self.save_devs: Dict[Any, Set] = {q: set() for q in list(self.comp_queues.values()) + list(self.copy_queues.values())}
+    for dev in self.devices: self.save_devs[self.comp_queues[dev]].add(dev)
+
+    self.graph_timeline = {dev: 0 for dev in self.devices} # Dict[dev, last graph sigval]
+    self.last_ji: Dict[Any, Any] = {q: None for q in list(self.comp_queues.values()) + list(self.copy_queues.values())}
+
+    for j,ji in enumerate(self.jit_cache):
+      enqueue_dev = ji.prg.device if isinstance(ji.prg, CompiledRunner) else Device[ji.bufs[1].device] #type:ignore
+      enqueue_queue = self.comp_queues[enqueue_dev] if isinstance(ji.prg, CompiledRunner) else self.copy_queues[enqueue_dev]
+      out_signal = self.signals[enqueue_queue]
+      writable_buffers = ji.prg.p.outcount if isinstance(ji.prg, CompiledRunner) else 1
+      deps = self.access_resources(enqueue_queue, ji.bufs[writable_buffers:], ji.bufs[:writable_buffers], j + 1)
+
+      if isinstance(ji.prg, CompiledRunner):
+        # Update signal on compute kernel to depend on the previous kernel.
+        if (last_j:=self.last_ji[enqueue_queue]) is not None: deps = [x for x in deps if id(x[0]) != id(out_signal)] + [(out_signal, last_j + 1)]
+
+        # Remove self-dependency for AMD or NV with only 1 same-queue dep, since NV chains 2+ execs in this case, eliminating dep need.
+        if (dname:=enqueue_dev.dname.split(":", 1)[0]) == "AMD" or (dname == "NV" and len(deps) == 1 and id(deps[0][0]) == id(out_signal)):
+          deps = [x for x in deps if id(x[0]) != id(out_signal)]
+      elif isinstance(ji.prg, BufferXfer): deps = [x for x in deps if id(x[0]) != id(out_signal)]
+
+      # Go through all dependencies and, if we need the signal from that ji, enable it by setting the signal value in the signal schedule.
+      for sig, val in deps:
+        if id(sig) in [id(x) for x in self.signals.values()]:
+          self.signal_sched[val - 1] = self.signal_sched[val - 1][:1] + (val,) + self.signal_sched[val - 1][2:]
+
+      prof_ji_desc = ji.prg.clprg.name if isinstance(ji.prg, CompiledRunner) else f"{ji.bufs[1].device} -> {ji.bufs[0].device}" # type: ignore
+      prof_info = ([enqueue_dev._get_signal() for _ in range(2)] + [enqueue_dev, prof_ji_desc, isinstance(ji.prg, BufferXfer)]) if PROFILE else None
+      self.signal_sched[j] = (deps, None if isinstance(ji.prg, CompiledRunner) else (j + 1), prof_info)
+      self.last_ji[enqueue_queue] = j
+
+    # Build hardware queues.
+    self.exec_ptrs: Dict[int, Tuple[Any, int]] = {}
+    self.copy_to_devs: Dict[Compiled, Set[Compiled]] = {dev: set() for dev in self.devices}
+    self.kickoff_wait_cmds: Dict[Any, List] = {q: list() for q in list(self.comp_queues.values()) + list(self.copy_queues.values())}
+
+    for dev in self.devices:
+      self.comp_queues[dev].memory_barrier().wait(dev.timeline_signal, dev.timeline_value - 1) \
+                           .wait(self.dev_kickoff_signal['CPU'], self.kickoff_value).signal(self.dev_kickoff_signal[dev], self.kickoff_value)
+
+    for j,ji in enumerate(self.jit_cache):
+      deps, signal_value, prof_info = self.signal_sched[j]
+      enqueue_queue = self.copy_queues[Device[ji.bufs[1].device]] if isinstance(ji.prg, BufferXfer) else self.comp_queues[ji.prg.device] #type:ignore
+
+      # Encode waits and start profile timestamp (if needed).
+      for sig, val in deps:
+        enqueue_queue.wait(sig, val)
+        if id(sig) in [id(x) for x in self.dev_kickoff_signal.values()]: self.kickoff_wait_cmds[enqueue_queue].append(len(enqueue_queue) - 1)
+      if prof_info: enqueue_queue.timestamp(prof_info[0])
+
+      # Encode main commands based on ji type.
+      if isinstance(ji.prg, CompiledRunner):
+        enqueue_queue.exec(ji.prg.clprg, self.kargs_addrs[j], *ji.prg.p.launch_dims(var_vals),
+                           signal=self.signals[enqueue_queue] if signal_value is not None else None, signal_value=signal_value)
+        self.exec_ptrs[j] = (enqueue_queue, len(enqueue_queue) - 1)
+      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
+        enqueue_queue.copy(dest._buf.va_addr, src._buf.va_addr, dest.nbytes).signal(self.signals[enqueue_queue], signal_value)
+        self.copy_to_devs[Device[dest.device]].add(Device[src.device])
+
+      # Encode finish profile timestamp (if needed).
+      if prof_info: enqueue_queue.timestamp(prof_info[1])
+
+    for dev in self.devices:
+      for dep_dev in list(self.copy_to_devs[dev]) + [dev]:
+        if (last_j:=self.last_ji[self.copy_queues[dep_dev]]) is None: continue
+        self.comp_queues[dev].wait(self.signals[self.copy_queues[dep_dev]], self.signal_sched[last_j][1])
+
+      self.comp_queues[dev].signal(dev.timeline_signal, dev.timeline_value)
+      if hasattr(self.comp_queues[dev], 'bind'): self.comp_queues[dev].bind(dev)
+      if hasattr(self.copy_queues[dev], 'bind') and self.last_ji[self.copy_queues[dev]] is not None: 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 queue in self.comp_queues.values(): self.devices[0]._set_signal(self.signals[queue], 0)
+    for queue in self.copy_queues.values(): self.devices[0]._set_signal(self.signals[queue], 0)
+    self.devices[0]._set_signal(self.dev_kickoff_signal['CPU'], self.kickoff_value)
+
+    if PROFILE and self.kickoff_value > 1:
+      for _,_,(st,en,dev,desc,is_cp) in self.signal_sched.values(): #type: ignore
+        dev.raw_prof_records += [(dev._read_timestamp(st), dev._read_timestamp(en), desc, is_cp)]
+
+    # Update rawbuffers
+    for (j,i),input_idx in self.input_replace.items(): self.ji_args_bufs[j][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_args_vars[j][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:
+      self.comp_queues[dev].update_wait(1, dev.timeline_signal, dev.timeline_value - 1).update_wait(2, value=self.kickoff_value) \
+                           .update_signal(3, value=self.kickoff_value) \
+                           .update_signal(len(self.comp_queues[dev]) - 1, dev.timeline_signal, dev.timeline_value).submit(dev)
+
+      if self.last_ji[(cp_queue:=self.copy_queues[dev])] is not None:
+        for cmd_idx in self.kickoff_wait_cmds[cp_queue]: cp_queue.update_wait(cmd_idx, value=self.kickoff_value)
+        cp_queue.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, queue, read, write, new_val):
+    deps = self._access_resources(read, write, (queue, new_val))
+
+    sync_signals = []
+    for dep_queue,_ in deps: self.save_devs[queue].update(self.save_devs[dep_queue])
+    for buf in read+write:
+      if buf.device not in self.save_devs[queue]:
+        self.save_devs[queue].add(buf.device)
+        sync_signals += [(self.dev_kickoff_signal[Device[buf.device]], self.kickoff_value)]
+
+    return [(self.signals[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()] + sync_signals
+
+  def __del__(self):
+    for dev in self.devices: dev._wait_signal(dev.timeline_signal, self.graph_timeline[dev])
+
+    # Graph is destructed. No need to keep signals any more, so return them as part of profiling.
+    if PROFILE and self.kickoff_value > 1:
+      for _,_,(st,en,dev,desc,is_cp) in self.signal_sched.values(): dev.sig_prof_records += [(st, en, desc, is_cp)] #type: ignore
+
+    self.devices[0].signals_pool += list(self.dev_kickoff_signal.values()) + list(self.signals.values()) # type: ignore
+    for dev, buf in self.kernargs_bufs.items(): dev.allocator._free(buf, BufferOptions(cpu_access=True))

tinygrad/runtime/graph/metal.py

@@ -1,22 +1,17 @@
 from typing import List, Any, Dict, cast, Optional
-import numpy as np
 import Metal
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import dedup, unwrap2
-from tinygrad.device import Buffer, CompiledASTRunner, update_stats
-from tinygrad.jit import JitItem, get_input_replace, get_jit_stats, get_jc_idxs_with_updatable_launch_dims, GraphException
+from tinygrad.helpers import dedup, unwrap2, GraphException
+from tinygrad.device import Buffer
+from tinygrad.engine.realize import ExecItem, CompiledRunner
+from tinygrad.engine.jit import GraphRunner
 from tinygrad.shape.symbolic import Variable
-from tinygrad.runtime.ops_metal import MetalDevice
+from tinygrad.runtime.ops_metal import wait_check
 
-class MetalGraph:
-  def __init__(self, device:MetalDevice, 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
-
-    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_idx_with_updatable_launch_dims = get_jc_idxs_with_updatable_launch_dims(jit_cache)
-    self.device: MetalDevice = device
+class MetalGraph(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
 
     # create metal batch exec
     icb_descriptor = Metal.MTLIndirectCommandBufferDescriptor.new()
@@ -24,56 +19,57 @@ class MetalGraph:
     icb_descriptor.setInheritBuffers_(False)
     icb_descriptor.setInheritPipelineState_(False)
     icb_descriptor.setMaxKernelBufferBindCount_(31)
-    self.icb = self.device.device.newIndirectCommandBufferWithDescriptor_maxCommandCount_options_(icb_descriptor, len(self.jit_cache), Metal.MTLResourceOptions(0))  # noqa: E501
+    self.icb = self.device.device.newIndirectCommandBufferWithDescriptor_maxCommandCount_options_(icb_descriptor, len(self.jit_cache),
+                                                                                                  Metal.MTLResourceOptions(0))
     if self.icb is None: raise GraphException("create indirect command buffer failed, does your system support this?")
 
-    if len(var_vals): self.int_buf = self.device.allocator.alloc(len(var_vals)*dtypes.int32.itemsize)
-    all_resources = [self.int_buf] if len(var_vals) else []
+    if len(self.vars): self.int_buf = self.device.allocator.alloc(len(self.vars)*dtypes.int32.itemsize)
+    all_resources = [self.int_buf] if len(self.vars) else []
+
     for j,ji in enumerate(self.jit_cache):
-      prg: CompiledASTRunner = cast(CompiledASTRunner, ji.prg)
+      prg: CompiledRunner = cast(CompiledRunner, ji.prg)
       descriptor = Metal.MTLComputePipelineDescriptor.new()
       descriptor.setComputeFunction_(prg.clprg.fxn)
       descriptor.setSupportIndirectCommandBuffers_(True)
-      pipeline_state = unwrap2(self.device.device.newComputePipelineStateWithDescriptor_options_reflection_error_(descriptor, Metal.MTLPipelineOption(0), None, None))  # noqa: E501
       icb_command = self.icb.indirectComputeCommandAtIndex_(j)
-      icb_command.setComputePipelineState_(pipeline_state)
-      for i,b in enumerate(ji.rawbufs):
+      icb_command.setComputePipelineState_(unwrap2(
+        self.device.device.newComputePipelineStateWithDescriptor_options_reflection_error_(descriptor, Metal.MTLPipelineOption(0), None, None)))
+      for i,b in enumerate(ji.bufs):
         if b is not None:
           icb_command.setKernelBuffer_offset_atIndex_(b._buf, 0, i)
           all_resources.append(b._buf)
-      var_vals_keys = list(var_vals.keys())
-      for i,v in enumerate(prg.vars):
-        icb_command.setKernelBuffer_offset_atIndex_(self.int_buf, var_vals_keys.index(v)*4, len(ji.rawbufs)+i)
+      for i,v in enumerate(prg.p.vars): icb_command.setKernelBuffer_offset_atIndex_(self.int_buf, self.vars.index(v)*4, len(ji.bufs)+i)
       if j not in self.jc_idx_with_updatable_launch_dims:
-        global_size, local_size = prg.launch_dims(var_vals)
+        global_size, local_size = prg.p.launch_dims(var_vals)
         icb_command.concurrentDispatchThreadgroups_threadsPerThreadgroup_(Metal.MTLSize(*global_size), Metal.MTLSize(*local_size))
       icb_command.setBarrier()
+
     self.all_resources = dedup(all_resources)
     self.command_buffer: Any = None
-    if len(var_vals): self.int_buf_view = np.frombuffer(self.int_buf.contents().as_buffer(self.int_buf.length()), np.int32)
+    if len(self.vars): self.int_buf_view = self.int_buf.contents().as_buffer(self.int_buf.length()).cast('i')
+
+  def __call__(self, input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int], wait=False) -> Optional[float]:
+    if self.command_buffer is not None and self.command_buffer in self.device.mtl_buffers_in_flight: wait_check(self.command_buffer)
+    all_resources = dedup(self.all_resources + [x._buf for x in input_rawbuffers])
 
-  def __call__(self, input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int], wait=False, jit=False) -> Optional[float]:
-    # NOTE: you at least can't update the ints if this is running
-    if self.command_buffer is not None and self.command_buffer in self.device.mtl_buffers_in_flight: self.command_buffer.waitUntilCompleted()
-    all_resources = self.all_resources + [x._buf for x in input_rawbuffers]
     for (j,i),input_idx in self.input_replace.items():
       self.icb.indirectComputeCommandAtIndex_(j).setKernelBuffer_offset_atIndex_(input_rawbuffers[input_idx]._buf, 0, i)
     for j in self.jc_idx_with_updatable_launch_dims:
-      global_size, local_size = cast(CompiledASTRunner, self.jit_cache[j].prg).launch_dims(var_vals)
-      self.icb.indirectComputeCommandAtIndex_(j).concurrentDispatchThreadgroups_threadsPerThreadgroup_(Metal.MTLSize(*global_size), Metal.MTLSize(*local_size))  # noqa: E501
-    if len(var_vals): self.int_buf_view[:] = list(var_vals.values())
+      global_size, local_size = cast(CompiledRunner, self.jit_cache[j].prg).p.launch_dims(var_vals)
+      self.icb.indirectComputeCommandAtIndex_(j).concurrentDispatchThreadgroups_threadsPerThreadgroup_(Metal.MTLSize(*global_size),
+                                                                                                       Metal.MTLSize(*local_size))
+    for j, var in enumerate(self.vars): self.int_buf_view[j] = var_vals[var]
+
     command_buffer = self.device.mtl_queue.commandBuffer()
     encoder = command_buffer.computeCommandEncoder()
     encoder.useResources_count_usage_(all_resources, len(all_resources), Metal.MTLResourceUsageRead | Metal.MTLResourceUsageWrite)
-    encoder.executeCommandsInBuffer_withRange_(self.icb, Metal.MTLIndirectCommandBufferExecutionRangeMake(0,len(self.jit_cache)))
+    encoder.executeCommandsInBuffer_withRange_(self.icb, Metal.MTLIndirectCommandBufferExecutionRangeMake(0, len(self.jit_cache)))
     encoder.endEncoding()
     command_buffer.commit()
     self.command_buffer = command_buffer
+
     if wait:
-      command_buffer.waitUntilCompleted()
-      et = command_buffer.GPUEndTime() - command_buffer.GPUStartTime()
-    else:
-      self.device.mtl_buffers_in_flight.append(command_buffer)
-      et = None
-    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
+      wait_check(command_buffer)
+      return command_buffer.GPUEndTime() - command_buffer.GPUStartTime()
+    self.device.mtl_buffers_in_flight.append(command_buffer)
+    return None