[Performance] `A@B.T` - The GEMM performance with the column major B matrix is not as good as row major B matrix.

[chengjunlu]

The performance gap is found in #2347

Need to investigate root cause of the performance drops of the column major B matrix case.
Roughly 1.5x worse than the row major B matrix case.

(I): Detected 7680 spills, recompiling the kernel using large GRF mode
(I): Kernel has now 0 spills
✅ Triton and Torch match
Time for torch: 0.31633758544921875 ms
Time for triton: 0.44517597556114197 ms
Compute A x B.T
OpenCL API not available for this operation
OpenCL API not available for this operation
OpenCL API not available for this operation
OpenCL API not available for this operation
(I): Detected 7680 spills, recompiling the kernel using large GRF mode
(I): Kernel has now 0 spills
✅ Triton and Torch match
Time for torch: 0.3375360071659088 ms
Time for triton: 0.6348815560340881 ms

[Egor-Krivov]

I think this issue is essential for GEMM perf. Very often weights are stored with K dimensions as the last. Even pytorch linear layer does that: weight torch.Tensor – the learnable weights of the module of shape : (out_features, in_features)

https://pytorch.org/docs/stable/generated/torch.nn.Linear.html

[alexbaden]

Adding to this, if the A matrix is column-major we have similar problems.

[Egor-Krivov]

We now have microbenchmarks to track this performance. Currently GeoMean for onednn is ~90-100TFLOPs for both cases of A.T@B and for A@B.T.

A@B.T for triton currently stands at ~60TFLOPs. Dashboard gemm-bt
A.T@B for triton currently stands at ~30TFLOPs, it significantly improved and was ~15TFLOPs recently. Dashboard gemm-at

So onednn is 1.5 times faster for B.T and 3 times faster for A.T

[Egor-Krivov]

@alexbaden Should we change the title to reflect issue with A.T as well or create separate issue for that case?

[alexbaden]

Current Triton tiling strategy for DPAS for AxBT:

oneDNN tiling strategy mapped to Triton (thanks to @Jianhui-Li and @chengjunlu ) :

I plan to try to implement the oneDNN strategy in Triton.