work

tinygrad/codegen/__init__.py

@@ -4,7 +4,6 @@ import itertools
 from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
 from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, panic
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, GroupOp
-from tinygrad.uop.ops import shape_to_shape_arg
 from tinygrad.uop.render import pyrender
 from tinygrad.uop.spec import type_verify, spec_tensor, spec_program
 from tinygrad.renderer import Renderer, Estimates
@@ -17,7 +16,7 @@ from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic, p
 from tinygrad.uop.decompositions import get_late_rewrite_patterns, get_transcendental_patterns, pm_dtype_decomps, get_simplifying_rewrite_patterns
 from tinygrad.codegen.late.expander import expander, pm_pre_expander, pm_group_for_reduce
 from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize_buf_and_index, devectorize_alu, pm_reduce, \
-  ReduceContext, correct_load_store, pm_render, pm_add_loads, pm_make_images
+  ReduceContext, correct_load_store, pm_render, pm_add_loads
 from tinygrad.codegen.opt.postrange import apply_opts
 from tinygrad.codegen.late.gater import pm_move_gates_from_index
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse
@@ -58,16 +57,6 @@ pm_no_weakints = PatternMatcher([
   (UPat(GroupOp.All, dtype=dtypes.weakint, name="x"), lambda x: x.replace(dtype=dtypes.int))
 ])
 
