[QST] What are other arguments in TiledMMA and TiledCopy, other than their thread layouts?

[hyhieu]

I am trying to learn about TiledMMA and TiledCopy, and have some questions, hoping to boost my understanding.

I (sort of) grasped the concepts of having MMA_Atoms and Copy_Atoms operating on a spatial "atomic" set of threads, and then tile those sets into TiledMMA and TiledCopy using thread layouts.

What I don't understand is how are these concepts map to code. Can someone please help answering my questions below? Thanks a ton in advance!!!

1. For TiledMMA:

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

It looks like the MMA_Atom will be tiled into AtomLayoutMNK.

But what do ValLayoutMNK and PermutationMNK do?

What are some realistic examples where we should use them?

2. Similarly, for TiledCopy:

template <class Copy_Atom,
          class LayoutCopy_TV,  // (tid,vid) -> coord   [Need not be 2D...]
          class ShapeTile_MN>   // coord space
struct TiledCopy : Copy_Atom

Is LayoutCopy_TV equivalent to AtomLayoutMNK in TiledMMA?

What does ShapeTiled_MN do? What are some realistic examples where we we should use it?

[ccecka]

These interfaces are under development, so don't get too married to them. But the concept remains consistent: We want some representation of partitioning patterns that can be applied to any data layouts (Tensors). In the case of MMA, there are three (projections) of these patterns: A, B, and C. In the case of Copy, there are two of these patterns: S (source), and D (destination). In both cases, the Atom may require transformations between threads and values, so these partitioning patterns must also account for the Atom.

I find that the best way to learn these objects is to experiment with them. It appears that we have two unreleased scripts that would help. I'll copy them below and try to include them in the next release.

/* mma_tiled_tex.cpp
 * $$ make mma_tiled_tex && ./mma_tiled_tex > mma_tiled.tex && pdflatex mma_tiled.pdf
 */

#include <cute/tensor.hpp>

using namespace cute;

