_threefry_random_bits to mixin (#15959)

start RandMixin

test/null/test_tensor_uop_mixin.py

@@ -137,6 +137,14 @@ class TestTensorUOpAllclose(unittest.TestCase):
     a, b = _t(4).float(), _t(4).float()
     self.assertIs(_strip_unique(a.allclose(b).uop), _strip_unique(a.uop.allclose(b.uop)))
 
+class TestTensorUOpRand(unittest.TestCase):
+  def test_random_bits(self):
+    k = UOp.empty((2,), dtype=dtypes.uint32)
+    c = UOp.zeros(2, dtype=dtypes.uint32)
+    for num in (1, 4, 7, 1024):
+      self.assertIs(_strip_unique(Tensor.random_bits(Tensor(k), Tensor(c), num).uop),
+                    _strip_unique(UOp.random_bits(k, c, num)))
+
 class TestTensorUOpGather(unittest.TestCase):
   def _check(self, t, dim, idx):
     self.assertIs(_strip_unique(t.gather(dim, idx).uop), _strip_unique(t.uop.gather(dim, idx.uop)))

tinygrad/mixin/__init__.py

@@ -3,6 +3,7 @@ import functools, itertools
 from typing import TYPE_CHECKING, Callable, Self, Sequence, Literal, get_args
 from tinygrad.mixin.elementwise import ElementwiseMixin
 from tinygrad.mixin.movement import MovementMixin
+from tinygrad.mixin.rand import RandMixin
 from tinygrad.mixin.reduce import ReduceMixin
 from tinygrad.uop import Ops
 from tinygrad.uop.ops import _broadcast_shape, resolve, smax, smin, identity_element
@@ -15,7 +16,7 @@ if TYPE_CHECKING:
 ReductionStr = Literal["mean", "sum", "none"]
 
 
-class OpMixin(ElementwiseMixin, ReduceMixin):
+class OpMixin(ElementwiseMixin, ReduceMixin, RandMixin):
   @staticmethod
   def unique_const(fill_value:ConstType, **kwargs): raise NotImplementedError("creation helpers are only supported on Tensor and UOp")
 

tinygrad/mixin/rand.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+from typing import Self
+from tinygrad.dtype import dtypes
+from tinygrad.helpers import ceildiv
+
+class RandMixin:
+  @staticmethod
+  def _threefry_random_bits(key, counts0, counts1):
+    x = (counts1.cast(dtypes.uint64) << 32) | counts0.cast(dtypes.uint64)
+    x = x.threefry((key[1]._broadcast_to(x.shape).cast(dtypes.uint64) << 32) | key[0]._broadcast_to(x.shape).cast(dtypes.uint64))
+    return (x & 0xffffffff).cast(dtypes.uint32).cat(((x >> 32) & 0xffffffff).cast(dtypes.uint32))
+
+  @classmethod
+  def random_bits(cls, key:Self, counter:Self, num:int) -> Self:
+    low, high = counter[0:1], counter[1:2]  # type: ignore[index]
+    bits = []
+    for i in range(0, num, dtypes.uint32.max):
+      chunk_num = min(num - i, dtypes.uint32.max)
+      c_low = low + (i & 0xffffffff)
+      c_high = high + (i >> 32) + (c_low < low).cast(dtypes.uint32)
+      new_key = cls._threefry_random_bits(key, c_low, c_high)
+      counts0 = cls.arange(ceildiv(chunk_num, 2), device=key.device, dtype=dtypes.uint32)  # type: ignore[attr-defined]
+      counts1 = counts0 + ceildiv(chunk_num, 2)
+      bits.append(cls._threefry_random_bits(new_key, counts0, counts1)[:chunk_num])
+    return bits[0].cat(*bits[1:])

tinygrad/tensor.py

@@ -551,13 +551,6 @@ class Tensor(OpMixin):
     """
     Tensor._seed, Tensor._device_seeds, Tensor._device_rng_counters = seed, {}, {}
 
-  @staticmethod
-  def _threefry_random_bits(key:Tensor, counts0:Tensor, counts1:Tensor) -> Tensor:
-    x = (counts1.cast(dtypes.uint64) << 32) | counts0.cast(dtypes.uint64)
-    x = x._apply_uop(UOp.threefry, (key[1]._broadcast_to(x.shape).cast(dtypes.uint64) << 32) | key[0]._broadcast_to(x.shape).cast(dtypes.uint64))
-    counts0, counts1 = (x & 0xffffffff).cast(dtypes.uint32), ((x >> 32) & 0xffffffff).cast(dtypes.uint32)
-    return counts0.cat(counts1)
-
   @staticmethod
   def rand(*shape, device:str|None=None, dtype:DTypeLike|None=None, contiguous:bool=True, **kwargs) -> Tensor:
     """
@@ -596,18 +589,7 @@ class Tensor(OpMixin):
 
     low = Tensor._device_rng_counters[device][0:1] - (num & 0xffffffff)
     high = Tensor._device_rng_counters[device][1:2] - (num >> 32) - (Tensor._device_rng_counters[device][0] < (num & 0xffffffff)).cast(dtypes.uint32)
-
-    # threefry random bits
-    bits_list = []
-    for i in range(0, num, dtypes.uint32.max):
-      chunk_num = min(num - i, dtypes.uint32.max)
-      c_low = low + (i & 0xffffffff)
-      c_high = high + (i >> 32) + (c_low < low).cast(dtypes.uint32)
-      new_key = Tensor._threefry_random_bits(Tensor._device_seeds[device], c_low, c_high)
-      counts0 = Tensor.arange(ceildiv(chunk_num, 2), device=device, dtype=dtypes.uint32, requires_grad=False)
-      counts1 = counts0 + ceildiv(chunk_num, 2)
-      bits_list.append(Tensor._threefry_random_bits(new_key, counts0, counts1)[:chunk_num])
-    bits = Tensor.cat(*bits_list)
+    bits = Tensor.random_bits(Tensor._device_seeds[device], low.cat(high), num)
 
     # bitcast to uint with same number of bits
     _, nmant = dtypes.finfo(dt)