mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-09-11 02:33:51 +08:00
firestar5683
171 lines
9.1 KiB
Python
171 lines
9.1 KiB
Python
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, param=True)
|
|
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 = ctx
|
|
if not IMAGE or ren.target.device not in {"QCOM", "CL", "PYTHON", "NULL"}: 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)
|
|
if valid.op is not Ops.CONST or valid.val is not True:
|
|
return buf.index(cidx.src[1].valid(valid), cidx.src[0].valid(valid))
|
|
else:
|
|
return buf.index(cidx.src[1], cidx.src[0])
|
|
|
|
pm_simplify_add_image = PatternMatcher([
|
|
(UPat(Ops.SHRINK, src=(UPat(Ops.PARAM, name="buf"), UPat(name="x"), UPat(arg=4))), transform_to_image),
|
|
# image load/store is always float
|
|
(UPat(Ops.INDEX, dtype=dtypes.float, name="x").store(UPat(name="d", dtype=dtypes.half)), lambda x,d: x.store(d.cast(dtypes.float))),
|
|
(UPat.var("x", dtype=dtypes.float).cast(dtypes.half).cast(dtypes.float), lambda x: x),
|
|
])
|
|
|
|
def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
|
|
if getenv("DMC"): return sink
|
|
|
|
# collect
|
|
memory: defaultdict[tuple[Ops, UOp, UOp|str, UOp, object], dict[int, list[UOp]]] = defaultdict(dict)
|
|
for u in sink.toposort():
|
|
# TODO: this should handle images too, it's just memory coalescing
|
|
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
|
|
idx, valid = idx_u.get_idx(), idx_u.get_valid()
|
|
root_src: UOp|str
|
|
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 and idx.val is Invalid: root_src, arg = "INVALID", 0
|
|
elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.val
|
|
else: root_src, arg = idx, 0
|
|
# loads/stores only coalesce with others carrying the same arg (e.g. the nontemporal flag)
|
|
memory[(u.op, buf, root_src, valid, u.arg)].setdefault(arg, []).append(u)
|
|
|
|
# build replacements
|
|
replacements = {}
|
|
for (op,buf,base,valid,ld_arg),offsets in memory.items():
|
|
# allowed lengths (copied in)
|
|
lengths = []
|
|
must_divide = True
|
|
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.int, dtypes.uint, *dtypes.fp8s) and not is_image_shape(buf._shape):
|
|
pass
|
|
elif buf.addrspace == AddrSpace.REG:
|
|
pass
|
|
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
|
|
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
|
|
grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])]
|
|
for full_grp in grouped_offsets:
|
|
while len(full_grp):
|
|
offset = (base+full_grp[0]) if isinstance(base, UOp) else UOp.const(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]
|
|
# NOTE: we apply the valid again after we determine the length
|
|
offset = offset.valid(valid) if valid is not None else offset
|
|
idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(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])
|
|
store = idx.store(UOp.stack(*datas) if len(datas) > 1 else datas[0])
|
|
for i,g in enumerate(grp): replacements[offsets[g][0]] = store
|
|
else:
|
|
ld = idx.load(arg=ld_arg)
|
|
for i,g in enumerate(grp):
|
|
for oo in offsets[g]:
|
|
replacements[oo] = ld.index(i) if len(grp) > 1 else ld
|
|
full_grp = full_grp[length:]
|
|
|
|
# apply
|
|
return sink.substitute(replacements, name="memory coalescing")
|