CuTe Tiled MMA

Introduction

Matrix multiplication and accumulation (MMA) is a key operation for general matrix multiplication (GEMM). In CUTLASS, CuTe provides APIs for configuring MMA from MMA atoms to MMA tiles so that larger MMA problems can be solved.

In this blog post, I would like to discuss the CuTe tiled MMA configurations, layouts, and API usages, using an example.

CuTe Tiled MMA Preview Example

The following CuTe tiled MMA preview example does not actually perform any MMA computation, because it is completely a host program. Instead, it demonstrates how to configure the MMA atom, MMA tile, and MMA layout using CuTe APIs.

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260

#include <cassert>#include <fstream>#include <iomanip>#include <iostream>#include <cute/layout.hpp>#include <cute/swizzle.hpp>#include <cute/tensor.hpp>#include <thrust/host_vector.h>int main(int argc, const char** argv){    // Tiled MMA requires everything to be static.    // Therefore, this preview program does not allow user to configure    // dynamically. To preview a new tiled MMA configuration, the user has to    // modify this program and recompile.    // Configure data type.    using TA = cute::half_t;    using TB = cute::half_t;    using TC = cute::half_t;    // Configure static "shared memory".    // The "shared memory" is actually on host for preview purpose.    // For tiled mma, the shared memory layout has to be static.    constexpr int bM{128 * 2 / sizeof(TA)};    constexpr int bN{128 * 2 / sizeof(TB)};    constexpr int bK{32};    auto const blk_M = cute::Int<bM>{};    auto const blk_N = cute::Int<bN>{};    auto const blk_K = cute::Int<bK>{};    auto const smem_shape_A{cute::make_shape(blk_M, blk_K)};    auto const smem_shape_B{cute::make_shape(blk_N, blk_K)};    auto const smem_shape_C{cute::make_shape(blk_M, blk_N)};    auto const smem_stride_A{        cute::make_stride(cute::Int<1>{}, blk_M)}; // Column-major    auto const smem_stride_B{        cute::make_stride(cute::Int<1>{}, blk_N)}; // Column-major    auto const smem_stride_C{        cute::make_stride(cute::Int<1>{}, blk_M)}; // Column-major    auto const smem_layout_A{        cute::make_layout(smem_shape_A, smem_stride_A)}; // (blk_M, blk_K)    auto const smem_layout_B{        cute::make_layout(smem_shape_B, smem_stride_B)}; // (blk_N, blk_K)    auto const smem_layout_C{        cute::make_layout(smem_shape_C, smem_stride_C)}; // (blk_M, blk_N)    auto const size_a{blk_M * blk_K};    auto const size_b{blk_N * blk_K};    auto const size_c{blk_M * blk_N};    auto h_A = thrust::host_vector<TA>(size_a);    auto h_B = thrust::host_vector<TB>(size_b);    auto h_C = thrust::host_vector<TC>(size_c);    // Make tensor for smem_A and smem_B.    auto smem_tensor_A{cute::make_tensor(h_A.data(), smem_layout_A)};    auto smem_tensor_B{cute::make_tensor(h_B.data(), smem_layout_B)};    auto smem_tensor_C{cute::make_tensor(h_C.data(), smem_layout_C)};    std::cout << "smem_tensor_A" << std::endl;    cute::print(smem_tensor_A);    std::cout << std::endl;    std::cout << "smem_tensor_B" << std::endl;    cute::print(smem_tensor_B);    std::cout << std::endl;    std::cout << "smem_tensor_C" << std::endl;    cute::print(smem_tensor_C);    std::cout << std::endl;    // Configure tiled MMA.    using MmaTraits = cute::MMA_Traits<cute::SM80_16x8x16_F16F16F16F16_TN>;    using MmaAtomShape = MmaTraits::Shape_MNK;    auto const mma_atom = cute::MMA_Atom<MmaTraits>{};    auto const mma_atom_shape = MmaAtomShape{};    // Repeating the mma atom along the M, N, and K dimensions.    // This increases the number of threads to process the tiled MMA.    constexpr int MMA_LAYOUT_M{2};    constexpr int MMA_LAYOUT_N{2};    constexpr int MMA_LAYOUT_K{1};    auto mma_layout{cute::make_layout(        cute::make_shape(cute::Int<MMA_LAYOUT_M>{}, cute::Int<MMA_LAYOUT_N>{},                         cute::Int<MMA_LAYOUT_K>{}))};    // Repeating the mma processing along the M, N, and K dimensions.    // This does not increase the number of threads to process the tiled MMA.    // But the number of registers required for processing the tiled MMA    // increases.    constexpr int NUM_MMA_TILE_M{1};    constexpr int NUM_MMA_TILE_N{2};    constexpr int NUM_MMA_TILE_K{1};    constexpr int MMA_TILE_M{cute::get<0>(mma_atom_shape) * MMA_LAYOUT_M *                             NUM_MMA_TILE_M};    constexpr int MMA_TILE_N{cute::get<1>(mma_atom_shape) * MMA_LAYOUT_N *                             NUM_MMA_TILE_N};    constexpr int MMA_TILE_K{cute::get<2>(mma_atom_shape) * MMA_LAYOUT_K *                             NUM_MMA_TILE_K};    auto mma_tile{cute::make_tile(cute::Int<MMA_TILE_M>{},                                  cute::Int<MMA_TILE_N>{},                                  cute::Int<MMA_TILE_K>{})};    auto tiled_mma{cute::make_tiled_mma(mma_atom, mma_layout, mma_tile)};    constexpr auto NUM_THREADS{cute::size(tiled_mma)};    CUTE_STATIC_ASSERT(NUM_THREADS ==                       MMA_LAYOUT_M * MMA_LAYOUT_N * MMA_LAYOUT_K *                           cute::size(decltype(mma_atom)::ThrID{}));    std::cout << "mma_atom" << std::endl;    cute::print(mma_atom);    std::cout << std::endl;    std::cout << "tiled_mma" << std::endl;    cute::print(tiled_mma);    std::cout << std::endl;    // ThrLayoutVMNK static asserts.    CUTE_STATIC_ASSERT_V(cute::shape<0>(decltype(tiled_mma)::ThrLayoutVMNK{}) ==                         cute::shape(decltype(mma_atom)::ThrID{}));    CUTE_STATIC_ASSERT_V(cute::shape<1>(decltype(tiled_mma)::ThrLayoutVMNK{}) ==                         cute::Int<MMA_LAYOUT_M>{});    CUTE_STATIC_ASSERT_V(cute::shape<2>(decltype(tiled_mma)::ThrLayoutVMNK{}) ==                         cute::Int<MMA_LAYOUT_N>{});    CUTE_STATIC_ASSERT_V(cute::shape<3>(decltype(tiled_mma)::ThrLayoutVMNK{}) ==                         cute::Int<MMA_LAYOUT_K>{});    // PermutationMNK static asserts.    CUTE_STATIC_ASSERT_V(tiled_mma.tile_size_mnk<0>() ==                         cute::Int<MMA_TILE_M>{});    CUTE_STATIC_ASSERT_V(tiled_mma.tile_size_mnk<1>() ==                         cute::Int<MMA_TILE_N>{});    CUTE_STATIC_ASSERT_V(tiled_mma.tile_size_mnk<2>() ==                         cute::Int<MMA_TILE_K>{});    // Partition via MMA.    // set an arbitrary thread index.    constexpr int THREAD_IDX{0};    CUTE_STATIC_ASSERT(THREAD_IDX < NUM_THREADS);    CUTE_STATIC_ASSERT(THREAD_IDX >= 0);    auto thread_mma{tiled_mma.get_slice(THREAD_IDX)};    // Register tensors used for MMA.    auto thread_layout_C_register_tensor_A{        thread_mma.partition_fragment_A(smem_tensor_A)}; // (MMA, MMA_M, MMA_K)    auto thread_layout_C_register_tensor_B{        thread_mma.partition_fragment_B(smem_tensor_B)}; // (MMA, MMA_N, MMA_K)    auto thread_layout_C_register_tensor_C{        thread_mma.partition_fragment_C(smem_tensor_C)}; // (MMA, MMA_M, MMA_N)    CUTE_STATIC_ASSERT_V(        cute::shape<1>(decltype(mma_atom)::LayoutA_TV{}) ==        cute::shape<0>(thread_layout_C_register_tensor_A)); // MMA_A    CUTE_STATIC_ASSERT_V(        cute::shape<1>(decltype(mma_atom)::LayoutB_TV{}) ==        cute::shape<0>(thread_layout_C_register_tensor_B)); // MMA_B    // Use no tiled copy from shared memory to register.    auto thread_layout_C_smem_tensor_A_no_tiled_copy{        thread_mma.partition_A(smem_tensor_A)}; // (MMA, MMA_M, MMA_K)    auto thread_layout_C_smem_tensor_B_no_tiled_copy{        thread_mma.partition_B(smem_tensor_B)}; // (MMA, MMA_N, MMA_K)    auto thread_layout_C_smem_tensor_C_no_tiled_copy{        thread_mma.partition_C(smem_tensor_C)}; // (MMA, MMA_M, MMA_N)    // thread_layout_C_smem_tensor_A_no_tiled_copy and    // thread_layout_C_register_tensor_A shall have the same shape.    CUTE_STATIC_ASSERT_V(        cute::shape(thread_layout_C_smem_tensor_A_no_tiled_copy) ==        cute::shape(thread_layout_C_register_tensor_A));    std::cout << "thread_layout_C_register_tensor_A" << std::endl;    cute::print(thread_layout_C_register_tensor_A);    std::cout << std::endl;    std::cout << "thread_layout_C_register_tensor_B" << std::endl;    cute::print(thread_layout_C_register_tensor_B);    std::cout << std::endl;    std::cout << "thread_layout_C_register_tensor_C" << std::endl;    cute::print(thread_layout_C_register_tensor_C);    std::cout << std::endl;    std::cout << "thread_layout_C_smem_tensor_A_no_tiled_copy" << std::endl;    cute::print(thread_layout_C_smem_tensor_A_no_tiled_copy);    std::cout << std::endl;    std::cout << "thread_layout_C_smem_tensor_B_no_tiled_copy" << std::endl;    cute::print(thread_layout_C_smem_tensor_B_no_tiled_copy);    std::cout << std::endl;    std::cout << "thread_layout_C_smem_tensor_C_no_tiled_copy" << std::endl;    cute::print(thread_layout_C_smem_tensor_C_no_tiled_copy);    std::cout << std::endl;    // Use tiled copy from shared memory to register.    auto copy_atom_A = cute::Copy_Atom<cute::SM75_U16x8_LDSM_T, TA>{};    auto copy_atom_B = cute::Copy_Atom<cute::SM75_U16x8_LDSM_T, TB>{};    auto smem_tiled_copy_A{cute::make_tiled_copy_A(copy_atom_A, tiled_mma)};    auto smem_tiled_copy_B{cute::make_tiled_copy_B(copy_atom_B, tiled_mma)};    CUTE_STATIC_ASSERT_V(        cute::shape<0>(decltype(smem_tiled_copy_A)::Tiler_MN{}) ==        tiled_mma.tile_size_mnk<0>()); // MMA_TILE_M    CUTE_STATIC_ASSERT_V(        cute::shape<1>(decltype(smem_tiled_copy_A)::Tiler_MN{}) ==        tiled_mma.tile_size_mnk<2>()); // MMA_TILE_K    CUTE_STATIC_ASSERT_V(        cute::shape<0>(decltype(smem_tiled_copy_B)::Tiler_MN{}) ==        tiled_mma.tile_size_mnk<1>()); // MMA_TILE_N    CUTE_STATIC_ASSERT_V(        cute::shape<1>(decltype(smem_tiled_copy_B)::Tiler_MN{}) ==        tiled_mma.tile_size_mnk<2>()); // MMA_TILE_K    auto smem_thread_copy_A{smem_tiled_copy_A.get_slice(THREAD_IDX)};    auto smem_thread_copy_B{smem_tiled_copy_B.get_slice(THREAD_IDX)};    auto thread_layout_C_smem_tensor_A_tiled_copy{        smem_thread_copy_A.partition_S(smem_tensor_A)};    auto thread_layout_C_smem_tensor_B_tiled_copy{        smem_thread_copy_B.partition_S(smem_tensor_B)};    auto thread_layout_C_register_tensor_A_copy_view{        smem_thread_copy_A.retile_D(thread_layout_C_register_tensor_A)};    auto thread_layout_C_register_tensor_B_copy_view{        smem_thread_copy_B.retile_D(thread_layout_C_register_tensor_B)};    CUTE_STATIC_ASSERT_V(        cute::shape(thread_layout_C_smem_tensor_A_tiled_copy) ==        cute::shape(thread_layout_C_register_tensor_A_copy_view));    CUTE_STATIC_ASSERT_V(        cute::shape(thread_layout_C_smem_tensor_B_tiled_copy) ==        cute::shape(thread_layout_C_register_tensor_B_copy_view));    std::cout << "copy_atom_A" << std::endl;    cute::print(copy_atom_A);    std::cout << std::endl;    std::cout << "copy_atom_B" << std::endl;    cute::print(copy_atom_B);    std::cout << std::endl;    std::cout << "smem_tiled_copy_A" << std::endl;    cute::print(smem_tiled_copy_A);    std::cout << std::endl;    std::cout << "smem_tiled_copy_B" << std::endl;    cute::print(smem_tiled_copy_B);    std::cout << std::endl;    std::cout << "thread_layout_C_smem_tensor_A_tiled_copy" << std::endl;    cute::print(thread_layout_C_smem_tensor_A_tiled_copy);    std::cout << std::endl;    std::cout << "thread_layout_C_smem_tensor_B_tiled_copy" << std::endl;    cute::print(thread_layout_C_smem_tensor_B_tiled_copy);    std::cout << std::endl;    std::cout << "thread_layout_C_register_tensor_A_copy_view" << std::endl;    cute::print(thread_layout_C_register_tensor_A_copy_view);    std::cout << std::endl;    std::cout << "thread_layout_C_register_tensor_B_copy_view" << std::endl;    cute::print(thread_layout_C_register_tensor_B_copy_view);    std::cout << std::endl;    return 0;}

