Files
dragonpilot/tinygrad_repo/extra/gemm/gemm.py
T
Vehicle Researcher dd778596b7 openpilot v0.9.8 release
date: 2025-03-15T21:10:51
master commit: fb7b9c0f9420d228f03362970ebcfb7237095cf3
2025-03-18 10:05:17 -07:00

29 lines
599 B
Python
Executable File

#!/usr/bin/env python3
import os
#os.environ['OMP_NUM_THREADS'] = '1'
import time
import numpy as np
N = 512
if __name__ == "__main__":
# N^2
A = np.random.randn(N, N).astype(np.float32)
# N^2
B = np.random.randn(N, N).astype(np.float32)
# 2N compute in N^2 output cells
flop = 2*N*N*N
#print(f"{flop / 1e9:.2f} GFLOP")
for i in range(10):
st = time.monotonic()
C = A @ B.T
et = time.monotonic()
s = et-st
print(f"{flop/s * 1e-9:.2f} GFLOP/S, {s*1e3:.2f} ms")
with open("/tmp/matmul", "wb") as f:
f.write(A.data)
f.write(B.data)
f.write(C.data)