int main()
{
#if 0
  {
    auto tiled_mma = make_tiled_mma(UniversalFMA<double>{});

    print_latex(tiled_mma);
  }
#endif

#if 0
  {
    auto tiled_mma = make_tiled_mma(SM61_DP4A{},
                                    Layout<Shape<_8,_8>>{});  // 8x8 layout of atoms (threads)

    print_latex(tiled_mma);
  }
#endif

#if 0
  {
    auto tiled_mma = make_tiled_mma(SM61_DP4A{},
                                    composition(Swizzle<3,0,-3>{}, Layout<Shape<_8,_8>>{}));

    print_latex(tiled_mma);
  }
#endif

#if 1
  {
    auto tiled_mma = make_tiled_mma(SM70_8x8x4_F32F16F16F32_NT{});

    print_latex(tiled_mma);
  }
#endif

#if 0
  {
    auto tiled_mma = make_tiled_mma(SM70_8x8x4_F32F16F16F32_NT{},
                                    Layout<Shape<_2,_2>>{},      // 2x2 layout of atoms (threads)
                                    Layout<Shape<_1,_1>>{},      // 1x1 layout of atoms (values)
                                    Tile<Layout<Shape <_4,_2>,   // Permutation in M
                                                Stride<_1,_8>>>{});

    print_latex(tiled_mma);
  }
#endif

#if 0
  {
    auto tiled_mma = make_tiled_mma(SM70_8x8x4_F32F16F16F32_NT{},
                                    composition(Swizzle<1,0,-1>{},
                                                Layout<Shape <Shape <_2,_2>,Shape <_2,_2>>,
                                                       Stride<Stride<_1,_4>,Stride<_2,_8>>>{}),  // 4x4 layout of atoms (threads)
                                    Layout<Shape<_2,_2>>{},                                      // 2x2 layout of atoms (values)
                                    Tile<Layout<Shape<_4,_2,_2>,Stride<_1,_32,_8>>,              // Permutation in M
                                         Layout<Shape<_4,_2,_2>,Stride<_1,_32,_8>>>{});          // Permutation in N

    print_latex(tiled_mma);
  }
#endif

#if 0
  {
    auto tiled_mma = make_tiled_mma(SM70_8x8x4_F32F16F16F32_NT{},
                                    Layout<Shape<_2,_2>>{},      // 2x2 tile of atoms (threads)
                                    Layout<Shape<_2,_2>>{},      // 2x2 tile of atoms (data)
                                    Tile<Layout<Shape <_4,_2>,
                                                Stride<_1,_16>>, // Permutation of M
                                         Layout<_1,
                                                _1>>{});         // Permutation of N

    print_latex(tiled_mma);
  }
#endif

#if 0
  {
    auto tiled_mma = make_tiled_mma(SM75_16x8x8_F32F16F16F32_TN{});

    print_latex(tiled_mma);
  }
#endif

#if 0
  {
    auto tiled_mma = make_tiled_mma(SM80_16x8x16_F32F16F16F32_TN{}, // HMMA.16816 warp-wide MmaAtom. 32 threads, 1 warp. Each thread owns 4 output values.
                                    Layout<Shape<_2,_2>>{},   // Tile in Threads/Warp
                                    Layout<Shape<_1,_1>>{});  // Tile in Output Val per thread.

    print_latex(tiled_mma);
  }
#endif

#if 0
  {
    auto tiled_mma = make_tiled_mma(SM80_8x8x4_F64F64F64F64_TN{},
                                    Layout<Shape<_2,_2>>{},  // Layout in Thr
                                    Layout<Shape<_2,_2>>{},  // Layout in Val
                                    Tile<Layout<_2,          // Permutation on M
                                                _16>,
                                         Layout<_2,          // Permutation on N
                                                _16>>{});

    print_latex(tiled_mma);
  }
#endif

#if 0
  {
    auto tiled_mma = make_tiled_mma(SM80_16x8x32_S32S8S8S32_TN{});

    print_latex(tiled_mma);
  }
#endif

#if 0
  {
    auto tiled_mma = make_tiled_mma(SM90_64x32x16_F16F16F16_SS<GMMA::Major::MN, GMMA::Major::MN>{});

    print_latex(tiled_mma);
  }
#endif
}

/* copy_tiled_tex.cpp
 * $$ make copy_tiled_tex && ./copy_tiled_tex > copy_tiled.tex && pdflatex copy_tiled.pdf
 */

#include <cute/tensor.hpp>

using namespace cute;

