nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

reduce const folding (#12344)

Browse Source
This commit is contained in:
chenyu
2025-09-29 23:08:56 -04:00
committed by GitHub
parent 3a480b858f
commit 8e508a9927
2 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/test.yml
+1 -1
View File
@@ -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 \
tinygrad/schedule/rangeify.py
+3 -1
View File
@@ -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