A high-performance example that uses almost the same tiled MMA configurations to perform the GEMM computation can be found on my GitHub.

CuTe Tiled MMA Configurations and Layouts

MMA Problem Size and Shared Memory Configuration

In one thread block, per one main loop iteration, the MMA problem size is $M \times N \times K = 128 \times 128 \times 32$. The static shared memory is used to store the $M \times K = 128 \times 32$ sub-matrix of matrix $A$ in a column-major layout and the $K \times N = 32 \times 128$ sub-matrix of matrix $B$ in a row-major layout. Using the convention of MMA, we typically describe the sub-matrix of matrix $B$ as $N \times K = 128 \times 32$ column-major.

1234

smem_tensor_Aptr[16b](0x57b7b93248c0) o (_128,_32):(_1,_128)smem_tensor_Bptr[16b](0x57b7b93268d0) o (_128,_32):(_1,_128)

The shared memory configuration has to be compatible with the tiled MMA we configured later. Otherwise, the tiled MMA will not be able to process the MMA problem correctly because the cute::gemm API takes no predicates (for performance reasons).

MMA Atom Configuration

The MMA atom processes an MMA problem of size $M^{\prime} \times N^{\prime} \times K^{\prime}$ using a certain number of threads.

In our case, the MMA atom cute::SM80_16x8x16_F16F16F16F16_TN processes an MMA problem of size $M^{\prime} \times N^{\prime} \times K^{\prime} = 16 \times 8 \times 16$, and this MMA atom consists of a warp of 32 threads. The MMA atom is responsible for processing the $M^{\prime} \times K^{\prime} = 16 \times 16$ sub-matrix of matrix $A$ and the $K^{\prime} \times N^{\prime} = 16 \times 8$ sub-matrix of matrix $B$.