int main()
{
#if 0
  {
    Copy_Atom<UniversalCopy<double>, double> copy_atom;

    auto tiled_copy = make_tiled_copy(copy_atom,
                                      Layout<Shape<_32,_1>>{},  // 32x1 threads
                                      Layout<Shape< _1,_4>>{}); //  1x4 values

    print_latex(tiled_copy);
  }
#endif

#if 1
  {
    // The canonical LDSM_N image
    Copy_Atom<SM75_U32x1_LDSM_N, uint32_t> copy_atom;

    auto tiled_copy = make_tiled_copy(copy_atom,
                                      Layout<Shape<_8,_4>, Stride<_4,_1>>{}); // 8x4 RowMajor threads

    print_latex(tiled_copy);
  }
#endif

#if 0
  {
    // The canonical LDSM_T image
    Copy_Atom<SM75_U16x2_LDSM_T, uint16_t> copy_atom;

    auto tiled_copy = make_tiled_copy(copy_atom,
                                      Layout<Shape<_8,_4>, Stride<_4,_1>>{},  // 8x4 RowMajor threads
                                      Layout<Shape<_1,_2>>{});                // 1x2 values per thread

    print_latex(tiled_copy);
  }
#endif

#if 0
  {
    // Generate a TiledCopy layout from a TiledMMA
    SM80_16x8x16_F32F16F16F32_TN mma;
    auto tiled_mma = make_tiled_mma(mma);   // HMMA.16816 warp-wide MmaAtom. 32 threads, 1 warp. Each thread owns 4 output values.

    // Pick an Atom
    //Copy_Atom<DefaultCopy, uint16_t> copy_atom;
    //Copy_Atom<SM75_U16x2_LDSM_T, uint16_t> copy_atom;
    //Copy_Atom<SM75_U16x4_LDSM_T, uint16_t> copy_atom;
    Copy_Atom<SM75_U16x8_LDSM_T, uint16_t> copy_atom;

    // Define the layout of threads and values from the tiled_mma
    auto tiled_copy = make_tiled_copy_A(copy_atom, tiled_mma);

    print_latex(tiled_copy);
  }
#endif

#if 0
  {
    Copy_Atom<SM75_U16x2_LDSM_T, uint16_t> copy_atom;

    auto tiled_copy = make_tiled_copy(copy_atom,
                                      Layout<Shape<_4,_8>, Stride<_1,_4>>{},
                                      Layout<Shape<_2,_1>>{});

    print_latex(tiled_copy);
  }
#endif

#if 0
  {
    Copy_Atom<SM75_U16x2_LDSM_N, uint16_t> copy_atom;

    auto tiled_copy = make_tiled_copy(copy_atom,
                                      Layout<Shape<_32,_1>>{},
                                      Layout<Shape< _1,_8>>{});

    print_latex(tiled_copy);
  }
#endif
}

[hyhieu]

Thank you for your reply @ccecka. Playing around with the objects and printing out what the are helped me a lot.

I also found it helpful to print multiple TiledCopy and TiledMMA layouts, each with a small change in their arguments, to see the differences.

I modified the script you shared to do so. Here's the modified script. I hope it helps more people like it did to me.

/* tiled_tex.cu

Usage:

nvcc tiled_tex.cu -w \
  -Xcompiler "-Wfatal-errors" \
  -lineinfo \
  -std=c++17 \
  -I/usr/local/cuda/include \
  -I/path/to/cutlass/include \
  -I/path/to/cutlass/tools/util/include

./a.out
*/


#include "cute/tensor.hpp"
#include "cutlass/numeric_types.h"


void print_header() {
  const char* latex_header =
    "\\documentclass{article}\n"
    "\\usepackage[a4paper, margin=0.5cm]{geometry}\n"
    "\\usepackage{adjustbox}\n"
    "\\usepackage{graphicx}\n"
    "\\usepackage{lipsum}\n"
    "\\usepackage{tikz}\n"
    "\n"
    "\\begin{document}\n";
  printf("%s", latex_header);
}


void print_footer() {
  const char* latex_footer = "\\end{document}\n";
  printf("%s", latex_footer);
}


// Copy from mma_atom.hpp
//
// Modified to remove printing header and footder, hence allows printing
// multiple MMAs per TEX file for easier comparisons.
template <class AtomLayoutMNK,
          class ValLayoutMNK,
          class PermutationsMNK,
          class LayoutC, class ThrIDC,
          class LayoutA, class ThrIDA,
          class LayoutB, class ThrIDB>
