You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
20 lines
638 B
20 lines
638 B
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")
|
|
|