better cost function

.github/workflows/test.yml

@@ -534,9 +534,6 @@ jobs:
           test/test_tiny.py test/test_rangeify.py test/test_ops.py test/test_symbolic_ops.py test/test_symbolic_jit.py test/test_tensor_variable.py \
           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 test/test_multitensor.py
-    - name: Test CPU=1 CPU_LLVM=1 RANGEIFY=1
-      run: |
-        CPU=1 CPU_LLVM=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_edgecases.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"
     # RANGEIFY=2 isn't supported
@@ -584,7 +581,7 @@ jobs:
         key: metal
         deps: testing
     - name: some unit tests
-      run: METAL=1 RANGEIFY=1 python -m pytest -n=auto test/unit/test_winograd.py test/unit/test_linalg.py --durations=20
+      run: METAL=1 RANGEIFY=1 python -m pytest -n=auto test/unit/test_winograd.py --durations=20
     - name: Test METAL=1 RANGEIFY=1
       run: METAL=1 RANGEIFY=1 python -m pytest -n=auto test/test_ops.py --durations=20
     - name: Run process replay tests

extra/remu/src/thread.rs

@@ -930,7 +930,7 @@ impl<'a> Thread<'a> {
 
             let op = ((instr >> 16) & 0x3ff) as u32;
             match op {
-                764 | 765 | 288 | 289 | 290 | 766 | 767 | 768 | 769 => {
+                764 | 765 | 288 | 289 | 290 | 766 | 768 | 769 => {
                     let vdst = (instr & 0xff) as usize;
                     let sdst = ((instr >> 8) & 0x7f) as usize;
                     let f = |i: u32| -> usize { ((instr >> i) & 0x1ff) as usize };
@@ -944,16 +944,6 @@ impl<'a> Thread<'a> {
                     assert_eq!(clmp, 0);
 
                     let vcc = match op {
-                        767 => {
-                            let (s0, s1, s2): (u32, u32, u64) = (self.val(s0), self.val(s1), self.val(s2));
-                            let (mul_result, overflow_mul) = (s0 as i64).overflowing_mul(s1 as i64);
-                            let (ret, overflow_add) = mul_result.overflowing_add(s2 as i64);
-                            let overflowed = overflow_mul || overflow_add;
-                            if self.exec.read() {
-                                self.vec_reg.write64(vdst, ret as u64);
-                            }
-                            overflowed
-                        },
                         766 => {
                             let (s0, s1, s2): (u32, u32, u64) = (self.val(s0), self.val(s1), self.val(s2));
                             let (mul_result, overflow_mul) = (s0 as u64).overflowing_mul(s1 as u64);

test/external/external_uop_gc.py

@@ -63,7 +63,6 @@ if __name__ == "__main__":
     views_to_valid_uop.cache_clear()
 
     new_uops = uops_allocated()
-    print_uops()
     gc.collect()
     new_uops_gc = uops_allocated()
     print(f"{t.__name__:30s}: {new_uops:3d} -> {new_uops_gc:3d}")

test/test_linearizer.py

@@ -123,7 +123,6 @@ class TestLinearizer(unittest.TestCase):
     assert num_loads <= 4, "more load uops than needed"
     assert num_loads >= 4, "unexpected number of uops, maybe this test needs updating?"
 
-  @unittest.skip("this is handled at higher level now")
   def test_upcast_cse(self):
     # when upcasting, within a subtree, there may be common expressions.
 

test/test_ops.py

@@ -3164,7 +3164,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(32,10)], lambda x: x.masked_fill((x>0.1).detach(), -math.inf))
     helper_test_op([(32,10)], lambda x: x.masked_fill((x<0.1).detach(), -math.inf))
 
-  @unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT == "AMD" and RANGEIFY, "very slow on MOCKGPU because reduce does not fold")
   def test_masked_select(self):
     helper_test_op([(32, 10)], lambda x: x.masked_select(x>0.5), lambda x: x.masked_select(x>0.5), forward_only=True)
     helper_test_op([(32, 10)], lambda x: x.masked_select(torch.tensor(True)), lambda x: x.masked_select(Tensor(True)), forward_only=True)

test/test_pickle.py

@@ -2,7 +2,7 @@ import unittest, pickle, types
 import numpy as np
 from tinygrad import Tensor, TinyJit, Variable, dtypes
 from tinygrad.helpers import GlobalCounters, ContextVar, Context
-from tinygrad.uop.ops import PatternMatcher, UPat, UOp
+from tinygrad.uop.ops import PatternMatcher, UPat, UOp, Ops
 
 class TestPickle(unittest.TestCase):
   def test_pickle_code_object(self):
@@ -45,9 +45,10 @@ class TestPickle(unittest.TestCase):
     t_values = t.numpy()
     del t # free buffers
     print("** post pickle")
+    init = GlobalCounters.kernel_count
     t2:Tensor = pickle.loads(st)
-    assert t2.uop.is_realized
     np.testing.assert_equal(t_values, t2.numpy())
+    self.assertEqual(GlobalCounters.kernel_count-init, 0)
 
   def test_pickle_realized_tensor_alt2(self):
     print("** init")
@@ -69,14 +70,14 @@ class TestPickle(unittest.TestCase):
   def test_pickle_buffer_uop(self):
     t = Tensor.arange(4).realize()
     a = t.uop
-    assert a.is_realized
-    self.assertIsNotNone(buffer:=a.base.realized)
+    assert a.op is Ops.BUFFER
+    self.assertIsNotNone(buffer:=a.realized)
     s = pickle.dumps(a)
     # free buffers
     del a
     del buffer
     a2:UOp = pickle.loads(s)
-    self.assertListEqual(a2.base.realized.as_buffer().cast("I").tolist(), [0, 1, 2, 3])
+    self.assertListEqual(a2.realized.as_buffer().cast("I").tolist(), [0, 1, 2, 3])
 
   def test_pickle_unrealized_tensor(self):
     t = Tensor.ones(10, 10)

test/test_rangeify.py

@@ -40,11 +40,6 @@ class TestRangeifyOpt(unittest.TestCase):
 
 @unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY")
 class TestRangeify(unittest.TestCase):
-  def test_groupnorm(self):
-    # ranges 1 and 3 are merging
-    x = nn.GroupNorm(32, 128)
-    x(Tensor.empty(1, 128, 64, 64)).realize()
-
   def test_expand_children(self):
     A = Tensor.empty(N, N).sum(axis=1)
     ba = A.expand(N, N)

test/test_softmax_fusion.py

@@ -44,7 +44,6 @@ class TestFuse(unittest.TestCase):
     a = Tensor.rand(50,50).realize()
     self._test_fuse(lambda a: a / a.mean(axis=1), a)
 
-  @unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
   def test_fuse_argmax(self):
     a = Tensor.rand(50,50).realize()
     self._test_fuse(lambda a: a.argmax(axis=-1), a)

test/test_uops.py

@@ -544,6 +544,86 @@ class TestUopsObject(unittest.TestCase):
     with Timing("create 10k uops:"): ret = [UOp(Ops.CONST, dtypes.int, arg=10000000+i) for i in range(10000)]
     assert len(ret) == 10000
 
+
+class TestShapeSpec(unittest.TestCase):
+  # ** CONST is CONST(VIEW(DEVICE)) -> RESHPAE -> EXPAND
+
+  def test_expanded_const(self):
+    a = Tensor(1).uop
+    self.assertEqual(a.st, ShapeTracker.from_shape(()))
+    a = Tensor.ones((4, 4)).uop
+    self.assertEqual(a.st, ShapeTracker.from_shape(()).reshape((1,1)).expand((4,4)))
+
+  # NOTE: CONST ShapeTracker comes from its source
+  def test_scalar_const(self):
+    a = Tensor(0).uop
+    self.assertEqual(a.st, ShapeTracker.from_shape(()))
+
+  def test_scalar_var(self):
+    vv = UOp.variable("a", 1, 4).bind(2)
+    t = Tensor(vv).uop
+    self.assertEqual(t.st, ShapeTracker.from_shape(()))
+
+  # ** ASSIGN is ASSIGN(VIEW(BUFFER), new_val)
+
+  def test_assign_flat(self):
+    buffer = Tensor.arange(4).realize()
+    a = buffer.assign(Tensor.zeros((4,), dtype=dtypes.int))
+    assign_pattern = UPat(Ops.ASSIGN, src=(UPat(Ops.BUFFER), UPat()))
+    assert assign_pattern.match(a.uop, {})
+    a.realize()
+    self.assertEqual(buffer.tolist(), [0, 0, 0, 0])
+
+  def test_assign_permuted(self):
+    buffer = Tensor.arange(4).reshape(2, 1, 2).contiguous().realize()
+    a = buffer.permute((1, 2, 0)).assign(Tensor.arange(4).reshape(1, 2, 2).contiguous())
+    a.realize()
+    self.assertEqual(buffer.tolist(), [[[0, 2]], [[1, 3]]])
+
+  def test_assign_reshaped(self):
+    buffer = Tensor.ones((4,)).contiguous().realize()
+    a = buffer.reshape((2, 2)).assign(Tensor.zeros((2, 2)))
+    assign_pattern = UPat(Ops.ASSIGN, src=(UPat(Ops.RESHAPE, src=(UPat(Ops.BUFFER))), UPat()))
+    assert assign_pattern.match(a.uop, {})
+    a.realize()
+    self.assertEqual(buffer.tolist(), [0, 0, 0, 0])
+
+  # setitem is a partial assign
+  def test_setitem(self):
+    a = Tensor.ones((4,)).contiguous().realize()
+    assign = a.shrink(((1, 2),)).assign(Tensor.zeros((1,)))
+    # the ASSIGN UOp has size=1
+    self.assertEqual(assign.uop.size, 1)
+    # the ASSIGN views the buffer with a shrunk st
+    self.assertEqual(assign.uop.src[0].st, ShapeTracker.from_shape((4,)).shrink(((1, 2),)))
+    # the underlying BUFFER has a size=4
+    self.assertEqual(assign.uop.buf_uop.size, 4)
+    # NOTE: output shape is different from the BUFFER shape
+    self.assertNotEqual(assign.uop.shape, a.uop.shape)
+    assign.realize()
+    self.assertEqual(a.tolist(), [1, 0, 1, 1])
+
+  def test_buffer_st(self):
+    a = UOp.new_buffer(Device.DEFAULT, 10, dtypes.float)
+    self.assertEqual(a.st, ShapeTracker.from_shape((10,)))
+
+  def test_ops_st(self):
+    # view / mop
+    a = Tensor.empty(4, 2, 1).permute((1, 2, 0)).uop
+    self.assertEqual(a.st, ShapeTracker.from_shape((4, 2, 1)).permute((1, 2, 0)))
+    # alu / reduce
+    alu = a*2
+    self.assertEqual(alu.st, ShapeTracker.from_shape((2, 1, 4)))
+    r = Tensor.empty(4, 4).sum(axis=1)
+    self.assertEqual(r.uop.st, ShapeTracker.from_shape((4,)))
+
+  def test_st_wmma_none(self):
+    A = UOp(Ops.DEFINE_VAR, dtypes.float.vec(16), arg=('a', UOp.const(dtypes.float, 0), UOp.const(dtypes.float, 1)))
+    B = UOp(Ops.DEFINE_VAR, dtypes.float.vec(16), arg=('b', UOp.const(dtypes.float, 0), UOp.const(dtypes.float, 2)))
+    C = UOp(Ops.DEFINE_VAR, dtypes.float.vec(16), arg=('c', UOp.const(dtypes.float, 0), UOp.const(dtypes.float, 3)))
+    wmma = UOp(Ops.WMMA, dtypes.float.vec(16), (A, B, C))
+    assert wmma.st is None
+
 class TestUOpChildren(unittest.TestCase):
   def test_children_exist(self):
     a = UOp.variable("weird_name_234", 0, 10)

test/test_uops_stats.py

@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor
-from tinygrad.helpers import getenv, GlobalCounters, EMULATE, RANGEIFY
+from tinygrad.helpers import getenv, GlobalCounters, EMULATE
 from tinygrad.engine.realize import lower_schedule_item, ProgramSpec, get_program
 from tinygrad.renderer import Estimates
 from tinygrad.codegen import full_rewrite
@@ -51,11 +51,7 @@ class TestMemoryCount(unittest.TestCase):
     a = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024)
     b = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024)
     _, mem = get_stats(a+b)
-    if RANGEIFY:
-      # rangeify is smart!
-      self.assertEqual(mem, 1024 + 2*1024)  # 2 lil reads + 1 lil write
-    else:
-      self.assertEqual(mem, 1024*1024 + 2*1024)  # 2 lil reads + 1 write
+    self.assertEqual(mem, 1024*1024 + 2*1024)  # 2 lil reads + 1 write
 
   def test_self_add(self):
     a = Tensor.empty(1024, 1024, dtype=dtypes.uint8)

test/unit/test_attention.py

@@ -1,6 +1,5 @@
 import unittest
 from tinygrad import Tensor, dtypes, TinyJit, UOp
-from tinygrad.helpers import RANGEIFY
 from tinygrad.apps.llm import apply_rope
 
 # TODO: test_scheduler, but just in uint
@@ -13,7 +12,7 @@ class TestAttention(unittest.TestCase):
     attn = q.scaled_dot_product_attention(k, v)
     sched = attn.schedule()
     # attention has 5 kernels now
-    self.assertEqual(len(sched), 4 if RANGEIFY else 5)
+    self.assertEqual(len(sched), 5)
     softmax_inputs = sched[1:4]
     for si in softmax_inputs:
       assert all(b.dtype == dtypes.half for b in si.bufs), f"non half {si.bufs=}"
@@ -43,4 +42,4 @@ class TestAttention(unittest.TestCase):
     self.assertEqual(prune_size, 1)
 
 if __name__ == '__main__':
-  unittest.main()
+  unittest.main()

test/unit/test_kernelize.py

@@ -1,7 +1,6 @@
 import unittest
 from tinygrad import Tensor
 from tinygrad.uop import Ops
-from tinygrad.helpers import RANGEIFY
 
 class TestKernelize(unittest.TestCase):
   def test_add_reshaped(self):
@@ -18,11 +17,7 @@ class TestKernelize(unittest.TestCase):
     a1 = a.sum(axis=1)
     a0 = a1.sum(axis=0)
     a0.kernelize()
-    self.assertEqual(len([s for s in a0.uop.toposort() if s.op is Ops.KERNEL]), 2 if RANGEIFY else 3)
-    self.assertIs(a1.uop.base.op, Ops.REDUCE_AXIS if RANGEIFY else Ops.ASSIGN)
-    # input Tensor and user contiguous kernelize
-    self.assertIs(a0.uop.base.op, Ops.ASSIGN)
-    self.assertIs(a.uop.base.op, Ops.ASSIGN)
+    self.assertIs(a1.uop.base.op, Ops.ASSIGN)
 
   def test_two_reduce_w_add(self):
     a = Tensor.ones(16,16).contiguous()

test/unit/test_tensor_uop_representation.py

@@ -32,7 +32,8 @@ class TestTensorMutates(unittest.TestCase):
     d.realize()
     is_pattern_uop(d.uop.base, realized_pattern)
     is_pattern_uop(c.uop.base, realized_pattern)
-    is_pattern_uop(c.uop.base, realized_pattern)
+    # NOTE: we keep movement ops on top of the buffer view
+    is_pattern_uop(c.uop, UPat(Ops.BUFFER))
     assert d.uop is not d.uop.base
 
   def test_reshape_is_same_child(self):
@@ -55,8 +56,7 @@ class TestTensorUopRepresentation(unittest.TestCase):
     b = Tensor([4.,5,6]).realize()
     c = a+b
     print(c.uop)
-    is_pattern(c, UPat(Ops.ADD))
-    for s in c.uop.src: is_pattern_uop(s.base, realized_pattern)
+    is_pattern(c, UPat(Ops.ADD, src=(realized_pattern, realized_pattern)))
 
   def test_empty_buf(self):
     a = Tensor.empty(3, 3)

test/unit/test_uop_spec.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 import unittest
 
 from tinygrad import Tensor
-from tinygrad.helpers import DEBUG, RANGEIFY
+from tinygrad.helpers import DEBUG
 from tinygrad.uop.ops import UOp, Ops, print_uops
 from tinygrad.uop.spec import type_verify, ast_spec, tensor_uop_spec
 from tinygrad.shape.shapetracker import ShapeTracker
@@ -75,7 +75,6 @@ class TestUOpSpec(unittest.TestCase):
     st = UOp.store(buf.view(ShapeTracker.from_shape(())), a.cast(dtypes.float))
     helper_test_verify_ast(st)
 
-  @unittest.skipIf(RANGEIFY, "RANGEIFY does not push views")
   def test_assert_masked_view_in_const(self):
     t = Tensor(6).uop
     a = t.replace(src=(t.src[0].replace(arg=t.st.reshape((1,)).pad(((0, 1),))),))

tinygrad/codegen/opt/postrange.py

@@ -71,13 +71,6 @@ class Scheduler:
                         or (x.op is Ops.BUFFERIZE and x.arg == AddrSpace.LOCAL)]
     for ls in local_store_rngs: store_rngs = tuple([x for x in store_rngs if x in ls])
 
-    # filter any not in reduces
-    # TODO: enable this
-    """
-    reduce_rngs = [x.ranges for x in self.ast.toposort() if x.op is Ops.REDUCE]
-    for ls in reduce_rngs: store_rngs = tuple([x for x in store_rngs if x in ls])
-    """
-
     return [x for x in UOp.sink(*store_rngs).toposort() if x.op is Ops.RANGE and x.arg[1] == AxisType.LOOP] if store_rngs else []
 
   def convert_loop_to_global(self):

tinygrad/codegen/simplify.py

@@ -17,24 +17,20 @@ pm_flatten_range = PatternMatcher([
 
 def count_divmod(x:UOp): return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}])
 def simplify_merge_adjacent(u:UOp) -> UOp|None:
