remove ImageDType (glm) (#16918)

* remove ImageDType (glm)

* gpt work

* better

* is_image_shape

* cleanups

* remove crap from index

* revert crap python change

* image index with float, this needs to be fixed

* simp

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])
 

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)

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):

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)

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

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) \

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", "_"))}

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

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:

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()

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}

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<u32>" 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:

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)))

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

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)

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))),
 ])

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()