simpler postrange works

tinygrad/codegen/gpudims.py

@@ -91,7 +91,8 @@ def add_gpudims(ctx:Renderer, s:UOp):
 def fix_reduce_unroll(x:UOp):
   reduce_range, reduce_expand = partition(x.src[1:], lambda y: y.op is Ops.RANGE)
   if len(reduce_expand) == 0: return None
-  assert all(x.op is Ops.UNROLL for x in reduce_expand), f"not all UNROLLS in {reduce_expand} for {x.axis_arg}"
+  reduce_expand = [x for x in reduce_expand if x.op is not Ops.CONST]
+  assert all(x.op is Ops.UNROLL for x in reduce_expand), f"not all UNROLLS in {reduce_expand}"
   ret = x.src[0]
   if len(contract_axis:=flatten(x.arg for x in reduce_expand)):
     ret = UOp(Ops.CONTRACT, x.dtype.vec(prod(x[1] for x in contract_axis)), (ret,), tuple(contract_axis), tag=1)

tinygrad/codegen/opt/__init__.py

@@ -39,13 +39,22 @@ def apply_opt(ast:UOp, renderer:Renderer, cls:type[Kernel]):
   k = cls(ast, opts=renderer)
   k.apply_opts(ast.arg.opts_to_apply)
   ret = k.get_optimized_ast()
-  if __debug__: type_verify(list(ret.toposort()))
+  if __debug__ and cls == Kernel: type_verify(list(ret.toposort()))
   return ret
 
 pm_do_optimize = PatternMatcher([
   (UPat(Ops.SINK, name="ast"), lambda ctx,ast: apply_opt(ast, ctx, Kernel) if ast.arg is not None and ast.arg.opts_to_apply is not None else None),
 ])
 
+def flatten_range(r:UOp):
+  off = 2 if r.op is Ops.STORE else 1
+  rngs = r.src[off:]
+  if not len(rngs): return None
+  new_rngs = [x for x in UOp.sink(*rngs).toposort() if x.op is Ops.RANGE]
+  return r.replace(src=r.src[:off]+tuple(new_rngs))
+
 pm_postrange_opt = PatternMatcher([
   (UPat(Ops.SINK, name="ast"), lambda ctx,ast: apply_opt(ast, ctx, RKernel) if ast.arg is not None and ast.arg.opts_to_apply is not None else None),
+  # real ranges only
+  (UPat((Ops.REDUCE, Ops.STORE), name="r"), flatten_range),
 ])

tinygrad/codegen/opt/kernel.py

@@ -86,7 +86,7 @@ class Kernel:
 
     # group simplifies
     self.simplify_ones()
-    self.simplify_merge_adjacent()
+    #self.simplify_merge_adjacent()
 
     # axis types
     global_loops = AxisType.GLOBAL if self.opts.has_local else AxisType.LOOP
@@ -122,7 +122,7 @@ class Kernel:
   @property
   def output_shape(self) -> tuple[sint, ...]: return self.sts[0].shape
   @property
-  def shape_len(self) -> int: return len(self.sts[0].shape)
+  def shape_len(self) -> int: return len(self.output_shape)
 
   def axes_of(self, *axis_type:AxisType) -> list[int]: return [i for i,t in enumerate(self.axis_types) if t in argfix(axis_type)]
   @property

tinygrad/codegen/opt/postrange.py

@@ -1,7 +1,52 @@
-from tinygrad.uop.ops import UOp, Ops
+from dataclasses import replace
+from tinygrad.uop.ops import UOp, Ops, sint, ssimplify, AxisType
 from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.renderer import Renderer
+from tinygrad.dtype import dtypes
 
 class RKernel(Kernel):
   def __init__(self, ast:UOp, opts:Renderer|None=None):
+    self.rng = sorted([u for u in ast.toposort() if u.op is Ops.RANGE and u.vmax > 0], key=lambda x: x.arg)
     super().__init__(ast, opts)
+    self.sts.clear()
+
+    # convert LOOP to GLOBAL
+    self.replaces = {}
+    if self.opts.has_local:
+      rng = [x.replace(arg=(x.arg[0], AxisType.GLOBAL)) if x.arg[1] == AxisType.LOOP else x for x in self.rng]
+      self.replaces.update(dict(zip(self.rng, rng)))
+      self.rng = rng
+
+  def shift_to(self, axis:int, amount:int, new_type:AxisType, top:bool=False, insert_at:int|None=None):
+    old_sz = self.rng[axis].src[0].arg // amount
+    assert old_sz > 0, f"bad old_sz on {axis} {amount} {self.rng[axis]}"
+
+    maxarg = max([x.arg[0] for x in self.rng])
+    new_rng = UOp.range(dtypes.int, amount, maxarg+1, new_type)
+
+    if old_sz == 1:
+      self.replaces[self.rng[axis]] = new_rng
+      self.rng.insert(insert_at if insert_at is not None else len(self.rng), new_rng)
+      del self.rng[axis]
+    else:
+      replaced_rng = self.rng[axis].replace(src=(UOp.const(dtypes.int, old_sz),))
+      self.replaces[self.rng[axis]] = replaced_rng * amount + new_rng
+      self.rng[axis] = replaced_rng
+      self.rng.insert(insert_at if insert_at is not None else len(self.rng), new_rng)
+
+  @property
+  def axis_types(self) -> list[AxisType]: return [x.arg[1] for x in self.rng]
+  @property
+  def shape_len(self): return len(self.rng)
+
+  @property
+  def full_shape(self) -> tuple[sint, ...]: return tuple([ssimplify(x.src[0]) for x in self.rng])
+  @property
+  def output_shape(self) -> tuple[sint, ...]: return tuple([ssimplify(x.src[0]) for x in self.ast.src[0].src[2:]])
+
+  def get_optimized_ast(self, name_override:str|None=None) -> UOp:
+    return self.ast.substitute(self.replaces).replace(arg=replace(self.ast.arg, name=self.name, opts_to_apply=None))
+
+  # does nothing
+  @axis_types.setter
+  def axis_types(self, value): pass