diff --git a/test/test_rangeify.py b/test/test_rangeify.py index 1d4f48abe5..9379390ae5 100644 --- a/test/test_rangeify.py +++ b/test/test_rangeify.py @@ -11,6 +11,11 @@ class TestRangeify(unittest.TestCase): ba = A.expand(N, N) ((ba+1).sum(axis=1) + (ba+2).sum(axis=0)).realize() + def test_partial_contig(self): + A = Tensor.empty(64, 64, 64) + ret = A.sum(axis=2).contiguous(arg=(1,)).sum(axis=1) + ret.realize() + def test_double_gemm(self): A = Tensor.empty(N, N) B = Tensor.empty(N, N) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index fe1949ddfa..b0c31c25a4 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -72,14 +72,14 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC # ** expander (expand_rewrite) ** ret.append(RewriteStep(sym+migrate_indexing, name="initial symbolic")) - # expand - ret.append(RewriteStep(sym+expander, name="expander")) + # add locals + ret.append(RewriteStep(pm_flatten_range+pm_add_buffers_local+rangeify_codegen, name="add local buffers")) # add gpu dims (late). this also handles UNROLL range ret.append(RewriteStep(pm_add_gpudims, lambda _: opts, name="add gpudims")) - # add locals - ret.append(RewriteStep(pm_flatten_range+pm_add_buffers_local+rangeify_codegen, name="add local buffers")) + # expand + ret.append(RewriteStep(sym+expander, name="expander")) # ** devectorizer (full_graph_rewrite) ** # remove reduce diff --git a/tinygrad/codegen/late/expander.py b/tinygrad/codegen/late/expander.py index bd88548d87..50dafcbdbf 100644 --- a/tinygrad/codegen/late/expander.py +++ b/tinygrad/codegen/late/expander.py @@ -86,7 +86,7 @@ expander = PatternMatcher([ (UPat(Ops.UNROLL, name="outer", src=(UPat(Ops.UNROLL, name="inner"),)), lambda outer, inner: UOp(Ops.UNROLL, outer.dtype, (inner.src[0],), inner.arg+outer.arg)), # do expansion - (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.GEP, Ops.WMMA, Ops.LOAD, Ops.STORE, Ops.INDEX, Ops.BUFFERIZE, + (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.GEP, Ops.WMMA, Ops.LOAD, Ops.STORE, Ops.INDEX, #Ops.BUFFERIZE, Ops.VECTORIZE, Ops.IF, Ops.REDUCE), name="root", custom_early_reject=set([Ops.UNROLL])), do_expand), (UPat(Ops.CONTRACT, name="con"), do_contract), # BARRIERs aren't actually expanded diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 041011d9c3..8ed9e107c8 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -207,8 +207,10 @@ class UOp(MathTrait, metaclass=UOpMetaClass): ret: dict[UOp, None] = {} if self.op in range_start.keys(): for s in self.src[:range_start[self.op]]: ret.update(s.ranges) - for s in self.src[range_start[self.op]:]: - if s in ret: del ret[s] + delete_ranges = self.src[range_start[self.op]:] + if len(delete_ranges): + for s in UOp.sink(*delete_ranges).ranges: + if s in ret: del ret[s] elif self.op in {Ops.BARRIER}: ret = {x:None for x in self.src[0].ranges if x.arg[1] != AxisType.LOCAL} else: