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

tinygrad/runtime/graph/hcq.py

@@ -1,7 +1,8 @@
-import collections, array, time
+import collections, 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.helpers import round_up, PROFILE, memsize_to_str
+from tinygrad.device import HCQCompiled, HCQAllocator, HCQSignal, HCQBuffer, HWCommandQueue, HWComputeQueue, HWCopyQueue, HCQArgsState, \
+                            Buffer, BufferOptions, Compiled, Device
 from tinygrad.shape.symbolic import Variable
 from tinygrad.engine.realize import ExecItem, BufferXfer, CompiledRunner
 from tinygrad.engine.jit import MultiGraphRunner
@@ -9,179 +10,191 @@ 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]))
+    self.devices = list(set(cast(HCQCompiled, 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()}
+    self.kernargs_bufs: Dict[Compiled, HCQBuffer] = {dev:dev.allocator._alloc(sz, BufferOptions(cpu_access=True)) for dev,sz in kernargs_size.items()}
 
     # Fill initial arguments.
-    self.kargs_addrs: Dict[int, int] = {}
-    self.ji_args_bufs: Dict[int, memoryview] = {}
-    self.ji_args_vars: Dict[int, memoryview] = {}
+    self.ji_args: Dict[int, HCQArgsState] = {}
+
+    kargs_ptrs: Dict[Compiled, int] = {dev:buf.va_addr for dev,buf in self.kernargs_bufs.items()}
     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)
+      kargs_ptrs[ji.prg.device] = (kargs_ptr:=kargs_ptrs[ji.prg.device]) + round_up(ji.prg.clprg.kernargs_alloc_size, 16)
+      self.ji_args[j] = ji.prg.clprg.fill_kernargs([cast(Buffer, b)._buf for b in ji.bufs], [var_vals[v] for v in ji.prg.p.vars], kargs_ptr)
 
     # 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.ji_schedule: Dict[int, Tuple[HCQCompiled, HWCommandQueue, List, List, HCQSignal, Optional[int]]] = {}
+
+    self.comp_queues: Dict[HCQCompiled, HWComputeQueue] = {dev: dev.hw_compute_queue_t() for dev in self.devices}
+    self.copy_queues: Dict[HCQCompiled, HWCopyQueue] = {} # lazy allocation
 
-    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.signals: Dict[Any, HCQSignal] = {**{dev: dev.signal_t(value=0) for dev in self.devices}, **{"CPU": self.devices[0].signal_t(value=0)}}
+    self.kickoff_value: int = 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.prof_signals: List[HCQSignal] = [self.devices[0].signal_t() for i in range(len(self.jit_cache) * 2)] if PROFILE else []
+    self.prof_records: List[Tuple[Tuple[int, bool], Tuple[int, bool], HCQCompiled, str, bool, List[int], Optional[Dict]]] = []
 
-    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())}
+    last_j: Dict[HWCommandQueue, Optional[int]] = collections.defaultdict(lambda: None)
+    queue_access: Dict[HWCommandQueue, Dict[HWCommandQueue, Optional[int]]] = collections.defaultdict(lambda: collections.defaultdict(lambda: None))
+    dev_access: Dict[HWCommandQueue, Set[HCQCompiled]] = collections.defaultdict(set)
+
+    for dev, queue in self.comp_queues.items(): dev_access[queue].add(dev)
 
     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)
+      enqueue_dev = ji.prg.device if (is_exec_prg:=isinstance(ji.prg, CompiledRunner)) else Device[ji.bufs[1].device] #type:ignore
+      enqueue_queue = self.comp_queues[enqueue_dev] if is_exec_prg else self.copy_queues.setdefault(enqueue_dev, enqueue_dev.hw_copy_queue_t())
+      out_signal = self.signals.setdefault(enqueue_queue, enqueue_dev.signal_t(value=0))
 
