From fa84e6ec4892cefcec3907fa1386680ca4328203 Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Tue, 13 Aug 2024 17:11:58 +0300 Subject: [PATCH] init hcq args state (#6046) * init hcq args state * cleaner * amd * fillargs * fixes * myoy * docs * fix * not needed * spacing --- docs/developer/hcq.md | 13 +++++++++++-- tinygrad/device.py | 30 ++++++++++++++++-------------- tinygrad/runtime/graph/hcq.py | 25 ++++++++++--------------- tinygrad/runtime/ops_amd.py | 30 +++++++++++++++++++----------- tinygrad/runtime/ops_nv.py | 30 +++++++++++++++++++----------- 5 files changed, 75 insertions(+), 53 deletions(-) diff --git a/docs/developer/hcq.md b/docs/developer/hcq.md index 5c896af2e5..982825b652 100644 --- a/docs/developer/hcq.md +++ b/docs/developer/hcq.md @@ -11,7 +11,7 @@ To interact with devices, there are 2 types of queues: `HWComputeQueue` and `HWC For example, the following Python code enqueues a wait, execute, and signal command on the HCQ-compatible device: ```python HWComputeQueue().wait(signal_to_wait, value_to_wait) \ - .exec(program, kernargs_ptr, global_dims, local_dims) \ + .exec(program, args_state, global_dims, local_dims) \ .signal(signal_to_fire, value_to_fire) \ .submit(your_device) ``` @@ -118,13 +118,22 @@ Backends must adhere to the `HCQBuffer` protocol when returning allocation resul ### HCQ Compatible Program -The `HCQProgram` is a helper base class for defining programs compatible with HCQ-compatible devices. Currently, the arguments consist of pointers to buffers, followed by `vals` fields. The convention expects a packed struct containing the passed pointers, followed by `vals` located at `kernargs_args_offset`. +`HCQProgram` is a base class for defining programs compatible with HCQ-enabled devices. It provides a flexible framework for handling different argument layouts (see `HCQArgsState`). ::: tinygrad.device.HCQProgram options: members: true show_source: false +#### Arguments State + +`HCQArgsState` is a base class for managing the argument state for HCQ programs. Backend implementations should create a subclass of `HCQArgsState` to manage arguments for the given program. + +::: tinygrad.device.HCQArgsState + options: + members: true + show_source: false + ### Synchronization HCQ-compatible devices use a global timeline signal for synchronizing all operations. This mechanism ensures proper ordering and completion of tasks across the device. By convention, `self.timeline_value` points to the next value to signal. So, to wait for all previous operations on the device to complete, wait for `self.timeline_value - 1` value. The following Python code demonstrates the typical usage of signals to synchronize execution to other operations on the device: diff --git a/tinygrad/device.py b/tinygrad/device.py index e75f7ae7f3..1e3cc84776 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -323,18 +323,18 @@ class HWComputeQueue(HWCommandQueue): def _memory_barrier(self): pass @hcq_command - def exec(self, prg:HCQProgram, kernargs:int, global_size:Tuple[int,int,int], local_size:Tuple[int,int,int]): + def exec(self, prg:HCQProgram, args_state:HCQArgsState, global_size:Tuple[int,int,int], local_size:Tuple[int,int,int]): """ Enqueues an execution command for a kernel program. Args: prg: The program to execute - kernargs: The pointer to kernel arguments + args_state: The args state to execute program with global_size: The global work size local_size: The local work size """ - self._exec(prg, kernargs, global_size, local_size) - def _exec(self, prg, kernargs, global_size, local_size): raise NotImplementedError("backend should overload this function") + self._exec(prg, args_state, global_size, local_size) + def _exec(self, prg, args_state, global_size, local_size): raise NotImplementedError("backend should overload this function") def update_exec(self, cmd_idx:int, global_size:Optional[Tuple[int,int,int]]=None, local_size:Optional[Tuple[int,int,int]]=None): """ @@ -431,25 +431,27 @@ def hcq_profile(dev, enabled, desc, queue_type=None, queue=None): if enabled and PROFILE: dev.sig_prof_records.append((st, en, desc, queue_type is dev.hw_copy_queue_t)) -class HCQProgram: - def __init__(self, device:HCQCompiled, name:str, kernargs_alloc_size:int, kernargs_args_offset:int=0): - self.device, self.name, self.kernargs_alloc_size, self.kernargs_args_offset = device, name, kernargs_alloc_size, kernargs_args_offset +class HCQArgsState: + def __init__(self, ptr:int, prg:HCQProgram, bufs:Tuple[HCQBuffer, ...], vals:Tuple[int, ...]=()): self.ptr, self.prg = ptr, prg + def update_buffer(self, index:int, buf:HCQBuffer): raise NotImplementedError("need update_buffer") + def update_var(self, index:int, val:int): raise NotImplementedError("need update_var") - def fill_kernargs(self, bufs:Tuple[HCQBuffer, ...], vals:Tuple[int, ...]=(), kernargs_ptr:Optional[int]=None) -> int: +class HCQProgram: + def __init__(self, args_state_t:Type[HCQArgsState], device:HCQCompiled, name:str, kernargs_alloc_size:int, kernargs_args_offset:int=0): + self.args_state_t, self.device, self.name = args_state_t, device, name + self.kernargs_alloc_size, self.kernargs_args_offset = kernargs_alloc_size, kernargs_args_offset + + def fill_kernargs(self, bufs:Tuple[HCQBuffer, ...], vals:Tuple[int, ...]=(), kernargs_ptr:Optional[int]=None) -> HCQArgsState: """ Fills arguments for the kernel, optionally allocating space from the device if `kernargs_ptr` is not provided. - Args: bufs: Buffers to be written to kernel arguments. vals: Values to be written to kernel arguments. kernargs_ptr: Optional pointer to pre-allocated kernel arguments memory. - Returns: - Pointer to the filled kernel arguments. + Arguments state with the given buffers and values set for the program. """ - self._fill_kernargs(ptr:=(kernargs_ptr or self.device._alloc_kernargs(self.kernargs_alloc_size)), bufs, vals) - return ptr - def _fill_kernargs(self, kernargs_ptr:int, bufs:Tuple[HCQBuffer, ...], vals:Tuple[int, ...]=()): raise NotImplementedError("need fill_kernargs") + return self.args_state_t(kernargs_ptr or self.device._alloc_kernargs(self.kernargs_alloc_size), self, bufs, vals=vals) def __call__(self, *bufs:HCQBuffer, global_size:Tuple[int,int,int]=(1,1,1), local_size:Tuple[int,int,int]=(1,1,1), vals:Tuple[int, ...]=(), wait:bool=False) -> Optional[float]: diff --git a/tinygrad/runtime/graph/hcq.py b/tinygrad/runtime/graph/hcq.py index b5a65ce5ea..cf0737081e 100644 --- a/tinygrad/runtime/graph/hcq.py +++ b/tinygrad/runtime/graph/hcq.py @@ -1,7 +1,7 @@ import collections, time from typing import List, Any, Dict, cast, Optional, Tuple, Set -from tinygrad.helpers import round_up, to_mv, PROFILE, memsize_to_str -from tinygrad.device import HCQCompiled, HCQAllocator, HCQSignal, HCQBuffer, HWCommandQueue, HWComputeQueue, HWCopyQueue, \ +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 @@ -18,20 +18,15 @@ class HCQGraph(MultiGraphRunner): 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, HCQBuffer] = {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] = {} + 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) - - ji.prg.clprg.fill_kernargs([cast(Buffer, b)._buf for b in ji.bufs], [var_vals[v] for v in ji.prg.p.vars], self.kargs_addrs[j]) - self.ji_args_bufs[j] = to_mv(self.kargs_addrs[j] + ji.prg.clprg.kernargs_args_offset, len(ji.bufs) * 8).cast('Q') - self.ji_args_vars[j] = to_mv(self.kargs_addrs[j] + ji.prg.clprg.kernargs_args_offset + len(ji.bufs) * 8, len(ji.prg.p.vars) * 4).cast('I') + 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 @@ -123,7 +118,7 @@ class HCQGraph(MultiGraphRunner): # Encode main commands based on ji type. if isinstance(ji.prg, CompiledRunner): - cast(HWComputeQueue, enqueue_queue).exec(ji.prg.clprg, self.kargs_addrs[j], *ji.prg.p.launch_dims(var_vals)) + 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]] cast(HCQAllocator, Device[src.device].allocator).map(dest._buf) @@ -157,11 +152,11 @@ class HCQGraph(MultiGraphRunner): # Update rawbuffers for (j,i),input_idx in self.input_replace.items(): - if j in self.ji_args_bufs: self.ji_args_bufs[j][i] = input_rawbuffers[input_idx]._buf.va_addr + 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, i, v in self.updated_vars(var_vals): self.ji_args_vars[j][i] = v + for j, i, v in self.updated_vars(var_vals): self.ji_args[j].update_var(i, v) # Update launch dims for j, global_dims, local_dims in self.updated_launch_dims(var_vals): diff --git a/tinygrad/runtime/ops_amd.py b/tinygrad/runtime/ops_amd.py index 4268180993..f71b55c04f 100644 --- a/tinygrad/runtime/ops_amd.py +++ b/tinygrad/runtime/ops_amd.py @@ -2,7 +2,7 @@ from __future__ import annotations from typing import Tuple, List, Any, cast import os, fcntl, ctypes, ctypes.util, functools, pathlib, mmap, errno, time, array, contextlib, decimal from dataclasses import dataclass -from tinygrad.device import HCQCompiled, HCQAllocator, HCQBuffer, HWComputeQueue, HWCopyQueue, \ +from tinygrad.device import HCQCompiled, HCQAllocator, HCQBuffer, HWComputeQueue, HWCopyQueue, HCQArgsState, \ HCQSignal, HCQProgram, BufferOptions from tinygrad.helpers import getenv, to_mv, round_up, data64_le, DEBUG, mv_address from tinygrad.renderer.cstyle import AMDRenderer @@ -101,18 +101,18 @@ class AMDComputeQueue(HWComputeQueue): nbioreg(regBIF_BX_PF1_GPU_HDP_FLUSH_DONE), 0xffffffff, 0xffffffff, 0x20] self._acquire_mem() - def _exec(self, prg, kernargs, global_size:Tuple[int,int,int]=(1,1,1), local_size:Tuple[int,int,int]=(1,1,1)): + def _exec(self, prg, args_state, global_size:Tuple[int,int,int]=(1,1,1), local_size:Tuple[int,int,int]=(1,1,1)): self._acquire_mem(gli=0, gl2=0) user_regs, cmd_idx = [], len(self) - 1 if prg.enable_dispatch_ptr: - dp = hsa.hsa_kernel_dispatch_packet_t.from_address(dp_addr:=kernargs + prg.kernargs_segment_size) + dp = hsa.hsa_kernel_dispatch_packet_t.from_address(dp_addr:=args_state.ptr + prg.kernargs_segment_size) dp.workgroup_size_x, dp.workgroup_size_y, dp.workgroup_size_z = local_size[0], local_size[1], local_size[2] dp.grid_size_x, dp.grid_size_y, dp.grid_size_z = global_size[0]*local_size[0], global_size[1]*local_size[1], global_size[2]*local_size[2] - dp.group_segment_size, dp.private_segment_size, dp.kernarg_address = prg.group_segment_size, prg.private_segment_size, kernargs + dp.group_segment_size, dp.private_segment_size, dp.kernarg_address = prg.group_segment_size, prg.private_segment_size, args_state.ptr user_regs += [*data64_le(dp_addr)] self.cmd_idx_to_dispatch_packet[cmd_idx] = dp - user_regs += [*data64_le(kernargs)] + user_regs += [*data64_le(args_state.ptr)] self.q += [amd_gpu.PACKET3(amd_gpu.PACKET3_SET_SH_REG, 6), gfxreg(amd_gpu.regCOMPUTE_PGM_LO), *data64_le(prg.prog_addr >> 8), *data64_le(0), *data64_le(prg.device.scratch.va_addr >> 8)] @@ -256,6 +256,19 @@ class AMDCopyQueue(HWCopyQueue): device.sdma_queue.write_ptr[0] = device.sdma_queue.put_value device.sdma_queue.doorbell[0] = device.sdma_queue.put_value +class AMDArgsState(HCQArgsState): + def __init__(self, ptr:int, prg:AMDProgram, bufs:Tuple[HCQBuffer, ...], vals:Tuple[int, ...]=()): + super().__init__(ptr, prg, bufs, vals=vals) + + self.bufs = to_mv(self.ptr, len(bufs) * 8).cast('Q') + self.vals = to_mv(self.ptr + len(bufs) * 8, len(vals) * 4).cast('I') + + self.bufs[:] = array.array('Q', [b.va_addr for b in bufs]) + self.vals[:] = array.array('I', vals) + + def update_buffer(self, index:int, buf:HCQBuffer): self.bufs[index] = buf.va_addr + def update_var(self, index:int, val:int): self.vals[index] = val + class AMDProgram(HCQProgram): def __init__(self, device:AMDDevice, name:str, lib:bytes): # TODO; this API needs the type signature of the function and global_size/local_size @@ -288,16 +301,11 @@ class AMDProgram(HCQProgram): self.enable_dispatch_ptr = code.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_DISPATCH_PTR additional_alloc_sz = ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if self.enable_dispatch_ptr else 0 - super().__init__(self.device, self.name, kernargs_alloc_size=self.kernargs_segment_size+additional_alloc_sz) + super().__init__(AMDArgsState, self.device, self.name, kernargs_alloc_size=self.kernargs_segment_size+additional_alloc_sz) def __del__(self): if hasattr(self, 'lib_gpu'): cast(AMDDevice, self.device)._gpu_free(self.lib_gpu) - def _fill_kernargs(self, kernargs_ptr:int, bufs:Tuple[Any, ...], vals:Tuple[int, ...]=()): - if (given:=len(bufs)*8 + len(vals)*4) != (want:=self.kernargs_segment_size): raise RuntimeError(f'incorrect args size {given=} != {want=}') - if len(bufs): to_mv(kernargs_ptr, len(bufs) * 8).cast('Q')[:] = array.array('Q', [b.va_addr for b in bufs]) - if len(vals): to_mv(kernargs_ptr + len(bufs) * 8, len(vals) * 4).cast('I')[:] = array.array('I', vals) - class AMDAllocator(HCQAllocator): def __init__(self, device:AMDDevice): super().__init__(device, batch_size=SDMA_MAX_COPY_SIZE) diff --git a/tinygrad/runtime/ops_nv.py b/tinygrad/runtime/ops_nv.py index 7c89bc7c2e..ab70152446 100644 --- a/tinygrad/runtime/ops_nv.py +++ b/tinygrad/runtime/ops_nv.py @@ -3,7 +3,7 @@ import os, ctypes, contextlib, re, fcntl, functools, mmap, struct, time, array, from typing import Tuple, List, Any, cast, Union, Dict, Type from dataclasses import dataclass from tinygrad.device import HCQCompiled, HCQAllocator, HCQBuffer, HWCommandQueue, HWComputeQueue, HWCopyQueue, hcq_command, \ - HCQProgram, HCQSignal, BufferOptions + HCQArgsState, HCQProgram, HCQSignal, BufferOptions from tinygrad.helpers import getenv, mv_address, init_c_struct_t, to_mv, round_up, data64, data64_le, DEBUG, prod from tinygrad.renderer.assembly import PTXRenderer from tinygrad.renderer.cstyle import NVRenderer @@ -142,17 +142,17 @@ class NVComputeQueue(NVCommandQueue, HWComputeQueue): def _memory_barrier(self): self.q += [nvmethod(1, nv_gpu.NVC6C0_INVALIDATE_SHADER_CACHES_NO_WFI, 1), (1 << 12) | (1 << 4) | (1 << 0)] - def _exec(self, prg, kernargs, global_size, local_size): + def _exec(self, prg, args_state, global_size, local_size): cmd_idx = len(self) - 1 - ctypes.memmove(qmd_addr:=(kernargs + round_up(prg.constbufs[0][1], 1 << 8)), ctypes.addressof(prg.qmd), 0x40 * 4) + ctypes.memmove(qmd_addr:=(args_state.ptr + round_up(prg.constbufs[0][1], 1 << 8)), ctypes.addressof(prg.qmd), 0x40 * 4) self.cmd_idx_to_qmd[cmd_idx] = qmd = qmd_struct_t.from_address(qmd_addr) # Save qmd for later update self.cmd_idx_to_global_dims[cmd_idx] = to_mv(qmd_addr + nv_gpu.NVC6C0_QMDV03_00_CTA_RASTER_WIDTH[1] // 8, 12).cast('I') self.cmd_idx_to_local_dims[cmd_idx] = to_mv(qmd_addr + nv_gpu.NVC6C0_QMDV03_00_CTA_THREAD_DIMENSION0[1] // 8, 6).cast('H') qmd.cta_raster_width, qmd.cta_raster_height, qmd.cta_raster_depth = global_size qmd.cta_thread_dimension0, qmd.cta_thread_dimension1, qmd.cta_thread_dimension2 = local_size - qmd.constant_buffer_addr_upper_0, qmd.constant_buffer_addr_lower_0 = data64(kernargs) + qmd.constant_buffer_addr_upper_0, qmd.constant_buffer_addr_lower_0 = data64(args_state.ptr) if (prev_qmd:=self.cmd_idx_to_qmd.get(cmd_idx - 1)) is None: self.q += [nvmethod(1, nv_gpu.NVC6C0_SEND_PCAS_A, 0x1), qmd_addr >> 8] @@ -210,6 +210,19 @@ class NVCopyQueue(NVCommandQueue, HWCopyQueue): def _submit(self, device): self._submit_to_gpfifo(device, cast(NVDevice, device).dma_gpfifo) +class NVArgsState(HCQArgsState): + def __init__(self, ptr:int, prg:NVProgram, bufs:Tuple[HCQBuffer, ...], vals:Tuple[int, ...]=()): + super().__init__(ptr, prg, bufs, vals=vals) + + if MOCKGPU: prg.constbuffer_0[0:2] = [len(bufs), len(vals)] + kernargs = [arg_half for arg in bufs for arg_half in data64_le(arg.va_addr)] + list(vals) + to_mv(self.ptr, (len(prg.constbuffer_0) + len(kernargs)) * 4).cast('I')[:] = array.array('I', prg.constbuffer_0 + kernargs) + self.bufs = to_mv(self.ptr + len(prg.constbuffer_0) * 4, len(bufs) * 8).cast('Q') + self.vals = to_mv(self.ptr + len(prg.constbuffer_0) * 4 + len(bufs) * 8, len(vals) * 4).cast('I') + + def update_buffer(self, index:int, buf:HCQBuffer): self.bufs[index] = buf.va_addr + def update_var(self, index:int, val:int): self.vals[index] = val + class NVProgram(HCQProgram): def __init__(self, device:NVDevice, name:str, lib:bytes): self.device, self.name, self.lib = device, name, lib @@ -266,17 +279,12 @@ class NVProgram(HCQProgram): self.max_threads = ((65536 // round_up(max(1, self.registers_usage) * 32, 256)) // 4) * 4 * 32 # NV's kernargs is constbuffer (size 0x160), then arguments to the kernel follows. Kernargs also appends QMD at the end of the kernel. - super().__init__(self.device, self.name, kernargs_alloc_size=round_up(self.constbufs[0][1], 1 << 8) + (8 << 8), kernargs_args_offset=0x160) + super().__init__(NVArgsState, self.device, self.name, + kernargs_alloc_size=round_up(self.constbufs[0][1], 1 << 8) + (8 << 8), kernargs_args_offset=0x160) def __del__(self): if hasattr(self, 'lib_gpu'): self.device.allocator.free(self.lib_gpu, self.lib_gpu.size, BufferOptions(cpu_access=True)) - def _fill_kernargs(self, kernargs_ptr:int, bufs:Tuple[Any, ...], vals:Tuple[int, ...]=()): - # HACK: Save counts of args and vars to "unused" constbuffer for later extraction in mockgpu to pass into gpuocelot. - if MOCKGPU: self.constbuffer_0[0:2] = [len(bufs), len(vals)] - kernargs = [arg_half for arg in bufs for arg_half in data64_le(arg.va_addr)] + list(vals) - to_mv(kernargs_ptr, (len(self.constbuffer_0) + len(kernargs)) * 4).cast('I')[:] = array.array('I', self.constbuffer_0 + kernargs) - def __call__(self, *bufs, global_size:Tuple[int,int,int]=(1,1,1), local_size:Tuple[int,int,int]=(1,1,1), vals:Tuple[int, ...]=(), wait=False): if prod(local_size) > 1024 or self.max_threads < prod(local_size): raise RuntimeError("Too many resources requsted for launch") if any(cur > mx for cur,mx in zip(global_size, [2147483647, 65535, 65535])) or any(cur > mx for cur,mx in zip(local_size, [1024, 1024, 64])):