Repository: tinygrad/tinygrad
Pull request: threefry_2x32
Closes: #2418


Problem statement:
Add fast PRNG support (remove LoadOps.RAND)
JAX uses an encryption function to encrypt arange, threefry2x32. I think this is the right idea

- [ ] Support doing math on integers
- [ ] Add BinaryOps.XOR and `bitwise_xor` (note, rotate_left is a mul+div+add)
- [ ] Implement threefry2x32 and test it
- [ ] Make arange fast? Ensure this is all fused.
- [ ] Make symbolics work as normal inputs (aka jitted symbolic arange)
- [ ] Add "seed" to the JIT (TODO: how does JAX do this?)

Hints from issue discussion:

Commit message: feat: numpy xor

Patch to review:
diff --git a/tinygrad/mlops.py b/tinygrad/mlops.py
index ee0766fbe964e..2309d6affb3fe 100644
--- a/tinygrad/mlops.py
+++ b/tinygrad/mlops.py
@@ -90,6 +90,10 @@ class Less(Function):
   def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer:
     return x.e(BinaryOps.CMPLT, y)
 
+class Xor(Function):
+  def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer:
+    return x.e(BinaryOps.XOR, y)
+
 class Add(Function):
   def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer:
     return x.e(BinaryOps.ADD, y)
diff --git a/tinygrad/ops.py b/tinygrad/ops.py
index ccfd454e53fbd..ec0e3775bd938 100644
--- a/tinygrad/ops.py
+++ b/tinygrad/ops.py
@@ -11,7 +11,7 @@
 # NOTE: MOD, CMPLT don't have to be implemented on vectors, just scalars
 # NOTE: rdna3 only has RECIP and not DIV. DIV and POW are on the chopping block
 class UnaryOps(Enum): EXP2 = auto(); LOG2 = auto(); CAST = auto(); SIN = auto(); SQRT = auto(); RECIP = auto(); NEG = auto() # noqa: E702
-class BinaryOps(Enum): ADD = auto(); SUB = auto(); MUL = auto(); DIV = auto(); MAX = auto(); MOD = auto(); CMPLT = auto() # noqa: E702
+class BinaryOps(Enum): ADD = auto(); SUB = auto(); MUL = auto(); DIV = auto(); MAX = auto(); MOD = auto(); CMPLT = auto(); XOR = auto() # noqa: E702
 class TernaryOps(Enum): MULACC = auto(); WHERE = auto() # noqa: E702
 class ReduceOps(Enum): SUM = auto(); MAX = auto() # noqa: E702
 class BufferOps(Enum): LOAD = auto(); CONST = auto(); STORE = auto() # noqa: E702
diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py
index c2d4043c6ffa2..bffb8a94b4631 100644
--- a/tinygrad/renderer/cstyle.py
+++ b/tinygrad/renderer/cstyle.py
@@ -35,8 +35,8 @@ class CStyleLanguage(NamedTuple):
     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})", TernaryOps.MULACC: lambda a,b,c,dtype: f"(({a}*{b})+{c})",
-    TernaryOps.WHERE: lambda a,b,c,dtype: f"({a}!=0?{b}:{c})"
+    BinaryOps.CMPLT: lambda a,b,dtype: f"({a}<{b})", BinaryOps.XOR: lambda a,b,dtype: f"({a}^{b})",
+    TernaryOps.MULACC: lambda a,b,c,dtype: f"(({a}*{b})+{c})", TernaryOps.WHERE: lambda a,b,c,dtype: f"({a}!=0?{b}:{c})"
   }
 
   # returns a str expression of the casted xs with the given type
@@ -154,7 +154,7 @@ def ssa(u, prefix="t"):
     elif uop == UOps.ALU:
       assert dtype is not None
       # remove parens if ALU types are the same. TODO: can do more here
-      if vin[0].uop == UOps.ALU and vin[0].arg == args and args in {BinaryOps.ADD, BinaryOps.SUB, BinaryOps.MUL}:
+      if vin[0].uop == UOps.ALU and vin[0].arg == args and args in {BinaryOps.ADD, BinaryOps.SUB, BinaryOps.MUL, BinaryOps.XOR}:
         val = lang.code_for_op[args](strip_parens(r[vin[0]]), *[r[x] for x in vin[1:]], dtype)
       else:
         val = lang.code_for_op[args](*[r[x] for x in vin] + [dtype])