-      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)]
+      # Get dependencies based on input and output buffers.
+      rdeps = self._access_resources(ji.bufs[(wb:=ji.prg.p.outcount if is_exec_prg else 1):], ji.bufs[:wb], (enqueue_queue, j + 1)) #type:ignore
+
+      # Update dependencies to include previous kernel in queue. This is required for timeline signals.
+      opt_deps, deps = [], rdeps + ([(enqueue_queue, prev_ji + 1)] if (prev_ji:=last_j[enqueue_queue]) is not None else [])
+
+      # Optimize dependencies by removing redundant ones. Remove waiting for the value of the queue which is known to be already
+      # synced with the current queue.
+      for dep_queue, dep_val in sorted(deps, key=lambda x: x[1], reverse=True):
+        if (qa:=queue_access[enqueue_queue][dep_queue]) is None or qa < dep_val:
+          opt_deps.append((self.signals[dep_queue], dep_val))
+          queue_access[enqueue_queue][dep_queue] = dep_val
 
-        # 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)]
+      # Ensure device is ready for use in current context: the graph has initialized the device and it's safe to operate on it within this graph.
+      for dep_queue, _ in opt_deps: dev_access[enqueue_queue].update(dev_access[dep_queue])
+      sync_signals = [(self.signals[d], self.kickoff_value) for b in ji.bufs if (d:=Device[cast(Buffer, b).device]) not in dev_access[enqueue_queue]]
+      dev_access[enqueue_queue].update(cast(HCQCompiled, Device[cast(Buffer, b).device]) for b in ji.bufs)
 
-      # 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:]
+      # Remove self-dependency for compute and copy queues.
+      # For compute, in case of NV, optimize when only 1 same-queue dependency exists, since NV chains 2+ executions in this case,
+      # eliminating dependency need.
+      dname = enqueue_dev.dname.split(":", 1)[0]
+      can_opt = (dname == "AMD" or (dname == "NV" and len(sync_signals) == 0 and len(opt_deps) == 1 and id(opt_deps[0][0]) == id(out_signal)))
+      if can_opt or isinstance(ji.prg, BufferXfer): opt_deps = [x for x in opt_deps if id(x[0]) != id(out_signal)]
 
-      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
+      # Enable necessary signals in the schedule by setting the signal value.
+      for sig, val in opt_deps: self.ji_schedule[val - 1] = self.ji_schedule[val - 1][:5] + (val,)
+      self.ji_schedule[j] = (enqueue_dev, enqueue_queue, sync_signals, opt_deps[::-1], out_signal, None if is_exec_prg else (j + 1))
+
+      # Collect profile information if profiling is enabled.
+      if PROFILE:
+        prof_ji_desc = ji.prg.clprg.name if is_exec_prg else f"{ji.bufs[1].device} -> {ji.bufs[0].device}" # type: ignore
+
+        sig_st, sig_en = (j * 2, True), (j * 2 + 1, True)
+        if len(opt_deps) == 0 and (prev_ji:=last_j[enqueue_queue]) is not None: sig_st = (prev_ji * 2 + 1, False)
+
+        if is_exec_prg: prof_args = None
+        else: prof_args = {"Size": memsize_to_str(ji.bufs[0].nbytes), "GB/S": lambda dur, b=ji.bufs[0].nbytes: f"{b/1e3/dur:.2f}"} # type: ignore
+
+        self.prof_records.append((sig_st, sig_en, enqueue_dev, prof_ji_desc, not is_exec_prg, [d - 1 for _, d in rdeps], prof_args))
+
+      last_j[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())}
+    self.op_cmd_idx: Dict[int, Tuple[Any, int]] = {}
+    self.copy_to_devs: Dict[HCQCompiled, Set[HCQCompiled]] = {dev: set() for dev in self.devices}
+    self.kickoff_wait_cmds: Dict[HWCommandQueue, 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)
+                           .wait(self.signals['CPU'], self.kickoff_value).signal(self.signals[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
+      enqueue_dev, enqueue_queue, sync_signals, deps, signal, signal_val = self.ji_schedule[j]
+
+      for i in range(len(sync_signals)): self.kickoff_wait_cmds[enqueue_queue].append(len(enqueue_queue) + i)
+      for sig, val in sync_signals + deps: enqueue_queue.wait(sig, val)
 
       # 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])
+      if PROFILE and self.prof_records[j][0][1]: enqueue_queue.timestamp(self.prof_signals[self.prof_records[j][0][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)
+        cast(HWComputeQueue, enqueue_queue).exec(ji.prg.clprg, self.ji_args[j], *ji.prg.p.launch_dims(var_vals))
       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])
+        cast(HCQAllocator, Device[src.device].allocator).map(dest._buf)
+        cast(HWCopyQueue, enqueue_queue).copy(dest._buf.va_addr, src._buf.va_addr, dest.nbytes)
+        self.copy_to_devs[cast(HCQCompiled, Device[dest.device])].add(cast(HCQCompiled, Device[src.device]))
+      self.op_cmd_idx[j] = (enqueue_queue, len(enqueue_queue) - 1)
 
       # Encode finish profile timestamp (if needed).
-      if prof_info: enqueue_queue.timestamp(prof_info[1])
+      if PROFILE and self.prof_records[j][1][1]: enqueue_queue.timestamp(self.prof_signals[self.prof_records[j][1][0]])
+
+      if signal_val is not None: enqueue_queue.signal(signal, signal_val)
 
     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])
