diff --git a/tinygrad/ops.py b/tinygrad/ops.py index d03ed818a2..c6eb1489c1 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -325,7 +325,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): return self.arg[0] def bind(self, val:int): assert self.op is UOps.DEFINE_VAR, f"op is {self.op}, need DEFINE_VAR" - assert self.arg[1] <= val and val <= self.arg[2], f"bind {val} not in range {self.arg[1]}-{self.arg[2]}" + assert self.arg[1] <= val and val <= self.arg[2], f"bind {val} not in range [{self.arg[1]}, {self.arg[2]}]" return UOp(UOps.BIND, self.dtype, (self, self.const_like(val))) def unbind(self) -> Tuple[Variable, int]: assert self.op is UOps.BIND and self.src[0].op is UOps.DEFINE_VAR and self.src[1].op is UOps.CONST, f"can't unbind {self}" @@ -392,7 +392,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): return (always_ne, sometimes_ne) # float has NAN issue and we use explicit NAN in transcendental if self.arg is TernaryOps.WHERE and dtypes.is_int(s1.dtype): return min(s1.vmin, s2.vmin), max(s1.vmax, s2.vmax) - if self.dtype is dtypes.bool: + if self.dtype == dtypes.bool: if self.arg is BinaryOps.OR: return s0.vmin or s1.vmin, s0.vmax or s1.vmax if self.arg is BinaryOps.AND: return s0.vmin and s1.vmin, s0.vmax and s1.vmax return dtypes.min(self.dtype), dtypes.max(self.dtype) @@ -430,11 +430,10 @@ python_alu: Dict[Op, Callable] = { UnaryOps.LOG2: lambda x: math.log2(x) if x > 0 else -math.inf if x == 0 else math.nan, UnaryOps.EXP2: hook_overflow(math.inf, lambda x: 2**x), UnaryOps.SQRT: lambda x: math.sqrt(x) if x >= 0 else math.nan, UnaryOps.RECIP: lambda x: 1/x if x != 0 else math.copysign(math.inf, x), UnaryOps.SIN: lambda x: math.sin(x) if not math.isinf(x) else math.nan, - UnaryOps.NEG: operator.neg, BinaryOps.ADD: operator.add, BinaryOps.SUB: operator.sub, - BinaryOps.SHR: operator.rshift, BinaryOps.SHL: operator.lshift, BinaryOps.MUL: operator.mul, - BinaryOps.XOR: operator.xor, BinaryOps.MAX: max, BinaryOps.CMPNE: operator.ne, BinaryOps.CMPLT: operator.lt, - BinaryOps.OR: operator.or_, BinaryOps.AND: operator.and_, + UnaryOps.NEG: operator.neg, BinaryOps.ADD: operator.add, BinaryOps.SUB: operator.sub, BinaryOps.MUL: operator.mul, BinaryOps.MOD: lambda x,y: abs(int(x))%abs(int(y))*(1,-1)[x<0], BinaryOps.IDIV: lambda x,y: abs(x)//abs(y)*(1,-1)[x*y<0] if y != 0 else x*math.inf, + BinaryOps.MAX: max, BinaryOps.CMPNE: operator.ne, BinaryOps.CMPLT: operator.lt, BinaryOps.XOR: operator.xor, + BinaryOps.OR: operator.or_, BinaryOps.AND: operator.and_, BinaryOps.SHR: operator.rshift, BinaryOps.SHL: operator.lshift, TernaryOps.MULACC: lambda x,y,z: (x*y)+z, TernaryOps.WHERE: lambda x,y,z: y if x else z} def exec_alu(op:Op, dtype:DType, operands, truncate_output=True): @@ -550,10 +549,9 @@ class UPat(MathTrait): return UPat(UOps.ALU, None if arg in {BinaryOps.CMPLT, BinaryOps.CMPNE} else asrc[-1].dtype, list(asrc) if arg in COMMUTATIVE else asrc, arg) def printable(self:UPat) -> str: - try: - return lines(self.location[0])[self.location[1]-1].strip() - except FileNotFoundError: - return "" + try: return lines(self.location[0])[self.location[1]-1].strip() + except FileNotFoundError: return "" + def __repr__(self): def rep(x): form = "UPat(%s, %s, name=%s, dtype=%s, allow_any_len=%s, src=%s)" @@ -788,7 +786,7 @@ spec = PatternMatcher([ def type_verify(uops:List[UOp]): for i,u in enumerate(uops): - if cast(bool, spec.rewrite(u)) is not True: + if not spec.rewrite(u): print_uops(uops) raise RuntimeError(f"UOp verification failed at {i} on {u.op} {u.dtype} {len(u.src)} {[x.op for x in u.src]} {u.arg}")