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

cleaner mxfp4 gemm prelude (#17796)

Browse Source 
* cleaner mxfp4 prelude

* rename sgprs

* min diff
This commit is contained in:
qazal
2026-08-28 15:59:35 +09:00
committed by GitHub
parent 92d92c70a9
commit 6242b09066
1 changed files with 68 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/gemm/gemm_mxfp4.py
+68 -115
View File
@@ -20,34 +20,39 @@ def v_mfma_fp4(dst, a, b, opsel, opsel_hi, scale_a, scale_b):
def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k = Kernel()
scale_k = K // 32
k.emit(s_and_b32(s[1], s[1], LIT, 65535))
if (tile_m, tile_n) == (128, 512):
k.emit(s_and_b32(s[1], s[1], LIT, 65535))
k.emit(s_mov_b32(s[47], s[2]))
k.emit(s_mov_b32(s[48], s[3]))
k.emit(s_load_dwordx2(s[4:5], s[0:1], s[0], 0, 0, 0, 0, 1))
k.emit(s_mov_b32(s[8], 0))
k.emit(s_mov_b32(s[9], 0))
k.emit(s_load_dwordx2(s[12:13], s[0:1], s[0], 8, 0, 0, 0, 1))
k.emit(s_load_dwordx2(s[16:17], s[0:1], s[0], 16, 0, 0, 0, 1))
k.emit(s_mov_b32(s[36], N))
k.emit(s_mov_b32(s[37], K))
k.emit(s_mov_b32(s[38], K))
k.emit(s_mov_b32(s[43], M))
k.emit(s_mov_b32(s[44], N))
k.emit(s_mov_b32(s[45], K))
k.emit(s_load_dwordx2(s[20:21], s[0:1], s[0], 24, 0, 0, 0, 1))
k.emit(s_load_dwordx2(s[24:25], s[0:1], s[0], 32, 0, 0, 0, 1))
k.emit(s_mov_b32(s[39], scale_k))
k.emit(s_mov_b32(s[40], scale_k))
k.emit(v_lshrrev_b32_e32(v[1], 10))
k.emit(v_lshrrev_b32_e32(v[2], 10, v[1]))
k.emit(v_and_b32_e32(v[2], LIT, v[2], 1023))
k.emit(v_and_b32_e32(v[1], LIT, v[1], 1023))
k.emit(v_and_b32_e32(v[0], LIT, v[0], 1023))
k.emit(v_lshrrev_b32_e32(v[3], 6))
k.emit(v_and_b32_e32(v[0], 63))
k.emit(v_readfirstlane_b32_e32(v[46], v[3]))
k.emit(s_waitcnt(49279))
k.emit(s_load_dwordx2(s[4:5], s[0:1], s[0], 0, 0, 0, 0, 1))
k.emit(s_mov_b32(s[8], 0))
k.emit(s_mov_b32(s[9], 0))
k.emit(s_load_dwordx2(s[12:13], s[0:1], s[0], 8, 0, 0, 0, 1))
k.emit(s_load_dwordx2(s[16:17], s[0:1], s[0], 16, 0, 0, 0, 1))
k.emit(s_mov_b32(s[36], N))
k.emit(s_mov_b32(s[37], K))
k.emit(s_mov_b32(s[38], K))
k.emit(s_mov_b32(s[43], M))
k.emit(s_mov_b32(s[44], N))
k.emit(s_mov_b32(s[45], K))
k.emit(s_load_dwordx2(s[20:21], s[0:1], s[0], 24, 0, 0, 0, 1))
k.emit(s_load_dwordx2(s[24:25], s[0:1], s[0], 32, 0, 0, 0, 1))
k.emit(s_mov_b32(s[39], scale_k))
k.emit(s_mov_b32(s[40], scale_k))
k.emit(v_lshrrev_b32_e32(v[1], 10))
k.emit(v_lshrrev_b32_e32(v[2], 10, v[1]))
k.emit(v_and_b32_e32(v[2], LIT, v[2], 1023))
k.emit(v_and_b32_e32(v[1], LIT, v[1], 1023))
k.emit(v_and_b32_e32(v[0], LIT, v[0], 1023))
k.emit(v_lshrrev_b32_e32(v[3], 6))
k.emit(v_and_b32_e32(v[0], 63))
if (tile_m, tile_n) == (256, 256):
k.emit(s_mov_b32(s[49], s[2]))
k.emit(s_mov_b32(s[47], s[3]))
k.emit(v_readfirstlane_b32_e32(v[46], v[3]))
k.emit(s_waitcnt(49279))
if (tile_m, tile_n) == (128, 512):
for i in range(2):
k.emit(s_mov_b32(s[6 + i * 8], -16))
k.emit(s_mov_b32(s[10 + i * 12], -16))
@@ -1213,31 +1218,6 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_waitcnt())
k.emit(s_endpgm())
elif (tile_m, tile_n) == (192, 256):
k.emit(s_and_b32(s[1], s[1], LIT, 65535))
k.emit(s_load_dwordx2(s[4:5], s[0:1], s[0], 0, 0, 0, 0, 1))
k.emit(s_mov_b32(s[8], 0))
k.emit(s_mov_b32(s[9], 0))
k.emit(s_load_dwordx2(s[12:13], s[0:1], s[0], 8, 0, 0, 0, 1))
k.emit(s_load_dwordx2(s[16:17], s[0:1], s[0], 16, 0, 0, 0, 1))
k.emit(s_mov_b32(s[36], N))
k.emit(s_mov_b32(s[37], K))
k.emit(s_mov_b32(s[38], K))
k.emit(s_mov_b32(s[43], M))
k.emit(s_mov_b32(s[44], N))
k.emit(s_mov_b32(s[45], K))
k.emit(s_load_dwordx2(s[20:21], s[0:1], s[0], 24, 0, 0, 0, 1))
k.emit(s_load_dwordx2(s[24:25], s[0:1], s[0], 32, 0, 0, 0, 1))
k.emit(s_mov_b32(s[39], scale_k))
k.emit(s_mov_b32(s[40], scale_k))
k.emit(v_lshrrev_b32_e32(v[1], 10))
k.emit(v_lshrrev_b32_e32(v[2], 10, v[1]))
k.emit(v_and_b32_e32(v[2], LIT, v[2], 1023))
k.emit(v_and_b32_e32(v[1], LIT, v[1], 1023))
k.emit(v_and_b32_e32(v[0], LIT, v[0], 1023))
k.emit(v_lshrrev_b32_e32(v[3], 6))
k.emit(v_and_b32_e32(v[0], 63))
k.emit(v_readfirstlane_b32_e32(v[46], v[3]))
k.emit(s_waitcnt(49279))
k.emit(s_mul_i32(s[63], LIT, 8, 192))
k.emit(v_cvt_f32_u32_e32(v[4], s[63]))
k.emit(s_sub_i32(s[62], 0, s[63]))
@@ -2234,49 +2214,22 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_waitcnt())
k.emit(s_endpgm())
elif (tile_m, tile_n) == (256, 256):
k.emit(s_and_b32(s[1], s[1], LIT, 65535))
k.emit(s_load_dwordx2(s[4:5], s[0:1], s[0], 0, 0, 0, 0, 1))
k.emit(s_mov_b32(s[8], 0))
k.emit(s_mov_b32(s[9], 0))
k.emit(s_load_dwordx2(s[12:13], s[0:1], s[0], 8, 0, 0, 0, 1))
k.emit(s_load_dwordx2(s[16:17], s[0:1], s[0], 16, 0, 0, 0, 1))
k.emit(s_mov_b32(s[40], N))
k.emit(s_mov_b32(s[41], K))
k.emit(s_mov_b32(s[42], K))
k.emit(s_mov_b32(s[43], M))
k.emit(s_mov_b32(s[44], N))
k.emit(s_mov_b32(s[45], K))
k.emit(s_load_dwordx2(s[20:21], s[0:1], s[0], 24, 0, 0, 0, 1))
k.emit(s_load_dwordx2(s[24:25], s[0:1], s[0], 32, 0, 0, 0, 1))
k.emit(s_mov_b32(s[36], scale_k))
k.emit(s_mov_b32(s[37], scale_k))
k.emit(v_lshrrev_b32_e32(v[1], 10))
k.emit(v_lshrrev_b32_e32(v[2], 10, v[1]))
k.emit(v_and_b32_e32(v[2], LIT, v[2], 1023))
k.emit(v_and_b32_e32(v[1], LIT, v[1], 1023))
k.emit(v_and_b32_e32(v[0], LIT, v[0], 1023))
k.emit(v_lshrrev_b32_e32(v[3], 6))
k.emit(v_and_b32_e32(v[0], 63))
k.emit(s_mov_b32(s[46], s[2]))
k.emit(s_mov_b32(s[47], s[3]))
k.emit(v_readfirstlane_b32_e32(v[49], v[3]))
k.emit(s_waitcnt(49279))
k.emit(s_add_u32(s[55], s[44], LIT, 255))
k.emit(s_lshr_b32(s[54], s[55], 8))
k.emit(s_mul_i32(s[48], s[54], s[47]))
k.emit(s_add_i32(s[48], s[48], s[46]))
k.emit(s_add_i32(s[48], s[48], s[49]))
k.emit(s_add_u32(s[55], s[43], LIT, 255))
k.emit(s_lshr_b32(s[52], s[55], 8))
k.emit(s_lshl_b32(s[52], s[52], 5))
k.emit(s_mov_b32(s[46], 0))
k.emit(s_mov_b32(s[49], 0))
k.label('L2_00E8')
k.emit(s_cmp_lt_i32(s[48], s[52]))
k.emit(s_cbranch_scc1(3), target='L2_00FC')
k.emit(s_sub_i32(s[48], s[48], s[52]))
k.emit(s_add_i32(s[46], s[46], 32))
k.emit(s_add_i32(s[49], s[49], 32))
k.emit(s_branch(65531), target='L2_00E8')
k.label('L2_00FC')
k.emit(s_sub_i32(s[54], s[54], s[46]))
k.emit(s_sub_i32(s[54], s[54], s[49]))
k.emit(s_cmp_lt_i32(s[54], 32))
k.emit(s_cbranch_scc1(3), target='L2_0114')
k.emit(s_lshr_b32(s[47], s[48], 5))
@@ -2311,7 +2264,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_mul_i32(s[52], s[54], s[47]))
k.emit(s_sub_i32(s[52], s[48], s[52]))
k.label('L2_0194')
k.emit(s_add_i32(s[46], s[52], s[46]))
k.emit(s_add_i32(s[49], s[52], s[49]))
k.emit(s_mov_b32(s[6], -16))
k.emit(s_mov_b32(s[10], -16))
k.emit(s_mov_b32(s[18], -16))
@@ -2328,18 +2281,18 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_or_b32(s[9], s[9], LIT, 262144))
k.emit(s_or_b32(s[17], s[17], LIT, 262144))
k.emit(s_or_b32(s[13], s[13], LIT, 262144))
k.emit(s_lshr_b32(s[41], s[41], 1))
k.emit(s_mul_i32(s[52], s[41], s[43]))
k.emit(s_lshr_b32(s[37], s[37], 1))
k.emit(s_mul_i32(s[52], s[37], s[43]))
k.emit(s_mov_b32(s[14], s[52]))
k.emit(s_lshr_b32(s[42], s[42], 1))
k.emit(s_mul_i32(s[52], s[42], s[44]))
k.emit(s_lshr_b32(s[38], s[38], 1))
k.emit(s_mul_i32(s[52], s[38], s[44]))
k.emit(s_mov_b32(s[18], s[52]))
k.emit(s_add_u32(s[52], s[43], 31))
k.emit(s_lshr_b32(s[52], s[52], 5))
k.emit(s_lshl_b32(s[52], s[52], 5))
k.emit(s_mul_i32(s[53], s[52], s[36]))
k.emit(s_mul_i32(s[53], s[52], s[39]))
k.emit(s_mov_b32(s[22], s[53]))
k.emit(s_mul_i32(s[53], s[44], s[37]))
k.emit(s_mul_i32(s[53], s[44], s[40]))
k.emit(s_mov_b32(s[26], s[53]))
k.emit(s_mov_b32(s[23], LIT, 131072))
k.emit(s_mov_b32(s[27], LIT, 131072))
@@ -2356,23 +2309,23 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_add_u32_e32(v[5], v[5], v[6]))
k.emit(v_and_b32_e32(v[4], 1, v[4]))
k.emit(v_add_u32_e32(v[5], v[5], v[4]))
k.emit(v_mul_lo_u32(v[212], s[41], v[5]))
k.emit(v_mul_lo_u32(v[212], s[37], v[5]))
k.emit(v_and_b32_e32(v[4], 7))
k.emit(v_lshlrev_b32_e32(v[4], 4, v[4]))
k.emit(v_add_u32_e32(v[212], v[212], v[4]))
k.emit(s_lshr_b32(s[52], s[49], 1))
k.emit(s_lshr_b32(s[52], s[46], 1))
k.emit(s_mul_i32(s[52], s[52], 8))
k.emit(s_and_b32(s[53], s[49], 1))
k.emit(s_and_b32(s[53], s[46], 1))
k.emit(s_mul_i32(s[53], s[53], 2))
k.emit(s_add_u32(s[52], s[52], s[53]))
k.emit(s_mul_i32(s[53], s[47], LIT, 256))
k.emit(s_add_u32(s[52], s[52], s[53]))
k.emit(s_mul_i32(s[52], s[41], s[52]))
k.emit(s_mul_i32(s[52], s[37], s[52]))
k.emit(v_add_u32_e32(v[212], s[52], v[212]))
k.emit(s_mul_i32(s[52], s[41], 32))
k.emit(s_mul_i32(s[52], s[37], 32))
for i in range(7):
k.emit(v_add_u32_e32(v[213 + i * 1], s[52], v[212 + i * 1]))
k.emit(s_mul_i32(s[59], LIT, s[49], 1056))
k.emit(s_mul_i32(s[59], LIT, s[46], 1056))
k.emit(s_add_u32(s[59], LIT, s[59], 4096))
k.emit(v_and_b32_e32(v[4], 15))
k.emit(v_lshrrev_b32_e32(v[5], 3, v[4]))
@@ -2396,35 +2349,35 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_add_u32_e32(v[221], LIT, v[220], 33792))
k.emit(v_lshlrev_b32_e32(v[222], 2))
k.emit(s_mul_i32(s[52], s[47], LIT, 256))
k.emit(s_mul_i32(s[53], s[49], 32))
k.emit(s_mul_i32(s[53], s[46], 32))
k.emit(s_add_i32(s[52], s[53], s[52]))
k.emit(s_mul_i32(s[53], s[52], s[36]))
k.emit(s_mul_i32(s[53], s[52], s[39]))
k.emit(v_add_u32_e32(v[222], s[53], v[222]))
k.emit(s_mul_i32(s[53], LIT, s[36], 128))
k.emit(s_mul_i32(s[53], LIT, s[39], 128))
k.emit(v_add_u32_e32(v[223], s[53], v[222]))
k.emit(s_mul_i32(s[60], s[49], LIT, 256))
k.emit(s_mul_i32(s[60], s[46], LIT, 256))
k.emit(s_add_i32(s[60], s[60], 0))
k.emit(v_lshlrev_b32_e32(v[224], 2))
k.emit(v_add_u32_e32(v[224], 0, v[224]))
k.emit(v_lshlrev_b32_e32(v[225], 4))
k.emit(s_mul_i32(s[52], s[46], LIT, 256))
k.emit(s_mul_i32(s[53], s[49], 64))
k.emit(s_mul_i32(s[52], s[49], LIT, 256))
k.emit(s_mul_i32(s[53], s[46], 64))
k.emit(s_add_u32(s[52], s[52], s[53]))
k.emit(s_mul_i32(s[52], s[52], s[42]))
k.emit(s_mul_i32(s[52], s[52], s[38]))
k.emit(v_add_u32_e32(v[225], s[52], v[225]))
k.emit(s_mul_i32(s[52], 16, s[42]))
k.emit(s_mul_i32(s[52], 16, s[38]))
k.emit(v_add_u32_e32(v[226], s[52], v[225]))
k.emit(v_add_u32_e32(v[227], s[52], v[226]))
k.emit(v_add_u32_e32(v[228], s[52], v[227]))
for i in range(4):
k.emit(v_add_u32_e32(v[229 + i * 1], LIT, v[225 + i * 1], 1024))
k.emit(v_lshlrev_b32_e32(v[233], 2))
k.emit(s_mul_i32(s[52], s[46], LIT, 256))
k.emit(s_mul_i32(s[53], s[49], 64))
k.emit(s_mul_i32(s[52], s[49], LIT, 256))
k.emit(s_mul_i32(s[53], s[46], 64))
k.emit(s_add_i32(s[52], s[53], s[52]))
k.emit(s_mul_i32(s[53], s[52], s[37]))
k.emit(s_mul_i32(s[53], s[52], s[40]))
k.emit(v_add_u32_e32(v[233], s[53], v[233]))
k.emit(s_mul_i32(s[52], 32, s[37]))
k.emit(s_mul_i32(s[52], 32, s[40]))
k.emit(v_add_u32_e32(v[234], s[52], v[233]))
k.emit(s_mov_b32(s[61], LIT, 128))
k.emit(s_mov_b32(s[62], LIT, 2048))
@@ -2510,18 +2463,18 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(ds_read_b32(v[201], v[224], v[0], v[0], 0, 0, 1))
k.emit(ds_read_b32(v[202], v[224], v[0], v[0], 0, 0, 2))
k.emit(ds_read_b32(v[203], v[224], v[0], v[0], 0, 0, 3))
k.emit(s_lshl_b32(s[40], s[40], 1))
k.emit(s_lshl_b32(s[36], s[36], 1))
k.emit(s_mul_i32(s[52], s[47], LIT, 256))
k.emit(s_mul_hi_u32(s[53], s[52], s[40]))
k.emit(s_mul_hi_u32(s[53], s[52], s[36]))
k.emit(s_add_u32(s[5], s[5], s[53]))
k.emit(s_mul_i32(s[53], s[52], s[40]))
k.emit(s_mul_i32(s[53], s[52], s[36]))
k.emit(s_add_u32(s[4], s[4], s[53]))
k.emit(s_addc_u32(s[5], 0, s[5]))
k.emit(s_sub_i32(s[52], s[43], s[52]))
k.emit(s_mul_i32(s[52], s[52], s[40]))
k.emit(s_mul_i32(s[52], s[52], s[36]))
k.emit(s_mov_b32(s[6], s[52]))
k.emit(v_and_b32_e64(v[235], v[0], 15))
k.emit(v_mul_lo_u32(v[235], v[235], s[40]))
k.emit(v_mul_lo_u32(v[235], v[235], s[36]))
k.emit(v_lshrrev_b32_e32(v[4], 5))
k.emit(v_mul_i32_i24_e32(v[4], 16, v[4]))
k.emit(v_add_u32_e32(v[235], v[4], v[235]))
@@ -2529,12 +2482,12 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(v_and_b32_e32(v[4], 1, v[4]))
k.emit(v_mul_i32_i24_e32(v[4], 32, v[4]))
k.emit(v_add_u32_e32(v[235], v[4], v[235]))
k.emit(s_mul_i32(s[52], s[46], LIT, 256))
k.emit(s_mul_i32(s[53], s[49], 64))
k.emit(s_mul_i32(s[52], s[49], LIT, 256))
k.emit(s_mul_i32(s[53], s[46], 64))
k.emit(s_add_i32(s[52], s[52], s[53]))
k.emit(s_lshl_b32(s[52], s[52], 1))
k.emit(v_add_u32_e32(v[235], s[52], v[235]))
k.emit(s_mul_i32(s[53], s[40], 16))
k.emit(s_mul_i32(s[53], s[36], 16))
for i in range(15):
k.emit(v_add_u32_e64(v[236 + i * 1], v[235 + i * 1], s[53]))
k.emit(s_mov_b32(s[50], 0))
@@ -2543,7 +2496,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
k.emit(s_cmp_lt_u32(LIT, s[51], 512 + i * -256))
k.emit(s_cselect_b32(s[61 + i * 1], s[61 + i * 1], 0))
k.emit(s_cselect_b32(s[63 + i * 1], s[63 + i * 1], 0))
k.emit(s_cmp_lt_i32(s[49], 2))
k.emit(s_cmp_lt_i32(s[46], 2))
k.emit(s_cbranch_scc0(1367), target='L2_25B8')
k.label('L2_105C')
k.emit(s_waitcnt(122))