[MFMA][FRONTEND] Add more options for forced mfma layout sizes - Githubissues

ROCm / triton

Development repository for the Triton language and compiler

MIT License

80 stars 22 forks source link

[MFMA][FRONTEND] Add more options for forced mfma layout sizes #538

Open binarman opened 3 months ago

binarman commented 3 months ago

This PR:

adds an matrix_instr_nonkdim options to force MFMA 64x4 and 4x64 layout: 464 corresponds 4(M)x64(N), 644 corresponds 64(M)x4(N)
adds tests for this option
fixes swizzling patter in some cases

MFMA size heuristic now looks like this:

If kernel specific option is set, pick it
If the result tile shape is larger than 32x32, pick mfma32
If the tile shape is smaller than 32x32 but larger than 16x16, pick mfma16
if the tile shape is smaller than 4x64 or 64x4, pick mfma4x4
Otherwise, pick mfma4x64 or mfma64x4, depending on what tile fits into matrices