move decomp dtypes

tinygrad/codegen/__init__.py

@@ -16,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
@@ -104,16 +104,15 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   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)
+  #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")
-  sink = graph_rewrite(sink, load_store_indexing+gep_pushing, name="load/store indexing")
 
   # lower the index dtype to a concrete int
-  sink = graph_rewrite(sink, pm_lower_index_dtype, name="lower all index dtypes")
+  sink = graph_rewrite(sink, pm_lower_index_dtype+load_store_indexing+gep_pushing, name="lower all index dtypes")
   sink = graph_rewrite(sink, symbolic, name="post index symbolic")
 
   # optional pre matcher
@@ -123,11 +122,13 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   supported_ops = tuple(ren.code_for_op.keys())
   pm_decomp = symbolic_simple+get_simplifying_rewrite_patterns(supported_ops)
   sink = graph_rewrite(sink, pm_decomp, name="early decompositions")
-  sink = graph_rewrite(sink, pm_dtype_decomps, ctx=(set(), ren), name="decomp dtypes")
 
   # do memory coalesing (late)
   sink = memory_coalesing(sink, ren)
 
+  # do dtype decomps
+  sink = graph_rewrite(sink, pm_dtype_decomps, ctx=(set(), ren), name="decomp dtypes")
+
   # instruction selection decompositions
   pm_decomp = pm_decomp+\
     get_late_rewrite_patterns(supported_ops, bool(DISABLE_FAST_IDIV))+\

tinygrad/codegen/late/coalese.py

@@ -2,9 +2,46 @@ from typing import Any
 import itertools
 from collections import defaultdict
 from tinygrad.dtype import dtypes, AddrSpace, Invalid, ImageDType
-from tinygrad.uop.ops import UOp, Ops
-from tinygrad.helpers import getenv
+from tinygrad.uop.ops import UOp, Ops, graph_rewrite, PatternMatcher, UPat
+from tinygrad.helpers import getenv, IMAGE
 from tinygrad.renderer import Renderer
+from tinygrad.codegen.late.devectorizer import image_valid_dims, _drop_valid_stmts, uop_given_valid
+
+def transform_to_image(ctx, buf:UOp, x:UOp, valid:UOp|None=None) -> UOp|None:
+  # search for dims that drop the most valid statements
+  best_drop, cands = -1, []
+  for ch, cw in image_valid_dims(buf.dtype.base, buf.max_numel(), ctx.target.arch):
+    cidx = UOp.vectorize((x//4)%cw, x//(4*cw))
+    dropped = 0
+    if valid is not None:
+      cidx = uop_given_valid(valid, UOp.vectorize((x//4)%cw, x//(4*cw)))
+      dropped = len(_drop_valid_stmts(valid, cidx, ch, cw))
+    else:
+      cidx = cidx.simplify()
+    if dropped > best_drop: best_drop, cands = dropped, [(ch, cw, cidx)]
+    elif dropped == best_drop: cands.append((ch, cw, cidx))
+  # if no candidates, we don't rewrite
+  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].gep(1).simplify().backward_slice))
+  idx = buf.replace(dtype=(dtypes.imageh if buf.dtype.itemsize == 2 else dtypes.imagef)((h, w, 4))).index(cidx.src[1], cidx.src[0])
+  if valid is not None:
+    # TODO: simplify valid here
+    idx = valid.where(idx, UOp(Ops.CONST, dtype=idx.dtype, arg=Invalid))
+  return idx
+
+pm_add_image = PatternMatcher([
+  (UPat(Ops.SHRINK, src=(UPat(Ops.PARAM, name="buf"), UPat(name="x"), UPat(arg=4))).where(UPat.var("valid"), UPat(arg=Invalid)), transform_to_image),
+  (UPat(Ops.SHRINK, src=(UPat(Ops.PARAM, name="buf"), UPat(name="x"), UPat(arg=4))), transform_to_image),
+])
+
+pm_new_gater = PatternMatcher([
+  # here we create the alt value for load to be 0s and remove the where Invalid
+  (UPat.var("gate").where(UPat.var("idx"), UPat(Ops.CONST, arg=Invalid)).load(),
+   lambda gate,idx: idx.load(idx.vconst_like(0), gate)),
+  (UPat.var("gate").where(UPat.var("idx"), UPat(Ops.CONST, arg=Invalid)).store(UPat.var("data")),
+   lambda gate,idx,data: idx.store(data, gate)),
+])
 
 def memory_coalesing(sink:UOp, ctx:Renderer) -> UOp:
   if getenv("DMC"): return sink
@@ -12,7 +49,7 @@ def memory_coalesing(sink:UOp, ctx:Renderer) -> UOp:
   # collect
   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
+    # TODO: this should already have the gates in the new style, it shouldn't be required
     if u.op in {Ops.LOAD, Ops.STORE} and not isinstance(u.src[0].src[0].dtype, ImageDType):
       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
@@ -53,21 +90,28 @@ 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]
-        idx = buf._mop(Ops.SHRINK, arg=[(offset, len(grp))]) if len(grp) > 1 else buf.index(offset)
+        idx = UOp(Ops.SHRINK, dtype=buf.dtype, src=(buf, offset, UOp.const(dtypes.int, len(grp)))) if len(grp) > 1 else buf.index(offset)
+        # broadcasting!
+        idx = valid.where(idx, UOp(Ops.CONST, idx.dtype, arg=Invalid)) if valid is not None else idx
         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)
