geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

method alu (#5711)

Browse Source 
Co-authored-by: chenyu <[email protected]>
This commit is contained in:
kormann
2024-07-25 13:25:38 -04:00
committed by GitHub
co-authored by chenyu
parent 76877df518
commit 92eefab4b0
3 changed files with 26 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/codegen/lowerer.py
+2 -2
View File
@@ -148,7 +148,7 @@ class IndependentLowerer:
has_valid = valid.op is not UOps.CONST or valid.arg is not True
if x.op is BufferOps.CONST:
dtype = x.arg.dtype.base if isinstance(x.arg.dtype, ImageDType) else x.arg.dtype
return UOp.alu(TernaryOps.WHERE, valid, UOp.const(dtype, x.arg.val), UOp.const(dtype, 0))
return valid.alu(TernaryOps.WHERE, UOp.const(dtype, x.arg.val), UOp.const(dtype, 0))
if x.arg.idx == -1:
buf = UOp(UOps.DEFINE_LOCAL, PtrDType(x.arg.dtype.base if isinstance(x.arg.dtype, ImageDType) else x.arg.dtype), (), ("temp", x.arg.st.size))
else:
@@ -179,6 +179,6 @@ class IndependentLowerer:
return UOp(UOps.EXPAND, dtype, tuple(UOp(UOps.GEP, dtype, (ret,), i) for i in range(wmma_sz[2])), arg=upcast_axis[2])
# NOTE: always using ridxs is fine here
return UOp(UOps.REDUCE, dtype, (in_uops[0],) + tuple(self.ridxs[i] for i in x.arg), x.op)
return UOp.alu(x.op, *in_uops)
return in_uops[0].alu(x.op, *in_uops[1:])
def lazyop_to_uop(ast:LazyOp, opts:Renderer) -> UOp: return IndependentLowerer().lower(ast, opts)
tinygrad/codegen/uopgraph.py
+4 -4
View File
@@ -19,7 +19,7 @@ def image_contract_load(buf, idx, idy, id4, ls_allow_any_len):
extra = (ls_allow_any_len.src[2], UOp(UOps.VECTORIZE, ls_allow_any_len.dtype.vec(4), (ls_allow_any_len.src[3],)*4))
else: extra = ls_allow_any_len.src[2:] # NOTE: image load shouldn't have barrier and this shouldn't matter
vec_load = UOp(UOps.LOAD, ls_allow_any_len.dtype.vec(4), (buf, UOp(UOps.VECTORIZE, dtypes.int.vec(2), (idx, idy))) + extra)
return functools.reduce(lambda ret, i: UOp.alu(TernaryOps.WHERE, id4.ne(i), ret, UOp(UOps.GEP, ls_allow_any_len.dtype, (vec_load,), i)), range(4),
return functools.reduce(lambda ret, i: id4.ne(i).alu(TernaryOps.WHERE, ret, UOp(UOps.GEP, ls_allow_any_len.dtype, (vec_load,), i)), range(4),
ls_allow_any_len.const(float('nan')))
def image_contract_store(buf, ex, idx, idy, ls_allow_any_len, var):
@@ -131,7 +131,7 @@ def loop_collapse(loop_start, loop_end, compval, idx, mval, multconst, rng, redu
return None
if idx2 is not None: idx = idx + idx2
if idx3 is not None: idx = idx + idx3
comprange = UOp.min(loop_end, UOp.max(UOp.alu(BinaryOps.IDIV, idx-compval-mval, mval) + (loop_end-loop_start), loop_start))
comprange = UOp.min(loop_end, UOp.max((idx-compval-mval).alu(BinaryOps.IDIV, mval) + (loop_end-loop_start), loop_start))
return UOp(UOps.REDUCE, reduce_allow_any_len.dtype, (comprange.cast(multconst.dtype) * multconst,) +
tuple(x for x in reduce_allow_any_len.src[1:] if x is not rng), reduce_allow_any_len.arg)
@@ -377,9 +377,9 @@ def do_reduce_with_expand(root):
assert root.src[0].op is UOps.EXPAND
expand_reduce_args = dedup(flatten([x.arg for x in expands_reduce]))
assert prod([y[1] for y in expand_reduce_args]) == len(root.src[0].src)
ret = functools.reduce(lambda x,y: UOp.alu(alu_op, x, y), (ret,)+root.src[0].src)
ret = functools.reduce(lambda x,y: x.alu(alu_op, y), (ret,)+root.src[0].src)
else:
ret = UOp.alu(alu_op, ret, root.src[0])
ret = ret.alu(alu_op, root.src[0])
ret = UOp(UOps.PHI, ret.dtype, (acc, ret))
if len(expands_non_reduce): ret = ret * prod([sz for _,sz in flatten([x.arg for x in expands_non_reduce])])
return ret
tinygrad/codegen/uops.py
+20 -20
View File
@@ -49,26 +49,26 @@ class UOp:
def cast(self, dtype=None): return UOp(UOps.CAST, dtype, (self,))
def bitcast(self, dtype=None): return UOp(UOps.BITCAST, dtype, (self,))
def name(self, name:Optional[str]): return UOp(UOps.VAR, src=(self,), arg=name)
def __neg__(self): return UOp.alu(UnaryOps.NEG, self)
def __add__(self, x): return UOp.alu(BinaryOps.ADD, self, self.ufix(x))
def __radd__(self, x): return UOp.alu(BinaryOps.ADD, self, self.ufix(x))
def __sub__(self, x): return UOp.alu(BinaryOps.ADD, self, self.ufix(-x))
def __mul__(self, x): return UOp.alu(BinaryOps.MUL, self, self.ufix(x))
def __rmul__(self, x): return UOp.alu(BinaryOps.MUL, self.ufix(x), self)
def __floordiv__(self, x): return UOp.alu(BinaryOps.IDIV, self, self.ufix(x))
def __truediv__(self, x): return UOp.alu(BinaryOps.MUL, self, UOp.alu(UnaryOps.RECIP, self.ufix(x)))
def __mod__(self, x): return UOp.alu(BinaryOps.MOD, self, self.ufix(x))
def __xor__(self, x): return UOp.alu(BinaryOps.XOR, self, self.ufix(x))
def __and__(self, x): return UOp.alu(BinaryOps.AND, self, self.ufix(x))
def __or__(self, x): return UOp.alu(BinaryOps.OR, self, self.ufix(x))
def ne(self, x): return UOp.alu(BinaryOps.CMPNE, self, self.ufix(x))
def __neg__(self): return self.alu(UnaryOps.NEG)
def __add__(self, x): return self.alu(BinaryOps.ADD, self.ufix(x))
def __radd__(self, x): return self.alu(BinaryOps.ADD, self.ufix(x))
def __sub__(self, x): return self.alu(BinaryOps.ADD, self.ufix(-x))
def __mul__(self, x): return self.alu(BinaryOps.MUL, self.ufix(x))
def __rmul__(self, x): return self.ufix(x).alu(BinaryOps.MUL, self)
def __floordiv__(self, x): return self.alu(BinaryOps.IDIV, self.ufix(x))
def __truediv__(self, x): return self.alu(BinaryOps.MUL, self.ufix(x).alu(UnaryOps.RECIP))
def __mod__(self, x): return self.alu(BinaryOps.MOD, self.ufix(x))
def __xor__(self, x): return self.alu(BinaryOps.XOR, self.ufix(x))
def __and__(self, x): return self.alu(BinaryOps.AND, self.ufix(x))
def __or__(self, x): return self.alu(BinaryOps.OR, self.ufix(x))
def ne(self, x): return self.alu(BinaryOps.CMPNE, self.ufix(x))
def eq(self, x): return -self.ne(x)
def lt(self, x): return UOp.alu(BinaryOps.CMPLT, self, self.ufix(x))
def lt(self, x): return self.alu(BinaryOps.CMPLT, self.ufix(x))
def ge(self, x): return -self.lt(x)
def max(self, x): return UOp.alu(BinaryOps.MAX, self, x)
def min(self, x): return -UOp.alu(BinaryOps.MAX, -self, -x)
def where(self, x, y): return UOp.alu(TernaryOps.WHERE, self, x, y)
def recip(self): return UOp.alu(UnaryOps.RECIP, self)
def max(self, x): return self.alu(BinaryOps.MAX, x)
def min(self, x): return -(-self).max(-x)
def where(self, x, y): return self.alu(TernaryOps.WHERE, x, y)
def recip(self): return self.alu(UnaryOps.RECIP)
def const(self:Union[UOp, DType, None], b:ConstType|Variable): return UOp._const(self.dtype if isinstance(self, UOp) else self, b)
@staticmethod
@functools.lru_cache(maxsize=None)
@@ -76,8 +76,8 @@ class UOp:
# TODO: fix dtype of b.max after Variable is just an UOp
if isinstance(b, Variable): return UOp(UOps.DEFINE_VAR, dtype, (UOp.const(dtypes.int, b.min), UOp.const(dtypes.int, cast(int,b.max))), b)
return UOp(UOps.CONST, dtype, arg=dtypes.as_const(b, dtype) if dtype is not None else b)
@staticmethod
def alu(arg, *src:UOp): return UOp(UOps.ALU, dtypes.bool if arg in {BinaryOps.CMPLT, BinaryOps.CMPNE} else src[-1].dtype, src, arg)
def alu(self, arg, *src:UOp):
return UOp(UOps.ALU, dtypes.bool if arg in {BinaryOps.CMPLT, BinaryOps.CMPNE} else (self, *src)[-1].dtype, (self,)+src, arg)
@staticmethod
def load(*src:UOp, dtype:Optional[DType]=None, **kwargs): return UOp(UOps.LOAD, dtype, tuple(src)+tuple(kwargs.values()))
@staticmethod