-pm_fix_image_shrink = PatternMatcher([
-  # proper image index
-  (UPat(Ops.SHRINK, src=(UPat(Ops.RESHAPE, src=(UPat(Ops.PARAM, name="img"), UPat())), UPat(name="idx"), UPat()), name="out"),
-   lambda img,out,idx: img.reshape(-1, 4).index(idx.src[0]//4) if len(img.shape) == 3 and out.shape == (4,) else None),
-  # demote this to non image
-  (UPat(Ops.INDEX, src=(UPat(Ops.RESHAPE, src=(UPat(Ops.PARAM, name="img"), UPat())), UPat(name="idx")), name="out"),
-   lambda img,out,idx: img.replace(src=(shape_to_shape_arg((img.max_numel(),)),),
-                                   dtype=dtypes.half if img.dtype.itemsize == 2 else dtypes.float).index(idx) if out.shape == () else None)
-])+pm_mops
-
 def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
   if DEBUG >= 5: print(pyrender(ast))
@@ -114,10 +103,6 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   # add loads and remove invalids
   sink = graph_rewrite(sink, pm_add_loads+pm_remove_invalid, name="** add loads (code)")
 
-  # create image buffers
-  #if IMAGE and ren.target.device in {"QCOM", "CL", "PYTHON", "NULL"}:
-  #  sink = graph_rewrite(sink, pm_make_images, name="create image buffers", bottom_up=True, ctx=ren.target.arch)
-
   # devectorize
   sink = graph_rewrite(sink, sym+devectorize_alu+devectorize_buf_and_index+load_store_folding+correct_load_store+load_store_indexing,
                        ctx=ren, name="devectorize")
@@ -138,10 +123,6 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   # do memory coalesing (late)
   sink = memory_coalesing(sink, ren)
 
-  # image fixup
-  #if IMAGE and ren.target.device in {"QCOM", "CL", "PYTHON", "NULL"}:
-  #  sink = graph_rewrite(sink, pm_fix_image_shrink, name="fix image shrink")
-
   # instruction selection decompositions
   pm_decomp = pm_decomp+\
     get_late_rewrite_patterns(supported_ops, bool(DISABLE_FAST_IDIV))+\

tinygrad/codegen/late/coalese.py

@@ -1,24 +1,32 @@
-from typing import Any
+from typing import Any, cast
 import itertools
 from collections import defaultdict
-from tinygrad.dtype import dtypes, AddrSpace, Invalid, ImageDType
-from tinygrad.uop.ops import UOp, Ops
+from tinygrad.dtype import dtypes, AddrSpace, Invalid, ImageDType, PtrDType
+from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp
 from tinygrad.helpers import getenv, IMAGE
 from tinygrad.renderer import Renderer
+from tinygrad.uop.symbolic import uop_given_valid
+from tinygrad.codegen.late.devectorizer import simplify_valid_image_load, _drop_valid_stmts
 
-# image now lives here
+pm_imageh_store = PatternMatcher([
+  # store<imageh>(idx, x) is actually store(idx, x.cast(half)) so we can pull the cast into the store
+  (UPat(Ops.CAST, src=(UPat.var("x"),), name="c"), lambda x,c: x if c.dtype.scalar() == dtypes.half and x.dtype.scalar() == dtypes.float else None),
+  # store(imageh, a.where(b.half(), c).float()) -> store(imageh, a.where(b, c.float()))
+  (UPat(Ops.WHERE, src=(UPat.var("a"), UPat.var("b", dtypes.float).cast(dtypes.half), UPat.var("c"))), lambda a,b,c: a.where(b,c.cast(dtypes.float))),
+  (UPat(Ops.STACK, name="x"), lambda x: UOp.vectorize(*(s.cast(dtypes.float) for s in x.src)) if x.dtype.scalar() == dtypes.half else None),
+  # otherwise, we cast to float
+  (UPat(GroupOp.All, name="x"), lambda x: x.cast(dtypes.float))
+])
 
 def memory_coalesing(sink:UOp, ctx:Renderer) -> UOp:
   if getenv("DMC"): return sink
-  using_image = IMAGE and ctx.target.device in {"QCOM", "CL", "PYTHON", "NULL"}
 
   # collect
-  memory: defaultdict[tuple[Ops, UOp, Any, Any], dict[int, list[UOp]]]  = defaultdict(dict)
+  memory: defaultdict[tuple[Ops, UOp, Any, Any], dict[int, list[UOp]]] = defaultdict(dict)
   for u in sink.toposort():
-    # TODO: this should handle images too, it's just memory coalesing
     if u.op in {Ops.LOAD, Ops.STORE}:
       assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalesing does not support gated loads/stores"
-      if u.src[0].op is not Ops.INDEX: continue
+      if u.src[0].op is not Ops.INDEX or len(u.src[0].src) != 2: continue
       buf, idx_u = u.src[0].src
       if buf.addrspace == AddrSpace.REG: continue
       idx: Any = idx_u.src[1] if idx_u.op is Ops.WHERE and idx_u.src[2].arg is Invalid else idx_u
@@ -30,24 +38,35 @@ def memory_coalesing(sink:UOp, ctx:Renderer) -> UOp:
       else: root_src, arg = idx, 0
       memory[(u.op, buf, root_src, valid)].setdefault(arg, []).append(u)
 
+  image_bufs = {}
+  if IMAGE and ctx is not None and ctx.target.device in {"QCOM", "CL", "PYTHON", "NULL"}:
+    for _,buf,_,_ in memory:
+      if buf in image_bufs: continue
+      ibuf = buf if isinstance(buf.dtype, ImageDType) else None
+      if ibuf is None and buf.op is Ops.PARAM and buf.addrspace is AddrSpace.GLOBAL and isinstance(buf.dtype, PtrDType) and \
+         (dims:=ImageDType.valid_dims(buf.dtype, ctx.target.arch)):
+        ibuf = buf.replace(dtype=(dtypes.imageh if buf.dtype.base == dtypes.half else dtypes.imagef)((*dims[0], 4)))
+      if ibuf is not None: image_bufs[buf] = ibuf
+
   # build replacements
   replacements = {}
   for (op,buf,base,valid),offsets in memory.items():
+    if isinstance(buf.dtype, ImageDType) and buf not in image_bufs: continue
     # allowed lengths (copied in)
     lengths = []
     must_divide = True
-    if ctx is not None and ctx.target.device == "DSP":
+    if buf in image_bufs:
+      lengths = [4]
+    elif ctx is not None and ctx.target.device == "DSP":
       lengths = [128,64,32,16,8,4]
       must_divide = False
     elif buf.dtype.base not in (dtypes.float, dtypes.half, *dtypes.fp8s) and not isinstance(buf.dtype, ImageDType):
       pass
     elif buf.addrspace == AddrSpace.REG:
       pass
-    elif isinstance(buf.dtype, ImageDType):
-      lengths = [4]
     elif ctx is not None and ctx.supports_float4:
       # TODO: a better way to get this than ctx
-      lengths = [8,4,2] if buf.dtype.base == dtypes.half and getenv("ALLOW_HALF8") and not using_image else [4,2]
+      lengths = [8,4,2] if buf.dtype.base == dtypes.half and getenv("ALLOW_HALF8") else [4,2]
     lengths.append(1)  # worst case, it's not folded
     # do the grouping
     grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])]