+        if dep_dev in self.copy_queues: self.comp_queues[dev].wait(self.signals[(copy_q:=self.copy_queues[dep_dev])], cast(int, last_j[copy_q]) + 1)
+
+      self.comp_queues[dev].signal(dev.timeline_signal, dev.timeline_value).bind(dev)
+      if dev in self.copy_queues: self.copy_queues[dev].bind(dev)
 
-      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)
+    self.last_timeline: Dict[HCQCompiled, Tuple[HCQSignal, int]] = {dev: (dev.timeline_signal, 0) for dev in self.devices}
+    self.queue_signals_to_reset = [self.signals[q] for q in list(self.comp_queues.values()) + list(self.copy_queues.values()) if q in self.signals]
 
   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)
+    for dev in self.devices: self.last_timeline[dev][0].wait(self.last_timeline[dev][1])
+    for sig in self.queue_signals_to_reset: sig.value = 0
+    self.signals['CPU'].value = 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)]
+    if PROFILE and self.kickoff_value > 1: self.collect_timestamps()
 
     # 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
+    for (j,i),input_idx in self.input_replace.items():
+      if j in self.ji_args: self.ji_args[j].update_buffer(i, input_rawbuffers[input_idx]._buf)
+      else: self.op_cmd_idx[j][0].update_copy(self.op_cmd_idx[j][1], **{('dest' if i == 0 else 'src'): 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, i, v in self.updated_vars(var_vals): self.ji_args[j].update_var(i, 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))
+    # Update launch dims
+    for j, global_dims, local_dims in self.updated_launch_dims(var_vals):
+      queue, cmd_ptr = self.op_cmd_idx[j]
+      queue.update_exec(cmd_ptr, global_dims, local_dims)
 
     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)
+      comp_queue, copy_queue, need_sig_upd = self.comp_queues[dev], self.copy_queues.get(dev, None), dev.timeline_signal != self.last_timeline[dev][0]
+      comp_queue.update_wait(1, dev.timeline_signal if need_sig_upd else None, dev.timeline_value - 1) \
+                .update_wait(2, value=self.kickoff_value).update_signal(3, value=self.kickoff_value) \
+                .update_signal(len(comp_queue)-1, dev.timeline_signal if need_sig_upd else None, 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)
+      if copy_queue is not None:
+        for cmd_idx in self.kickoff_wait_cmds[copy_queue]: copy_queue.update_wait(cmd_idx, value=self.kickoff_value)
+        copy_queue.submit(dev)
 
