From ab8fb191b2fee527258eb50108fad64101fd6587 Mon Sep 17 00:00:00 2001 From: Christopher Milan Date: Mon, 27 Jul 2026 23:56:09 -0400 Subject: [PATCH] 64-bit UOp.variable support (#17246) --- test/backend/test_tensor_variable.py | 16 ++++++++++++---- test/mockgpu/nv/nvgpu.py | 2 +- tinygrad/device.py | 8 +++++++- tinygrad/renderer/amd/elf.py | 7 +++---- tinygrad/renderer/cstyle.py | 12 +++++++----- tinygrad/renderer/nir.py | 18 +++++++++++------- tinygrad/runtime/graph/cuda.py | 2 +- tinygrad/runtime/graph/metal.py | 19 +++++++++++-------- tinygrad/runtime/ops_amd.py | 3 +-- tinygrad/runtime/ops_cl.py | 2 +- tinygrad/runtime/ops_cpu.py | 4 ++-- tinygrad/runtime/ops_cuda.py | 11 ++++++----- tinygrad/runtime/ops_dsp.py | 22 +++++++++++++--------- tinygrad/runtime/ops_hip.py | 7 ++++--- tinygrad/runtime/ops_metal.py | 5 +++-- tinygrad/runtime/ops_nv.py | 8 +++++--- tinygrad/runtime/ops_qcom.py | 8 +++++--- tinygrad/runtime/support/hcq.py | 13 +++++++------ tinygrad/uop/ops.py | 2 +- tinygrad/uop/spec.py | 3 ++- 20 files changed, 103 insertions(+), 69 deletions(-) diff --git a/test/backend/test_tensor_variable.py b/test/backend/test_tensor_variable.py index f7b92295ec..8544d04e2b 100644 --- a/test/backend/test_tensor_variable.py +++ b/test/backend/test_tensor_variable.py @@ -1,6 +1,6 @@ import unittest import numpy as np -from tinygrad import Tensor, Variable, dtypes +from tinygrad import Device, Tensor, Variable, TinyJit, dtypes from tinygrad.helpers import CHECK_OOB class TestTensorVariable(unittest.TestCase): @@ -18,10 +18,18 @@ class TestTensorVariable(unittest.TestCase): self.assertListEqual((vv * t).tolist(), [2, 2, 2]) except RuntimeError: pass - # TODO: a Variable PARAM lowers to int32, so a bound value that doesn't fit int32 truncates or fails to bind - @unittest.expectedFailure + @unittest.skipUnless(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires long support") def test_large_range_variable(self): - self.assertEqual(Tensor(Variable("b", 0, 2**40).bind(2**35)).item(), 2**35) + self.assertEqual(Tensor(Variable("b", 0, 2**40, dtype=dtypes.long).bind(2**35)).clone(Device.DEFAULT).item(), 2**35) + + @unittest.skipUnless(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires long support") + def test_large_range_variable_jit(self): + @TinyJit + def f(a,b): return (Tensor(a+b).clone(Device.DEFAULT) * 2).realize() + for i in range(3): + a = Variable("a", 0, 2**10, dtype=dtypes.int).bind(i) + b = Variable("b", 0, 2**40, dtype=dtypes.long).bind(2**35) + self.assertEqual(f(a,b).item(), (2**35 + i) * 2) def test_variable_defers_like_a_literal(self): vv = Variable("a", 1, 10).bind(2) diff --git a/test/mockgpu/nv/nvgpu.py b/test/mockgpu/nv/nvgpu.py index eb6cf7f7ff..84893f581b 100644 --- a/test/mockgpu/nv/nvgpu.py +++ b/test/mockgpu/nv/nvgpu.py @@ -91,7 +91,7 @@ class GPFIFO: args_cnt, vals_cnt = const0[80], const0[81] args_addr = qmd.constant_buffer_addr_lower_0 + (qmd.constant_buffer_addr_upper_0 << 32) + 0x160 args = to_mv(args_addr, args_cnt*8).cast('Q') - vals = to_mv(args_addr + args_cnt*8, vals_cnt*4).cast('I') + vals = to_mv(args_addr + args_cnt*8, vals_cnt*8).cast('Q') cargs = [ctypes.cast(args[i], ctypes.c_void_p) for i in range(args_cnt)] + [ctypes.cast(vals[i], ctypes.c_void_p) for i in range(vals_cnt)] gx, gy, gz = qmd.cta_raster_width, qmd.cta_raster_height, qmd.cta_raster_depth lx, ly, lz = qmd.cta_thread_dimension0, qmd.cta_thread_dimension1, qmd.cta_thread_dimension2 diff --git a/tinygrad/device.py b/tinygrad/device.py index 265e0a507b..570cd79813 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -5,7 +5,7 @@ from typing import Any, Generic, TypeVar, Iterator, Generator, Self, TYPE_CHECKI import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal from tinygrad.helpers import LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, PROFILE, temp, colored from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing -from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize, Target, unwrap +from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize, Target, unwrap, round_up from tinygrad.dtype import DType, _to_np_dtype if TYPE_CHECKING: from tinygrad.renderer import Renderer @@ -319,6 +319,12 @@ class TinyELF: # tuple of (name, slot, dtype, shape) signature: tuple[tuple[str|None, int, DType, tuple], ...] + @staticmethod + def iter_sig(signature:tuple[tuple[str|None, int, DType, tuple], ...], offset:int=0) -> Generator[tuple[int, DType], None, None]: + for _,_,dt,_ in signature: + yield (offset:=round_up(offset, dt.itemsize)), dt + offset += dt.itemsize + class Program(Generic[DeviceType]): def __init__(self, dev:DeviceType, obj:TinyELF): pass 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, ...]=(), diff --git a/tinygrad/renderer/amd/elf.py b/tinygrad/renderer/amd/elf.py index 18f0aed8a8..718700cf84 100644 --- a/tinygrad/renderer/amd/elf.py +++ b/tinygrad/renderer/amd/elf.py @@ -35,10 +35,9 @@ def assemble_linear(prg:UOp, lin:UOp, arch:str) -> bytes: elif val.offset < 106: max_sgpr = max(max_sgpr, val.offset + val.sz) # ** scan sink for metadata - sink, n_bufs, n_vars, lds_size, gids = prg.src[0], 0, 0, 0, set() + sink, param_sizes, lds_size, gids = prg.src[0], {}, 0, set() for u in sink.toposort(): - if u.op is Ops.PARAM and u.addrspace is AddrSpace.ALU: n_vars += 1 - elif u.op is Ops.PARAM: n_bufs += 1 + if u.op is Ops.PARAM: param_sizes[u.arg.slot] = u.dtype.itemsize if u.addrspace is AddrSpace.ALU else 8 elif u.op is Ops.BUFFER and u.addrspace is AddrSpace.LOCAL: lds_size += u.max_numel() * u.dtype.itemsize elif u.op is Ops.SPECIAL and u.arg.startswith("gidx"): gids.add(int(u.arg[-1])) code_bytes = b"".join(inst.to_bytes() for inst in insts) @@ -60,7 +59,7 @@ def assemble_linear(prg:UOp, lin:UOp, arch:str) -> bytes: sgpr_granule = max(0, ceildiv(next_free_sgpr + 6, 8) - 1) if is_cdna else 0 desc = amdgpu_kd.llvm_amdhsa_kernel_descriptor_t() desc.group_segment_fixed_size = lds_size - desc.kernarg_size = n_bufs * 8 + n_vars * 4 + for sz in (param_sizes[i] for i in sorted(param_sizes)): desc.kernarg_size = round_up(desc.kernarg_size, sz) + sz desc.kernel_code_entry_byte_offset = -len(text) # https://llvm.org/docs/AMDGPUUsage.html#amdgpu-amdhsa-compute-pgm-rsrc1-gfx6-gfx12-table diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 1ce2024c0d..fb56cd1b53 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -123,7 +123,8 @@ class CStyleLanguage(Renderer): smem_align: str = "" smem_prefix: str = "" smem_prefix_for_cast: bool = True - arg_int_prefix: str = "const int" + var_prefix: str = "const " + var_suffix: str = "" barrier: str = "" code_for_workitem: dict[Literal["g", "l", "i"], Callable] = {} extra_args: list[str] = [] @@ -149,9 +150,9 @@ class CStyleLanguage(Renderer): tmp = "" if any(is_image_shape(u._shape) for _,(u,_) in bufs): tmp = "const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;\n" - buftypes = [(name, ("volatile " if u.arg.volatile else "")+ - self._render_dtype(u.dtype, sz=1, addrspace=u.addrspace, mutable=mutable, shape=u._shape)+self.buffer_suffix \ - if u.addrspace == AddrSpace.GLOBAL else self.arg_int_prefix if u.dtype == dtypes.int else None) for name,(u,mutable) in bufs] + buftypes = [(name, ("volatile " if u.arg.volatile else "")+(self.var_prefix if u.addrspace == AddrSpace.ALU else "")+ + self._render_dtype(u.dtype, sz=1, addrspace=u.addrspace, mutable=mutable, shape=u._shape)+ + (self.var_suffix if u.addrspace == AddrSpace.ALU else self.buffer_suffix)) for name,(u,mutable) in bufs] local_dims = [u.src[0] for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"] launch_bounds = prod([d.vmax for d in local_dims]) prg = ''.join([f"{self.kernel_typedef.format(launch_bounds=launch_bounds)} {function_name}(",] + @@ -349,7 +350,8 @@ class MetalRenderer(CStyleLanguage): kernel_typedef = "kernel void" buffer_prefix = "device " smem_prefix = "threadgroup __attribute__((aligned(16))) " - arg_int_prefix = "constant int&" + var_prefix = "constant " + var_suffix = "&" barrier = "threadgroup_barrier(mem_flags::mem_threadgroup);" float4 = "float4" code_for_workitem = {"g": lambda x: f"gid.{chr(120+int(x))}", "l": lambda x: f"lid.{chr(120+int(x))}"} diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index c45a293ae8..15352a1dda 100644 --- a/tinygrad/renderer/nir.py +++ b/tinygrad/renderer/nir.py @@ -1,6 +1,6 @@ from typing import Callable, Any from tinygrad.dtype import AddrSpace, DType, dtypes, truncate -from tinygrad.helpers import DEBUG, OSX, unwrap, fromimport, Target, is_image_shape +from tinygrad.helpers import DEBUG, OSX, unwrap, fromimport, Target, is_image_shape, round_up from tinygrad.renderer import Renderer from tinygrad.renderer.cstyle import CUDARenderer from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str @@ -246,9 +246,11 @@ class NIRRenderer(Renderer): def supported_dtypes(self): return {d for d in Renderer.supported_dtypes(self) if d not in dtypes.fp8s+(dtypes.bfloat16,)} +def padded_idx(param_idx:int, size:int): return round_up(param_idx, size) + size + class NAKRenderer(NIRRenderer): - param = nir_instr(nc=1, num_components=1, bs=lambda sz:sz*8, also=lambda self,sz: setattr(self, "param_idx", self.param_idx + sz), - intrins={"ALIGN_MUL":lambda sz:sz}, srcs=lambda self,b: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, self.param_idx, dtypes.int))])( + param = nir_instr(nc=1, num_components=1, bs=lambda sz:sz*8, also=lambda self,sz: setattr(self, "param_idx", padded_idx(self.param_idx, sz)), + intrins={"ALIGN_MUL":lambda sz:sz}, srcs=lambda self,b,sz: [nsrc(nimm(b,0,dtypes.int)), nsrc(nimm(b, round_up(self.param_idx,sz), dtypes.int))])( lambda self, b, x, sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_ldc_nv)) def supported_dtypes(self): return {d for d in super().supported_dtypes() if (d != dtypes.half or int(self.target.arch[3:]) >= 53)} @@ -263,12 +265,13 @@ class LVPRenderer(NIRRenderer): code_for_op = {k:v for k,v in NIRRenderer.code_for_op.items() if k != Ops.EXP2} param = nir_instr(nc=1, bs=lambda sz: sz * 8, num_components=1, intrins={"ALIGN_MUL":lambda sz: sz, "RANGE":lambda self: self.param_sz}, - srcs=lambda b, self: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, self.param_idx, dtypes.int))], also=lambda self, sz: - setattr(self, "param_idx", self.param_idx+sz))(lambda self,b,x,sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_ubo)) + srcs=lambda b,self,sz: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, round_up(self.param_idx, sz), dtypes.int))], also=lambda self, sz: + setattr(self, "param_idx", padded_idx(self.param_idx, sz)))(lambda self,b,x,sz: + mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_ubo)) def prerender(self, uops:list[UOp]): super().prerender(uops) - self.param_sz = sum([u.dtype.itemsize if u.addrspace is AddrSpace.ALU else 8 for u in uops if u.op is Ops.PARAM]) + self.param_sz = functools.reduce(padded_idx, (u.element_size() if u.addrspace is AddrSpace.ALU else 8 for u in uops if u.op is Ops.PARAM), 0) def tovec(b, idx_y, idx_x): return nalu(b, "vec4", idx_x, idx_y, nundef(b, dtypes.int), nundef(b, dtypes.int)) def nfloat(dtype): return mesa.nir_type_float16 if dtype == dtypes.half else mesa.nir_type_float32 @@ -306,7 +309,8 @@ class IR3Renderer(NIRRenderer): super().prerender(uops) self.texs:set[UOp] = set() self.img_idx = 0 - self.param_sz = sum([u.dtype.itemsize if u.addrspace is AddrSpace.ALU else 8 for u in uops if u.op is Ops.PARAM]) + self.param_sz = functools.reduce(padded_idx, (u.element_size() if u.addrspace is AddrSpace.ALU else 8 + for u in uops if u.op is Ops.PARAM and not is_image_shape(u._shape)), 0) def postrender(self, uops:list[UOp]): bufs = [u for u in uops if u.op is Ops.PARAM and u.addrspace is not AddrSpace.ALU] diff --git a/tinygrad/runtime/graph/cuda.py b/tinygrad/runtime/graph/cuda.py index 11035888c0..ed228d1350 100644 --- a/tinygrad/runtime/graph/cuda.py +++ b/tinygrad/runtime/graph/cuda.py @@ -20,7 +20,7 @@ class CUDAGraph(MultiGraphRunner): global_size, local_size = ast.arg.launch_dims({v: 0 for v in self.vars}) c_deps, new_node = self.new_node([b.base for b in bufs], ast.arg.outs) - c_args, vargs = encode_args([b._buf for b in bufs], [device_vars.get(x.expr, 0) for x in ast.arg.vars]) + c_args, vargs = encode_args([b._buf for b in bufs], [device_vars.get(x.expr, 0) for x in ast.arg.vars], runtime.signature) kern_params = cuda.CUDA_KERNEL_NODE_PARAMS_v1(runtime.prg, *global_size, *local_size, runtime.smem, ctypes.cast(0, ctypes.POINTER(ctypes.c_void_p)), vargs) check(cuda.cuGraphAddKernelNode(ctypes.byref(new_node), self.graph, c_deps, len(c_deps or []), ctypes.byref(kern_params))) diff --git a/tinygrad/runtime/graph/metal.py b/tinygrad/runtime/graph/metal.py index 228275d787..8340023d40 100644 --- a/tinygrad/runtime/graph/metal.py +++ b/tinygrad/runtime/graph/metal.py @@ -1,7 +1,6 @@ from typing import Any, cast -import ctypes, decimal -from tinygrad.dtype import dtypes -from tinygrad.helpers import dedup, getenv, PROFILE +import ctypes, decimal, struct +from tinygrad.helpers import dedup, getenv, unwrap, PROFILE from tinygrad.device import Buffer, Device, ProfileGraphEntry, ProfileGraphEvent from tinygrad.uop.ops import UOp, Ops from tinygrad.engine.jit import GraphRunner, GraphException @@ -25,9 +24,11 @@ class MetalGraph(GraphRunner): if self.icb.value is None: raise GraphException("create indirect command buffer failed, does your system support this?") self.needs_icb_fix = int(not self.dev.arch.startswith("Apple") or int(self.dev.arch[5:]) < 9) # ICB fix not required on M3+ (Apple9+) - if len(self.vars): self.int_buf = self.dev.allocator.alloc(len(self.vars)*dtypes.int32.itemsize) + if len(self.vars): + self.var_buf = self.dev.allocator.alloc(sum(dt.itemsize for r in self.runtimes for (_,_,dt,s) in unwrap(r).signature if s == ())) + self.var_bind_data, self.var_buf_view, var_buf_offset = [], cast(MetalAllocator, self.dev.allocator)._as_buffer(self.var_buf), 0 - all_pipelines, all_resources = [], [self.int_buf.buf] if len(self.vars) else [] + all_pipelines, all_resources = [], [self.var_buf.buf] if len(self.vars) else [] for j, ((_, ast, bufs, _), runtime, replace) in enumerate(zip(self.calls, self.runtimes, self.uop_replace)): assert runtime is not None icb_command = self.icb.indirectComputeCommandAtIndex(j).retained() @@ -37,7 +38,10 @@ class MetalGraph(GraphRunner): if not any(pos == i for pos, _ in replace): icb_command.setKernelBuffer_offset_atIndex(b._buf.buf, b._buf.offset, i) all_resources.append(b._buf.buf) - for i, v in enumerate(ast.arg.vars): icb_command.setKernelBuffer_offset_atIndex(self.int_buf.buf, self.vars.index(v.expr)*4, len(bufs)+i) + for nm,i,dt,_ in runtime.signature[len(bufs):]: + icb_command.setKernelBuffer_offset_atIndex(self.var_buf.buf, var_buf_offset, i) + self.var_bind_data.append((nm, var_buf_offset, dt.fmt)) + var_buf_offset += dt.itemsize global_size, local_size = ast.arg.launch_dims({v: 0 for v in self.vars}) icb_command.concurrentDispatchThreadgroups_threadsPerThreadgroup(metal.MTLSize(*global_size), metal.MTLSize(*local_size)) icb_command.setBarrier() @@ -45,7 +49,6 @@ class MetalGraph(GraphRunner): self.all_resources = dedup(all_resources) self.all_pipelines = dedup(all_pipelines) self.command_buffer: Any = None - if len(self.vars): self.int_buf_view = cast(MetalAllocator, self.dev.allocator)._as_buffer(self.int_buf).cast('i') self.range = metal.NSRange(0, len(self.calls)) self.updatable = sorted({j for j,r in enumerate(self.uop_replace) if r} | self.var_vals_replace.keys() | self.launch_dims_replace.keys()) @@ -66,7 +69,7 @@ class MetalGraph(GraphRunner): for j, global_dims, local_dims in self.updated_launch_dims(var_vals): self.icb.indirectComputeCommandAtIndex(j).concurrentDispatchThreadgroups_threadsPerThreadgroup(metal.MTLSize(*global_dims), metal.MTLSize(*local_dims)) - for i, var in enumerate(self.vars): self.int_buf_view[i] = var_vals[var] + for nm,ofs,fmt in self.var_bind_data: struct.pack_into(fmt, self.var_buf_view, ofs, var_vals[nm]) command_buffer = self.dev.mtl_queue.commandBuffer().retained() encoder = command_buffer.computeCommandEncoder().retained() diff --git a/tinygrad/runtime/ops_amd.py b/tinygrad/runtime/ops_amd.py index 0a7ceaaa49..a9a1d610f3 100644 --- a/tinygrad/runtime/ops_amd.py +++ b/tinygrad/runtime/ops_amd.py @@ -602,8 +602,7 @@ class AMDProgram(HCQProgram['AMDDevice']): if dev.sqtt_enabled: self.libhash: tuple[int, int] = struct.unpack(' float|None: for i, (_, slot, dt, shape) in enumerate(self.signature): - b = bufs[slot] if slot < len(bufs) else ctypes.c_int32(vals[slot-len(bufs)]) + b = bufs[slot] if slot < len(bufs) else getattr(ctypes, f"c_int{dt.bitsize}")(vals[slot-len(bufs)]) if is_image_shape(shape): pitch = (round_up(shape[1], 256) if OSX else shape[1]) * 4 * dt.itemsize fmt = cl.cl_image_format(cl.CL_RGBA, {2:cl.CL_HALF_FLOAT, 4:cl.CL_FLOAT}[dt.itemsize]) diff --git a/tinygrad/runtime/ops_cpu.py b/tinygrad/runtime/ops_cpu.py index d43b2c1c53..173e907712 100644 --- a/tinygrad/runtime/ops_cpu.py +++ b/tinygrad/runtime/ops_cpu.py @@ -105,7 +105,7 @@ class CPUProgram(HCQProgram['CPUDevice']): except OSError: pass def __init__(self, dev:CPUDevice, obj:TinyELF): - self.runtimevars = {name:slot for name,slot,*_ in obj.signature if name == 'core_id'} + self.signature, self.runtimevars = obj.signature, {name:slot for name,slot,*_ in obj.signature if name == 'core_id'} LVP = obj.target.renderer == "LVP" if sys.platform == "win32": # mypy doesn't understand when WIN is used here @@ -140,7 +140,7 @@ class CPUProgram(HCQProgram['CPUDevice']): self.fxn = ctypes.CFUNCTYPE(None)(self.addr) - super().__init__(LVPArgsState if LVP else HCQArgsState, dev, obj.name, kernargs_alloc_size=12+256 if LVP else 0) + super().__init__(LVPArgsState if LVP else HCQArgsState, dev, obj, kernargs_alloc_size=12+256 if LVP else 0) @suppress_finalizing def __del__(self): diff --git a/tinygrad/runtime/ops_cuda.py b/tinygrad/runtime/ops_cuda.py index 772d1962b8..eed8996b36 100644 --- a/tinygrad/runtime/ops_cuda.py +++ b/tinygrad/runtime/ops_cuda.py @@ -15,9 +15,10 @@ def check(status): error = ctypes.string_at(init_c_var(ctypes.POINTER(ctypes.c_char), lambda x: cuda.cuGetErrorString(status, ctypes.byref(x)))).decode() raise RuntimeError(f"CUDA Error {status}, {error}") -def encode_args(args, vals) -> tuple[ctypes.Structure, ctypes.Array]: - c_args = init_c_struct_t(len(args) * 8 + len(vals) * 4, tuple([(f'f{i}', cuda.CUdeviceptr_v2, i*8) for i in range(len(args))] + - [(f'v{i}', ctypes.c_int, len(args)*8 + i*4) for i in range(len(vals))]))(*args, *vals) +def encode_args(args, vals, signature) -> tuple[ctypes.Structure, ctypes.Array]: + fields = ([(f'f{i}', cuda.CUdeviceptr_v2, i*8) for i in range(len(args))] + + [(f'v{i}', getattr(ctypes, f"c_int{dt.bitsize}"), off) for i,(off,dt) in enumerate(TinyELF.iter_sig(signature[len(args):], len(args)*8))]) + c_args = init_c_struct_t(fields[-1][2] + ctypes.sizeof(fields[-1][1]) if len(fields) else 0, tuple(fields))(*args, *vals) vargs = (ctypes.c_void_p * 5)(ctypes.c_void_p(1), ctypes.cast(ctypes.byref(c_args), ctypes.c_void_p), ctypes.c_void_p(2), ctypes.cast(ctypes.pointer(ctypes.c_size_t(ctypes.sizeof(c_args))), ctypes.c_void_p), ctypes.c_void_p(0)) return c_args, vargs @@ -35,7 +36,7 @@ def cu_time_execution(cb, enable=False) -> float|None: class CUDAProgram(Program['CUDADevice']): def __init__(self, dev:CUDADevice, obj:TinyELF, smem:int=0): - self.dev, self.name, self.lib, self.smem = dev, obj.name, obj.lib, smem + self.dev, self.name, self.lib, self.signature, self.smem = dev, obj.name, obj.lib, obj.signature, smem if DEBUG >= 5: print("\n".join([f"{i+1:>3} {line}" for i, line in enumerate(pretty_ptx(obj.lib.decode('utf-8')).split("\n"))])) check(cuda.cuCtxSetCurrent(self.dev.context)) @@ -54,7 +55,7 @@ class CUDAProgram(Program['CUDADevice']): def __call__(self, *args, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw): check(cuda.cuCtxSetCurrent(self.dev.context)) if not hasattr(self, "vargs"): - self.c_args, self.vargs = encode_args(args, vals) + self.c_args, self.vargs = encode_args(args, vals, self.signature) # HACK: For MOCKGPU send the args struct itself. if MOCKGPU: self.vargs = self.c_args # type: ignore[assignment] diff --git a/tinygrad/runtime/ops_dsp.py b/tinygrad/runtime/ops_dsp.py index c508dda051..943c3b3c4e 100644 --- a/tinygrad/runtime/ops_dsp.py +++ b/tinygrad/runtime/ops_dsp.py @@ -4,7 +4,7 @@ assert sys.platform != 'win32' from tinygrad.device import BufferSpec, Compiled, Allocator, Compiler, Program, TinyELF from tinygrad.dtype import dtypes, AddrSpace from tinygrad.uop.ops import Ops, UOp -from tinygrad.helpers import getenv, round_up, mv_address, to_mv, cpu_objdump, system, DEBUG, suppress_finalizing, Target +from tinygrad.helpers import getenv, round_up, mv_address, to_mv, cpu_objdump, system, DEBUG, suppress_finalizing, Target, unwrap from tinygrad.renderer.cstyle import ClangRenderer from tinygrad.runtime.autogen import libc, qcom_dsp if getenv("IOCTL"): import extra.dsp.run # noqa: F401 # pylint: disable=unused-import @@ -49,7 +49,9 @@ class DSPRenderer(ClangRenderer): 'struct dcvs_v2_req req = {.type=7, .dcvs_enable=0, .set_latency=1, .latency=100, .set_dcvs_params=1, .target_corner = 6 /* TURBO */};', 'HAP_power_set((void*)handle, (void*)&req);'] msrc += ['if ((sc>>24) != 2) return 0;'] - msrc += [f'int sz_or_val_{i} = ((int*)pra[0].buf.pv)[{i}];' for i,b in enumerate(bufs)] + msrc += [f'{self._render_dtype(b[1][0].dtype) if b[1][0].addrspace == AddrSpace.ALU else "int"} sz_or_val_{i} = ' + f'*({self._render_dtype(b[1][0].dtype) if b[1][0].addrspace == AddrSpace.ALU else "int"}*)((char*)pra[0].buf.pv+{i*8});' + for i,b in enumerate(bufs)] msrc += [f'int off{i} = ((int*)pra[1].buf.pv)[{i}];' for i,b in enumerate(bufs) if b[1][0].addrspace == AddrSpace.GLOBAL] msrc += [f'void *buf_{i} = HAP_mmap(0,sz_or_val_{i},3,0,pra[{i+3}].dma.fd,0)+off{i};' for i,b in enumerate(bufs) if b[1][0].addrspace == AddrSpace.GLOBAL] @@ -75,15 +77,15 @@ def rpc_prep_args(ins=None, outs=None, in_fds=None): return pra, fds, attrs, (ins, outs) class DSPProgram(Program['DSPDevice']): - def __init__(self, dev:DSPDevice, obj:TinyELF): - self.dev, self.lib = dev, obj.lib + def __init__(self, dev:DSPDevice, obj:TinyELF): self.dev, self.lib, self.signature = dev, obj.lib, obj.signature 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, **kw): if len(bufs) >= 16: raise RuntimeError(f"Too many buffers to execute: {len(bufs)}") - pra, fds, attrs, _ = rpc_prep_args(ins=[var_vals_mv:=memoryview(bytearray((len(bufs)+len(vals))*4)), off_mv:=memoryview(bytearray(len(bufs)*4))], + pra, fds, attrs, _ = rpc_prep_args(ins=[var_vals_mv:=memoryview(bytearray((len(bufs)+len(vals))*8)), off_mv:=memoryview(bytearray(len(bufs)*4))], outs=[timer:=memoryview(bytearray(8)).cast('Q')], in_fds=[b.share_info.fd for b in bufs]) - var_vals_mv.cast('i')[:] = array.array('i', tuple(b.size for b in bufs) + vals) + for i,b in enumerate(bufs): struct.pack_into('i', var_vals_mv, i*8, b.size) + for i,(v,(_,_,dt,_)) in enumerate(zip(vals, self.signature[len(bufs):]), start=len(bufs)): struct.pack_into(unwrap(dt.fmt), var_vals_mv, i*8, v) off_mv.cast('I')[:] = array.array('I', tuple(b.offset for b in bufs)) self.dev.exec_lib(self.lib, rpc_sc(method=2, ins=2, outs=1, fds=len(bufs)), pra, fds, attrs) return timer[0] / 1e6 @@ -277,7 +279,7 @@ class MockDSPRenderer(DSPRenderer): # for loop for big reads msrc.append(f"void *buf{i} = mmap2(0, {sz}, 3, 0x21, -1, 0); for(int rd = 0; rd < {sz}; rd += read(0, buf{i}+rd, {sz}-rd));") else: - msrc.append(f"unsigned int val{i}; read(0, &val{i}, 4);") + msrc.append(f"{self._render_dtype(b[1][0].dtype)} val{i}; read(0, &val{i}, {b[1][0].dtype.itemsize});") msrc.append("unsigned int st = inscount();") params = [(f'(void*)buf{i}' if b[1][0].addrspace == AddrSpace.GLOBAL else f'val{i}') for i,b in enumerate(bufs)] msrc.append(f"{function_name}({', '.join(params)});") @@ -288,14 +290,16 @@ class MockDSPRenderer(DSPRenderer): return '\n'.join(msrc) class MockDSPProgram(Program[DSPDevice]): - def __init__(self, dev:DSPDevice, obj:TinyELF): self.lib = obj.lib + def __init__(self, dev:DSPDevice, obj:TinyELF): self.lib, self.signature = obj.lib, obj.signature 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, **kw): with tempfile.NamedTemporaryFile(suffix=".out") as dsp_lib: dsp_lib.write(self.lib) dsp_lib.flush() os.chmod(dsp_lib.name, 0o0777) proc = subprocess.run(["qemu-hexagon-static", *(['-strace'] if DEBUG >= 5 else []), dsp_lib.name], - input=b''.join([bytes(to_mv(x.va_addr, x.size)) for x in bufs] + [struct.pack("I", x) for x in vals]), stdout=subprocess.PIPE, check=True) + input=b''.join([bytes(to_mv(x.va_addr, x.size)) for x in bufs] + + [struct.pack(unwrap(dt.fmt), x) for x,(_,_,dt,_) in zip(vals, self.signature[len(bufs):])]), + stdout=subprocess.PIPE, check=True) offset = 4 for x in bufs: to_mv(x.va_addr, x.size)[:] = proc.stdout[offset:offset+x.size] diff --git a/tinygrad/runtime/ops_hip.py b/tinygrad/runtime/ops_hip.py index 62c11c8b35..81480f13e9 100644 --- a/tinygrad/runtime/ops_hip.py +++ b/tinygrad/runtime/ops_hip.py @@ -25,7 +25,7 @@ class HIPDevice(Compiled): class HIPProgram(Program[HIPDevice]): def __init__(self, dev:HIPDevice, obj:TinyELF): - self.dev, self.name, self.lib = dev, obj.name, obj.lib + self.dev, self.name, self.lib, self.signature = dev, obj.name, obj.lib, obj.signature check(hip.hipSetDevice(self.dev.device_id)) self.module = init_c_var(hip.hipModule_t, lambda x: check(hip.hipModuleLoadData(ctypes.byref(x), obj.lib))) self.prg = init_c_var(hip.hipFunction_t, lambda x: check(hip.hipModuleGetFunction(ctypes.byref(x), self.module, obj.name.encode("utf-8")))) @@ -37,8 +37,9 @@ class HIPProgram(Program[HIPDevice]): def __call__(self, *args, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw): check(hip.hipSetDevice(self.dev.device_id)) if not hasattr(self, "vargs"): - fields = [(f'f{i}', hip.hipDeviceptr_t, i*8) for i in range(len(args))] + [(f'v{i}', ctypes.c_int, len(args)*8+i*4) for i in range(len(vals))] - self.c_args = init_c_struct_t(len(args)*8+len(vals)*4, tuple(fields))(*args, *vals) + fields = ([(f'f{i}', hip.hipDeviceptr_t, i*8) for i in range(len(args))] + + [(f'v{i}', getattr(ctypes, f"c_int{dt.bitsize}"), o) for i,(o,dt) in enumerate(TinyELF.iter_sig(self.signature[len(args):], len(args)*8))]) + self.c_args = init_c_struct_t(fields[-1][2] + ctypes.sizeof(fields[-1][1]) if len(fields) else 0, tuple(fields))(*args, *vals) self.vargs = (ctypes.c_void_p * 5)(1, ctypes.cast(ctypes.byref(self.c_args), ctypes.c_void_p), 2, ctypes.cast(ctypes.pointer(ctypes.c_size_t(ctypes.sizeof(self.c_args))), ctypes.c_void_p), 3) diff --git a/tinygrad/runtime/ops_metal.py b/tinygrad/runtime/ops_metal.py index 06a9b3aca7..3a0fe42d2c 100644 --- a/tinygrad/runtime/ops_metal.py +++ b/tinygrad/runtime/ops_metal.py @@ -113,7 +113,7 @@ class MetalCompiler(Compiler): class MetalProgram(Program[MetalDevice]): def __init__(self, dev:MetalDevice, obj:TinyELF): - self.dev, self.name, self.lib = dev, obj.name, obj.lib + self.dev, self.name, self.lib, self.signature = dev, obj.name, obj.lib, obj.signature data = objc.dispatch_data_create(obj.lib, len(obj.lib), None, None) self.library = self.dev.sysdevice.newLibraryWithData_error(data, ctypes.byref(error_lib:=metal.NSError().retained())).retained() error_check(error_lib) @@ -138,7 +138,8 @@ class MetalProgram(Program[MetalDevice]): encoder = command_buffer.computeCommandEncoder().retained() encoder.setComputePipelineState(self.pipeline_state) for i,a in enumerate(bufs): encoder.setBuffer_offset_atIndex(a.buf, a.offset, i) - for i,a in enumerate(vals, start=len(bufs)): encoder.setBytes_length_atIndex(bytes(ctypes.c_int(a)), 4, i) + for a,(_,i,dt,_) in zip(vals, self.signature[len(bufs):]): + encoder.setBytes_length_atIndex(bytes(getattr(ctypes, f"c_int{dt.bitsize}")(a)), dt.itemsize, i) encoder.dispatchThreadgroups_threadsPerThreadgroup(metal.MTLSize(*global_size), metal.MTLSize(*local_size)) encoder.endEncoding() command_buffer.setLabel(to_ns_str(self.name)) # TODO: is this always needed? diff --git a/tinygrad/runtime/ops_nv.py b/tinygrad/runtime/ops_nv.py index 3e781e1fb1..2c274a6dd7 100644 --- a/tinygrad/runtime/ops_nv.py +++ b/tinygrad/runtime/ops_nv.py @@ -240,8 +240,10 @@ class NVVideoQueue(NVCommandQueue): class NVArgsState(CLikeArgsState): def __init__(self, buf:HCQBuffer, prg:NVProgram, bufs:tuple[HCQBuffer, ...], vals:tuple[int, ...]=()): - if isinstance(prg.dev.iface, MOCKIface): prg.cbuf_0[80:82] = [len(bufs), len(vals)] - super().__init__(buf, prg, bufs, vals=vals, prefix=prg.cbuf_0 or None) + if (is_mock:=isinstance(prg.dev.iface, MOCKIface)): prg.cbuf_0[80:82] = [len(bufs), len(vals)] + super().__init__(buf, prg, bufs, vals=() if is_mock else vals, prefix=prg.cbuf_0 or None) + # mock expects all vars to be 64 bit + if is_mock and vals: self.bind_sints_to_buf(*vals, buf=self.buf, fmt='q', offset=len(prg.cbuf_0)*4 + len(bufs)*8) class NVProgram(HCQProgram['NVDevice']): def __init__(self, dev:NVDevice, obj:TinyELF): @@ -314,7 +316,7 @@ class NVProgram(HCQProgram['NVDevice']): self.max_threads = ((65536 // round_up(max(1, self.regs_usage) * 32, 256)) // 4) * 4 * 32 # NV's kernargs is constbuffer, then arguments to the kernel follows. Kernargs also appends QMD at the end of the kernel. - super().__init__(NVArgsState, self.dev, self.name, kernargs_alloc_size=round_up(self.constbufs[0][1], 1 << 8) + (8 << 8)) + super().__init__(NVArgsState, self.dev, obj, kernargs_alloc_size=round_up(self.constbufs[0][1], 1 << 8) + (8 << 8)) weakref.finalize(self, self._fini, self.dev, self.lib_gpu, buf_spec) def _parse_elf_info(self, sh, start_off=0): diff --git a/tinygrad/runtime/ops_qcom.py b/tinygrad/runtime/ops_qcom.py index ecdeb7b768..224f925762 100644 --- a/tinygrad/runtime/ops_qcom.py +++ b/tinygrad/runtime/ops_qcom.py @@ -210,10 +210,12 @@ class QCOMArgsState(HCQArgsState): if prg.samp_cnt > 0: to_mv(int(self.buf.va_addr) + prg.samp_off, len(prg.samplers) * 4).cast('I')[:] = array.array('I', prg.samplers) if prg.NIR: self.bind_sints_to_buf(*[b.va_addr for b in ubos], buf=self.buf, fmt='Q', offset=prg.buf_off) - self.bind_sints_to_buf(*vals, buf=self.buf, fmt='I', offset=prg.buf_off + len(ubos) * 8) + for v,(o,dt) in zip(vals, TinyELF.iter_sig(prg.signature[len(bufs):], len(ubos)*8)): + self.bind_sints_to_buf(v, buf=self.buf, fmt=dt.fmt, offset=prg.buf_off + o) else: for i, b in enumerate(ubos): self.bind_sints_to_buf(b.va_addr, buf=self.buf, fmt='Q', offset=prg.buf_offs[i]) - for i, v in enumerate(vals): self.bind_sints_to_buf(v, buf=self.buf, fmt='I', offset=prg.buf_offs[i+len(ubos)]) + for i,(v,(_,_,dt,_)) in enumerate(zip(vals, prg.signature[len(bufs):])): + self.bind_sints_to_buf(v, buf=self.buf, fmt=dt.fmt, offset=prg.buf_offs[i+len(ubos)]) def _tex(b, ibo=False): imgdt, shape, buf = b @@ -265,7 +267,7 @@ class QCOMProgram(HCQProgram['QCOMDevice']): dev._ensure_stack_size(self.hw_stack_offset * 4) kernargs_alloc_size = round_up(2048 + (self.tex_cnt + self.ibo_cnt) * 0x40 + len(self.samplers) * 4, 0x100) - super().__init__(QCOMArgsState, self.dev, self.name, kernargs_alloc_size=kernargs_alloc_size) + super().__init__(QCOMArgsState, self.dev, obj, kernargs_alloc_size=kernargs_alloc_size) weakref.finalize(self, self._fini, self.dev, self.lib_gpu, buf_spec) def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), diff --git a/tinygrad/runtime/support/hcq.py b/tinygrad/runtime/support/hcq.py index 8f8df4cf84..e22a436b0a 100644 --- a/tinygrad/runtime/support/hcq.py +++ b/tinygrad/runtime/support/hcq.py @@ -6,7 +6,7 @@ try: import fcntl # windows misses that except ImportError: fcntl = None #type:ignore[assignment] from tinygrad.helpers import DEV, PROFILE, getenv, to_mv, from_mv, cpu_profile, ProfileRangeEvent, select_first_inited, select_by_name, unwrap from tinygrad.helpers import suppress_finalizing, pluralize, TracingKey -from tinygrad.device import Device, BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent, Program +from tinygrad.device import Device, BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent, Program, TinyELF from tinygrad.uop.ops import sym_infer, sint, UOp from tinygrad.runtime.autogen import libc from tinygrad.runtime.support.memory import BumpAllocator @@ -326,14 +326,15 @@ class CLikeArgsState(HCQArgsState[ProgramType]): if prefix is not None: self.buf.cpu_view().view(size=len(prefix) * 4, fmt='I')[:] = array.array('I', prefix) self.bind_sints_to_buf(*[b.va_addr for b in bufs], buf=self.buf, fmt='Q', offset=len(prefix or []) * 4) - assert None not in vals - self.bind_sints_to_buf(*cast(tuple[sint, ...], vals), buf=self.buf, fmt='I', offset=len(prefix or []) * 4 + len(bufs) * 8) + for v,(val_offset,dt) in zip(vals, TinyELF.iter_sig(prg.signature[-len(vals):], len(bufs) * 8)): + assert v is not None + self.bind_sints_to_buf(v, buf=self.buf, fmt=dt.fmt, offset=len(prefix or []) * 4 + val_offset) class HCQProgram(Program[HCQDeviceType]): - def __init__(self, args_state_t:Type[HCQArgsState], dev:HCQDeviceType, name:str, kernargs_alloc_size:int, lib:bytes|None=None, base:int|None=None): - self.args_state_t, self.dev, self.name, self.kernargs_alloc_size = args_state_t, dev, name, kernargs_alloc_size + def __init__(self, args_state_t:Type[HCQArgsState], dev:HCQDeviceType, obj:TinyELF, kernargs_alloc_size:int, base:int|None=None): + self.args_state_t, self.dev, self.name, self.signature, self.kernargs_alloc_size = args_state_t, dev, obj.name, obj.signature, kernargs_alloc_size self.prof_prg_counter = next(self.dev.prof_prg_counter) - if PROFILE: Compiled.profile_events += [ProfileProgramEvent(dev.device, name, lib, base, self.prof_prg_counter)] + if PROFILE: Compiled.profile_events += [ProfileProgramEvent(dev.device, obj.name, obj.lib, base, self.prof_prg_counter)] @staticmethod def _fini(dev, buf, spec): dev.allocator.free(buf, buf.size, spec) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 6b48d94b07..4c454ffef2 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1734,7 +1734,7 @@ pm_lower_weak = PatternMatcher([ # a weakfloat Unary (sin/exp2/...) must resolve here, before the transcendental decomposition (UPat(GroupOp.Binary|GroupOp.Unary|{Ops.WHERE, Ops.RANGE, Ops.STACK, Ops.SPECIAL}, name="u"), lower_weak_node), (UPat(Ops.PARAM, dtype=dtypes.weakint, name="u"), - lambda u: u.replace(dtype=None, arg=replace(u.arg, dtype=dtypes.int)).cast(dtypes.weakint) if u.addrspace == AddrSpace.ALU else None), + lambda u: u.replace(dtype=None, arg=replace(u.arg, dtype=select_dtype(u))).cast(dtypes.weakint) if u.addrspace == AddrSpace.ALU else None), ]) def lower_weak_srcs(ctx:dict[UOp, UOp]|None, u:UOp) -> UOp|None: if ctx is None: ctx = {} diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index c1f6897db9..7d32444341 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -142,7 +142,8 @@ spec_tensor = PatternMatcher([ if isinstance(buf.arg, ParamArg) and buf.addrspace is AddrSpace.GLOBAL else None), # Tensor variable bindings - (UPat(Ops.BIND, (dtypes.int, dtypes.weakint,), (UPat(Ops.PARAM), UPat.cvar(dtype=(dtypes.int,dtypes.weakint,))), arg=None), lambda: True), + (UPat(Ops.BIND, (dtypes.int, dtypes.long, dtypes.weakint,), (UPat(Ops.PARAM), UPat.cvar(dtype=(dtypes.int,dtypes.long,dtypes.weakint,))), arg=None), + lambda: True), # custom function (UPat(Ops.CUSTOM_FUNCTION, name="x"), lambda x: isinstance(x.arg, str)),