1234567

mma_atomMMA_Atom  ThrID:      _32:_1  Shape_MNK:  (_16,_8,_16)  LayoutA_TV: ((_4,_8),(_2,_2,_2)):((_32,_1),(_16,_8,_128))  LayoutB_TV: ((_4,_8),(_2,_2)):((_16,_1),(_8,_64))  LayoutC_TV: ((_4,_8),(_2,_2)):((_32,_1),(_16,_8))

To process the MMA problem of size $M \times N \times K = 128 \times 128 \times 32$, theoretically we could tile the MMA atoms in one of the following configurations:

Note that we did not configure parallelism in the $K$ dimension in the above configurations. But theoretically it’s also possible to do it, especially when $\frac{K}{K^{\prime}}$ is large.

The MMA atom also defines the layouts of the MMA matrices it works on. The thread-value layouts of the MMA matrices are usually very complicated. But fortunately, we could usually visualize them using the CuTe cute::print_latex or cute::print_svg functions.

In our case, the MMA atom cute::SM80_16x8x16_F16F16F16F16_TN defines the thread-value layouts of the MMA matrices as follows:

123

LayoutA_TV: ((_4,_8),(_2,_2,_2)):((_32,_1),(_16,_8,_128))LayoutB_TV: ((_4,_8),(_2,_2)):((_16,_1),(_8,_64))LayoutC_TV: ((_4,_8),(_2,_2)):((_32,_1),(_16,_8))

Apparently, each thread in the MMA atom in one process will access $2 \times 2 \times 2 = 8$ elements in the sub-matrix of matrix $A$ and $2 \times 2 = 4$ elements in the sub-matrix of matrix $B$, and produce $2 \times 2 = 4$ elements in the sub-matrix of matrix $C$.

Tiled MMA Configuration

The MMA tile configuration is a trade-off between resource and performance. The more MMA atoms we use, the higher parallelism we can achieve, at a cost of more threads to use and higher pressure for memory access. The fewer MMA atoms we use, the lower parallelism we can achieve, but we can save more threads and reduce the pressure for memory access. To achieve the best performance, we need to find the sweet spot between the two extremes.

123456789101112131415161718192021222324252627282930

// Configure tiled MMA.using MmaTraits = cute::MMA_Traits<cute::SM80_16x8x16_F16F16F16F16_TN>;using MmaAtomShape = MmaTraits::Shape_MNK;auto const mma_atom = cute::MMA_Atom<MmaTraits>{};auto const mma_atom_shape = MmaAtomShape{};// Repeating the mma atom along the M, N, and K dimensions.// This increases the number of threads to process the tiled MMA.constexpr int MMA_LAYOUT_M{2};constexpr int MMA_LAYOUT_N{2};constexpr int MMA_LAYOUT_K{1};auto mma_layout{cute::make_layout(    cute::make_shape(cute::Int<MMA_LAYOUT_M>{}, cute::Int<MMA_LAYOUT_N>{},                     cute::Int<MMA_LAYOUT_K>{}))};// Repeating the mma processing along the M, N, and K dimensions.// This does not increase the number of threads to process the tiled MMA.// But the number of registers required for processing the tiled MMA// increases.constexpr int NUM_MMA_TILE_M{1};constexpr int NUM_MMA_TILE_N{2};constexpr int NUM_MMA_TILE_K{1};constexpr int MMA_TILE_M{cute::get<0>(mma_atom_shape) * MMA_LAYOUT_M *                         NUM_MMA_TILE_M};constexpr int MMA_TILE_N{cute::get<1>(mma_atom_shape) * MMA_LAYOUT_N *                         NUM_MMA_TILE_N};constexpr int MMA_TILE_K{cute::get<2>(mma_atom_shape) * MMA_LAYOUT_K *                         NUM_MMA_TILE_K};auto mma_tile{cute::make_tile(cute::Int<MMA_TILE_M>{},                              cute::Int<MMA_TILE_N>{},                              cute::Int<MMA_TILE_K>{})};auto tiled_mma{cute::make_tiled_mma(mma_atom, mma_layout, mma_tile)};