-      self.graph_timeline[dev] = dev.timeline_value
+      self.last_timeline[dev] = (dev.timeline_signal, 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])
+      for dev in self.devices: self.last_timeline[dev][0].wait(self.last_timeline[dev][1])
       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))
+  def collect_timestamps(self):
+    timestamps = [s.timestamp for s in self.prof_signals]
 
-    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)]
+    for (st,_), (en,_), dev, desc, is_cp, deps, args in self.prof_records:
+      dev.raw_prof_records += [(timestamps[st], timestamps[en], desc, is_cp, args)]
 
-    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
+      for x in deps:
+        (b_st,_), (b_en,_), b_dev, _, b_is_cp, _, _ = self.prof_records[x]
+        dev.dep_prof_records += [(timestamps[b_st], timestamps[b_en], b_dev, b_is_cp, timestamps[st], timestamps[en], dev, is_cp)]
 
   def __del__(self):
-    for dev in self.devices: dev._wait_signal(dev.timeline_signal, self.graph_timeline[dev])
+    for dev in self.devices: self.last_timeline[dev][0].wait(self.last_timeline[dev][1])
 
-    # 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
+    if PROFILE and self.kickoff_value >= 1: self.collect_timestamps()
 
-    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))
+    for fdev, buf in self.kernargs_bufs.items(): fdev.allocator._free(buf, BufferOptions(cpu_access=True))

tinygrad/runtime/graph/metal.py

@@ -1,10 +1,10 @@
 from typing import List, Any, Dict, cast, Optional
 import Metal
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import dedup, unwrap2, GraphException
+from tinygrad.helpers import dedup, unwrap2
 from tinygrad.device import Buffer
 from tinygrad.engine.realize import ExecItem, CompiledRunner
-from tinygrad.engine.jit import GraphRunner
+from tinygrad.engine.jit import GraphRunner, GraphException
 from tinygrad.shape.symbolic import Variable
 from tinygrad.runtime.ops_metal import wait_check
 
@@ -24,7 +24,7 @@ class MetalGraph(GraphRunner):
     if self.icb is None: raise GraphException("create indirect command buffer failed, does your system support this?")
 
     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 []
+    all_resources = [self.int_buf.buf] if len(self.vars) else []
 
     for j,ji in enumerate(self.jit_cache):
       prg: CompiledRunner = cast(CompiledRunner, ji.prg)
@@ -35,27 +35,30 @@ class MetalGraph(GraphRunner):
       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)
-      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.p.launch_dims(var_vals)
-        icb_command.concurrentDispatchThreadgroups_threadsPerThreadgroup_(Metal.MTLSize(*global_size), Metal.MTLSize(*local_size))
+        if b is not None and b not in input_rawbuffers:
+          icb_command.setKernelBuffer_offset_atIndex_(b._buf.buf, b._buf.offset, i)
+          all_resources.append(b._buf.buf)
+      for i,v in enumerate(prg.p.vars): icb_command.setKernelBuffer_offset_atIndex_(self.int_buf.buf, self.vars.index(v)*4, len(ji.bufs)+i)
+
+      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(self.vars): self.int_buf_view = self.int_buf.contents().as_buffer(self.int_buf.length()).cast('i')
+    if len(self.vars): self.int_buf_view = self.int_buf.buf.contents().as_buffer(self.int_buf.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])
+    all_resources = dedup(self.all_resources + [x._buf.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(CompiledRunner, self.jit_cache[j].prg).p.launch_dims(var_vals)
+      self.icb.indirectComputeCommandAtIndex_(j).setKernelBuffer_offset_atIndex_(input_rawbuffers[input_idx]._buf.buf,
+                                                                                 input_rawbuffers[input_idx]._buf.offset, i)
+
+    for j, global_dims, local_dims in self.updated_launch_dims(var_vals):
+      prg = cast(CompiledRunner, self.jit_cache[j].prg)
+      global_size, local_size = global_dims or prg.p.global_size, local_dims or prg.p.local_size
       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]