[XLA] Cublaslt fp8 matmul restriction work-around
/CC @philipphack  @reedwm This is largely independent of the requirement of transposing A but not B. The column-major/row-major layout returned by GemmConfig::For doesn’t fully describe the configuration of a GEMM. Not considering batch dimensions, A and B each can have two possible contracting dimensions (0, 1) and two possible minor-to-major orders ({0,1}, {1,0}).
 
The idea is to bitcast/transpose the operands into the configuration supported by cuBLASLt (A and B have contracting dimensions 1 and 0 and both have {1, 0} orders). Since this only affects FP8 GEMMs, I think it can make sense to introduce this here without changing the general logic. > @reedwm This is largely independent of the requirement of transposing A but not B. The column-major/row-major layout returned by GemmConfig::For doesn’t fully describe the configuration of a GEMM. Not considering batch dimensions, A and B each can have two possible contracting dimensions (0, 1) and two possible minor-to-major orders ({0,1}, {1,0}).
> 
> The idea is to bitcast/transpose the operands into the configuration supported by cuBLASLt (A and B have contracting dimensions 1 and 0 and both have {1, 0} orders). Since this only affects FP8 GEMMs, I think it can make sense to introduce this here without changing the general logic.

@philipphack @wenscarl @reedwm what is the next step here?
- reedwm was going to give the non-TN invocation of cuBLAS lt  a try
- wenscarl/philipphack were to look into reedwm's suggestion about moving some of the transposing logic to the matmul_util

Is that still the plan? I confirmed the non-TN invocation works (the NN invocation, where neither input is transposed).

I'm working on modifying this PR to run with NN, then I'll share it and ask if you think it's clearer.

@philipphack @wenscarl do you plan on adding more test cases? If not I can also add some. Eventually can we please streamline and shorten the tests by perhaps parametrizing them? This would also help with readability. https://google.github.io/googletest/reference/testing.html#TEST_P Thanks for the comments.  The monotonic layout is not necessary but only serves as a means to simplify the logic thereafter. This will be fixed. This HLO is a batched matmul and the batch dim being the first dim which falls in to the cases currently supported. But if batch dim is not first dim, e.g. [16,2,32], no matter fp8 gemm rewrite succeed or not, this case is not supported anyway.  So in this PR, we don't plan to check for batch dim being first dim.  Can you resolve branch conflicts?
As of 2/1/2023, cublasLt f8 matmul only support col major input(default to cublas). But calling from TF/XLA, input can be in all kinds of storage type. This PR aims to "canonicalize" fp8 matmuls by having lrs/rhs_contracting_dim={1,0} and adding necessary transposes to inputs. A reproducer of this bug restriction is located at [here](https://github.com/wenscarl/fp8_gemm_test/blob/main/fp8_gemm_backward_fail.py).

A remaining restriction is the batch dimension still needs to be a leading dimension.