@@ -56,23 +75,42 @@ def memory_coalesing(sink:UOp, ctx:Renderer) -> UOp:
         offset = (base+full_grp[0]) if isinstance(base, UOp) else UOp.const(dtypes.int, full_grp[0])
         length = [l for l in lengths if l <= len(full_grp) and (not must_divide or offset.divides(l) is not None)][0]
         grp = full_grp[:length]
-        if len(grp) == 4 and using_image:
-          print("IMAGE OPTION")
-          pass
-        idx = buf._mop(Ops.SHRINK, arg=[(offset, len(grp))]) if len(grp) > 1 else buf.index(offset)
+        ibuf, lane0 = image_bufs.get(buf) if length == 4 else None, None
+        if ibuf is None and op is Ops.LOAD and buf in image_bufs and (lane:=(offset%4).simplify()).op is Ops.CONST:
+          ibuf, lane0 = image_bufs[buf], lane.arg
+        idx_valid = valid
+        if ibuf is not None:
+          if valid is not None and not isinstance(buf.dtype, ImageDType) and isinstance(buf.dtype, PtrDType):
+            best_drop, cands = -1, []
+            for h,w in ImageDType.valid_dims(buf.dtype, ctx.target.arch):
+              cidx = uop_given_valid(valid, UOp.vectorize((offset // 4) % w, offset // (4*w)))
+              if (dropped:=len(_drop_valid_stmts(valid, cidx, h, w))) > best_drop: best_drop, cands = dropped, [(h, w, cidx)]
+              elif dropped == best_drop: cands.append((h, w, cidx))
+            if cands:
+              h, w, _ = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].gep(1).simplify().backward_slice))
+              ibuf = buf.replace(dtype=(dtypes.imageh if buf.dtype.base == dtypes.half else dtypes.imagef)((h, w, 4)))
+          idt = cast(ImageDType, ibuf.dtype)
+          idx_y, idx_x = (offset // (4*idt.shape[1])).simplify(), ((offset // 4) % idt.shape[1]).simplify()
+          idx = simplify_valid_image_load(ibuf, idx_y, idx_x, valid) if valid is not None else None
+          if idx is None: idx = ibuf.index(idx_y, idx_x, ptr=True)
+          idx_valid = None
+        else:
+          idx = buf._mop(Ops.SHRINK, arg=[(offset, len(grp))]) if len(grp) > 1 else buf.index(offset)
         if op == Ops.STORE:
           datas = []
           for i,g in enumerate(grp):
             assert len(offsets[g]) == 1, f"attempting multiple stores: {len(offsets[g])}"
             datas.append(offsets[g][0].src[1])
           data = UOp.vectorize(*datas) if len(datas) > 1 else datas[0]
-          store = idx.store(data, valid) if valid is not None else idx.store(data)
+          if ibuf is not None and ibuf.dtype.itemsize == 2: data = pm_imageh_store.rewrite(data)
+          store = idx.store(data, idx_valid) if idx_valid is not None else idx.store(data)
           for i,g in enumerate(grp): replacements[offsets[g][0]] = store
         else:
-          ld = idx.load(idx.vconst_like(0), valid) if valid is not None else idx.load()
+          ld = idx.load(idx.vconst_like(0), idx_valid) if idx_valid is not None else idx.load()
           for i,g in enumerate(grp):
             for oo in offsets[g]:
-              replacements[oo] = ld.index(UOp.const(dtypes.int, i)) if len(grp) > 1 else ld
+              ret = ld.index(UOp.const(dtypes.int, lane0 if lane0 is not None else i)) if ibuf is not None or len(grp) > 1 else ld
+              replacements[oo] = ret.cast(buf.dtype.base) if ibuf is not None and buf.dtype.base != dtypes.float else ret
         full_grp = full_grp[length:]
 
   # apply

tinygrad/codegen/late/devectorizer.py

@@ -5,7 +5,7 @@ from dataclasses import dataclass
 from tinygrad.dtype import dtypes, ImageDType, DType, AddrSpace, Invalid, PtrDType
 from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, identity_element
 from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
-from tinygrad.helpers import getenv, flatten, prod, OSX, ceildiv
+from tinygrad.helpers import getenv, flatten, prod
 from tinygrad.renderer import Renderer
 
 # ***** image load valid simplification *****
@@ -60,15 +60,6 @@ load_store_indexing = PatternMatcher([
 
 # ***** load/store grouping *****
 
-# get list of (height, width) that do not require pitch padding
-def image_valid_dims(base:DType, size:int, arch:str) -> list[tuple[int,int]]:
-  if (ALIGN:=next((int(p.split('=')[1]) for p in arch.split(',') if p.startswith("IMAGE_PITCH_ALIGNMENT=")), 0)) == 0: return []
-  MAXW, pxls = 16384, size // 4
-  if base not in (dtypes.half, dtypes.float) or size > 4*MAXW*MAXW: return []
-  # height=1 images just need to abide by alignment requirements in bytes, not pixels!
-  if size % (ALIGN * 4) != 0: return [] if (base.itemsize * size) % (64 if OSX else ALIGN) != 0 or pxls > MAXW else [(1, pxls)]
-  return [(pxls//ALIGN//k, ALIGN*k) for k in range(ceildiv(pxls//ALIGN, MAXW), min(pxls//ALIGN, MAXW//ALIGN)+1) if (pxls//ALIGN)%k == 0]
-
 def expand_index(ctx, buf:UOp, vec:UOp):
   # determine optimal image shapes
   if isinstance(dt:=buf.dtype, ImageDType):
@@ -77,7 +68,7 @@ def expand_index(ctx, buf:UOp, vec:UOp):
     x, valid = idxs.gep(lane), valids.gep(lane)
     # search for dims that drop the most valid statements
     best_drop, cands = -1, []
-    for ch, cw in image_valid_dims(dt.base, dt.size, ctx.target.arch):
+    for ch, cw in ImageDType.valid_dims(dt, ctx.target.arch):
       if (dropped:=len(_drop_valid_stmts(valid, cidx:=uop_given_valid(valid, UOp.vectorize((x//4)%cw, x//(4*cw))), ch, cw))) > best_drop:
         best_drop, cands = dropped, [(ch, cw, cidx)]
       elif dropped == best_drop: cands.append((ch, cw, cidx))
@@ -172,7 +163,7 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
   lengths = []
   must_divide = True
   # TODO: this belongs in coalese
-  #if isinstance(buf.dtype, ImageDType): lengths = [4]
+  if isinstance(buf.dtype, ImageDType): lengths = [4]
   lengths.append(1)  # worst case, it's not folded
 
   # filter fold lengths that don't divide
@@ -369,28 +360,3 @@ pm_add_loads = PatternMatcher([
   (UPat(Ops.STORE, src=(UPat(Ops.LOAD),), allow_any_len=True, name="s"), lambda s: s.replace(src=(s.src[0].src[0],)+s.src[1:])),
   (UPat(Ops.LOAD, src=(UPat(Ops.LOAD),), allow_any_len=True, name="l"), lambda l: l.replace(src=(l.src[0].src[0],)+l.src[1:])),
 ])
-
-# make images
-
-pm_imageh_store = PatternMatcher([
-  # store<imageh>(idx, x) is actually store(idx, x.cast(half)) so we can pull the cast into the store
-  (UPat.var("x", dtypes.float).cast(dtypes.half), lambda x: x),
-  # store(imageh, a.where(b.half(), c).float()) -> store(imageh, a.where(b, c.float()))
-  (UPat(Ops.WHERE, src=(UPat.var("a"), UPat.var("b", dtypes.float).cast(dtypes.half), UPat.var("c"))), lambda a,b,c: a.where(b,c.cast(dtypes.float))),
-  # otherwise, we cast to float
-  (UPat(GroupOp.All, name="x"), lambda x: x.cast(dtypes.float))
-])
-
-def make_image(ctx, ls, buf, off):
-  if (vcount:=buf.dtype.vcount) != 1: buf = buf.src[0]
-  if buf.op == Ops.PARAM and not isinstance(dt:=buf.dtype, ImageDType) and (dims:=image_valid_dims(dt.base, dt.size, ctx)):
-    buf = buf.replace(dtype=(dtypes.imageh if dt.base == dtypes.half else dtypes.imagef)((*dims[0], 4))).flatten()
-    if vcount != 1: buf = UOp.vectorize(*([buf] * vcount))
-    if ls.op is Ops.LOAD: return ls.replace(src=(buf.index(off, ptr=True),), dtype=dtypes.float.vec(ls.dtype.vcount)).cast(dt.base)
-    return buf.index(off, ptr=True).store(pm_imageh_store.rewrite(ls.src[1]) if dt.base == dtypes.half else ls.src[1])
-
-pm_make_images = PatternMatcher([
-  (UPat((Ops.LOAD, Ops.STORE), src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))),), allow_any_len=True, name="ls"), make_image),
-  # load<imageh> is actually load<half>.cast(float), so load<imageh>.half().float() -> load<half>.float().half().float() -> load<half>.float()
-  (UPat(Ops.LOAD, name="li").cast(dtypes.half).cast(dtypes.float), lambda li: li if isinstance(li.src[0].dtype, ImageDType) else None),
-])

tinygrad/codegen/opt/heuristic.py

@@ -1,9 +1,8 @@
 import itertools
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.helpers import getenv, DEBUG, prod, NOLOCALS, TC_OPT, TC_SELECT, USE_TC, IMAGE
-from tinygrad.dtype import PtrDType
+from tinygrad.dtype import PtrDType, ImageDType
 from tinygrad.uop.ops import Ops, resolve, AxisType
-from tinygrad.codegen.late.devectorizer import image_valid_dims
 from tinygrad.codegen.opt.postrange import Scheduler
 
 def hand_coded_optimizations(k:Scheduler) -> Scheduler:
@@ -51,7 +50,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
   # upcast float4 images, this must be early so we don't accidentally add locals before the upcast
   if IMAGE:
     for buf_index,buf in enumerate(k.bufs):
-      if isinstance(buf.src[0].dtype, PtrDType) and image_valid_dims(buf.src[0].dtype.base, buf.src[0].dtype.size, k.ren.target.arch):
+      if isinstance(buf.src[0].dtype, PtrDType) and ImageDType.valid_dims(buf.src[0].dtype, k.ren.target.arch):
         # part of is_expanded
         unit_stride_axes_mul_4 = [k.rngs.index(c) for c in k.bufs[buf_index].src[1].get_idx().split_uop(Ops.ADD) if
           c.op is Ops.RANGE and (c.vmax+1)%4 == 0]

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 ceildiv, getenv, prod, round_up, OSX
 from enum import IntEnum, auto
 
 class ConstFloat(float):
@@ -136,6 +136,15 @@ class ImageDType(PtrDType):
   @property
   def pitch(self): return (round_up(self.shape[1], 256) if OSX else self.shape[1]) * 4 * self.itemsize
 
+  # get list of (height, width) that do not require pitch padding
+  @staticmethod
+  def valid_dims(ptr:PtrDType, arch:str) -> list[tuple[int,int]]:
+    if (ALIGN:=next((int(p.split('=')[1]) for p in arch.split(',') if p.startswith("IMAGE_PITCH_ALIGNMENT=")), 0)) == 0: return []
+    MAXW, pxls = 16384, ptr.size // 4
+    if ptr.base not in (dtypes.half, dtypes.float) or ptr.size > 4*MAXW*MAXW: return []
+    # height=1 images just need to abide by alignment requirements in bytes, not pixels!
+    if ptr.size % (ALIGN * 4) != 0: return [] if ptr.nbytes() % (64 if OSX else ALIGN) != 0 or pxls > MAXW else [(1, pxls)]
+    return [(pxls//ALIGN//k, ALIGN*k) for k in range(ceildiv(pxls//ALIGN, MAXW), min(pxls//ALIGN, MAXW//ALIGN)+1) if (pxls//ALIGN)%k == 0]
 
 class dtypes:
   @staticmethod

tinygrad/runtime/ops_python.py

@@ -105,7 +105,7 @@ class PythonProgram:
           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)}"
+            assert len(src_values) == 3, "image index must be 3 srcs"
             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))

tinygrad/schedule/indexing.py

@@ -141,7 +141,7 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO
         symbolic+pm_simplify_valid, name="pad").where(r-off, UOp.invalid()) for r,sh,(off,sz) in zip(rngs, in_shape, arg))
     case Ops.RESHAPE:
       sink = UOp.sink(*rngs).simplify() # NOTE: this applies any commutative flips to the rngs early
-      sub_array = {r:UOp.range(r.src[0], i, AxisType.PLACEHOLDER, dtype=r.dtype) for i,r in enumerate(sink.ranges)}
+      sub_array = {r:UOp.range(r.src[0], i, AxisType.PLACEHOLDER) for i,r in enumerate(sink.ranges)}
       rngs = _apply_reshape(in_shape, arg, sink.substitute(sub_array)).substitute({v:k for k,v in sub_array.items()}).src
     case _: raise RuntimeError(f"{op} is not a MovementOp")
   return rngs