move all int (#1903)

tinygrad/helpers.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 import os, functools, platform, time, re, contextlib, operator
 import numpy as np
 from typing import Dict, Tuple, Union, List, NamedTuple, Final, Iterator, ClassVar, Optional, Iterable, Any, TypeVar
+from typing_extensions import TypeGuard
 
 T = TypeVar("T")
 # NOTE: it returns int 1 if x is empty regardless of the type of x
@@ -15,6 +16,7 @@ def dedup(x): return list(dict.fromkeys(x))   # retains list order
 def argfix(*x): return tuple(x[0]) if x and x[0].__class__ in (tuple, list) else x
 def argsort(x): return type(x)(sorted(range(len(x)), key=x.__getitem__)) # https://stackoverflow.com/questions/3382352/equivalent-of-numpy-argsort-in-basic-python
 def all_same(items): return all(x == items[0] for x in items)
+def all_int(t: Tuple[Any, ...]) -> TypeGuard[Tuple[int, ...]]: return all(isinstance(s, int) for s in t)
 def colored(st, color, background=False): return f"\u001b[{10*background+60*(color.upper() == color)+30+['black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'].index(color.lower())}m{st}\u001b[0m" if color is not None else st  # replace the termcolor library with one line
 def ansilen(s): return len(re.sub('\x1b\\[(K|.*?m)', '', s))
 def make_pair(x:Union[int, Tuple[int, ...]], cnt=2) -> Tuple[int, ...]: return (x,)*cnt if isinstance(x, int) else x

tinygrad/lazy.py

@@ -5,10 +5,10 @@ from weakref import ref, WeakSet, WeakValueDictionary
 
 import numpy as np
 from tinygrad.graph import log_op
-from tinygrad.helpers import GRAPH, DEBUG, prod, getenv, DType, dtypes, flatten, ImageDType, partition
+from tinygrad.helpers import GRAPH, DEBUG, prod, getenv, DType, dtypes, flatten, ImageDType, partition, all_int
 from tinygrad.ops import Device, Compiled, UnaryOps, BinaryOps, TernaryOps, ReduceOps, MovementOps, LoadOps, OpType, LazyOp
 from tinygrad.shape.shapetracker import ShapeTracker, View, get_contraction
-from tinygrad.shape.symbolic import Variable, sint, all_int
+from tinygrad.shape.symbolic import Variable, sint
 
 from tinygrad.runtime.lib import RawConst, RawBuffer, RawBufferMapped, RawBufferTransfer
 from tinygrad.runtime.ops_cpu import RawNumpyBuffer

tinygrad/nn/__init__.py

@@ -1,8 +1,7 @@
 import math
 from typing import Optional, Union, Tuple
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import prod
-from tinygrad.shape.symbolic import all_int
+from tinygrad.helpers import prod, all_int
 
 class BatchNorm2d:
   def __init__(self, sz, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1):

tinygrad/renderer/triton.py

@@ -3,7 +3,7 @@ from collections import defaultdict
 from tinygrad.ops import UnaryOps, BinaryOps, TernaryOps, Op
 from tinygrad.helpers import dtypes, ImageDType, DEBUG, getenv
 from tinygrad.codegen.linearizer import  UOp, UOps
-from triton.compiler import compile as triton_compile
+from triton.compiler import compile as triton_compile  # type: ignore
 import hashlib
 import math
 import re
@@ -49,7 +49,7 @@ def uops_to_triton(function_name:str, uops:List[UOp]):
   for ru in uops:
     for v in ru.vin:
       child_count[v] += 1
-  
+
   def kk(s): kernel.append("  "*depth+s)
   code_for_op: Final[Dict[Op, Callable]] = {
     UnaryOps.EXP2: lambda x,: f"tl.math.exp2({x})",
@@ -80,7 +80,7 @@ def uops_to_triton(function_name:str, uops:List[UOp]):
       if child_count[u] <=1 or dtypes.is_int(dtype): r[u] = int_div(*[r[x] for x in vin]) if args == BinaryOps.DIV and dtypes.is_int(dtype) else val
       else: kk(f"{ssa(u, 'alu')} = ({val}).to({triton_dtypes[dtype]})")
     elif uop == UOps.LOAD:
-      assert dtype is not None 
+      assert dtype is not None
       if len(vin) == 2: kk(f"{ssa(u, 'val')} = tl.load({r[vin[0]]} + { fill_dims_for_idx(r[vin[1]], dims)}, mask = {render_valid(valid)}).to({triton_dtypes[vin[0].dtype]})")# type: ignore
       else: kk(f"{ssa(u, 'val')} = tl.where({r[vin[2]]}, tl.load({r[vin[0]]}+{fill_dims_for_idx(r[vin[1]],dims)} , mask={render_valid(valid+[r[vin[2]]])}), 0.0).to({triton_dtypes[vin[0].dtype]})")# type: ignore
     elif uop == UOps.DEFINE_ACC: kk(f"{ssa(u, 'acc')} = {define_scalar(local_size, triton_dtypes[dtype], args).replace('//', '/')}") # type: ignore
@@ -101,15 +101,15 @@ def uops_to_triton(function_name:str, uops:List[UOp]):
         kk(f"{args[1]} = tl.arange({0}, {next_power_of_2(args[2])})")
         local_size.append(args[2])
       r[u] = args[1]
-    else: raise NotImplementedError(f"unimplemented: {uop}")  
-  
+    else: raise NotImplementedError(f"unimplemented: {uop}")
+
   prg = f"import triton\nimport triton.language as tl\ntl.core.TRITON_MAX_TENSOR_NUMEL = float('inf')\n@triton.jit\ndef {function_name}("+','.join(f"{buf[0]}" for buf in bufs)+"):\n"
   for i, line in enumerate(list(filter(lambda line: "tl.arange" in line, kernel))): kernel[kernel.index(line)] +=  f"[{', '.join([':' if i == j else 'None' for j in range(len(local_size))])}]"
   prg += "\n".join(kernel)
-  
+
   acc_local_size = 1
   for x in local_size: acc_local_size *= next_power_of_2(x)
-  local_size = [acc_local_size] + [1] * (len(local_size) - 1)  
+  local_size = [acc_local_size] + [1] * (len(local_size) - 1)
 
   if DEBUG >=4: print(prg)
   hsh = hashlib.md5(prg.encode('utf-8')).hexdigest()

tinygrad/shape/symbolic.py

@@ -5,7 +5,6 @@ from math import gcd
 from itertools import product
 from tinygrad.helpers import partition
 from typing import List, Dict, Callable, Tuple, Type, Union, Optional, Any, Iterator
-from typing_extensions import TypeGuard
 
 # NOTE: Python has different behavior for negative mod and floor div than c
 # symbolic matches the Python behavior, but the code output is agnostic, and will never have negative numbers in div or mod
@@ -323,9 +322,6 @@ def sym_infer(a: Union[Node, int], var_vals: Dict[Variable, int]) -> int:
 
 # symbolic int
 sint = Union[Node, int]
-
-def all_int(t: Tuple[sint, ...]) -> TypeGuard[Tuple[int, ...]]: return all(isinstance(s, int) for s in t)
-
 VariableOrNum = Union[Variable, NumNode]
 
 render_python: Dict[Type, Callable] = {

tinygrad/shape/view.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 import functools
 from typing import Tuple, List, Optional, NamedTuple
-from tinygrad.helpers import prod
-from tinygrad.shape.symbolic import NumNode, is_sym_int, sint, all_int
+from tinygrad.helpers import prod, all_int
+from tinygrad.shape.symbolic import NumNode, is_sym_int, sint
 
 @functools.lru_cache(maxsize=None)
 def filter_strides(shape:Tuple[int, ...], strides:Tuple[int, ...]) -> Tuple[int, ...]:

tinygrad/tensor.py

@@ -7,10 +7,10 @@ from itertools import accumulate
 import numpy as np
 from typing import List, Tuple, Callable, Optional, ClassVar, Type, Union, Sequence
 
-from tinygrad.helpers import ImageDType, argfix, make_pair, getenv, IMAGE, DEBUG, flatten, DType, dtypes, prod
+from tinygrad.helpers import ImageDType, argfix, make_pair, getenv, IMAGE, DEBUG, flatten, DType, dtypes, prod, all_int
 from tinygrad.lazy import LazyBuffer
 from tinygrad.ops import Device, LoadOps
-from tinygrad.shape.symbolic import sint, all_int
+from tinygrad.shape.symbolic import sint
 
 # An instantiation of the Function is the Context
 class Function: