Inject desired pattern for handling Transpose for fp8 gemm rewrite

[wenscarl]

Related to #17276 and #16975.
This PR updates the GemmRewriter to handle the transpose of non-descending layouts directly, eliminating the need for the layout_normalization pass to correct this error-prone pattern post-rewrite. The desired transformation is now injected into GemmRewriter, ensuring the problematic layout is handled internally. This PR transforms the following error-prone pattern, where the transpose of a non-descending layout is the issue:

a = f8e4m3fn[x,y]{0,1} xxx
transpose.0 = f8e4m3fn[y,x]{0,1} transpose(a), dimensions=(1,0)
custom-call(a,...)

to

a = f8e4m3fn[x,y]{0,1} xxx
bt = f8e4m3fn[y,x]{1,0} bitcast(a)
transpose.1 = f8e4m3fn[x,y]{1,0} transpose(bt), dimensions=(1,0)
bt.1= f8e4m3fn[y,x]{0,1} bitcast(transpose.1)
custom-call(bt.1,...)

[elfiegg]

Also cc @akuegel

[akuegel]

Can you please also add a test?

[akuegel]

I don't know how to mark a conversation again as unresolved, but for visibility I am copying my answer to a resolved conversation here as well:

I am not sure whether I understand your argument correctly, but to me it seems you assume that RowMajor means that the XLA layout of the operand will be the default layout? Why can we assume that? And even if it turns out that we don't need bitcast-transpose-bitcast, because the operand already has default layout, AlgebraicSimplifier would simplify the bitcasts away. So my suggestion is to always generate the bitcast - transpose - bitcast sequence, and let AlgebraicSimplifier clean it up later. Alternatively, don't distinguish between whether the call was for the RowMajor->ColumnMajor swap or for the ColumnMajor -> RowMajor swap, and instead check whether the operand has default layout or not (and in case of default layout, use the if (!col_maj) code path. That would be cleaner, as we don't need the col_maj parameter, and we are checking exactly the precondition for whether we need bitcast-transpose-bitcast or a single transpose will be enough.

[akuegel]

Internal testing revealed some issues. Please see my comments.

[philipphack]

Would it be conceptually simpler to insert a copy before the transpose to change the layout of the input? (This assumes that the copy -> transpose sequence is optimized by another pass which I haven't verified.)

Also, can we pick a more descriptive variable name here?

[akuegel]

This pass runs after layout normalization which turns copies into bitcast + transpose, so it should not produce any Copy ops that change the layout (otherwise we would have to run layout normalization again).