diff --git a/test/unit/test_keccak.py b/test/unit/test_hashing.py similarity index 75% rename from test/unit/test_keccak.py rename to test/unit/test_hashing.py index c6dce04282..803883b78f 100644 --- a/test/unit/test_keccak.py +++ b/test/unit/test_hashing.py @@ -2,6 +2,21 @@ from typing_extensions import Callable import hashlib, random, unittest from tinygrad import Tensor, Device, getenv, dtypes from tinygrad.device import is_dtype_supported +from tinygrad.helpers import CI + +@unittest.skipUnless(is_dtype_supported(dtypes.uint8) and is_dtype_supported(dtypes.uint64), "Device must support uint8 and uint64") +@unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT == "NV", "crashes in NV CI") +class TestHashing(unittest.TestCase): + def _python_hash_1mb(self, data:bytes): + chunks = [data[i:i+4096] for i in range(0, len(data), 4096)] + chunk_hashes = [hashlib.shake_128(chunk).digest(16) for chunk in chunks] + return hashlib.shake_128(b''.join(chunk_hashes)).digest(16) + + @unittest.skipIf(CI, "very slow") + def test_abc(self): + expected = self._python_hash_1mb(b"abc" + b"\x00" * (2**20 - 3)) + out = Tensor(b"abc").hash() + self.assertEqual(bytes(out.data()), expected) @unittest.skipUnless(is_dtype_supported(dtypes.uint8) and is_dtype_supported(dtypes.uint64), "Device must support uint8 and uint64") @unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT == "NV", "crashes in NV CI") @@ -50,10 +65,8 @@ class TestKeccak(unittest.TestCase): data = b"\x00" * 4 self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16)) - data = b"\x00" * 4096 - with self.assertRaises(RecursionError): - # TODO: fix - self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16)) + data = b"\x00" * (1000 if CI else 4096) + self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16)) if __name__ == "__main__": unittest.main() diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index f614db8ddc..53a7790a62 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -1976,7 +1976,9 @@ class Tensor(MathTrait): # https://keccak.team/keccak_specs_summary.html - def ctensor(l: Sequence[ConstType], dtype: DType = dtypes.uint64): return Tensor.stack(*(Tensor(v, dtype=dtype, device=self.device) for v in l)) + def ctensor(l: Sequence[ConstType], dtype: DType = dtypes.uint64): + # TODO: contiguous is here for compile speed + return Tensor.stack(*(Tensor(v, dtype=dtype, device=self.device) for v in l)).contiguous() rot_offsets = [44, 43, 21, 14, 28, 20, 3, 45, 61, 1, 6, 25, 8, 18, 27, 36, 10, 15, 56, 62, 55, 39, 41, 2] rot_offsets_v0, rot_offsets_v1 = ctensor([0] + [1 << v for v in rot_offsets]), ctensor([1] + [1 << (64 - v) for v in rot_offsets]) @@ -2013,8 +2015,42 @@ class Tensor(MathTrait): # χ and ι step state = state.bitwise_xor(~state.roll(shifts=-1, dims=2) & state.roll(shifts=-2, dims=2)) state = state.flatten(1) ^ rnd_const_masks[i] + # NOTE: kernelize here to prevent internal stack from growing propotional to data size + state = state.kernelize() return state.bitcast(dtypes.uint8)[:,:(obytes:=(200 - rate) // 2)].reshape(*self.shape[:-1], obytes) + def _hash_1mb(self) -> Tensor: + assert self.dtype == dtypes.uint8, "only support uint8 tensors for hashing" + assert self.ndim == 2, "only support batched 1d tensors" + assert self.shape[1] == 1024 * 1024, "only support messages of 1mb" + + blocks = self.shape[0] * self.shape[1] // 4096 + data = self.reshape(blocks, 4096) + block_hashes = data.keccak("shake_128").reshape(self.shape[0], 4096) + return block_hashes.keccak("shake_128").reshape(self.shape[0], 16) + + def hash(self) -> Tensor: + """ + Calculates a 16-byte hash of the tensor. + ```python exec="false source="above" session="tensor" result="python" + t = Tensor(b"Hello World!").hash() + print(t.data().hex()) + ``` + """ + + data = self.flatten().bitcast(dtypes.uint8) + if (tsize := data.shape[0]) % 2**20 != 0: data = data.pad((0, 2**20 - tsize % 2**20)) + base_chunks = ceildiv(data.shape[0], 2**20) + tree_depth = math.ceil(math.log(base_chunks, 65536)) if base_chunks > 1 else 0 + + level_chunks = base_chunks + for _ in range(tree_depth + 1): + data = data.reshape(level_chunks, 2**20)._hash_1mb().flatten() + if (tsize := data.shape[0]) % 2**20 != 0: data = data.pad((0, 2**20 - tsize % 2**20)) + level_chunks = ceildiv(data.shape[0], 2**20) + + return data[:16] + def _softmax(self, axis, dtype:DTypeLike|None=None) -> tuple[Tensor, Tensor, Tensor]: m = self - self.max(axis=axis, keepdim=True).detach() if dtype is not None: m = m.cast(dtype)