nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

tf32 tc for nv and ptx (#8635)

Browse Source 
Co-authored-by: George Hotz <[email protected]>
This commit is contained in:
ignaciosica
2025-01-17 17:43:57 -08:00
committed by GitHub
co-authored by George Hotz
parent 5afb0a4a81
commit d2234e308a
7 changed files with 23 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/runtime/ops_python.py
+6 -1
View File
@@ -154,11 +154,16 @@ class PythonProgram:
def b_elem(x, col, k, goff): return x[k%2 + (k//8)*2][goff + (k//2)%4 + col*4]
ul[i] = wmma_helper(32, 16, 8, 4, 4, a_elem, b_elem, c_map)
elif arg[1] == (8,16,8):
elif arg[1] == (8,16,8) and arg[2] == dtypes.half:
def a_elem(x, k, row, goff): return x[k%2 + (row//8)*2][goff + k//2 + (row%8)*4]
def b_elem(x, col, k, goff): return x[k%2][goff + k//2 + col*4]
ul[i] = wmma_helper(32, 8, 4, 2, 4, a_elem, b_elem, c_map)
elif arg[1] == (8,16,8) and arg[2] == dtypes.float:
def a_elem(x, k, row, goff): return x[(k//4)*2 + row//8][goff + k%4 + (row%8)*4]
def b_elem(x, col, k, goff): return x[k//4][goff + k%4 + col*4]
ul[i] = wmma_helper(32, 8, 4, 2, 4, a_elem, b_elem, c_map)
else: raise NotImplementedError(f"unimplemented tensor core {arg}")
elif arg[4] == "INTEL":
# A (16 elements on 8 threads)