Variable.num -> NumNode (#2354)

test/external/fuzz_symbolic.py

@@ -1,7 +1,7 @@
 import itertools
 import random
 from tinygrad.helpers import DEBUG
-from tinygrad.shape.symbolic import Variable
+from tinygrad.shape.symbolic import Variable, NumNode
 random.seed(42)
 
 def add_v(expr, rng=None):
@@ -51,7 +51,7 @@ if __name__ == "__main__":
     tape = [random.choice(ops) for _ in range(random.randint(2, 30))]
     # 10% of the time, add one of lt, le, gt, ge
     if random.random() < 0.1: tape.append(random.choice([lt, le, gt, ge]))
-    expr = Variable.num(0)
+    expr = NumNode(0)
     rngs = []
     for t in tape:
       expr, rng = t(expr)

test/test_symbolic_shapetracker.py

@@ -1,6 +1,6 @@
 import unittest
 from tinygrad.shape.shapetracker import ShapeTracker, View
-from tinygrad.shape.symbolic import Variable
+from tinygrad.shape.symbolic import Variable, NumNode
 from tinygrad.tensor import Tensor
 
 class TestSymbolic(unittest.TestCase):
@@ -161,7 +161,7 @@ class TestSymbolicShapeExpr(unittest.TestCase):
     i = Variable("i", 1, 120)
     gidx0 = Variable("gidx0", 0, i)
     lidx1 = Variable("lidx1", 0, 7)
-    idx = (gidx0, lidx1, Variable.num(1))
+    idx = (gidx0, lidx1, NumNode(1))
     shape = (i+1, 8, 4)
     strides = (1, (i*4)+4, i+1)
     st = ShapeTracker((View.create(shape, strides), ))

test/unit/test_shapetracker.py

@@ -3,7 +3,7 @@ import unittest
 import numpy as np
 from tinygrad.helpers import prod, DEBUG
 from tinygrad.shape.shapetracker import ShapeTracker, View, get_contraction
-from tinygrad.shape.symbolic import Variable
+from tinygrad.shape.symbolic import Variable, NumNode
 from itertools import product
 
 def shapetracker_getitem(st, val):
@@ -136,7 +136,7 @@ class TestIndexExpressions2d(unittest.TestCase):
     shapes = [(30, 5), (15, 10), (15, 1), (5, 10), (5, 1)] # Make sure dim0 is a multiple of 5, one of the tests divides this dimension by 5
     offsets = [0, 1, 15, 28, 10000]
     self.sts = [ShapeTracker((View.create(base_shape, offset=offset),)) for base_shape in shapes for offset in offsets]
-    self.offset = [Variable.num(offset) for base_shape in shapes for offset in offsets]
+    self.offset = [NumNode(offset) for base_shape in shapes for offset in offsets]
     self.shapes = [shape for shape in shapes for offset in offsets]
     self.node_exprs = []
     self.idxs_exprs = []

test/unit/test_symbolic.py

@@ -74,13 +74,13 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(Variable("a", 0, 8)+1, 1, 9, "(1+a)")
 
   def test_add_num_1(self):
-    self.helper_test_variable(Variable("a", 0, 8)+Variable.num(1), 1, 9, "(1+a)")
+    self.helper_test_variable(Variable("a", 0, 8)+NumNode(1), 1, 9, "(1+a)")
 
   def test_sub_1(self):
     self.helper_test_variable(Variable("a", 0, 8)-1, -1, 7, "(-1+a)")
 
   def test_sub_num_1(self):
-    self.helper_test_variable(Variable("a", 0, 8)-Variable.num(1), -1, 7, "(-1+a)")
+    self.helper_test_variable(Variable("a", 0, 8)-NumNode(1), -1, 7, "(-1+a)")
 
   def test_mul_0(self):
     self.helper_test_variable(Variable("a", 0, 8)*0, 0, 0, "0")
@@ -120,7 +120,7 @@ class TestSymbolic(unittest.TestCase):
 
   def test_sum_div_some_partial_factor(self):
     self.helper_test_variable(Variable.sum([Variable("a", 0, 7)*6, Variable("b", 0, 7)*6]) // 16, 0, 5, "(((a*3)+(b*3))//8)")
-    self.helper_test_variable(Variable.sum([Variable.num(16), Variable("a", 0, 7)*6, Variable("b", 0, 7)*6]) // 16, 1, 6, "((((a*3)+(b*3))//8)+1)")
+    self.helper_test_variable(Variable.sum([NumNode(16), Variable("a", 0, 7)*6, Variable("b", 0, 7)*6]) // 16, 1, 6, "((((a*3)+(b*3))//8)+1)")
 
   def test_sum_div_no_factor(self):
     self.helper_test_variable(Variable.sum([Variable("a", 0, 7)*5, Variable("b", 0, 3)*5]) // 2, 0, 25, "(((a*5)+(b*5))//2)")
@@ -134,10 +134,10 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((1+Variable("a",1,2))%2, 0, 1, (Variable("a",1,2)-1).render())
 
   def test_sum_div_const(self):
-    self.helper_test_variable(Variable.sum([Variable("a", 0, 7)*4, Variable.num(3)]) // 4, 0, 7, "a")
+    self.helper_test_variable(Variable.sum([Variable("a", 0, 7)*4, NumNode(3)]) // 4, 0, 7, "a")
 
   def test_sum_div_const_big(self):
-    self.helper_test_variable(Variable.sum([Variable("a", 0, 7)*4, Variable.num(3)]) // 16, 0, 1, "(a//4)")
+    self.helper_test_variable(Variable.sum([Variable("a", 0, 7)*4, NumNode(3)]) // 16, 0, 1, "(a//4)")
 
   def test_sum_lt_fold(self):
     self.helper_test_variable(Variable.sum([Variable("a", 0, 7) * 4, Variable("b", 0, 3)]) < 16, 0, 1, "(a<4)")
@@ -195,23 +195,23 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((Variable("a", 0, 6) + 2) < 3, 0, 1, "(a<1)")
 
   def test_and_fold(self):
-    self.helper_test_variable(Variable.ands([Variable.num(0), Variable("a", 0, 1)]), 0, 0, "0")
+    self.helper_test_variable(Variable.ands([NumNode(0), Variable("a", 0, 1)]), 0, 0, "0")
 
   def test_and_remove(self):
-    self.helper_test_variable(Variable.ands([Variable.num(1), Variable("a", 0, 1)]), 0, 1, "a")
+    self.helper_test_variable(Variable.ands([NumNode(1), Variable("a", 0, 1)]), 0, 1, "a")
 
   def test_mod_factor_negative(self):
-    self.helper_test_variable(Variable.sum([Variable.num(-29), Variable("a", 0, 10), Variable("b", 0, 10)*28]) % 28, 0, 27, "((27+a)%28)")
-    self.helper_test_variable(Variable.sum([Variable.num(-29), Variable("a", 0, 100), Variable("b", 0, 10)*28]) % 28, 0, 27, "((27+a)%28)")
+    self.helper_test_variable(Variable.sum([NumNode(-29), Variable("a", 0, 10), Variable("b", 0, 10)*28]) % 28, 0, 27, "((27+a)%28)")
+    self.helper_test_variable(Variable.sum([NumNode(-29), Variable("a", 0, 100), Variable("b", 0, 10)*28]) % 28, 0, 27, "((27+a)%28)")
 
   def test_sum_combine_num(self):
-    self.helper_test_variable(Variable.sum([Variable.num(29), Variable("a", 0, 10), Variable.num(-23)]), 6, 16, "(6+a)")
+    self.helper_test_variable(Variable.sum([NumNode(29), Variable("a", 0, 10), NumNode(-23)]), 6, 16, "(6+a)")
 
   def test_sum_num_hoisted_and_factors_cancel_out(self):
     self.helper_test_variable(Variable.sum([Variable("a", 0, 1) * -4 + 1, Variable("a", 0, 1) * 4]), 1, 1, "1")
 
   def test_div_factor(self):
-    self.helper_test_variable(Variable.sum([Variable.num(-40), Variable("a", 0, 10)*2, Variable("b", 0, 10)*40]) // 40, -1, 9, "(-1+b)")
+    self.helper_test_variable(Variable.sum([NumNode(-40), Variable("a", 0, 10)*2, Variable("b", 0, 10)*40]) // 40, -1, 9, "(-1+b)")
 
   def test_mul_div(self):
     self.helper_test_variable((Variable("a", 0, 10)*4)//4, 0, 10, "a")
@@ -238,7 +238,7 @@ class TestSymbolicNumeric(unittest.TestCase):
     MIN, MAX = 0, 10
     # one number
     for i in range(MIN, MAX):
-      v = f(Variable.num(i))
+      v = f(NumNode(i))
       #print(i, f(i), v.min, v.max)
       self.assertEqual(v.min, v.max)
       self.assertEqual(v.min, f(i))

tinygrad/codegen/linearizer.py

@@ -82,7 +82,7 @@ class Linearizer(Kernel):
     ret = []
     invalid_value = 0 if dtypes.is_int(buf.dtype) else 0.0
     for idx, valid, rep_idx in zip(e_idxs, e_valids, Node.iter_idxs(expand_vars)):
-      this_const, idx, valid = (invalid_value, Variable.num(0), Variable.num(1)) if valid.max == 0 else (const, idx, valid)
+      this_const, idx, valid = (invalid_value, NumNode(0), NumNode(1)) if valid.max == 0 else (const, idx, valid)
       key = f"{acc}{localtype}{this_const if this_const is not None and acc is None else (buf.idx if isinstance(buf, MemBuffer) else cast(LocalBuffer, buf).name)}{idx.render()}{valid.render()}"
       if key not in self.load_cache:
         if acc is not None:
@@ -240,7 +240,7 @@ class Linearizer(Kernel):
         def calc_tc_idxs(local_size: int, aliases: List[List[int]]):
           replace_idxs = []
           for alias in aliases:
-            full_var, full_var_sz = Variable.num(0), 1
+            full_var, full_var_sz = NumNode(0), 1
             if alias[0] != 0:
               for i in alias:
                 next_var = local_idxs[-i] if i > 0 else Variable(None, 0, local_size-1)
@@ -318,7 +318,7 @@ class Linearizer(Kernel):
         stores = self.global_store(-1, fake_global_idxs+local_idxs+fake_reduce_idxs+upcast_idxs, acc)  # store accumulators
         barrier = self.uop(UOps.BARRIER, None, tuple(stores), cachable=False)
         if self.opts.has_local:
-          fake_idxs = [Variable.num(0)]*len(self.sts[-1].shape)
+          fake_idxs = [NumNode(0)]*len(self.sts[-1].shape)
           fake_idxs[self.global_dims+self.local_dims:self.global_dims+len(local_idxs)] = local_idxs[self.local_dims:]
           if_cond: UOp = (self.sts[-1].expr_idxs(fake_idxs)[0]<1).render(self.render_ops, self)
           barrier = self.uop(UOps.IF, None, (if_cond, barrier), cachable=False)

tinygrad/graph.py

@@ -9,7 +9,7 @@ from tinygrad.ops import ScheduleItem, UnaryOps, BinaryOps, ReduceOps, MovementO
 from tinygrad.helpers import GRAPH, GRAPHPATH, DEBUG, GlobalCounters, getenv, dedup
 from tinygrad.codegen.linearizer import UOps
 from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.symbolic import Variable
+from tinygrad.shape.symbolic import NumNode
 
 # **** debugging and graphing ****
 
@@ -53,7 +53,7 @@ def add_st_node(nmx, nmo, label, st:ShapeTracker):
   global node_count
   inter_node = node_count
   node_count += 1
-  offset = st.expr_node(Variable.num(0))[0]
+  offset = st.expr_node(NumNode(0))[0]
   G.add_node(inter_node, style='filled', fillcolor="#80ff8080", color="black", label=f"{st.shape}\n{st.real_strides()}" + (f"\n{offset}" if offset != 0 else ""))
   G.add_edge(nmx, inter_node, color='#00000060')
   G.add_edge(inter_node, nmo, label=label, color='#00000060')

tinygrad/shape/shapetracker.py

@@ -5,7 +5,7 @@ from dataclasses import dataclass
 from typing import Tuple, List, Optional, Dict, cast
 from tinygrad.ops import MovementOps
 from tinygrad.helpers import prod, DEBUG, dedup, merge_dicts
-from tinygrad.shape.symbolic import Variable, MulNode, Node, SumNode, sint
+from tinygrad.shape.symbolic import Variable, MulNode, Node, SumNode, NumNode, sint
 from tinygrad.shape.view import View
 
 @functools.lru_cache(maxsize=None)
@@ -35,7 +35,7 @@ def expr_node_mask(view:View, idx, valid=None) -> Node:
 # generate an expression if you have a single idx variable
 def expr_node(view:View, idx=None) -> Node:
   if idx is None: idx = Variable('idx', 0, prod(view.shape)-1)
-  ret: List[Node] = [Variable.num(view.offset) if isinstance(view.offset, int) else view.offset] if view.offset else []
+  ret: List[Node] = [NumNode(view.offset) if isinstance(view.offset, int) else view.offset] if view.offset else []
   acc = 1
   for d,s in reversed(to_shape_strides(view.shape, view.strides)):
     ret.append(((idx//acc)%d)*s)
@@ -45,7 +45,7 @@ def expr_node(view:View, idx=None) -> Node:
 # generate an expression if you have a variable or expression for each index
 def expr_idxs(view:View, idxs) -> Node:
   assert len(idxs) == len(view.shape), f"need an idx for all dimensions {idxs} vs {view.shape}"
-  return Variable.sum([Variable.num(view.offset) if isinstance(view.offset, int) else view.offset] + [idx*st for idx,sh,st in zip(idxs, view.shape, view.strides) if sh != 1 and st != 0])
+  return Variable.sum([NumNode(view.offset) if isinstance(view.offset, int) else view.offset] + [idx*st for idx,sh,st in zip(idxs, view.shape, view.strides) if sh != 1 and st != 0])
 
 @functools.lru_cache(maxsize=None)
 def merge_views(vm2:View, vm1:View) -> Optional[View]:
@@ -129,7 +129,7 @@ class ShapeTracker:
 
   def _expr_idx(self, idx, valid) -> Tuple[Node, Node]:
     for v in reversed(self.views[0:-1]):
-      if valid.max == 0: return Variable.num(-1), valid
+      if valid.max == 0: return NumNode(-1), valid
       valid = expr_node_mask(v, idx, valid)
       idx = expr_node(v, idx)
     return idx, valid

tinygrad/shape/symbolic.py

@@ -46,7 +46,7 @@ class Node:
     if not isinstance(other, Node): return NotImplemented
     return self.key == other.key
   def __neg__(self): return self*-1
-  def __add__(self, b:Union[Node,int]): return Variable.sum([self, b if isinstance(b, Node) else Variable.num(b)])
+  def __add__(self, b:Union[Node,int]): return Variable.sum([self, b if isinstance(b, Node) else NumNode(b)])
   def __radd__(self, b:int): return self+b
   def __sub__(self, b:Union[Node,int]): return self+-b
   def __rsub__(self, b:int): return -self+b
@@ -101,9 +101,6 @@ class Node:
     if self.min < 0: return (self - ((self.min//b)*b)) % b
     return create_node(ModNode(self, b))
 
-  @staticmethod
-  def num(num:int) -> NumNode: return NumNode(num)
-
   @staticmethod
   def sum(nodes:List[Node]) -> Node:
     nodes = [x for x in nodes if x.max or x.min]
@@ -274,7 +271,7 @@ class SumNode(RedNode):
     if isinstance(b, Node) and (b - self).min > 0: return self # b - self simplifies the node
     new_nodes: List[Node] = []
     for x in self.nodes:
-      if x.__class__ is NumNode: new_nodes.append(Variable.num(x.b%b))
+      if x.__class__ is NumNode: new_nodes.append(NumNode(x.b%b))
       elif isinstance(x, MulNode): new_nodes.append(x.a * (x.b%b))
       else: new_nodes.append(x)
     return Node.__mod__(Node.sum(new_nodes), b)
@@ -327,7 +324,7 @@ def sym_rename(s) -> str: return f"s{sym_rename.cache_info().currsize}"
 def sym_render(a: Union[Node, int], ops=None, ctx=None) -> str: return str(a) if isinstance(a, int) else a.render(ops, ctx)
 def sym_infer(a: Union[Node, int], var_vals: Dict[Variable, int]) -> int:
   if isinstance(a, (int, float)): return a
-  ret = a.substitute({k:Variable.num(v) for k, v in var_vals.items()})
+  ret = a.substitute({k:NumNode(v) for k, v in var_vals.items()})
   assert isinstance(ret, NumNode), f"sym_infer didn't produce NumNode from {a} with {var_vals}"
   return ret.b