In our tiled MMA configuration, we have 2 MMA atoms along the $M$ dimension, 2 MMA atoms along the $N$ dimension, and $1$ MMA along the $K$ dimension, i.e., no parallelism in the $K$ dimension. Therefore, we have $2 \times 2 \times 1 = 4$ MMA atoms in total. Because each MMA atom consists of a warp of 32 threads, we have the tiled MMA ThrLayoutVMNK = (_32,_2,_2,_1):(_1,_32,_64,_0). The number of threads involved in this tiled MMA is $32 \times 2 \times 2 \times 1 = 128$. This tiled MMA in process can solve an MMA problem of size $(2 \times M^{\prime}) \times (2 \times N^{\prime}) \times (1 \times K^{\prime}) = 32 \times 16 \times 16$. The same tiled MMA can process multiple times along different dimensions to solve larger MMA problems. In our case, we configured the tiled MMA to process along the $N$ dimension 2 times. As a result, with such permutation, we have the tiled MMA PermutationMNK: (_32,_32,_16) that solves an MMA problem of size $32 \times 32 \times 16$.

12345678910

tiled_mmaTiledMMA  ThrLayoutVMNK:  (_32,_2,_2,_1):(_1,_32,_64,_0)  PermutationMNK: (_32,_32,_16)MMA_Atom  ThrID:      _32:_1  Shape_MNK:  (_16,_8,_16)  LayoutA_TV: ((_4,_8),(_2,_2,_2)):((_32,_1),(_16,_8,_128))  LayoutB_TV: ((_4,_8),(_2,_2)):((_16,_1),(_8,_64))  LayoutC_TV: ((_4,_8),(_2,_2)):((_32,_1),(_16,_8))

The tiled MMA layouts for the MMA matrices can also be visualized using the CuTe cute::print_latex or cute::print_svg functions.

Tiled MMA Memory Copy Partition

The tiled MMA could then be used as the building block for solving MMA problems of even larger size. Given large MMA matrix tensors, the tiled MMA can be decomposed into thread MMAs and each thread MMA has the methods partition_A, partition_B, and partition_C to partition the MMA matrix tensors into sub-matrices needed for the thread that can be processed by the tiled MMA. The partitioned MMA matrix tensors are then used as the input to the tiled MMA and larger MMA problems are solved by repeating the tiled MMA along different dimensions using the cute::gemm API.

123456789

auto thread_mma{tiled_mma.get_slice(THREAD_IDX)};// Use no tiled copy from shared memory to register.auto thread_layout_C_smem_tensor_A_no_tiled_copy{    thread_mma.partition_A(smem_tensor_A)}; // (MMA, MMA_M, MMA_K)auto thread_layout_C_smem_tensor_B_no_tiled_copy{    thread_mma.partition_B(smem_tensor_B)}; // (MMA, MMA_N, MMA_K)auto thread_layout_C_smem_tensor_C_no_tiled_copy{    thread_mma.partition_C(smem_tensor_C)}; // (MMA, MMA_M, MMA_N)

In our case, the tiled MMA partitions the matrix A, matrix B, and matrix C as follows:

123456

thread_layout_C_smem_tensor_A_no_tiled_copyptr[16b](0x5c6c073e78c0) o ((_2,_2,_2),_4,_2):((_128,_8,_1024),_32,_2048)thread_layout_C_smem_tensor_B_no_tiled_copyptr[16b](0x5c6c073e98d0) o ((_2,_2),_8,_2):((_128,_1024),_16,_2048)thread_layout_C_smem_tensor_C_no_tiled_copyptr[16b](0x5c6c073eb8e0) o ((_2,_2),_4,_8):((_128,_8),_32,_2048)

Note that, again, each thread in the MMA atom in one process will access $2 \times 2 \times 2 = 8$ elements in the sub-matrix of matrix $A$ and $2 \times 2 = 4$ elements in the sub-matrix of matrix $B$, and produce $2 \times 2 = 4$ elements in the sub-matrix of matrix $C$. Such patterns are repeated $4$ times along the $M$ dimension and $2$ times along the $K$ dimension for matrix $A$, $8$ times along the $N$ dimension and $2$ times along the $K$ dimension for matrix $B$, and $4$ times along the $M$ dimension and $8$ times along the $N$ dimension for matrix $C$. The cute::gemm API will be responsible for accessing the desired data using the correct indices from MMA iterations.

Instead of accessing data from global or shared memory, sometimes we would like to access data from register for data reuse and better performance. The thread MMA also provides the methods partition_fragment_A, partition_fragment_B, and partition_fragment_C that configure the minimum amount of registers needed for the tiled MMA operations in each thread.

1234567

// Register tensors used for MMA.auto thread_layout_C_register_tensor_A{    thread_mma.partition_fragment_A(smem_tensor_A)}; // (MMA, MMA_M, MMA_K)auto thread_layout_C_register_tensor_B{    thread_mma.partition_fragment_B(smem_tensor_B)}; // (MMA, MMA_N, MMA_K)auto thread_layout_C_register_tensor_C{    thread_mma.partition_fragment_C(smem_tensor_C)}; // (MMA, MMA_M, MMA_N)

We could see that the shape of the register MMA tensors are exactly the same as the shape of the corresponding ones on shared memory or global memory. However, because of its compact striding comparing to the ones on shared memory or global memory, no registers are wasted. This is important because the number of registers is limited for each thread and the performance can be compromised if the number of registers configured is too large or there are wasted registers (the compiler might not be able to identify the wasted registers and optimize them out).

123456

thread_layout_C_register_tensor_Aptr[16b](0x7ffc34e465f0) o ((_2,_2,_2),_4,_2):((_1,_2,_4),_8,_32)thread_layout_C_register_tensor_Bptr[16b](0x7ffc34e46670) o ((_2,_2),_8,_2):((_1,_2),_4,_32)thread_layout_C_register_tensor_Cptr[16b](0x7ffc34e466f0) o ((_2,_2),_4,_8):((_1,_2),_4,_16)

Tiled MMA Memory Tiled Copy Partition

