diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index c1de2a02c3..8ea786f2b2 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -239,6 +239,8 @@ pm_debuf = PatternMatcher([ (UPat(Ops.BUFFER, name="b"), debuf), # HACK: consts shouldn't have srcs by here (UPat(Ops.CONST, name="x"), lambda x: x.replace(src=()) if len(x.src) else None), + # no movement ops + (UPat(GroupOp.Movement, name="x"), lambda x: x.src[0]), ]) @track_rewrites(name=lambda sink,ret: f"Schedule {pluralize('Kernel',len([u for u in ret[sink].toposort() if u.op is Ops.KERNEL]))}", replay=True) @@ -253,9 +255,9 @@ def get_kernelize_map(sink:UOp) -> dict[UOp, UOp]: rsink = tensor_map[sink] from tinygrad.codegen.devectorizer import pm_reduce, ReduceContext - rsink = graph_rewrite(rsink, pm_reduce, ctx=ReduceContext(), name="remove reduce") rsink = graph_rewrite(rsink, pm_add_buffers, name="buffer", bottom_up=True) rsink = graph_rewrite(rsink, pm_debuf, ctx=[0], name="debuf", bottom_up=True) + rsink = graph_rewrite(rsink, pm_reduce, ctx=ReduceContext(), name="remove reduce") from tinygrad.codegen import rewrites_for_linearizer, apply_rewrites rsink = apply_rewrites(rsink, rewrites_for_linearizer) from tinygrad.renderer.cstyle import CStyleLanguage diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 37322ae6c1..86dc92f8cb 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -188,10 +188,15 @@ class UOp(MathTrait, metaclass=UOpMetaClass): @functools.cached_property def ranges(self) -> dict[UOp, None]: if self.op is Ops.RANGE: return {self:None} - if self.op in {Ops.BUFFERIZE, Ops.REDUCE, Ops.STORE}: + if self.op in {Ops.BUFFERIZE, Ops.REDUCE}: ret = self.src[0].ranges.copy() for s in self.src[1:]: if s in ret: del ret[s] + elif self.op in {Ops.STORE}: + ret = self.src[0].ranges.copy() + ret.update(self.src[1].ranges) + for s in self.src[2:]: + if s in ret: del ret[s] else: ret = {} for s in self.src: ret.update(s.ranges)