tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 08:36:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

remove args from min/max [run_process_replay] (#6430)

Browse Source 
* remove args from min/max [run_process_replay]

* it's a ConstType

* sconst_like unused

* any const is fine
This commit is contained in:
George Hotz
2024-09-09 18:18:20 +08:00
committed by GitHub
parent f4e83b30b4
commit 42e5c8335e
5 changed files with 29 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_linearizer_dumb.py
+2 -2
View File
@@ -216,8 +216,8 @@ class TestLinearizerDumb(unittest.TestCase):
print(prg.src)
store_idxs = [x.src[1] for x in k.uops if x.op is UOps.STORE]
for i in range(len(store_idxs) - 1):
first_bounds = store_idxs[i].vmin.arg+store_idxs[i].vmax.arg
next_bounds = store_idxs[i+1].vmin.arg+store_idxs[i+1].vmax.arg
first_bounds = store_idxs[i].vmin+store_idxs[i].vmax
next_bounds = store_idxs[i+1].vmin+store_idxs[i+1].vmax
assert first_bounds < next_bounds, f"first stored (max) idx {first_bounds} then {next_bounds}!"
if __name__ == '__main__':
test/unit/test_uop_symbolic.py
+1 -1
View File
@@ -23,7 +23,7 @@ def render(self) -> Tuple[str, ConstType, ConstType]:
code_for_op = {**CStyleLanguage().code_for_op, BinaryOps.IDIV: lambda a,b,dtype: f"({a}//{b})"}
rewritten_uop = [uop for uop in uops if uop.op is UOps.STORE][0].src[-1]
fxn = TestRenderer().render("", uops)
return fxn.split("data0[0] = ")[1].split(";")[0], rewritten_uop.vmin.arg, rewritten_uop.vmax.arg
return fxn.split("data0[0] = ")[1].split(";")[0], rewritten_uop.vmin, rewritten_uop.vmax
def NumNode(val): return UOp.const(dtypes.int, val)
def Variable(expr, nmin, nmax):
tinygrad/codegen/uopgraph.py
+5 -5
View File
@@ -99,7 +99,7 @@ def mod_folding(x:UOp, c:int) -> Optional[UOp]:
# simplify x % c, None means no change
# simple cancel mod case
if 0 < c and 0 <= x.vmin.arg and (quotient:=x.vmin.arg//c) == x.vmax.arg//c: return x-quotient*c
if 0 < c and 0 <= x.vmin and (quotient:=x.vmin//c) == x.vmax//c: return x-quotient*c
remainder, something_changed = [], False
for u in _get_add_chain(x):
@@ -117,7 +117,7 @@ def div_folding(x:UOp, c:int) -> Optional[UOp]:
# simplify x // c, None means no change
# simple cancel div case
if 0 <= x.vmin.arg and x.vmax.arg < c: return x.const_like(0)
if 0 <= x.vmin and x.vmax < c: return x.const_like(0)
quotient, remainder, rem_const, something_changed, gcd, divisor = [], [], 0, False, c, 1
for u in _get_add_chain(x):
@@ -253,7 +253,7 @@ constant_folder = PatternMatcher([
NOp(UOps.LOAD, src=(NOp.var("buf"), NOp.var('add')+NOp.var('mul')*NOp(UOps.RANGE, name="rng")), name="ld"), NOp.const(None, 0.0)),),
arg=BinaryOps.ADD, name="reduce", allow_any_len=True), index_collapse),
# max folding
(NOp.max(NOp.var('x'), NOp.var('y')), lambda x,y: x if x.vmin.arg >= y.vmax.arg else y if x.vmax.arg <= y.vmin.arg else None),
(NOp.max(NOp.var('x'), NOp.var('y')), lambda x,y: x if x.vmin >= y.vmax else y if x.vmax <= y.vmin else None),
# GEP/CAST const rules
(NOp(UOps.GEP, src=(NOp.cvar("c"),), name="root"), lambda root, c: root.const_like(c.arg)),
(UPat(UOps.CAST, name="root", src=UPat(UOps.CONST, name="c")), lambda root, c: root.const_like(c.arg)),
@@ -281,7 +281,7 @@ constant_folder = PatternMatcher([
# NOTE: this can be wrong for loaded NaN
(NOp.var('x') * 0, lambda x: x.const_like(float('nan') if isinstance(x.arg, float) and (math.isnan(x.arg) or math.isinf(x.arg)) else 0)),
# min==max -> CONST (slow!)
(UPat({UOps.ALU, UOps.DEFINE_VAR}, name='x'), lambda x: x.const_like(x.vmin.arg) if x.vmin.arg == x.vmax.arg else None),
(UPat({UOps.ALU, UOps.DEFINE_VAR}, name='x'), lambda x: x.const_like(x.vmin) if x.vmin == x.vmax else None),
# ** load/store folding **
(NOp.store(NOp.var("buf"), NOp.var("idx"), NOp.load(NOp.var("buf"), NOp.var("idx"))), lambda buf,idx:UOp(UOps.NOOP)),
# ** two stage add/mul folding **
@@ -294,7 +294,7 @@ constant_folder = PatternMatcher([
lambda x,c0,c1: x.lt(math.ceil(c1.arg/c0.arg)) if dtypes.is_int(x.dtype) and c0.arg > 0 and c1.arg > 0 else None),
# mul add lt
(((NOp.cvar('c0')*NOp.var('x'))+NOp.var('x2')).lt(NOp.cvar('c1')),
lambda x,x2,c0,c1: x.lt(c1//c0) if c1.arg % c0.arg == 0 and c0.arg > x2.vmax.arg and x2.vmin.arg >= 0 else None),
lambda x,x2,c0,c1: x.lt(c1//c0) if c1.arg % c0.arg == 0 and c0.arg > x2.vmax and x2.vmin >= 0 else None),
# generic lt folding (using div)
(NOp.var('x').lt(NOp.cvar('c')),
lambda x,c: lt_folding(x, c.arg) if 0 < c.arg and dtypes.is_int(x.dtype) and not dtypes.is_unsigned(x.dtype) else None),
tinygrad/ops.py
+18 -23
View File
@@ -325,11 +325,6 @@ BUFFER_UOPS = {UOps.LOAD, UOps.STORE, UOps.VALID}
END_FOR_UOP = {UOps.IF:(UOps.STORE, UOps.ENDIF), UOps.RANGE:(UOps.ASSIGN, UOps.ENDRANGE)}
@functools.lru_cache(None)
def _min_bound(dtype:DType): return UOp.const(dtype.scalar(), dtypes.min(dtype))
@functools.lru_cache(None)
def _max_bound(dtype:DType): return UOp.const(dtype.scalar(), dtypes.max(dtype))
@dataclass(frozen=True, eq=False)
class UOp(MathTrait):
op: UOps
@@ -366,7 +361,6 @@ class UOp(MathTrait):
def bitcast(self, dtype=None): return type(self)(UOps.BITCAST, dtype, (self,))
def gep(self, i:int): return type(self)(UOps.GEP, self.dtype.scalar() if self.dtype is not None else None, (self,), i)
def const_like(self, b:ConstType|Variable): return type(self).const(self.dtype, b)
def sconst_like(self, b:ConstType|Variable): return type(self).const(self.dtype.scalar() if self.dtype is not None else None, b)
@classmethod
@functools.lru_cache(None)
def const(cls, dtype:Optional[DType], b:ConstType|Variable): return cls._const(dtype, b)
@@ -414,33 +408,34 @@ class UOp(MathTrait):
if (d1:=self.src[1].divides(v)) is not None: return self.src[0] * d1
return None # generic None if we aren't sure
@property
def vmin(self) -> UOp: return self._min_max[0]
def vmin(self) -> ConstType: return self._min_max[0]
@property
def vmax(self) -> UOp: return self._min_max[1]
def vmax(self) -> ConstType: return self._min_max[1]
@functools.cached_property
def _min_max(self) -> Tuple[UOp, UOp]:
def _min_max(self) -> Tuple[ConstType, ConstType]:
# NOTE: returned UOp is assumed to be CONST
if self.op is UOps.DEFINE_VAR and self.arg: return self.arg[1], self.arg[2] if isinstance(self.arg[2].arg, int) else _max_bound(self.dtype)
if self.op is UOps.DEFINE_VAR and self.arg:
return self.arg[1].arg, self.arg[2].arg if self.arg[2].op is UOps.CONST else dtypes.max(cast(DType, self.dtype))
if self.op is UOps.RANGE: return self.src[0].vmin, (self.src[1]-1).vmax
# TODO: UOps.SPECIAL is UOps.DEFINE_VAR
if self.op is UOps.SPECIAL: return self.const_like(0), self.const_like(self.arg[1]-1) if isinstance(self.arg[1], int) else _max_bound(self.dtype)
if self.op is UOps.CONST: return self, self
if self.op is UOps.SPECIAL: return 0, self.arg[1]-1 if isinstance(self.arg[1], int) else dtypes.max(cast(DType, self.dtype))
if self.op is UOps.CONST: return self.arg, self.arg
if self.op is UOps.ALU and cast(DType, self.dtype).count == 1:
s0,s1 = [cast(UOp, self.src[i] if i < len(self.src) else None) for i in range(2)]
if self.arg is BinaryOps.ADD: return self.sconst_like(s0.vmin.arg+s1.vmin.arg), self.sconst_like(s0.vmax.arg+s1.vmax.arg)
if self.arg is BinaryOps.MUL and (s0.vmin.arg >= 0 or s1.vmin.arg >= 0):
if self.arg is BinaryOps.ADD: return s0.vmin+s1.vmin, s0.vmax+s1.vmax
if self.arg is BinaryOps.MUL and (s0.vmin >= 0 or s1.vmin >= 0):
# handle at lease one is non-negative
Lmin, Lmax = (s0.vmin.arg, s0.vmax.arg) if s1.vmin.arg >= 0 else (s0.vmax.arg, s0.vmin.arg)
Rmin, Rmax = (s1.vmin.arg, s1.vmax.arg) if s0.vmin.arg >= 0 else (s1.vmax.arg, s1.vmin.arg)
Lmin, Lmax = (s0.vmin, s0.vmax) if s1.vmin >= 0 else (s0.vmax, s0.vmin)
Rmin, Rmax = (s1.vmin, s1.vmax) if s0.vmin >= 0 else (s1.vmax, s1.vmin)
assert math.isnan(Lmax*Rmax) or math.isnan(Lmin*Rmin) or Lmax*Rmax >= Lmin*Rmin, f"{Lmax=}, {Lmin=}, {Rmax=}, {Rmin=}"
return self.sconst_like(Lmin*Rmin), self.sconst_like(Lmax*Rmax)
if self.arg is BinaryOps.MOD and s1.vmin.arg > 0: return self.sconst_like(0), self.sconst_like(s1.vmax.arg-1)
return Lmin*Rmin, Lmax*Rmax
if self.arg is BinaryOps.MOD and s1.vmin > 0: return 0, s1.vmax-1
if self.arg is BinaryOps.IDIV and s1.op is UOps.CONST:
if s1.arg > 0: return self.sconst_like(s0.vmin.arg//s1.arg), self.sconst_like(s0.vmax.arg//s1.arg)
if s1.arg < 0: return self.sconst_like(-(s0.vmax.arg//-s1.arg)), self.sconst_like(-(s0.vmin.arg//-s1.arg))
if self.arg is BinaryOps.MAX: return self.sconst_like(max(s0.vmin.arg, s1.vmin.arg)), self.sconst_like(max(s0.vmax.arg, s1.vmax.arg))
if self.arg is BinaryOps.CMPLT: return (UOp.const(dtypes.bool, s0.vmax.arg<s1.vmin.arg), UOp.const(dtypes.bool, s0.vmin.arg<s1.vmax.arg))
return _min_bound(self.dtype), _max_bound(self.dtype)
if s1.arg > 0: return s0.vmin//s1.arg, s0.vmax//s1.arg
if s1.arg < 0: return -(s0.vmax//-s1.arg), -(s0.vmin//-s1.arg)
if self.arg is BinaryOps.MAX: return max(s0.vmin, s1.vmin), max(s0.vmax, s1.vmax)
if self.arg is BinaryOps.CMPLT: return (s0.vmax<s1.vmin, s0.vmin<s1.vmax)
return dtypes.min(cast(DType, self.dtype)), dtypes.max(cast(DType, self.dtype))
@dataclass(frozen=True)
class KernelInfo:
tinygrad/shape/shapetracker.py
+3 -3
View File
@@ -85,9 +85,9 @@ class ShapeTracker:
def real_size(self) -> int:
if 0 in self.shape: return 0
idx, valid = self.to_indexed_uops()
if not valid.vmax.arg: return 0
assert idx.vmax.arg < 1e12, f"real_size broken for {self}"
return idx.vmax.arg+1
if not valid.vmax: return 0
assert idx.vmax < 1e12, f"real_size broken for {self}"
return int(idx.vmax+1)
def vars(self) -> Set[Variable]: return set().union(*[v.vars() for v in self.views])