void
print_mma(const char* name,
          const AtomLayoutMNK& atom_layout_mnk,
          const ValLayoutMNK& val_layout_mnk,
          const PermutationsMNK& permutations_mnk,
          LayoutC const& C, ThrIDC const& TC,    // (m,n) -> (tid,vid)  and  tid -> thr_idx
          LayoutA const& A, ThrIDA const& TA,    // (m,k) -> (tid,vid)  and  tid -> thr_idx
          LayoutB const& B, ThrIDB const& TB) {  // (n,k) -> (tid,vid)  and  tid -> thr_idx
  using namespace cute;

  printf("\\begin{verbatim}\n");
  printf("\n%s\n\n", name);

  printf("  AtomLayoutMNK: "); print(atom_layout_mnk);   printf("\n");
  printf("   ValLayoutMNK: "); print(val_layout_mnk);    printf("\n");
  printf("PermutationsMNK: "); print(permutations_mnk); printf("\n\n");

  printf("LayoutC: "); print(C);  printf("\n");
  printf(" ThrIDC: "); print(TC); printf("\n");
  printf("LayoutA: "); print(A);  printf("\n");
  printf(" ThrIDA: "); print(TA); printf("\n");
  printf("LayoutB: "); print(B);  printf("\n");
  printf(" ThrIDB: "); print(TB); printf("\n");
  printf("\\end{verbatim}\n");

  printf("\\begin{adjustbox}{max height=0.7\\textheight,max width=\\textwidth}%");
  printf("\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},box/"
         ".style={rectangle,draw=black,thick,minimum size=1cm,anchor=center}]\n\n");
  char const* color_map[8] = {"{rgb,255:red,175;green,175;blue,255}",
                              "{rgb,255:red,175;green,255;blue,175}",
                              "{rgb,255:red,255;green,255;blue,175}",
                              "{rgb,255:red,255;green,175;blue,175}",
                              "{rgb,255:red,210;green,210;blue,255}",
                              "{rgb,255:red,210;green,255;blue,210}",
                              "{rgb,255:red,255;green,255;blue,210}",
                              "{rgb,255:red,255;green,210;blue,210}"};

  // C starting at 0,0
  for (int m = 0; m < size<0>(C); ++m) {
    for (int n = 0; n < size<1>(C); ++n) {
      int thrid   = C(m,n) % size(TC);
      int val_idx = C(m,n) / size(TC);
      int thr_idx = TC(thrid);

      printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
             color_map[thr_idx % 8],
             m, n,
             thr_idx, val_idx);
    }
  }

  // A starting at 0,-size<1>(A)-1
  for (int m = 0; m < size<0>(A); ++m) {
    for (int k = 0; k < size<1>(A); ++k) {
      int thrid   = A(m,k) % size(TA);
      int val_idx = A(m,k) / size(TA);
      int thr_idx = TA(thrid);

      printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
             color_map[thr_idx % 8],
             m, k-1-size<1>(A),
             thr_idx, val_idx);
    }
  }

  // B starting at -size<1>(B)-1,0
  for (int n = 0; n < size<0>(B); ++n) {
    for (int k = 0; k < size<1>(B); ++k) {
      int thrid   = B(n,k) % size(TB);
      int val_idx = B(n,k) / size(TB);
      int thr_idx = TB(thrid);

      printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
             color_map[thr_idx % 8],
             k-1-size<1>(B), n,
             thr_idx, val_idx);
    }
  }

  // A labels
  for (int m = 0, k = -1; m < size<0>(A); ++m) {
    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", m, k-1-size<1>(A), m);
  }
  for (int k = 0, m = -1; k < size<1>(A); ++k) {
    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", m, k-1-size<1>(A), k);
  }
  // B labels
  for (int n = 0, k = -1; n < size<0>(B); ++n) {
    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", k-1-size<1>(B), n, n);
  }
  for (int k = 0, n = -1; k < size<1>(B); ++k) {
    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", k-1-size<1>(B), n, k);
  }

  printf("\\end{tikzpicture}\n\\end{adjustbox}%\n");
}


template <class TiledCopy,
          class LayoutS, class ThrIDS,
          class LayoutD, class ThrIDD>