+          store = idx.store(UOp.vectorize(*datas) if len(datas) > 1 else datas[0])
           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()
           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
         full_grp = full_grp[length:]
 
   # apply
-  return sink.substitute(replacements, name="memory coalesing")
+  sink = sink.substitute(replacements, name="memory coalesing")
+  # image
+  if IMAGE and ctx.target.device in {"QCOM", "CL", "PYTHON", "NULL"}:
+    sink = graph_rewrite(sink, pm_add_image, name="add image", ctx=ctx, bottom_up=True)
+  # new gater
+  sink = graph_rewrite(sink, pm_new_gater, name="new gater")
+  return sink

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
+from tinygrad.helpers import getenv, flatten, prod, OSX, ceildiv
 from tinygrad.renderer import Renderer
 
 # ***** image load valid simplification *****
@@ -60,6 +60,15 @@ 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):
@@ -68,7 +77,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 ImageDType.valid_dims(dt, ctx.target.arch):
+    for ch, cw in image_valid_dims(dt.base, dt.size, 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))
@@ -374,7 +383,7 @@ pm_imageh_store = PatternMatcher([
 
 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:=ImageDType.valid_dims(dt, ctx)):
+  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)

tinygrad/codegen/opt/heuristic.py

@@ -1,8 +1,9 @@
 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, ImageDType
+from tinygrad.dtype import PtrDType
 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:
@@ -50,7 +51,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 ImageDType.valid_dims(buf.src[0].dtype, k.ren.target.arch):
+      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):
         # 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 ceildiv, getenv, prod, round_up, OSX
+from tinygrad.helpers import getenv, prod, round_up, OSX
 from enum import IntEnum, auto
 
 class ConstFloat(float):
@@ -136,15 +136,6 @@ 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/helpers.py

@@ -275,7 +275,7 @@ SCACHE = ContextVar("SCACHE", 1)
 # allow use of atomics for embedding backward
 USE_ATOMICS = ContextVar("USE_ATOMICS", 0)
 # don't allow broadcast
-DISALLOW_BROADCAST = ContextVar("DISALLOW_BROADCAST", 1)
+DISALLOW_BROADCAST = ContextVar("DISALLOW_BROADCAST", 0)
 
 @dataclass(frozen=True)
 class Metadata:

tinygrad/renderer/cstyle.py

@@ -314,12 +314,17 @@ class OpenCLRenderer(CStyleLanguage):
       lambda ctx,x: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(x.arg)))[0] >> 16)}u"),
     # load/store image (OpenCL)
     (UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), lambda ctx,buf,idx_y,idx_x: f"IMAGE<{ctx[buf]}, {ctx[idx_y]}, {ctx[idx_x]}>"),
-    (UPat(Ops.LOAD, dtype=dtypes.float, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("var"), UPat.var("gate"))),
-      lambda ctx,buf,idx_y,idx_x,var,gate: f"({ctx[gate]}?read_imagef({ctx[buf]}, smp, (int2)({ctx[idx_x]},{ctx[idx_y]})):{ctx[var]})"),
-    (UPat(Ops.LOAD, dtype=dtypes.float, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')),)),
-      lambda ctx,buf,idx_y,idx_x: f"read_imagef({ctx[buf]}, smp, (int2)({ctx[idx_x]},{ctx[idx_y]}))"),
-    (UPat(Ops.STORE, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("var", dtypes.float))),
-      lambda ctx,buf,idx_y,idx_x,var: f"write_imagef({ctx[buf]}, (int2)({ctx[idx_x]},{ctx[idx_y]}), {ctx[var]});"),
+    (UPat(Ops.LOAD, name="l", dtype=(dtypes.float,dtypes.half),
+          src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("var"), UPat.var("gate"))),
+      lambda ctx,buf,idx_y,idx_x,var,gate,l:
+        f"({ctx[gate]}?read_image{'f' if l.dtype==dtypes.float else 'h'}({ctx[buf]}, smp, (int2)({ctx[idx_x]},{ctx[idx_y]})):{ctx[var]})"),
+    (UPat(Ops.LOAD, name="l", dtype=(dtypes.float,dtypes.half),
+          src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')),)),
+      lambda ctx,buf,idx_y,idx_x,l:
+        f"read_image{'f' if l.dtype==dtypes.float else 'h'}({ctx[buf]}, smp, (int2)({ctx[idx_x]},{ctx[idx_y]}))"),
+    (UPat(Ops.STORE, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("var", (dtypes.float,dtypes.half)))),
+      lambda ctx,buf,idx_y,idx_x,var:
+        f"write_image{'f' if var.dtype==dtypes.float else 'h'}({ctx[buf]}, (int2)({ctx[idx_x]},{ctx[idx_y]}), {ctx[var]});"),
   ]) + base_rewrite
 
   def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str:

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, "image index must be 3 srcs"
+            assert len(src_values) == 3, f"image index must be 3 srcs, not {len(src_values)}"
             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) for i,r in enumerate(sink.ranges)}
+      sub_array = {r:UOp.range(r.src[0], i, AxisType.PLACEHOLDER, dtype=r.dtype) 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