diff --git a/extra/datasets/sops.gz b/extra/datasets/sops.gz index 321f83c413..8411b0af42 100644 Binary files a/extra/datasets/sops.gz and b/extra/datasets/sops.gz differ diff --git a/test/test_linearizer.py b/test/test_linearizer.py index 4cad943bc2..f5ec562258 100644 --- a/test/test_linearizer.py +++ b/test/test_linearizer.py @@ -1234,7 +1234,7 @@ class TestKernelOpts(unittest.TestCase): @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores") def test_buf_index_not_found_tensor_core(self): - ast = LazyOp(op=BufferOps.STORE, src=(LazyOp(op=ReduceOps.SUM, src=(LazyOp(op=BinaryOps.MUL, src=(LazyOp(op=UnaryOps.CAST, src=(LazyOp(op=BinaryOps.CMPEQ, src=(LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=1, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(1243, 256), strides=(0, 1), offset=0, mask=None, contiguous=False),)))), LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=2, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(1243, 256), strides=(1, 0), offset=0, mask=None, contiguous=False),))))), arg=None),), arg=dtypes.float), LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=3, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(1243, 256), strides=(1, 0), offset=0, mask=None, contiguous=False),))))), arg=None),), arg=(0,)),), arg=MemBuffer(idx=0, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(1, 256), strides=(0, 1), offset=0, mask=None, contiguous=True),)))) # noqa: E501 + ast = LazyOp(op=BufferOps.STORE, src=(LazyOp(op=ReduceOps.SUM, src=(LazyOp(op=BinaryOps.MUL, src=(LazyOp(op=UnaryOps.CAST, src=(LazyOp(op=BinaryOps.CMPNE, src=(LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=1, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(1243, 256), strides=(0, 1), offset=0, mask=None, contiguous=False),)))), LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=2, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(1243, 256), strides=(1, 0), offset=0, mask=None, contiguous=False),))))), arg=None),), arg=dtypes.float), LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=3, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(1243, 256), strides=(1, 0), offset=0, mask=None, contiguous=False),))))), arg=None),), arg=(0,)),), arg=MemBuffer(idx=0, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(1, 256), strides=(0, 1), offset=0, mask=None, contiguous=True),)))) # noqa: E501 k = Linearizer(ast, opts=Device[Device.DEFAULT].renderer) with self.assertRaises(KernelOptError): k.apply_opt(Opt(OptOps.TC, 0, 1)) diff --git a/test/test_ops.py b/test/test_ops.py index 9631c3ef12..a4cdee7e93 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -243,17 +243,17 @@ class TestOps(unittest.TestCase): def test_cmp_lt(self): self._test_cmp(lambda x,y: x Set[UOp]: return set.union(set(self.vin), *[x.parents for x in self.vin]) @@ -424,7 +424,7 @@ class UOpGraph: if uop is UOps.ALU: if arg in UnaryOps: assert dtype == vin[0].dtype, f"{arg} dtype mismatch {dtype=} != {vin[0].dtype=}" - elif arg in (BinaryOps.CMPLT, BinaryOps.CMPEQ): + elif arg in (BinaryOps.CMPLT, BinaryOps.CMPNE): assert dtype == dtypes.bool, f"{arg} output dtype mismatch {dtype=} != {dtypes.bool}" assert vin[0].dtype == vin[1].dtype, f"{arg} dtype mismatch {dtype=} != {vin[0].dtype=} != {vin[1].dtype=}" elif arg in BinaryOps: diff --git a/tinygrad/function.py b/tinygrad/function.py index 189f82bc26..8390dd52a9 100644 --- a/tinygrad/function.py +++ b/tinygrad/function.py @@ -88,8 +88,8 @@ class Sigmoid(Function): class Sign(Function): def forward(self, x:LazyBuffer) -> LazyBuffer: - return x.e(BinaryOps.CMPEQ, x.const(0)).e(TernaryOps.WHERE, x.const(0), - x.e(BinaryOps.CMPLT, x.const(0)).e(TernaryOps.WHERE, x.const(-1), x.const(1))) + return x.e(BinaryOps.CMPNE, x.const(0)).e( + TernaryOps.WHERE, x.e(BinaryOps.CMPLT, x.const(0)).e(TernaryOps.WHERE, x.const(-1), x.const(1)), x.const(0)) # backward always return 0 to match torch def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output.const(0) @@ -99,8 +99,8 @@ class Less(Function): def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer: return x.e(BinaryOps.CMPLT, y) def backward(self, grad_output:LazyBuffer) -> Tuple[Optional[LazyBuffer], Optional[LazyBuffer]]: return None, None -class Eq(Function): - def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer: return x.e(BinaryOps.CMPEQ, y) +class Neq(Function): + def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer: return x.e(BinaryOps.CMPNE, y) def backward(self, grad_output:LazyBuffer) -> Tuple[Optional[LazyBuffer], Optional[LazyBuffer]]: return None, None class Xor(Function): @@ -166,7 +166,7 @@ class Max(Function): def backward(self, grad_output:LazyBuffer) -> LazyBuffer: # 1s in locations where the max was chosen (can be two locations) - max_is_1s = self.x.e(BinaryOps.CMPEQ, self.ret.expand(self.x.shape)).cast(dtypes.float) + max_is_1s = self.x.const(1.0).cast(dtypes.float).e(BinaryOps.SUB, self.x.e(BinaryOps.CMPNE, self.ret.expand(self.x.shape)).cast(dtypes.float)) div = max_is_1s.r(ReduceOps.SUM, self.axis).expand(self.x.shape) return max_is_1s.e(BinaryOps.DIV, div).cast(grad_output.dtype).e(BinaryOps.MUL, grad_output.expand(self.x.shape)) diff --git a/tinygrad/lazy.py b/tinygrad/lazy.py index 889f469b24..020e6817b7 100644 --- a/tinygrad/lazy.py +++ b/tinygrad/lazy.py @@ -146,7 +146,7 @@ class LazyBuffer: if op is TernaryOps.WHERE: assert srcs[0].dtype == dtypes.bool, "TernaryOps.WHERE must have the first arg be bool" if op is UnaryOps.NEG: assert srcs[0].dtype != dtypes.bool, "UnaryOps.NEG does not accept dtype bool" - out_dtype = dtypes.bool if op in (BinaryOps.CMPLT, BinaryOps.CMPEQ) else srcs[-1].dtype + out_dtype = dtypes.bool if op in (BinaryOps.CMPLT, BinaryOps.CMPNE) else srcs[-1].dtype # const folding if op in python_alu and all(s.is_unrealized_unmasked_const() for s in srcs): diff --git a/tinygrad/ops.py b/tinygrad/ops.py index bdfa8d8d81..fc818cc939 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -17,7 +17,7 @@ class UnaryOps(Enum): EXP2 = auto(); LOG2 = auto(); CAST = auto(); BITCAST = auto(); SIN = auto(); SQRT = auto(); NEG = auto() # noqa: E702 class BinaryOps(Enum): """A + A -> A (elementwise)""" - ADD = auto(); SUB = auto(); MUL = auto(); DIV = auto(); MAX = auto(); MOD = auto(); CMPLT = auto(); CMPEQ = auto(); XOR = auto() # noqa: E702 + ADD = auto(); SUB = auto(); MUL = auto(); DIV = auto(); MAX = auto(); MOD = auto(); CMPLT = auto(); CMPNE = auto(); XOR = auto() # noqa: E702 SHR = auto(); SHL = auto() # noqa: E702 class TernaryOps(Enum): """A + A + A -> A (elementwise)""" @@ -31,7 +31,7 @@ class LoadOps(Enum): EMPTY = auto(); CONST = auto(); COPY = auto(); CONTIGUOUS = Op = Union[UnaryOps, BinaryOps, ReduceOps, LoadOps, TernaryOps, BufferOps] # do not preserve f(0) = 0 -UNSAFE_PAD_OPS = {BinaryOps.DIV, BinaryOps.CMPEQ, UnaryOps.LOG2, UnaryOps.EXP2} +UNSAFE_PAD_OPS = {BinaryOps.DIV, UnaryOps.LOG2, UnaryOps.EXP2} @dataclass(frozen=True) class MemBuffer: @@ -62,7 +62,7 @@ class LazyOp: def dtype(self) -> DType: if self.op in BufferOps: return self.arg.dtype if self.op in [UnaryOps.CAST, UnaryOps.BITCAST]: return self.arg - return dtypes.bool if self.op in {BinaryOps.CMPLT, BinaryOps.CMPEQ} else self.src[-1].dtype + return dtypes.bool if self.op in {BinaryOps.CMPLT, BinaryOps.CMPNE} else self.src[-1].dtype @functools.cached_property def key(self) -> bytes: @@ -122,7 +122,7 @@ python_alu = { UnaryOps.NEG: lambda x: (not x) if isinstance(x, bool) else -x, BinaryOps.SHR: operator.rshift, BinaryOps.SHL: operator.lshift, BinaryOps.MUL: operator.mul, BinaryOps.ADD: operator.add, BinaryOps.SUB: operator.sub, BinaryOps.XOR: operator.xor, - BinaryOps.MAX: max, BinaryOps.CMPEQ: operator.eq, BinaryOps.CMPLT: operator.lt, + BinaryOps.MAX: max, BinaryOps.CMPNE: operator.ne, BinaryOps.CMPLT: operator.lt, BinaryOps.MOD: lambda x,y: abs(int(x))%abs(int(y))*(1,-1)[x<0], BinaryOps.DIV: lambda x,y: int(x/y) if isinstance(x, int) else (x/y if y != 0 else x*math.inf), TernaryOps.WHERE: lambda x,y,z: y if x else z} diff --git a/tinygrad/renderer/assembly.py b/tinygrad/renderer/assembly.py index 5759668cbd..93724597c2 100644 --- a/tinygrad/renderer/assembly.py +++ b/tinygrad/renderer/assembly.py @@ -48,7 +48,7 @@ class PTXRenderer(Renderer): BinaryOps.DIV: lambda d,a,b,dt,name: f"div{'.approx' if dtypes.is_float(dt) else ''}.{name} {d}, {a}, {b};", BinaryOps.MAX: lambda d,a,b,dt,name: f"max.{name} {d}, {a}, {b};", BinaryOps.MOD: lambda d,a,b,dt,name: f"rem.{name} {d}, {a}, {b};", BinaryOps.CMPLT: lambda d,a,b,dt,name: f"setp.lt.{name} {d}, {a}, {b};", - BinaryOps.CMPEQ: lambda d,a,b,dt,name: f"setp.eq.{name} {d}, {a}, {b};", + BinaryOps.CMPNE: lambda d,a,b,dt,name: f"setp.ne.{name} {d}, {a}, {b};", TernaryOps.MULACC: lambda d,a,b,c,dt,name: f"{'fma.rn' if dtypes.is_float(dt) else 'mad.lo'}.{name} {d}, {a}, {b}, {c};", TernaryOps.WHERE: lambda d,a,b,c,dt,name: f"@{a} mov.{name} {d}, {b};\n@!{a} mov.{name} {d}, {c};" if name == "pred" else f"selp.{'b16' if name == 'f16' else name} {d}, {b}, {c}, {a};" @@ -167,7 +167,7 @@ class PTXRenderer(Renderer): if uop is UOps.RANGE: kk(*self.render_loop(ssa('ridx', u), r[vin[0]], ssa_label('loop', u))) elif uop is UOps.ALU: assert vin[0].dtype is not None - if args is BinaryOps.CMPLT or args is BinaryOps.CMPEQ: + if args is BinaryOps.CMPLT or args is BinaryOps.CMPNE: # pass in the other dtype here kk(self.asm_for_op[args](ssa("alu", u), *[r[x] for x in vin], vin[0].dtype, self.types[vin[0].dtype])) else: @@ -243,8 +243,8 @@ ptx_matcher = PatternMatcher([ ({"__name__": "root", "uop": UOps.ALU, "arg": BinaryOps.DIV, "dtype": set([dt for dt in dtypes.fields().values() if dtypes.is_int(dt)]), "vin": [{"__name__": "const", "uop": UOps.CONST, "arg": set([2**i for i in range(64)])}, {"__name__": "div"}]}, lambda root, div, const: UOp(UOps.ALU, root.dtype, (div, UOp.const(root.dtype, int(math.log2(const.arg)))), BinaryOps.SHR)), - ({"__name__": "root", "uop": UOps.ALU, "arg": BinaryOps.CMPEQ, "vin": ({"dtype": dtypes.bool},{})}, - lambda root: UOp(UOps.ALU, dtypes.bool, (UOp(root.uop, root.dtype, root.vin, BinaryOps.XOR),), UnaryOps.NEG)), + ({"__name__": "root", "uop": UOps.ALU, "arg": BinaryOps.CMPNE, "vin": ({"dtype": dtypes.bool},{})}, + lambda root: UOp(root.uop, root.dtype, root.vin, BinaryOps.XOR)), ({"__name__": "root", "uop": UOps.ALU, "arg": BinaryOps.CMPLT, "vin": ({"__name__": "x", "dtype": dtypes.bool},{"__name__": "y"})}, lambda root,x,y: UOp(root.uop, root.dtype, (UOp(UOps.ALU, dtypes.bool, (x,), UnaryOps.NEG), y), BinaryOps.MUL)), ({"__name__": "root", "uop": UOps.ALU, "arg": BinaryOps.ADD, diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 7df8c01e8c..bc534d625a 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -28,7 +28,7 @@ class CStyleLanguage(Renderer): UnaryOps.EXP2: lambda x,dtype: f"exp2({x})", UnaryOps.LOG2: lambda x,dtype: f"log2({x})", UnaryOps.SIN: lambda x,dtype: f"sin({x})", BinaryOps.ADD: lambda a,b,dtype: f"({a}+{b})", BinaryOps.SUB: lambda a,b,dtype: f"({a}-{b})", BinaryOps.MUL: lambda a,b,dtype: f"({a}*{b})", BinaryOps.DIV: lambda a,b,dtype: f"({a}/{b})", BinaryOps.MAX: lambda a,b,dtype: f"max({a},{b})", BinaryOps.MOD: lambda a,b,dtype: f"({a}%{b})", - BinaryOps.CMPLT: lambda a,b,dtype: f"({a}<{b})", BinaryOps.CMPEQ: lambda a,b,dtype: f"({a}=={b})", BinaryOps.XOR: lambda a,b,dtype: f"({a}^{b})", + BinaryOps.CMPLT: lambda a,b,dtype: f"({a}<{b})", BinaryOps.CMPNE: lambda a,b,dtype: f"({a}!={b})", BinaryOps.XOR: lambda a,b,dtype: f"({a}^{b})", TernaryOps.WHERE: lambda a,b,c,dtype: f"({a}?{b}:{c})"} # returns a str expression of the casted xs with the given type diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index 4fc6f868e8..2a07abe1e0 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -22,7 +22,7 @@ code_for_op: Final[Dict[Op, Callable]] = { BinaryOps.MUL: lambda builder, x, y, dtype: builder.mul(x, y) if is_bool_or_unsigned(dtype) or dtypes.is_int(dtype) else builder.fmul(x, y, flags=MFLAGS), # noqa: E501 BinaryOps.DIV: lambda builder, x, y, dtype: builder.udiv(x, y) if is_bool_or_unsigned(dtype) else builder.sdiv(x, y) if dtypes.is_int(dtype) else builder.fdiv(x, y, flags=MFLAGS), # noqa: E501 BinaryOps.CMPLT: lambda builder, x, y, dtype: builder.icmp_unsigned("<", x, y) if is_bool_or_unsigned(dtype) else builder.icmp_signed("<", x, y) if dtypes.is_int(dtype) else builder.fcmp_unordered("<", x, y, flags=MFLAGS), # noqa: E501 - BinaryOps.CMPEQ: lambda builder, x, y, dtype: builder.icmp_unsigned("==", x, y) if is_bool_or_unsigned(dtype) else builder.icmp_signed("==", x, y) if dtypes.is_int(dtype) else builder.fcmp_unordered("==", x, y, flags=MFLAGS), # noqa: E501 + BinaryOps.CMPNE: lambda builder, x, y, dtype: builder.icmp_unsigned("!=", x, y) if is_bool_or_unsigned(dtype) else builder.icmp_signed("!=", x, y) if dtypes.is_int(dtype) else builder.fcmp_unordered("!=", x, y, flags=MFLAGS), # noqa: E501 BinaryOps.MAX: lambda builder, x, y, dtype: builder.select(builder.icmp_unsigned(">", x, y) if is_bool_or_unsigned(dtype) else builder.icmp_signed(">", x, y) if dtypes.is_int(dtype) else builder.fcmp_unordered(">", x, y, flags=MFLAGS), x, y), # noqa: E501 BinaryOps.MOD: lambda builder, x, y, dtype: builder.urem(x, y) if is_bool_or_unsigned(dtype) else builder.srem(x, y) if dtypes.is_int(dtype) else builder.frem(x, y), # noqa: E501 BinaryOps.XOR: lambda builder, x, y, dtype: builder.xor(x, y), @@ -149,7 +149,7 @@ class LLVMRenderer(Renderer): while backward.uop is UOps.PHI: backward = backward.vin[0] lvars[backward] = lvars[u] elif uop is UOps.ALU: - lvars[u] = code_for_op[args](bb[-1], *[lvars[x] for x in vin], dtype if args not in (BinaryOps.CMPLT, BinaryOps.CMPEQ) else vin[0].dtype) + lvars[u] = code_for_op[args](bb[-1], *[lvars[x] for x in vin], dtype if args not in (BinaryOps.CMPLT, BinaryOps.CMPNE) else vin[0].dtype) elif uop in {UOps.CAST, UOps.BITCAST}: lvars[u] = cast(bb, lvars[vin[0]], vin[0].dtype, dtype, bitcast=uop is UOps.BITCAST) elif uop in {UOps.DEFINE_GLOBAL, UOps.DEFINE_VAR}: lvars[u] = func.args[buf_index[args]] elif uop is UOps.SPECIAL: lvars[u] = lvars[args.expr] diff --git a/tinygrad/runtime/ops_python.py b/tinygrad/runtime/ops_python.py index 30ef5e7cef..3d42a13e91 100644 --- a/tinygrad/runtime/ops_python.py +++ b/tinygrad/runtime/ops_python.py @@ -174,7 +174,7 @@ class PythonProgram: else: raise NotImplementedError(f"unimplemented tensor core {arg}") elif uop is UOps.ALU: assert all_same([len(x) for x in inp]), f"{[len(x) for x in inp]} doesn't match on {arg}" - assert all_same([dtype] + dtp) or arg in {BinaryOps.CMPEQ, BinaryOps.CMPLT, TernaryOps.WHERE}, f"dtype mismatch on {arg}" + assert all_same([dtype] + dtp) or arg in {BinaryOps.CMPNE, BinaryOps.CMPLT, TernaryOps.WHERE}, f"dtype mismatch on {arg}" ul[i] = [exec_alu(arg, dtype, p) for p in zip(*inp)] assert i in ul, (uop, dtype, idp, arg) i += 1 diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 96cea73c02..32486368c4 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -1793,7 +1793,7 @@ class Tensor: print(Tensor([False, True]).logical_not().numpy()) ``` """ - return F.Eq.apply(*self._broadcasted(False)) + return F.Neq.apply(*self.cast(dtypes.bool)._broadcasted(True)) def neg(self): """ Negates the tensor element-wise. @@ -2542,8 +2542,8 @@ class Tensor: def __gt__(self, x) -> Tensor: return F.Less.apply(*self._broadcasted(x, True)) def __ge__(self, x) -> Tensor: return (self Tensor: return (self>x).logical_not() - def __eq__(self, x) -> Tensor: return F.Eq.apply(*self._broadcasted(x, True)) # type: ignore[override] - def __ne__(self, x) -> Tensor: return (self==x).logical_not() # type: ignore[override] + def __ne__(self, x) -> Tensor: return F.Neq.apply(*self._broadcasted(x)) # type: ignore[override] + def __eq__(self, x) -> Tensor: return (self!=x).logical_not() # type: ignore[override] # ***** functional nn ops *****