void
print_copy(const char* name,
           TiledCopy& copy,
           LayoutS const& S, ThrIDS const& TS,   // (m,n) -> (tid,vid)  and  tid -> thr_idx
           LayoutD const& D, ThrIDD const& TD) { // (m,n) -> (tid,vid)  and  tid -> thr_idx
  using namespace cute;

  CUTE_STATIC_ASSERT_V(rank(S) == Int<2>{});
  CUTE_STATIC_ASSERT_V(rank(D) == Int<2>{});

  assert(size<0>(S) == size<0>(D));
  assert(size<1>(S) == size<1>(D));

  printf("\\begin{verbatim}\n");
  printf("\n%s\n\n", name);
  printf("LayoutCopy_TV: "); print(typename TiledCopy::TiledLayout_TV{});  printf("\n");
  printf(" ShapeTile_MN: "); print(typename TiledCopy::Tiler_MN{});      printf("\n\n");

  printf("      LayoutS: "); print(S);  printf("\n");
  printf("       ThrIDS: "); print(TS); printf("\n");
  printf("      LayoutD: "); print(D);  printf("\n");
  printf("       ThrIDD: "); print(TD); printf("\n");
  printf("\\end{verbatim}\n");

  printf("\\begin{adjustbox}{max height=0.7\\textheight,max width=\\textwidth}%");
  printf("\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},box/"
         ".style={rectangle,draw=black,thick,minimum size=1cm,anchor=center}]\n\n");
  char const* color_map[8] = {"{rgb,255:red,175;green,175;blue,255}",
                              "{rgb,255:red,175;green,255;blue,175}",
                              "{rgb,255:red,255;green,255;blue,175}",
                              "{rgb,255:red,255;green,175;blue,175}",
                              "{rgb,255:red,210;green,210;blue,255}",
                              "{rgb,255:red,210;green,255;blue,210}",
                              "{rgb,255:red,255;green,255;blue,210}",
                              "{rgb,255:red,255;green,210;blue,210}"};

  // S starting at 0,0
  for (int i = 0; i < size<0>(S); ++i) {
    for (int j = 0; j < size<1>(S); ++j) {
      int thrid   = S(i,j) % size(TS);
      int val_idx = S(i,j) / size(TS);
      int thr_idx = TS(thrid);

      printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
             color_map[thr_idx % 8],
             i, j,
             thr_idx, val_idx);
    }
  }

  // D starting at 0,size<1>(S)+3
  for (int i = 0; i < size<0>(D); ++i) {
    for (int j = 0; j < size<1>(D); ++j) {
      int thrid   = D(i,j) % size(TD);
      int val_idx = D(i,j) / size(TD);
      int thr_idx = TD(thrid);

      printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
             color_map[thr_idx % 8],
             i + size<0>(S) + 3, j,
             thr_idx, val_idx);
    }
  }

  // S Labels
  for (int i = 0, j = -1; i < size<0>(S); ++i) {
    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, i);
  }
  for (int j = 0, i = -1; j < size<1>(S); ++j) {
    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, j);
  }
  // D Labels
  for (int i = 0, j = -1; i < size<0>(S); ++i) {
    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i + size<0>(S) + 3, j, i);
  }
  for (int j = 0, i = -1; j < size<1>(D); ++j) {
    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i + size<0>(S) + 3, j, j);
  }

  printf("\\end{tikzpicture}\n\\end{adjustbox}%\n");
}


template <class MMA_Atom_Arch,
          class AtomLayoutMNK,
          class ValLayoutMNK,
          class PermutationsMNK>
void print_mma_content(
  const char* name,
  cute::TiledMMA<MMA_Atom_Arch, AtomLayoutMNK, ValLayoutMNK, PermutationsMNK> const& mma) {
  printf("\n\\newpage\n");

  auto layout_and_thrid_C = mma.get_layoutC_MN();
  auto layoutC_MN = cute::get<0>(layout_and_thrid_C);
  auto thrID_C    = cute::get<1>(layout_and_thrid_C);

  auto layout_and_thrid_A = mma.get_layoutA_MK();
  auto layoutA_MK = cute::get<0>(layout_and_thrid_A);
  auto thrID_A    = cute::get<1>(layout_and_thrid_A);

  auto layout_and_thrid_B = mma.get_layoutB_NK();
  auto layoutB_NK = cute::get<0>(layout_and_thrid_B);
  auto thrID_B    = cute::get<1>(layout_and_thrid_B);

  print_mma(name,
            AtomLayoutMNK{},
            ValLayoutMNK{},
            PermutationsMNK{},
            layoutC_MN, thrID_C,
            layoutA_MK, thrID_A,
            layoutB_NK, thrID_B);
}


