diff --git a/test/null/test_tensor_uop_mixin.py b/test/null/test_tensor_uop_mixin.py index 02fbafb9ae..3ec3e7414b 100644 --- a/test/null/test_tensor_uop_mixin.py +++ b/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))) diff --git a/tinygrad/mixin/__init__.py b/tinygrad/mixin/__init__.py index 437ad81768..353046dade 100644 --- a/tinygrad/mixin/__init__.py +++ b/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") diff --git a/tinygrad/mixin/rand.py b/tinygrad/mixin/rand.py new file mode 100644 index 0000000000..a66d016ef5 --- /dev/null +++ b/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:]) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index e7833e03fc..17f3973f62 100644 --- a/tinygrad/tensor.py +++ b/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)