From c1db62ff7c30ac6437c5bccdb6d9ff51746f3bee Mon Sep 17 00:00:00 2001 From: Sieds Lykles <93992551+S-Lykles@users.noreply.github.com> Date: Thu, 23 Oct 2025 15:44:17 +0200 Subject: [PATCH] move reduce collapse to rangeify (#12845) --- test/test_const_folding.py | 2 +- test/test_uop_graph.py | 34 +++++++++++++-------------- tinygrad/codegen/__init__.py | 6 +++-- tinygrad/codegen/late/devectorizer.py | 3 --- tinygrad/codegen/simplify.py | 33 ++++++++++++++++---------- tinygrad/schedule/rangeify.py | 7 ++++-- 6 files changed, 48 insertions(+), 37 deletions(-) diff --git a/test/test_const_folding.py b/test/test_const_folding.py index f0dd3054cf..184bbf274a 100644 --- a/test/test_const_folding.py +++ b/test/test_const_folding.py @@ -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. diff --git a/test/test_uop_graph.py b/test/test_uop_graph.py index 5c2135b621..f26dcf9705 100644 --- a/test/test_uop_graph.py +++ b/test/test_uop_graph.py @@ -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): diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 7758adac74..1c1f8267f9 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -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) diff --git a/tinygrad/codegen/late/devectorizer.py b/tinygrad/codegen/late/devectorizer.py index 95831a5532..f5f76a28c7 100644 --- a/tinygrad/codegen/late/devectorizer.py +++ b/tinygrad/codegen/late/devectorizer.py @@ -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 ***** diff --git a/tinygrad/codegen/simplify.py b/tinygrad/codegen/simplify.py index 54b5b26b5a..a61762c2dd 100644 --- a/tinygrad/codegen/simplify.py +++ b/tinygrad/codegen/simplify.py @@ -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")) 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")