template <class TiledCopy>
void print_copy_content(const char* name, TiledCopy& copy) {
  printf("\n\\newpage\n");

  auto [layoutS_MN, thrID_S] = copy.get_layoutS_MN();
  auto [layoutD_MN, thrID_D] = copy.get_layoutD_MN();

  print_copy(name, copy,
             layoutS_MN, thrID_S,
             layoutD_MN, thrID_D);
}


void print_layouts_for_mma() {
  using namespace cute;
  using _X = cute::Underscore;

  {
    auto tiled_mma = make_tiled_mma(SM70_8x8x4_F32F16F16F32_NT{});
    print_mma_content("SM70_8x8x4_F32F16F16F32_NT", tiled_mma);
  }

  {
    auto tiled_mma = make_tiled_mma(SM70_8x8x4_F32F16F16F32_NT{},
                                    Layout<Shape<_1,_2>>{}  // AtomLayoutMNK
    );
    print_mma_content("SM70_8x8x4_F32F16F16F32_NT", tiled_mma);
  }

  {
    auto tiled_mma = make_tiled_mma(SM70_8x8x4_F32F16F16F32_NT{},
                                    Layout<Shape<_2,_1>>{}  // AtomLayoutMNK
    );
    print_mma_content("SM70_8x8x4_F32F16F16F32_NT", tiled_mma);
  }

  {
    auto tiled_mma = make_tiled_mma(SM70_8x8x4_F32F16F16F32_NT{},
                                    Layout<Shape<_2,_2>>{}  // AtomLayoutMNK
    );
    print_mma_content("SM70_8x8x4_F32F16F16F32_NT", tiled_mma);
  }

  {
    auto tiled_mma = make_tiled_mma(SM70_8x8x4_F32F16F16F32_NT{},
                                    Layout<Shape<_2,_2>>{},  // AtomLayoutMNK
                                    Layout<Shape<_2,_1>>{}   // ValLayoutMNK
    );
    print_mma_content("SM70_8x8x4_F32F16F16F32_NT", tiled_mma);
  }

  {
    auto tiled_mma = make_tiled_mma(SM70_8x8x4_F32F16F16F32_NT{},
                                    Layout<Shape<_2,_2>>{},  // AtomLayoutMNK
                                    Layout<Shape<_1,_2>>{}   // ValLayoutMNK
    );
    print_mma_content("SM70_8x8x4_F32F16F16F32_NT", tiled_mma);
  }

  {
    auto tiled_mma = make_tiled_mma(SM70_8x8x4_F32F16F16F32_NT{},
                                    Layout<Shape<_1,_1,_2>>{},  // AtomLayoutMNK
                                    Layout<Shape<_1,_1>>{}      // ValLayoutMNK
    );
    print_mma_content("SM70_8x8x4_F32F16F16F32_NT", tiled_mma);
  }

  {
    auto tiled_mma = make_tiled_mma(SM70_8x8x4_F32F16F16F32_NT{},
                                    Layout<Shape<_1,_1>>{},     // AtomLayoutMNK
                                    Layout<Shape<_1,_1,_2>>{}   // ValLayoutMNK
    );
    print_mma_content("SM70_8x8x4_F32F16F16F32_NT", tiled_mma);
  }

  // {
  //   auto tiled_mma = make_tiled_mma(SM70_8x8x4_F32F16F16F32_NT{},
  //                                   Layout<Shape<_2,_2>>{},        // 2x2 layout of atoms (threads)
  //                                   Layout<Shape<_1,_1>>{},        // 1x1 layout of atoms (values)
  //                                   Tile<Layout<Shape <_4,_2>,     // Permutation in M
  //                                               Stride<_1,_8>>>{});
  //   print_mma_content("SM70_8x8x4_F32F16F16F32_NT", tiled_mma);
  // }

#if 0
  {
    auto tiled_mma = make_tiled_mma(SM70_8x8x4_F32F16F16F32_NT{},
                                    Layout<Shape<_2,_2>>{},      // 2x2 layout of atoms (threads)
                                    Layout<Shape<_1,_1>>{},      // 1x1 layout of atoms (values)
                                    Tile<Layout<Shape <_4,_2>,   // Permutation in M
                                                Stride<_1,_8>>>{});

    print_latex(tiled_mma);
  }
#endif

#if 0
  {
    auto tiled_mma = make_tiled_mma(SM70_8x8x4_F32F16F16F32_NT{},
                                    composition(Swizzle<1,0,-1>{},
                                                Layout<Shape <Shape <_2,_2>,Shape <_2,_2>>,
                                                       Stride<Stride<_1,_4>,Stride<_2,_8>>>{}),  // 4x4 layout of atoms (threads)
                                    Layout<Shape<_2,_2>>{},                                      // 2x2 layout of atoms (values)
                                    Tile<Layout<Shape<_4,_2,_2>,Stride<_1,_32,_8>>,              // Permutation in M
                                         Layout<Shape<_4,_2,_2>,Stride<_1,_32,_8>>>{});          // Permutation in N

    print_latex(tiled_mma);
  }
#endif

#if 0
  {
    auto tiled_mma = make_tiled_mma(SM70_8x8x4_F32F16F16F32_NT{},
                                    Layout<Shape<_2,_2>>{},      // 2x2 tile of atoms (threads)
                                    Layout<Shape<_2,_2>>{},      // 2x2 tile of atoms (data)
                                    Tile<Layout<Shape <_4,_2>,
                                                Stride<_1,_16>>, // Permutation of M
                                         Layout<_1,
                                                _1>>{});         // Permutation of N

    print_latex(tiled_mma);
  }
#endif

#if 0
  {
    auto tiled_mma = make_tiled_mma(SM75_16x8x8_F32F16F16F32_TN{});

    print_latex(tiled_mma);
  }
#endif

#if 0
  {
    auto tiled_mma = make_tiled_mma(SM80_16x8x16_F32F16F16F32_TN{}, // HMMA.16816 warp-wide MmaAtom. 32 threads, 1 warp. Each thread owns 4 output values.
                                    Layout<Shape<_2,_2>>{},   // Tile in Threads/Warp
                                    Layout<Shape<_1,_1>>{});  // Tile in Output Val per thread.

    print_latex(tiled_mma);
  }
#endif

#if 0
  {
    auto tiled_mma = make_tiled_mma(SM80_8x8x4_F64F64F64F64_TN{},
                                    Layout<Shape<_2,_2>>{},  // Layout in Thr
                                    Layout<Shape<_2,_2>>{},  // Layout in Val
                                    Tile<Layout<_2,          // Permutation on M
                                                _16>,
                                         Layout<_2,          // Permutation on N
                                                _16>>{});

    print_latex(tiled_mma);
  }
#endif

#if 0
  {
    auto tiled_mma = make_tiled_mma(SM80_16x8x32_S32S8S8S32_TN{});

    print_latex(tiled_mma);
  }
#endif

#if 0
  {
    auto tiled_mma = make_tiled_mma(SM90_64x32x16_F16F16F16_SS<GMMA::Major::MN, GMMA::Major::MN>{});

    print_latex(tiled_mma);
  }
#endif
}