diff --git a/tinygrad/runtime/ops_cpu.py b/tinygrad/runtime/ops_cpu.py
index e4c0c56c0d5d6..63d1d7ba2c226 100644
--- a/tinygrad/runtime/ops_cpu.py
+++ b/tinygrad/runtime/ops_cpu.py
@@ -30,7 +30,7 @@ def mulacc(a, b, new_shape):
   UnaryOps.CAST: lambda x,y: x.view(y[0].np) if y[1] else x.astype(y[0].np, copy=False), UnaryOps.NEG: lambda x: np.logical_not(x) if x.dtype == np.bool_ else np.negative(x),
   BinaryOps.MAX: np.maximum, BinaryOps.CMPLT: lambda x,y: (x<y).astype(output_type(x,y)), BinaryOps.ADD: lambda x, y: np.add(*match_types(x, y)),
   BinaryOps.SUB: lambda x, y: np.subtract(*match_types(x, y)), BinaryOps.MUL: lambda x, y: np.multiply(*match_types(x, y)),
-  BinaryOps.DIV: lambda x, y: np.divide(*match_types(x, y)).astype(output_type(x, y), copy=False), UnaryOps.SQRT: np.sqrt,
+  BinaryOps.DIV: lambda x, y: np.divide(*match_types(x, y)).astype(output_type(x, y), copy=False), BinaryOps.XOR: lambda x, y: np.bitwise_xor(*match_types(x, y)), UnaryOps.SQRT: np.sqrt,
   ReduceOps.SUM: lambda x, new_shape: x.sum(shape_to_axis(x.shape, new_shape), dtype=x.dtype, keepdims=True) if x.shape != new_shape else x,
   ReduceOps.MAX: lambda x, new_shape: x.max(shape_to_axis(x.shape, new_shape), keepdims=True) if x.shape != new_shape else x,
   MovementOps.AS_STRIDED: lambda x, arg: np.ndarray(arg[0], buffer=np.require(x, requirements='C'), dtype=x.dtype, offset=arg[2]*x.dtype.itemsize, strides=tuple(y*x.dtype.itemsize for y in arg[1])),
diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
index 835fa17254b76..dabdb29da5f76 100644
--- a/tinygrad/tensor.py
+++ b/tinygrad/tensor.py
@@ -10,7 +10,7 @@
 from tinygrad.helpers import ImageDType, argfix, make_pair, getenv, IMAGE, DEBUG, flatten, DType, dtypes, prod, all_int, round_up
 from tinygrad.lazy import LazyBuffer
 from tinygrad.ops import LoadOps
-from tinygrad.device import Device, Buffer
+from tinygrad.device import Device
 from tinygrad.shape.symbolic import sint
 from tinygrad.realize import run_schedule
 
@@ -156,9 +156,27 @@ def empty(*shape, **kwargs):
   @staticmethod
   def manual_seed(seed=0): Tensor._seed = seed
 
+  _rng_counter: int = 0
   @staticmethod
   def rand(*shape, **kwargs):
-    return Tensor._loadop(LoadOps.CUSTOM, prod((shape:=argfix(*shape))), arg=custom_random, **kwargs).reshape(shape)
+    num = prod((shape:=argfix(*shape)))
+    if (odd_counts := num % 2): num += 1
+    counts = Tensor.arange(num, dtype=dtypes.uint32) + Tensor._rng_counter
+    Tensor._rng_counter += num
+
+    rotations = [[13, 15, 26, 6], [17, 29, 16, 24]]
+    ks = [0x0, Tensor._seed ^ 0x0 ^ 0x1BD11BDA, Tensor._seed]
+
+    x = counts.chunk(2)
+    for i in range(5):
+      for r in rotations[0]:
+        x[0] = x[0] + x[1]
+        x[1] = x[0] ^ ((x[1] * (2 ** r)) + (x[1] / (2 ** (32 - r))))
+      x = [(x[0] + ks[0]).realize(), (x[1] + ks[1] + i + 1).realize()]
+      rotations = rotations[1:] + rotations[:1]
+      ks = ks[1:] + ks[:1]
+    out = x[0].cat(x[1])
+    return out[odd_counts:].reshape(shape).cast(Tensor.default_type if kwargs.get("dtype") is None else kwargs.get("dtype")) / (2**32 - 1)
 
   # ***** creation helper functions *****
 
