diff --git a/extra/qcom_gpu_driver/qcom_opencl_interop.py b/extra/qcom_gpu_driver/qcom_opencl_interop.py index 1be64fafeb..7de9336849 100644 --- a/extra/qcom_gpu_driver/qcom_opencl_interop.py +++ b/extra/qcom_gpu_driver/qcom_opencl_interop.py @@ -78,9 +78,7 @@ hexdump(to_mv(cl_buf_desc_ptr, 0x100)) rawbuf_ptr = to_mv(cl_buf_desc_ptr, 0x100).cast('Q')[20] # offset 0xA0 is a raw gpu pointer. # create QCOM tensor with the externally managed buffer -# dtypes.imageh = cl.cl_image_format(cl.CL_RGBA, cl.CL_HALF_FLOAT) -# dtypes.imagef = cl.cl_image_format(cl.CL_RGBA, cl.CL_FLOAT) -x = Tensor.from_blob(rawbuf_ptr, (h*w*4,), dtype=dtypes.imagef((h,w)), device='QCOM') +x = Tensor.from_blob(rawbuf_ptr, (h,w,4), dtype=dtypes.float, device='QCOM') y = (x + 1).tolist() print(y[:10]) diff --git a/test/external/external_benchmark_op_conv.py b/test/external/external_benchmark_op_conv.py index 8f84b8b8c2..2ab421cb81 100644 --- a/test/external/external_benchmark_op_conv.py +++ b/test/external/external_benchmark_op_conv.py @@ -12,7 +12,7 @@ from tinygrad.dtype import Invalid # PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx def vision_conv_143(): - c0 = UOp.param(0, dtypes.imageh((16, 1024, 4))) + c0 = UOp.param(0, dtypes.half, shape=(16, 1024, 4)) c2 = UOp.range(32, 3, AxisType.LOOP) c5 = UOp.range(128, 4, AxisType.LOOP) c8 = UOp.range(16, 2, AxisType.LOOP) @@ -22,11 +22,11 @@ def vision_conv_143(): c26 = UOp.range(7, 1, AxisType.REDUCE) c27 = c2*2+c26 c32 = ((c27<3)!=True)&(c27<67) - c34 = UOp.param(1, dtypes.imageh((32, 1024, 4))) + c34 = UOp.param(1, dtypes.half, shape=(32, 1024, 4)) c38 = c5//2 c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.weakint, Invalid)) c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0)) - c49 = UOp.param(2, dtypes.imageh((64, 49, 4))) + c49 = UOp.param(2, dtypes.half, shape=(64, 49, 4)) c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196)) c63 = UOp.param(3, dtypes.float, (128,)) c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5) @@ -38,7 +38,7 @@ def vision_conv_143(): return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts)) def vision_conv_153(): - c0 = UOp.param(0, dtypes.imageh((8, 1024, 4))) + c0 = UOp.param(0, dtypes.half, shape=(8, 1024, 4)) c2 = UOp.range(16, 3, AxisType.LOOP) c5 = UOp.range(256, 4, AxisType.LOOP) c8 = UOp.range(8, 2, AxisType.LOOP) @@ -48,11 +48,11 @@ def vision_conv_153(): c26 = UOp.range(7, 1, AxisType.REDUCE) c27 = c2*2+c26 c32 = ((c27<3)!=True)&(c27<35) - c34 = UOp.param(1, dtypes.imageh((16, 1024, 4))) + c34 = UOp.param(1, dtypes.half, shape=(16, 1024, 4)) c38 = c5//2 c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.weakint, Invalid)) c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0)) - c49 = UOp.param(2, dtypes.imageh((128, 49, 4))) + c49 = UOp.param(2, dtypes.half, shape=(128, 49, 4)) c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196)) c63 = UOp.param(3, dtypes.float, (256,)) c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5) @@ -64,14 +64,14 @@ def vision_conv_153(): return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts)) def dm_conv_172(): - c0 = UOp.param(0, dtypes.imageh((1, 240, 4))) + c0 = UOp.param(0, dtypes.half, shape=(1, 240, 4)) c2 = UOp.range(960, 4, AxisType.LOOP) - c5 = UOp.param(1, dtypes.imageh((8, 384, 4))) + c5 = UOp.param(1, dtypes.half, shape=(8, 384, 4)) c7 = UOp.range(32, 0, AxisType.REDUCE) c10 = UOp.range(4, 1, AxisType.REDUCE) c13 = UOp.range(12, 3, AxisType.REDUCE) c18 = UOp.range(8, 2, AxisType.REDUCE) - c23 = UOp.param(2, dtypes.imageh((240, 128, 4))) + c23 = UOp.param(2, dtypes.half, shape=(240, 128, 4)) c35 = c5.index((c7*4+c10+c13*128+c18*1536))*c23.index((c10*4+c2%4+c7*16+c2//4*512)) c37 = UOp.param(3, dtypes.float, (960,)) c39 = c35.reduce(c7, c10, arg=Ops.ADD)+c37.index(c2) diff --git a/test/null/test_dtype.py b/test/null/test_dtype.py index 63850b01c3..35c157b5f9 100644 --- a/test/null/test_dtype.py +++ b/test/null/test_dtype.py @@ -1,25 +1,10 @@ import unittest, pickle from tinygrad.tensor import Tensor -from tinygrad.dtype import dtypes, DType, ImageDType, to_dtype, Invalid, InvalidType - -class TestImageDType(unittest.TestCase): - def test_image_scalar(self): - assert dtypes.imagef((10,10)).base.scalar() == dtypes.float32 - assert dtypes.imageh((10,10)).base.scalar() == dtypes.float32 - def test_image_vec(self): - assert dtypes.imagef((10,10)).base.vec(4) == dtypes.float32.vec(4) - assert dtypes.imageh((10,10)).base.vec(4) == dtypes.float32.vec(4) +from tinygrad.dtype import dtypes, DType, to_dtype, Invalid, InvalidType class TestEqStrDType(unittest.TestCase): - def test_image_ne(self): - if ImageDType is None: raise unittest.SkipTest("no ImageDType support") - assert dtypes.float == dtypes.float32, "float doesn't match?" - assert dtypes.imagef((1,2,4)) != dtypes.imageh((1,2,4)), "different image dtype doesn't match" - assert dtypes.imageh((1,2,4)) != dtypes.imageh((1,4,2)), "different shape doesn't match" - assert dtypes.imageh((1,2,4)) == dtypes.imageh((1,2,4)), "same shape matches" - assert isinstance(dtypes.imageh((1,2,4)), ImageDType) def test_strs(self): - self.assertEqual(str(dtypes.imagef((1,2,4))), "dtypes.imagef((1, 2, 4))") + self.assertEqual(str(dtypes.float32), "dtypes.float") class TestToDtype(unittest.TestCase): def test_dtype_to_dtype(self): diff --git a/test/null/test_simplify_valid_idx.py b/test/null/test_simplify_valid_idx.py index 7ca642e7df..e3fab1495b 100644 --- a/test/null/test_simplify_valid_idx.py +++ b/test/null/test_simplify_valid_idx.py @@ -20,7 +20,7 @@ def get_gated_load_uop(valid:UOp, idx:UOp): def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]): return UOp(Ops.LOAD, dtypes.float, ( - UOp.param(0, dtypes.imagef(image_shape)).index(idx[1].valid(valid), idx[0].valid(valid)), + UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)), )) def Special(expr, nmax): return UOp(Ops.SPECIAL, dtypes.weakint, (UOp.const(dtypes.weakint, nmax),), expr) diff --git a/tinygrad/codegen/late/coalese.py b/tinygrad/codegen/late/coalese.py index 245fa22f85..244c9c8bd3 100644 --- a/tinygrad/codegen/late/coalese.py +++ b/tinygrad/codegen/late/coalese.py @@ -1,10 +1,10 @@ from typing import Any import itertools, functools from collections import defaultdict -from tinygrad.dtype import dtypes, AddrSpace, Invalid, ImageDType, DType +from tinygrad.dtype import dtypes, AddrSpace, Invalid, DType from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, shape_to_shape_arg from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate -from tinygrad.helpers import getenv, IMAGE, OSX, ceildiv +from tinygrad.helpers import getenv, IMAGE, OSX, ceildiv, is_image_shape from tinygrad.renderer import Renderer # ***** image load valid simplification ***** @@ -40,15 +40,16 @@ def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None: return None if idx is start_idx or idx is start_idx.simplify() else buf.index(idx.valid(valid)) def simplify_valid_image_load(buf:UOp, idx_y:UOp, idx_x:UOp, valid:UOp) -> UOp|None: - if not isinstance(buf.dtype, ImageDType): return None + if not is_image_shape(buf._shape): return None start_idx = idx_x._stack(idx_y) idx = uop_given_valid(valid, start_idx) - drop_stmt = _drop_valid_stmts(valid, idx, buf.dtype.shape[0], buf.dtype.shape[1]) + drop_stmt = _drop_valid_stmts(valid, idx, buf._shape[0], buf._shape[1]) if not drop_stmt and idx is start_idx: return None new_valid = UOp.uprod(*ss) if (ss:=[s for s in valid.split_uop(Ops.AND) if s not in drop_stmt]) else None idx_y, idx_x = idx.index(1), idx.index(0) - return buf.index(idx_y.valid(new_valid), idx_x.valid(new_valid)) if new_valid is not None else buf.index(idx_y, idx_x) + if new_valid is not None: return buf.index(idx_y.valid(new_valid), idx_x.valid(new_valid), dtype=dtypes.float) + return buf.index(idx_y, idx_x, dtype=dtypes.float) indexing_simplify = PatternMatcher([ # image load valid idx simplification @@ -82,12 +83,12 @@ def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None: if len(cands) == 0: return None # and tiebreak with indexing complexity (ie. number of nodes) h, w, cidx = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].index(1).simplify().backward_slice)) - buf = buf.replace(src=(shape_to_shape_arg((h, w, 4)),), dtype=(dtypes.imageh if buf.dtype.itemsize == 2 else dtypes.imagef)((h, w, 4))) + buf = buf.replace(src=(shape_to_shape_arg((h, w, 4)),)) shapes[buf.arg.slot] = (h, w) if valid.op is not Ops.CONST or valid.arg is not True: - return buf.index(cidx.src[1].valid(valid), cidx.src[0].valid(valid)) + return buf.index(cidx.src[1].valid(valid), cidx.src[0].valid(valid), dtype=dtypes.float) else: - return buf.index(cidx.src[1], cidx.src[0]) + return buf.index(cidx.src[1], cidx.src[0], dtype=dtypes.float) pm_simplify_add_image = PatternMatcher([ (UPat(Ops.SHRINK, src=(UPat(Ops.PARAM, name="buf"), UPat(name="x"), UPat(arg=4))), transform_to_image), @@ -127,11 +128,11 @@ def memory_coalesing(sink:UOp, ctx:Renderer) -> UOp: if ctx is not None and ctx.target.device == "DSP": lengths = [128,64,32,16,8,4] must_divide = False - elif buf.dtype not in (dtypes.float, dtypes.half, *dtypes.fp8s) and not isinstance(buf.dtype, ImageDType): + elif buf.dtype not in (dtypes.float, dtypes.half, *dtypes.fp8s) and not is_image_shape(buf._shape): pass elif buf.addrspace == AddrSpace.REG: pass - elif isinstance(buf.dtype, ImageDType): + elif is_image_shape(buf._shape): lengths = [4] elif ctx is not None and ctx.supports_float4: # TODO: a better way to get this than ctx diff --git a/tinygrad/codegen/late/gater.py b/tinygrad/codegen/late/gater.py index cd88f5b608..97b105720d 100644 --- a/tinygrad/codegen/late/gater.py +++ b/tinygrad/codegen/late/gater.py @@ -6,10 +6,10 @@ pm_move_gates_from_index = PatternMatcher([ # for image idx (must be first) (UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx_y"), UPat(arg=Invalid)), UPat.var("gate").where(UPat.var("idx_x"), UPat(arg=Invalid))).load(name="l"), - lambda buf,gate,idx_y,idx_x,l: buf.index(idx_y, idx_x).load(l.vconst_like(0), gate)), + lambda buf,gate,idx_y,idx_x,l: buf.index(idx_y, idx_x, dtype=dtypes.float).load(l.vconst_like(0), gate)), (UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx_y"), UPat(arg=Invalid)), UPat.var("gate").where(UPat.var("idx_x"), UPat(arg=Invalid))).store(UPat.var("data")), - lambda buf,gate,idx_y,idx_x,data: buf.index(idx_y, idx_x).store(data, gate)), + lambda buf,gate,idx_y,idx_x,data: buf.index(idx_y, idx_x, dtype=dtypes.float).store(data, gate)), # here we create the alt value for load to be 0s and remove the where Invalid (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat(), UPat.var("gate").where(UPat.var("idx"), UPat(arg=Invalid)),), name="mop", allow_any_len=True) \ diff --git a/tinygrad/dtype.py b/tinygrad/dtype.py index 40f4e0d27d..2653ae6e91 100644 --- a/tinygrad/dtype.py +++ b/tinygrad/dtype.py @@ -2,7 +2,7 @@ from __future__ import annotations from typing import Final, ClassVar, Callable, Literal import math, struct, ctypes, functools from dataclasses import dataclass, fields -from tinygrad.helpers import getenv, prod, round_up, OSX +from tinygrad.helpers import getenv from enum import IntEnum, auto class ConstFloat(float): @@ -97,36 +97,11 @@ class DType(metaclass=DTypeMetaClass): # int is the default. wrap floats in ConstFloat to distinguish -0.0 from 0.0 in cache return ConstFloat(float(val)) if dtypes.is_float(self) else bool(val) if dtypes.is_bool(self) else int(val) -@dataclass(frozen=True, eq=False) -class ImageDType(DType): - _base: DType - addrspace: AddrSpace - v: int - size: int = -1 # -1 is unlimited size - shape: tuple[int, ...] = () # shape of the Image - @property - def base(self): return self._base - @functools.cache # pylint: disable=method-cache-max-size-none - def vec(self, sz:int) -> DType: - assert self.v == 1, f"can't vectorize image {self} with size {sz}" - if sz == 1: return self # sz=1 is a scalar - return ImageDType(self.priority, self.bitsize, self.name, self.fmt, self.count, self, self._base, self.addrspace, sz, self.size, self.shape) - def nbytes(self) -> int: - if self.size == -1: raise RuntimeError("can't get nbytes of a pointer with unlimited size") - return self.size*self.itemsize - @property - def vcount(self): return self.v - def __repr__(self): return f"dtypes.{self.name}({self.shape})" + (f'.vec({self.v})' if self.v != 1 else '') - - # for 1d images on macos, we need to round pitch up to 256 pixels to make CL happy - @property - def pitch(self): return (round_up(self.shape[1], 256) if OSX else self.shape[1]) * 4 * self.itemsize - class dtypes: @staticmethod @functools.cache - def is_float(x: DType) -> bool: return x.scalar() in dtypes.floats or isinstance(x, ImageDType) + def is_float(x: DType) -> bool: return x.scalar() in dtypes.floats @staticmethod # static methods on top, or bool in the type info will refer to dtypes.bool @functools.cache def is_int(x: DType) -> bool: return x.scalar() in (dtypes.ints + (dtypes.weakint,)) @@ -178,12 +153,6 @@ class dtypes: uchar = uint8; ushort = uint16; uint = uint32; ulong = uint64 # noqa: E702 char = int8; short = int16; int = int32; long = int64 # noqa: E702 - # NOTE: these are image dtypes - @staticmethod - def imageh(shp): return ImageDType(100, 16, "imageh", 'e', 1, None, dtypes.float32, AddrSpace.GLOBAL, 1, prod(shp), shp) - @staticmethod - def imagef(shp): return ImageDType(100, 32, "imagef", 'f', 1, None, dtypes.float32, AddrSpace.GLOBAL, 1, prod(shp), shp) - default_float: ClassVar[DType] = float32 default_int: ClassVar[DType] = int32 @@ -222,8 +191,7 @@ def _get_recursive_parents(dtype:DType) -> set[DType]: return set.union(*[_get_recursive_parents(d) for d in promo_lattice[dtype]], {dtype}) if dtype != dtypes.float64 else {dtypes.float64} @functools.cache def least_upper_dtype(*ds:DType) -> DType: - return min(set.intersection(*[_get_recursive_parents(d.scalar()) for d in ds])) \ - if not (images:=[d for d in ds if isinstance(d, ImageDType)]) else images[0] + return min(set.intersection(*[_get_recursive_parents(d.scalar()) for d in ds])) def least_upper_float(dt:DType) -> DType: return dt if dtypes.is_float(dt) else least_upper_dtype(dt, dtypes.default_float) DTYPES_DICT = {k: v for k, v in dtypes.__dict__.items() if isinstance(v, DType) and not k.startswith(("default", "void", "weakint", "_"))} diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 2fbca2e70d..ec5e9c95a5 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -34,6 +34,7 @@ def get_shape(x) -> tuple[int, ...]: if not hasattr(x, "__len__") or isinstance(x, str) or getattr(x, "shape", None) == (): return () if not all_same(subs:=[get_shape(xi) for xi in x]): raise ValueError(f"inhomogeneous shape from {x}") return (len(subs),) + (subs[0] if subs else ()) +def is_image_shape(shape): return shape is not None and len(shape) == 3 and shape[-1] == 4 def all_int(t: Sequence[Any]) -> TypeGuard[tuple[int, ...]]: return all(isinstance(s, int) for s in t) def colored(st, color:str|None, background=False): # replace the termcolor library if NO_COLOR: return st diff --git a/tinygrad/mixin/op.py b/tinygrad/mixin/op.py index dcd00fbad3..a6ed3ff06b 100644 --- a/tinygrad/mixin/op.py +++ b/tinygrad/mixin/op.py @@ -6,7 +6,7 @@ from tinygrad.mixin.movement import MovementMixin from tinygrad.mixin.reduce import ReduceMixin from tinygrad.uop import Ops from tinygrad.uop.ops import _broadcast_shape, resolve, smax, smin, identity_element -from tinygrad.dtype import ConstType, DType, DTypeLike, Invalid, ImageDType, PyConst, dtypes, least_upper_dtype, sum_acc_dtype, to_dtype +from tinygrad.dtype import ConstType, DType, DTypeLike, Invalid, PyConst, dtypes, least_upper_dtype, sum_acc_dtype, to_dtype from tinygrad.helpers import all_int, argfix, argsort, ceildiv, flatten, flat_to_grouped, fully_flatten, get_shape, make_tuple, merge_dicts, prod from tinygrad.helpers import resolve_pool_pads, round_up, IMAGE, FLOAT16, WINO @@ -365,8 +365,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin): out_shape = _broadcast_shape(x.shape, y.shape) x, y = x._broadcast_to(out_shape), y._broadcast_to(out_shape) except (RuntimeError, ValueError): pass - # ptr dtypes aren't in the promo lattice - if x.dtype == y.dtype or any(isinstance(d, ImageDType) for d in (x.dtype, y.dtype)): return x, y + if x.dtype == y.dtype: return x, y return x.cast(out_dtype := least_upper_dtype(x.dtype, y.dtype)), y.cast(out_dtype) def dot(self, w:Self, dtype:DTypeLike|None=None) -> Self: diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 0071e57cac..f429340346 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -3,8 +3,8 @@ import math, sys, struct from collections import defaultdict, Counter from tinygrad.codegen.opt import tc from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str, axis_letters -from tinygrad.helpers import strip_parens, getenv, prod, dedup, Target, CPU_COUNT, IMAGE, FLOAT16 -from tinygrad.dtype import ImageDType, dtypes, DType, AddrSpace, truncate, float_to_bf16 +from tinygrad.helpers import strip_parens, getenv, prod, dedup, Target, CPU_COUNT, IMAGE, FLOAT16, is_image_shape +from tinygrad.dtype import dtypes, DType, AddrSpace, truncate, float_to_bf16 from tinygrad.renderer import Renderer @@ -140,9 +140,9 @@ class CStyleLanguage(Renderer): def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str,tuple[UOp,bool]]], uops:list[UOp], prefix=None) -> str: tmp = "" - if any(isinstance(u.dtype, ImageDType) for _,(u,_) in bufs): + 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, self._render_dtype(u.dtype, sz=1, addrspace=u.addrspace, mutable=mutable)+self.buffer_suffix \ + buftypes = [(name, 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] 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]) @@ -164,8 +164,8 @@ class CStyleLanguage(Renderer): suffix = f"[{x.max_numel()}]" return f"{prefix}{self._render_dtype(x.dtype, sz=lanes)} {self[x]}{suffix};" - def _render_dtype(self, dtype:DType, sz:int=1, addrspace=AddrSpace.ALU, mutable=True, override_ptr=False): - if isinstance(dtype, ImageDType): return f"{'write_only' if mutable else 'read_only'} image2d_t" + def _render_dtype(self, dtype:DType, sz:int=1, addrspace=AddrSpace.ALU, mutable=True, override_ptr=False, shape=None): + if is_image_shape(shape): return f"{'write_only' if mutable else 'read_only'} image2d_t" prefix, suffix = "", "" if addrspace in (AddrSpace.LOCAL, AddrSpace.GLOBAL): if addrspace == AddrSpace.LOCAL and self.smem_prefix_for_cast: prefix = self.smem_prefix @@ -176,16 +176,16 @@ class CStyleLanguage(Renderer): return prefix + self.type_map.get(scalar:=dtype.scalar(), scalar.name).replace(" ", "_") + str(sz) + suffix return prefix + self.type_map.get(scalar:=dtype.scalar(), scalar.name) + suffix - def render_type(self, u:UOp): return self._render_dtype(u.dtype, u.max_numel(), u.addrspace) + def render_type(self, u:UOp): return self._render_dtype(u.dtype, u.max_numel(), u.addrspace, shape=u._shape) def render_access(self, u:UOp): if u.max_numel() > 1 or u.dtype != u.src[0].dtype: - return f"*(({self._render_dtype(u.dtype, u.max_numel(), u.addrspace, override_ptr=True)})({self[u]}))" + return f"*(({self._render_dtype(u.dtype, u.max_numel(), u.addrspace, override_ptr=True, shape=u._shape)})({self[u]}))" else: return f"*{self[u]}" def render_cast(self, u:UOp, val:str) -> str: return f"({self.render_type(u)})({val})" # LEGACY def render_dtype(self, dt:DType, mutable=True) -> str: - return self._render_dtype(dt, dt.count, dt.addrspace if isinstance(dt, ImageDType) else AddrSpace.REG) + return self._render_dtype(dt, dt.count, AddrSpace.REG) def __getitem__(self, key): return self.r[key] # hacky helper def _render(self, uops:list[UOp]) -> tuple[str, list[str], list[tuple[str,tuple[UOp,bool]]]]: @@ -327,10 +327,10 @@ class OpenCLRenderer(CStyleLanguage): return super().render_kernel(function_name, kernel, bufs, uops, prefix) def aux(self, uops:list[UOp]): - arg_dtypes:list[list[tuple[int, DType]]] = [] + arg_dtypes:list[list[tuple[int, DType, tuple|None]]] = [] for i,u in enumerate(u for u in uops if u.op is Ops.PARAM): while len(arg_dtypes) <= u.arg.slot: arg_dtypes.append([]) - arg_dtypes[u.arg.slot].append((i, u.dtype)) + arg_dtypes[u.arg.slot].append((i, u.dtype, u._shape)) return tuple(tuple(a) for a in arg_dtypes), def supported_dtypes(self): return {d for d in super().supported_dtypes() diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index 9f87d8ffec..e7e6012aad 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, ImageDType, dtypes, truncate -from tinygrad.helpers import DEBUG, OSX, unwrap, fromimport, Target +from tinygrad.dtype import AddrSpace, DType, dtypes, truncate +from tinygrad.helpers import DEBUG, OSX, unwrap, fromimport, Target, is_image_shape from tinygrad.renderer import Renderer from tinygrad.renderer.cstyle import CUDARenderer, OpenCLRenderer from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str @@ -137,7 +137,7 @@ class NIRRenderer(Renderer): (UPat(Ops.CAST, (dtypes.uchar, dtypes.ushort), src=(UPat.var("x", dtypes.floats),), name="c"), lambda x,c: x.cast(dtypes.int32).cast(c.dtype)), # load/store use pointer arithmetic, and the cast does nothing. NOTE: this doesn't apply to image indexing cause it's 1-D (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"), UPat.var("off")), allow_any_len=True, name="x"), lambda x,buf,off: x.replace( - src=(buf,off.cast(dtypes.long))+x.src[2:]) if buf.addrspace != AddrSpace.REG and not isinstance(buf.dtype, ImageDType) else None), + src=(buf,off.cast(dtypes.long))+x.src[2:]) if buf.addrspace != AddrSpace.REG and not is_image_shape(buf._shape) else None), # images need index to be int for nir (UPat.var("buf").index(UPat.var("idx_y"), UPat.var("idx_x")), lambda buf,idx_y,idx_x: buf.index(idx_y.cast(dtypes.int), idx_x.cast(dtypes.int))), @@ -290,7 +290,7 @@ class IR3Renderer(NIRRenderer, OpenCLRenderer): self.img_idx += 1 return nimm(self.b, self.img_idx - 1, dtypes.int) - def param(self, b, x, sz): return self._param_img(x) if isinstance(x.dtype, ImageDType) else self._param(b, x, sz) + def param(self, b, x, sz): return self._param_img(x) if is_image_shape(x._shape) else self._param(b, x, sz) def prerender(self, uops:list[UOp]): super().prerender(uops) @@ -301,9 +301,10 @@ class IR3Renderer(NIRRenderer, OpenCLRenderer): 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] texs, imgs = itertools.count().__next__, itertools.count().__next__ - for b in filter(lambda b: isinstance(b.dtype, ImageDType), bufs): nimm_set(self.r[b], texs() if b in self.texs else imgs(), dtypes.int) + for b in filter(lambda b: is_image_shape(b._shape), bufs): + nimm_set(self.r[b], texs() if b in self.texs else imgs(), dtypes.int) - self.b.shader.contents.info.num_ubos = len([u for u in bufs if not isinstance(u.dtype, ImageDType)]) + self.b.shader.contents.info.num_ubos = len([u for u in bufs if not is_image_shape(u._shape)]) self.b.shader.contents.info.num_images = texs() + imgs() def supported_dtypes(self): return {d for d in NIRRenderer.supported_dtypes(self) if d != dtypes.double} diff --git a/tinygrad/renderer/wgsl.py b/tinygrad/renderer/wgsl.py index ac687ea031..88a160dffe 100644 --- a/tinygrad/renderer/wgsl.py +++ b/tinygrad/renderer/wgsl.py @@ -93,7 +93,7 @@ class WGSLRenderer(CStyleLanguage): ]) + base_rewrite def render_cast(self, u:UOp, val: str) -> str: return f"{self.type_map[u.dtype]}({val})" - def _render_dtype(self, dtype:DType, sz:int=1, addrspace=AddrSpace.REG, mutable=True, override_ptr=False): return "var" + def _render_dtype(self, dtype:DType, sz:int=1, addrspace=AddrSpace.REG, mutable=True, override_ptr=False, shape=None): return "var" def render_load(self, x:str, u:UOp) -> str: return f"atomicLoad(&{x})" if is_packed(u) else x def buf_map(self, u:UOp) -> str: return "atomic" if is_packed(u) else self.type_map[u.dtype.base] def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str,tuple[UOp,bool]]], uops:list[UOp], prefix=None) -> str: diff --git a/tinygrad/runtime/ops_cl.py b/tinygrad/runtime/ops_cl.py index 3e29997498..2848b17c31 100644 --- a/tinygrad/runtime/ops_cl.py +++ b/tinygrad/runtime/ops_cl.py @@ -3,10 +3,9 @@ from typing import cast import ctypes, functools, hashlib from tinygrad.runtime.autogen import opencl as cl from tinygrad.runtime.support import c -from tinygrad.helpers import to_char_p_p, from_mv, OSX, DEBUG, mv_address, suppress_finalizing, unwrap +from tinygrad.helpers import to_char_p_p, from_mv, OSX, DEBUG, mv_address, suppress_finalizing, unwrap, round_up, is_image_shape from tinygrad.renderer.cstyle import OpenCLRenderer from tinygrad.device import BufferSpec, LRUAllocator, Compiled, Compiler, CompileError -from tinygrad.dtype import ImageDType CC_CB = c.CFUNCTYPE[None, [c.POINTER[ctypes.c_char], c.POINTER[None], cl.size_t, c.POINTER[None]]] BP_CB = c.CFUNCTYPE[None, [cl.cl_program, c.POINTER[None]]] @@ -57,10 +56,11 @@ class CLProgram: wait=False, **kw) -> float|None: i = 0 for i,b in enumerate(bufs): - for real_i, dt in self.arg_dtypes[i]: - if isinstance(dt, ImageDType): + for real_i, dt, shape in self.arg_dtypes[i]: + 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]) - desc = cl.cl_image_desc(cl.CL_MEM_OBJECT_IMAGE2D, dt.shape[1], dt.shape[0], image_row_pitch=dt.pitch, buffer=b) + desc = cl.cl_image_desc(cl.CL_MEM_OBJECT_IMAGE2D, shape[1], shape[0], image_row_pitch=pitch, buffer=b) img = checked(cl.clCreateImage(self.dev.context, cl.CL_MEM_READ_WRITE, fmt, desc, None, status:=ctypes.c_int32()), status) check(cl.clSetKernelArg(self.kernel, real_i, ctypes.sizeof(img), ctypes.byref(img))) else: check(cl.clSetKernelArg(self.kernel, real_i, ctypes.sizeof(b), ctypes.byref(b))) diff --git a/tinygrad/runtime/ops_python.py b/tinygrad/runtime/ops_python.py index 1112b00eb1..bc6429af62 100644 --- a/tinygrad/runtime/ops_python.py +++ b/tinygrad/runtime/ops_python.py @@ -5,8 +5,8 @@ from typing import Any, TYPE_CHECKING import pickle, base64, itertools, time, sys, functools from dataclasses import replace -from tinygrad.dtype import DType, dtypes, ImageDType, AddrSpace, truncate, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar -from tinygrad.helpers import all_same, getenv, flatten, Target, IMAGE +from tinygrad.dtype import DType, dtypes, AddrSpace, truncate, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar +from tinygrad.helpers import all_same, getenv, flatten, Target, IMAGE, is_image_shape from tinygrad.device import Compiled, Compiler, Allocator from tinygrad.codegen.opt import tc from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp, bitcast @@ -104,11 +104,10 @@ class PythonProgram: ret:list = [] if u.src[0].addrspace == AddrSpace.ALU: ret = [src_values[0][i][t] for t,i in enumerate(src_values[1])] - elif isinstance(src_dtypes[0], ImageDType): - assert len(src_values) == 3, f"image index must be 3 srcs, not {len(src_values)}" + elif is_image_shape(u.src[0]._shape): for m,oy,ox in zip(*src_values): - if ox < 0 or ox >= src_dtypes[0].shape[1] or oy < 0 or oy >= src_dtypes[0].shape[0]: ret.append((m, None)) - else: ret.append((m, ox*4 + oy*src_dtypes[0].shape[1]*4)) + if ox < 0 or ox >= u.src[0]._shape[1] or oy < 0 or oy >= u.src[0]._shape[0]: ret.append((m, None)) + else: ret.append((m, ox*4 + oy*u.src[0]._shape[1]*4)) else: for m,o in zip(src_values[0], src_values[1]): ret.append((m,o)) values[u] = ret diff --git a/tinygrad/runtime/ops_qcom.py b/tinygrad/runtime/ops_qcom.py index dac94ffe8f..d6e2fc0b96 100644 --- a/tinygrad/runtime/ops_qcom.py +++ b/tinygrad/runtime/ops_qcom.py @@ -8,9 +8,9 @@ from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface from tinygrad.runtime.autogen import kgsl, mesa from tinygrad.renderer.cstyle import QCOMCLRenderer from tinygrad.renderer.nir import IR3Renderer -from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, ceildiv, prod, cpu_profile, lo32, suppress_finalizing +from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, ceildiv, prod, cpu_profile, lo32, suppress_finalizing, is_image_shape from tinygrad.helpers import next_power2, flatten, PROFILE, IMAGE -from tinygrad.dtype import ImageDType, dtypes +from tinygrad.dtype import dtypes from tinygrad.runtime.support.system import System if getenv("IOCTL"): import extra.qcom_gpu_driver.opencl_ioctl # noqa: F401 # pylint: disable=unused-import @@ -200,8 +200,8 @@ class QCOMArgsState(HCQArgsState): super().__init__(buf, prg, bufs, vals=vals) ctypes.memset(int(self.buf.va_addr), 0, prg.kernargs_alloc_size) - ubos = [b for i,b in enumerate(bufs) for _,dt in prg.buf_dtypes[i] if not isinstance(dt, ImageDType)] - uavs = [(dt,b) for i,b in enumerate(bufs) for _,dt in prg.buf_dtypes[i] if isinstance(dt, ImageDType)] + ubos = [b for i,b in enumerate(bufs) for _,dt,shape in prg.buf_dtypes[i] if not is_image_shape(shape)] + uavs = [(dt,shape,b) for i,b in enumerate(bufs) for _,dt,shape in prg.buf_dtypes[i] if is_image_shape(shape)] # NIR can reorder images to different texture slots ibos, texs = uavs[:prg.ibo_cnt], [uavs[prg.ibo_cnt + (prg.tex_to_image[i] if prg.NIR else i)] for i in range(prg.tex_cnt)] for cnst_val,cnst_off,cnst_sz in prg.consts_info: @@ -216,11 +216,12 @@ class QCOMArgsState(HCQArgsState): for i, v in enumerate(vals): self.bind_sints_to_buf(v, buf=self.buf, fmt='I', offset=prg.buf_offs[i+len(ubos)]) def _tex(b, ibo=False): - imgdt, buf = b + imgdt, shape, buf = b + pitch = shape[1] * 4 * imgdt.itemsize fmt = mesa.FMT6_32_32_32_32_FLOAT if imgdt.itemsize == 4 else mesa.FMT6_16_16_16_16_FLOAT return [qreg.a6xx_tex_const_0(fmt=fmt) if ibo else qreg.a6xx_tex_const_0(0x8, swiz_x=0, swiz_y=1, swiz_z=2, swiz_w=3, fmt=fmt), - qreg.a6xx_tex_const_1(width=imgdt.shape[1], height=imgdt.shape[0]), - qreg.a6xx_tex_const_2(type=mesa.A6XX_TEX_2D, pitch=imgdt.pitch, pitchalign=ctz(imgdt.pitch)-6), 0, *data64_le(buf.va_addr), + qreg.a6xx_tex_const_1(width=shape[1], height=shape[0]), + qreg.a6xx_tex_const_2(type=mesa.A6XX_TEX_2D, pitch=pitch, pitchalign=ctz(pitch)-6), 0, *data64_le(buf.va_addr), qreg.a6xx_tex_const_6(plane_pitch=0x400000), qreg.a6xx_tex_const_7(13), 0, 0, 0, 0, 0, 0, 0, 0] self.bind_sints_to_buf(*flatten(map(_tex, texs)), buf=self.buf, fmt='I', offset=prg.tex_off) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 57d36c92ff..4086d8e3b1 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -4,7 +4,7 @@ import sys, time, functools, itertools, math, operator, hashlib, os, types, pick from dataclasses import dataclass from enum import Enum, auto from tinygrad.uop import Ops, GroupOp -from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, DTypeLike, to_dtype, truncate, least_upper_dtype, Invalid, AddrSpace +from tinygrad.dtype import ConstType, dtypes, DType, DTypeLike, to_dtype, truncate, least_upper_dtype, Invalid, AddrSpace from tinygrad.dtype import ConstFloat, PyConst, InvalidType, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar from tinygrad.device import Buffer, MultiBuffer, canonicalize_device from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA @@ -828,7 +828,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return ret assert self.op is Ops.BUFFER, f"must be BUFFER {self.op}" if (cret:=buffers.get(self)) is not None: return cret - rdtype = self.dtype if isinstance(self.dtype, ImageDType) else self.dtype.base + rdtype = self.dtype.base if isinstance(self.device, tuple): ret = MultiBuffer(self.device, self.max_numel(), rdtype).ref(1) else: ret = Buffer(self.device, self.max_numel(), rdtype).ref(1) buffers[self] = ret @@ -1629,11 +1629,11 @@ pm_lower_index_dtype = PatternMatcher([ # remove hanging casts for images (UPat(Ops.PARAM, src=(UPat.var("shape").cast(),), name="p"), lambda p,shape: p.replace(src=(shape,))), (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx_y", dtypes.ints).cast(), UPat.var("idx_x", dtypes.ints).cast()),), - lambda buf,idx_x,idx_y: buf.index(idx_y, idx_x)), + lambda buf,idx_x,idx_y: buf.index(idx_y, idx_x, dtype=dtypes.float)), (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("gate").where(UPat.var("idx_y", dtypes.ints).cast(), UPat(Ops.CONST, arg=Invalid)), UPat.var("gate").where(UPat.var("idx_x", dtypes.ints).cast(), UPat(Ops.CONST, arg=Invalid)))), - lambda buf,idx_x,idx_y,gate: buf.index(idx_y.valid(gate), idx_x.valid(gate))), + lambda buf,idx_x,idx_y,gate: buf.index(idx_y.valid(gate), idx_x.valid(gate), dtype=dtypes.float)), (UPat((Ops.SINK, Ops.NOOP, Ops.END), name="n"), lambda n: n.replace(src=tuple(s.src[0] if s.op is Ops.CAST and s.dtype == dtypes.weakint else s for s in n.src))), ]) diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 84743a6bc0..654e022d2a 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -2,8 +2,8 @@ import math from typing import Any from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, AxisType, KernelInfo, ParamArg from tinygrad.uop.render import print_uops, pyrender -from tinygrad.dtype import DType, ImageDType, dtypes, AddrSpace, Invalid, ConstFloat -from tinygrad.helpers import DEBUG, Context, SPEC, Metadata, panic, CHECK_OOB, all_same +from tinygrad.dtype import DType, dtypes, AddrSpace, Invalid, ConstFloat +from tinygrad.helpers import DEBUG, Context, SPEC, Metadata, panic, CHECK_OOB, all_same, is_image_shape # ***** uop helpers ***** @@ -13,7 +13,7 @@ def validate_index(uidx:UOp, gate:UOp|None=None): if idx.op is Ops.CONST and idx.arg is Invalid: return True if gate is None: gate = UOp.const(dtypes.bool, True) # TODO: check for overflow - if not CHECK_OOB or isinstance(buf.dtype, ImageDType): return True + if not CHECK_OOB or is_image_shape(buf._shape): return True # buffer size sz = buf.max_numel()