In this case, there are some performance issues when the threads tries to access the data from shared memory or global memory. When each thread in the MMA atom tries to access the $2 \times 2 \times 2 = 8$ elements in the sub-matrix of matrix $A$ and the $2 \times 2 = 4$ elements in the sub-matrix of matrix $B$, and the $2 \times 2 = 4$ elements in the sub-matrix of matrix $C$, multiple transactions have to be performed because those data are not contiguous in memory, not even any two elements.

CUDA has special warp-level matrix load instruction ldmatrix that specifically addresses this problem and it is wrapped into the CuTe copy atoms. In our case, for cute::half_t data type, the copy atoms we used are cute::SM75_U16x8_LDSM_T for both matrix $A$ and matrix $B$. This copy atom consists of a warp of 32 threads from the ValLayoutSrc: (_32,_8):(_8,_1) we learned that each thread will copy $8$ contiguous elements from the source memory to the destination memory by abstraction. It is by abstraction because under the hood ldmatrix does not work in this way. The copy atom will copy $32 \times 8 = 128$ elements in total.

123456789101112131415

copy_atom_ACopy_Atom  ThrID:        _32:_1  ValLayoutSrc: (_32,_8):(_8,_1)  ValLayoutDst: ((_4,_8),(_1,_2,_4)):((_16,_1),(_1,_8,_64))  ValLayoutRef: ((_4,_8),(_1,_2,_4)):((_16,_1),(_1,_8,_64))  ValueType:    16bcopy_atom_BCopy_Atom  ThrID:        _32:_1  ValLayoutSrc: (_32,_8):(_8,_1)  ValLayoutDst: ((_4,_8),(_1,_2,_4)):((_16,_1),(_1,_8,_64))  ValLayoutRef: ((_4,_8),(_1,_2,_4)):((_16,_1),(_1,_8,_64))  ValueType:    16b

The copy atoms could be tiled with the tiled MMA for matrix $A$ sub-matrix and matrix $B$ sub-matrix tiled copy using cute::make_tiled_copy_A and cute::make_tiled_copy_B functions. In our case, because the tiled MMA (that has permutations) will solve an MMA problem of size $32 \times 32 \times 16$ (not $32 \times 16 \times 16$), the tiled copy will be responsible for a tile of size $32 \times 16$ for matrix $A$ and $32 \times 16$ for matrix $B$. In fact, without the permutation in the tiled MMA, the tile size for matrix $B$ would become $16 \times 16$ and $16 \times 16$ is smaller than $128$, the number of elements in total that the copy atom will copy, and such tiled copy will be prohibited by CuTe.

123456

// Use tiled copy from shared memory to register.auto copy_atom_A = cute::Copy_Atom<cute::SM75_U16x8_LDSM_T, TA>{};auto copy_atom_B = cute::Copy_Atom<cute::SM75_U16x8_LDSM_T, TB>{};auto smem_tiled_copy_A{cute::make_tiled_copy_A(copy_atom_A, tiled_mma)};auto smem_tiled_copy_B{cute::make_tiled_copy_B(copy_atom_B, tiled_mma)};

123456789101112131415161718192021

smem_tiled_copy_ATiledCopy  Tiler_MN:       (_32,_16)  TiledLayout_TV: ((_4,_8,_2,_2),((_2,_2,_2),(_1,_1))):((_64,_1,_16,_0),((_32,_8,_256),(_0,_0)))Copy_Atom  ThrID:        _32:_1  ValLayoutSrc: (_32,_8):(_8,_1)  ValLayoutDst: ((_4,_8),(_1,_2,_4)):((_16,_1),(_1,_8,_64))  ValLayoutRef: ((_4,_8),(_1,_2,_4)):((_16,_1),(_1,_8,_64))  ValueType:    16bsmem_tiled_copy_BTiledCopy  Tiler_MN:       (_32,_16)  TiledLayout_TV: ((_4,_8,_2,_2),((_2,_2),(_2,_1))):((_64,_1,_0,_8),((_32,_256),(_16,_0)))Copy_Atom  ThrID:        _32:_1  ValLayoutSrc: (_32,_8):(_8,_1)  ValLayoutDst: ((_4,_8),(_1,_2,_4)):((_16,_1),(_1,_8,_64))  ValLayoutRef: ((_4,_8),(_1,_2,_4)):((_16,_1),(_1,_8,_64))  ValueType:    16b

The tiled copy layouts could be printed using the cute::print_latex function as well. In our case, the left layout is the source layout and actually it’s misleading in our particular case. The value mapping between the left source layout and the right destination layout is incorrect. For example, the $T0V1$ value from the left layout will not be copied to the $T0V1$ value in the right layout. The left source layout shows that a thread copies $8$ elements from the same column because we need the threads to be located at the beginning of the column so that the column address can be correctly passed for ldmatrix.

The tiled copy can be decomposed into (abstracted) thread copies and each thread copy has the method partition_S to partition the source layout. Similarly, the thread copy source layout is misleading as one thread will not actually load $8$ contiguous elements from the source memory using the cute::SM75_U16x8_LDSM_T copy atom. But the CuTe tiled copy abstraction will at least ensure the consequent copy behavior is correct.

1234

thread_layout_C_smem_tensor_A_tiled_copyptr[16b](0x57b7b93248c0) o ((_8,_1),_4,_2):((_1,_0),_32,_2048)thread_layout_C_smem_tensor_B_tiled_copyptr[16b](0x57b7b93268d0) o ((_8,_1),_4,_2):((_1,_0),_32,_2048)

To perform the tiled copy, there is still one problem. The layouts for the destination register tensors for tiled MMA are not immediately compatible with the source tensor layouts.

1234

thread_layout_C_register_tensor_Aptr[16b](0x7ffc34e465f0) o ((_2,_2,_2),_4,_2):((_1,_2,_4),_8,_32)thread_layout_C_register_tensor_Bptr[16b](0x7ffc34e46670) o ((_2,_2),_8,_2):((_1,_2),_4,_32)

However, we realized that the sub-layout of the destination register tensor A (_2,_2,_2):(_1,_2,_4) is equivalent as the sub-layout of the source shared memory tensor A (_8,_1):(_1,_0). The sub-layout of the destination register tensor B ((_2,_2),_8:(_1,_2),_4) is equivalent as the sub-layout of the source shared memory tensor B ((_8,_1),_4):((_1,_0),_8). So before running tiled copy using the cute::copy API, the destination register tensors should be retiled using the retile_D method from thread copy.

After retiling the destination register tensors, the layouts become compatible with the source shared memory tensors for tiled copy.