@@ -719,6 +737,9 @@ def pow(self, x:Union[Tensor, float], reverse=False) -> Tensor:
     inject_nan = ((((-to_nan) * 2) + 1)).log().add(1) if isinstance(to_nan, Tensor) else 1 if not to_nan else float("nan")
     return ar.mul(sign * base_sign + (1 - base_sign)).mul(inject_nan)
   def matmul(self, x:Tensor, reverse=False) -> Tensor: return x.dot(self) if reverse else self.dot(x)
+  def xor(self, x:Union[Tensor, float], reverse=False) -> Tensor:
+    x = self._to_float(x)
+    return mlops.Xor.apply(*self._broadcasted(x, reverse)) if x.__class__ is Tensor or x else self
 
   def maximum(self, x:Union[Tensor, float]) -> Tensor: return (self<x).detach().where(x, (self>x).detach().where(self, (self+x)/2))
   def minimum(self, x:Union[Tensor, float]) -> Tensor: return -((-self).maximum(-x))
@@ -738,6 +759,7 @@ def __mul__(self, x) -> Tensor: return self.mul(x)
   def __pow__(self, x) -> Tensor: return self.pow(x)
   def __truediv__(self, x) -> Tensor: return self.div(x)
   def __matmul__(self, x) -> Tensor: return self.matmul(x)
+  def __xor__(self, x) -> Tensor: return self.xor(x)
 
   def __radd__(self, x) -> Tensor: return self.add(x, True)
   def __rsub__(self, x) -> Tensor: return self.sub(x, True)
@@ -745,6 +767,7 @@ def __rmul__(self, x) -> Tensor: return self.mul(x, True)
   def __rpow__(self, x) -> Tensor: return self.pow(x, True)
   def __rtruediv__(self, x) -> Tensor: return self.div(x, True)
   def __rmatmul__(self, x) -> Tensor: return self.matmul(x, True)
+  def __rxor__(self, x) -> Tensor: return self.xor(x, True)
 
   def __iadd__(self, x) -> Tensor: return self.assign(self.add(x))
   def __isub__(self, x) -> Tensor: return self.assign(self.sub(x))
@@ -752,6 +775,7 @@ def __imul__(self, x) -> Tensor: return self.assign(self.mul(x))
   def __ipow__(self, x) -> Tensor: return self.assign(self.pow(x))
   def __itruediv__(self, x) -> Tensor: return self.assign(self.div(x))
   def __imatmul__(self, x) -> Tensor: return self.assign(self.matmul(x))
+  def __ixor__(self, x) -> Tensor: return self.assign(self.xor(x))
 
   def __lt__(self, x) -> Tensor: return mlops.Less.apply(*self._broadcasted(x, False))
   def __gt__(self, x) -> Tensor: return mlops.Less.apply(*self._broadcasted(x, True))
@@ -829,11 +853,3 @@ def is_floating_point(self) -> bool: return dtypes.is_float(self.dtype)
   from tinygrad.features.image import image_conv2d, image_dot
   setattr(Tensor, "conv2d", image_conv2d)
   setattr(Tensor, "dot", image_dot)
-
-# TODO: remove the custom op and replace with threefry
-def custom_random(out:Buffer):
-  Tensor._seed += 1
-  if DEBUG >= 2: print(f"***      rand {out.device} seed {Tensor._seed} size {out.size:<16d} dtype {out.dtype}")
-  rng = np.random.default_rng(Tensor._seed)
-  rng_np_buffer = rng.random(size=out.size, dtype=np.float32).astype(dtype=out.dtype.np, copy=False)
-  out.copyin(rng_np_buffer.data)


Code review comments:
diff_hunk: @@ -156,9 +156,27 @@ def empty(*shape, **kwargs):
   @staticmethod
   def manual_seed(seed=0): Tensor._seed = seed
 
+  _rng_counter: int = 0
path: tinygrad/tensor.py
text: @user1:
from what i read, this will need to be device specific. 

how i was planning to do it in my PR was:
`_rng_counter = defaultdict(int)` and then hash device like `device_id = int(hashlib.sha256('CPU'.encode()).hexdigest(), 16) & 0xFFFFFFFF` to get the current counter for the device

also i think below the `0x0` in `ks` should be the `device_id` 

that way the key/counter contain:
- seed
- device id
- device specific counter

which should allow parallel use without getting the same numbers twice