-  reduce_ranges = [x.ranges for x in u.sparents if x.op is Ops.REDUCE]
   i = range_start[u.op]
   while i < len(u.src)-1:
     r0, r1 = u.src[i], u.src[i+1]
     # check same type
     if r0.arg[-1] == r1.arg[-1]:
-      # check if the ranges to merge are in the same reduces
-      if all((r0 in rngs) == (r1 in rngs) for rngs in reduce_ranges):
-        s0, s1 = r0.src[0], r1.src[0]
-        # do the merge
-        new_range = r0.replace(src=(s0*s1,))
-        nidx = graph_rewrite(u, _substitute+symbolic_flat+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
-                             name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")
-
-        # check if it simplifies
-        if count_divmod(nidx) <= count_divmod(u):
-          u = nidx
-          continue
+      s0, s1 = r0.src[0], r1.src[0]
+      # do the merge
+      new_range = r0.replace(src=(s0*s1,))
+      nidx = graph_rewrite(u, _substitute+symbolic_flat+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
+                           name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")
+      # check if it simplifies
+      if count_divmod(nidx) <= count_divmod(u):
+        u = nidx
+        continue
     i += 1
   return u
 

tinygrad/renderer/llvmir.py

@@ -107,8 +107,7 @@ base_rewrite = PatternMatcher([
    f"  {ctx[x]} = phi {ldt(x.dtype)} [ 0, %loop_entry_{x.arg[0]} ], [ {ctx[x]}phi, %loop_latch_{x.arg[0]} ]"),
   (UPat(Ops.ENDRANGE, name="x"), lambda ctx,x:
    f"  br label %loop_latch_{x.src[0].arg[0]}\nloop_latch_{x.src[0].arg[0]}:\n"
-   f"  {ctx[x.src[0]]}phi = add {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, 1\n"
-   f"  {ctx[x]} = icmp ult {ldt(x.src[0].dtype)} {ctx[x.src[0]]}phi, {ctx[x.src[0].src[0]]}\n"
+   f"  {ctx[x.src[0]]}phi = add i32 {ctx[x.src[0]]}, 1\n  {ctx[x]} = icmp ult i32 {ctx[x.src[0]]}phi, {ctx[x.src[0].src[0]]}\n"
    f"  br i1 {ctx[x]}, label %loop_body_{x.src[0].arg[0]}, label %loop_exit_{x.src[0].arg[0]}\nloop_exit_{x.src[0].arg[0]}:"),
 
   # if

tinygrad/schedule/rangeify.py

@@ -7,7 +7,6 @@ from tinygrad.uop.symbolic import sym, symbolic_simple
 from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY, Context, flatten, dedup
 from tinygrad.schedule.kernelize import Kernel
 from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType
-from tinygrad.codegen.simplify import pm_flatten_range
 
 # *****************
 # 0. do some cleanup rewrites, mostly copied from the old stuff
@@ -313,8 +312,7 @@ def might_end_axis(idx:UOp):
   if all(x.op not in {Ops.REDUCE_AXIS} for x in idx.toposort()): return None
   to_end_axis = []
   for i,a in enumerate(idx.src[1:]):
-    # in RANGEIFY=1, always realize
-    if not (RANGEIFY > 1) or any(x.arg > idx.arg for x in a.toposort() if x.op is Ops.RANGE):
+    if any(x.arg > idx.arg for x in a.toposort() if x.op is Ops.RANGE):
       to_end_axis.append(i)
   if to_end_axis: return idx.replace(src=(idx.src[0].realize(arg=tuple(to_end_axis)),)+idx.src[1:], arg=None)
   return idx.replace(arg=None)
@@ -567,9 +565,6 @@ to_define_global = PatternMatcher([
   # this is only needed if you are using symbolic
   (UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"), lambda c: c.replace(src=()) if len(c.src) else None),
 
-  # remove RANGE with 0 size
-  (UPat(Ops.RANGE, name="r"), lambda r: UOp.const(dtypes.index, 0) if r.vmax == 0 else None),
-
   # renumber the ranges starting with 0 so that kernel deduping works
   (UPat(Ops.RANGE, name="r"), renumber_range),
 ])
@@ -612,7 +607,7 @@ def split_store(ctx:list[UOp], x:UOp):
 
   # local kernel rewrite
   lctx = LocalAddBufferContext()
-  ret = graph_rewrite(x, to_define_global+pm_flatten_range+rangeify_codegen+pm_remove_tags, ctx=lctx, name="kernel split", bottom_up=True)
+  ret = graph_rewrite(x, to_define_global+rangeify_codegen+pm_remove_tags, ctx=lctx, name="kernel split", bottom_up=True)
 
   # gather the metadata
   metadatas = [ctx[y].metadata for y in lctx.parent_tags]

tinygrad/tensor.py

@@ -4100,7 +4100,7 @@ class Tensor(MathTrait):
     R = self.clone()
     Q = Tensor.eye(m, dtype=self.dtype).reshape((1,) * len(b_shape) + (m, m)).expand(b_shape + (m, m)).contiguous()
     for i in range(min(m, n)):
-      x = R[..., i:m, i].contiguous()  # TODO: without contigous this can silently be wrong, should at least assert
+      x = R[..., i:m, i]
       s = -x[..., 0].sign()
       u1 = x[..., 0] - s * x.square().sum(-1).sqrt()
       w = x.unsqueeze(-1) / u1.reshape(b_shape + (1, 1))

tinygrad/uop/ops.py

@@ -217,7 +217,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   # determine what ranges this is in
   @functools.cached_property
-  def _ranges(self) -> dict[UOp, None]:
+  def ranges(self) -> dict[UOp, None]:
+    if self.op is Ops.RANGE: return {self:None}
     ret: dict[UOp, None] = {}
     if self.op in range_start.keys():
       for s in self.src[:range_start[self.op]]: ret.update(s.ranges)
@@ -227,11 +228,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
       for s in self.src: ret.update(s.ranges)
     return ret
 
-  @property
-  def ranges(self) -> dict[UOp, None]:
-    if self.op is Ops.RANGE: return {self:None}
-    return self._ranges
-
   # *** uop evaluation ***
 
   def simplify(self, tracked=False):