move is_increasing to ops [pr] (#9134)

test/unit/test_simplify_valid_idx.py

@@ -1,6 +1,6 @@
 import unittest, itertools
 
-from tinygrad.codegen.rewriter import full_graph_rewrite, is_increasing
+from tinygrad.codegen.rewriter import full_graph_rewrite
 from tinygrad.dtype import dtypes
 from tinygrad.ops import UOp, Ops
 from tinygrad.codegen.symbolic import simplify_valid
@@ -34,14 +34,14 @@ class TestHelpers(unittest.TestCase):
     f2 = (idx2*2)+ridx1+((idx1+((ridx2+7)//8)+31)//32)+(-2)
     f3 = (idx2*2)+ridx1+(-1)
 
-    self.assertFalse(is_increasing(f0))
-    self.assertTrue(is_increasing(f1))
-    self.assertTrue(is_increasing(f2))
-    self.assertTrue(is_increasing(f3))
+    self.assertFalse(f0.is_increasing())
+    self.assertTrue(f1.is_increasing())
+    self.assertTrue(f2.is_increasing())
+    self.assertTrue(f3.is_increasing())
 
     rng = UOp(Ops.RANGE, dtypes.int, arg=(2, True), src=(UOp(Ops.CONST, dtypes.int, arg=0, src=()), UOp(Ops.CONST, dtypes.int, arg=5, src=()),))
-    self.assertTrue(is_increasing(rng))
-    self.assertTrue(is_increasing(rng+2))
+    self.assertTrue(rng.is_increasing())
+    self.assertTrue((rng+2).is_increasing())
 
 class TestValidIdxSimplification(unittest.TestCase):
   def check(self, load, sidx, svalid):

tinygrad/codegen/lowerer.py

@@ -1,5 +1,5 @@
 # the job of the lowerer is to do indexing
-import functools, itertools, operator
+import functools, itertools, operator, math
 from dataclasses import dataclass
 from typing import cast
 from tinygrad.dtype import dtypes, PtrDType
@@ -7,7 +7,6 @@ from tinygrad.ops import KernelInfo, UOp, Ops, graph_rewrite, PatternMatcher, UP
 from tinygrad.renderer import Renderer
 from tinygrad.helpers import all_int, prod, partition, flatten, unwrap
 from tinygrad.codegen.expander import expand_rewrite
-import math
 
 # returns the axes to create new_shape if new_shape can be created by combining axis from old_shape
 def get_contraction(old_shape:tuple[sint, ...], new_shape:tuple[sint, ...]) -> list[list[int]]|None:

tinygrad/codegen/rewriter.py

@@ -3,7 +3,7 @@ import functools, operator
 from collections import defaultdict
 from tinygrad.dtype import dtypes, ImageDType, PtrDType
 from tinygrad.ops import UOp, Ops, UPat, PatternMatcher, resolve
-from tinygrad.ops import graph_rewrite, is_increasing, GroupOp
+from tinygrad.ops import graph_rewrite, GroupOp
 from tinygrad.codegen.symbolic import symbolic_simple, split_uop, uop_given_valid, parse_valid, simplify_valid, sym, mulacc_unrolled
 from tinygrad.helpers import getenv, flatten, dedup, TRANSCENDENTAL, AMX, prod, DEVECTORIZE
 from tinygrad.codegen.transcendental import xexp2, xlog2, xsin, xpow, TRANSCENDENTAL_SUPPORTED_DTYPES
@@ -109,7 +109,7 @@ def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
     # if X >= c, check if it's out of bound when X = c-1
     test_value = c + 1 if is_upper_bound else c - 1
     for i,b in zip(idx.src, (buf.dtype.shape[1], buf.dtype.shape[0])):
-      if is_increasing(i):
+      if i.is_increasing():
         rw = i.substitute({X:X.const_like(test_value)}).simplify()
         if rw.vmin >= b or rw.vmax < 0:
           drop_stmt.append(stmt)

tinygrad/ops.py

@@ -588,6 +588,12 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   # *** uop symbolic stuff ***
 
+  def is_increasing(self:UOp) -> bool:
+    # is f a monotonically increasing function regards its input
+    if self.op in GroupOp.Irreducible: return True
+    if self.op is Ops.ADD: return self.src[0].is_increasing() and self.src[1].is_increasing()
+    if self.op in (Ops.MUL, Ops.IDIV) and self.src[1].op is Ops.CONST and self.src[1].arg >= 0: return self.src[0].is_increasing()
+    return False  # False if not sure
   def const_factor(self) -> int:
     """largest known int that divides self"""
     if self.op is Ops.CONST: return self.arg
@@ -938,13 +944,6 @@ def graph_rewrite_map(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False) ->
   rewrite_ctx = RewriteContext(pm, ctx)
   return {k:(rewrite_ctx.bottom_up_rewrite(k) if bottom_up else rewrite_ctx.top_down_rewrite(k)) for k in list(sink.toposort)[::-1]}
 
-def is_increasing(f:UOp) -> bool:
-  # is f a monotonically increasing function regards its input
-  if f.op in GroupOp.Irreducible: return True
-  if f.op is Ops.ADD: return is_increasing(f.src[0]) and is_increasing(f.src[1])
-  if f.op in (Ops.MUL, Ops.IDIV) and f.src[1].op is Ops.CONST and f.src[1].arg >= 0: return is_increasing(f.src[0])
-  return False  # False if not sure
-
 def sint_to_uop(x:sint, dtype:DType=dtypes.int) -> UOp: return UOp.const(dtype, x) if isinstance(x, int) else x
 
 _substitute = PatternMatcher([(UPat(tuple(Ops), name="x"), lambda ctx,x: ctx.get(x,None))])