diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index cd14be81d9..46adf80ad2 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -532,7 +532,7 @@ jobs: - name: Test multitensor run: | CPU=1 RANGEIFY=1 python3 test/test_multitensor.py TestMultiTensor.test_matmul_shard_1_1 TestMultiTensor.test_simple_add_W TestMultiTensor.test_simple_reduce \ - TestMultiTensor.test_elementwise_dtype TestMultiTensor.test_shard_no_recompile + TestMultiTensor.test_elementwise_dtype TestMultiTensor.test_shard_no_recompile TestHandleData.test_copied_to_device TestMultiRamUsage CPU=1 RANGEIFY=1 python3 -m pytest test/test_multitensor.py::TestMultiAssign -k 'not (multi_assign_piece_noncontig or multi_assign_var_offset)' - name: Test CPU=1 RANGEIFY=2 run: CPU=1 CPU_LLVM=0 RANGEIFY=2 python3 -m pytest -n auto test/test_tiny.py test/test_rangeify.py test/test_ops.py --durations 20 diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 8ef9bb53c8..340e11cab1 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -455,7 +455,7 @@ pm_add_buffers = pm_mops+PatternMatcher([ # move RESHAPEs through MSELECT/MSTACK (UPat((Ops.MSELECT, Ops.MSTACK), src=UPat(Ops.RESHAPE), name="m"), - lambda m: m.replace(src=tuple([x.src[0] for x in m.src])).reshape(m.src[0].arg)), + lambda m: m.replace(src=tuple([x.src[0] for x in m.src]), tag=None).reshape(m.src[0].arg).rtag(m.tag)), ]) # ***************** @@ -581,8 +581,9 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]: tsink = graph_rewrite(tsink, pm_cleanups, bottom_up=True, name="remove costly buffers") # rebuild the sink with all the BUFFERIZEs with tags, this is what's ending up in the tensor graph + # MSTACK stacks multiple BUFFERIZEs in one tagged tensor # if it's not tagged by here, it's out - tsink = UOp.sink(*[x for x in tsink.parents if x.base.op in {Ops.BUFFERIZE, Ops.CONST} and x.tag is not None]) + tsink = UOp.sink(*[x for x in tsink.parents if x.base.op in {Ops.BUFFERIZE, Ops.MSTACK, Ops.CONST} and x.tag is not None]) if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Tagged Rangeify")