void print_layouts_for_sm80_cp_async_cachealways() {
  using namespace cute;

  {
    using ST = float;
    using DT = float;
    using Copy_Atom_Arch = Copy_Atom<Copy_Traits<SM80_CP_ASYNC_CACHEALWAYS<ST, DT>>, DT>;
    auto tiled_copy = TiledCopy<
      Copy_Atom_Arch,
      Layout<Shape<_16, _1>>,
      Layout<Shape<_4, _8>>
    >{};
    print_copy_content("SM80_CP_ASYNC_CACHEALWAYS", tiled_copy);
  }

  {
    using ST = float;
    using DT = float;
    using Copy_Atom_Arch = Copy_Atom<Copy_Traits<SM80_CP_ASYNC_CACHEALWAYS<ST, DT>>, DT>;
    auto tiled_copy = TiledCopy<
      Copy_Atom_Arch,
      Layout<Shape<_32, _1>>,
      Layout<Shape<_4, _8>>
    >{};
    print_copy_content("SM80_CP_ASYNC_CACHEALWAYS", tiled_copy);
  }

  {
    using ST = float;
    using DT = float;
    using Copy_Atom_Arch = Copy_Atom<Copy_Traits<SM80_CP_ASYNC_CACHEALWAYS<ST, DT>>, DT>;
    auto tiled_copy = TiledCopy<
      Copy_Atom_Arch,
      Layout<Shape<_2, _1>>,
      Layout<Shape<_8, _4>>
    >{};
    print_copy_content("SM80_CP_ASYNC_CACHEALWAYS", tiled_copy);
  }

  {
    using ST = float;
    using DT = float;
    using Copy_Atom_Arch = Copy_Atom<Copy_Traits<SM80_CP_ASYNC_CACHEALWAYS<ST, DT>>, DT>;
    auto tiled_copy = TiledCopy<
      Copy_Atom_Arch,
      Layout<Shape<_16, _1>>,
      Layout<Shape<_8, _8>>
    >{};
    print_copy_content("SM80_CP_ASYNC_CACHEALWAYS", tiled_copy);
  }

  {
    using ST = float;
    using DT = float;
    using Copy_Atom_Arch = Copy_Atom<Copy_Traits<SM80_CP_ASYNC_CACHEALWAYS<ST, DT>>, DT>;
    auto tiled_copy = TiledCopy<
      Copy_Atom_Arch,
      Layout<Shape<_32, _1>>,
      Layout<Shape<_8, _8>>
    >{};
    print_copy_content("SM80_CP_ASYNC_CACHEALWAYS", tiled_copy);
  }

  {
    using ST = float;
    using DT = float;
    using Copy_Atom_Arch = Copy_Atom<Copy_Traits<SM80_CP_ASYNC_CACHEALWAYS<ST, DT>>, DT>;
    auto tiled_copy = TiledCopy<
      Copy_Atom_Arch,
      Layout<Shape<_4, _4>>,
      Layout<Shape<_8, _8>>
    >{};
    print_copy_content("SM80_CP_ASYNC_CACHEALWAYS", tiled_copy);
  }

  {
    using ST = float;
    using DT = float;
    using Copy_Atom_Arch = Copy_Atom<Copy_Traits<SM80_CP_ASYNC_CACHEALWAYS<ST, DT>>, DT>;
    auto tiled_copy = TiledCopy<
      Copy_Atom_Arch,
      Layout<Shape<_4, _1>>,
      Layout<Shape<_8, _32>, Stride<_32, _32>>
    >{};
    print_copy_content("SM80_CP_ASYNC_CACHEALWAYS", tiled_copy);
  }
}


int main() {
  print_header();

  print_layouts_for_mma();
  // print_layouts_for_sm80_cp_async_cachealways();

  print_footer();
  return 0;
}

[Rya-Sanovar]

Thank you SO MUCH! This helped me a ton and will probably save me days of work just figuring out how copy and tiled layouts work :)))

[yangjianfengo1]

very thanks