mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-08-23 09:13:43 +08:00
The smallest promotion
This commit is contained in:
firestar5683
@@ -1,7 +1,7 @@
|
||||
import itertools
|
||||
from typing import Callable
|
||||
from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start, AxisType
|
||||
from tinygrad.uop.symbolic import symbolic
|
||||
from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const, invalid_gate
|
||||
from tinygrad.helpers import partition
|
||||
from tinygrad.dtype import dtypes
|
||||
|
||||
@@ -9,8 +9,9 @@ def flatten_range(r:UOp) -> UOp|None:
|
||||
off = range_start[r.op]
|
||||
rngs = r.src[off:]
|
||||
if not len(rngs): return None
|
||||
new_rngs = [x for x in UOp.sink(*rngs).toposort() if x.op is Ops.RANGE]
|
||||
return r.replace(src=r.src[:off]+tuple(new_rngs))
|
||||
# ranges in the cond should not be ended
|
||||
backedge = tuple(x for x in rngs if x.dtype in (dtypes.void, dtypes.bool))
|
||||
return r.replace(src=r.src[:off]+tuple(UOp.sink(*[x for x in rngs if x not in backedge]).ranges)+backedge)
|
||||
|
||||
pm_flatten_range = PatternMatcher([
|
||||
# real ranges only
|
||||
@@ -20,6 +21,7 @@ pm_flatten_range = PatternMatcher([
|
||||
# index/range arithmetic uses FLOORDIV/FLOORMOD prior to late rewrite
|
||||
def count_divmod(x:UOp) -> int: return sum(u.op in {Ops.FLOORDIV, Ops.FLOORMOD} for u in x.backward_slice)
|
||||
def simplify_merge_adjacent(u:UOp) -> UOp|None:
|
||||
if not all(r.op is Ops.RANGE for r in u.ended_ranges): return None
|
||||
reduce_ranges = [x.ranges for x in u.backward_slice_with_self if x.op is Ops.REDUCE]
|
||||
# on END we only want to merge adjacent ranges, on REDUCE we want to try all combinations
|
||||
for r0, r1 in (zip(u.ended_ranges, u.ended_ranges[1:]) if u.op is Ops.END else itertools.permutations(u.ended_ranges, 2)):
|
||||
@@ -30,7 +32,7 @@ def simplify_merge_adjacent(u:UOp) -> UOp|None:
|
||||
s0, s1 = r0.src[0], r1.src[0]
|
||||
# do the merge
|
||||
new_range = r0.replace(src=(s0*s1,))
|
||||
nidx = graph_rewrite(u, _substitute+symbolic+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
|
||||
nidx = graph_rewrite(u, _substitute+symbolic+pm_fold_cast_const+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
|
||||
name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")
|
||||
|
||||
# check if it simplifies
|
||||
@@ -39,13 +41,13 @@ def simplify_merge_adjacent(u:UOp) -> UOp|None:
|
||||
return u
|
||||
|
||||
def mark_gated(ctx, idx):
|
||||
if idx.src[1].op is Ops.WHERE:
|
||||
if len(idx.src) > 1 and idx.src[1].op is Ops.WHERE:
|
||||
x, cond = idx.src[1].get_idx(), idx.src[1].get_valid()
|
||||
# get all ranges r with guards "r < c" for some const c
|
||||
guards = {r:c for v in cond.split_uop(Ops.AND) if v.op is Ops.CMPLT and (r:=v.src[0]).op is Ops.RANGE and (c:=v.src[1]).op is Ops.CONST}
|
||||
else: x, guards = idx, {}
|
||||
# ensure that we choose max(c_i) for all i where r < c_i
|
||||
ctx |= {r:c for r,c in guards.items() if (r not in ctx or ctx[r].arg < c.arg)}
|
||||
ctx |= {r:c for r,c in guards.items() if (r not in ctx or ctx[r].val < c.val)}
|
||||
# but if a range is ever ungated, we cannot shrink it
|
||||
ctx |= {r:r.src[0] for r in x.ranges if r not in guards}
|
||||
|
||||
@@ -58,7 +60,9 @@ pm_simplify_ranges = PatternMatcher([
|
||||
])
|
||||
|
||||
def mark_range_mod(ctx:dict[UOp, UOp|None], r:UOp, c:UOp) -> None:
|
||||
if r not in ctx and r.arg[-1] is not AxisType.WARP and r.src[0].op is Ops.CONST and r.src[0].divides(c.arg) is not None: ctx[r] = c
|
||||
# ranges that aren't looped over can't be split
|
||||
if r not in ctx and r.arg[-1] not in {AxisType.WARP, AxisType.DEVICE} \
|
||||
and r.src[0].op is Ops.CONST and r.src[0].divides(c.val) is not None: ctx[r] = c
|
||||
|
||||
def do_substitute(ctx:dict, x: UOp, sub_fxn:Callable[[UOp, UOp], UOp]) -> UOp|None:
|
||||
ret = x.substitute({k:sub_fxn(k,v) for k,v in ctx.items() if v is not None})
|
||||
@@ -82,9 +86,9 @@ def reduce_unparented(red:UOp) -> UOp|None:
|
||||
if len(reduce_unparented) == 0: return None
|
||||
ret = red.replace(src=(red.src[0],)+tuple(reduce_parented)) if len(reduce_parented) or red.dtype != red.src[0].dtype else red.src[0]
|
||||
if red.arg[0] is Ops.ADD:
|
||||
for r in reduce_unparented: ret = ret * r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
|
||||
for r in reduce_unparented: ret = ret * r.src[0]
|
||||
if red.arg[0] is Ops.MUL:
|
||||
for r in reduce_unparented: ret = ret ** r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
|
||||
for r in reduce_unparented: ret = ret ** r.src[0]
|
||||
return ret
|
||||
|
||||
pm_reduce_unparented = PatternMatcher([
|
||||
@@ -94,27 +98,25 @@ pm_reduce_unparented = PatternMatcher([
|
||||
|
||||
pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
|
||||
# lift x+y out of reduce on lt
|
||||
((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"), lambda x,y,c: (x < (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
|
||||
((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"), lambda x,y,c: (x < (c-y)) if no_range(y) and no_range(c) else None),
|
||||
# lift x*y out of reduce
|
||||
((UPat.var("x")*UPat.var("y")) < UPat.var("c"),
|
||||
lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and dtypes.is_int(y.dtype) and y.vmin > 0 else None),
|
||||
# fold the range
|
||||
# bound from below
|
||||
((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(0, UPat.var("val")).reduce(UPat.var("r"), arg=Ops.ADD),
|
||||
lambda r,cut,val: (r.src[0]-cut).maximum(0).minimum(r.src[0]).cast(val.dtype) * val if no_range(val) else None),
|
||||
# bound from two sides
|
||||
(((UPat.var("r")<UPat.var("lower")).logical_not()&(UPat(Ops.RANGE, name="r")<UPat.var("upper"))).where(UPat.var("val"), 0).reduce(UPat.var("r"),
|
||||
arg=Ops.ADD), lambda r,lower,upper,val:
|
||||
(upper.minimum(r.src[0])-lower.maximum(0)).maximum(0).minimum(r.src[0]).cast(val.dtype) * val if no_range(val) else None),
|
||||
# bound from above
|
||||
((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(UPat.var("val"), 0).reduce(UPat.var("r"), arg=Ops.ADD),
|
||||
lambda r,cut,val: cut.maximum(0).minimum(r.src[0]).cast(val.dtype) * val if no_range(val) else None),
|
||||
# REDUCE on ADD
|
||||
# sum over r in [0,N) of [lower<=r<upper]*val -> clamp(min(upper,N) - max(lower,0), 0, N) * val
|
||||
(UPat.any(
|
||||
(UPat(Ops.RANGE, name="r") < UPat.var("upper")).where(UPat.var("val"), 0),
|
||||
(UPat(Ops.RANGE, name="r") < UPat.var("lower")).where(0, UPat.var("val")),
|
||||
((UPat.var("r")<UPat.var("lower")).logical_not()&(UPat(Ops.RANGE, name="r")<UPat.var("upper"))).where(UPat.var("val"), 0),
|
||||
).reduce(UPat.var("r"), arg=Ops.ADD), lambda r,val,lower=None,upper=None:
|
||||
((upper.minimum(r.src[0]) if upper is not None else r.src[0]) -
|
||||
(lower.maximum(0) if lower is not None else r.const_like(0))).maximum(0).minimum(r.src[0]) * val if no_range(val) else None),
|
||||
(invalid_gate.reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
|
||||
lambda cond,x,i,r: cond.where(x.reduce(*r.src[1:], arg=Ops.ADD), i) if no_range(cond) else None),
|
||||
((UPat.var("x")+UPat.var("y")).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
|
||||
lambda x,y,r: x.reduce(*r.src[1:], arg=Ops.ADD) + y.reduce(*r.src[1:],arg=Ops.ADD)),
|
||||
# AND on WHERE
|
||||
((UPat(Ops.DEFINE_VAR, name="x") & UPat.var("y")).where(UPat.var("c"), 0).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
|
||||
lambda x,y,c,r: y.where(c, 0).reduce(*r.src[1:], arg=Ops.ADD)*x.cast(c.dtype)),
|
||||
((UPat(Ops.PARAM, name="x") & UPat.var("y")).where(UPat.var("c"), 0).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
|
||||
lambda x,y,c,r: y.where(c, 0).reduce(*r.src[1:], arg=Ops.ADD)*x),
|
||||
# MUL casted bool
|
||||
((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast()), lambda x,gate: gate.where(x, 0)),
|
||||
])+symbolic
|
||||
@@ -134,7 +136,7 @@ def reduce_collapse(red:UOp, u:UOp, pm:PatternMatcher=pm_reduce_collapse) -> UOp
|
||||
replaces: dict[UOp, UOp] = {}
|
||||
for u in included:
|
||||
for s in u.src:
|
||||
if s in included or s in replaces or s.op in {Ops.CONST, Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}: continue
|
||||
if s in included or s in replaces or s.op in {Ops.CONST, Ops.PARAM, Ops.BUFFER}: continue
|
||||
replaces[s] = UOp.variable(f'in{len(replaces)}', s.vmin, s.vmax, s.dtype)
|
||||
collapse_fxn = u.substitute(replaces).reduce(r, arg=Ops.ADD)
|
||||
sink = graph_rewrite(collapse_fxn, pm, name="reduce_collapse")
|
||||
@@ -146,12 +148,12 @@ def reduce_load_collapse(red:UOp, u:UOp) -> UOp|None: return reduce_collapse(red
|
||||
|
||||
# remove REDUCE without loads (generic arange opt / indexing).
|
||||
pm_reduce_simplify = pm_reduce_unparented + PatternMatcher([
|
||||
(UPat(Ops.REDUCE, src=(UPat.var("u"),), allow_any_len=True, arg=(Ops.ADD, ()), name="red"), reduce_collapse),
|
||||
(UPat(Ops.REDUCE, src=(UPat.var("u"),), allow_any_len=True, arg=(Ops.ADD, 0), name="red"), reduce_collapse),
|
||||
])
|
||||
# remove REDUCE on load, comes from indexing a tensor with another tensor
|
||||
def no_load(u:UOp) -> bool: return not any(x.op is Ops.INDEX for x in u.backward_slice_with_self)
|
||||
pm_load_collapse = PatternMatcher([
|
||||
(UPat(Ops.REDUCE, arg=(Ops.ADD, ()), src=(UPat.var("u"), UPat()), name="red"), reduce_load_collapse),
|
||||
(UPat(Ops.REDUCE, arg=(Ops.ADD, 0), src=(UPat.var("u"), UPat()), name="red"), reduce_load_collapse),
|
||||
# we want to make sure we dont do math on a loaded index since that can cause overflow, this undoes the rule in pm_reduce_load_collapse
|
||||
((UPat.var("x", dtypes.weakint)+UPat.var("y"))<UPat.var("c"), lambda x,y,c: x < c-y if no_load(y) and no_load(c) and not no_load(x) else None),
|
||||
])
|
||||
|
||||
Reference in New Issue
Block a user