StarPilot/tinygrad_repo/tinygrad/codegen/late/coalesce.py

import itertools, functools
from collections import defaultdict
from dataclasses import replace
from tinygrad.dtype import dtypes, AddrSpace, Invalid, DType
from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, graph_rewrite
from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate, sym
from tinygrad.helpers import getenv, IMAGE, OSX, ceildiv, is_image_shape
from tinygrad.renderer import Renderer

# ***** image load valid simplification *****

@functools.cache
def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]:
  # can drop valid if idx is out of bound when valid is False
  drop_stmt = []
  for i,stmt in enumerate(valid.split_uop(Ops.AND)):
    if (res:=parse_valid(stmt)) is None: continue
    X, is_upper_bound, c = res

    # for X0 + X1 + ... >= 1, check if it's out of bound when Xi = 0 for all i
    if not is_upper_bound and c == 1 and all(u.op in GroupOp.Irreducible and u.vmin == 0 for u in X.split_uop(Ops.ADD)):
      testidx = functools.reduce(lambda nowidx,u: nowidx.substitute({u:u.const_like(0)}), X.split_uop(Ops.ADD), idx)
      if testidx.index(0).vmax < 0 or testidx.index(1).vmax < 0:
        drop_stmt.append(stmt)
        continue

    # check if idx is out of bound when X is on the wrong side of the bound: X in [c+1, vmax] or [vmin, c-1]
    lo, hi = (c + 1, X.vmax) if is_upper_bound else (X.vmin, c - 1)
    if lo <= hi:
      fake = UOp.variable(f"fake{i}", lo, hi, X.dtype)
      subs = [{X: fake}]
      # idx may not have X itself, so also substitute a term of X: v -> fake - (X - v)
      terms = list(X.split_uop(Ops.ADD))
      v = next((u for u in terms if u.op in GroupOp.Irreducible and u.op is not Ops.CONST), None)
      if v is not None and (rest:=[u for u in terms if u is not v]): subs.append({v: fake - UOp.usum(*rest)})
      if any((testidx:=graph_rewrite(coord.substitute(sub), sym)).vmin >= b or testidx.vmax < 0
             for sub in subs for coord,b in zip(idx.src, (width, height))):
        drop_stmt.append(stmt)
  return drop_stmt

def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
  idx = uop_given_valid(valid, start_idx)
  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 is_image_shape(buf._shape): return None
  if idx_x.dtype != idx_y.dtype: idx_x, idx_y = idx_x.cast(dtypes.int), idx_y.cast(dtypes.int)
  start_idx = idx_x.stack(idx_y)
  idx = uop_given_valid(valid, start_idx)
  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)
  if new_valid is not None: return buf.index(idx_y.valid(new_valid), idx_x.valid(new_valid))
  return buf.index(idx_y, idx_x)

indexing_simplify = PatternMatcher([
  # image load valid idx simplification
  (UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)),
  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("valid").where(UPat.var("idx_y"), UPat(arg=Invalid)),
                                         UPat.var("valid").where(UPat.var("idx_x"), UPat(arg=Invalid)))), simplify_valid_image_load),
])

# 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 transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
  shapes, ren, read_write = ctx
  if not IMAGE or ren.target.device not in {"QCOM", "CL", "PYTHON", "NULL"} or buf in read_write: return None
  valid, x = x.get_valid(), x.get_idx()
  # search for dims that drop the most valid statements
  best_drop, cands = -1, []
  for ch, cw in [shapes[buf.arg.slot]] if buf.arg.slot in shapes else image_valid_dims(buf.dtype, buf.max_numel(), ren.target.arch):
    cidx = uop_given_valid(valid, ((x//4)%cw).stack(x//(4*cw)))
    dropped = len(_drop_valid_stmts(valid, cidx, ch, cw))
    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].index(1).simplify().backward_slice))
  # the image dims are stored in the param's arg, the size stays the flat buffer len
  buf = buf.replace(arg=replace(buf.arg, image=(h, w)))
  shapes[buf.arg.slot] = (h, w)
  return buf.index(cidx.src[1].valid(valid), cidx.src[0].valid(valid))

