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

tinygrad/runtime/graph/hsa.py

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

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