IQ.Pilot Prebuilt Release @ 27f668a
This commit is contained in:
IQ.Lvbs CI [bot]
165
artifacts/package_runtime/tinygrad/codegen/late/coalesce.py
Normal file
165
artifacts/package_runtime/tinygrad/codegen/late/coalesce.py
Normal file
@@ -0,0 +1,165 @@
|
||||
import itertools, functools
|
||||
from collections import defaultdict
|
||||
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, 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)
|
||||
for coord,b in zip(idx.src, (width, height)):
|
||||
rw = coord.substitute({X:fake}).simplify()
|
||||
if rw.vmin >= b or rw.vmax < 0:
|
||||
drop_stmt.append(stmt)
|
||||
break
|
||||
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), dtype=dtypes.float)
|
||||
return buf.index(idx_y, idx_x, dtype=dtypes.float)
|
||||
|
||||
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))
|
||||
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.val is not True:
|
||||
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], 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),
|
||||
# image load/store is always float
|
||||
(UPat(Ops.INDEX, dtype=dtypes.float, name="x").load(dtype=dtypes.half), lambda x: x.load().cast(dtypes.half)),
|
||||
(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], 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
|
||||
memory[(u.op, buf, root_src, valid)].setdefault(arg, []).append(u)
|
||||
|
||||
# build replacements
|
||||
replacements = {}
|
||||
for (op,buf,base,valid),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()
|
||||
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")
|
||||
Reference in New Issue
Block a user