no const end

tinygrad/schedule/rangeify2.py

@@ -5,6 +5,8 @@ from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, GroupOp, KernelInfo
 from tinygrad.uop.ops import graph_rewrite, AxisType, rewrite_group, remove_all_tags, resolve
 from tinygrad.helpers import all_int, VIZ, SPEC, Context, panic
 from tinygrad.schedule.indexing import BufferizeOpts, apply_movement_op
+from tinygrad.uop.symbolic import symbolic
+from tinygrad.codegen.simplify import pm_reduce_simplify
 
 # *** preparation ***
 
@@ -227,6 +229,8 @@ def get_kernel_graph(sink:UOp) -> UOp:
 
   # ***** MERGING AND SPLITTING (should be totally optional) *****
 
+  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify, name="reduce simplify")
+
   while 1:
     staged: dict[UOp, list[UOp]] = {}
     for u in tsink.toposort():
@@ -246,7 +250,7 @@ def get_kernel_graph(sink:UOp) -> UOp:
   if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Merged Rangeify")
 
   next_buffer_num = itertools.count(1000)
-  tsink = graph_rewrite(tsink, pm_remove_stage, ctx=next_buffer_num, bottom_up=True, name="remove stage")
+  tsink = graph_rewrite(tsink, symbolic+pm_remove_stage, ctx=next_buffer_num, bottom_up=True, name="remove stage")
   tsink = graph_rewrite(tsink, split_kernels, bottom_up=True, name="split kernels")
   tsink = graph_rewrite(tsink, pm_no_indexing_calls, name="remove indexing from call args")
 

tinygrad/uop/symbolic.py

@@ -312,6 +312,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
                         else y.src for y in x.src[1:]]))))),
   # after/end with 1 src is just src[0]
   (UPat((Ops.AFTER, Ops.END), src=(UPat.var("s"),)), lambda s: s),
+  # ranges can be subbed for CONSTs, remove them from ENDs
+  (UPat(Ops.END, name="x"), lambda x: x.replace(src=(x.src[0],)+tuple([r for r in x.src[1:] if r.op is not Ops.CONST]))),
   # the rules above key on bare CONSTs, so a redundantly committed const has to be uncast in the same fixpoint
 ])+div_and_mod_symbolic+pm_uncast_const