Efficient GEMV

[dhernandez0]

We got the request from migraphx to improve our vector-matrix multiplication. So, if M=1, N>>1 and K>>1, if we use (for example) wmma instructions, we are padding M to 16 with zeros and doing extra work, using only 1/16 of the instruction output.

Possible implementation

I think we could use the non-accel path of rocmlir for this, this would be a strategy:

1. Make sure we can tune with smaller m/n per block for example: 1, 2, 4, 8 (maybe with 8 it makes sense to use wmma?)
2. it might make sense to skip LDS of the matrix (then maybe we can make kperblock bigger?)
3. make sure we issue V_DUAL_DOT2ACC_F32_F16 instructions so that we do two dot products in one cycle (kPack=2?)
4. splitk is going to be important, as a second stage, try to improve it, can we use the LDS reduction trick (blockwise_broadcast_reduce)?

Steps/Tickets

0. write HIP code to verify if this makes sense
1. refactor non-accel code into an accelemitter (optional but would be nice)
2. changes in tuning params when n/m < 8 (or 4?)
3. changes to speed up GEMV: skip LDS for the matrix (not the vector), make sure we issue V_DUAL_DOT2ACC_F32_F16
4. further splitk investigation (fourth point in "Possible implementation")

[dhernandez0]

input from @manupak: there's batched 4x4 mfmas which could be used for this (if we have a batch dimension).

[umangyadav]

WMMA only has 16x16x16

[dhernandez0]

yes, sorry MFMA

[manupak]

in the absense of batching, one could technically do 4x64 or 64x4 using broadcast modifier of 4x4 mfma where it will broadcast d=4 dimension implicitly

[dhernandez0]

After doing some experiments and analysis, I realized that the arithmetic intensity is almost 1. So, the kernel is bandwidth bound. Using https://github.com/ROCm/rocm_bandwidth_test I get an estimate of effective device bandwidth (not sure if this is the best way to do it):
gfx942 -> 3847 GB/s
gfx1101 -> 329 GB/s

So, just loading A and B and storing C should take about 63k ns (gfx942) or 742k ns (gfx1101). Using tuningRunner+perfRunner (exhaustive tuning, transB=true) we get 58k ns for gfx942 and 729k ns for gfx1101. So, it looks like there's no improvement possible given that the kernel is limited by bandwith. We can focus on fusing this kernel with previous or next operations.

Additionally, I've coded a specialized HIP kernel to do GEMV and we get the same run-time.