commaai
/
tinygrad
mirror of https://github.com/commaai/tinygrad.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

feat: make buffer (#6745)

This commit is contained in:
wozeparrot 2024-09-25 18:31:03 +08:00 committed by GitHub
parent c100f3d406
commit 4ebc9589a6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 15 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
extra/threefry.py
View File

@ -1,10 +1,16 @@
import os
if "DEBUG" not in os.environ: os.environ["DEBUG"] = "2"
if __name__ == "__main__":
import os
if "DEBUG" not in os.environ: os.environ["DEBUG"] = "2"
from tinygrad import Tensor, GlobalCounters
from tinygrad import Tensor, GlobalCounters
from tinygrad.helpers import getenv
for N in [10_000_000, 100_000_000, 1_000_000_000]:
GlobalCounters.reset()
t = Tensor.rand(N)
t.realize()
print(f"N {N:>20_}, global_ops {GlobalCounters.global_ops:>20_}, global_mem {GlobalCounters.global_mem:>20_}")
if (seed := getenv("SEED", 0)) != 0:
Tensor.manual_seed(seed)
print(f"using seed {Tensor._seed}")
for N in [10_000_000, 100_000_000, 1_000_000_000]:
GlobalCounters.reset()
t = Tensor.rand(N)
t.realize()
print(f"N {N:>20_}, global_ops {GlobalCounters.global_ops:>20_}, global_mem {GlobalCounters.global_mem:>20_}")

2
tinygrad/tensor.py
View File

@ -416,7 +416,7 @@ class Tensor:
@staticmethod
def _threefry_random_bits(key0, key1, counts0, counts1):
x = (counts1.cast(dtypes.uint64) << 32) | counts0.cast(dtypes.uint64)
key = (Tensor(key0, device=x.device, dtype=dtypes.uint64, requires_grad=False) << 32) | key1
key = (Tensor([key0], device=x.device, dtype=dtypes.uint64, requires_grad=False) << 32) | key1
x = F.Threefry.apply(*x._broadcasted(key))
counts0, counts1 = (x & 0xffffffff).cast(dtypes.uint32), ((x >> 32) & 0xffffffff).cast(dtypes.uint32)
return counts0.cat(counts1)