Accept axis mapping when defining MmaOp

[jacobhinkle]

This keeps the default interface of fusedMultiplySum but also adds an option to provide an MmaOp::AxisMapping object. This mapping defines, for each output dimension, which axis in each operand (if any) corresponds to that output dimension.

This PR does not alter the behavior of mma_utils::MatmulPattern::translateToMmaOp meaning we still have BroadcastOp in translations for Hopper matmuls, but that change should be relatively simpler.

Fixes #3372

The included test only checks that dimensions are properly mapped in an MmaOp defined without broadcast axes. In followup PRs I plan to do the following:

1. Demonstrate scheduling a Hopper matmul with unbroadcasted inputs manually. This should surface any bugs in the lowering of the MmaOp instruction when broadcasts are absent.
2. Ensure that we don't depend on having broadcast dims in the Hopper matmul scheduler. For example, we will handle this case in moveInnerBroadcastLeft and we may also need to adjust the swizzling of the TMA smem load TensorView. At this point we will be able to automatically schedule an MmaOp without broadcasted inputs that has been manually defined using our automatic scheduler.
3. Add an option MatmulPattern::translateToMmaOp(/*avoid_intermediates=*/true) and enable that in the Hopper matmul scheduler. At this point it will be safe for us to accept MatmulOp and LinearOp in the Hopper matmul scheduler.

[jacobhinkle]

This was only needed for defining MmaOp::mAxes() and friends, but:

1. Those methods are never used so I removed them and
2. We can reconstruct that information easily using mma->axisMapping().

[jacobhinkle]

!test

[jacobhinkle]

!test

[jacobhinkle]

!test

[jacobhinkle]

!test

[jacobhinkle]

After this PR, one thing we can do is specify the dimension order of the output of the MmaOp independently from the inputs. When we translate MatmulOp and LinearOp, the output already has logical order M, N and we are free to place K wherever we want, so I'll place it last. I think this will let us avoid using commitLeafToLogical like is done here:

Fuser/tests/cpp/test_mma.cpp

Line 603 in d34553f

tv2->commitLeafToLogical();

So in that case we can see how the AxisMapping is standing in for a root->logical reordering. Since there is one for each input operand this feels like another nice use case for read/write/compute domains as suggested by @zasdfgbnm for indexing ldmatrix.

[zasdfgbnm]

Could you provide an example? For example, what does it mean if we have:

a_axes = {0, 2, 1, 3}
b_axes = {1, 0 , 3, 2}

does it mean a.axis(0) is mapped to b.axis(1)?

[jacobhinkle]

I'll put this in a comment but yes, this would mean we would map the following sets:

{a.axis(0), b.axis(1), out.axis(0)}
{a.axis(2), b.axis(0), out.axis(1)}
{a.axis(1), b.axis(3), out.axis(2)}
{a.axis(3), b.axis(2), out.axis(3)}

[zasdfgbnm]

Let's use broadcast graph instead of permissive graph.

[jacobhinkle]

I am actually planning to do this, but it means I need to update mma_utils since permissive is still used there, so I will do that in a follow-up.

[jacobhinkle]

#3397

[jacobhinkle]

!build

[jacobhinkle]

!build