40 lines
1.5 KiB
Python
40 lines
1.5 KiB
Python
from tinygrad import UOp, dtypes
|
|
from tinygrad.uop.ops import AxisType, KernelInfo, AddrSpace
|
|
from extra.gemm.amd_uop_matmul import test_matmul
|
|
|
|
N = 2048
|
|
|
|
# metal has an 8x8 tensor core. this is the indexing
|
|
def mat_idx(buf, g0, g1, warp, u):
|
|
l = [(warp//2**i)%2 for i in range(5)]
|
|
return buf[g0, l[4]*4 + l[2]*2 + l[1], g1, l[3]*4 + l[0]*2 + u]
|
|
|
|
def hand_spec_tc_cores():
|
|
gx = UOp.special(N // 8, "gidx0")
|
|
gy = UOp.special(N // 8, "gidx1")
|
|
warp = UOp.special(32, "lidx0")
|
|
|
|
c = UOp.placeholder((N, N), dtypes.float, slot=0).reshape((N//8, 8, N//8, 8))
|
|
a = UOp.placeholder((N, N), dtypes.float, slot=1).reshape((N//8, 8, N//8, 8))
|
|
b = UOp.placeholder((N, N), dtypes.float, slot=2).reshape((N//8, 8, N//8, 8))
|
|
|
|
gk = UOp.range(N // 8, 0, AxisType.REDUCE)
|
|
|
|
a_tc = UOp.stack(*[mat_idx(a, gx, gk, warp, i) for i in range(2)])
|
|
b_tc = UOp.stack(*[mat_idx(b, gk, gy, warp, i) for i in range(2)])
|
|
|
|
acc = UOp.placeholder((2,), dtypes.float, slot=0, addrspace=AddrSpace.REG)
|
|
acc = acc[0].set(0.0)
|
|
acc = acc[1].set(0.0)
|
|
|
|
acc_load = UOp.stack(acc.after(gk)[0], acc.after(gk)[1])
|
|
out = UOp.wmma(a_tc, b_tc, acc_load, (8, 8, 8), 'METAL', 32)
|
|
|
|
end_loop = UOp.group(*[acc[i].store(out.index(i)) for i in range(2)]).end(gk)
|
|
|
|
sink = UOp.group(*[mat_idx(c.after(end_loop), gx, gy, warp, i).store(acc[i]) for i in range(2)])
|
|
return sink.sink(arg=KernelInfo(name="custom_metal_matmul", opts_to_apply=())).simplify()
|
|
|
|
if __name__ == "__main__":
|
|
test_matmul(hand_spec_tc_cores(), N=N)
|