21 lines
638 B
Python
21 lines
638 B
Python
|
import time
|
||
|
|
from tinygrad import Tensor, Device, TinyJit
|
||
|
|
from tinygrad.helpers import getenv
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
DEVS = [f"NV:{i}" for i in range(getenv("GPUS", 2))]
|
||
|
|
N = getenv("N", 8192)
|
||
|
|
A = Tensor.rand(N, N).shard(DEVS, 0).realize()
|
||
|
|
B = Tensor.rand(N, N).shard(DEVS, 1).realize()
|
||
|
|
print("***** MUL *****")
|
||
|
|
jmatmul = TinyJit(Tensor.dot)
|
||
|
|
for i in range(10):
|
||
|
|
Device["NV:0"].synchronize()
|
||
|
|
Device["NV:1"].synchronize()
|
||
|
|
st = time.perf_counter()
|
||
|
|
jmatmul(A, B)
|
||
|
|
Device["NV:0"].synchronize()
|
||
|
|
Device["NV:1"].synchronize()
|
||
|
|
et = time.perf_counter()
|
||
|
|
print(f"{(N*N*N*2*1e-12)/(et-st):.2f} TFLOPS")
|