diff --git a/tinygrad/codegen/lowerer.py b/tinygrad/codegen/lowerer.py index d6ebb47f9d..bf00a30a4d 100644 --- a/tinygrad/codegen/lowerer.py +++ b/tinygrad/codegen/lowerer.py @@ -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) diff --git a/tinygrad/codegen/uopgraph.py b/tinygrad/codegen/uopgraph.py index 124c17aa89..8b95c29f01 100644 --- a/tinygrad/codegen/uopgraph.py +++ b/tinygrad/codegen/uopgraph.py @@ -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 diff --git a/tinygrad/codegen/uops.py b/tinygrad/codegen/uops.py index 32e9d83efd..74989703f8 100644 --- a/tinygrad/codegen/uops.py +++ b/tinygrad/codegen/uops.py @@ -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