1234

thread_layout_C_register_tensor_A_copy_viewptr[16b](0x7ffd7a129350) o ((_8,_1),_4,_2):((_1,_0),_8,_32)thread_layout_C_register_tensor_B_copy_viewptr[16b](0x7ffd7a1293d0) o ((_8,_1),_4,_2):((_1,_0),_8,_32)

Without using tiled copy, to load $8$ elements from matrix $A$ for tiled MMA in each thread, 8 memory access instructions have to be performed. But with tiled copy, only 1 memory access instruction is needed. Thus, tiled copy from shared memory to register could usually improve the performance.

Of course, when the tiled copy is performed, there could be shared memory bank conflicts. We could try minimizing the shared memory bank conflicts using CuTe swizzle.

References

CuTe Tiled MMA

https://leimao.github.io/blog/CuTe-Tiled-MMA/

Author

Lei Mao

Posted on

01-09-2025

Updated on

01-09-2025

Licensed under

Like this article? Support the author with

Comments

Lei Mao

Artificial Intelligence
Machine Learning
Computer Science

Santa Clara, California

Posts

1008

Categories

7

Tags

662

Advertisement

Catalogue

© 2017-2025 Lei Mao  Powered by Hexo & IcarusSite UV:  Site PV:


cuBLAS GEMM API Usages for Column-Major and Row-Major Matrices

Introduction

The cuBLAS GEMM API has very strict requirements on the storage format of the input and output matrices. If all the matrices are stored in column-major format, the cuBLAS GEMM API can be used straightforwardly. But if some of the matrices are stored in row-major format, setting the parameters for the cuBLAS GEMM API for such matrix multiplications can be error-prone.

In this blog post, we will discuss the relationship between the transpose and column-major storage of matrices and how cuBLAS GEMM API should be used for different cases.

cuBLAS GEMM

cuBLAS GEMM API

The cuBLAS single-precision GEMM API is declared as follows.

12345678

cublasStatus_t cublasSgemm(cublasHandle_t handle,                           cublasOperation_t transa, cublasOperation_t transb,                           int m, int n, int k,                           const float *alpha,                           const float *A, int lda,                           const float *B, int ldb,                           const float *beta,                           float *C, int ldc)

This function performs the general matrix-matrix multiplication

$$\begin{align}C = \alpha \text{op}(A) \text{op}(B) + \beta C\end{align}$$

where $\alpha$ and $\beta$ are scalars, and $A$, $B$, and $C$ are matrices stored in column-major format with dimensions of $\text{op}(A)$ being $m \times k$, $\text{op}(B)$ being $k \times n$, and $C$ being $m \times n$, respectively. Also for matrix $A$

$$\begin{align}\text{op}(A) =\begin{cases}A & \text{if transa = CUBLAS_OP_N} \\A^{\top} & \text{if transa = CUBLAS_OP_T} \\A^{\dagger} & \text{if transa = CUBLAS_OP_C} \\\end{cases}\end{align}$$

cuBLAS GEMM and Row-Major Matrices

But what if some of the matrices are stored in row-major format? Let’s see a few examples.

Suppose $m^{\prime} \times k^{\prime}$ matrix $A^{\prime}$ is stored in row-major format, and $k^{\prime} \times n^{\prime}$ matrix $B^{\prime}$ and $m^{\prime} \times n^{\prime}$ matrix $C^{\prime}$ are stored in column-major format. The transpose of $A^{\prime}$, $k^{\prime} \times m^{\prime}$ matrix $A^{\prime\top}$, stored in column-major format, is equivalent to the original $A^{\prime}$ stored in row-major format. But in order to perform the general matrix-matrix multiplication using cuBLAS, $A^{\prime\top}$ has to be transposed to $A^{\prime}$. In this case, transa = CUBLAS_OP_T, transb = CUBLAS_OP_N, m = m', n = n', k = k', A = A', B = B', and C = C'.

Suppose $m^{\prime} \times k^{\prime}$ matrix $A^{\prime}$ and $k^{\prime} \times n^{\prime}$ matrix $B^{\prime}$ are stored in column-major format, and $m^{\prime} \times n^{\prime}$ matrix $C^{\prime}$ is stored in row-major format. In this case, there is no way to transpose $C^{\prime}$ via the cuBLAS API.

We notice that we could transpose matrix $C$ in the formula first before performing the general matrix-matrix multiplication.

$$\begin{align}C^{\top} &= \alpha \left(\text{op}(A) \text{op}(B)\right)^{\top} + \beta C^{\top} \\&= \alpha \text{op}(B)^{\top} \text{op}(A)^{\top} + \beta C^{\top} \\&= \alpha \text{op}(B^{\top}) \text{op}(A^{\top}) + \beta C^{\top}\end{align}$$

So if $B^{\top}$, $A^{\top}$, and $C^{\top}$ are stored in column-major format, we could still perform the general matrix-matrix multiplication using the existing cuBLAS API.

In this case, the transpose of $C^{\prime}$, $n^{\prime} \times m^{\prime}$ matrix $C^{\prime\top}$, stored in column-major format, is equivalent to the original $C^{\prime}$ stored in row-major format. In addition, the matrix $A^{\prime}$ and $B^{\prime}$ have to be transposed as well. The transpose of $A^{\prime}$, $k^{\prime} \times m^{\prime}$ matrix $A^{\prime\top}$, stored in row-major format, is equivalent to the original $A^{\prime}$ stored in column-major format. The transpose of $B^{\prime}$, $n^{\prime} \times k^{\prime}$ matrix $B^{\prime\top}$, stored in row-major format, is equivalent to the original $B^{\prime}$ stored in column-major format. In this case, transa = CUBLAS_OP_T, transb = CUBLAS_OP_T, m = n', n = m', k = k', A = B', B = A', and C = C'.

Conclusions

Suppose we want to perform matrix multiplication $C^{\prime} = \alpha A^{\prime} B^{\prime} + \beta C^{\prime}$, where $A^{\prime}$, $B^{\prime}$, and $C^{\prime}$ are matrices of shapes $m^{\prime} \times k^{\prime}$, $k^{\prime} \times n^{\prime}$, and $m^{\prime} \times n^{\prime}$, respectively, using cuBLAS API. The following table summarizes the relationship between the transpose and column-major storage of matrices $A^{\prime}$, $B^{\prime}$, and $C^{\prime}$, and how cuBLAS API should be used.

References

cuBLAS GEMM API Usages for Column-Major and Row-Major Matrices

https://leimao.github.io/blog/cuBLAS-Transpose-Column-Major-Relationship/

