Enable A^T GEMM for BF16

[itaraban]

Could you please enable such case in GEMM for BF16?
This is related to DGL project optimizations

[hfp]

Is there a source file in DGL referring to what's needed or a model/case that runs into this gap?

[itaraban]

In RCGN model DGL use segment_mm operation, which is currently implemented via torch(so it is super slow for cpu right now),
I prepared branch with LIBXSMM implementation - itaraban/dgl@cc03905
Which use same algorithm as CUDA version - https://github.com/dmlc/dgl/blob/master/src/array/cuda/gather_mm.cu#L202

It works pretty well for float and double, performance gain is up to 3x times for full model.
But I cannot use it for BF16, I got https://github.com/libxsmm/libxsmm/blob/main_stable/src/generator_gemm.c#L344.
In such case we will still use torch for BF16 and model will be 2 times slower than float version .

[hfp]

I checked, we do not have Bf16 TN-case or at least not tested/exercised, i.e., this seems to be a valid issue (beside of TN/A-transpose being an unfortunate case in the hot-path ;-).