nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

64-bit UOp.variable support (#17246)

Browse Source
This commit is contained in:
sirhcm
2026-07-27 23:56:09 -04:00
committed by GitHub
parent f837ca3587
commit ab8fb191b2
20 changed files with 103 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/backend/test_tensor_variable.py
+12 -4
View File
@@ -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)
test/mockgpu/nv/nvgpu.py
+1 -1
View File
@@ -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
tinygrad/device.py
+7 -1
View File
@@ -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, ...]=(),
tinygrad/renderer/amd/elf.py
+3 -4
View File
@@ -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
tinygrad/renderer/cstyle.py
+7 -5
View File
@@ -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))}"}
tinygrad/renderer/nir.py
+11 -7
View File
@@ -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]
tinygrad/runtime/graph/cuda.py
+1 -1
View File
@@ -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)))
tinygrad/runtime/graph/metal.py
+11 -8
View File
@@ -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()
tinygrad/runtime/ops_amd.py
+1 -2
View File
@@ -602,8 +602,7 @@ class AMDProgram(HCQProgram['AMDDevice']):
if dev.sqtt_enabled: self.libhash: tuple[int, int] = struct.unpack('<Q', hashlib.md5(self.lib).digest()[:8])*2
super().__init__(CLikeArgsState, self.dev, self.name, kernargs_alloc_size=self.kernargs_segment_size+additional_alloc_sz, lib=self.lib,
base=self.lib_gpu.va_addr)
super().__init__(CLikeArgsState, self.dev, obj, kernargs_alloc_size=self.kernargs_segment_size+additional_alloc_sz, base=self.lib_gpu.va_addr)
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), vals:tuple[int|None, ...]=(),
tinygrad/runtime/ops_cl.py
+1 -1
View File
@@ -55,7 +55,7 @@ class CLProgram(Program['CLDevice']):
def __call__(self, *bufs:cl.cl_mem, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]|None=None, vals:tuple[int, ...]=(),
wait=False, **kw) -> 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])
tinygrad/runtime/ops_cpu.py
+2 -2
View File
@@ -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):
tinygrad/runtime/ops_cuda.py
+6 -5
View File
@@ -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]
tinygrad/runtime/ops_dsp.py
+13 -9
View File
@@ -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]
tinygrad/runtime/ops_hip.py
+4 -3
View File
@@ -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)
tinygrad/runtime/ops_metal.py
+3 -2
View File
@@ -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?
tinygrad/runtime/ops_nv.py
+5 -3
View File
@@ -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):
tinygrad/runtime/ops_qcom.py
+5 -3
View File
@@ -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),
tinygrad/runtime/support/hcq.py
+7 -6
View File
@@ -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)
tinygrad/uop/ops.py
+1 -1
View File
@@ -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 = {}
tinygrad/uop/spec.py
+2 -1
View File
@@ -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)),