def store_image(x:UOp, d:UOp) -> UOp:
  # image load/store is always float, a half image converts on store
  def as_float(s:UOp):
    return s.src[0] if x.src[0].dtype is dtypes.half and s.op is Ops.CAST and s.src[0].dtype is dtypes.float else s.cast(dtypes.float)
  lanes = d.src if d.op is Ops.STACK else tuple(d.index(i) for i in range(d.shape[0]))
  return x.store(UOp.stack(*[as_float(s) for s in lanes]))

pm_simplify_add_image = PatternMatcher([
  (UPat(Ops.SHRINK, src=(UPat(Ops.PARAM, name="buf"), UPat(name="x"), UPat(arg=4))), transform_to_image),
  (UPat(Ops.INDEX, dtype=dtypes.float, name="x").store(UPat(name="d", dtype=dtypes.half)), store_image),
])

def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
  if getenv("DMC"): return sink

  # collect
  memory: defaultdict[tuple[Ops, UOp, UOp|None, UOp, object, int|tuple[int, int]], dict[int, list[UOp]]] = defaultdict(dict)
  # depth[u] is the lexicographically largest (GLOBAL stores, LOCAL stores) count on any chain u waits on, u included
  depth: dict[UOp, tuple[int, int]] = {}
  for u in sink.toposort():
    depth[u] = waits = max([depth[s] for s in u.src_without_body], default=(0, 0))
    if u.op is Ops.STORE: depth[u] = (waits[0]+(u.src[0].addrspace is AddrSpace.GLOBAL), waits[1]+(u.src[0].addrspace is AddrSpace.LOCAL))
    if u.op in {Ops.LOAD, Ops.STORE}:
      assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalescing does not support gated loads/stores"
      assert u.src[0].op is Ops.INDEX, f"memory coalescing should be on INDEX, not {u.src[0].op}"
      buf, idx_u = u.src[0].src
      if buf.addrspace == AddrSpace.REG: continue
      if buf.buf_uop.op is Ops.PARAM and buf.buf_uop.arg.volatile: continue # volatile accesses never merge
      idx, valid = idx_u.get_idx(), idx_u.get_valid()
      if idx.is_invalid: continue
      root_src: UOp|None
      if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].val
      elif idx.op is Ops.ADD and idx.src[0].op is Ops.CONST: root_src, arg = idx.src[1], idx.src[0].val
      elif idx.op is Ops.CONST: root_src, arg = None, idx.val
      else: root_src, arg = idx, 0
      # loads/stores only coalesce with others carrying the same arg (e.g. the nontemporal flag)
      # and the same depth (GLOBAL count only for GLOBAL accesses), so no merge closes a dependency cycle
      memory[(u.op, buf, root_src, valid, u.arg, waits[0] if buf.addrspace is AddrSpace.GLOBAL else waits)].setdefault(arg, []).append(u)

  # build replacements
  replacements = {}
  for (op,buf,_,_,ld_arg,_),offsets in memory.items():
    # allowed lengths (copied in)
    lengths = []
    must_divide = True
    if ctx.target.device == "DSP":
      lengths = [128,64,32,16,8,4]
      must_divide = False
    elif buf.dtype in (dtypes.float, dtypes.half, dtypes.int, dtypes.uint, *dtypes.fp8s) and ctx.supports_float4:
      # TODO: a better way to get this than ctx
      lengths = [8,4,2] if buf.dtype == dtypes.half and getenv("ALLOW_HALF8") else [4,2]
    lengths.append(1)  # worst case, it's not folded
    # do the grouping
    for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0]):
      full_grp = [offsets[x] for _,x in group]
      while len(full_grp):
        # the group starts at its first access's index, valid included
        offset = full_grp[0][0].src[0].src[1]
        length = [l for l in lengths if l <= len(full_grp) and (not must_divide or offset.get_idx().divides(l) is not None)][0]
        grp, full_grp = full_grp[:length], full_grp[length:]
        assert op is not Ops.STORE or all(len(g) == 1 for g in grp), "attempting multiple stores"
        if length == 1: continue
        idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(length)))
        if op == Ops.STORE:
          datas = [g[0].src[1] for g in grp]
          store = idx.store(UOp.stack(*datas))
          for g in grp: replacements[g[0]] = store
        else:
          ld = idx.load(arg=ld_arg)
          for i,g in enumerate(grp):
            for oo in g: replacements[oo] = ld.index(i)

  # apply
  return sink.substitute(replacements, name="memory coalescing")