From 0571dd76279980d6f68ff4c8f76a588cec4bfb91 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Sat, 23 Sep 2023 14:43:45 +0800 Subject: [PATCH] move all int (#1903) --- tinygrad/helpers.py | 2 ++ tinygrad/lazy.py | 4 ++-- tinygrad/nn/__init__.py | 3 +-- tinygrad/renderer/triton.py | 14 +++++++------- tinygrad/shape/symbolic.py | 4 ---- tinygrad/shape/view.py | 4 ++-- tinygrad/tensor.py | 4 ++-- 7 files changed, 16 insertions(+), 19 deletions(-) diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 7873bf40ec..82aa9d0b40 100644 --- a/tinygrad/helpers.py +++ b/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 diff --git a/tinygrad/lazy.py b/tinygrad/lazy.py index f863855505..f9d511ef57 100644 --- a/tinygrad/lazy.py +++ b/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 diff --git a/tinygrad/nn/__init__.py b/tinygrad/nn/__init__.py index 05f206c03c..edc23417ed 100644 --- a/tinygrad/nn/__init__.py +++ b/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): diff --git a/tinygrad/renderer/triton.py b/tinygrad/renderer/triton.py index e541886651..0a8ca7cdb4 100644 --- a/tinygrad/renderer/triton.py +++ b/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() diff --git a/tinygrad/shape/symbolic.py b/tinygrad/shape/symbolic.py index 02d368b13d..1fc33d0053 100644 --- a/tinygrad/shape/symbolic.py +++ b/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] = { diff --git a/tinygrad/shape/view.py b/tinygrad/shape/view.py index 7dc43812ba..c13be65bfa 100644 --- a/tinygrad/shape/view.py +++ b/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, ...]: diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index fa4d9bfe2a..d747ab1504 100644 --- a/tinygrad/tensor.py +++ b/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: