tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 09:36:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix metal matmul example

Browse Source
This commit is contained in:
geohot
2023-03-22 23:42:51 -07:00
parent f5aea472a3
commit c3a3db75c7
1 changed files with 4 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/gemm/metal_matmul.py
+4 -9
View File
@@ -1,19 +1,14 @@
import numpy as np
from tinygrad.helpers import dtypes
from tinygrad.runtime.ops_metal import RawMetalBuffer, MetalProgram
N = 2048
a = RawMetalBuffer(N*N*4)
b = RawMetalBuffer(N*N*4)
c = RawMetalBuffer(N*N*4)
a = RawMetalBuffer(N*N, dtypes.float32)
nb = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32) #.astype(np.int32).astype(np.float32)
nc = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32) #.astype(np.int32).astype(np.float32)
#nb = np.eye(N)
#nc = np.eye(N)
#nb = np.ones((N,N))
#nc = np.ones((N,N))
b.copyin(nb)
c.copyin(nc)
b = RawMetalBuffer.fromCPU(nb)
c = RawMetalBuffer.fromCPU(nc)
FLOPS = N*N*N*2