diff --git a/test/test_schedule.py b/test/test_schedule.py index 7238bea5d1..b419589083 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -1503,6 +1503,18 @@ class TestSchedule(unittest.TestCase): run_schedule(sched) np.testing.assert_allclose(dx.numpy(), [[[[0.,3.,9.],[0,1.,3.],[0.,0.,0.]]]*3]*3) + def test_fuse_arange_avg_pool2d_ceil_mode(self): + x = Tensor.avg_pool2d(Tensor.empty(1,1,6,6), kernel_size=(3,3), padding=1, stride=3, ceil_mode=True) + sched = check_schedule(x, 1) + self.assertEqual(len([x for x in sched[0].ast.backward_slice_with_self if x.op is Ops.REDUCE]), 1) + + def test_fuse_arange_pad_circular_mode_bw(self): + x = Tensor.empty(1,1,5,5,5) + out = x.pad((1,2,3,5,1,2), mode="circular") + g = out.sum().gradient(x)[0] + sched = check_schedule(g, 1) + self.assertEqual(len([x for x in sched[0].ast.backward_slice_with_self if x.op is Ops.REDUCE]), 0) + # TODO like openpilot with imagef @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") def test_base_change_expand_expand(self): diff --git a/tinygrad/codegen/simplify.py b/tinygrad/codegen/simplify.py index 13b67606d1..a625dc2cf2 100644 --- a/tinygrad/codegen/simplify.py +++ b/tinygrad/codegen/simplify.py @@ -91,47 +91,59 @@ pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([ # 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), # 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), - (((UPat.var("r")= 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, 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 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, name="reduce_collapse") - return sink.substitute({v:k for k,v in replaces.items()}) if no_range(sink) else None +def reduce_collapse(red:UOp, u:UOp, pm=pm_reduce_collapse): + for r in red.src[1:]: + included = u.toposort(gate=lambda x: r in x.ranges) + 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 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 = u.substitute(replaces).reduce(r, arg=Ops.ADD) + sink = graph_rewrite(collapse_fxn, pm, name="reduce_collapse") + if not no_range(sink): return None + u = sink.substitute({v:k for k,v in replaces.items()}) + return u -def reduce_load_collapse(red:UOp): return reduce_collapse(red, pm=pm_reduce_load_collapse) +def reduce_load_collapse(red:UOp, u:UOp): return reduce_collapse(red, u, 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 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), +]) # 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, src=(UPat(), UPat()), name="red"), reduce_load_collapse), + (UPat(Ops.REDUCE, 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.index)+UPat.var("y"))