forked from IQ.Lvbs/IQ.Pilot
IQ.Pilot Prebuilt Release @ ab07000
This commit is contained in:
IQ.Lvbs history cleanup
147
tinygrad_repo/tinygrad/uop/__init__.py
Normal file
147
tinygrad_repo/tinygrad/uop/__init__.py
Normal file
@@ -0,0 +1,147 @@
|
||||
# flake8: noqa: E702
|
||||
# allow semicolons to put multiple ops on one line
|
||||
from enum import auto, IntEnum, Enum
|
||||
|
||||
# wrapper around IntEnum that preserves Enum.__str__ and makes auto() unique across all FastEnum subclasses
|
||||
class FastEnum(IntEnum):
|
||||
def __str__(self): return Enum.__str__(self)
|
||||
def __repr__(x): return str(x)
|
||||
@staticmethod
|
||||
def _generate_next_value_(_, __, ___, last_values): return 1 + max([0, *last_values, *[max(c) for c in FastEnum.__subclasses__()]])
|
||||
|
||||
# the order of these Ops controls the order of the toposort
|
||||
class Ops(FastEnum):
|
||||
# ** 1 -- defines/special **
|
||||
|
||||
# define GLOBAL/VAR are ptrs to outside the Kernel
|
||||
DEFINE_VAR = auto(); BIND = auto()
|
||||
|
||||
# this is a RANGE for GPU dimensions, similar to symbolic shapes but not exactly
|
||||
SPECIAL = auto()
|
||||
|
||||
# define LOCAL/REG allocate things
|
||||
DEFINE_LOCAL = auto(); DEFINE_REG = auto()
|
||||
|
||||
# ** 2 -- non op uops **
|
||||
|
||||
# uops that aren't rendered
|
||||
NOOP = auto(); REWRITE_ERROR = auto()
|
||||
# FUNCTION has a TUPLE body and is gradient-able; CALL is an opaque kernel invocation
|
||||
PARAM = auto(); FUNCTION = auto(); CALL = auto()
|
||||
|
||||
# renderer
|
||||
# LINEAR is a list of UOps, SOURCE has a str arg that's human readable, BINARY has bytes arg that's compiled
|
||||
PROGRAM = auto(); LINEAR = auto(); SOURCE = auto(); BINARY = auto()
|
||||
|
||||
# AFTER passes src[0] through and promises in the toposort that any consumers of the AFTER run after src[1:]
|
||||
# GROUP is a NOOP that just merges things together
|
||||
SINK = auto(); AFTER = auto(); GROUP = auto()
|
||||
|
||||
# vector creation / item selection
|
||||
GEP = auto(); STACK = auto()
|
||||
|
||||
# tuple/gettuple for function with multiple returns
|
||||
TUPLE = auto(); GETTUPLE = auto()
|
||||
|
||||
# hcq specific
|
||||
GETADDR = auto()
|
||||
|
||||
# ** 3 -- load/store **
|
||||
|
||||
# INDEX is a BinaryOp similar to ADD, but it operates on pointers
|
||||
INDEX = auto()
|
||||
|
||||
# load/store before math
|
||||
LOAD = auto(); STORE = auto()
|
||||
|
||||
# ** 4 -- math **
|
||||
|
||||
# tensor core math op, not elementwise
|
||||
WMMA = auto(); SHAPED_WMMA = auto()
|
||||
|
||||
# UnaryOps
|
||||
CAST = auto(); BITCAST = auto(); EXP2 = auto(); LOG2 = auto(); SIN = auto()
|
||||
SQRT = auto(); RECIPROCAL = auto(); NEG = auto(); TRUNC = auto()
|
||||
|
||||
# BinaryOps
|
||||
ADD = auto(); MUL = auto(); SHL = auto(); SHR = auto(); CDIV = auto(); MAX = auto(); CMOD = auto()
|
||||
CMPLT = auto(); CMPNE = auto(); CMPEQ = auto()
|
||||
XOR = auto(); OR = auto(); AND = auto()
|
||||
THREEFRY = auto(); SUB = auto(); FDIV = auto(); POW = auto()
|
||||
FLOORDIV = auto(); FLOORMOD = auto()
|
||||
|
||||
# TernaryOps
|
||||
WHERE = auto(); MULACC = auto()
|
||||
|
||||
# ** 5 -- control flow / consts / custom **
|
||||
|
||||
# control flow ops
|
||||
BARRIER = auto(); RANGE = auto(); IF = auto(); END = auto(); ENDIF = auto(); WAIT = auto()
|
||||
|
||||
# const.
|
||||
CONST = auto()
|
||||
|
||||
# CUSTOM/CUSTOMI are used to output strings into codegen. the I makes the string inline
|
||||
CUSTOM = auto(); CUSTOMI = auto()
|
||||
|
||||
# INS is a machine instruction
|
||||
INS = auto()
|
||||
|
||||
# ** 6 -- ops that don't exist in programs **
|
||||
|
||||
# tensor graph ops
|
||||
UNIQUE = auto(); DEVICE = auto()
|
||||
|
||||
# local unique
|
||||
LUNIQUE = auto()
|
||||
|
||||
# ops that adjust the behavior of the scheduler
|
||||
CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto()
|
||||
|
||||
# buffer ops
|
||||
STAGE = auto(); COPY = auto(); BUFFER = auto(); SLICE = auto(); MSELECT = auto(); MSTACK = auto(); CUSTOM_FUNCTION = auto()
|
||||
|
||||
# the core 6 movement ops! these only exist in the tensor graph
|
||||
RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); FLIP = auto()
|
||||
MULTI = auto() # MULTI is really a movement op
|
||||
|
||||
# reduce
|
||||
REDUCE = auto(); ALLREDUCE = auto()
|
||||
|
||||
# expander ops
|
||||
UNROLL = auto(); CONTRACT = auto(); VCAT = auto(); PTRCAT = auto()
|
||||
|
||||
class GroupOp:
|
||||
Unary = {Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.SQRT, Ops.RECIPROCAL, Ops.NEG, Ops.TRUNC}
|
||||
Binary = {Ops.ADD, Ops.MUL, Ops.CDIV, Ops.MAX, Ops.CMOD, Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ,
|
||||
Ops.XOR, Ops.SHL, Ops.SHR, Ops.OR, Ops.AND, Ops.THREEFRY, Ops.SUB, Ops.FDIV, Ops.POW, Ops.FLOORDIV, Ops.FLOORMOD}
|
||||
Ternary = {Ops.WHERE, Ops.MULACC}
|
||||
ALU = set.union(Unary, Binary, Ternary)
|
||||
Broadcastable = set.union(Binary, Ternary, {Ops.GROUP, Ops.STORE})
|
||||
|
||||
# TODO: is BITCAST always Elementwise if it's shape changing?
|
||||
Elementwise = set.union(ALU, {Ops.CAST, Ops.BITCAST})
|
||||
|
||||
Defines = {Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}
|
||||
|
||||
Irreducible = {Ops.CONST, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.RANGE}
|
||||
Movement = {Ops.RESHAPE, Ops.EXPAND, Ops.PERMUTE, Ops.PAD, Ops.SHRINK, Ops.FLIP}
|
||||
|
||||
Buffer = {Ops.LOAD, Ops.STORE, Ops.CONST, Ops.DEFINE_VAR}
|
||||
|
||||
# BinaryOps that can be flipped
|
||||
Commutative = {Ops.ADD, Ops.MUL, Ops.MAX, Ops.CMPNE, Ops.CMPEQ, Ops.XOR, Ops.AND, Ops.OR}
|
||||
|
||||
# BinaryOps where f(f(a,b),c) = f(a,f(b,c))
|
||||
Associative = {Ops.ADD, Ops.MUL, Ops.AND, Ops.OR, Ops.MAX}
|
||||
|
||||
# BinaryOps that satisfy f(x,x)=x see https://en.wikipedia.org/wiki/Idempotence
|
||||
Idempotent = {Ops.OR, Ops.AND, Ops.MAX}
|
||||
|
||||
# These can change the dtype to bool
|
||||
Comparison = {Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ}
|
||||
|
||||
# do not preserve f(0) = 0
|
||||
UnsafePad = {Ops.RECIPROCAL, Ops.LOG2, Ops.EXP2, Ops.CDIV, Ops.POW, Ops.FLOORDIV}
|
||||
|
||||
All = set(Ops)
|
||||
572
tinygrad_repo/tinygrad/uop/decompositions.py
Normal file
572
tinygrad_repo/tinygrad/uop/decompositions.py
Normal file
@@ -0,0 +1,572 @@
|
||||
from typing import Callable
|
||||
import math, functools
|
||||
from tinygrad.dtype import dtypes, DType, promo_lattice, truncate
|
||||
from tinygrad.helpers import flatten, polyN, DEBUG, EMULATED_DTYPES
|
||||
from tinygrad.uop import GroupOp
|
||||
from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite
|
||||
from tinygrad.renderer import Renderer
|
||||
|
||||
TRANSCENDENTAL_DTYPES = (dtypes.float16, dtypes.float32, dtypes.float64)
|
||||
|
||||
def _lazy_map_numbers(x:UOp, inf:UOp, _inf:UOp, nan:UOp, ratio:UOp):
|
||||
"""replace inf -> inf, -inf -> _inf, nan -> nan, otherwise -> ratio"""
|
||||
return x.ne(math.inf).where(x.ne(x).where(nan, x.ne(-math.inf).where(ratio, _inf)), inf)
|
||||
|
||||
# *** helper functions for bit manipulation ***
|
||||
def mantissa_bits(d:DType) -> int: return dtypes.finfo(d.scalar())[1]
|
||||
def exponent_bias(d:DType) -> int: return (1 << (dtypes.finfo(d.scalar())[0] - 1)) - (0 if d.scalar() in dtypes.fp8_fnuz else 1)
|
||||
def exponent_mask(d:DType) -> int: return (1 << dtypes.finfo(d.scalar())[0]) - 1
|
||||
|
||||
# **** utils ****
|
||||
def shr(x:UOp|int, y:UOp|int) -> UOp: return x // (2**(y.simplify().arg) if isinstance(y, UOp) else 2**y)
|
||||
def shl(x:UOp|int, y:UOp|int) -> UOp: return x * (2**(y.simplify().arg) if isinstance(y, UOp) else 2**y)
|
||||
|
||||
def rintk(d:UOp) -> UOp:
|
||||
"""round d:float to int away from 0"""
|
||||
out_dtype = {dtypes.float64: dtypes.int64, dtypes.float32: dtypes.int32, dtypes.float16: dtypes.int16}[d.dtype.scalar()].vec(d.dtype.vcount)
|
||||
return (d + (d<0.0).where(d.const_like(-0.5), d.const_like(0.5))).cast(out_dtype)
|
||||
|
||||
def pow2if(q:UOp, float_dtype:DType):
|
||||
"""cast(2^q, float_dtype) where q is any integer in the range of [-126, 127]"""
|
||||
out_dtype = {dtypes.int64: dtypes.float64, dtypes.int32: dtypes.float32, dtypes.int16: float_dtype.scalar()}[q.dtype.scalar()].vec(q.dtype.vcount)
|
||||
return shl(q + exponent_bias(out_dtype), mantissa_bits(out_dtype)).bitcast(out_dtype)
|
||||
|
||||
def ilogb2k(d:UOp) -> UOp:
|
||||
"""calculate the integer part of log2(d), where d is normalized fp value in the range of [0, +inf)."""
|
||||
assert d.dtype.scalar() in TRANSCENDENTAL_DTYPES
|
||||
dint = d.bitcast({dtypes.float64: dtypes.int64, dtypes.float32: dtypes.int32, dtypes.float16: dtypes.int16}[d.dtype.scalar()].vec(d.dtype.vcount))
|
||||
# -1 <= ilog2bk(d) <= 128
|
||||
return (shr(dint, mantissa_bits(d.dtype)) & exponent_mask(d.dtype)) - exponent_bias(d.dtype)
|
||||
|
||||
def ldexp3k(d:UOp, e:UOp) -> UOp:
|
||||
"""d*2^e. e is a number obtained by casting an integer in the range [-127, 127] to a float. d is any float number."""
|
||||
assert d.dtype.scalar() in TRANSCENDENTAL_DTYPES and e.dtype.scalar() in TRANSCENDENTAL_DTYPES
|
||||
dtype = {dtypes.float64: dtypes.int64, dtypes.float32: dtypes.int32, dtypes.float16: dtypes.int16}[d.dtype.scalar()].vec(d.dtype.count)
|
||||
m1 = d.bitcast(dtype)
|
||||
m2 = shl(e.cast(dtype), mantissa_bits(d.dtype))
|
||||
return (m1 + m2).bitcast(d.dtype)
|
||||
|
||||
def ldexp2k(d:UOp, e:UOp) -> UOp:
|
||||
"""d*2^e. much faster than ldexp3k but risky. d > 0 and d is not denormal."""
|
||||
assert d.dtype.scalar() in TRANSCENDENTAL_DTYPES and e.dtype.scalar() in (dtypes.int16, dtypes.int32, dtypes.int64)
|
||||
return (d * pow2if(shr(e, 1), d.dtype)) * pow2if(e - shr(e, 1), d.dtype)
|
||||
|
||||
def frexp(v:UOp) -> tuple[UOp, UOp]:
|
||||
"""frexp(v) -> (mantissa, exponent) assuming v != 0"""
|
||||
assert v.dtype.scalar() in TRANSCENDENTAL_DTYPES
|
||||
# m1 = masks for mantissa, m2 = masks to normalize the mantissa.
|
||||
m1 = {dtypes.float64: 0x000FFFFFFFFFFFFF, dtypes.float32: 0x807FFFFF, dtypes.float16: 0x83FF}[v.dtype.scalar()]
|
||||
m2 = {dtypes.float64: 0x3FE0000000000000, dtypes.float32: 0x3F000000, dtypes.float16: 0x3800}[v.dtype.scalar()]
|
||||
bits = v.bitcast({dtypes.float64: dtypes.uint64, dtypes.float32: dtypes.uint32, dtypes.float16: dtypes.uint16}[v.dtype.scalar()].vec(v.dtype.count))
|
||||
exponent = shr(bits, mantissa_bits(v.dtype)) & exponent_mask(v.dtype)
|
||||
# Set the exponent bits appropriately to normalize the mantissa into the range of [0.5, 1.0).
|
||||
mantissa = ((bits & m1) | m2).bitcast(v.dtype)
|
||||
exp = exponent - exponent_bias(v.dtype) + 1
|
||||
return mantissa, exp
|
||||
|
||||
# *** reduction algorithms for sine ***
|
||||
def payne_hanek_reduction(d:UOp) -> tuple[UOp, UOp]:
|
||||
"""
|
||||
Performs Payne-Hanek Reduction: computes the remainder of `d` modulo pi/2 for the values `d` where
|
||||
39800.0 <= d <= +Inf
|
||||
Returns a tuple of `(r, q)`:
|
||||
- `r`[d.dtype] is the reminder value corresponding to `round_to_nearest(x % pi/2)`.
|
||||
- `q`[int32] is an integer, and q % 4 is corresponding to the quadrant of the original angle `d`.
|
||||
"""
|
||||
assert d.dtype.scalar() in TRANSCENDENTAL_DTYPES
|
||||
# https://stackoverflow.com/questions/30463616/payne-hanek-algorithm-implementation-in-c/30465751#30465751
|
||||
# 190 bits of 2/pi for Payne-Hanek style argument reduction
|
||||
two_over_pi_f = [0x00000000, 0x28be60db, 0x9391054a, 0x7f09d5f4, 0x7d4d3770, 0x36d8a566, 0x4f10e410]
|
||||
|
||||
intermediate_dtype = dtypes.float32.vec(d.dtype.count) if d.dtype.base.scalar() == dtypes.float16 else d.dtype
|
||||
|
||||
f, e = frexp(d)
|
||||
ia = (f.cast(intermediate_dtype) * 4.294967296e9).cast(dtypes.uint64)
|
||||
# extract 96 relevant bits of 2/pi based on magnitude of argument
|
||||
i = shr(e.cast(dtypes.uint64), 5)
|
||||
e = e.cast(dtypes.int32) & 31
|
||||
offset = 32 - e
|
||||
|
||||
def _take(an:UOp, offset:int, count:int=0) -> UOp:
|
||||
"""an = two_over_pi_f[i+offset]"""
|
||||
if count+offset < len(two_over_pi_f) - 1:
|
||||
an = i.ne(count).where(_take(an, offset, count=count+1), an.const_like(two_over_pi_f[count+offset]))
|
||||
return an
|
||||
def _shl_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) * pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32)
|
||||
def _shr_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) // pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32)
|
||||
|
||||
a = [_take(UOp.const(dtypes.uint32.vec(d.dtype.count), 0), i) for i in range(4)]
|
||||
# (two_over_pi_f[Int(i) + n] << e) | (two_over_pi_f[Int(i) + n+1] >> (nbits - e))
|
||||
# Note: e >= 1 for all numbers d >= 1.0. assume e != 0
|
||||
hi = _shl_lazy(a[0], e) | _shr_lazy(a[1], offset)
|
||||
mi = _shl_lazy(a[1], e) | _shr_lazy(a[2], offset)
|
||||
lo = _shl_lazy(a[2], e) | _shr_lazy(a[3], offset)
|
||||
|
||||
def _hp_mul(x:UOp, y:UOp) -> UOp: return x.cast(dtypes.uint64) * y.cast(dtypes.uint64)
|
||||
# compute x * 2/pi
|
||||
p = shl(_hp_mul(ia, hi), 32) + _hp_mul(ia, mi) + shr(_hp_mul(ia, lo), 32)
|
||||
|
||||
# round quotient to nearest
|
||||
q = shr(p, 62).cast(dtypes.int32)
|
||||
p = p & 0x3fffffffffffffff
|
||||
r = (p.cast(intermediate_dtype) * (3.4061215800865545e-19)).cast(d.dtype)
|
||||
|
||||
# if fraction >= 0.5, r -= pi/2, q += 1
|
||||
return (f<0.5).where(r, r - math.pi/2), (f<0.5).where(q, q + 1)
|
||||
|
||||
def cody_waite_reduction(d:UOp) -> tuple[UOp, UOp]:
|
||||
"""
|
||||
Performs Cody-Waite Reduction: computes the reminder of `d` modulo pi/2 for the values `d` where
|
||||
0 <= abs(d) <= 39800.0
|
||||
Returns a tuple of `(r, q)`, where the output format is the same as that of `payne_hanek_reduction`.
|
||||
"""
|
||||
def _reduce_d(x:UOp, q:UOp):
|
||||
# https://github.com/shibatch/sleef/blob/4e08851f59fc2b545f9c393c6a23dfd311a26308/src/libm/sleefdp.c#L789-L823
|
||||
if x.dtype.scalar() == dtypes.float64:
|
||||
# https://github.com/shibatch/sleef/blob/f6d8a841fbfddd26ce712834d4da220cd76048fb/src/common/misc.h#L77
|
||||
PI_A, PI_B, PI_C, PI_D = 3.1415926218032836914, 3.1786509424591713469e-08, 1.2246467864107188502e-16, 1.2736634327021899816e-24
|
||||
d = qdh * -PI_A + x
|
||||
d = q * -PI_A + d
|
||||
d = qdh * -PI_B + d
|
||||
d = q * -PI_B + d
|
||||
d = qdh * -PI_C + d
|
||||
d = q * -PI_C + d
|
||||
d = (qdh + q) * -PI_D + d
|
||||
elif x.dtype.scalar() == dtypes.float16:
|
||||
# [FIXME] when reducing `d`, FP16 needs FP32 precision to achieve 1.0 ULP precision.
|
||||
d = _reduce_d(x.cast(dtypes.float32), q.cast(dtypes.float32)).cast(dtypes.float16)
|
||||
else:
|
||||
# https://github.com/shibatch/sleef/blob/4e08851f59fc2b545f9c393c6a23dfd311a26308/src/libm/sleefsp.c#L464-L503
|
||||
d = q * -3.1414794921875 + x
|
||||
d = q * -0.00011315941810607910156 + d
|
||||
d = q * -1.9841872589410058936e-09 + d
|
||||
d = q * -1.2154201256553420762e-10 + d
|
||||
return d
|
||||
|
||||
m_1_pi = 0.318309886183790671537767526745028724
|
||||
qdh = (d * (m_1_pi / 2.0**24)).cast(dtypes.int64).cast(d.dtype) * (2.0**24)
|
||||
quadrant = rintk(d * m_1_pi -qdh) if d.dtype.base.scalar() == dtypes.float64 else rintk(d * m_1_pi)
|
||||
return _reduce_d(d, quadrant.cast(d.dtype)), quadrant.cast(dtypes.int32)
|
||||
|
||||
# *** approximate sine on small angle. ***
|
||||
def trig_poly(d:UOp, coeff32, coeff64): return d * (polyN(d*d, coeff64) if d.dtype.scalar() == dtypes.float64 else polyN(d*d, coeff32))
|
||||
# approximate sine on [-pi/2, pi/2]
|
||||
def sin_poly(d:UOp) -> UOp:
|
||||
return trig_poly(d, [2.6083159809786593541503e-06, -0.0001981069071916863322258, 0.00833307858556509017944336, -0.166666597127914428710938, 1.0],
|
||||
[-7.97255955009037868891952e-18, 2.81009972710863200091251e-15, -7.64712219118158833288484e-13, 1.60590430605664501629054e-10,
|
||||
-2.50521083763502045810755e-08, 2.75573192239198747630416e-06, -0.000198412698412696162806809, 0.00833333333333332974823815,
|
||||
-0.166666666666666657414808, 1.0])
|
||||
|
||||
def _ifand(q:UOp, n:int): return (q & n).ne(0)
|
||||
|
||||
def sin_poly_small(d:UOp, q:UOp) -> UOp:
|
||||
r = sin_poly(d)
|
||||
return r * _ifand(q, 1).where(r.const_like(-1), r.const_like(1))
|
||||
|
||||
def sin_poly_large(d:UOp, q:UOp) -> UOp:
|
||||
r = sin_poly(d + _ifand(q, 1).where(d.const_like(math.pi / 2), d.const_like(0)))
|
||||
return r * _ifand(q, 2).where(r.const_like(-1), r.const_like(1))
|
||||
|
||||
# *** toplevel functions for xsin/xlog2/xexp2 ***
|
||||
|
||||
def xsin(d:UOp, fast:bool=False, switch_over:float=30.0) -> UOp:
|
||||
"""
|
||||
Implements a 1.0 ULP approximation for Ops.SIN.
|
||||
- fast=True assumes x <= switch_over.
|
||||
- switch_over is the threshold for switching to payne_hanek_reduction.
|
||||
"""
|
||||
assert d.dtype.scalar() in TRANSCENDENTAL_DTYPES
|
||||
# mask +-inf/nan as zero
|
||||
x = _lazy_map_numbers(d, d.const_like(0.0), d.const_like(0.0), d.const_like(0.0), d)
|
||||
# x_sign = sign(x)
|
||||
x_sign = x.ne(0).where((x<0).where(x.const_like(-1), x.const_like(1)), x.const_like(0))
|
||||
x_abs = x * x_sign
|
||||
r, q = (cody_waite_reduction if fast else payne_hanek_reduction)(x_abs)
|
||||
if fast: result = sin_poly_small(r, q)
|
||||
else:
|
||||
# Payne Hanek Reduction assumes abs(x) >= pi/4, so for smaller values, use cody_waite_reduction.
|
||||
r_small, q_small = cody_waite_reduction(x_abs)
|
||||
result = (x_abs<switch_over).where(sin_poly_small(r_small, q_small), sin_poly_large(r, q))
|
||||
# adjusts the sign for abs(x)
|
||||
result = result * x_sign
|
||||
# sin(Inf) = NaN, sin(-Inf) = NaN, sin(NaN) = NaN
|
||||
return _lazy_map_numbers(d, d.const_like(math.nan), d.const_like(math.nan), d.const_like(math.nan), result)
|
||||
|
||||
def xexp2(d:UOp) -> UOp:
|
||||
"""
|
||||
Implements a 1.0 ULP approximation for Ops.EXP2
|
||||
- Paper: https://arxiv.org/pdf/2001.09258
|
||||
"""
|
||||
assert d.dtype.scalar() in TRANSCENDENTAL_DTYPES
|
||||
# mask +=inf/nan as zero.
|
||||
x = _lazy_map_numbers(d, d.const_like(0.0), d.const_like(0.0), d.const_like(0.0), d)
|
||||
q = rintk(x)
|
||||
# s = d - round(d)
|
||||
s = x - q.cast(x.dtype)
|
||||
# a polynomial approximation with 13 non-zero terms in the range of [−(log 2)/2,(log 2)/2].
|
||||
if d.dtype.scalar() == dtypes.float64:
|
||||
u = polyN(s, [0.4434359082926529454e-9, 0.7073164598085707425e-8, 0.1017819260921760451e-6, 0.1321543872511327615e-5, 0.1525273353517584730e-4,
|
||||
0.1540353045101147808e-3, 0.1333355814670499073e-2, 0.9618129107597600536e-2, 0.5550410866482046596e-1, 0.2402265069591012214e+0,
|
||||
0.6931471805599452862e+0, 0.1000000000000000000e+1])
|
||||
else: u = polyN(s, [0.1535920892e-3, 0.1339262701e-2, 0.9618384764e-2, 0.5550347269e-1, 0.2402264476e+0, 0.6931471825e+0, 1.0])
|
||||
u = ldexp2k(u, q) # u*2^q
|
||||
upper, lower = {dtypes.float64: (1024, -2000), dtypes.float32: (128, -150), dtypes.float16: (23, -22)}[d.dtype.scalar()]
|
||||
# Replace x >= upper with +inf
|
||||
u = (d >= upper).where(d.const_like(math.inf), u)
|
||||
# Replace x < lower with zero.
|
||||
u = (d<lower).where(d.const_like(0.0), u)
|
||||
# exp2(NaN) = NaN
|
||||
return d.ne(d).where(d.const_like(math.nan), u)
|
||||
|
||||
def xlog2(d:UOp) -> UOp:
|
||||
"""
|
||||
Implements a 1.0 ULP approximation for Ops.LOG2
|
||||
Paper: https://arxiv.org/pdf/2001.09258 5.5
|
||||
"""
|
||||
assert d.dtype.scalar() in TRANSCENDENTAL_DTYPES
|
||||
# float16 uses 2^10 for denormal scaling (2^64 overflows), float32/64 use 2^64
|
||||
denormal_exp = 10 if d.dtype.scalar() == dtypes.float16 else 64
|
||||
FLT_MIN = d.const_like({dtypes.float16: 6.1e-5, dtypes.float32: 1e-4, dtypes.float64: 1e-4}[d.dtype.scalar()])
|
||||
is_denormal = d<FLT_MIN
|
||||
a = is_denormal.where(d * (2 ** denormal_exp), d)
|
||||
|
||||
e = ilogb2k(a * (1.0 / 0.75)).cast(a.dtype)
|
||||
m = ldexp3k(a, -e)
|
||||
e = is_denormal.where(e - denormal_exp, e)
|
||||
|
||||
x = (m - 1.0) / (m + 1.0)
|
||||
x2 = x * x
|
||||
if d.dtype.scalar() == dtypes.float64:
|
||||
t = polyN(x2, [0.2211941750456081490e+0, 0.2200768693152277689e+0, 0.2623708057488514656e+0, 0.3205977477944495502e+0,
|
||||
0.4121985945485324709e+0, 0.5770780162997058982e+0, 0.96179669392608091449])
|
||||
r = t * (x * x2) + e + x * 2.885390081777926774
|
||||
else:
|
||||
t = polyN(x2, [0.4374550283e+0, 0.5764790177e+0, 0.9618012905120])
|
||||
# s_lo term (x*3.27e-08) only for float32 - underflows in float16
|
||||
r = t * (x * x2) + e + x * 2.8853900432586669922 + (x * 3.2734474483568488616e-08 if d.dtype.scalar() == dtypes.float32 else 0)
|
||||
|
||||
# log2(Inf) = Inf
|
||||
r = d.ne(math.inf).where(r, r.const_like(math.inf))
|
||||
# log2(0) = -Inf (handle both +0.0 and -0.0)
|
||||
r = d.ne(0.0).where(r, r.const_like(-math.inf))
|
||||
# log2(x) = NaN for x < 0
|
||||
r = (d<-0.0).where(r.const_like(math.nan), r)
|
||||
# log2(NaN) = NaN
|
||||
r = d.ne(d).where(r.const_like(math.nan), r)
|
||||
# log2(-0.0) = -Inf. In certain devices like PTX, x == -0.0 won't be true. so making reciprocal.
|
||||
return d.reciprocal().ne(-math.inf).where(r, r.const_like(-math.inf))
|
||||
|
||||
def xpow(base:UOp, exponent:UOp) -> UOp:
|
||||
# start with b ** e = exp2(e * log2(b))
|
||||
ret = (base < 0).where(-base, base).log2().mul(exponent).exp2()
|
||||
# negative base: nan for non-integer exponent, negate for odd integer exponent
|
||||
non_int = exponent != exponent.cast(dtypes.int32).cast(exponent.dtype)
|
||||
is_odd = (exponent < 0).where(-exponent, exponent).cast(dtypes.int32).mod(2).cast(dtypes.bool)
|
||||
neg_base = non_int.where(ret.const_like(math.nan), is_odd.where(-ret, ret))
|
||||
# fix 0 ** 0 = 1
|
||||
return (base.eq(0) & exponent.eq(0)).where(ret.const_like(1), (base < 0).where(neg_base, ret))
|
||||
|
||||
# *** integer division ***
|
||||
|
||||
@functools.lru_cache(None)
|
||||
def magicgu(vmax:int, d:int) -> tuple[int,int]:
|
||||
# calculate m,s such that x//d == (x*m) >> s for all 0 <= x <= vmax, d>0; adapted from Hacker's Delight, Chapter 10
|
||||
nc = (vmax+1)//(d) * d - 1
|
||||
nbits = vmax.bit_length()
|
||||
for s in range(0, 2*nbits + 1):
|
||||
if 2**s > nc*(d - 1 - (2**s - 1) % d):
|
||||
m = (2**s + d - 1 - (2**s - 1) % d)//d
|
||||
return m, s
|
||||
assert False
|
||||
|
||||
def fast_idiv(ren: Renderer, x: UOp, d: int, dont_cast=False) -> UOp|None:
|
||||
from tinygrad.renderer.cstyle import MetalRenderer
|
||||
# NOTE: disable for METAL due to compiler bug. keccak with -O0 works but not with optimization
|
||||
if isinstance(ren, MetalRenderer): return None
|
||||
# If d is a power of two this is not valid for signed ints!
|
||||
is_unsigned = x.vmin>=0 or x.dtype in dtypes.uints
|
||||
assert d>0, "Sign should have been taken out of divisor"
|
||||
vmin,vmax = max(x.vmin, x.dtype.min), min(x.vmax, x.dtype.max)
|
||||
if vmin > -d and vmax < d: return x.const_like(0)
|
||||
m,s = magicgu(max(vmax, abs(vmin)), d)
|
||||
if m*vmin >= x.dtype.min and m*vmax <= x.dtype.max:
|
||||
return ((x*m) >> s) if is_unsigned else ((x*m) >> s) + (x<0).where(x.ufix(1), 0)
|
||||
# before we try casting to a larger dtype (slow), we see if there are powers of two in d we can shift to make x smaller
|
||||
# use explicit Ops.CDIV (trunc) since the recursion assumes trunc semantics throughout
|
||||
if (largest_factor_of_two_in_d := (d & -d)) > 1:
|
||||
if (ret:=fast_idiv(ren, x.alu(Ops.CDIV, x.const_like(largest_factor_of_two_in_d)),
|
||||
d//largest_factor_of_two_in_d, dont_cast=True)) is not None: return ret
|
||||
if dont_cast: return None
|
||||
# promo_lattice needs to return an unsigned type if the type is unsigned
|
||||
if dtypes.is_int(next_dtype := promo_lattice[x.dtype.scalar()][-1]) and next_dtype in ren.supported_dtypes():
|
||||
if m*vmin >= next_dtype.min and m*vmax <= next_dtype.max:
|
||||
return ((x.cast(next_dtype)*m) >> s).cast(x.dtype) if is_unsigned else ((x.cast(next_dtype)*m) >> s).cast(x.dtype) + (x<0).where(x.ufix(1), 0)
|
||||
return None
|
||||
|
||||
# ***** threefry *****
|
||||
|
||||
def threefry2x32(x: UOp, key: UOp):
|
||||
# split x and key from uint64 to two uint32
|
||||
x0, x1 = (x & 0xffffffff).cast(dtypes.uint32), ((x // 2**32) & 0xffffffff).cast(dtypes.uint32)
|
||||
key0, key1 = (key & 0xffffffff).cast(dtypes.uint32), ((key // 2**32) & 0xffffffff).cast(dtypes.uint32)
|
||||
|
||||
rotations = [[13, 15, 26, 6], [17, 29, 16, 24]]
|
||||
ks = [key1, key0 ^ key1 ^ 0x1BD11BDA, key0]
|
||||
xr:list[UOp] = [x0 + ks[-1], x1 + ks[0]]
|
||||
for i in range(5):
|
||||
for r in rotations[i % 2]: xr[0], xr[1] = (x0 := xr[0] + xr[1]), x0 ^ ((xr[1] * 2**r) + (xr[1] // 2**(32 - r)))
|
||||
xr = [(xr[0] + ks[i % 3]), (xr[1] + ks[(i + 1) % 3] + i + 1)]
|
||||
|
||||
return xr[1].cast(dtypes.uint64) * 2**32 | xr[0].cast(dtypes.uint64)
|
||||
|
||||
# ***** long as 2 ints *****
|
||||
|
||||
l2i_dt = {dtypes.long: dtypes.int, dtypes.ulong: dtypes.uint}
|
||||
def unpack32(v:UOp) -> tuple[UOp, UOp]: return v.bitcast(dtypes.uint) & 0xFFFF, shr(v.bitcast(dtypes.uint), 16)
|
||||
def reindex(idx:UOp, off:int, mul=2) -> UOp: return idx.replace(src=(idx.src[0], idx.src[1]*mul+off, *idx.src[2:]))
|
||||
|
||||
# 4.3.1 is the relevant section in TAOCP
|
||||
def l2i(op: Ops, dt: DType, *uops:UOp):
|
||||
zero = UOp.const(dt, 0)
|
||||
if len(uops) == 2: a0, a1 = uops
|
||||
elif len(uops) == 4: a0, a1, b0, b1 = uops
|
||||
match op:
|
||||
case Ops.NEG: return l2i(Ops.SUB, dt, zero, zero, *uops)
|
||||
case Ops.CAST if dt in (dtypes.long, dtypes.ulong) and uops[0].dtype not in dtypes.floats:
|
||||
return uops[0].cast(l2i_dt[dt]), (uops[0] < 0).where(UOp.const(l2i_dt[dt], -1), UOp.const(l2i_dt[dt], 0))
|
||||
case Ops.CAST if dt in (dtypes.long, dtypes.ulong):
|
||||
return (lo:=uops[0].cast(l2i_dt[dt])), (uops[0] / 2**32).cast(l2i_dt[dt]) - ((uops[0] < 0) & lo.ne(0)).cast(l2i_dt[dt])
|
||||
case Ops.CAST if dt in dtypes.floats:
|
||||
small = (a1.eq(0) & (a0 >= 0)) | (a1.eq(-1) & (a0 < 0))
|
||||
return small.where(a0.cast(dt), ((a1.cast(dtypes.float32) * (2**32)) + a0.bitcast(dtypes.uint).cast(dtypes.float32)).cast(dt))
|
||||
case Ops.CAST: return a0.bitcast(dtypes.uint).cast(dt)
|
||||
case Ops.BITCAST: return a0.bitcast(dt), a1.bitcast(dt)
|
||||
case Ops.SHL:
|
||||
lo, hi = shl(a0, b0_mod:=b0 & 31), shl(a1, b0_mod) | shr(shr(a0, 1), 31 - b0_mod)
|
||||
return (b0 >= 32).where(zero, lo), (b0 >= 32).where(lo, hi)
|
||||
case Ops.SHR:
|
||||
lo, hi = shr(a0, b0_mod:=b0 & 31) | shl(shl(a1, 1), 31 - b0_mod), shr(a1, b0_mod)
|
||||
return (b0 >= 32).where(hi, lo), (b0 >= 32).where(zero, hi)
|
||||
case Ops.ADD: return (low:=a0+b0), (a1 + b1).replace(dtype=dt) + (low.bitcast(dtypes.uint) < a0.bitcast(dtypes.uint)).cast(dt)
|
||||
case Ops.SUB: return a0 - b0, a1 - b1 - (a0.bitcast(dtypes.uint) < b0.bitcast(dtypes.uint)).cast(dt)
|
||||
case Ops.MUL:
|
||||
(a00, a01), (b00, b01) = unpack32(a0), unpack32(b0)
|
||||
mid = l2i(Ops.ADD, dt, shl(a00*b01, 16).bitcast(dt), shr(a00*b01, 16).bitcast(dt), shl(a01*b00, 16).bitcast(dt), shr(a01*b00, 16).bitcast(dt))
|
||||
return l2i(Ops.ADD, dt, *mid, (a00*b00).bitcast(dt), (a01*b01).bitcast(dt) + a0*b1 + a1*b0)
|
||||
case Ops.CDIV | Ops.CMOD:
|
||||
# TAOCP Algorithm 4.3.1D could be faster here, but must be parameterized over the width of b
|
||||
if dt == dtypes.int:
|
||||
ua0, ua1, ub0, ub1 = a0.bitcast(dtypes.uint), a1.bitcast(dtypes.uint), b0.bitcast(dtypes.uint), b1.bitcast(dtypes.uint)
|
||||
a0, a1 = (a_neg:=a1 < zero).where((n:=l2i(Ops.NEG, dtypes.uint, ua0, ua1))[0], ua0), a_neg.where(n[1], ua1)
|
||||
b0, b1 = (b_neg:=b1 < zero).where((n:=l2i(Ops.NEG, dtypes.uint, ub0, ub1))[0], ub0), b_neg.where(n[1], ub1)
|
||||
q, r = (z:=UOp.const(dtypes.uint, 0), z), (z, z)
|
||||
for i in range(63, -1, -1):
|
||||
r = l2i(Ops.SHL, dtypes.uint, *r, UOp.const(dtypes.uint, 1), z)
|
||||
r = (r[0] | l2i(Ops.SHR, dtypes.uint, a0, a1, UOp.const(dtypes.uint, i), z)[0] & 1), r[1]
|
||||
cond = l2i(Ops.CMPLT, dtypes.uint, *r, b0, b1).logical_not()
|
||||
diff = l2i(Ops.SUB, dtypes.uint, *r, b0, b1)
|
||||
q = ((q[0] | shl(cond.cast(dtypes.uint), i % 32), q[1]) if i < 32 else (q[0], q[1] | shl(cond.cast(dtypes.uint), i % 32)))
|
||||
r = l2i(Ops.WHERE, dtypes.uint, cond, *diff, *r)
|
||||
if dt == dtypes.int:
|
||||
(nq0, nq1), (nr0, nr1) = l2i(Ops.BITCAST, dt, *l2i(Ops.NEG, dtypes.uint, *q)), l2i(Ops.BITCAST, dt, *l2i(Ops.NEG, dtypes.uint, *r))
|
||||
(q0, q1), (r0, r1) = l2i(Ops.BITCAST, dt, *q), l2i(Ops.BITCAST, dt, *r)
|
||||
return (a_neg.where(nr0, r0), a_neg.where(nr1, r1)) if op == Ops.CMOD else ((a_neg^b_neg).where(nq0, q0), (a_neg^b_neg).where(nq1, q1))
|
||||
return (r[0].bitcast(dt), r[1].bitcast(dt)) if op == Ops.CMOD else (q[0].bitcast(dt), q[1].bitcast(dt))
|
||||
case Ops.CMPLT: return (a1 < b1) | ((a1.eq(b1)) & (a0.bitcast(dtypes.uint) < b0.bitcast(dtypes.uint)))
|
||||
case Ops.CMPEQ: return a0.eq(b0) & a1.eq(b1)
|
||||
case Ops.CMPNE: return a0.ne(b0) | a1.ne(b1)
|
||||
case Ops.XOR | Ops.OR | Ops.AND: return UOp(op, dt, src=(a0, b0)), UOp(op, dt, src=(a1, b1))
|
||||
case Ops.WHERE: return uops[0].where(uops[1], uops[3]), uops[0].where(uops[2], uops[4])
|
||||
case Ops.MAX: return l2i(Ops.WHERE, dt, l2i(Ops.CMPLT, dt, *uops), b0, b1, a0, a1)
|
||||
case _: raise NotImplementedError(f"long decomposition of {op} unsupported")
|
||||
|
||||
# ***** floats *****
|
||||
f2f_dt = { f:getattr(dtypes, f"uint{f.bitsize}") for f in dtypes.floats }
|
||||
|
||||
def rne(v: UOp, s) -> UOp: return shr(v, s) + ((shr(v, s - 1) & 1) & ((v & ((1 << (s - 1)) - 1)).ne(0).cast(v.dtype) | (shr(v, s) & 1)))
|
||||
|
||||
def f2f(v, fr:DType, to:DType, sat=True):
|
||||
fs, fb, (fe, fm), ts, tb, (te, tm) = fr.bitsize, exponent_bias(fr), dtypes.finfo(fr), to.bitsize, exponent_bias(to), dtypes.finfo(to)
|
||||
# NB: denormals are zero!
|
||||
if fe <= te and fm < tm:
|
||||
sign, nosign = shl((v & shl(1, fs-1)).cast(f2f_dt[to]), ts - fs), (v & (shl(1, fs-1) - 1)).cast(f2f_dt[to])
|
||||
exp, norm = shr(nosign, fm), shl(nosign, tm - fm) + shl(tb - fb, tm)
|
||||
nan = shl(nosign, tm - fm) | shl((shl(1, te) - 1), tm)
|
||||
if fr in dtypes.fp8_fnuz:
|
||||
fnuz_nan = sign.ne(0) & nosign.eq(0)
|
||||
qnan = shl(shl(1, te) - 1, tm) | shl(1, tm - 1)
|
||||
return fnuz_nan.where(qnan, sign | exp.eq(0).where(0, norm)).bitcast(to)
|
||||
# fp8e4m3 has only one nan
|
||||
is_nan = (nosign.eq(shl(1, fm + fe) - 1) if fr == dtypes.fp8e4m3 else exp.eq(shl(1, fe) - 1))
|
||||
return (sign | exp.eq(0).where(0, is_nan.where(nan, norm))).bitcast(to)
|
||||
elif fe >= te and fm > tm:
|
||||
v = f2f_clamp(v.bitcast(fr), to, sat).bitcast(f2f_dt[fr])
|
||||
sign, nosign = shr(v, fs - ts) & shl(1, ts - 1), v & (shl(1, fs - 1) - 1)
|
||||
norm = (rne(nosign, fm - tm) - shl(fb - tb, tm)).cast(f2f_dt[to])
|
||||
underflow = (shr(v, fm) & (shl(1, fe) - 1)) < (1 + fb - tb)
|
||||
nan_mantissa = (shl(1, tm) - 1) if to == dtypes.fp8e4m3 else (shr(nosign, fm - tm) & (shl(1, tm) - 1))
|
||||
nan = (sign | nan_mantissa | shl(shl(1, te) - 1, tm)).cast(f2f_dt[to])
|
||||
is_nan = (shr(v, fm) & (shl(1, fe) - 1)).eq(shl(1, fe) - 1)
|
||||
if to in dtypes.fp8_fnuz: return is_nan.where(shl(1, ts - 1), underflow.where(0, sign.cast(f2f_dt[to]) | norm))
|
||||
return is_nan.where(nan, sign.cast(f2f_dt[to]) | underflow.where(0, norm))
|
||||
else: raise NotImplementedError(f"unsupported decomp {fr} -> {to}")
|
||||
|
||||
def f2f_clamp(val:UOp, dt:DType, sat=True) -> UOp:
|
||||
e, m = dtypes.finfo(dt)
|
||||
if dt in dtypes.fp8_fnuz: max_exp, max_man = (1 << e) - 1, (1 << m) - 1
|
||||
else: max_exp, max_man = ((1 << e) - 1, (1 << m) - 2) if dt == dtypes.fp8e4m3 else ((1 << e) - 2, (1 << m) - 1)
|
||||
mx = val.const_like(2.0**(max_exp - exponent_bias(dt)) * (1.0 + max_man / (1 << m)))
|
||||
sat = mx if dt in dtypes.fp8s and sat else val.const_like(float('inf'))
|
||||
# FIXME: CMPLT of nan is undefined
|
||||
return val.ne(val).where(val, (val < -mx).where(-sat, (mx < val).where(sat, val)))
|
||||
|
||||
def f2f_load(x: UOp, fr:DType, to:DType) -> UOp:
|
||||
if (n:=x.dtype.count) == 1: return f2f(x.replace(dtype=f2f_dt[fr]), fr, to)
|
||||
return UOp.vectorize(*(f2f(x.replace(dtype=f2f_dt[fr], src=(reindex(x.src[0].src[0], i, 1),)), fr, to) for i in range(n)))
|
||||
|
||||
def f2f_store(st, idx, val, fr:DType, to:DType):
|
||||
if (n:=val.dtype.count) == 1: return st.replace(src=(idx, f2f(val.bitcast(f2f_dt[to]), to, fr)))
|
||||
return UOp.group(*(st.replace(src=(reindex(idx, i, 1), f2f(val.gep(i).bitcast(f2f_dt[to]), to, fr))) for i in range(n)))
|
||||
|
||||
# ***** decomposition patterns *****
|
||||
|
||||
@functools.cache
|
||||
def get_transcendental_patterns(ops:tuple[Ops, ...], force_transcendental:bool) -> PatternMatcher:
|
||||
pat: list[tuple[UPat, Callable]] = []
|
||||
for op,f in ((Ops.EXP2, xexp2), (Ops.LOG2, xlog2), (Ops.SIN, xsin)):
|
||||
if op not in ops or force_transcendental:
|
||||
pat += [(UPat(op, dtype=TRANSCENDENTAL_DTYPES, src=(UPat.var("d"),)), f),
|
||||
(UPat(op, dtype=tuple(dt for dt in dtypes.floats if dt not in TRANSCENDENTAL_DTYPES), src=(UPat.var("d"),), name="x"),
|
||||
lambda x,d: d.cast(dtypes.float32).alu(x.op).cast(x.dtype))]
|
||||
# rewrite SQRT to xpow 0.5
|
||||
if Ops.SQRT not in ops or force_transcendental: pat.append((UPat(Ops.SQRT, src=UPat.var("d")), lambda d: xpow(d, d.const_like(0.5))))
|
||||
return PatternMatcher(pat)
|
||||
|
||||
def floordiv_to_idiv(a:UOp, b:UOp) -> UOp:
|
||||
if (a.vmin >= 0 and b.vmin > 0) or (a.vmax <= 0 and b.vmax < 0): return a.alu(Ops.CDIV, b)
|
||||
return a.alu(Ops.CDIV, b) - (a.alu(Ops.CMOD, b).ne(0) & (a<0).ne(b<0)).cast(a.dtype)
|
||||
|
||||
def floormod_to_mod(a:UOp, b:UOp) -> UOp:
|
||||
if (a.vmin >= 0 and b.vmin > 0) or (a.vmax <= 0 and b.vmax < 0): return a.alu(Ops.CMOD, b)
|
||||
r = a.alu(Ops.CMOD, b)
|
||||
# use where instead of mul to avoid being fused into MULACC (which int64 long-decomp doesn't handle)
|
||||
return r + (r.ne(0) & (a<0).ne(b<0)).where(b, b.const_like(0))
|
||||
|
||||
powers_of_two: dict[int, int] = {2**i:i for i in range(64)}
|
||||
@functools.cache
|
||||
def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> PatternMatcher:
|
||||
pat: list[tuple[UPat, Callable]] = [(UPat.var("a")//UPat.var("b"), floordiv_to_idiv)]
|
||||
# FLOORMOD by 2**y -> x & (2**y-1) (correct floor mod for any sign in two's complement); fires before floormod_to_mod
|
||||
if Ops.AND in ops: pat.append((UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.arg-1) if c.arg in powers_of_two else None))
|
||||
pat.append((UPat.var("a")%UPat.var("b"), floormod_to_mod))
|
||||
# no real hardware supports THREEFRY, but NullRenderer does
|
||||
if Ops.THREEFRY not in ops: pat.append((UPat(Ops.THREEFRY, dtype=dtypes.uint64, src=(UPat.var("x"), UPat.var("key"))), threefry2x32))
|
||||
# MAX can be rewritten as CMPLT + WHERE (max function is annoying on many cstyle backends)
|
||||
if Ops.MAX not in ops and Ops.CMPLT in ops: pat.append((UPat(Ops.MAX, name="m"), lambda m: (m.src[0] < m.src[1]).where(m.src[1], m.src[0])))
|
||||
if Ops.OR in ops: pat += [(UPat.var("x", dtypes.bool).logical_not()&UPat.var("y", dtypes.bool).logical_not(),
|
||||
lambda x,y: (x | y).logical_not())]
|
||||
# rewrite MUL/CDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y)
|
||||
if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.arg, 0)) else None)]
|
||||
if Ops.SHR in ops:
|
||||
# uint CDIV by 2**v -> x >> v (FLOORDIV is lowered to CDIV by the rule above before reaching here)
|
||||
pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.uints), UPat.cvar("c"))),
|
||||
lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]
|
||||
# signed CDIV (trunc) by 2**v -> (x + (x<0 ? c-1 : 0)) >> v
|
||||
pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("c"))),
|
||||
lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(c-1, 0)) >> v
|
||||
if (v:=powers_of_two.get(c.arg, 0)) else None)]
|
||||
if not disable_fast_idiv:
|
||||
# fast_idiv handles non-pow2: only fire on non-negative inputs (signed magic-mul is unreliable for x<0)
|
||||
pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("d"))),
|
||||
lambda ctx, x, d: fast_idiv(ctx, x, d.arg) if x.vmin >= 0 or x.dtype in dtypes.uints else None)]
|
||||
# rewrite raw CMOD -> x - d*CDIV(x,d) so fast_idiv can pick up the CDIV. only on non-negative inputs;
|
||||
# avoids disturbing floormod_to_mod's general-path output (which uses a trunc Ops.CMOD as an implementation detail)
|
||||
pat += [(UPat(Ops.CMOD, src=(UPat.var("x", dtypes.ints), UPat.var("d"))),
|
||||
lambda x, d: x - d * x.alu(Ops.CDIV, d) if x.vmin >= 0 or x.dtype in dtypes.uints else None)]
|
||||
if Ops.NEG in ops:
|
||||
pat += [(UPat.var('x')*-1, lambda ctx,x: x.alu(Ops.NEG))]
|
||||
if Ops.SUB in ops: pat += [(UPat.var('x')+UPat.var('y').alu(Ops.NEG), lambda ctx,x,y: x.alu(Ops.SUB, y))]
|
||||
if Ops.CMPLT in ops:
|
||||
# These are late rewrites because simplex expects equalities to be a certain format
|
||||
pat += [
|
||||
((UPat.var("x", dtypes.sints) < UPat.cvar("c", dtypes.sints)).logical_not(), lambda x,c: c-1<x),
|
||||
((UPat.cvar("c", dtypes.sints) < UPat.var("x", dtypes.sints)).logical_not(), lambda x,c: x<c+1),
|
||||
(UPat.var("x", dtypes.sints)*-1 < UPat.var("y", dtypes.sints)*UPat.cvar("c"), lambda x,y,c: y*(-c)<x),
|
||||
(UPat.var("x", dtypes.sints)*-1 < UPat.cvar("c"), lambda x,c:-c<x),
|
||||
((UPat.cvar("c1")<UPat.var("x", dtypes.sints)) & (UPat.var("x", dtypes.sints)<UPat.cvar("c2")),
|
||||
lambda x,c1,c2: x.eq(c1+1) if c1.arg+1==c2.arg-1 else None), # (c-1)<x & x<(c+1) -> x==c
|
||||
]
|
||||
if Ops.CMPEQ in ops: pat += [(UPat.var('x').ne(UPat.var('y')).logical_not(), lambda x,y: x.alu(Ops.CMPEQ, y))]
|
||||
if Ops.MULACC in ops:
|
||||
pat += [(UPat.var('a')*UPat.var('b')+UPat.var('c'), lambda a,b,c: a.alu(Ops.MULACC, b, c))]
|
||||
# also fuse (x << n) + c → MULACC(x, 2^n, c) since MUL→SHL may run first
|
||||
if Ops.SHL in ops: pat += [(UPat.var('x').alu(Ops.SHL, UPat.cvar('n'))+UPat.var('c'), lambda x,n,c: x.alu(Ops.MULACC, x.const_like(1<<n.arg), c))]
|
||||
# some backends emit FDIV for RECIP, in that case: a*(1/b) -> a/b
|
||||
if Ops.FDIV in ops:
|
||||
pat += [(UPat.var("x").reciprocal(), lambda x: x.const_like(1).alu(Ops.FDIV, x))]
|
||||
pat += [(UPat.var("a", dtypes.floats) * UPat.const(dtypes.floats, 1).alu(Ops.FDIV, UPat.var("b")), lambda a,b: a.alu(Ops.FDIV, b))]
|
||||
return PatternMatcher(pat)
|
||||
|
||||
pm_long_decomp = PatternMatcher([
|
||||
(UPat((*GroupOp.Defines, Ops.INDEX), name="x"), lambda x:
|
||||
x.replace(dtype=l2i_dt[x.dtype.base].ptr(x.dtype.size * 2)) if hasattr(x.dtype, 'size') and x.dtype.base in l2i_dt else None),
|
||||
(UPat(Ops.INDEX, tuple(l2i_dt.keys()), name='x'), lambda x: reindex(x, x.tag).replace(dtype=l2i_dt[x.dtype])),
|
||||
(UPat(Ops.STORE, src=(UPat.var('idx'), UPat.var('val', tuple(l2i_dt.keys()))), name='st'), lambda st,idx,val:
|
||||
st.replace(src=(reindex(idx, 0), val.rtag(0))).group(st.replace(src=(reindex(idx, 1), val.rtag(1)))) if val.tag is None else None),
|
||||
(UPat(GroupOp.Comparison, src=(UPat.var('a', tuple(l2i_dt.keys())), UPat.var('b', tuple(l2i_dt.keys()))), name="x"), lambda a,b,x:
|
||||
l2i(x.op, dt:=l2i_dt[a.dtype], a.rtag(0).cast(dt), a.rtag(1).cast(dt), b.rtag(0).cast(dt), b.rtag(1).cast(dt))),
|
||||
(UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a'),), name="x"), lambda a,x:
|
||||
l2i(x.op, x.dtype, a)[x.tag] if x.tag is not None and a.dtype not in l2i_dt else None),
|
||||
(UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda a,x:
|
||||
(a.rtag(0).cast(dt:=l2i_dt[a.dtype]).bitcast(xdt:=l2i_dt[x.dtype]), a.rtag(1).cast(dt).bitcast(xdt))[x.tag]),
|
||||
(UPat(Ops.CAST, src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda a,x:
|
||||
l2i(x.op, x.dtype, a.rtag(0).cast(dt:=l2i_dt[a.dtype]), a.rtag(1).cast(dt)) if x.dtype not in l2i_dt and a.tag is None else None),
|
||||
(UPat((*(GroupOp.ALU - GroupOp.Comparison), Ops.BITCAST), tuple(l2i_dt.keys()), name="x"), lambda x:
|
||||
l2i(x.op, l2i_dt[x.dtype], *flatten((a.rtag(0).cast(dt:=l2i_dt[x.src[-1].dtype]), a.rtag(1).cast(dt))
|
||||
if a.dtype in l2i_dt else (a,) for a in x.src))[x.tag] if x.tag is not None else None),
|
||||
(UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda x,idx: x.replace(dtype=l2i_dt[x.dtype],src=(reindex(idx, x.tag),))),
|
||||
(UPat(Ops.CONST, tuple(l2i_dt.keys()), name='x'), lambda x:
|
||||
UOp.const(dt:=l2i_dt[x.dtype], truncate[dt]((x.arg >> 32) if x.tag == 1 else (x.arg & 0xFFFFFFFF))))
|
||||
])
|
||||
|
||||
# float decomposition patterns - ctx is (fr, to) tuple
|
||||
pm_float_decomp = PatternMatcher([
|
||||
(UPat((*GroupOp.Defines, Ops.INDEX), name="x"), lambda ctx,x:
|
||||
x.replace(dtype=f2f_dt[ctx[0]].ptr(x.dtype.size), tag=ctx[0]) if x.dtype.base == ctx[0] else None),
|
||||
(UPat(Ops.LOAD, dtypes.floats, name="x"), lambda ctx,x: f2f_load(x, *ctx) if x.dtype.scalar() == ctx[0] else None),
|
||||
# bitcasted load should just replace load
|
||||
(UPat(Ops.BITCAST, src=(UPat(Ops.LOAD, name="ld"),), name="bc"), lambda ctx,bc,ld:
|
||||
ld.replace(dtype=f2f_dt[ctx[0]]).bitcast(bc.dtype) if ld.dtype == ctx[0] else None),
|
||||
# bitcast from
|
||||
(UPat(Ops.BITCAST, src=(UPat.var("x", dtypes.floats),), name="bc"), lambda ctx,bc,x:
|
||||
bc.replace(src=(f2f(x.bitcast(f2f_dt[ctx[1]]), ctx[1], ctx[0]),)) if x.dtype == ctx[1] and bc.dtype.bitsize == ctx[0].bitsize else None),
|
||||
# bitcast to
|
||||
(UPat(Ops.BITCAST, src=(UPat.var("x"),), name="bc"), lambda ctx,bc,x:
|
||||
f2f(x.bitcast(f2f_dt[ctx[0]]), ctx[0], ctx[1]) if bc.dtype == ctx[0] else None),
|
||||
(UPat(Ops.CAST, dtypes.floats, src=(UPat.var("val"),), name="x"), lambda ctx,x,val:
|
||||
f2f_clamp(val.cast(ctx[1]), ctx[0]) if x.dtype.scalar() == ctx[0] else None),
|
||||
(UPat(GroupOp.All-{Ops.BITCAST}, dtypes.floats, name="x"), lambda ctx,x:
|
||||
x.replace(dtype=ctx[1].vec(x.dtype.count), src=tuple(s.cast(ctx[1]) if s.dtype == ctx[0] else s for s in x.src))
|
||||
if x.dtype.scalar() == ctx[0] else None),
|
||||
(UPat(Ops.STORE, src=(UPat.var("idx"), UPat(Ops.BITCAST, dtypes.floats, name="val")), name='st'), lambda ctx,st,idx,val:
|
||||
st.replace(src=(idx, val.replace(dtype=f2f_dt[ctx[0]]))) if val.dtype == ctx[0] and idx.tag == ctx[0] else None),
|
||||
(UPat(Ops.STORE, src=(UPat.var("idx"), UPat.var("val", dtypes.floats)), name='st'), lambda ctx,st,idx,val:
|
||||
f2f_store(st, idx, val, *ctx) if val.dtype.scalar() == ctx[1] and (idx:=idx.src[0] if idx.op == Ops.CAST else idx).tag == ctx[0] else None),
|
||||
])
|
||||
|
||||
def do_dtype_decomps(sink:UOp, ctx:tuple[set[DType], Renderer]) -> UOp:
|
||||
def _should_emulate(dt): return dt in EMULATED_DTYPES.tolist(dtypes) or dt not in ctx[1].supported_dtypes()
|
||||
for fr in sorted(filter(_should_emulate, ctx[0])):
|
||||
to = dtypes.int if fr == dtypes.long else dtypes.half if not _should_emulate(dtypes.half) and fr in dtypes.fp8s else dtypes.float
|
||||
if DEBUG >= 2: print(f"emulating {fr} as {to}")
|
||||
sink = graph_rewrite(sink, pm_float_decomp if fr in dtypes.floats else pm_long_decomp, name=f"decomp {fr} -> {to}", ctx=(fr, to), bottom_up=True)
|
||||
ctx[0].clear()
|
||||
return sink
|
||||
|
||||
pm_dtype_decomps = PatternMatcher([
|
||||
# detect dtypes to decompose
|
||||
(UPat(GroupOp.All, (*dtypes.fp8s, dtypes.bfloat16, dtypes.half, dtypes.long, dtypes.ulong), name="x"), lambda x,ctx:
|
||||
ctx[0].add({dtypes.ulong:dtypes.long}.get(dt:=x.dtype.base.scalar(), dt))),
|
||||
# do the rewrites
|
||||
(UPat(Ops.SINK, name="sink"), do_dtype_decomps),
|
||||
])
|
||||
116
tinygrad_repo/tinygrad/uop/divandmod.py
Normal file
116
tinygrad_repo/tinygrad/uop/divandmod.py
Normal file
@@ -0,0 +1,116 @@
|
||||
import functools, itertools, math
|
||||
from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.helpers import floordiv, floormod, unwrap
|
||||
|
||||
# NOTE: this cache is only on index UOps
|
||||
@functools.cache
|
||||
def fold_divmod_general(d: UOp) -> UOp|None:
|
||||
x, y = d.src
|
||||
|
||||
# cancel_divmod: simple cancel div/mod case when the range of the numerator lies within a single denominator interval
|
||||
x_min, x_max, y_min, y_max = x.vmin, x.vmax, y.vmin, y.vmax
|
||||
assert isinstance(x_min, int) and isinstance(x_max, int) and isinstance(y_min, int) and isinstance(y_max, int)
|
||||
if y_min==y_max==0: raise ZeroDivisionError(f"{'Division' if d.op is Ops.FLOORDIV else 'Mod'} by zero trying to rewrite {x.alu(d.op, y)}")
|
||||
if y_min*y_max > 0 and (qv:=floordiv(x_min,y_min)) == floordiv(x_min,y_max) == floordiv(x_max,y_min) == floordiv(x_max,y_max):
|
||||
return x - qv*y if d.op is Ops.FLOORMOD else d.const_like(qv)
|
||||
|
||||
# split uops for the rest of the processing
|
||||
x_peeled, const = x.pop_const()
|
||||
uops_no_const = list(x_peeled.split_uop(Ops.ADD))
|
||||
|
||||
# ** Constant Denominator Rules **
|
||||
# these rules strictly require y to be a scalar constant > 0
|
||||
if y.op is Ops.CONST and (c := y.arg) > 0:
|
||||
# nested_div_mod: (x%(k*c))//c -> (x//c)%k (requires k>0), and (x%(k*c))%c -> x%c
|
||||
if x.op is Ops.FLOORMOD and (k := x.src[1].divides(c)) is not None:
|
||||
if d.op is Ops.FLOORMOD: return x.src[0] % y
|
||||
if k > 0: return x.src[0] // y % k
|
||||
|
||||
# remove_nested_mod in sum: (a%4 + b)%2 -> (a+b)%2
|
||||
if d.op is Ops.FLOORMOD:
|
||||
new_xs, changed = [], False
|
||||
for u in uops_no_const:
|
||||
if u.op is Ops.FLOORMOD and u.src[1].divides(c) is not None:
|
||||
u = u.src[0]
|
||||
changed = True
|
||||
new_xs.append(u)
|
||||
if changed: return (UOp.usum(*new_xs) + const) % y
|
||||
|
||||
# Shared decomposition for folding rules
|
||||
decomp = [(u.divides(f:=u.const_factor()),f) for u in uops_no_const]
|
||||
terms, factors = zip(*decomp)
|
||||
|
||||
# fold_binary_numerator: fold if expression has one non-constant term that takes on two values
|
||||
if len(terms)==1 and (v:=terms[0]).vmax-v.vmin == 1:
|
||||
y1 = (floormod if d.op is Ops.FLOORMOD else floordiv)(factors[0]*v.vmin+const, c)
|
||||
y2 = (floormod if d.op is Ops.FLOORMOD else floordiv)(factors[0]*v.vmax+const, c)
|
||||
return (y2-y1)*(v-v.vmin) + y1
|
||||
|
||||
# fold_divmod_congruence: fold if a is congruent to an expression whose range is between 0 and c
|
||||
# when f%c == c//2, abs(r) == abs(r-c) is a tie, try both signs since either may fit in one period
|
||||
rem_choices = [(r, r-c) if (r:=f%c)*2 == c else (min(r, r-c, key=abs),) for f in factors]
|
||||
for rems in itertools.product(*rem_choices):
|
||||
if (rem:=sum(r*v for r,v in zip(rems,terms))+const%c).vmin//c==rem.vmax//c:
|
||||
if d.op is Ops.FLOORMOD: return rem - rem.vmin//c*c
|
||||
return sum((f-r)//c * v for f,r,v in zip(factors,rems,terms)) + const//c + rem.vmin//c
|
||||
|
||||
# gcd_with_remainder: factor out common gcd from numerator
|
||||
if (g:=math.gcd(*factors, c)) > 1:
|
||||
new_x = unwrap(x_peeled.divides(g)).simplify() + (const//g)%(c//g)
|
||||
if new_x.vmin >= 0:
|
||||
if d.op is Ops.FLOORMOD: return new_x % (c//g) * g + const%g
|
||||
return new_x // (c//g) + const//c
|
||||
|
||||
# nest_by_factor: x//c -> (x//f)//(c//f), x%c -> (x//f%(c//f))*f + b where b=x%f
|
||||
# FLOORDIV identity holds for any sign of x; FLOORMOD reconstruction needs x.vmin>=0
|
||||
results = []
|
||||
for div in {abs(f) for u, f in zip(uops_no_const, factors) if u.op is not Ops.CONST and 1 < abs(f) < c and (c%f)==0}:
|
||||
if (newxs := fold_divmod_general(x//div)) is not None:
|
||||
if d.op is Ops.FLOORDIV:
|
||||
results.append((len(newxs.backward_slice), newxs // (c // div)))
|
||||
elif x.vmin >= 0 and newxs.vmin >= 0:
|
||||
b_parts = [f%div*t for f, t in zip(factors, terms) if f%div]
|
||||
if const % div: b_parts.append(x.const_like(const % div))
|
||||
b = UOp.usum(*b_parts) if b_parts else x.const_like(0)
|
||||
if 0 <= b.vmin and b.vmax < div:
|
||||
results.append((len((r:=(newxs % x.ufix(c//div))*div + b).backward_slice), r))
|
||||
if results: return min(results, key=lambda r: r[0])[1]
|
||||
|
||||
# ** Variable Denominator / Fallback Rules **
|
||||
# These rules apply to variables OR constants that failed the checks above.
|
||||
# Reconstruct all uops including const for these checks.
|
||||
all_uops = list(x.split_uop(Ops.ADD))
|
||||
|
||||
# divide_by_gcd: x//y -> (x//gcd)//(y//gcd)
|
||||
gcd = UOp.gcd(*all_uops, y).simplify()
|
||||
if not (gcd.op is Ops.CONST and gcd.arg==1):
|
||||
ret = unwrap(x.divide_exact(gcd)).alu(d.op, unwrap(y.divide_exact(gcd)))
|
||||
return ret*gcd if d.op is Ops.FLOORMOD else ret
|
||||
|
||||
# factor_remainder: (d*x+y)//d -> x+y//d
|
||||
if y.vmin<0 or x.vmin<0: return None
|
||||
quo, rem = [], []
|
||||
for u in all_uops:
|
||||
if (q:=u.divide_exact(y)) is not None: quo.append(q)
|
||||
elif y.op is Ops.CONST and (c:=u.const_factor())%y.arg!=c:
|
||||
rem.append(u.divides(c)*(c%y.arg))
|
||||
quo.append(u.divides(c)*(c//y.arg) if d.op is Ops.FLOORDIV else u.const_like(0))
|
||||
else: rem.append(u)
|
||||
|
||||
if not quo: return None
|
||||
new_x = sum(rem)+x.const_like(0)
|
||||
if new_x.vmin<0: return None
|
||||
return new_x%y if d.op is Ops.FLOORMOD else new_x//y+sum(quo)
|
||||
|
||||
div_and_mod_symbolic = PatternMatcher([
|
||||
# ** 1. Fast Inline Rules **
|
||||
# (x//c+a)//d -> (x+a*c)//(c*d) for c>0, d>0
|
||||
((UPat.var("x")//UPat.cvar("c") + UPat.cvar("a"))//UPat.cvar("d"), lambda x,c,a,d: (x+a*c)//(c*d) if c.vmin>0 and d.vmin>0 else None),
|
||||
# (x+c)//d -> (x+c%d)//d + c//d for d>0 (split out the multiple of d in the constant)
|
||||
((UPat.var("x", dtypes.weakint)+UPat.cvar("c"))//UPat.cvar("d"),
|
||||
lambda x,c,d: (x+c.arg%d.arg)//d + c.arg//d.arg if c.arg%d.arg!=c.arg and d.arg>0 else None),
|
||||
|
||||
# ** 2. Slow Rules **
|
||||
(UPat((Ops.FLOORDIV, Ops.FLOORMOD), dtypes.weakint, name="d"), lambda d: fold_divmod_general(d)),
|
||||
])
|
||||
1676
tinygrad_repo/tinygrad/uop/ops.py
Normal file
1676
tinygrad_repo/tinygrad/uop/ops.py
Normal file
File diff suppressed because it is too large
Load Diff
166
tinygrad_repo/tinygrad/uop/render.py
Normal file
166
tinygrad_repo/tinygrad/uop/render.py
Normal file
@@ -0,0 +1,166 @@
|
||||
from typing import cast
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.uop import Ops, GroupOp
|
||||
from tinygrad.uop.ops import UOp, PatternMatcher, UPat, multirange_str, range_str, consumer_map_from_toposort
|
||||
from tinygrad.helpers import strip_parens, all_same
|
||||
|
||||
def pretty_print(x:UOp, cache=None, d=0)->str:
|
||||
def dfs(x:UOp, cache:dict):
|
||||
for s in x.src:
|
||||
cache.setdefault(s, [len(cache), 0, False])[1] += 1
|
||||
if cache[s][1] == 1: dfs(s, cache)
|
||||
if cache is None: dfs(x, cache:={})
|
||||
if (cx:=cache.setdefault(x, [0,0,False]))[2]: return f"{' '*d}x{cx[0]}"
|
||||
cx[2], srcs = True, (''.join(f'\n{pretty_print(s, cache, d+2)},' for s in x.src))
|
||||
return f"{' '*d}{f'x{cx[0]}:=' * (cx[1]>1)}{type(x).__name__}({x.op}, {x.dtype}, arg={x.argstr()}{x.tagstr()}, src=({srcs}))"
|
||||
|
||||
# ***** uop helpers *****
|
||||
|
||||
def print_uops(uops:list[UOp]):
|
||||
uops_index = {u:i for i,u in enumerate(uops)}
|
||||
for i,u in enumerate(uops):
|
||||
formatted_srcs = [(uops_index[x] if x.op is not Ops.CONST else f"{x.arg}") if x in uops else "--" for x in u.src]
|
||||
print(f"{i:4d} {str(u.op):20s}: {multirange_str(u.ranges, color=True, pad=10)} {str(u.dtype):40s} " f"{str(formatted_srcs):32s} {u.arg}")
|
||||
|
||||
# for debug
|
||||
syms = { Ops.ADD: "+", Ops.SUB: "-", Ops.FLOORDIV: "//", Ops.FLOORMOD: "%", Ops.SHL: "<<", Ops.SHR: ">>",
|
||||
Ops.MUL: "*", Ops.CMPLT: "<", Ops.CMPNE: "!=", Ops.AND: "&", Ops.OR: "|", Ops.XOR: "^"}
|
||||
# comparison operators are not in here because they are chained in python, not left-associative
|
||||
precedence = {Ops.MUL:1, Ops.FLOORDIV:1, Ops.FLOORMOD:1, Ops.ADD:2, Ops.SUB:2, Ops.SHL:3, Ops.SHR:3, Ops.AND:4, Ops.XOR:5, Ops.OR:6}
|
||||
def strip_binary_parens(x:UOp, left:str, right:str, code_for_op) -> str:
|
||||
if x.op not in precedence: return code_for_op(left, right)
|
||||
return code_for_op(strip_parens(left) if precedence.get(x.src[0].op,99)<=precedence[x.op] else left, strip_parens(right) if
|
||||
precedence.get(x.src[1].op,99)<precedence[x.op] else right)
|
||||
|
||||
renderer = PatternMatcher([
|
||||
(UPat((Ops.DEFINE_VAR,), name="x"), lambda x: x.expr),
|
||||
(UPat(Ops.PARAM, name="x"), lambda x: x.arg.name if x.arg.name is not None else f"p{x.arg.slot}"),
|
||||
(UPat((Ops.SPECIAL), name="x"), lambda x: x.arg),
|
||||
(UPat(Ops.RANGE, name="x"), lambda x: f"r{range_str(x)}"),
|
||||
(UPat(Ops.CONST, name="x"), lambda x: str(x.arg)),
|
||||
(UPat(Ops.UNROLL, name="x"), lambda ctx,x,u: f"UNROLL({ctx[x.src[0]]}, {u.arg})"),
|
||||
(UPat(Ops.CAST, name="x"), lambda ctx,x: f"({str(x.dtype)[7:]})({ctx[x.src[0]]})"),
|
||||
(UPat(Ops.BIND, name="x"), lambda ctx,x: ctx[x.src[0]]),
|
||||
(UPat(Ops.NEG, name="x"), lambda ctx,x: f"(-{ctx[x.src[0]]})"),
|
||||
(UPat(Ops.RECIPROCAL, name="x"), lambda ctx,x: f"(1/{ctx[x.src[0]]})"),
|
||||
(UPat(Ops.MAX, name="x"), lambda ctx,x: f"max({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
|
||||
(UPat(Ops.MULACC, name="x"), lambda ctx,x: f"({ctx[x.src[0]]}*{ctx[x.src[1]]}+{ctx[x.src[2]]})"),
|
||||
(UPat(Ops.WHERE, name="x"), lambda ctx,x: f"({ctx[x.src[1]]} if {ctx[x.src[0]]} else {ctx[x.src[2]]})"),
|
||||
(UPat(Ops.CDIV, name="x"), lambda ctx,x: f"cdiv({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
|
||||
(UPat(Ops.CMOD, name="x"), lambda ctx,x: f"cmod({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
|
||||
(UPat(set(syms.keys()), name="x"), lambda ctx,x: strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")),
|
||||
(UPat((Ops.INDEX, Ops.STAGE), name="x"), lambda x, ctx: ''.join([f"[{strip_parens(ctx[y])}]" for y in x.src[1:]])),
|
||||
(UPat(Ops.STACK, name="x"),
|
||||
lambda ctx,x: f"{{{','.join([ctx[y] for y in x.src])}}}" if not x.src or not all_same(x.src) else f"{{{ctx[x.src[0]]}, ...}}"),
|
||||
(UPat(GroupOp.All, name="x"), lambda x: str(x)),
|
||||
])
|
||||
|
||||
renderer_infer = PatternMatcher([
|
||||
(UPat(Ops.CMOD, name="x"), lambda ctx,x: f"cmod({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
|
||||
(UPat(Ops.CDIV, name="x"), lambda ctx,x: f"cdiv({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
|
||||
(UPat(Ops.FLOORMOD, name="x"), lambda ctx,x: f"floormod({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
|
||||
(UPat(Ops.FLOORDIV, name="x"), lambda ctx,x: f"floordiv({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
|
||||
(UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"bitcast({ctx[x.src[0]]}, {x.src[0].dtype!r}, {x.dtype!r})"),
|
||||
]) + renderer
|
||||
|
||||
# *** pyrender ***
|
||||
|
||||
def srcs(ctx, src): return f"({ctx[src[0]]},)" if len(src) == 1 else f"({', '.join([ctx[x] for x in src])})"
|
||||
def render_marg(ctx,x:UOp):
|
||||
if x.op is Ops.PERMUTE: return str(x.marg)
|
||||
if x.op is Ops.FLIP: return str(tuple([i for i,x in enumerate(x.marg) if x]))
|
||||
pieces = []
|
||||
if x.op in {Ops.RESHAPE, Ops.EXPAND}:
|
||||
pieces = [f"{ctx[a] if isinstance(a, UOp) else str(a)}" for a in x.marg]
|
||||
if x.op in {Ops.PAD, Ops.SHRINK}:
|
||||
pieces = [f"({ctx[a[0]] if isinstance(a[0], UOp) else str(a[0])}, {ctx[a[1]] if isinstance(a[1], UOp) else str(a[1])})" for a in x.marg]
|
||||
return f"({','.join(pieces)})" if len(pieces) != 1 else f"({pieces[0]},)"
|
||||
|
||||
sugar = {Ops.SINK, Ops.END, Ops.STORE, Ops.LOAD, Ops.UNIQUE, Ops.SQRT, Ops.INDEX, Ops.REDUCE, Ops.AFTER, Ops.THREEFRY,
|
||||
Ops.WHERE, Ops.RECIPROCAL, Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.CONTIGUOUS, Ops.BARRIER, Ops.DETACH}
|
||||
pm_pyrender_extra = PatternMatcher([
|
||||
(UPat(Ops.CONST, src=(UPat(Ops.UNIQUE, name="u"), UPat(Ops.DEVICE, name="d")), name="x"),
|
||||
lambda x,u,d: f"UOp.unique_const({x.arg}, dtype={x.dtype}, device={repr(d.arg)}, unique={u.arg})"),
|
||||
(UPat(Ops.CONST, src=(UPat(Ops.DEVICE, name="d"),), name="x"), lambda x,d: f"UOp.const({x.dtype}, {x.arg}, device={repr(d.arg)})"),
|
||||
(UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.dtype}, {x.arg})"),
|
||||
(UPat(Ops.DEFINE_VAR, src=(), name="x"), lambda x:
|
||||
f"UOp.variable(\"{x.arg[0]}\", {x.arg[1]}, {x.arg[2]}{', dtype='+str(x.dtype) if x.dtype is not dtypes.weakint else ''})"),
|
||||
(UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})"),
|
||||
(UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].arg}, {repr(x.arg)}, dtype={x.dtype})"),
|
||||
(UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE, name="u"), UPat(Ops.DEVICE, name="d")), name="x"), lambda x,u,d:
|
||||
f"UOp.new_buffer({repr(d.arg)}, {x.arg}, {x.dtype}, {u.arg})"),
|
||||
(UPat(Ops.COPY, src=(UPat(name="x"), UPat(Ops.DEVICE, name="d"))), lambda ctx,x,d: f"{ctx[x]}.copy_to_device({repr(d.arg)})"),
|
||||
(UPat(Ops.CUSTOM_FUNCTION, name="x"), lambda ctx,x: f"UOp(Ops.CUSTOM_FUNCTION, {x.dtype}, src={srcs(ctx, x.src)}, arg={x.arg!r})"),
|
||||
(UPat(Ops.REDUCE, name="r"), lambda ctx,r: f"{ctx[r.src[0]]}._rop({r.arg[0]}, {r.arg[1]})" if len(r.arg[1]) else None),
|
||||
# NOTE: range has srcs sometimes after control flow
|
||||
(UPat(Ops.RANGE, src=(UPat(Ops.CONST, name="c"),), allow_any_len=True, name="x"), lambda ctx,x,c:
|
||||
"UOp.range("+', '.join([str(c.arg)] + [repr(y) for y in x.arg])+
|
||||
(f', src={srcs(ctx, x.src[1:])}' if len(x.src) > 1 else '')+(', dtype='+str(x.dtype) if x.dtype is not dtypes.weakint else '')+")"),
|
||||
# TODO: index shouldn't mismatch dtype
|
||||
(UPat(Ops.INDEX, src=(UPat(), UPat()), allow_any_len=True, name="x"), lambda ctx,x:
|
||||
f"{ctx[x.src[0]]}.index({ctx[x.src[1]]}, "+''.join([f"{ctx[xx]}, " for xx in x.src[2:]])+
|
||||
(f"dtype={x.dtype})" if x.src[0].dtype != x.dtype else "ptr=True)") if x.src[0].dtype.base != x.dtype else None),
|
||||
# TODO: movement ops simplify stuff, this can break SPEC=2
|
||||
#(UPat(GroupOp.Movement, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({render_marg(ctx,x)})"),
|
||||
# NOTE: CMPNE doesn't work cause there's no __rne__
|
||||
# explicit trunc ops: `//` and `%` parse as FLOORDIV/FLOORMOD, so render CDIV/CMOD via .alu()
|
||||
(UPat(Ops.CDIV, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.CDIV, {ctx[x.src[1]]})"),
|
||||
(UPat(Ops.CMOD, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.CMOD, {ctx[x.src[1]]})"),
|
||||
# NOTE: only match CONSTs without UNIQUE (len(src)==1), unique_const needs explicit rendering
|
||||
(UPat(set(syms.keys())-{Ops.SUB, Ops.CMPNE, Ops.CDIV, Ops.CMOD}, src=(UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="y"), UPat(name="z")),
|
||||
name="x"), lambda ctx,x,y,z: strip_binary_parens(x, str(y.arg), ctx[z], lambda a,b: f"({a}{syms[x.op]}{b})") if y.device==z.device else None),
|
||||
# NOTE: sub doesn't work cause it's written as add/mul
|
||||
(UPat(set(syms.keys())-{Ops.SUB, Ops.CDIV, Ops.CMOD}, src=(UPat(name="y"), UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="z")), name="x"),
|
||||
lambda ctx,x,y,z: strip_binary_parens(x, ctx[y], str(z.arg), lambda a,b: f"({a}{syms[x.op]}{b})") if y.device==z.device else None),
|
||||
(UPat(set(syms.keys())-{Ops.SUB, Ops.CDIV, Ops.CMOD}, name="x"), lambda ctx,x:
|
||||
strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")),
|
||||
(UPat(sugar, src=(), name="x"), lambda x: f"UOp.{x.op.name.lower()}("+', '.join(([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),
|
||||
(UPat(sugar, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}("+', '.join([ctx[y] for y in x.src[1:]] + \
|
||||
([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),
|
||||
])
|
||||
|
||||
# NOTE: you can remove pm_pyrender_extra and it'll still be correct
|
||||
pm_pyrender = pm_pyrender_extra+PatternMatcher([
|
||||
(UPat(GroupOp.All, name="u"), lambda ctx,u: f"UOp({u.op}, {u.dtype}, {srcs(ctx,u.src)}"+(f", {repr(u.arg)})" if u.arg is not None else ")")),
|
||||
])
|
||||
|
||||
def _render_with_splits(lst:list[UOp], pm:PatternMatcher, to_render:set[UOp], split_depth:int=100) -> dict[str, str]:
|
||||
r: dict[UOp, str] = {}
|
||||
ret: dict[str, str] = {}
|
||||
depth: dict[UOp, int] = {}
|
||||
for i,u in enumerate(lst):
|
||||
# limit inline depth to avoid "too many nested parentheses" in Python parser
|
||||
op_depth = 1 + max([depth.get(s, 0) for s in u.src], default=0)
|
||||
if op_depth > split_depth: to_render.add(u)
|
||||
depth[u] = 0 if u in to_render else op_depth
|
||||
ren = cast(str, pm.rewrite(u, ctx=r))
|
||||
assert isinstance(ren, str)
|
||||
if u.tag is not None: ren += f".rtag({repr(u.tag)})"
|
||||
if u not in to_render: r[u] = ren
|
||||
else:
|
||||
r[u] = f"c{i}" if u is not lst[-1] else "ast"
|
||||
ret[r[u]] = ren
|
||||
return ret
|
||||
|
||||
def pyrender(ast:UOp) -> str:
|
||||
lst = list(ast.toposort())
|
||||
|
||||
cmap = consumer_map_from_toposort(lst)
|
||||
not_rendered = {Ops.CONST, Ops.DEVICE}
|
||||
always_rendered = {Ops.PARAM, Ops.LOAD, Ops.SPECIAL, Ops.RANGE, Ops.CONTIGUOUS, Ops.STACK,
|
||||
Ops.BUFFER, Ops.COPY, Ops.CALL, Ops.FUNCTION, Ops.WHERE, Ops.END}
|
||||
|
||||
to_render: set[UOp] = {ast}
|
||||
for u in lst:
|
||||
if u.op in {Ops.SINK}:
|
||||
for s in u.src: to_render.add(s)
|
||||
if u.op is Ops.STORE: to_render.add(u.src[1])
|
||||
if u.op is Ops.REDUCE: to_render.add(u.src[0])
|
||||
if u.op in {Ops.CALL, Ops.FUNCTION}: raise NotImplementedError("call can't be pyrendered")
|
||||
if u.op in not_rendered: continue
|
||||
# checking the consumers is not enough, you have to make sure it's not used twice by the one consumer
|
||||
if len(cmap[u]) == 1 and len([x for x in list(cmap[u].keys())[0].src if x is u]) == 1 and u.op not in always_rendered: continue
|
||||
to_render.add(u)
|
||||
|
||||
ret = _render_with_splits(lst, pm_pyrender, to_render)
|
||||
return '\n'.join([f"{k} = {strip_parens(v)}" for k,v in ret.items()])
|
||||
273
tinygrad_repo/tinygrad/uop/spec.py
Normal file
273
tinygrad_repo/tinygrad/uop/spec.py
Normal file
@@ -0,0 +1,273 @@
|
||||
import math
|
||||
from typing import cast, Any
|
||||
from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, AxisType, KernelInfo, ParamArg
|
||||
from tinygrad.uop.render import print_uops, pyrender
|
||||
from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace, Invalid, ConstFloat
|
||||
from tinygrad.helpers import DEBUG, Context, prod, SPEC, Metadata, panic, CHECK_OOB
|
||||
|
||||
# ***** uop helpers *****
|
||||
|
||||
def validate_index(uidx:UOp, gate:UOp|None=None):
|
||||
if len(uidx.src) != 2: return True # skip for non final index. TODO: check more complex index with shape
|
||||
buf,idx = uidx.src
|
||||
if idx.op is Ops.CONST and idx.arg is Invalid: return True
|
||||
if gate is None: gate = UOp.const(dtypes.bool, True)
|
||||
# TODO: check for overflow
|
||||
if not CHECK_OOB or isinstance(buf.dtype, ImageDType) or (sz := buf.ptrdtype.size) == -1: return True
|
||||
|
||||
# We can use UOp min/max to do a faster check, but it can give false positive since its not an exact bound and doesn't consider the mask
|
||||
if 0<=idx.vmin and idx.vmax<sz: return True
|
||||
|
||||
# TODO: validate these
|
||||
# WEBGPU has a BITCAST in the index, PTX casts pointer to long
|
||||
# VECTORIZE/GEP can't be properly modeled in z3 since it doesn't support vectors
|
||||
for x in idx.toposort() | gate.toposort():
|
||||
if x.op in {Ops.BITCAST, Ops.STACK, Ops.GEP} or (x.op is Ops.CAST and isinstance(x.src[0].dtype, PtrDType)): return True
|
||||
|
||||
# if all is good and CHECK_OOB=1, validate with z3
|
||||
from tinygrad.uop.validate import validate_index_with_z3
|
||||
return validate_index_with_z3(sz, idx, gate)
|
||||
|
||||
def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher):
|
||||
lst = list(ast.toposort()) if isinstance(ast, UOp) else ast
|
||||
if SPEC > 1: test_pyrender(lst[-1]) # assume this is the sink
|
||||
|
||||
with Context(TRACK_MATCH_STATS=0):
|
||||
for i,u in enumerate(lst):
|
||||
ret = check_spec.rewrite(u)
|
||||
if cast(bool|None, ret) is not True:
|
||||
if DEBUG >= 3: print_uops(lst)
|
||||
raise RuntimeError(f"UOp verification failed at {i} on {u.op} {u.dtype} {len(u.src)} {[(x.op, x.dtype, x.arg) for x in u.src]} {u.arg}")
|
||||
|
||||
# ***** new specs *****
|
||||
|
||||
# these ops can be used in the tensor graph and programs
|
||||
spec_shared = PatternMatcher([
|
||||
(UPat(Ops.SINK, dtypes.void), lambda: True), # NOTE: for testing, we let sinks be anything
|
||||
|
||||
# NOOP. TODO: remove this
|
||||
(UPat(Ops.NOOP), lambda: True),
|
||||
|
||||
# CONST/DEFINE_VAR are everywhere
|
||||
(UPat(Ops.CONST, src=(), name="x"), lambda x: type(x.arg) is type(x.dtype.const(x.arg))),
|
||||
(UPat(Ops.DEFINE_VAR, name="x"), lambda x: len(x.arg) == 3 and isinstance(x.arg[0], str)),
|
||||
|
||||
# ALUs: most ALUs have all matching dtypes, except CMPLT, CMPNE, and WHERE
|
||||
(UPat(Ops.WHERE, name="w", src=(UPat(dtype=dtypes.bool), UPat.var("x"), UPat.var("y"))), lambda w,x,y: w.dtype == x.dtype == y.dtype),
|
||||
(UPat((Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ), dtype=dtypes.bool, src=(UPat.var("x"), UPat.var("y"))), lambda x,y: x.dtype.base == y.dtype.base),
|
||||
# and SHL/SHR, the shift distance can be an int
|
||||
(UPat((Ops.SHL, Ops.SHR), src=(UPat.var("x"), UPat.var("y")), name="a"), lambda a,x,y: a.dtype == x.dtype and y.dtype in (x.dtype, dtypes.uint)),
|
||||
(UPat((Ops.CDIV, Ops.CMOD, Ops.FLOORDIV, Ops.FLOORMOD), name="x"), lambda x: None if dtypes.is_int(x.dtype) else False),
|
||||
(UPat(GroupOp.ALU, name="x"), lambda x: all(x.dtype.base == y.dtype.base for y in x.src)),
|
||||
|
||||
# CAST
|
||||
(UPat((Ops.BITCAST, Ops.CAST), src=(UPat(),), name="x"), lambda x: x.arg is None),
|
||||
|
||||
# RANGE can be in the big graph now
|
||||
(UPat(Ops.RANGE, src=(UPat.var("x"),), allow_any_len=True, name="rng"), lambda rng,x:
|
||||
rng.dtype == x.dtype and isinstance(rng.arg, tuple) and len(rng.arg) >= 2 and \
|
||||
all(isinstance(ra, int) for ra in rng.arg[0:-1]) and isinstance(rng.arg[-1], AxisType)),
|
||||
(UPat(Ops.INDEX, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(dtypes.is_int(y.dtype) for y in x.src[1:]) or None),
|
||||
(UPat(Ops.END, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(u.op is Ops.RANGE for u in x.src[1:])),
|
||||
|
||||
# PARAM
|
||||
(UPat(Ops.PARAM, name="x"), lambda x: isinstance(x.arg, ParamArg)),
|
||||
|
||||
# GROUP of stores (or groups, or NOOPs)
|
||||
# TODO: remove UNROLL here, it's for SPEC=2
|
||||
(UPat(Ops.GROUP, dtypes.void, src=UPat((Ops.GROUP, Ops.STORE, Ops.NOOP, Ops.UNROLL, Ops.INS))), lambda: True),
|
||||
|
||||
# TOOD: these should be buffer with different addrspace
|
||||
(UPat(Ops.DEFINE_LOCAL, name="x"), lambda x: isinstance(x.dtype, PtrDType) and x.dtype.addrspace == AddrSpace.LOCAL),
|
||||
(UPat(Ops.DEFINE_REG, src=()), lambda: True),
|
||||
|
||||
# AFTER on Movement Op, PARAM, BUFFER, CONTIGUOUS, or another AFTER
|
||||
(UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.PARAM, Ops.BUFFER, Ops.CONTIGUOUS, Ops.DEFINE_REG, Ops.DEFINE_LOCAL, Ops.AFTER, Ops.MULTI,
|
||||
Ops.BITCAST, Ops.INS})),),
|
||||
allow_any_len=True), lambda: True),
|
||||
|
||||
# CUSTOM (inline and non inline)
|
||||
(UPat((Ops.CUSTOMI, Ops.CUSTOM)), lambda: True),
|
||||
|
||||
# BARRIER (on any length). TODO: this should only be in spec_program
|
||||
(UPat(Ops.BARRIER, dtypes.void), lambda: True),
|
||||
|
||||
# SPECIAL. TODO: this should only be in spec_program
|
||||
(UPat(Ops.SPECIAL, src=(UPat.var("x", (dtypes.weakint, dtypes.int32)),), name="s"), lambda s,x: s.dtype == x.dtype and isinstance(s.arg, str)),
|
||||
|
||||
# assembly instruction
|
||||
(UPat(Ops.INS), lambda: True),
|
||||
|
||||
# LOAD(idx) / STORE(idx, val) with gates on the LOAD/STORE
|
||||
(UPat(Ops.INDEX, name="uidx").or_casted().load(), validate_index),
|
||||
(UPat(Ops.INDEX, name="uidx").or_casted().load(UPat.var("alt"), UPat.var("gate", dtype=dtypes.bool), name="load"),
|
||||
lambda uidx,gate,alt,load: validate_index(uidx, gate) if alt.dtype == load.dtype else False),
|
||||
(UPat(Ops.INDEX, name="uidx").or_casted().store(UPat()), validate_index),
|
||||
(UPat(Ops.INDEX, name="uidx").or_casted().store(UPat(), UPat.var("gate", dtype=dtypes.bool)), validate_index),
|
||||
|
||||
# STORE in tensor graph: store a value into a target
|
||||
(UPat(Ops.STORE, dtypes.void, (UPat(name="x"), UPat())), lambda x: True),
|
||||
|
||||
# WMMA has a <a, b, acc>
|
||||
(UPat(Ops.WMMA, src=(UPat(), UPat(), UPat()), name="x"), lambda x: isinstance(x.arg, tuple) and len(x.arg) == 8),
|
||||
])
|
||||
|
||||
# these ops can exist in tensor but not programs. example: movement
|
||||
spec_tensor = PatternMatcher([
|
||||
# DEVICE
|
||||
(UPat(Ops.DEVICE, dtypes.void, (), name="d"), lambda d:
|
||||
isinstance(d.arg, str) or (isinstance(d.arg, tuple) and all(isinstance(s, str) for s in d.arg))),
|
||||
|
||||
# UNIQUE
|
||||
(UPat(Ops.UNIQUE, dtypes.void, ()), lambda: True),
|
||||
(UPat(Ops.LUNIQUE, dtypes.void, ()), lambda: True),
|
||||
|
||||
# CONST with a UNIQUE or DEVICE
|
||||
(UPat(Ops.CONST, src=(UPat(Ops.DEVICE),)), lambda: True),
|
||||
(UPat(Ops.CONST, src=(UPat((Ops.UNIQUE, Ops.LUNIQUE)), UPat(Ops.DEVICE)), name="c"), lambda c: c.arg is Invalid),
|
||||
|
||||
# BUFFER
|
||||
(UPat(Ops.BUFFER, src=(UPat((Ops.UNIQUE, Ops.LUNIQUE)), UPat(Ops.DEVICE)), name="buf"),
|
||||
lambda buf: isinstance(buf.arg, int) and isinstance(buf.dtype, DType)),
|
||||
|
||||
# Tensor variable bindings
|
||||
(UPat(Ops.BIND, (dtypes.int, dtypes.weakint,), (UPat(Ops.DEFINE_VAR), UPat.cvar(dtype=(dtypes.int,dtypes.weakint,))), arg=None), lambda: True),
|
||||
|
||||
# custom function
|
||||
(UPat(Ops.CUSTOM_FUNCTION, name="x"), lambda x: isinstance(x.arg, str)),
|
||||
|
||||
# CALL
|
||||
(UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.LINEAR, Ops.PROGRAM, Ops.COPY, Ops.CUSTOM_FUNCTION)),), allow_any_len=True), lambda: True),
|
||||
|
||||
# FUNCTION + TUPLE must have void dtype, GETTUPLE can only appear on FUNCTION or TUPLE
|
||||
(UPat(Ops.FUNCTION, dtypes.void, src=(UPat(Ops.TUPLE),), allow_any_len=True), lambda: True),
|
||||
(UPat(Ops.TUPLE, dtypes.void), lambda: True),
|
||||
(UPat(Ops.GETTUPLE, src=(UPat((Ops.FUNCTION, Ops.TUPLE)),), name="g"), lambda g: isinstance(g.arg, int)),
|
||||
|
||||
# inputs to movement ops
|
||||
(UPat(Ops.STACK), lambda: True),
|
||||
(UPat({Ops.ADD, Ops.MUL, Ops.CDIV, Ops.FLOORDIV}, dtype=dtypes.weakint), lambda: True),
|
||||
|
||||
# movement ops
|
||||
(UPat((Ops.RESHAPE, Ops.EXPAND), src=(UPat(), UPat(dtype=dtypes.weakint))), lambda: True),
|
||||
(UPat((Ops.PAD, Ops.SHRINK), src=(UPat(), UPat(dtype=dtypes.weakint), UPat(dtype=dtypes.weakint)), name="x"),
|
||||
lambda x: x.src[1].dtype.count == x.src[2].dtype.count),
|
||||
(UPat((Ops.PERMUTE, Ops.FLIP), name="mv", src=(UPat(),)), lambda mv: isinstance(mv.arg, tuple)),
|
||||
|
||||
# REDUCE has arg=(op, axis_tuple), src[1:] are ranges after lowering
|
||||
(UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"),
|
||||
lambda x: isinstance(x.arg, tuple) and len(x.arg) == 2 and x.arg[0] in {Ops.ADD, Ops.MUL, Ops.MAX}
|
||||
and isinstance(x.arg[1], tuple) and all(y.dtype in (dtypes.weakint, dtypes.int) for y in x.src[1:])),
|
||||
|
||||
# COPY. TODO: this should not have allow_any_len, but something is adding ranges
|
||||
(UPat(Ops.COPY, name="copy", src=(UPat.var("x"), UPat(Ops.DEVICE)), allow_any_len=True, arg=None), lambda copy,x: copy.dtype == x.dtype),
|
||||
(UPat(Ops.ALLREDUCE, name="red", src=(UPat.var("x"), UPat(Ops.DEVICE))), lambda red,x: red.dtype == x.dtype and isinstance(red.arg, Ops)),
|
||||
|
||||
# MULTI/MSELECT/MSTACK
|
||||
(UPat(Ops.MULTI, name="multi"), lambda multi: all(x.dtype == multi.dtype for x in multi.src) and isinstance(multi.arg, int)),
|
||||
(UPat(Ops.MSELECT, name="x"), lambda x: isinstance(x.src[0].device, tuple) and x.arg < len(x.src[0].device)),
|
||||
(UPat(Ops.MSTACK, name="x"), lambda x: all(isinstance(x.device, str) for x in x.src)),
|
||||
|
||||
# CONTIGUOUS ensures the source UOp realizes
|
||||
(UPat((Ops.DETACH, Ops.CONTIGUOUS, Ops.CONTIGUOUS_BACKWARD), name="root", src=(UPat.var("x"),), arg=None),
|
||||
lambda root,x: root.dtype == x.dtype),
|
||||
|
||||
# TODO: this should not be here. STAGE is transformed to DEFINE_LOCAL later
|
||||
(UPat(Ops.STAGE, src=(UPat(),), allow_any_len=True), lambda: True),
|
||||
|
||||
# codegen: PROGRAM with progressive sources through the pipeline (SINK, DEVICE, LINEAR?, SOURCE?, BINARY?)
|
||||
(UPat(Ops.LINEAR, dtypes.void), lambda: True),
|
||||
(UPat(Ops.SOURCE, dtypes.void, src=()), lambda: True),
|
||||
(UPat(Ops.BINARY, dtypes.void, src=()), lambda: True),
|
||||
(UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE))), lambda: True),
|
||||
(UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR))), lambda: True),
|
||||
(UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR), UPat(Ops.SOURCE))), lambda: True),
|
||||
(UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR), UPat(Ops.SOURCE), UPat(Ops.BINARY))), lambda: True),
|
||||
|
||||
# UNROLL/CONTRACT is used here for WMMA
|
||||
(UPat(Ops.CONTRACT, name="x"), lambda x: x.dtype.count == prod(y[1] for y in x.arg)),
|
||||
(UPat(Ops.UNROLL, name="x"), lambda x: x.src[0].dtype.count == prod(y[1] for y in x.arg)),
|
||||
])+spec_shared
|
||||
|
||||
# these ops can exist in programs but not the tensor spec. example: LOAD
|
||||
spec_program = PatternMatcher([
|
||||
# weakint is not allowed in programs
|
||||
(UPat(GroupOp.All, dtypes.weakint), lambda: False),
|
||||
|
||||
# movement ops are not allowed in programs
|
||||
(UPat(GroupOp.Movement), lambda: False),
|
||||
|
||||
# Invalid is not allowed in program
|
||||
(UPat(Ops.CONST, arg=Invalid), lambda: False),
|
||||
|
||||
# shape of uop must match dtype.count in program
|
||||
(UPat(GroupOp.All-{Ops.INS, Ops.NOOP}, name="x"),
|
||||
lambda x: False if x.dtype.count > 1 and (x.dtype.count,) != x.shape else None),
|
||||
|
||||
# STACK/GEP in program. TODO: this should match Tensor
|
||||
(UPat(Ops.STACK, name="x"), lambda x: len(x.src)>1 and len(x.src) == x.dtype.vcount and all(x.dtype == y.dtype.vec(len(x.src)) for y in x.src)),
|
||||
(UPat(Ops.GEP, src=(UPat.var("src"),), name="gep"), lambda gep,src: gep.dtype == src.dtype.scalar()),
|
||||
|
||||
# if has a <gate, index_for_dedup>
|
||||
(UPat(Ops.IF, dtype=dtypes.void, src=(UPat(dtype=dtypes.bool), UPat((Ops.CAST, Ops.INDEX)))), lambda: True),
|
||||
(UPat(Ops.ENDIF, dtype=dtypes.void, src=(UPat(Ops.IF),)), lambda: True),
|
||||
])+spec_shared
|
||||
|
||||
# these are intermediate ops. everything should be deleted from here
|
||||
spec_full = PatternMatcher([
|
||||
# SLICE on BUFFER is allowed if BUFFER is
|
||||
(UPat(Ops.SLICE, src=(UPat(GroupOp.Movement.union({Ops.BUFFER, Ops.PARAM, Ops.STAGE, Ops.AFTER})),
|
||||
UPat(Ops.CONST, dtype=dtypes.weakint)), allow_any_len=True, name="bv"),
|
||||
lambda bv: isinstance(bv.arg, int)),
|
||||
|
||||
(UPat(Ops.CALL, src=(UPat((Ops.SLICE,)),), allow_any_len=True), lambda: True),
|
||||
|
||||
# codegen may end ranges after gpudims has replaced RANGE with SPECIAL.
|
||||
(UPat(Ops.END, src=(UPat(), UPat()), allow_any_len=True), lambda: True),
|
||||
|
||||
# allow any AFTER
|
||||
(UPat(Ops.AFTER, src=(UPat(),), allow_any_len=True), lambda: True),
|
||||
|
||||
# expander: unroll/contract/gep/ptrcat/cat
|
||||
(UPat((Ops.UNROLL, Ops.CONTRACT), src=(UPat(),)), lambda: True),
|
||||
|
||||
# GEP multi is supported here
|
||||
(UPat(Ops.GEP, name="gep"), lambda gep: gep.dtype is dtypes.void or gep.dtype.vcount == len(gep.arg)),
|
||||
|
||||
# all loads/stores
|
||||
(UPat((Ops.LOAD, Ops.STORE)), lambda: True),
|
||||
|
||||
# while BIND is being casted
|
||||
(UPat(Ops.BIND, (dtypes.int, dtypes.weakint), (UPat(), UPat()), arg=None), lambda: True),
|
||||
|
||||
# TODO: PTRCAT and VCAT need to be deleted
|
||||
|
||||
# PTRCAT is like VECTORIZE, but it functions on ptrs
|
||||
(UPat(Ops.PTRCAT, name="x"), lambda x: x.dtype.vcount == sum([y.dtype.base.count for y in x.src])),
|
||||
# VCAT is like VECTORIZE, but the srcs can be vectors
|
||||
(UPat(Ops.VCAT, name="x"), lambda x: x.dtype.vcount == sum([y.dtype.vcount for y in x.src])),
|
||||
])+spec_tensor+spec_program
|
||||
|
||||
# **** pyrender (move this) ****
|
||||
|
||||
# late imports to avoid circular import
|
||||
from tinygrad.codegen.opt import Opt, OptOps
|
||||
from tinygrad.schedule.rangeify import BufferizeOpts
|
||||
glbls:dict[str, Any] = {"inf": math.inf, "nan": math.nan, "KernelInfo": KernelInfo, "Metadata": Metadata,
|
||||
"UOp": UOp, "dtypes": dtypes, "Ops": Ops, "AxisType": AxisType, "Invalid": Invalid,
|
||||
"Opt": Opt, "OptOps": OptOps, "BufferizeOpts": BufferizeOpts, "AddrSpace": AddrSpace, "panic": panic,
|
||||
"ConstFloat": ConstFloat, "ParamArg": ParamArg}
|
||||
def eval_pyrender(code:str) -> UOp:
|
||||
lcls:dict[str, Any] = {}
|
||||
exec(code, glbls, lcls)
|
||||
return lcls['ast']
|
||||
|
||||
def test_pyrender(test_ast:UOp, assert_parents=True):
|
||||
try: code = pyrender(test_ast)
|
||||
except NotImplementedError: return None # this is okay, not all ops can be pyrendered
|
||||
ast:UOp = eval_pyrender(code)
|
||||
if ast is not test_ast:
|
||||
if assert_parents:
|
||||
for u in test_ast.toposort(): test_pyrender(u, assert_parents=False)
|
||||
raise RuntimeError(f"PYRENDER ISSUE:\nSTR MATCH: {str(test_ast) == str(ast)}\nUOP:\n{test_ast}\nPRODUCED:\n{ast}\nCODE:\n{code}")
|
||||
return code
|
||||
485
tinygrad_repo/tinygrad/uop/symbolic.py
Normal file
485
tinygrad_repo/tinygrad/uop/symbolic.py
Normal file
@@ -0,0 +1,485 @@
|
||||
# all of symbolic lives here now
|
||||
import math, struct
|
||||
from collections import defaultdict
|
||||
from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu
|
||||
from tinygrad.dtype import ConstType, dtypes, PtrDType, can_lossless_cast, Invalid
|
||||
from tinygrad.helpers import partition, all_same, prod, flatten, get_single_element, unwrap, IMAGE, dedup
|
||||
from tinygrad.uop.decompositions import threefry2x32, xpow
|
||||
from tinygrad.uop.divandmod import div_and_mod_symbolic
|
||||
|
||||
# ******** phase 1 of symbolic used to live in ops, it's the most generic folding rules ********
|
||||
|
||||
def simplify_pow(x:UOp, c:UOp) -> UOp|None:
|
||||
if c.arg < 0: return x.reciprocal().pow(-c)
|
||||
if c.arg == 0: return x.const_like(1)
|
||||
if int(c.arg-0.5)+0.5 == c.arg: return x.pow(c.const_like(c.arg-0.5)) * x.sqrt()
|
||||
if int(c.arg) == c.arg: return (y := x.pow(c.const_like(c.arg//2))) * y * (x if c.arg%2 == 1 else 1)
|
||||
return None
|
||||
|
||||
def fold_bitcast(root:UOp, c:UOp) -> UOp|None:
|
||||
if (from_fmt:=c.dtype.scalar().fmt) is None or (to_fmt:=root.dtype.scalar().fmt) is None: return None
|
||||
if c.dtype.itemsize != root.dtype.itemsize: return None
|
||||
def convert(v:ConstType) -> ConstType: return struct.unpack(to_fmt, struct.pack(from_fmt, v))[0]
|
||||
return root.const_like(convert(c.arg) if root.dtype.count == 1 else tuple(map(convert, c.arg)))
|
||||
|
||||
def const_arg(u:UOp) -> ConstType|tuple[ConstType, ...]|None:
|
||||
if u.op is Ops.CONST: return u.arg
|
||||
if u.op is Ops.STACK and all(s.op is Ops.CONST for s in u.src): return tuple(s.arg for s in u.src)
|
||||
return None
|
||||
|
||||
def fold_const_alu(a:UOp) -> UOp|None:
|
||||
vals = [const_arg(s) for s in a.src]
|
||||
return None if any(v is None for v in vals) else a.const_like(exec_alu(a.op, a.dtype, vals, False))
|
||||
|
||||
invalid_pat = UPat(Ops.CONST, arg=Invalid, name="i")
|
||||
invalid_gate = UPat.var("cond").where(UPat.var("x"), invalid_pat)
|
||||
|
||||
def fold_add_divmod_recombine(x:UOp) -> UOp|None:
|
||||
terms = list(x.split_uop(Ops.ADD))
|
||||
for i,u in enumerate(terms):
|
||||
if u.op is Ops.FLOORMOD and u.src[1].op is Ops.CONST: base, div, mul = u.src[0], u.src[1].arg, 1
|
||||
elif u.op is Ops.MUL and u.src[1].op is Ops.CONST and (m:=u.src[0]).op is Ops.FLOORMOD and m.src[1].op is Ops.CONST:
|
||||
base, div, mul = m.src[0], m.src[1].arg, u.src[1].arg
|
||||
else: continue
|
||||
for j,v in enumerate(terms):
|
||||
if i == j: continue
|
||||
if v.op is not Ops.MUL or v.src[1].op is not Ops.CONST or v.src[1].arg != div*mul: continue
|
||||
q, exact = v.src[0], False
|
||||
# (base%div)*mul + (base//div)*(div*mul) -> base*mul
|
||||
if q.op is Ops.FLOORDIV and q.src[1].op is Ops.CONST and q.src[1].arg == div: exact = q.src[0] is base
|
||||
# ((base//d)%div)*mul + (base//(d*div))*(div*mul) -> (base//d)*mul if div>0
|
||||
if not exact and div > 0 and base.op is Ops.FLOORDIV and base.src[1].op is Ops.CONST:
|
||||
exact = q.op is Ops.FLOORDIV and q.src[1].op is Ops.CONST and q.src[0] is base.src[0] and q.src[1].arg == base.src[1].arg*div
|
||||
if exact: return (base*mul).usum(*[t for k,t in enumerate(terms) if k not in (i,j)])
|
||||
# ((base//div)%d)*(div*mul) + (base%div)*mul -> (base%(div*d))*mul
|
||||
if div > 0 and q.op is Ops.FLOORMOD and q.src[1].op is Ops.CONST and (d:=q.src[1].arg) > 0 and q.src[0].op is Ops.FLOORDIV:
|
||||
if q.src[0].src[0] is base and q.src[0].src[1].op is Ops.CONST and q.src[0].src[1].arg == div:
|
||||
return ((base % (div*d))*mul).usum(*[t for k,t in enumerate(terms) if k not in (i,j)])
|
||||
return None
|
||||
|
||||
# this needs to be before symbolic so that 0*something_that_might_be_invalid doesnt become 0
|
||||
propagate_invalid = PatternMatcher([
|
||||
# propagate invalid, push it past children
|
||||
(invalid_gate.cast(name="cast"), lambda i,x,cond,cast: x.cast(cast.dtype) if i.dtype is dtypes.weakint else None),
|
||||
(UPat(GroupOp.Unary, src=(invalid_gate,), name="alu"), lambda cond,x,alu,i: cond.where(x.alu(alu.op), i)),
|
||||
(UPat(GroupOp.Binary-GroupOp.Comparison, src=(invalid_gate, UPat.var("y")), name="alu"), lambda cond,x,y,alu,i: cond.where(x.alu(alu.op,y), i)),
|
||||
(UPat(GroupOp.Binary-GroupOp.Comparison, src=(UPat.var("y"), invalid_gate), name="alu"), lambda cond,x,y,alu,i: cond.where(y.alu(alu.op,x), i)),
|
||||
# TODO: when can this happen? and is it always safe to just drop invalid?
|
||||
(UPat(GroupOp.Comparison, src=(invalid_gate, UPat.var("y")), name="alu"), lambda cond,x,y,alu,i:
|
||||
x.alu(alu.op,y) if i.dtype is dtypes.weakint else cond.where(x.alu(alu.op,y), i.cast(dtypes.bool))),
|
||||
(UPat(GroupOp.Comparison, src=(UPat.var("y"), invalid_gate), name="alu"), lambda cond,x,y,alu,i:
|
||||
y.alu(alu.op,x) if i.dtype is dtypes.weakint else cond.where(y.alu(alu.op,x), i.cast(dtypes.bool))),
|
||||
# alu with invalid -> invalid
|
||||
(UPat(GroupOp.Unary, src=(invalid_pat,)), lambda i: i),
|
||||
(UPat(GroupOp.Binary-GroupOp.Comparison, src=[invalid_pat, UPat()]), lambda i: i),
|
||||
# normalize where(cond, Invalid, val) -> where(~cond, val, Invalid)
|
||||
(UPat.var("cond").where(invalid_pat, UPat.var("val")), lambda cond, i, val: cond.logical_not().where(val, i) if val.arg != Invalid else i),
|
||||
# lift Invalid out # TODO: this `a is cond` is asymmetric to preserve the pattern
|
||||
(UPat.var("a").where(invalid_gate, UPat.var("c")), lambda cond,i,x,a,c:
|
||||
(cond if a is cond else (a.logical_not()|cond)).where(a.where(x,c), i) if c.arg != Invalid else None),
|
||||
(UPat.var("a").where(UPat.var("b"), invalid_gate), lambda cond,i,x,a,b: (a|cond).where(a.where(b, x), i) if b.arg != Invalid else None),
|
||||
(UPat(Ops.BITCAST, src=(invalid_pat,), name="bc"), lambda bc,i: i.cast(bc.dtype)),
|
||||
(UPat(Ops.BITCAST, src=(invalid_gate,), name="bc"), lambda bc,cond,x,i: cond.where(x.bitcast(bc.dtype), i.bitcast(bc.dtype))),
|
||||
# fold gated LOAD/STORE
|
||||
(UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat(), invalid_pat), allow_any_len=True).or_casted(), UPat())), lambda i: UOp(Ops.NOOP)),
|
||||
(UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), invalid_pat), allow_any_len=True).or_casted(),), allow_any_len=True, name="x"),
|
||||
lambda x,i: x.src[1] if len(x.src) > 1 else x.const_like(0)),
|
||||
])
|
||||
|
||||
symbolic_simple = propagate_invalid + PatternMatcher([
|
||||
# ** self folding **
|
||||
(UPat.var("x") + 0, lambda x: x), # x+0 -> x
|
||||
(UPat.var("x") * 1, lambda x: x), # x*1 -> x
|
||||
(UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)) ^ 0, lambda x: x), # x^0 -> x
|
||||
(UPat.var("x") // UPat.var("x"), lambda x: x.const_like(1)), # x//x -> 1
|
||||
(UPat.var("x") // 1, lambda x: x), # x//1 -> x
|
||||
(UPat.var("x") // -1, lambda x: -x), # x//-1 -> -x
|
||||
((UPat.var("x") ^ UPat.var("y")) ^ UPat.var("y"), lambda x,y: x), # (x^y)^y -> x
|
||||
((UPat.var() % UPat.var("y")).named("base") % UPat.var("y"), lambda base,y: base), # (x%y)%y = -> x%y (rewritten with base for speed)
|
||||
# variations of (x%c)+(x//c)*c = x
|
||||
(UPat(Ops.ADD, dtype=dtypes.weakint, name="x"), fold_add_divmod_recombine),
|
||||
(UPat.var("x", dtype=dtypes.bool) & UPat.cvar("c"), lambda x,c: x if c.arg else c),
|
||||
(UPat.var("x", dtype=dtypes.bool) | UPat.cvar("c"), lambda x,c: c if c.arg else x),
|
||||
(UPat(GroupOp.Idempotent, src=(UPat.var("x"), UPat.var("x"))), lambda x: x),
|
||||
(UPat.var("x", dtype=dtypes.bool).logical_not().logical_not(), lambda x: x),
|
||||
(UPat.var("x", dtype=dtypes.bool).where(UPat.const(dtypes.bool, True), UPat.const(dtypes.bool, False)), lambda x: x),
|
||||
(UPat.var("x", dtype=dtypes.bool).where(UPat.const(dtypes.bool, False), UPat.const(dtypes.bool, True)), lambda x: x.logical_not()),
|
||||
# CAST(bool -> int) != const — CAST(True)=1, CAST(False)=0, so fold based on const value
|
||||
(UPat.var("x", dtype=dtypes.bool).cast(dtypes.ints+(dtypes.weakint,)) != UPat.cvar("c"),
|
||||
lambda x,c: x if c.arg == 0 else x.logical_not() if c.arg == 1 else x.const_like(True)),
|
||||
(UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)).trunc(), lambda x: x),
|
||||
# ** zero folding **
|
||||
(UPat.var("x") < UPat.var("x"), lambda x: x.const_like(False).cast(dtypes.bool.vec(x.dtype.count))), # x < x -> False
|
||||
(UPat.var("x") % UPat.var("x"), lambda x: x.const_like(0)), # x%x -> 0
|
||||
(UPat.var("x") ^ UPat.var("x"), lambda x: x.const_like(0)), # x^x -> 0
|
||||
(UPat.var("x") & 0, lambda x: x.const_like(0)), # x&0 -> 0
|
||||
# (x&mask)>>k -> x>>k when mask only clears bits below k
|
||||
# TODO: combine this with "# rules for threefry" below
|
||||
((UPat.var("x") & UPat.cvar("mask")) >> UPat.cvar("k"),
|
||||
lambda x,mask,k: x >> k.arg if mask.arg | ((1 << k.arg) - 1) == -1 else None),
|
||||
(UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)) != UPat.var("x"),
|
||||
lambda x: x.const_like(False).cast(dtypes.bool.vec(x.dtype.count))), # x != x -> False (only ints)
|
||||
# ** constant folding **
|
||||
(UPat(GroupOp.Unary, src=(UPat((Ops.CONST, Ops.STACK)),), name="a"), fold_const_alu),
|
||||
(UPat(GroupOp.Binary-{Ops.THREEFRY}, src=(UPat((Ops.CONST, Ops.STACK)),)*2, name="a"), fold_const_alu),
|
||||
(UPat(Ops.THREEFRY, src=(UPat.cvar("x"), UPat.cvar("key")), name="a"),
|
||||
lambda a, x, key: a.const_like(threefry2x32(x, key).simplify().arg)),
|
||||
(UPat(GroupOp.Ternary, src=(UPat((Ops.CONST, Ops.STACK)),)*3, name="a"), fold_const_alu),
|
||||
# bool MUL is AND, ADD/MAX is OR. prevents other rules to rewrite bool ADD/MUL incorrectly
|
||||
(UPat.var('x', dtype=dtypes.bool) * UPat.var('y', dtype=dtypes.bool), lambda x,y: x&y),
|
||||
(UPat.var('x', dtype=dtypes.bool) + UPat.var('y', dtype=dtypes.bool), lambda x,y: x|y),
|
||||
(UPat.var('x', dtype=dtypes.bool).maximum(UPat.var('y', dtype=dtypes.bool)), lambda x,y: x|y),
|
||||
# *** div rules ***
|
||||
(UPat.cvar('x', arg=0) / 0, lambda x: x.const_like(float('nan'))), # 0/0 -> nan
|
||||
((UPat.var("x") * 0) / 0, lambda x: x.const_like(float('nan'))), # (x*0)/0 -> nan
|
||||
# can be wrong if x or x2 is 0
|
||||
(UPat.var("x") / UPat.var("x"), lambda x: x.const_like(1)), # x/x -> 1
|
||||
((UPat.var("x") * UPat.var("x2")) / UPat.var("x2"), lambda x,x2: x), # (x*x2)/x2 -> x
|
||||
# x*0 -> 0 or 0*x -> 0
|
||||
# if x is nan or inf it should render the nan value.
|
||||
# NOTE: this can be wrong for loaded NaN
|
||||
(UPat.var("x") * 0, lambda x: x.const_like(float("nan") if x.op is Ops.CONST
|
||||
and isinstance(x.arg, float) and (math.isnan(x.arg) or math.isinf(x.arg)) else 0)),
|
||||
# *** cast/bitcast ***
|
||||
(UPat(Ops.CAST, name="root", src=(UPat.cvar("c"),)), lambda root, c: root.const_like(c.arg)),
|
||||
(UPat((Ops.CAST, Ops.BITCAST), name="root"), lambda root: root.src[0] if root.dtype == root.src[0].dtype else None),
|
||||
(UPat(Ops.BITCAST, name="root", src=(UPat.cvar("c"),)), fold_bitcast),
|
||||
# b.cast(a).cast(b) -> b if a preserves all values in b
|
||||
(UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x if x.dtype == b.dtype and can_lossless_cast(b.dtype, a.dtype) else None),
|
||||
(UPat.var("x").cast(dtypes.bool), lambda x: x != 0),
|
||||
# ** pow **
|
||||
(UPat.var("x").alu(Ops.POW, UPat.cvar("c")), simplify_pow),
|
||||
# positive const ** x
|
||||
(UPat.cvar("c").alu(Ops.POW, UPat.var("x")), lambda c,x: c if c.arg == 1 else (x*math.log2(c.arg)).exp2() if c.arg > 0 else None),
|
||||
# rules for threefry
|
||||
((UPat.var('x', dtypes.uint64)&0xFFFFFFFF).cast(dtypes.uint32), lambda x: x.cast(dtypes.uint32)),
|
||||
(((UPat.var(None, dtypes.uint64)*(1<<32)) | UPat.var('y', dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y),
|
||||
(((UPat.var('x', dtypes.uint64)*(1<<32)) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))//(1<<32), lambda x: x),
|
||||
(((UPat.var(None, dtypes.uint64)<<32) | UPat.var('y', dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y),
|
||||
(((UPat.var('x', dtypes.uint64)<<32) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))>>32, lambda x: x),
|
||||
# ** simple where folding **
|
||||
# a conditional with the same results either way is a noop, also fold const conditionals
|
||||
(UPat.var().where(UPat.var("val"), UPat.var("val")), lambda val: val),
|
||||
(UPat.cvar("gate").where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.arg else c1),
|
||||
# a.where(b.where(c, d), d) -> (a & b).where(c, d)
|
||||
(UPat.var("a").where(UPat.var("b").where(UPat.var("c"), UPat.var("d")), UPat.var("d")), lambda a,b,c,d: (a&b).where(c,d)),
|
||||
])
|
||||
|
||||
# ******** phase 2 builds on phase 1, it includes the old "symbolic", rules that match deeper ********
|
||||
|
||||
def lt_folding(x:UOp, c:int) -> UOp|None:
|
||||
p, np = partition(x.split_uop(Ops.ADD), lambda u: u.const_factor() == 1)
|
||||
if np and (d:=math.gcd(*[u.const_factor() for u in np], c)) > 1 and 0 <= sum(u.vmin for u in p) and sum(u.vmax for u in p) < d:
|
||||
return unwrap(UOp.usum(*np).divides(d))<(c//d)
|
||||
return None
|
||||
|
||||
def canonicalize_simplex(X:UOp) -> UOp|None:
|
||||
# (X := a0*x0 + a1*x1 + ...) > 0 is equivalent to x0 + x1 + ... > 0 if xi >= 0 and ai > 0 for ints.
|
||||
# returns x0 + x1 + ... in such case, or None if not
|
||||
changed, ret = False, []
|
||||
for u in X.split_uop(Ops.ADD):
|
||||
# assumed the const is the last src of MUL
|
||||
if u.op is Ops.MUL and u.src[1].op is Ops.CONST and u.src[1].arg > 0:
|
||||
changed = True
|
||||
u = u.src[0]
|
||||
if not (u.op in GroupOp.Irreducible and u.vmin >= 0): return None
|
||||
ret.append(u)
|
||||
return UOp.usum(*ret) if changed else None
|
||||
|
||||
def gep_through_wmma(gep:UOp, wmma:UOp) -> UOp|None:
|
||||
out_sz = prod(x[1] for x in wmma.arg[6][-1])
|
||||
wmma_idxs = gep.arg[::out_sz]
|
||||
for i in range(out_sz):
|
||||
if tuple(x-i for x in gep.arg[i::out_sz]) != wmma_idxs: return None
|
||||
tsrcs = []
|
||||
for s,sz in zip(wmma.src, wmma.arg[6]):
|
||||
src_args = []
|
||||
ssz = prod(x[1] for x in sz)
|
||||
for w in wmma_idxs: src_args += list(range((w//out_sz)*ssz, (w//out_sz)*ssz + ssz))
|
||||
tsrcs.append(s.gep(tuple(src_args)))
|
||||
return UOp(Ops.WMMA, gep.dtype, tuple(tsrcs), wmma.arg)
|
||||
|
||||
gep_pushing = PatternMatcher([
|
||||
# GEP/VECTORIZE, GEP/GEP, GEP/CONST
|
||||
(UPat(Ops.GEP, name='g2').f(Ops.GEP, name='g1'),
|
||||
lambda g1, g2: g2.src[0].gep(tuple(g2.arg[g1.arg[i]] for i in range(len(g1.arg))))),
|
||||
(UPat(Ops.STACK, name='vec').f(Ops.GEP, name='gep'),
|
||||
lambda gep, vec: UOp(Ops.STACK, gep.dtype, tuple(vec.src[i] for i in gep.arg)) if len(gep.arg) > 1 else vec.src[gep.arg[0]]),
|
||||
(UPat.cvar("c").f(Ops.GEP, name="gep"), lambda gep, c: gep.const_like(c.arg)),
|
||||
# GEP on void is skipped
|
||||
(UPat(Ops.GEP, src=(UPat(dtype=dtypes.void, name="x"),)), lambda x: x),
|
||||
# GEP in order is removed
|
||||
(UPat(Ops.GEP, name="g"), lambda g: g.src[0] if not isinstance(g.dtype, PtrDType) and g.arg == tuple(range(g.src[0].dtype.count)) else None),
|
||||
# push all GEPs through ALUs for index (TODO: remove this)
|
||||
(UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name='alu').f(Ops.GEP, dtype=dtypes.weakint, name='gep'),
|
||||
lambda gep,alu: UOp(alu.op, alu.dtype.scalar().vec(gep.dtype.count), tuple(x.gep(gep.arg) for x in alu.src), alu.arg) \
|
||||
if not isinstance(gep.dtype, PtrDType) and not isinstance(alu.dtype, PtrDType) else None),
|
||||
# CAT can't be rendered. it's a VECTORIZE on vectors, we expand to a single VECTORIZEs with GEPs (TODO: move this later)
|
||||
(UPat(Ops.VCAT, name="x"), lambda x: UOp(Ops.STACK, x.dtype, tuple(y.gep(i) for y in x.src for i in range(y.dtype.count))) \
|
||||
if not isinstance(x.dtype, PtrDType) else None),
|
||||
# VECTORIZE on same GEP
|
||||
(UPat(Ops.STACK, name="v", src=UPat(Ops.GEP, src=(UPat.var("x"),))), lambda v,x: x.gep(tuple(get_single_element(i.arg) for i in v.src))),
|
||||
# push some GEPs through WMMAs
|
||||
(UPat(Ops.WMMA, name="wmma").f(Ops.GEP, name="gep"), gep_through_wmma),
|
||||
])
|
||||
|
||||
commutative = PatternMatcher([
|
||||
# ** COMMUTATIVE flipping (only for index) **
|
||||
# NOTE: this can break merging vector math by only flipping some of them
|
||||
(UPat(GroupOp.Commutative, dtype=dtypes.weakint, name='x'), lambda x:
|
||||
x.replace(src=x.src[::-1]) if x.src[1].tuplize < x.src[0].tuplize and not x.src[0].tuplize < x.src[1].tuplize else None),
|
||||
])
|
||||
|
||||
symbolic = symbolic_simple+commutative+PatternMatcher([
|
||||
# ** boolean algebra **
|
||||
# TODO: make a more general or folder like simplify_valid
|
||||
(UPat.var("x", dtype=dtypes.bool) | UPat.var("x", dtype=dtypes.bool).logical_not(), lambda x: x.const_like(True)), # x|!x -> True
|
||||
# ** combine terms **
|
||||
(UPat.var("x") * UPat.cvar("c0") + UPat.var("x") * UPat.cvar("c1"), lambda x,c0,c1: x*(c0+c1)), # (x*c0)+(x*c1) -> x*(c0+c1)
|
||||
((UPat.var("y") + UPat.var("x") * UPat.cvar("c0")) + UPat.var("x") * UPat.cvar("c1"), lambda x,y,c0,c1: y+x*(c0+c1)),
|
||||
(UPat.var("x") + UPat.var("x") * UPat.cvar("c"), lambda x,c: x*(c+1)), # (x+x*c)-> x*(c+1)
|
||||
((UPat.var("y") + UPat.var("x")) + UPat.var("x") * UPat.cvar("c"), lambda x,y,c: y+x*(c+1)),
|
||||
((UPat.var("y") + UPat.var("x") * UPat.cvar("c")) + UPat.var("x"), lambda x,y,c: y+x*(c+1)),
|
||||
(UPat.var("x") + UPat.var("x"), lambda x: x*2), # (x+x)-> x*2
|
||||
((UPat.var("y") + UPat.var("x")) + UPat.var("x"), lambda y,x: y+x*2),
|
||||
((UPat.var("x") / UPat.var("x2")) / UPat.var("x3"), lambda x,x2,x3: x/(x2*x3) if x2 is not x3 else None), # (x/x2)/x3 -> x/(x2*x3)
|
||||
(-1 * (UPat.var("x") + UPat.cvar("c")), lambda x,c: (-x)+(-c)), # -(x+c) -> -x + -c
|
||||
(UPat.cvar("y") * (UPat.var("x", dtype=dtypes.weakint) + UPat.cvar("c")), lambda x,y,c: (y*x)+(y*c)), # y*(x+c) -> y*x + y*c
|
||||
# ** where folding **
|
||||
(UPat.var("cond", dtype=dtypes.bool).logical_not().where(UPat.var("t"), UPat.var("f")),
|
||||
lambda cond, t, f: cond.where(f,t) if f.arg is not Invalid else None),
|
||||
# alu of two where with same conds can combine, only do if true branch or false branch is const
|
||||
(UPat(GroupOp.Binary, name="alu", src=(UPat.var("c").where(UPat.var("t"), UPat.var("f")), UPat.var("c").where(UPat.var("tt"), UPat.var("ff")))), \
|
||||
lambda alu,c,t,tt,f,ff: c.where(t.alu(alu.op, tt), f.alu(alu.op, ff)) if t.op == tt.op == Ops.CONST or f.op == ff.op == Ops.CONST else None),
|
||||
# if its a plus we add the associative variation too
|
||||
((UPat.var("y")+UPat.var("c").where(UPat.var("t"), UPat.var("f"))) + UPat.var("c").where(UPat.var("tt"), UPat.var("ff")), \
|
||||
lambda y,c,t,tt,f,ff: y+c.where(t+tt, f+ff) if t.op == tt.op == Ops.CONST or f.op == ff.op == Ops.CONST else None),
|
||||
# ALU/variable min==max -> CONST
|
||||
(UPat({Ops.CMPLT, Ops.CMPNE, Ops.FLOORDIV, Ops.FLOORMOD, Ops.DEFINE_VAR, Ops.BIND, Ops.SPECIAL}, name="x"),
|
||||
lambda x: x.const_like(x.vmin) if x.vmin == x.vmax else None),
|
||||
(UPat(Ops.RANGE, src=(UPat(Ops.CONST,)), name="x"), lambda x: x.const_like(x.vmin) if x.vmin == x.vmax else None),
|
||||
# max folding
|
||||
(UPat.maximum(UPat.var("x"), UPat.var("y")), lambda x,y: x if x.vmin >= y.vmax else y if x.vmax <= y.vmin else None),
|
||||
# TODO: why does this rule break beautiful_mnist?
|
||||
#((UPat.var("x")+UPat.var("z")).maximum(UPat.var("y")+UPat.var("z")), lambda x,y,z: x.maximum(y) + z),
|
||||
# ** two stage ALU folding **
|
||||
*((UPat.var("x").alu(op, UPat.cvar("c1")).alu(op, UPat.cvar("c2")).named("f"),
|
||||
lambda f,x,c1,c2: x.alu(f.op,c1.alu(f.op,c2))) for op in GroupOp.Associative),
|
||||
((UPat.cvar("c0") + UPat.var("x")) < UPat.cvar("c1"), lambda x,c0,c1: x<(c1-c0)), # c0 + x < c1 -> x < c1 - c0
|
||||
# (x//c1)//c2 -> x//(c1*c2) for c2>0
|
||||
((UPat.var("x") // UPat.cvar("c1")) // UPat.cvar("c2"), lambda x,c1,c2: x//(c1*c2) if c2.vmin>0 else None),
|
||||
# ** lt **
|
||||
# c0*x<c1 for positive int c0,c1
|
||||
((UPat.cvar("c0")*UPat.var("x", dtype=dtypes.weakint))<UPat.cvar("c1"),
|
||||
lambda x,c0,c1: x<math.ceil(c1.arg/c0.arg) if c0.arg > 0 and c1.arg > 0 else None),
|
||||
# c0*x<c1 for negative int c0 and non-positive c1
|
||||
((UPat.cvar("c0")*UPat.var("x", dtype=dtypes.weakint))<UPat.cvar("c1"),
|
||||
lambda x,c0,c1: (-x)<(-(math.floor(-c1.arg/-c0.arg))) if c0.arg < 0 and c0.arg != -1 and c1.arg <= 0 else None),
|
||||
# x//d<c -> x<c*d for d>0
|
||||
((UPat.var("x", dtype=dtypes.weakint)//UPat.cvar("d"))<UPat.cvar("c"),
|
||||
lambda x,d,c: x<(c.arg*d.arg) if d.arg > 0 else None),
|
||||
# ** move add/mul consts to end (NOTE: this is still happening before constant folding) **
|
||||
((UPat.var("x") + UPat.cvar("c1")) + UPat.var("y"), lambda x,c1,y: (x+y)+c1),
|
||||
((UPat.var("x") * UPat.cvar("c1")) * UPat.var("y"), lambda x,c1,y: (x*y)*c1),
|
||||
# *** rules from symbolic ***
|
||||
# generic lt folding
|
||||
(UPat.var("x", dtypes.weakint)<UPat.cvar("c"), lambda x,c: lt_folding(x, c.arg) if 0 < c.arg else None),
|
||||
(UPat.var("x", dtypes.weakint)*-1 < UPat.var("y")*-1, lambda x,y: y<x),
|
||||
# canonicalize a simplex with positive coefficients > 0. NOTE: not x < 1 means x > 0
|
||||
((UPat.var("x", dtypes.weakint)<1).ne(True), lambda x: (newx<1).ne(True) if (newx:=canonicalize_simplex(x)) is not None else None),
|
||||
# a range mod its own upper bound is just the range
|
||||
(UPat(Ops.RANGE, src=UPat.var("end"), name="r")%UPat.var("end"), lambda r,end: r),
|
||||
(UPat(Ops.RANGE, src=UPat.var("end"), name="r")//UPat.var("end"), lambda r,end: r.const_like(0)),
|
||||
# cast/long folding
|
||||
# if the intermediate cast doesnt narrow we can do it in one cast
|
||||
(UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x.cast(b.dtype) if can_lossless_cast(x.dtype, a.dtype) else None),
|
||||
(UPat.var('x', dtypes.ints+(dtypes.weakint,)).cast(dtypes.ints+(dtypes.weakint,), name="a").cast(name="b"),
|
||||
lambda x,a,b: x.cast(b.dtype) if a.dtype.min<=x.vmin and x.vmax<=a.dtype.max else None),
|
||||
# try to do math in int instead of long
|
||||
(UPat(GroupOp.Binary, src=(UPat.var("x", dtypes.long), UPat.var("y", dtypes.long)), name="u"), lambda u,x,y:
|
||||
x.cast(dtypes.int).alu(u.op, y.cast(dtypes.int)).cast(u.dtype) if not any(v.overflows(dtypes.int) for v in (u,x,y)) else None),
|
||||
((UPat.var("x", dtypes.weakint) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+c.cast(cast.dtype)),
|
||||
# only RANGE/IF/STORE/KERNEL have side effects
|
||||
(UPat(Ops.AFTER, name="x"), lambda x: x.replace(src=(x.src[0],)+
|
||||
tuple(dedup(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.FUNCTION, Ops.BARRIER, Ops.END, Ops.UNROLL, Ops.LINEAR, Ops.STAGE}
|
||||
else y.src for y in x.src[1:]]))))),
|
||||
# after with 1 src is just src[0]
|
||||
(UPat(Ops.AFTER, src=(UPat.var("s"),)), lambda s: s),
|
||||
# VECTORIZE/CONST
|
||||
(UPat(Ops.STACK, src=UPat(Ops.CONST), name="vec"),
|
||||
lambda vec: UOp.const(vec.dtype, tuple(x.arg for x in vec.src)) if len(vec.src) > 0 else None),
|
||||
])+div_and_mod_symbolic+gep_pushing
|
||||
|
||||
# ******** we take a small aside to "simplify_valid" to rewrite valids ********
|
||||
|
||||
def parse_valid(v:UOp) -> tuple[UOp, bool, int]|None:
|
||||
# if it's X <= c, returns X, True, c
|
||||
# if it's X >= c, returns X, False, c
|
||||
|
||||
if v.op is Ops.CMPNE and v.src[1].op is Ops.CONST and v.src[1].arg == 1 and (s0:=v.src[0]).op is Ops.CMPLT and dtypes.is_int(s0.src[0].dtype):
|
||||
# (X < c).ne(True) -> X >= c
|
||||
return s0.src[0], False, int(s0.src[1].vmin)
|
||||
if v.op is Ops.CMPLT and dtypes.is_int(v.src[0].dtype):
|
||||
# X < c -> X <= c-1
|
||||
return v.src[0], True, int((v.src[1]).vmax)-1
|
||||
return None
|
||||
|
||||
def uop_given_valid(valid:UOp, uop:UOp, try_simplex=True) -> UOp:
|
||||
# return simplified uop (might be the same as input)
|
||||
|
||||
# first, parse valid into {expr: (lower_bound, upper_bound)}
|
||||
bounds:defaultdict[UOp, list[ConstType|None]] = defaultdict(lambda: [None, None])
|
||||
for stmt in valid.split_uop(Ops.AND):
|
||||
if (res:=parse_valid(stmt)) is None: continue
|
||||
expr, is_upper, c = res
|
||||
bounds[expr][int(is_upper)] = c
|
||||
|
||||
# simplify uop given that valid is True
|
||||
all_candidates = []
|
||||
for i,(expr,v) in enumerate(bounds.items()):
|
||||
v0, v1 = (expr.vmin if v[0] is None else v[0], expr.vmax if v[1] is None else v[1])
|
||||
# try checking the whole clause
|
||||
all_candidates.append((expr, UOp.variable(f"fake{i}", v0, v1, expr.dtype)))
|
||||
|
||||
if try_simplex:
|
||||
# every candidate is a set of constrained UOp based on valid, and if every item in a set simplifies the uop into a same output, we rewrite uop
|
||||
candidates = [[all_candidates[-1]]]
|
||||
if expr.op is Ops.ADD and v0 == 1 and all(u.op in GroupOp.Irreducible for u in expr.split_uop(Ops.ADD)):
|
||||
# if the constraint is a simplex: X0 + X1 + ... > 0, we can check if all Xi > 0 simplify into the same output
|
||||
candidates.append([(Xi, UOp.variable(f"fake{i}", 1, Xi.vmax, Xi.dtype)) for Xi in expr.split_uop(Ops.ADD)])
|
||||
|
||||
for candidate in candidates:
|
||||
# if every branch in candidate gives the same simplified uop, we can rewrite the uop
|
||||
newuops = [uop.substitute({X:newX}) for X,newX in candidate]
|
||||
if any(u is uop for u in newuops): continue # if any branch doesnt appear in uop, skip
|
||||
newuops = [u.simplify().substitute({newX:X}).simplify() for (X,newX),u in zip(candidate,newuops)]
|
||||
if all_same(newuops): uop = newuops[0]
|
||||
elif uop.op is Ops.STACK and len(uop.src) == 2:
|
||||
if all_same([uops.src[0] for uops in newuops]): uop = uop.replace(src=(newuops[0].src[0], uop.src[1]))
|
||||
if all_same([uops.src[1] for uops in newuops]): uop = uop.replace(src=(uop.src[0], newuops[0].src[1]))
|
||||
|
||||
# try all the valids together (but only the whole expressions)
|
||||
if (s_uop:=uop.substitute(sub_dict:=dict(all_candidates))) is not uop:
|
||||
uop = s_uop.simplify().substitute({newX:X for X,newX in sub_dict.items()}).simplify()
|
||||
return uop
|
||||
|
||||
def _valid_priority(v: UOp, valids:list[UOp]) -> int:
|
||||
# we want valid that's in other valids' parents to be first, so it's more likely the other valids get simplified
|
||||
return sum(-1 if (res:=parse_valid(v)) is not None and res[0] in other.toposort() else 0 for other in valids)
|
||||
|
||||
def simplify_valid(valid:UOp) -> UOp|None:
|
||||
if valid.op_in_backward_slice_with_self(Ops.INDEX): return None # this should only be for indexing, skip if there's a INDEX
|
||||
ret:list[UOp] = []
|
||||
valids = list(valid.split_uop(Ops.AND))
|
||||
valids = sorted(valids, key=lambda v: _valid_priority(v, valids))
|
||||
for stmt in dedup(valids):
|
||||
if ret: stmt = uop_given_valid(UOp.uprod(*ret), stmt)
|
||||
ret.append(stmt)
|
||||
return UOp.uprod(*ret) if ret != valids else None
|
||||
|
||||
# ******** phase 3 is the complete symbolic ********
|
||||
|
||||
def reduce_mul_chain(r:UOp) -> UOp|None:
|
||||
if r.arg[0] not in {Ops.ADD, Ops.MAX}: return None
|
||||
if r.dtype != r.src[0].dtype: return None
|
||||
inside, outside = [], []
|
||||
for m in r.src[0].split_uop(Ops.MUL):
|
||||
m_parents = m.backward_slice
|
||||
if m not in r.src[1:] and all(r not in m_parents for r in r.src[1:]) and (r.arg[0] != Ops.MAX or m.vmin >= 0): outside.append(m)
|
||||
else: inside.append(m)
|
||||
if len(outside) == 0: return None
|
||||
return r.replace(src=(prod(inside) if len(inside) else r.src[0].const_like(1),)+r.src[1:])*prod(outside)
|
||||
|
||||
def drop_and_clauses(cond:UOp, x:UOp, i:UOp) -> UOp|None:
|
||||
keep, drop = partition(cond.split_uop(Ops.AND), lambda c: any(r in x.ranges for r in c.ranges))
|
||||
return UOp.const(dtypes.bool, True).uprod(*keep).where(x, i) if drop else None
|
||||
pm_drop_and_clauses = PatternMatcher([(invalid_gate, drop_and_clauses)])
|
||||
|
||||
# move conditions from where to load's valid, drop clauses already in load
|
||||
def where_on_load(cond:UOp, buf:UOp, idx:UOp, or_cast:UOp) -> UOp|None:
|
||||
where_clauses, load_valid = list(cond.split_uop(Ops.AND)), idx.get_valid()
|
||||
in_load = set(load_valid.split_uop(Ops.AND))
|
||||
idx_index = {u for u in idx.backward_slice_with_self if u.op is Ops.INDEX}
|
||||
# can move if: condition's ranges are subset of idx's ranges, and no data dependent INDEX (only idx's INDEX allowed)
|
||||
def can_move(c:UOp) -> bool:
|
||||
return c.ranges.keys() <= idx.ranges.keys() and all(u in idx_index for u in c.backward_slice_with_self if u.op is Ops.INDEX)
|
||||
moved, keep = partition([c for c in where_clauses if c not in in_load], can_move)
|
||||
if len(keep) == len(where_clauses): return None
|
||||
idx = buf.index(idx.get_idx().valid(load_valid.uprod(*moved)))
|
||||
ret_idx = idx.cast(or_cast.dtype) if or_cast.op is Ops.CAST else idx
|
||||
return UOp.const(dtypes.bool, True).uprod(*keep).where(ret_idx, ret_idx.const_like(0))
|
||||
|
||||
# where after gated load becomes alt value, TODO: this is sort of duplicated with rules in devectorizer
|
||||
pm_move_where_on_load = PatternMatcher([
|
||||
(UPat.var("cond").where(UPat.var("buf").index(UPat.var("idx")).or_casted("or_cast"), 0), where_on_load),
|
||||
(UPat.var("cond").where(0, UPat.var("buf").index(UPat.var("idx")).or_casted("or_cast")),
|
||||
lambda cond,buf,idx,or_cast: where_on_load(cond.logical_not(),buf,idx,or_cast)),
|
||||
])
|
||||
|
||||
def gated_given_valid(cond:UOp, x:UOp, i:UOp) -> UOp|None:
|
||||
if x.dtype.scalar() is not dtypes.weakint: return None
|
||||
# Skip if x contains DIV/MOD AND IMAGE mode is enabled -> image index e.g. openpilot
|
||||
if IMAGE.value > 0 and x.op_in_backward_slice_with_self(Ops.CDIV, Ops.CMOD, Ops.FLOORDIV, Ops.FLOORMOD): return None
|
||||
return cond.where(uop_given_valid(cond, x, try_simplex=False), i)
|
||||
|
||||
# TODO: this is O(number of WHERE * number of node)
|
||||
# def fold_where_closure(cond:UOp, t:UOp, f:UOp) -> UOp|None:
|
||||
# """In cond.where(t, f), fold nested cond.where(a, b) -> a in t, -> b in f"""
|
||||
# def is_valid_where(u:UOp) -> bool: return u.op is Ops.WHERE and u.src[0] is cond and Invalid not in (u.src[1].arg, u.src[2].arg)
|
||||
# t_subs, f_subs = {u: u.src[1] for u in t.toposort() if is_valid_where(u)}, {u: u.src[2] for u in f.toposort() if is_valid_where(u)}
|
||||
# if not t_subs and not f_subs: return None
|
||||
# new_t, new_f = t.substitute(t_subs).simplify() if t_subs else t, f.substitute(f_subs).simplify() if f_subs else f
|
||||
# return None if new_t is t and new_f is f else cond.where(new_t, new_f)
|
||||
|
||||
pm_simplify_valid = PatternMatcher([
|
||||
# simplify valid
|
||||
(UPat(Ops.AND, name="valid"), simplify_valid),
|
||||
(invalid_gate, gated_given_valid),
|
||||
])
|
||||
|
||||
# this is symbolic 2.0
|
||||
REMOVE_FROM_SINK_LIKE = {Ops.UNROLL, Ops.NOOP, Ops.STACK, Ops.SINK, Ops.GROUP}
|
||||
pm_clean_up_group_sink = PatternMatcher([
|
||||
# clean up GROUP/SINK
|
||||
(UPat(Ops.GROUP, src=(UPat.var("x"),)), lambda x: x),
|
||||
(UPat((Ops.SINK, Ops.GROUP), name="root"),
|
||||
lambda root: UOp(root.op, root.dtype, tuple(flatten(x.src if x.op in REMOVE_FROM_SINK_LIKE else (x,) for x in root.src)), root.arg)
|
||||
if any(x.op in REMOVE_FROM_SINK_LIKE for x in root.src) else None),
|
||||
])
|
||||
|
||||
sym = symbolic+pm_simplify_valid+PatternMatcher([
|
||||
# reorder ALU/VECTORIZE
|
||||
(UPat(GroupOp.ALU, src=(UPat(Ops.STACK, src=UPat(name='x')), UPat(Ops.STACK, src=UPat(name='y'))), name='alu'),
|
||||
lambda x,y,alu: UOp(Ops.STACK, alu.dtype, (UOp(alu.op, alu.dtype.scalar(), (x,y)),)*alu.dtype.count)),
|
||||
# ** where **
|
||||
# # fold nested where with same condition: in cond.where(t,f), cond.where(a,b)->a in t, ->b in f
|
||||
# (UPat.var("cond").where(UPat.var("t"), UPat.var("f")), fold_where_closure),
|
||||
# push cast to branches
|
||||
(UPat.var("s").where(UPat.var("a"), UPat.var("b")).cast().named("cast"), lambda s,a,b,cast: s.where(a.cast(cast.dtype), b.cast(cast.dtype))),
|
||||
# ** pow **
|
||||
((UPat(Ops.POW, name="p"), lambda p: xpow(*p.src))),
|
||||
# ** load/store folding **
|
||||
(UPat.store(UPat(Ops.INDEX, name="index"), UPat.load(UPat(Ops.INDEX, name="index"))), lambda index: UOp(Ops.NOOP)),
|
||||
(UPat.store(UPat(Ops.INDEX, name="index"), UPat.var("gate").where(UPat.var("alt"),
|
||||
UPat.load(UPat(Ops.INDEX, name="index")))),
|
||||
lambda index, gate, alt: UOp.store(index.src[0].index(gate.where(index.src[1], UOp.invalid())), alt)),
|
||||
# fold gated LOAD/STORE
|
||||
(UPat(Ops.STORE, src=(UPat(), invalid_pat)), lambda i: UOp(Ops.NOOP)),
|
||||
# store of where with invalid -> gated store
|
||||
(UPat(Ops.STORE, src=(UPat(Ops.INDEX, name="index"), UPat.var("cond").where(UPat.var("val"), invalid_pat))),
|
||||
lambda index, cond, val, i: UOp.store(index.src[0].index(cond.where(index.src[1], UOp.invalid())), val)),
|
||||
((UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()), # 1/(x^c) -> (1/x)^c
|
||||
((UPat.var("x") * UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()*x.reciprocal()),
|
||||
((UPat.var("x") * UPat.cvar("c")).reciprocal(), lambda x,c: x.reciprocal()*c.reciprocal()), # 1/(x*c) -> (1/c)*(1/x)
|
||||
(UPat.var("x") * ((1+UPat.var("x")).reciprocal().named("d")), lambda x,d: 1-d), # x*/(1+x) -> 1-1/(1+x)
|
||||
(UPat.var("x") * ((1+UPat.var("x")).reciprocal().named("d")*UPat.var("y")), lambda x,y,d: y*(1-d)),
|
||||
(UPat.var("x") * ((1+UPat.var("x")).reciprocal().named("d")+UPat.var("y")), lambda x,y,d: (1-d)+x*y),
|
||||
# move const multiply after REDUCE (NOTE: the mul chain can do this, but only if it's a same dtype reduce)
|
||||
((UPat.var("x")*UPat.cvar("c")).reduce(arg=Ops.ADD, name="r", allow_any_len=True), lambda x,c,r: r.replace(src=(x,)+r.src[1:])*c.arg),
|
||||
# reduce mul chain, move muls after the reduce
|
||||
(UPat(Ops.MUL).reduce(name="r", allow_any_len=True), reduce_mul_chain),
|
||||
# ** combine terms (opinionated) **
|
||||
(-1 * (UPat.var("x") + UPat.var("y")), lambda x,y: (-x)+(-y)), # -(x+y) -> -x + -y
|
||||
# (x+y)*c -> x*c+y*c. only for int, float has inf*0=nan issue
|
||||
((UPat.var("x", dtypes.weakint) + UPat.var("y")) * UPat.cvar("c"), lambda x,y,c: x*c+y*c),
|
||||
])+pm_clean_up_group_sink
|
||||
168
tinygrad_repo/tinygrad/uop/upat.py
Normal file
168
tinygrad_repo/tinygrad/uop/upat.py
Normal file
@@ -0,0 +1,168 @@
|
||||
from typing import Any, Callable
|
||||
import itertools, inspect, functools, types
|
||||
from tinygrad.helpers import partition, dedup, Context
|
||||
from tinygrad.uop.ops import UPat, UOp, Ops, PatternMatcher, graph_rewrite, deconstruct_function
|
||||
|
||||
class UPatCompileError(Exception): pass
|
||||
|
||||
# **** UPat compiled ****
|
||||
|
||||
def _get_clause(self:UPat, base:UOp, depth=0) -> UOp:
|
||||
if self.is_any:
|
||||
assert len(self.src) == 1
|
||||
return UOp(Ops.AND, src=(UOp(Ops.OR, src=tuple(_get_clause(s, base, depth) for s in self.src[0])),))
|
||||
# build the and_clause for acceptance
|
||||
and_clause:list[UOp] = []
|
||||
if self.op is not None:
|
||||
if len(self.op) > 1: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=tuple(int(x) for x in self.op))), arg="{0}.op in {1}"))
|
||||
else: and_clause.append(UOp(Ops.CUSTOM, src=(base,), arg="{0}.op == "+str(self.op[0].value)))
|
||||
if self.arg is not None:
|
||||
if isinstance(self.arg, int): and_clause.append(UOp(Ops.CUSTOM, src=(base,), arg="{0}.arg == "+str(int(self.arg))))
|
||||
else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=self.arg)), arg="{0}.arg == {1}"))
|
||||
if self.strict_length or self.required_len > 0:
|
||||
and_clause.append(UOp(Ops.CUSTOM, src=(base,), arg=("len({0}.src)"+(" == " if self.strict_length else " >= ")+str(self.required_len))))
|
||||
if self.name is not None: and_clause.append(UOp(Ops.STORE, src=(UOp(Ops.DEFINE_VAR, arg=self.name), base)))
|
||||
if self.match_dtype is not None:
|
||||
if len(self.match_dtype) > 1:
|
||||
and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=tuple(self.match_dtype))), arg="({0}.dtype in {1} or {0}.dtype._scalar in {1})"))
|
||||
else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=self.match_dtype[0])), arg="({0}.dtype == {1} or {0}.dtype._scalar == {1})"))
|
||||
if self.match_tag is not None:
|
||||
if len(self.match_tag) > 1:
|
||||
and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=tuple(self.match_tag))), arg="{0}.tag in {1}"))
|
||||
else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=self.match_tag[0])), arg="{0}.tag == {1}"))
|
||||
if self.src is not None:
|
||||
# single match
|
||||
if len(self.src) == 1 and isinstance(self.src[0], tuple):
|
||||
and_clause += [_get_clause(s, base.gep(i), depth) for i,s in enumerate(self.src[0])]
|
||||
# repeat match
|
||||
elif len(self.src) == 1 and isinstance(self.src[0], itertools.repeat):
|
||||
it = UOp(Ops.NOOP, arg=f"ituop{depth}")
|
||||
match = _get_clause(next(self.src[0]), it, depth+1)
|
||||
and_clause.append(UOp(Ops.RANGE, src=(match, it, base), arg="all([{0} for {1} in {2}.src])"))
|
||||
# multi match (fork)
|
||||
elif len(self.src) > 1 and all(isinstance(x, tuple) for x in self.src):
|
||||
fork_cond = [UOp(Ops.AND, src=tuple([_get_clause(s, base.gep(i), depth) for i,s in enumerate(ss)])) for ss in self.src]
|
||||
and_clause.append(UOp(Ops.OR, src=tuple(fork_cond)))
|
||||
else: raise RuntimeError("broken")
|
||||
return UOp(Ops.AND, src=tuple(and_clause))
|
||||
|
||||
# *** pattern matcher ***
|
||||
|
||||
def do_process_and(a:UOp) -> UOp|None:
|
||||
found = False
|
||||
new_src:list[UOp] = []
|
||||
or_clause:list[UOp] = []
|
||||
|
||||
# remove any nested ANDs, extract or clauses
|
||||
for x in a.src:
|
||||
if x.op is Ops.AND:
|
||||
new_src.extend(x.src)
|
||||
found = True
|
||||
elif x.op is Ops.OR: or_clause.append(x)
|
||||
else: new_src.append(x)
|
||||
|
||||
# too big to compile
|
||||
if len(or_clause) >= 4: raise UPatCompileError("too big to compile")
|
||||
|
||||
# one or clause max
|
||||
if len(or_clause) > 1:
|
||||
# need the product of the or clauses
|
||||
or_clause = [UOp(Ops.OR, src=tuple([UOp(Ops.AND, src=x) for x in itertools.product(*[x.src for x in or_clause])]))]
|
||||
found = True
|
||||
|
||||
# handle stores
|
||||
stores, new_src = partition(new_src, lambda x: x.op is Ops.STORE)
|
||||
if len(stores):
|
||||
if len(or_clause):
|
||||
# push stores to the top if we have an or_clause
|
||||
assert len(or_clause) == 1 and all(x.op is Ops.AND for x in or_clause[0].src)
|
||||
or_clause = [UOp(Ops.OR, src=tuple([x.replace(src=x.src+tuple(stores)) for x in or_clause[0].src]))]
|
||||
found = True
|
||||
else:
|
||||
# check for duplicate stores
|
||||
dict_stores: dict[UOp, UOp] = {}
|
||||
for a in stores:
|
||||
if a.src[0] in dict_stores:
|
||||
# duplicate store is a compare
|
||||
new_src.append(UOp(Ops.CMPNE, src=(dict_stores[a.src[0]], a.src[1])))
|
||||
found = True
|
||||
else:
|
||||
dict_stores[a.src[0]] = a.src[1]
|
||||
# put the stores back
|
||||
for k,v in dict_stores.items(): new_src.append(UOp(Ops.STORE, src=(k,v)))
|
||||
|
||||
# reassemble, if there's any deduping to do, do it
|
||||
if len(dretand:=dedup(new_src+or_clause)) != len(new_src)+len(or_clause): found = True
|
||||
return UOp(Ops.AND, src=tuple(dretand)) if found else None
|
||||
|
||||
# processor
|
||||
pm_proc = PatternMatcher([(UPat(Ops.AND, name="a"), do_process_and)], compiled=False)
|
||||
|
||||
# renderer
|
||||
def wrap(ctx, x) -> UOp:
|
||||
ctx[ret:=f"a{len(ctx)}"] = x.arg
|
||||
return UOp(Ops.NOOP, arg=ret)
|
||||
|
||||
pm_renderer = PatternMatcher([
|
||||
(UPat(Ops.BIND, name="x"), wrap),
|
||||
|
||||
# CMPNE is actually equal
|
||||
(UPat(Ops.CMPNE, name="x"), lambda x: UOp(Ops.CUSTOM, src=x.src, arg="{0} is {1}")),
|
||||
|
||||
# RANGE can't have OR inside it
|
||||
(UPat(Ops.RANGE, src=(UPat(Ops.AND, src=UPat(Ops.NOOP), name="x"), UPat(), UPat()), name="r"),
|
||||
lambda r,x: r.replace(op=Ops.CUSTOM, src=(UOp(Ops.NOOP, arg="(" + ' and '.join(y.arg for y in x.src) + ")"),)+r.src[1:])),
|
||||
|
||||
(UPat(Ops.CUSTOM, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=x.arg.format(*[y.arg for y in x.src]))),
|
||||
(UPat(Ops.GEP, src=UPat(Ops.NOOP, name="x"), name="g"), lambda x,g: x.replace(arg=x.arg+f".src[{g.arg[0]}]"))
|
||||
], compiled=False)
|
||||
|
||||
def _final_render(x:UOp, has_ctx:bool, depth=1) -> list[str]:
|
||||
assert x.op is Ops.AND
|
||||
and_pieces, store_pieces = [], []
|
||||
or_pieces: list[str] = []
|
||||
for s in x.src:
|
||||
if s.op is Ops.OR:
|
||||
assert len(or_pieces) == 0 and len(s.src) >= 1
|
||||
for ss in s.src: or_pieces.extend(_final_render(ss, has_ctx, depth+1))
|
||||
elif s.op is Ops.STORE:
|
||||
assert s.src[0].op is Ops.DEFINE_VAR and s.src[1].op is Ops.NOOP
|
||||
store_pieces.append(f"{s.src[0].arg}={s.src[1].arg}")
|
||||
elif s.op is Ops.NOOP: and_pieces.append(s.arg)
|
||||
else: raise UPatCompileError(f"can't compile this {s}")
|
||||
# if we have an or, render it
|
||||
if len(or_pieces):
|
||||
assert len(store_pieces) == 0
|
||||
and_clause = ' and '.join(and_pieces)
|
||||
return [f"{' '*depth}if {and_clause if len(and_clause) else 'True'}:"] + or_pieces
|
||||
# if we don't, this is a final return
|
||||
store_clause = ', '.join((["ctx=ctx"] if has_ctx else [])+store_pieces)
|
||||
and_clause = ' and '.join(and_pieces + [f"(_ret:=_fxn({store_clause})) is not None"])
|
||||
return [f"{' '*depth}if {and_clause}: return _ret"]
|
||||
|
||||
def _get_code(self:UPat, has_ctx:bool):
|
||||
ret = _get_clause(self, UOp(Ops.NOOP, arg="uop"))
|
||||
try:
|
||||
# TODO: this should be tracked in a "system" rewrite, not untracked or tracked with kernel
|
||||
with Context(TRACK_MATCH_STATS=0):
|
||||
ret = graph_rewrite(ret, pm_proc, name="process UPat")
|
||||
dyn_lookup: dict[str, Any] = {}
|
||||
out = graph_rewrite(ret, pm_renderer, ctx=dyn_lookup, name="compile UPat")
|
||||
rendered = _final_render(out, has_ctx)
|
||||
except UPatCompileError:
|
||||
#print("FAILED", self, self.location)
|
||||
return None
|
||||
return '\n'.join([f"# match for {self.location}", "def compiled_match(uop, ctx):"] + rendered + [" return None"]), dyn_lookup
|
||||
|
||||
@functools.cache
|
||||
def upat_compile(self:UPat, fxn) -> Callable|None:
|
||||
real_fxn = types.FunctionType(*deconstruct_function(fxn))
|
||||
# UOps used here don't follow the spec
|
||||
with Context(SPEC=0): code = _get_code(self, 'ctx' in inspect.signature(real_fxn).parameters)
|
||||
if code is None: return None
|
||||
code_str, dyn_lookup = code
|
||||
globs = dyn_lookup.copy()
|
||||
globs["_fxn"] = real_fxn
|
||||
namespace: dict = {}
|
||||
exec(code_str, globs, namespace) # pylint: disable=W0122
|
||||
return namespace["compiled_match"]
|
||||
76
tinygrad_repo/tinygrad/uop/validate.py
Normal file
76
tinygrad_repo/tinygrad/uop/validate.py
Normal file
@@ -0,0 +1,76 @@
|
||||
from typing import Callable, cast
|
||||
from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, python_alu
|
||||
from tinygrad.dtype import dtypes, Invalid
|
||||
from tinygrad.helpers import cpu_profile
|
||||
import z3
|
||||
|
||||
# older versions of z3 dont have some operators like & overloaded
|
||||
if z3.get_version() < (4, 12, 4, 0):
|
||||
raise ImportError("bounds checking requires z3 >= 4.12.4, use CHECK_OOB=0 to disable, or \"pip install 'z3-solver>=4.12.4\"")
|
||||
|
||||
# IDIV is truncated division but z3 does euclidian division (floor if b>0 ceil otherwise); mod by power of two sometimes uses Ops.AND
|
||||
def z3_cdiv(a:z3.ArithRef, b:z3.ArithRef) -> z3.ArithRef:return z3.If((a<0), z3.If(0<b, (a+(b-1))/b, (a-(b+1))/b), a/b)
|
||||
def z3_floordiv(a:z3.ArithRef, b:z3.ArithRef) -> z3.ArithRef: return z3.If(b > 0, a/b, (-a)/(-b))
|
||||
def z3_xor(a:z3.ExprRef, b:z3.ExprRef) -> z3.ExprRef:
|
||||
if isinstance(a, z3.BoolRef): return a^b
|
||||
# x ^ -1 = -(x+1), i.e. bitwise NOT
|
||||
if isinstance(b, z3.IntNumRef) and b.as_long() == -1: return -(a+1)
|
||||
if isinstance(a, z3.IntNumRef) and a.as_long() == -1: return -(b+1)
|
||||
raise RuntimeError(f"z3 int XOR only supports XOR with -1, got {a=} {b=}")
|
||||
z3_alu: dict[Ops, Callable[..., z3.ExprRef]] = python_alu | {Ops.CMOD: lambda a,b: a-z3_cdiv(a,b)*b, Ops.CDIV: z3_cdiv, Ops.FLOORDIV: z3_floordiv,
|
||||
Ops.FLOORMOD: lambda a,b: a-z3_floordiv(a,b)*b,
|
||||
Ops.SHR: lambda a,b: a/(2**b.as_long()), Ops.SHL: lambda a,b: a*(2**b.as_long()),
|
||||
Ops.AND: lambda a,b: a%(b+1) if isinstance(b, z3.ArithRef) else a&b, Ops.WHERE: z3.If, Ops.XOR: z3_xor, Ops.MAX: lambda a,b: z3.If(a<b, b, a),}
|
||||
def create_bounded(name:str, vmin:int, vmax:int, z3ctx:z3.Context) -> tuple[z3.ArithRef, z3.BoolRef]:
|
||||
return (s:=z3.Int(name, ctx=z3ctx)), (vmin <= s)&(s <= vmax)
|
||||
|
||||
z3_renderer = PatternMatcher([
|
||||
(UPat.var("cond").where(UPat.var("x"), UPat.const(dtypes.weakint, Invalid)), lambda x,cond,ctx: (ctx[1][x], ctx[1][cond])),
|
||||
# variables
|
||||
(UPat(Ops.SPECIAL, name="x"), lambda x,ctx: create_bounded(x.arg, 0, ctx[1][x.src[0]]-1, ctx[0])),
|
||||
(UPat(Ops.DEFINE_VAR, name="x"), lambda x,ctx: create_bounded(x.arg[0], x.arg[1], x.arg[2], ctx[0])),
|
||||
(UPat(Ops.RANGE, name="x"), lambda x,ctx: create_bounded(x.render(simplify=False), 0, ctx[1][x.src[0]]-1, ctx[0])),
|
||||
# loads are variables bounded by the min/max of the dtype. non-pointer INDEX is also a LOAD
|
||||
(UPat((Ops.LOAD, Ops.INDEX), dtypes.ints+(dtypes.weakint,), name="x"), lambda x,ctx:
|
||||
create_bounded(f"load{len(ctx[1])}", x.dtype.min, x.dtype.max, ctx[0])),
|
||||
(UPat((Ops.LOAD, Ops.INDEX), dtypes.bool), lambda ctx: (z3.Bool(f"load{len(ctx[1])}", ctx=ctx[0]), None)),
|
||||
# constants
|
||||
(UPat(Ops.CONST, arg=Invalid), lambda ctx: (z3.Int("Invalid", ctx=ctx[0]), None)),
|
||||
(UPat(Ops.CONST, dtypes.ints+(dtypes.weakint,), name="x"), lambda x,ctx: (z3.IntVal(x.arg, ctx=ctx[0]), None)),
|
||||
(UPat(Ops.CONST, dtypes.bool, name="x"), lambda x,ctx: (z3.BoolVal(x.arg, ctx=ctx[0]), None)),
|
||||
# casts from floats create new variables
|
||||
(UPat(Ops.CAST, dtypes.ints+(dtypes.weakint,), src=(UPat(dtype=dtypes.floats),), name="x"), lambda x,ctx:
|
||||
create_bounded(f"cast{len(ctx[1])}", x.dtype.min, x.dtype.max, ctx[0])),
|
||||
# A comparison between floats introduces a new bool variable
|
||||
(UPat(GroupOp.Comparison, src=UPat(dtype=dtypes.floats)), lambda ctx: (z3.Bool(f"float_cmp{len(ctx[1])}", ctx=ctx[0]), None)),
|
||||
# casts from bool/int to int/bool
|
||||
(UPat(Ops.CAST, dtypes.ints+(dtypes.weakint,),src=(UPat.var("x", dtypes.bool),)), lambda x,ctx: (z3.If(ctx[1][x], 1, 0), None)),
|
||||
(UPat(Ops.CAST, dtypes.ints+(dtypes.weakint,), src=(UPat.var("x", dtypes.ints+(dtypes.weakint,)),)), lambda x,ctx: (ctx[1][x], None)),
|
||||
(UPat(Ops.CAST, dtypes.bool, name="x"), lambda x,ctx: (ctx[1][x.src[0]]!=0, None)),
|
||||
(UPat(GroupOp.ALU, name="x"), lambda x,ctx: (z3_alu[x.op](*(ctx[1][s] for s in x.src)), None)),
|
||||
])
|
||||
|
||||
def uops_to_z3(solver:z3.Solver, *uops: UOp) -> list[z3.ExprRef]:
|
||||
lst = list(UOp.sink(*uops).toposort(gate=lambda x: x.dtype.scalar() in dtypes.ints+(dtypes.bool, dtypes.weakint) or x.op is Ops.SINK))[:-1]
|
||||
z3map: dict[UOp, z3.ExprRef] = {}
|
||||
for u in lst:
|
||||
z3_rewritten = z3_renderer.rewrite(u, ctx=(solver.ctx, z3map))
|
||||
if z3_rewritten is None: raise NotImplementedError(f"{u.op} is not supported by z3")
|
||||
new_u, constraint = cast(tuple[z3.ArithRef, z3.BoolRef|None], z3_rewritten)
|
||||
if constraint is not None: solver.add(constraint)
|
||||
z3map[u] = new_u
|
||||
assert all(u in z3map for u in uops), "UOp failed to rewrite to z3!"
|
||||
return [z3map[u] for u in uops]
|
||||
|
||||
def validate_index_with_z3(sz:int, idx:UOp, gate:UOp) -> bool:
|
||||
solver = z3.Solver(ctx=z3.Context())
|
||||
z3_idx, z3_mask = uops_to_z3(solver, idx, gate)
|
||||
solver.add(z3_mask)
|
||||
with cpu_profile("validate index with z3", "TINY"):
|
||||
match solver.check((z3_idx<0)|(sz<=z3_idx)):
|
||||
case z3.unsat: return True
|
||||
case z3.sat: print(f"# OUT OF BOUNDS ACCESS: at {solver.model()} INDEX not in 0 - {sz}\nconstraints = {solver}")
|
||||
case z3.unknown: print(f"# UNKNOWN RESULT FROM Z3: {solver.reason_unknown()}\nconstraints = {solver}")
|
||||
print(f"idx={idx.render(simplify=False)}")
|
||||
print(f"mask={gate.render(simplify=False)}")
|
||||
return False
|
||||
Reference in New Issue
Block a user