diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 7e335c2712..af58f821b2 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -536,7 +536,7 @@ jobs: test/test_outerworld_range.py test/test_randomness.py test/test_nn.py test/test_arange.py test/test_tensor.py test/test_optim.py \ test/test_setitem.py test/test_assign.py - name: Test const folding - run: CPU=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_const_folding.py -k "not test_cast_padded and not TestReduceOpsConstFolding and not TestMultiConstFolding" + run: CPU=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_const_folding.py -k "not test_cast_padded and not TestMultiConstFolding" - 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 \ diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 8df51d966e..0b5061f4e8 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -5,6 +5,8 @@ from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, RewriteNotReady, _substitute, ssimplify from tinygrad.uop.symbolic import sym, symbolic_simple from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY, Context, flatten, dedup +from tinygrad.codegen.simplify import pm_reduce_simplify + from tinygrad.schedule.kernelize import Kernel from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType @@ -606,7 +608,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]: # rangeify tsink = graph_rewrite(tsink, pm_rangeify, ctx=RangeifyContext(), bottom_up=True, name="rangeify") # NOTE: sym (vs symbolic_simple) breaks things here because ranges with len 1 aren't handled right - tsink = graph_rewrite(tsink, symbolic_simple, name="symbolic") # this supports const folding + tsink = graph_rewrite(tsink, symbolic_simple+pm_reduce_simplify, name="symbolic") # this supports const folding 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