[QST] What is PermutationMNK in TiledMMA in CUTLASS 3.4 changes?

[hyhieu]

Hi everyone,

I noticed that in CUTLASS 3.4 (#1286), there was a change in cute/atom/mma_atom.hpp that altered the signature of make_tiled_mma and the template of TiledMMA.

Specifically:

CUTLASS 3.4:

// @tparam MMA_Atom The MMA_Atom to use in the TiledMMA
// @tparam AtomLayoutMNK The MNK-tiling of the Atom to be performed.
// @tparam PermuationsMNK Permutations to apply to each MNK-mode before tiling for the Atom.
template <class MMA_Atom,
          class AtomLayoutMNK,
          class PermutationMNK = Tile<Underscore,Underscore,Underscore>>
struct TiledMMA : MMA_Atom { ... }

Before:

template <class MMA_Atom,
          class AtomLayoutMNK   = Layout<Shape<_1,_1,_1>>,
          class ValLayoutMNK    = Layout<Shape<_1,_1,_1>>,
          class PermutationsMNK = Tile<Underscore,Underscore,Underscore>>
struct TiledMMA : MMA_Atom { ... }

So, what is the PermutationMNK in the new interface? How does it achieve the effect of ValLayoutMNK and PermutationsMNK in the older version?

To be clear, thanks to this Zhihu post, I kind of understood that in the old version:

• AtomLayoutMNK means that the MMA_Atom is replicated into more threads, and so the CTA becomes accordingly.
• ValLayoutMNK means that the TiledMMA performs more compute by repeating the MMA_Atom, and so the CTA does not become bigger, even though the TiledShape computed by the TiledMMA still becomes bigger.

Now that ValLayoutMNK is gone, I noticed that this still does the job for me:

using TiledMma = decltype(
  make_tiled_mma(
    MmaAtom{},
    Layout<Shape<_1, _1, _1>>{},  // this should be AtomLayoutMNK
    Layout<Shape<_2, _2, _2>>{}   // this should be PermutationMNK, but what does it do now?
  )
);

but I don't know what is going on with what I passed to make_tiled_mma as PermutationsMNK.

Hoping for an explanation.

Thanks!

[ccecka]

Thanks for the question, my update to the CuTe documentation and examples is still pending but should be available soon.

The easiest way to think about it is that the Permutation parameter is a Tiler for the MNK modes of the MMA. That is, it is a set of three layouts that are applied to the M-mode, N-mode, and K-mode individually before applying the TV-layouts and projective slicing. These layouts can act to permute the M-mode, N-mode, and K-mode individually to make the TV-partitioning patterns more manageable/intuitive and effectively interleave individual MMAs.

I'll start with some examples and then get to your case. Let's start with SM80_8x8x4_F64F64F64F64_TN

    TiledMMA tiled_mma = make_tiled_mma(SM80_8x8x4_F64F64F64F64_TN{});
    print_latex(tiled_mma);

The above is equivalent to

    TiledMMA tiled_mma = make_tiled_mma(SM80_8x8x4_F64F64F64F64_TN{},
                                        Layout<Shape<_1,_1,_1>>{},
                                        Tile<_8,_8,_4>{});
    print_latex(tiled_mma);

as the atom already has a natural tile size of 8x8x4. But we can immediately expand this "tile size" up to 8x16x8 instead:

    TiledMMA tiled_mma = make_tiled_mma(SM80_8x8x4_F64F64F64F64_TN{},
                                        Layout<Shape<_1,_1,_1>>{},     // AtomLayout
                                        Tile<_8,_16,_8>{});            // Tiler
    print_latex(tiled_mma);

This doesn't actually affect the partitioning of input/output tensors because, by convention, only a single atom is ever partitioned out. It will affect the output of tile_size and get_layoutC_MN and get_layoutC_TV etc, which could affect any TiledCopy that rely on those partitioning patterns by being built on this TiledMMA. Regardless, you'll find the resulting tensors from partition_C etc to be exactly the same since the atom partitioning is exactly the same.

Continuing, I see those four values that T0 write to in C are separate, but I would prefer them to be next to each other. That is, I would like to permute the N-mode.

    TiledMMA tiled_mma = make_tiled_mma(SM80_8x8x4_F64F64F64F64_TN{},
                                        Layout<Shape<_1,_1,_1>>{},     // AtomLayout
                                        Tile<_8,                       // Permutation on M, equivalent to 8:1, identity
                                             Layout<Shape <_2,_4,_2>, 
                                                    Stride<_1,_4,_2>>, // Permutation on N, size 16
                                             _8>{});                   // Permutation on K, equivalent to 8:1, identity
    print_latex(tiled_mma);

That layout (2,4,2):(1,4,2) is read like a scatter permutation, telling the n-coords of the original image where to go in the new image: The first 2 elements stay put, the next 4 of those 2-element groups get sent to n-coord 4, and the next 2 of those 8-element groups get sent to n-coord 2. This permutes only the n-mode (in C and B accordingly) and brings the access of all threads to be contiguous in n-coordinates. This is convenient when designing layouts for shared memory or registers, for example. The two MMA instructions contained within the image above are now effectively interleaved in the logical n-coordinates.

Your example is doing something a bit silly, but perfectly functional.

using TiledMma = decltype(
  make_tiled_mma(
    MmaAtom{},
    Layout<Shape<_1, _1, _1>>{},  // this should be AtomLayoutMNK
    Layout<Shape<_2, _2, _2>>{}   // this should be PermutationMNK, but what does it do now?
  )
);

Because make_tiled_mma massages the inputs to get reasonable defaults by:

make_tiled_mma(MMA_Atom<MMA_Op> const& mma_atom,
               MMAThrLayout     const& thr_layout   = {},
               Permutations     const& permutations = {})
{
  auto thr_layout_mnk  = append<3>(thr_layout, Layout<_1,_0>{});
  auto permutation_mnk = append<3>(permutations, _);

  return TiledMMA<MMA_Atom<MMA_Op>,
                  decltype(thr_layout_mnk),
                  decltype(permutation_mnk)>{mma_atom, thr_layout_mnk};
}

You're getting the following as a result

TiledMMA<AtomType,
         Layout<Shape<_1,_1,_1>>,       // AtomLayoutMNK
         Tile<Layout<Shape<_2,_2,_2>>,  // Permutation on M, (2,2,2):(1,2,4), identity of size 8
              Underscore,               // Permutation on N, noop identity
              Underscore>>              // Permutation on K, noop identity

So you can see that your Layout is becoming a rank-3 Tiler, but the permutations on each mode are all identities anyway and have little to no effect.

Hope that helps! I'll try to get explanations and examples like these into the documentation as soon as possible.

[hyhieu]

This is an amazing explanation. Thank you, Cris. Cannot wait to see more examples and read the documentation from you.

[MARD1NO]

Thanks for the question, my update to the CuTe documentation and examples is still pending but should be available soon.

The easiest way to think about it is that the Permutation parameter is a Tiler for the MNK modes of the MMA. That is, it is a set of three layouts that are applied to the M-mode, N-mode, and K-mode individually before applying the TV-layouts and projective slicing. These layouts can act to permute the M-mode, N-mode, and K-mode individually to make the TV-partitioning patterns more manageable/intuitive and effectively interleave individual MMAs.

I'll start with some examples and then get to your case. Let's start with SM80_8x8x4_F64F64F64F64_TN

    TiledMMA tiled_mma = make_tiled_mma(SM80_8x8x4_F64F64F64F64_TN{});
    print_latex(tiled_mma);

The above is equivalent to

    TiledMMA tiled_mma = make_tiled_mma(SM80_8x8x4_F64F64F64F64_TN{},
                                        Layout<Shape<_1,_1,_1>>{},
                                        Tile<_8,_8,_4>{});
    print_latex(tiled_mma);

as the atom already has a natural tile size of 8x8x4. But we can immediately expand this "tile size" up to 8x16x8 instead:

    TiledMMA tiled_mma = make_tiled_mma(SM80_8x8x4_F64F64F64F64_TN{},
                                        Layout<Shape<_1,_1,_1>>{},     // AtomLayout
                                        Tile<_8,_16,_8>{});            // Tiler
    print_latex(tiled_mma);

This doesn't actually affect the partitioning of input/output tensors because, by convention, only a single atom is ever partitioned out. It will affect the output of tile_size and get_layoutC_MN and get_layoutC_TV etc, which could affect any TiledCopy that rely on those partitioning patterns by being built on this TiledMMA. Regardless, you'll find the resulting tensors from partition_C etc to be exactly the same since the atom partitioning is exactly the same.

Continuing, I see those four values that T0 write to in C are separate, but I would prefer them to be next to each other. That is, I would like to permute the N-mode.

    TiledMMA tiled_mma = make_tiled_mma(SM80_8x8x4_F64F64F64F64_TN{},
                                        Layout<Shape<_1,_1,_1>>{},     // AtomLayout
                                        Tile<_8,                       // Permutation on M, equivalent to 8:1, identity
                                             Layout<Shape <_2,_4,_2>, 
                                                    Stride<_1,_4,_2>>, // Permutation on N, size 16
                                             _8>{});                   // Permutation on K, equivalent to 8:1, identity
    print_latex(tiled_mma);

That layout (2,4,2):(1,4,2) is read like a scatter permutation, telling the n-coords of the original image where to go in the new image: The first 2 elements stay put, the next 4 of those 2-element groups get sent to n-coord 4, and the next 2 of those 8-element groups get sent to n-coord 2. This permutes only the n-mode (in C and B accordingly) and brings the access of all threads to be contiguous in n-coordinates. This is convenient when designing layouts for shared memory or registers, for example. The two MMA instructions contained within the image above are now effectively interleaved in the logical n-coordinates.

Your example is doing something a bit silly, but perfectly functional.

using TiledMma = decltype(
  make_tiled_mma(
    MmaAtom{},
    Layout<Shape<_1, _1, _1>>{},  // this should be AtomLayoutMNK
    Layout<Shape<_2, _2, _2>>{}   // this should be PermutationMNK, but what does it do now?
  )
);

Because make_tiled_mma massages the inputs to get reasonable defaults by:

make_tiled_mma(MMA_Atom<MMA_Op> const& mma_atom,
               MMAThrLayout     const& thr_layout   = {},
               Permutations     const& permutations = {})
{
  auto thr_layout_mnk  = append<3>(thr_layout, Layout<_1,_0>{});
  auto permutation_mnk = append<3>(permutations, _);

  return TiledMMA<MMA_Atom<MMA_Op>,
                  decltype(thr_layout_mnk),
                  decltype(permutation_mnk)>{mma_atom, thr_layout_mnk};
}

You're getting the following as a result

TiledMMA<AtomType,
         Layout<Shape<_1,_1,_1>>,       // AtomLayoutMNK
         Tile<Layout<Shape<_2,_2,_2>>,  // Permutation on M, (2,2,2):(1,2,4), identity of size 8
              Underscore,               // Permutation on N, noop identity
              Underscore>>              // Permutation on K, noop identity

So you can see that your Layout is becoming a rank-3 Tiler, but the permutations on each mode are all identities anyway and have little to no effect.

Hope that helps! I'll try to get explanations and examples like these into the documentation as soon as possible.

Nice explanation! But I still have problem about this:

Layout<Shape <_2,_4,_2>, 
              Stride<_1,_4,_2>>

it seems it takes element idx is: 0 1 4 5 8 9 12 13 2 3 6 7 10 11 14 15

But in your case it need to take in n axis like: 0 1 8 9 2 3 10 11 4 5 12 13 6 7 14 15？。。。

[ccecka]

The layout you write is

Layout<Shape <_2,_2,_4>,
       Stride<_1,_8,_2>>
// 0 1 8 9 2 3 10 11 4 5 12 13 6 7 14 15

which I would describe as the "gather permutation" for the above, mapping the n-coords of the new image to their n-coords of the old image.

The layout I wrote is

Layout<Shape <_2,_4,_2>, 
       Stride<_1,_4,_2>>
// 0 1 4 5 8 9 12 13 2 3 6 7 10 11 14 15

which I would describe as the "scatter permutation" for the above, mapping the n-coords of the old image to their n-coords of the new image.

As you might expect from the descriptions, these are inverses of one another:

auto A = Layout<Shape<_2,_4,_2>,Stride<_1,_4,_2>>{};
auto B = Layout<Shape<_2,_2,_4>,Stride<_1,_8,_2>>{};
CUTE_STATIC_ASSERT_V(A == right_inverse(B));
CUTE_STATIC_ASSERT_V(B == right_inverse(A));

The TiledMMA takes the "scatter" version of these two.

[MARD1NO]

Oh! I understand it, so much thanks! :D

[wanghongyu2001]

Hi, AtomLayoutMNK mean what in this code

// @tparam MMA_Atom The MMA_Atom to use in the TiledMMA
// @tparam AtomLayoutMNK The MNK-tiling of the Atom to be performed.
// @tparam PermuationsMNK Permutations to apply to each MNK-mode before tiling for the Atom.
template <class MMA_Atom,
          class AtomLayoutMNK,
          class PermutationMNK = Tile<Underscore,Underscore,Underscore>>
struct TiledMMA : MMA_Atom { ... }

I didn't understand this sentence “AtomLayoutMNK means that the MMA_Atom is replicated into more threads, and so the CTA becomes accordingly. ”

[Ddd195]

How does this upper-right matrix represent its Layout in N dimension as Layout<Shape <_2,_4,_2>,Stride<_1,_4,_2>>

Thank you very much for your reply！
@ccecka

[thakkarV]

Have you read the mma atom readme in cute docs? It explains how these are transcribed

[Ddd195]

Thank you for reminding me! I found the answer I was looking for.
But it seems to use "new matrix" according to Layout<Shape <_2,_4,_2>, Stride<_1,_4,_2>> to get "old matrix", but normally, why not design the make_tiled_mma parameter to be an arrangement of the "old" to "new" matrix.
For instance, currently, given that I desire the "new" matrix where all the values of the T0 thread are arranged together, I am obliged to determine the permutation from the "old" to the "new" matrix by myself, namely 0 1 8 9 2 3 10 11 4 5 12 13 6 7 14 15. Then, via the inverse operation, obtain the 0, 1, 4, 5, 8, 9, 13, 2, 3, 6, 7, 10, 11, 14 15 layout, and subsequently fill this layout into make_tiled_mma.
Thanks for your reply.

[jcao-ai]

Is it right to think AtomLayoutMNK as intra-warp tiling, and PermutationMNK as inside-warp tiling ? Further more, permutation is intrinsic benefit brought by it.

[thakkarV]

That's a very Ampere centric view. Atom layout is just a tiling of atoms over "threads". In hopper the atom spans 128 threads.

[jcao-ai]

thanks, you are right.