forked from tinygrad/tinygrad
move reduce collapse to rangeify (#12845)
This commit is contained in:
Sieds Lykles
@@ -182,7 +182,7 @@ class TestReduceOpsConstFolding(unittest.TestCase):
|
||||
np.testing.assert_equal(Tensor(4).sum().numpy(), 4)
|
||||
|
||||
def test_padded_const_sum(self):
|
||||
_check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).sum())
|
||||
_check_ast_count(0, Tensor.ones(4).pad(((1, 1),)).sum())
|
||||
np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).sum().numpy(), 4)
|
||||
|
||||
# NOTE: cannot just count the non-padded area because some Ops f do not have f(0) = 0.
|
||||
|
||||
+17
-17
@@ -2,7 +2,7 @@ import unittest, pytest
|
||||
from tinygrad import dtypes, Variable
|
||||
from tinygrad.dtype import AddrSpace
|
||||
from tinygrad.helpers import DEBUG, Context
|
||||
from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp
|
||||
from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp, AxisType
|
||||
from tinygrad.uop.symbolic import sym
|
||||
from tinygrad.codegen import full_rewrite_to_sink
|
||||
from tinygrad.codegen.late.expander import expander
|
||||
@@ -460,23 +460,23 @@ class TestUOpGraph(unittest.TestCase):
|
||||
if u.op is Ops.STORE: assert u.src[1].arg==5
|
||||
|
||||
def test_load_idx_becomes_int(self):
|
||||
# These loads wont overflow int since we know from the gate that the value is bounded
|
||||
r0 = UOp.range(10, 0)
|
||||
d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0)
|
||||
d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 1)
|
||||
l0 = UOp(Ops.LOAD, dtypes.long, (d0.index(UOp.const(dtypes.int, 0)),)).cast(dtypes.index)
|
||||
idx = l0 * 600
|
||||
valid = (l0<-1).ne(True)&(l0<3000)
|
||||
l1 = valid.where(UOp(Ops.LOAD, dtypes.long, (d1.index(idx),)),0)
|
||||
uops = to_uops_list([l1])
|
||||
# mnist indexing with split reduceop
|
||||
# Make sure we are not doign math on the loaded index, which would promote it to long
|
||||
c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(128000), arg=0, src=())
|
||||
c1 = UOp.range(UOp.const(dtypes.index, 512), 1, AxisType.LOOP)
|
||||
c2 = UOp.range(UOp.const(dtypes.index, 250), 2, AxisType.LOOP)
|
||||
c3 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(512), arg=1, src=())
|
||||
c4 = c3.index(c1).load()
|
||||
c5 = UOp.range(UOp.const(dtypes.index, 240), 0, AxisType.REDUCE)
|
||||
c6 = ((c2*UOp.const(dtypes.index, 240))+c5)
|
||||
c7 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(60000), arg=2, src=())
|
||||
c8 = c7.index(c6).load()
|
||||
c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD)
|
||||
c10 = c0.index(((c1*UOp.const(dtypes.index, 250))+c2)).store(c9, c1, c2)
|
||||
ast = c10.sink()
|
||||
uops = to_uops_list([ast])
|
||||
for u in uops:
|
||||
if u.op is Ops.INDEX: self.assertEqual(u.src[1].dtype, dtypes.int)
|
||||
|
||||
valid = (10*r0<5-l0).ne(True)&(l0<3000)
|
||||
l2 = UOp(Ops.LOAD, dtypes.long, (d1.index(idx.valid(valid)),))
|
||||
uops = to_uops_list([l2])
|
||||
for u in uops:
|
||||
if u.op is Ops.INDEX: self.assertEqual(u.src[1].dtype, dtypes.int)
|
||||
self.assertNotEqual(u.dtype, dtypes.long)
|
||||
|
||||
def test_in_out_of_bounds_access(self):
|
||||
with Context(IGNORE_OOB=0):
|
||||
|
||||
@@ -12,7 +12,7 @@ from tinygrad.codegen.late.expander import migrate_indexing, expander, pm_pre_ex
|
||||
from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
|
||||
ReduceContext, correct_load_store, pm_render
|
||||
from tinygrad.codegen.opt.postrange import apply_opts
|
||||
from tinygrad.codegen.simplify import pm_simplify_ranges, pm_reduce_simplify, pm_flatten_range, pm_split_ranges
|
||||
from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse
|
||||
from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen
|
||||
from tinygrad.codegen.late.control_flow import CFGContext, pm_add_ends, pm_split_ends, pm_add_control_flow, linearize
|
||||
|
||||
@@ -26,6 +26,9 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
|
||||
# TODO: fix expander and remove this
|
||||
sink = graph_rewrite(sink, pm_add_buffers_local, name="add locals early")
|
||||
|
||||
# collapse loads reduce (indexing by a tensor)
|
||||
sink = graph_rewrite(sink, pm_load_collapse, name="load collapse")
|
||||
|
||||
# split ranges
|
||||
sink = graph_rewrite(sink, pm_split_ranges+pm_flatten_range, ctx={}, name="split ranges")
|
||||
|
||||
@@ -34,7 +37,6 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
|
||||
|
||||
# optimize (schedule) the AST
|
||||
sink = graph_rewrite(sink, pm_simplify_ranges, name="simplify ranges")
|
||||
sink = graph_rewrite(sink, pm_reduce_simplify, name="simplify reduces")
|
||||
|
||||
# do postrange optimization, BEAM or hand_coded_optimizations
|
||||
sink = apply_opts(sink, ren)
|
||||
|
||||
@@ -50,7 +50,6 @@ def delete_redundant_gates(store:UOp, buf:UOp, idx:UOp, val:UOp, store_gate:UOp,
|
||||
# remove the gate from the index
|
||||
return UOp.store(buf.index(idx).cast(cast.dtype) if cast is not None else buf.index(idx), val, *store.src[2:])
|
||||
|
||||
def no_load(u:UOp) -> bool: return not any(x.op is Ops.LOAD for x in u.backward_slice_with_self)
|
||||
load_store_indexing = 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)),
|
||||
@@ -61,8 +60,6 @@ load_store_indexing = PatternMatcher([
|
||||
# delete_redundant_gates (after expand)
|
||||
(UPat(Ops.STORE, src=(UPat.any(stidx:=UPat.var("buf").index(UPat.var("idx"), UPat.var("store_gate")), stidx.cast().named("cast")),
|
||||
UPat.var("val")), name="store", allow_any_len=True), delete_redundant_gates),
|
||||
# we want to make sure we dont do math on a loaded index since that can cause overflow, this undoes a pattern in reduce_collapse
|
||||
(UPat.var("c")<(UPat.var("x", dtypes.index)+UPat.var("y")), lambda x,y,c: (-x < -(c-y)) if no_load(y) and no_load(c) and not no_load(x) else None),
|
||||
])
|
||||
|
||||
# ***** load/store grouping *****
|
||||
|
||||
@@ -90,10 +90,7 @@ 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),
|
||||
# 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 y.vmin > 0 else None),
|
||||
# lift x+y out of reduce on ne
|
||||
((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")) < UPat.var("c"), lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and y.vmin > 0 else None),
|
||||
# fold the range
|
||||
((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(0, UPat.cvar("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),
|
||||
@@ -104,26 +101,38 @@ pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
|
||||
# REDUCE on ADD
|
||||
((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)),
|
||||
])+symbolic_flat
|
||||
|
||||
pm_reduce_load_collapse = PatternMatcher([
|
||||
# MUL casted bool
|
||||
((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast()), lambda x,gate: gate.where(x, 0)),
|
||||
# lift x+y out of reduce on ne
|
||||
((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),
|
||||
# reduce on gated load becomes can substitute the range and remove the reduce
|
||||
((UPat.var("idx")!=(UPat(Ops.RANGE, name="r").or_casted())).where(0, UPat.var("expr")).reduce(UPat.var("r"), arg=Ops.ADD),
|
||||
lambda r,idx,expr: (v:=(idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])).where(expr.substitute({r:idx.cast(r.dtype).valid(v)}),0)),
|
||||
])+symbolic_flat
|
||||
|
||||
def reduce_collapse(red:UOp):
|
||||
included, not_included = partition(red.backward_slice, lambda x: any(y in x.backward_slice_with_self for y in red.src[1:]))
|
||||
def reduce_collapse(red:UOp, pm=pm_reduce_collapse):
|
||||
included = red.src[0].toposort(gate=lambda x: any(y in x.ranges for y in red.src[1:]))
|
||||
if any(x.op in {Ops.STORE, Ops.REDUCE} for x in included): return None
|
||||
replaces: dict[UOp, UOp] = {}
|
||||
for u in included:
|
||||
for s in u.src:
|
||||
if s in not_included and s not in replaces and s.op not in {Ops.CONST, Ops.VCONST, Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}:
|
||||
replaces[s] = UOp(Ops.DEFINE_VAR, dtype=s.dtype, arg=(f'in{len(replaces)}', s.vmin, s.vmax))
|
||||
if s in included or s in replaces or s.op in {Ops.CONST, Ops.VCONST, Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}: continue
|
||||
replaces[s] = UOp(Ops.DEFINE_VAR, dtype=s.dtype, arg=(f'in{len(replaces)}', s.vmin, s.vmax))
|
||||
collapse_fxn = red.substitute(replaces)
|
||||
sink = graph_rewrite(collapse_fxn, pm_reduce_collapse, name="reduce_collapse")
|
||||
sink = graph_rewrite(collapse_fxn, pm, name="reduce_collapse")
|
||||
return sink.substitute({v:k for k,v in replaces.items()}) if no_range(sink) else None
|
||||
|
||||
pm_reduce_simplify = pm_reduce_unparented + PatternMatcher([
|
||||
# remove REDUCE without loads (generic arange opt / indexing). TODO: support multi range
|
||||
(UPat(Ops.REDUCE, src=(UPat(), UPat()), name="red"), reduce_collapse),
|
||||
def reduce_load_collapse(red:UOp): return reduce_collapse(red, pm=pm_reduce_load_collapse)
|
||||
|
||||
# remove REDUCE without loads (generic arange opt / indexing). TODO: support multi range
|
||||
pm_reduce_simplify = pm_reduce_unparented + PatternMatcher([(UPat(Ops.REDUCE, src=(UPat(), UPat()), 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.LOAD for x in u.backward_slice_with_self)
|
||||
pm_load_collapse = PatternMatcher([
|
||||
(UPat(Ops.REDUCE, src=(UPat(), 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.index)+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),
|
||||
])
|
||||
|
||||
@@ -6,7 +6,7 @@ from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, si
|
||||
from tinygrad.uop.symbolic import symbolic_flat
|
||||
from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, Metadata, DEBUG_RANGEIFY
|
||||
from tinygrad.helpers import PCONTIG, partition
|
||||
from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_unparented
|
||||
from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
|
||||
from tinygrad.codegen.opt import Opt
|
||||
from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, ALWAYS_CONTIGUOUS, IndexingContext, apply_movement_op
|
||||
|
||||
@@ -218,6 +218,9 @@ pm_const_buffer_folding = pm_mops+PatternMatcher([
|
||||
and (resolve(prod(x.dtype.shape)!=prod(x.shape)) or x.shape[-1]%4!=0) else None),
|
||||
# remove noop buffers. if we look at the next index we can remove even more of these
|
||||
(UPat(Ops.INDEX, name="idx").f(Ops.BUFFERIZE, allow_any_len=True, name="b2"), remove_noop_bufferize),
|
||||
# dont bufferize an arange
|
||||
(UPat.any((r:=UPat(dtype=dtypes.index).cast()).named("src"), r.eq(UPat()).named("src")).f(Ops.BUFFERIZE,
|
||||
allow_any_len=True, name="buf").f(Ops.INDEX, allow_any_len=True, name="idx"), remove_bufferize),
|
||||
# no buffers for const
|
||||
(UPat(Ops.CONST, name='c').f(Ops.BUFFERIZE, allow_any_len=True, name="b"), lambda c,b: b.const_like(c.arg).rtag(b.tag)),
|
||||
# indexing a const is a const
|
||||
@@ -516,7 +519,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
|
||||
# convert movement ops to ranges
|
||||
tsink, rctx = run_rangeify(tsink, DEBUG_RANGEIFY)
|
||||
|
||||
tsink = graph_rewrite(tsink, symbolic_flat+pm_reduce_unparented+pm_const_buffer_folding, name="symbolic") # this supports const folding
|
||||
tsink = graph_rewrite(tsink, symbolic_flat+pm_reduce_simplify+pm_const_buffer_folding, name="symbolic+reduce_collapse") # this does const folding
|
||||
tsink = graph_rewrite(tsink, pm_remove_bufferize, bottom_up=True, name="remove bufferize with cost function")
|
||||
tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user