Author

Lei Mao

Posted on

12-12-2024

Updated on

12-12-2024

Licensed under

Like this article? Support the author with

Comments

Lei Mao

Artificial Intelligence
Machine Learning
Computer Science

Santa Clara, California

Posts

1008

Categories

7

Tags

662

Advertisement

Catalogue

© 2017-2025 Lei Mao  Powered by Hexo & IcarusSite UV:  Site PV:


CuTe Swizzle

Introduction

In my article “CUDA Shared Memory Swizzling”, we have discussed how to use swizzling to avoid bank conflicts when a warp of threads accesses shared memory in a strided pattern. Because the swizzle operation and the mathematical proof involve both integer and bit operations, it might not be straightforward to understand and could be error-prone to implement.

CuTe provided a shared memory swizzling abstraction class Swizzle to simplify the shared memory swizzling implementation. Its implementation only involves bit operations, therefore it is more readable and somewhat easier to prove.

In this blog post, I would like to quickly discuss the implementation of the CuTe shared memory swizzling abstraction class Swizzle and its configurations in practice.

CuTe Swizzle

CuTe Swizzle Implementation

The CuTe shared memory swizzling abstraction class Swizzle from the source code is as follows. Only three parameters are used for the swizzle configuration: BBits, MBase, and SShift, where BBits is the number of bits in the mask, MBase is the number of least-significant bits to keep constant, and SShift is the distance to shift the mask. This might be obscure at first glance, let’s walk through a quick example.

For simplicity, suppose we have a 16-bit integer offset whose value is 65 and a swizzle configuration Swizzle<5, 0, 6>. The bit representation of offset is 0b0000000001000001. The bit_msk is 0b0000000000011111, the yyy_msk is 0b0000011111000000, the zzz_msk is 0b0000000000011111, and msk_sft is 6. To swizzle the offset, offset & yyy_msk{} is 0b0000000001000000 where only the bits in the masked region are kept, and shiftr(offset & yyy_msk{}, msk_sft{}) is 0b0000000000000001 where the masked bits are shifted to the right. The final result is offset ^ shiftr(offset & yyy_msk{}, msk_sft{}) which is 0b0000000001000001 xor 0b0000000000000001 equals 0b0000000001000000 whose value is 64. This means the swizzle operation Swizzle<5, 0, 6> projects the offset from 65 to 64.

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455

// A generic Swizzle functor/* 0bxxxxxxxxxxxxxxxYYYxxxxxxxZZZxxxx *                               ^--^ MBase is the number of least-sig bits to keep constant *                  ^-^       ^-^     BBits is the number of bits in the mask *                    ^---------^     SShift is the distance to shift the YYY mask *                                       (pos shifts YYY to the right, neg shifts YYY to the left) * * e.g. Given * 0bxxxxxxxxxxxxxxxxYYxxxxxxxxxZZxxx * the result is * 0bxxxxxxxxxxxxxxxxYYxxxxxxxxxAAxxx where AA = ZZ xor YY */template <int BBits, int MBase, int SShift = BBits>struct Swizzle{  static constexpr int num_bits = BBits;  static constexpr int num_base = MBase;  static constexpr int num_shft = SShift;  static_assert(num_base >= 0,             "MBase must be positive.");  static_assert(num_bits >= 0,             "BBits must be positive.");  static_assert(abs(num_shft) >= num_bits, "abs(SShift) must be more than BBits.");  // using 'int' type here to avoid unintentially casting to unsigned... unsure.  using bit_msk = cute::constant<int, (1 << num_bits) - 1>;  using yyy_msk = cute::constant<int, bit_msk{} << (num_base + max(0,num_shft))>;  using zzz_msk = cute::constant<int, bit_msk{} << (num_base - min(0,num_shft))>;  using msk_sft = cute::constant<int, num_shft>;  static constexpr uint32_t swizzle_code = uint32_t(yyy_msk{} | zzz_msk{});  template <class Offset>  CUTE_HOST_DEVICE constexpr static  auto  apply(Offset const& offset)  {    return offset ^ shiftr(offset & yyy_msk{}, msk_sft{});   // ZZZ ^= YYY  }  template <class Offset>  CUTE_HOST_DEVICE constexpr  auto  operator()(Offset const& offset) const  {    return apply(offset);  }  template <int B, int M, int S>  CUTE_HOST_DEVICE constexpr  auto  operator==(Swizzle<B,M,S> const&) const  {    return B == BBits && M == MBase && S == SShift;  }};

Offset Bijection

Given an integer $m$, a domain $X = [0, 2^m - 1]$, and a constant $c \in X$, the function $f: X \to X$ defined as $f(x) = x \oplus c$, where $\oplus$ is the XOR operation, is a bijection.

Proof

It is trivial to see the XOR operation is commutative and associative. Suppose there exist two different values $x_1, x_2 \in X$ such that $f(x_1) = f(x_2)$, then $x_1 \oplus c = x_2 \oplus c$, which implies $x_1 = x_2$. Therefore, the function $f$ is injective. Because $X = [0, 2^m - 1]$ and the XOR operation cannot produce a value outside of $X$, the function $f$ is surjective.

This concludes the proof. $\square$

Therefore, assuming MBase is zero, and BBits = m, for the offsets $x$ in the range $X = [k, k \cdot 2^m]$, because $c$ is a constant (as it’s in the bit position higher than $m$), the swizzle operation will permute the offsets in $X$ bijectively.

Shared Memory 2D Layout and Shared Memory Bank Conflicts

On devices of compute capability 5.x or newer, each bank has a bandwidth of 32 bits every clock cycle, and successive 32-bit words are assigned to successive banks. So element size matters for shared memory bank conflicts.

In many use cases, the shared memory layout is a 2D row-major matrix whose row size is a multiple of 32 and element size is 32-bit. When strided access from a warp of threads is performed on the column of the matrix, severe 32-way shared memory bank conflicts will occur. If the element in each column are mapped to different shared memory banks, the shared memory bank conflicts can be mitigated.

Assuming MBase is zero, element size is 32-bit, and the matrix row size is $n$, we could design the swizzle operation such that there is free of shared memory bank conflicts when accessing each column of the matrix, based on the offset bijection property we just proved above. To configure BBits and SShift such that offset % 32 is distinct when a warp of threads accesses the column of the matrix, we have to set SShift to be $\log_2 n$ and BBits to be $\log_2 32 = 5$, so that the $c$ used in $f(x) = x \oplus c$ are different for each row, resulting in distinct offset % 32 when a warp of threads accesses the column of the matrix.

