[DRAFT][s4xbf16] JoinOp vectorization patch

[ggengnv]

Do not merge until rebased on upstream Triton up to triton-lang@c1ed673

This is 1 of the 2 patches needed to improve int4xbf16 GEMM perf.

This is needed because joinOp by default interleaves every element of the two input matrices. In the case of bf16, this means Triton will extract the 2x bf16 values out of the 32-bit register and re-insert them into a new register. This results in many mov instructions before MMA. On certain shapes, this could mean a ~10% perf penalty.

This PR addresses the above by situationally "vectorizing" the interleaving; namely, join every two elements instead of one. This avoids the need to extract values out of registers. Of course, this would also require one to modify the inline_asm logic before the join to produce the correct layout.

cc @gflegar

[ggengnv]

For small-M shapes, for best perf, we'll additionally want XLA to swap A/B so that the LHS of dot is quantized, and then set envvar DISABLE_MMA_V3 to force Ampere-MMA.

[loislo]

I think this cl needs a rebase

[ggengnv]

@loislo, by rebase, do you mean to rebase on openxla/triton's llvm-head, or OAI's triton's main?
If it's the former case, this PR should be up-to-date.
If it's the latter, I can create a new PR targeting OAI triton's main.

[ggengnv]

@loislo please see rebased PR #25