Vectorized Memory Access

In CuTe, when we want to perform vectorized access, which is common in CUDA kernels such as matrix transpose, all the elements in the vector from both source and target must be contiguous on the memory. If MBase is zero, the swizzle operation will not be able to guarantee the contiguous memory access because of the XOR operation. However, when we know the number of elements in the vector, we could set MBase to the log2 of the number of elements in the vector, and as we will prove later, the swizzle operation will guarantee the contiguous memory access.

Suppose a vector holds $n$ values where $n$ is a power of 2, then all the bits of offsets $kn, kn + 1, \ldots, kn + n - 1$, where $k$ is an integer, are the same except the least significant $\log_2 n$ bits.

Proof

It is trivial to see the least significant $\log_2 n$ bits of $kn, kn + 1, \ldots, kn + n - 1$ are $0, 1, \ldots, n - 1$ in decimal. Suppose there exist one value among $kn, kn + 1, \ldots, kn + n - 1$ whose non-least significant bits are different from one of its neighbors, then the difference between the value and its neighbor must be different from $1$, which contradicts the fact that the difference between any two consecutive values is $1$ in decimal.

This concludes the proof. $\square$

Therefore, no matter how the swizzle operation is performed, as long as MBase is set to $\log_2 n$ and the memory access starting offset is a multiple of $n$, the contiguous memory access is guaranteed.

If the element size is less than or equal to 32-bit, such as 1-bit, 2-bit, 4-bit, 8-bit, 16-bit, and 32-bit, we could set the number of values in the vector to be 32, 16, 8, 4, 2, and 1, and set MBase to be $\log_2 32 = 5$, $\log_2 16 = 4$, $\log_2 8 = 3$, $\log_2 4 = 2$, $\log_2 2 = 1$, and $\log_2 1 = 0$, respectively, such that each vectorized memory access transaction is 32-bit. Consequently, the vector is just treated as if it were just a 32-bit element, and the offset can be re-indexed by offset >>= MBase before the swizzle operation. Then all the swizzle properties we just proved above still hold and the swizzle configurations can be reused.

If the element size is greater than 32-bit, such as 64-bit, 128-bit, 256-bit, etc., we could treat it as a vectorized memory access of multiple 32-bit elements. The offset can be re-indexed by offset <<= log2(element_size / 32) before the swizzle operation, and the swizzle configurations can be reused. This, however, will bring a consequence that there can never be free of shared memory bank conflicts when accessing the column of the matrix, because the image of shared memory bank ids offset % 32 will become smaller than $[0, 32)$. According to the pigeon hole principle, there will always be at least a log2(element_size / 32)-way shared memory bank conflict.

Universal Swizzle Equations and Configurations

Suppose the element size is $S$-byte, the number of elements in a vector is $N$, the number of elements in the fast dimension of the shared memory is $X$. The MBase should be set to $\log_2 N$, the BBits should be set to $\log_2 (32 \times 4 / S) - \text{MBase}$, and the SShift should be set to $\log_2 X - \text{MBase}$.

An example of the universal swizzle configurations could be implemented as follows.

123456789101112131415

constexpr int constexpr_log2(int n){    return ((n < 2) ? 0 : 1 + constexpr_log2(n / 2));}using VectorType = cute::uint128_t;CUTE_STATIC_ASSERT(sizeof(VectorType) % sizeof(T) == 0,                    "sizeof(VectorType) must be a multiple of sizeof(T)");constexpr unsigned int NUM_VECTOR_ELEMENTS{sizeof(VectorType) / sizeof(T)};using TileSizeX = cute::Int<128>; // Fast dimension size on shared memory.using TileSizeY = cute::Int<32>;  // Slow dimension size on shared memory.constexpr int NUM_BASE_BITS{constexpr_log2(NUM_VECTOR_ELEMENTS)};constexpr int NUM_MASK_BITS{constexpr_log2(32 * 4 / sizeof(T)) - NUM_BASE_BITS};constexpr int NUM_SHIFT_BITS{constexpr_log2(TileSizeX::value) - NUM_BASE_BITS};

Examples

Let’s see a few more complicated examples in which the data type is not of 32-bit size and vectorized memory access is used.

Suppose we have an INT8 8 x 128 row-major matrix, and we want to use 128-bit vectorized memory access. In this case, there are 16 elements in a vector, and MBase should be set to $\log_2 16 = 4$. Because there are 32 shared memory banks of 32-bit size, each 32-bit word contains 4 elements, so BBits should be set to $\log_2 (32 \times 4) - \text{MBase} = 7 - 4 = 3$. The SShift is set to $\log_2 128 - \text{MBase} = 7 - 4 = 3$ to ensure the constant $c$ used for the XOR operation is different for each row. Therefore, the swizzle configuration is Swizzle<3, 4, 3>.

Suppose we have an FP16 8 x 64 row-major matrix, and we want to use 128-bit vectorized memory access. In this case, there are 8 elements in a vector, and MBase should be set to $\log_2 8 = 3$. Because there are 32 shared memory banks of 32-bit size, each 32-bit word contains 2 elements, so BBits should be set to $\log_2 (32 \times 2) - \text{MBase} = 6 - 3 = 3$. The SShift is set to $\log_2 64 - \text{MBase} = 6 - 3 = 3$ to ensure the constant $c$ used for the XOR operation is different for each row. Therefore, the swizzle configuration is Swizzle<3, 3, 3>.

CuTe Swizzle Preview

The shared memory bank ids of CuTe swizzled layout can be previewed using the CuTe Swizzle Preview App I created. The app saves the shared memory bank ids of the swizzled layout to a LaTeX file, and the LaTeX file can be compiled to a PDF file for previewing.

For example, given a 2D row-major $32 \times 64$ matrix consisting of elements of 32-bit size, we printed the shared memory bank ids of the swizzled layouts, including Swizzle<5, 0, 6>, Swizzle<5, 0, 8>, Swizzle<5, 2, 6>, and Swizzle<5, 2, 8>. Only Swizzle<5, 0, 6> results in free of shared memory bank conflicts when accessing each column of the matrix.

References

CuTe Swizzle

https://leimao.github.io/blog/CuTe-Swizzle/

Author

Lei Mao

Posted on

12-01-2024

Updated on

12-26-2024

Licensed under

Like this article? Support the author with

Comments

Lei Mao

Artificial Intelligence
Machine Learning
Computer Science

Santa Clara, California

Posts

1008

Categories

7

Tags

662

Advertisement

Catalogue

© 2017-2025 Lei Mao  Powered by Hexo & IcarusSite UV:  Site PV: