CuTe Matrix Transpose

Introduction

CuTe is a C++ template library that provides a high-level abstraction for layout and tensor operations in CUDA kernels. CUTLASS 3.0 and beyond adopts CuTe throughout the GEMM hierarchy in its templates, allowing the implementation to be more readable and maintainable.

Previously, I have created an article “CuTe Layout Algebra” on the mathematical foundations of CuTe. In this blog post, we will have some hands-on experience and have a better understanding of CuTe by implementing matrix transpose CUDA kernels.

CuTe Matrix Transpose

Matrix transpose CUDA kernels are probably the most example CUDA kernels that I have ever implemented. In my previous examples, the thread and data index mappings in the CUDA kernels are completely manual. There were also some hard-coded assumptions, such as each CUDA thread will only process one single element in matrix transpose to make the implementation easier and more human readable. To have both the configuration complexity and human readability in the implementation, we can create matrix transpose CUDA kernels using CuTe.

To transpose a matrix in a CUDA kernel, performing strided memory reads or writes in a warp is inevitable and it will lead to uncoalesced memory accesses, resulting in performance degradation. To mitigate the performance degradation, the strided memory reads or writes could be performed on shared memory instead of global memory. When the strided memory reads or writes are performed on shared memory, special optimizations have also to be performed to avoid shared memory bank conflicts.

All the CuTe matrix transpose CUDA kernels implemented in this article and their unit tests could be found from my CUTLASS Examples GitHub repository.

CuTe Naive Matrix Transpose

In the CuTe naive matrix transpose CUDA kernel implementation, we will not use shared memory. Two slightly different CUDA kernel variants have been implemented. One performs coalesced global memory reads and strided global memory writes, and the other performs strided global memory reads and coalesced global memory writes. It turns out that the difference between the implementations of the two variants is just one line of code.

CuTe Naive Matrix Transpose Implementation

The CuTe naive matrix transpose CUDA kernel implementation could also be found from my CUTLASS Examples GitHub repository.

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190

#include <cuda_runtime.h>#include <cute/tensor.hpp>#include "cute_matrix_transpose.hpp"template <class TensorSrc, class TensorDst, class ThreadLayout>static __global__ void matrix_transpose_naive(TensorSrc tensor_src,                                              TensorDst tensor_dst_transposed,                                              ThreadLayout){    using Element = typename TensorSrc::value_type;    auto global_tile_src{tensor_src(cute::make_coord(cute::_, cute::_),                                    blockIdx.y,                                    blockIdx.x)}; // (TileSizeY, TileSizeX)    auto global_tile_dst_transposed{        tensor_dst_transposed(cute::make_coord(cute::_, cute::_), blockIdx.y,                              blockIdx.x)}; // (TileSizeY, TileSizeX)    auto thread_global_tile_src{cute::local_partition(        global_tile_src, ThreadLayout{},        threadIdx.x)}; // (ThreadValueSizeY, ThreadValueSizeX)    auto thread_global_tile_dst_transposed{cute::local_partition(        global_tile_dst_transposed, ThreadLayout{},        threadIdx.x)}; // (ThreadValueSizeY, ThreadValueSizeX)    // A 2D array of tuples that maps (x, y) to (x, y).    auto const identity_tensor{cute::make_identity_tensor(cute::make_shape(        cute::size<0>(global_tile_src), cute::size<1>(global_tile_src)))};    auto const thread_identity_tensor{        cute::local_partition(identity_tensor, ThreadLayout{}, threadIdx.x)};    auto fragment{cute::make_tensor_like(thread_global_tile_src)};    auto predicator{cute::make_tensor<bool>(        cute::make_shape(cute::size<0>(fragment), cute::size<1>(fragment)))};    auto const num_max_columns{cute::stride<0>(global_tile_src)};    auto const num_max_rows{cute::stride<1>(global_tile_dst_transposed)};    constexpr auto global_tile_columns{cute::size<1>(global_tile_src)};    constexpr auto global_tile_rows{cute::size<0>(global_tile_src)};    CUTE_UNROLL    for (unsigned int i{0}; i < cute::size<0>(predicator); ++i)    {        CUTE_UNROLL        for (unsigned int j{0}; j < cute::size<1>(predicator); ++j)        {            auto const thread_identity{thread_identity_tensor(i, j)};            bool const is_row_in_bound{cute::get<0>(thread_identity) +                                           blockIdx.y * global_tile_rows <                                       num_max_rows};            bool const is_column_in_bound{cute::get<1>(thread_identity) +                                              blockIdx.x * global_tile_columns <                                          num_max_columns};            predicator(i, j) = is_row_in_bound && is_column_in_bound;        }    }    cute::copy_if(predicator, thread_global_tile_src, fragment);    cute::copy_if(predicator, fragment, thread_global_tile_dst_transposed);    // Alternatively, we could just do the following instead.    // cute::copy_if(predicator, thread_global_tile_src,    // thread_global_tile_dst_transposed);}enum class GlobalMemoryCoalescedAccessMode{    Read,    Write};template <typename T>static cudaError_t launch_matrix_transpose_naive_base(    T const* input_matrix, T* output_matrix, unsigned int M, unsigned int N,    GlobalMemoryCoalescedAccessMode coalesced_access_mode, cudaStream_t stream){    auto const tensor_shape{cute::make_shape(M, N)};    auto const tensor_shape_transposed{cute::make_shape(N, M)};    // Input matrix: row-major M x N matrix.    auto const global_memory_layout_src{cute::make_layout(        tensor_shape, cute::GenRowMajor{})}; // (M, N) : (N, 1)    // Output matrix: row-major N x M matrix.    auto const global_memory_layout_dst{cute::make_layout(        tensor_shape_transposed, cute::GenRowMajor{})}; // (N, M) : (M, 1)    // Same output matrix, but different view: column-major M x N matrix.    auto const global_memory_layout_dst_transposed{cute::make_layout(        tensor_shape, cute::GenColMajor{})}; // (M, N) : (1, M)    auto const tensor_src{cute::make_tensor(cute::make_gmem_ptr(input_matrix),                                            global_memory_layout_src)};    auto const tensor_dst{cute::make_tensor(cute::make_gmem_ptr(output_matrix),                                            global_memory_layout_dst)};    auto const tensor_dst_transposed{        cute::make_tensor(cute::make_gmem_ptr(output_matrix),                          global_memory_layout_dst_transposed)};    using TileSizeX = cute::Int<64>; // bN    using TileSizeY = cute::Int<32>; // bM    constexpr auto block_shape{cute::make_shape(TileSizeY{}, TileSizeX{})};    auto const tiled_tensor_src{cute::tiled_divide(        tensor_src, block_shape)}; // ((TileSizeY, TileSizeX), M /                                   // TileSizeY, N / TileSizeX)    auto const tiled_tensor_dst_transposed{cute::tiled_divide(        tensor_dst_transposed, block_shape)}; // ((TileSizeY, TileSizeX), M                                              // / TileSizeY, N / TileSizeX)    using ThreadBlockSizeX = cute::Int<32>; // tN    using ThreadBlockSizeY = cute::Int<8>;  // tM    constexpr auto thread_block_shape{        cute::make_shape(ThreadBlockSizeY{}, ThreadBlockSizeX{})};    constexpr auto thread_block_shape_transposed{        cute::make_shape(ThreadBlockSizeX{}, ThreadBlockSizeY{})};    // Coalesced memory read.    constexpr auto thread_layout{        cute::make_layout(thread_block_shape, cute::GenRowMajor{})};    // Coalesced memory write.    constexpr auto thread_layout_transposed{        cute::make_layout(thread_block_shape_transposed, cute::GenColMajor{})};    dim3 const grid_dim{cute::size<2>(tiled_tensor_src),                        cute::size<1>(tiled_tensor_src)};    dim3 const thread_dim{        cute::size(ThreadBlockSizeX::value * ThreadBlockSizeY::value)};    if (coalesced_access_mode == GlobalMemoryCoalescedAccessMode::Read)    {        CUTE_STATIC_ASSERT_V(TileSizeX{} % ThreadBlockSizeX{} == cute::Int<0>{},                             "TileSizeX must be divisible by ThreadBlockSizeX");        CUTE_STATIC_ASSERT_V(TileSizeY{} % ThreadBlockSizeY{} == cute::Int<0>{},                             "TileSizeY must be divisible by ThreadBlockSizeY");        matrix_transpose_naive<<<grid_dim, thread_dim, 0, stream>>>(            tiled_tensor_src, tiled_tensor_dst_transposed, thread_layout);    }    else    {        CUTE_STATIC_ASSERT_V(TileSizeX{} % ThreadBlockSizeY{} == cute::Int<0>{},                             "TileSizeX must be divisible by ThreadBlockSizeY");        CUTE_STATIC_ASSERT_V(TileSizeY{} % ThreadBlockSizeX{} == cute::Int<0>{},                             "TileSizeY must be divisible by ThreadBlockSizeX");        matrix_transpose_naive<<<grid_dim, thread_dim, 0, stream>>>(            tiled_tensor_src, tiled_tensor_dst_transposed,            thread_layout_transposed);    }    return cudaGetLastError();}template <typename T>cudaError_t launch_matrix_transpose_naive_coalesced_read(T const* input_matrix,                                                         T* output_matrix,                                                         unsigned int M,                                                         unsigned int N,                                                         cudaStream_t stream){    return launch_matrix_transpose_naive_base(        input_matrix, output_matrix, M, N,        GlobalMemoryCoalescedAccessMode::Read, stream);}template <typename T>cudaError_t launch_matrix_transpose_naive_coalesced_write(T const* input_matrix,                                                          T* output_matrix,                                                          unsigned int M,                                                          unsigned int N,                                                          cudaStream_t stream){    return launch_matrix_transpose_naive_base(        input_matrix, output_matrix, M, N,        GlobalMemoryCoalescedAccessMode::Write, stream);}// Explicit instantiation.template cudaError_t launch_matrix_transpose_naive_coalesced_read<float>(    float const* input_matrix, float* output_matrix, unsigned int M,    unsigned int N, cudaStream_t stream);template cudaError_t launch_matrix_transpose_naive_coalesced_read<double>(    double const* input_matrix, double* output_matrix, unsigned int M,    unsigned int N, cudaStream_t stream);template cudaError_t launch_matrix_transpose_naive_coalesced_write<float>(    float const* input_matrix, float* output_matrix, unsigned int M,    unsigned int N, cudaStream_t stream);template cudaError_t launch_matrix_transpose_naive_coalesced_write<double>(    double const* input_matrix, double* output_matrix, unsigned int M,    unsigned int N, cudaStream_t stream);

Matrix Layout

There are typically two ways to describe a matrix stored in linear storages: row-major and column-major. In row-major layout, the elements of each row are contiguous in memory, while in column-major layout, the elements of each column are contiguous in memory. Given a $M \times N$ input matrix $A$ with $M$ rows and $N$ columns in row-major layout, typically we want to transpose the input matrix $A$ to a $N \times M$ output matrix $A^{\top}$ with $N$ rows and $M$ columns that is also in row-major layout. In this case, the output matrix $A^{\top}$ can not only be viewed as a row-major $N \times M$ matrix but also as a column-major $M \times N$ matrix.

The matrix transpose operation maps the element $A_{i, j}$ from the input matrix $A$ to the element $A^{\top}_{j, i}$ in the output matrix $A^{\top}$. On one hand, if the input matrix is described using row-major layout, the input matrix $A$ is of shape $(M, N)$, and the element $A_{i, j}$ is stored at the coordinate $(i, j)$ in the row-major layout of $A$. The output matrix $A^{\top}$ is of shape $(N, M)$, and the element $A^{\top}_{j, i}$ is stored at the coordinate $(j, i)$ in the row-major layout of $A^{\top}$. On the other hand, if the input matrix is described using column-major layout, the same input matrix $A$ is of shape $(M, N)$, and the element $A_{i, j}$ is stored at the coordinate $(j, i)$ in the column-major layout of $A$. The output matrix $A^{\top}$ is of shape $(M, N)$, and the element $A^{\top}_{j, i}$ is stored at the coordinate $(i, j)$ in the column-major layout of $A^{\top}$.

Although being a little bit brain-twisting, matrix transpose maps an element at the coordinate $(i, j)$ in the row-major layout of a matrix to the element at the coordinate $(i, j)$ in the column-major layout of the output matrix. In CuTe, given an 1D input coordinate and the input matrix and the output matrix both have a shape of $(M, N)$ in the layout, the 1D input coordinate will always be mapped to the same natural coordinate in both the input matrix and the output matrix. When CuTe iterates over $M \times N$ 1D coordinates, the corresponding elements in the input matrix and the output matrix is in a relationship of transpose. This is the key reason why we have to use row-major layout and column-major layout to describe the input matrix and the output matrix, respectively, in CuTe. Otherwise, if the input matrix and the output matrix are both described using the same layout, when CuTe iterates over $M \times N$ 1D coordinates, the corresponding elements in the input matrix and the output matrix will not be in a relationship of transpose.

1234567891011121314151617181920

auto const tensor_shape{cute::make_shape(M, N)};auto const tensor_shape_transposed{cute::make_shape(N, M)};// Input matrix: row-major M x N matrix.auto const global_memory_layout_src{cute::make_layout(    tensor_shape, cute::GenRowMajor{})}; // (M, N) : (N, 1)// Output matrix: row-major N x M matrix.auto const global_memory_layout_dst{cute::make_layout(    tensor_shape_transposed, cute::GenRowMajor{})}; // (N, M) : (M, 1)// Same output matrix, but different view: column-major M x N matrix.auto const global_memory_layout_dst_transposed{cute::make_layout(    tensor_shape, cute::GenColMajor{})}; // (M, N) : (1, M)auto const tensor_src{cute::make_tensor(cute::make_gmem_ptr(input_matrix),                                        global_memory_layout_src)};auto const tensor_dst{cute::make_tensor(cute::make_gmem_ptr(output_matrix),                                        global_memory_layout_dst)};auto const tensor_dst_transposed{    cute::make_tensor(cute::make_gmem_ptr(output_matrix),                      global_memory_layout_dst_transposed)};

Divide and Conquer and Matrix Tiling

To accelerate matrix transpose for large problems, we will have to divide the input matrix and the output matrix into smaller tiles and compute the transpose of each tile in parallel. In this example, we divide the input matrix and the output matrix into tiles of shape $(bM, bN)$, where $bM$ and $bN$ are the number of rows and columns in each tile, respectively. Both the input matrix and the output matrix will be divided into $\left\lceil \frac{M}{bM} \right\rceil \times \left\lceil \frac{N}{bN} \right\rceil$ tiles. The matrix transpose in each tile is independent and can be processed in parallel.

The divided input matrix and the divided output matrix now have new layouts, whose shapes are both $\left((bM, bN), \left\lceil \frac{M}{bM} \right\rceil \times \left\lceil \frac{N}{bN} \right\rceil\right)$. The row-major and column-major notations are no longer applicable for describing the divided matrices, because the shapes now have 3 modes, i.e., a rank of 3. To describe the storage layout of a tensor that has higher rank (any rank), CuTe uses stride. In our particular problem, it’s not too important, because CuTe automatically handles those concepts for us. In other problems, it might not be the case though.

1234567891011

using TileSizeX = cute::Int<64>; // bNusing TileSizeY = cute::Int<32>; // bMconstexpr auto block_shape{cute::make_shape(TileSizeY{}, TileSizeX{})};auto const tiled_tensor_src{cute::tiled_divide(    tensor_src, block_shape)}; // ((TileSizeY, TileSizeX), M /                               // TileSizeY, N / TileSizeX)auto const tiled_tensor_dst_transposed{cute::tiled_divide(    tensor_dst_transposed, block_shape)}; // ((TileSizeY, TileSizeX), M                                          // / TileSizeY, N / TileSizeX)

CUDA Thread Block Layout and Coalesced Memory Access

Each tile of the input matrix and the output matrix is processed by a CUDA thread block that consists of multiple CUDA threads. In our case, we use a thread block of shape $(tM, tN)$ and row-major layout or $(tN, tM)$ and column-major layout. The number of threads in a CUDA thread block is $tM \times tN$. The number of CUDA thread blocks to issue is, obviously, just the number of tiles, i.e., $\left\lceil \frac{M}{bM} \right\rceil \times \left\lceil \frac{N}{bN} \right\rceil$. This is feasible because $bM$ and $bN$, $tM$ and $tN$, are all compile-time constants.

Because the input matrix and its tiles are of row-major layout, and the output matrix and its tiles are of column-major layout, when the thread block is of row-major layout, each warp in the thread block will read from the input matrix on global memory in a coalesced fashion but write to the output matrix on global memory in a strided fashion. Similarly, when the thread block is of column-major layout, each warp in the thread block will read from the input matrix on global memory in a strided fashion but write to the output matrix on global memory in a coalesced fashion.

1234567891011121314151617181920212223242526272829303132333435363738

using ThreadBlockSizeX = cute::Int<32>; // tNusing ThreadBlockSizeY = cute::Int<8>;  // tMconstexpr auto thread_block_shape{    cute::make_shape(ThreadBlockSizeY{}, ThreadBlockSizeX{})};constexpr auto thread_block_shape_transposed{    cute::make_shape(ThreadBlockSizeX{}, ThreadBlockSizeY{})};// Coalesced memory read.constexpr auto thread_layout{    cute::make_layout(thread_block_shape, cute::GenRowMajor{})};// Coalesced memory write.constexpr auto thread_layout_transposed{    cute::make_layout(thread_block_shape_transposed, cute::GenColMajor{})};dim3 const grid_dim{cute::size<2>(tiled_tensor_src),                    cute::size<1>(tiled_tensor_src)};dim3 const thread_dim{    cute::size(ThreadBlockSizeX::value * ThreadBlockSizeY::value)};if (coalesced_access_mode == GlobalMemoryCoalescedAccessMode::Read){    CUTE_STATIC_ASSERT_V(TileSizeX{} % ThreadBlockSizeX{} == cute::Int<0>{},                         "TileSizeX must be divisible by ThreadBlockSizeX");    CUTE_STATIC_ASSERT_V(TileSizeY{} % ThreadBlockSizeY{} == cute::Int<0>{},                         "TileSizeY must be divisible by ThreadBlockSizeY");    matrix_transpose_naive<<<grid_dim, thread_dim, 0, stream>>>(        tiled_tensor_src, tiled_tensor_dst_transposed, thread_layout);}else{    CUTE_STATIC_ASSERT_V(TileSizeX{} % ThreadBlockSizeY{} == cute::Int<0>{},                         "TileSizeX must be divisible by ThreadBlockSizeY");    CUTE_STATIC_ASSERT_V(TileSizeY{} % ThreadBlockSizeX{} == cute::Int<0>{},                         "TileSizeY must be divisible by ThreadBlockSizeX");    matrix_transpose_naive<<<grid_dim, thread_dim, 0, stream>>>(        tiled_tensor_src, tiled_tensor_dst_transposed,        thread_layout_transposed);}

Tensor Partitions

There are three major partitions in CuTe, inner-partition, outer-partition, and thread-value partition.

Inner-partition has been performed previously, where we divide the input matrix and the output matrix into tiles. Usually inner-partition is performed at the CUDA thread block level that distributes the large problems into smaller problems that can be solved by a single CUDA thread block.

Outer-partition is usually performed at the CUDA thread level that distributes the smaller problems into even smaller problems that can be solved by a single CUDA thread. There is a different between inner-partition and outer-partition, without understanding which the implementation can work correctly.

Suppose we have a CuTe layout $(8, 4) : (4, 1)$ and a tile layout $(4, 2) : (2, 1)$. Inner-partition will result in a layout of shape $\left((4, 2), \frac{8}{4}, \frac{4}{2}\right) = \left((4, 2), 2, 2\right)$, and outer-partition will result in a a layout of shape $\left(\left(\frac{8}{4}, \frac{4}{2}\right), 4, 2\right) = \left((2, 2), 4, 2\right)$. Assuming the partitions are accessed using the last two modes layout, the inner-partition layout has 4 partitions whereas the out-partition layout has 8 partitions. The starting coordinates of inner-partition and outer-partition are also different. In this case, the starting coordinates of inner-partition is $(0, 0)$, $(4, 0)$, $(0, 2)$, and $(4, 2)$, whereas the starting coordinates of outer-partition is $(0, 0)$, $(1, 0)$, $(2, 0)$, $(3, 0)$, $(0, 1)$, $(1, 1)$, $(2, 1)$, and $(3, 1)$. Outer-partition is usually performed at the CUDA thread level because all the consecutive threads in a warp, if accessing a piece of contiguous data synergistically especially on global memory, can have a better performance because of the CUDA coalesced memory access.

In each partition, the partition tensor will follow the layout algebra and apply the correct strides to access the data during iteration.

12345678910111213

auto global_tile_src{tensor_src(cute::make_coord(cute::_, cute::_),                                blockIdx.y,                                blockIdx.x)}; // (TileSizeY, TileSizeX)auto global_tile_dst_transposed{    tensor_dst_transposed(cute::make_coord(cute::_, cute::_), blockIdx.y,                          blockIdx.x)}; // (TileSizeY, TileSizeX)auto thread_global_tile_src{cute::local_partition(    global_tile_src, ThreadLayout{},    threadIdx.x)}; // (ThreadValueSizeY, ThreadValueSizeX)auto thread_global_tile_dst_transposed{cute::local_partition(    global_tile_dst_transposed, ThreadLayout{},    threadIdx.x)}; // (ThreadValueSizeY, ThreadValueSizeX)

Predicates and Boundary Checking

CUDA memory access boundary checking is critical in a CUDA kernel in practice, when the problem distribution is not perfect. In CuTe, CUDA memory access boundary checking is performed via predicates. In our particular, we could query the matrix sizes $M$ and $N$ from the strides of the tiled tensor. It’s also more common to just pass these two values to the CUDA kernel.

During the iteration of the CuTe tensor, the iterator has to know its 2D coordinate and check if the element it is about to access is within the boundary. So we will create a 2D identity tensor (it’s an 2D array of tuples though) whose shape is exactly the same as the partition tensor from the global memory. The 2D identity tensor takes a 2D coordinate as input and produces the same 2D coordinate as output. If the partition tensor abd the identity tensor are iterated together, the iterator could get the information of its current coordinate within the partition tensor, making boundary checking possible. At the CUDA thread level, the 2D identity tensor is further partitioned into a 2D thread identity tensor according the same thread layout that is used for partitioning the data. Then the predicates used for accessing the partitioned input tensor and the output tensor can be prepared using the 2D coordinates of the iterator, the partitioned tensor index, the partitioned tensor and the original full tensor shape information.

123456789101112131415161718192021222324252627282930313233

// A 2D array of tuples that maps (x, y) to (x, y).auto const identity_tensor{cute::make_identity_tensor(cute::make_shape(    cute::size<0>(global_tile_src), cute::size<1>(global_tile_src)))};auto const thread_identity_tensor{    cute::local_partition(identity_tensor, ThreadLayout{}, threadIdx.x)};auto fragment{cute::make_tensor_like(thread_global_tile_src)};auto predicator{cute::make_tensor<bool>(    cute::make_shape(cute::size<0>(fragment), cute::size<1>(fragment)))};auto const num_max_columns{cute::stride<0>(global_tile_src)};auto const num_max_rows{cute::stride<1>(global_tile_dst_transposed)};constexpr auto global_tile_columns{cute::size<1>(global_tile_src)};constexpr auto global_tile_rows{cute::size<0>(global_tile_src)};CUTE_UNROLLfor (unsigned int i{0}; i < cute::size<0>(predicator); ++i){    CUTE_UNROLL    for (unsigned int j{0}; j < cute::size<1>(predicator); ++j)    {        auto const thread_identity{thread_identity_tensor(i, j)};        bool const is_row_in_bound{cute::get<0>(thread_identity) +                                       blockIdx.y * global_tile_rows <                                   num_max_rows};        bool const is_column_in_bound{cute::get<1>(thread_identity) +                                          blockIdx.x * global_tile_columns <                                      num_max_columns};        predicator(i, j) = is_row_in_bound && is_column_in_bound;    }}cute::copy_if(predicator, thread_global_tile_src, fragment);cute::copy_if(predicator, fragment, thread_global_tile_dst_transposed);

Using predicates and performing boundary checking can have degrade CUDA kernel performance, because the warp instruction, such as load from global memory, has to stall before all the predicates from all the threads in the warp are evaluated. To accelerate the computing the most, usually specialized kernels are used for each of the problem configurations and boundary checking are eliminated from the CUDA kernel.

So instead of using cute::copy_if, cute::copy should be used.

12

cute::copy(thread_global_tile_src, fragment);cute::copy(fragment, thread_global_tile_dst_transposed);

CuTe Matrix Transpose Using Shared Memory

In the CuTe matrix transpose CUDA kernel implementation using shared memory, we will perform strided memory reads and writes on shared memory instead of global memory. Using shared memory naively will result in shared memory bank conflicts when performing strided memory reads or writes on shared memory, which will degrade the performance. To mitigate the shared memory bank conflicts, we will also perform special optimizations, such as shared memory padding and swizzling.

CuTe Matrix Transpose Using Shared Memory Implementation

The CuTe matrix transpose using shared memory CUDA kernel implementation could also be found from my CUTLASS Examples GitHub repository.

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931

#include <iomanip>#include <iostream>#include <cuda_runtime.h>#include <cute/tensor.hpp>#include "cute_matrix_transpose.hpp"template <class TensorSrc, class TensorDst, class SharedMemoryLayoutSrc,          class SharedMemoryLayoutDst, class ThreadLayoutSrc,          class ThreadLayoutDst>__global__ void static matrix_transpose_shared_memory(    TensorSrc tensor_src, TensorDst tensor_dst, SharedMemoryLayoutSrc,    SharedMemoryLayoutDst, ThreadLayoutSrc, ThreadLayoutDst){    using Element = typename TensorSrc::value_type;    CUTE_STATIC_ASSERT_V(cute::size(SharedMemoryLayoutSrc{}) ==                             cute::size(SharedMemoryLayoutDst{}),                         "SharedMemoryLayoutSrc and SharedMemoryLayoutDst "                         "must have the same size.");    __shared__ Element shared_memory[cute::cosize(SharedMemoryLayoutSrc{})];    auto tensor_cache_src{cute::make_tensor(cute::make_smem_ptr(shared_memory),                                            SharedMemoryLayoutSrc{})};    auto tensor_cache_dst{cute::make_tensor(cute::make_smem_ptr(shared_memory),                                            SharedMemoryLayoutDst{})};    auto global_tile_src{tensor_src(cute::make_coord(cute::_, cute::_),                                    blockIdx.y,                                    blockIdx.x)}; // (TileSizeY, TileSizeX)    auto global_tile_dst{tensor_dst(cute::make_coord(cute::_, cute::_),                                    blockIdx.y,                                    blockIdx.x)}; // (TileSizeY, TileSizeX)    auto thread_global_tile_src{cute::local_partition(        global_tile_src, ThreadLayoutSrc{},        threadIdx.x)}; // (ThreadValueSizeY, ThreadValueSizeX)    auto thread_global_tile_dst{cute::local_partition(        global_tile_dst, ThreadLayoutDst{},        threadIdx.x)}; // (ThreadValueSizeX, ThreadValueSizeY)    auto thread_shared_tile_src{cute::local_partition(        tensor_cache_src, ThreadLayoutSrc{},        threadIdx.x)}; // (ThreadValueSizeY, ThreadValueSizeX)    auto thread_shared_tile_dst{cute::local_partition(        tensor_cache_dst, ThreadLayoutDst{},        threadIdx.x)}; // (ThreadValueSizeX, ThreadValueSizeY)    // A 2D array of tuples that maps (x, y) to (x, y).    auto const identity_tensor_src{cute::make_identity_tensor(cute::make_shape(        cute::size<0>(global_tile_src), cute::size<1>(global_tile_src)))};    auto const thread_identity_tensor_src{cute::local_partition(        identity_tensor_src, ThreadLayoutSrc{}, threadIdx.x)};    auto predicator_src{cute::make_tensor<bool>(        cute::make_shape(cute::size<0>(thread_global_tile_src),                         cute::size<1>(thread_global_tile_src)))};    auto const identity_tensor_dst{cute::make_identity_tensor(cute::make_shape(        cute::size<0>(global_tile_dst), cute::size<1>(global_tile_dst)))};    auto const thread_identity_tensor_dst{cute::local_partition(        identity_tensor_dst, ThreadLayoutDst{}, threadIdx.x)};    auto predicator_dst{cute::make_tensor<bool>(        cute::make_shape(cute::size<0>(thread_global_tile_dst),                         cute::size<1>(thread_global_tile_dst)))};    auto const num_max_columns{cute::stride<0>(global_tile_src)};    auto const num_max_rows{cute::stride<1>(global_tile_dst)};    constexpr auto global_tile_columns{cute::size<1>(global_tile_src)};    constexpr auto global_tile_rows{cute::size<0>(global_tile_src)};    CUTE_UNROLL    for (unsigned int i{0}; i < cute::size<0>(predicator_src); ++i)    {        CUTE_UNROLL        for (unsigned int j{0}; j < cute::size<1>(predicator_src); ++j)        {            auto const thread_identity{thread_identity_tensor_src(i, j)};            bool const is_row_in_bound{cute::get<0>(thread_identity) +                                           blockIdx.y * global_tile_rows <                                       num_max_rows};            bool const is_column_in_bound{cute::get<1>(thread_identity) +                                              blockIdx.x * global_tile_columns <                                          num_max_columns};            predicator_src(i, j) = is_row_in_bound && is_column_in_bound;        }    }    CUTE_UNROLL    for (unsigned int i{0}; i < cute::size<0>(predicator_dst); ++i)    {        CUTE_UNROLL        for (unsigned int j{0}; j < cute::size<1>(predicator_dst); ++j)        {            auto const thread_identity{thread_identity_tensor_dst(i, j)};            bool const is_row_in_bound{cute::get<0>(thread_identity) +                                           blockIdx.y * global_tile_rows <                                       num_max_rows};            bool const is_column_in_bound{cute::get<1>(thread_identity) +                                              blockIdx.x * global_tile_columns <                                          num_max_columns};            predicator_dst(i, j) = is_row_in_bound && is_column_in_bound;        }    }    cute::copy_if(predicator_src, thread_global_tile_src,                  thread_shared_tile_src);    cute::cp_async_fence();    cute::cp_async_wait<0>();    __syncthreads();    cute::copy_if(predicator_dst, thread_shared_tile_dst,                  thread_global_tile_dst);}template <class TensorSrc, class TensorDst, class SharedMemoryLayoutSrc,          class SharedMemoryLayoutDst, class ThreadLayoutSrc,          class ThreadLayoutDst, class VectorLayout>__global__ void static matrix_transpose_shared_memory_vectorized(    TensorSrc tensor_src, TensorDst tensor_dst, SharedMemoryLayoutSrc,    SharedMemoryLayoutDst, ThreadLayoutSrc, ThreadLayoutDst, VectorLayout){    using Element = typename TensorSrc::value_type;    CUTE_STATIC_ASSERT_V(cute::size(SharedMemoryLayoutSrc{}) ==                             cute::size(SharedMemoryLayoutDst{}),                         "SharedMemoryLayoutSrc and SharedMemoryLayoutDst "                         "must have the same size.");    __shared__ Element shared_memory[cute::cosize(SharedMemoryLayoutSrc{})];    auto tensor_cache_src{cute::make_tensor(cute::make_smem_ptr(shared_memory),                                            SharedMemoryLayoutSrc{})};    auto tensor_cache_dst{cute::make_tensor(cute::make_smem_ptr(shared_memory),                                            SharedMemoryLayoutDst{})};    auto global_tile_src{tensor_src(cute::make_coord(cute::_, cute::_),                                    blockIdx.y,                                    blockIdx.x)}; // (TileSizeY, TileSizeX)    auto global_tile_dst{tensor_dst(cute::make_coord(cute::_, cute::_),                                    blockIdx.y,                                    blockIdx.x)}; // (TileSizeY, TileSizeX)    using AccessType =        cutlass::AlignedArray<Element, cute::size(VectorLayout{})>;    using CopyAtom = cute::Copy_Atom<cute::UniversalCopy<AccessType>, Element>;    auto tiled_copy_src{        cute::make_tiled_copy(CopyAtom{}, ThreadLayoutSrc{}, VectorLayout{})};    auto thread_copy_src{tiled_copy_src.get_thread_slice(threadIdx.x)};    auto thread_global_tile_src{thread_copy_src.partition_S(        global_tile_src)}; // (CopyAtomShape, NumCopyTile)    auto thread_shared_tile_src{thread_copy_src.partition_D(        tensor_cache_src)}; // (CopyAtomShape, NumCopyTile)    auto thread_global_tile_dst{cute::local_partition(        global_tile_dst, ThreadLayoutDst{},        threadIdx.x)}; // (ThreadValueSizeX, ThreadValueSizeY)    auto thread_shared_tile_dst{cute::local_partition(        tensor_cache_dst, ThreadLayoutDst{},        threadIdx.x)}; // (ThreadValueSizeX, ThreadValueSizeY)    auto const num_max_columns{cute::stride<0>(global_tile_src)};    auto const num_max_rows{cute::stride<1>(global_tile_dst)};    constexpr auto global_tile_columns{cute::size<1>(global_tile_src)};    constexpr auto global_tile_rows{cute::size<0>(global_tile_src)};    // A 2D array of tuples that maps (x, y) to (x, y).    auto const identity_tensor_src{cute::make_identity_tensor(cute::make_shape(        cute::size<0>(global_tile_src), cute::size<1>(global_tile_src)))};    auto thread_identity_tensor_src{thread_copy_src.partition_S(        identity_tensor_src)}; // (CopyAtomShape, NumCopyTile)    auto predicator_src{cute::make_tensor<bool>(        cute::make_shape(cute::size<1>(thread_global_tile_src),                         cute::size<2>(thread_global_tile_src)))};    CUTE_UNROLL    for (unsigned int i{0}; i < cute::size<0>(predicator_src); ++i)    {        CUTE_UNROLL        for (unsigned int j{0}; j < cute::size<1>(predicator_src); ++j)        {            auto const thread_identity{thread_identity_tensor_src(0, i, j)};            bool const is_row_in_bound{cute::get<0>(thread_identity) +                                           blockIdx.y * global_tile_rows <                                       num_max_rows};            bool const is_column_in_bound{cute::get<1>(thread_identity) +                                              blockIdx.x * global_tile_columns <                                          num_max_columns};            predicator_src(i, j) = is_row_in_bound && is_column_in_bound;        }    }    auto const identity_tensor_dst{cute::make_identity_tensor(cute::make_shape(        cute::size<0>(global_tile_dst), cute::size<1>(global_tile_dst)))};    auto const thread_identity_tensor_dst{cute::local_partition(        identity_tensor_dst, ThreadLayoutDst{}, threadIdx.x)};    auto predicator_dst{cute::make_tensor<bool>(        cute::make_shape(cute::size<0>(thread_global_tile_dst),                         cute::size<1>(thread_global_tile_dst)))};    CUTE_UNROLL    for (unsigned int i{0}; i < cute::size<0>(predicator_dst); ++i)    {        CUTE_UNROLL        for (unsigned int j{0}; j < cute::size<1>(predicator_dst); ++j)        {            auto const thread_identity{thread_identity_tensor_dst(i, j)};            bool const is_row_in_bound{cute::get<0>(thread_identity) +                                           blockIdx.y * global_tile_rows <                                       num_max_rows};            bool const is_column_in_bound{cute::get<1>(thread_identity) +                                              blockIdx.x * global_tile_columns <                                          num_max_columns};            predicator_dst(i, j) = is_row_in_bound && is_column_in_bound;        }    }    cute::copy_if(tiled_copy_src, predicator_src, thread_global_tile_src,                  thread_shared_tile_src);    cute::cp_async_fence();    cute::cp_async_wait<0>();    __syncthreads();    cute::copy_if(predicator_dst, thread_shared_tile_dst,                  thread_global_tile_dst);}enum class SharedMemoryBankConflictAccessMode{    Read,    Write};template <typename T>static cudaError_t launch_matrix_transpose_shared_memory_bank_conflict_base(    T const* input_matrix, T* output_matrix, unsigned int M, unsigned int N,    SharedMemoryBankConflictAccessMode bank_conflict_access_mode,    cudaStream_t stream){    auto const tensor_shape{cute::make_shape(M, N)};    auto const tensor_shape_transposed{cute::make_shape(N, M)};    // Input matrix: row-major M x N matrix.    auto const global_memory_layout_src{cute::make_layout(        tensor_shape, cute::GenRowMajor{})}; // (M, N) : (N, 1)    // Output matrix: row-major N x M matrix.    auto const global_memory_layout_dst{cute::make_layout(        tensor_shape_transposed, cute::GenRowMajor{})}; // (N, M) : (M, 1)    // Same output matrix, but different view: column-major M x N matrix.    auto const global_memory_layout_dst_transposed{cute::make_layout(        tensor_shape, cute::GenColMajor{})}; // (M, N) : (1, M)    auto const tensor_src{cute::make_tensor(cute::make_gmem_ptr(input_matrix),                                            global_memory_layout_src)};    auto const tensor_dst{cute::make_tensor(cute::make_gmem_ptr(output_matrix),                                            global_memory_layout_dst)};    auto const tensor_dst_transposed{        cute::make_tensor(cute::make_gmem_ptr(output_matrix),                          global_memory_layout_dst_transposed)};    using TileSizeX = cute::Int<128>; // bN    using TileSizeY = cute::Int<32>;  // bM    constexpr auto block_shape{cute::make_shape(TileSizeY{}, TileSizeX{})};    constexpr auto block_shape_transposed{        cute::make_shape(TileSizeX{}, TileSizeY{})};    auto const shared_memory_layout_src{cute::make_layout(        block_shape, cute::GenRowMajor{})}; // (bM, bN) : (bN, 1)    auto const shared_memory_layout_dst{cute::make_layout(        block_shape_transposed, cute::GenRowMajor{})}; // (bN, bM) : (bM, 1)    auto const shared_memory_layout_dst_transposed{cute::make_layout(        block_shape, cute::GenColMajor{})}; // (bM, bN) : (1, bM)    auto const tiled_tensor_src{cute::tiled_divide(        tensor_src, block_shape)}; // ((TileSizeY, TileSizeX), M /                                   // TileSizeY, N / TileSizeX)    auto const tiled_tensor_dst{cute::tiled_divide(        tensor_dst, block_shape_transposed)}; // ((TileSizeX, TileSizeY), N                                              // / TileSizeX, M / TileSizeY)    auto const tiled_tensor_dst_transposed{cute::tiled_divide(        tensor_dst_transposed, block_shape)}; // ((TileSizeY, TileSizeX), M                                              // / TileSizeY, N / TileSizeX)    using ThreadBlockSizeX = cute::Int<32>; // tN    using ThreadBlockSizeY = cute::Int<8>;  // tM    CUTE_STATIC_ASSERT_V(TileSizeX{} % ThreadBlockSizeX{} == cute::Int<0>{},                         "TileSizeX must be divisible by ThreadBlockSizeX");    CUTE_STATIC_ASSERT_V(TileSizeY{} % ThreadBlockSizeY{} == cute::Int<0>{},                         "TileSizeY must be divisible by ThreadBlockSizeY");    constexpr auto thread_block_shape{        cute::make_shape(ThreadBlockSizeY{}, ThreadBlockSizeX{})};    constexpr auto thread_block_shape_transposed{        cute::make_shape(ThreadBlockSizeX{}, ThreadBlockSizeY{})};    constexpr auto thread_layout{        cute::make_layout(thread_block_shape, cute::GenRowMajor{})};    constexpr auto thread_layout_transposed{        cute::make_layout(thread_block_shape_transposed, cute::GenColMajor{})};    dim3 const grid_dim{cute::size<2>(tiled_tensor_src),                        cute::size<1>(tiled_tensor_src)};    dim3 const thread_dim{ThreadBlockSizeX::value * ThreadBlockSizeY::value};    if (bank_conflict_access_mode == SharedMemoryBankConflictAccessMode::Read)    {        matrix_transpose_shared_memory<<<grid_dim, thread_dim, 0, stream>>>(            tiled_tensor_src, tiled_tensor_dst_transposed,            shared_memory_layout_src, shared_memory_layout_src, thread_layout,            thread_layout_transposed);    }    else    {        matrix_transpose_shared_memory<<<grid_dim, thread_dim, 0, stream>>>(            tiled_tensor_src, tiled_tensor_dst_transposed,            shared_memory_layout_dst_transposed,            shared_memory_layout_dst_transposed, thread_layout,            thread_layout_transposed);    }    return cudaGetLastError();}template <typename T>static cudaError_tlaunch_matrix_transpose_shared_memory_vectorized_bank_conflict_base(    T const* input_matrix, T* output_matrix, unsigned int M, unsigned int N,    SharedMemoryBankConflictAccessMode bank_conflict_access_mode,    cudaStream_t stream){    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)};    if (N % NUM_VECTOR_ELEMENTS != 0)    {        return cudaErrorInvalidValue;    }    auto const tensor_shape{cute::make_shape(M, N)};    auto const tensor_shape_transposed{cute::make_shape(N, M)};    // Input matrix: row-major M x N matrix.    auto const global_memory_layout_src{cute::make_layout(        tensor_shape, cute::GenRowMajor{})}; // (M, N) : (N, 1)    // Output matrix: row-major N x M matrix.    auto const global_memory_layout_dst{cute::make_layout(        tensor_shape_transposed, cute::GenRowMajor{})}; // (N, M) : (M, 1)    // Same output matrix, but different view: column-major M x N matrix.    auto const global_memory_layout_dst_transposed{cute::make_layout(        tensor_shape, cute::GenColMajor{})}; // (M, N) : (1, M)    auto const tensor_src{cute::make_tensor(cute::make_gmem_ptr(input_matrix),                                            global_memory_layout_src)};    auto const tensor_dst{cute::make_tensor(cute::make_gmem_ptr(output_matrix),                                            global_memory_layout_dst)};    auto const tensor_dst_transposed{        cute::make_tensor(cute::make_gmem_ptr(output_matrix),                          global_memory_layout_dst_transposed)};    using TileSizeX = cute::Int<128>; // bN    using TileSizeY = cute::Int<32>;  // bM    constexpr auto block_shape{cute::make_shape(TileSizeY{}, TileSizeX{})};    constexpr auto block_shape_transposed{        cute::make_shape(TileSizeX{}, TileSizeY{})};    auto const shared_memory_layout_src{cute::make_layout(        block_shape, cute::GenRowMajor{})}; // (bM, bN) : (bN, 1)    auto const shared_memory_layout_dst{cute::make_layout(        block_shape_transposed, cute::GenRowMajor{})}; // (bN, bM) : (bM, 1)    auto const shared_memory_layout_dst_transposed{cute::make_layout(        block_shape, cute::GenColMajor{})}; // (bM, bN) : (1, bM)    auto const tiled_tensor_src{cute::tiled_divide(        tensor_src, block_shape)}; // ((TileSizeY, TileSizeX), M /                                   // TileSizeY, N / TileSizeX)    auto const tiled_tensor_dst{cute::tiled_divide(        tensor_dst, block_shape_transposed)}; // ((TileSizeX, TileSizeY), N                                              // / TileSizeX, M / TileSizeY)    auto const tiled_tensor_dst_transposed{cute::tiled_divide(        tensor_dst_transposed, block_shape)}; // ((TileSizeY, TileSizeX), M                                              // / TileSizeY, N / TileSizeX)    using ThreadBlockSizeX = cute::Int<32>; // tN    using ThreadBlockSizeY = cute::Int<8>;  // tM    CUTE_STATIC_ASSERT_V(TileSizeX{} % ThreadBlockSizeX{} == cute::Int<0>{},                         "TileSizeX must be divisible by ThreadBlockSizeX");    CUTE_STATIC_ASSERT_V(TileSizeY{} % ThreadBlockSizeY{} == cute::Int<0>{},                         "TileSizeY must be divisible by ThreadBlockSizeY");    constexpr auto thread_block_shape{        cute::make_shape(ThreadBlockSizeY{}, ThreadBlockSizeX{})};    constexpr auto thread_block_shape_transposed{        cute::make_shape(ThreadBlockSizeX{}, ThreadBlockSizeY{})};    constexpr auto thread_layout{        cute::make_layout(thread_block_shape, cute::GenRowMajor{})};    constexpr auto thread_layout_transposed{        cute::make_layout(thread_block_shape_transposed, cute::GenColMajor{})};    using VECTOR_SIZE_X = cute::Int<NUM_VECTOR_ELEMENTS>;    constexpr auto vector_shape{        cute::make_shape(cute::Int<1>{}, VECTOR_SIZE_X{})};    // Copy atom vector layout.    constexpr auto vector_layout{        cute::make_layout(vector_shape, cute::GenRowMajor{})};    dim3 const grid_dim{cute::size<2>(tiled_tensor_src),                        cute::size<1>(tiled_tensor_src)};    dim3 const thread_dim{ThreadBlockSizeX::value * ThreadBlockSizeY::value};    if (bank_conflict_access_mode == SharedMemoryBankConflictAccessMode::Read)    {        matrix_transpose_shared_memory_vectorized<<<grid_dim, thread_dim, 0,                                                    stream>>>(            tiled_tensor_src, tiled_tensor_dst_transposed,            shared_memory_layout_src, shared_memory_layout_src, thread_layout,            thread_layout_transposed, vector_layout);    }    else    {        return cudaErrorInvalidValue;    }    return cudaGetLastError();}template <typename T>cudaError_t launch_matrix_transpose_shared_memory_bank_conflict_read(    T const* input_matrix, T* output_matrix, unsigned int M, unsigned int N,    cudaStream_t stream){    return launch_matrix_transpose_shared_memory_bank_conflict_base(        input_matrix, output_matrix, M, N,        SharedMemoryBankConflictAccessMode::Read, stream);}template <typename T>cudaError_t launch_matrix_transpose_shared_memory_bank_conflict_write(    T const* input_matrix, T* output_matrix, unsigned int M, unsigned int N,    cudaStream_t stream){    return launch_matrix_transpose_shared_memory_bank_conflict_base(        input_matrix, output_matrix, M, N,        SharedMemoryBankConflictAccessMode::Write, stream);}template <typename T>cudaError_t launch_matrix_transpose_shared_memory_vectorized_bank_conflict_read(    T const* input_matrix, T* output_matrix, unsigned int M, unsigned int N,    cudaStream_t stream){    return launch_matrix_transpose_shared_memory_vectorized_bank_conflict_base<        T>(input_matrix, output_matrix, M, N,           SharedMemoryBankConflictAccessMode::Read, stream);}template <typename T>static cudaError_t launch_matrix_transpose_shared_memory_padded(    T const* input_matrix, T* output_matrix, unsigned int M, unsigned int N,    cudaStream_t stream){    auto const tensor_shape{cute::make_shape(M, N)};    auto const tensor_shape_transposed{cute::make_shape(N, M)};    // Input matrix: row-major M x N matrix.    auto const global_memory_layout_src{cute::make_layout(        tensor_shape, cute::GenRowMajor{})}; // (M, N) : (N, 1)    // Output matrix: row-major N x M matrix.    auto const global_memory_layout_dst{cute::make_layout(        tensor_shape_transposed, cute::GenRowMajor{})}; // (N, M) : (M, 1)    // Same output matrix, but different view: column-major M x N matrix.    auto const global_memory_layout_dst_transposed{cute::make_layout(        tensor_shape, cute::GenColMajor{})}; // (M, N) : (1, M)    auto const tensor_src{cute::make_tensor(cute::make_gmem_ptr(input_matrix),                                            global_memory_layout_src)};    auto const tensor_dst{cute::make_tensor(cute::make_gmem_ptr(output_matrix),                                            global_memory_layout_dst)};    auto const tensor_dst_transposed{        cute::make_tensor(cute::make_gmem_ptr(output_matrix),                          global_memory_layout_dst_transposed)};    using TileSizeX = cute::Int<64>;          // bN    using TILE_SIZE_X_PADDED = cute::Int<65>; // bN + 1    using TileSizeY = cute::Int<32>;          // bM    constexpr auto block_shape{cute::make_shape(TileSizeY{}, TileSizeX{})};    constexpr auto block_shape_transposed{        cute::make_shape(TileSizeX{}, TileSizeY{})};    auto const shared_memory_layout_src{cute::make_layout(        block_shape, cute::GenRowMajor{})}; // (bM, bN) : (bN, 1)    auto const shared_memory_layout_src_padded{cute::make_layout(        block_shape,        cute::make_stride(TILE_SIZE_X_PADDED{},                          cute::Int<1>{}))}; // (bM, bN) : (bN + 1, 1)    auto const shared_memory_layout_dst{cute::make_layout(        block_shape_transposed, cute::GenRowMajor{})}; // (bN, bM) : (bM, 1)    auto const shared_memory_layout_dst_transposed{cute::make_layout(        block_shape, cute::GenColMajor{})}; // (bM, bN) : (1, bM)    auto const tiled_tensor_src{cute::tiled_divide(        tensor_src, block_shape)}; // ((TileSizeY, TileSizeX), M /                                   // TileSizeY, N / TileSizeX)    auto const tiled_tensor_dst{cute::tiled_divide(        tensor_dst, block_shape_transposed)}; // ((TileSizeX, TileSizeY), N                                              // / TileSizeX, M / TileSizeY)    auto const tiled_tensor_dst_transposed{cute::tiled_divide(        tensor_dst_transposed, block_shape)}; // ((TileSizeY, TileSizeX), M                                              // / TileSizeY, N / TileSizeX)    using ThreadBlockSizeX = cute::Int<32>; // tN    using ThreadBlockSizeY = cute::Int<8>;  // tM    CUTE_STATIC_ASSERT_V(TileSizeX{} % ThreadBlockSizeX{} == cute::Int<0>{},                         "TileSizeX must be divisible by ThreadBlockSizeX");    CUTE_STATIC_ASSERT_V(TileSizeY{} % ThreadBlockSizeY{} == cute::Int<0>{},                         "TileSizeY must be divisible by ThreadBlockSizeY");    constexpr auto thread_block_shape{        cute::make_shape(ThreadBlockSizeY{}, ThreadBlockSizeX{})};    constexpr auto thread_block_shape_transposed{        cute::make_shape(ThreadBlockSizeX{}, ThreadBlockSizeY{})};    constexpr auto thread_layout{        cute::make_layout(thread_block_shape, cute::GenRowMajor{})};    constexpr auto thread_layout_transposed{        cute::make_layout(thread_block_shape_transposed, cute::GenColMajor{})};    dim3 const grid_dim{cute::size<2>(tiled_tensor_src),                        cute::size<1>(tiled_tensor_src)};    dim3 const thread_dim{ThreadBlockSizeX::value * ThreadBlockSizeY::value};    matrix_transpose_shared_memory<<<grid_dim, thread_dim, 0, stream>>>(        tiled_tensor_src, tiled_tensor_dst_transposed,        shared_memory_layout_src_padded, shared_memory_layout_src_padded,        thread_layout, thread_layout_transposed);    return cudaGetLastError();}template <typename T>static cudaError_t launch_matrix_transpose_shared_memory_vectorized_padded(    T const* input_matrix, T* output_matrix, unsigned int M, unsigned int N,    cudaStream_t stream){    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)};    if (N % NUM_VECTOR_ELEMENTS != 0)    {        return cudaErrorInvalidValue;    }    auto const tensor_shape{cute::make_shape(M, N)};    auto const tensor_shape_transposed{cute::make_shape(N, M)};    // Input matrix: row-major M x N matrix.    auto const global_memory_layout_src{cute::make_layout(        tensor_shape, cute::GenRowMajor{})}; // (M, N) : (N, 1)    // Output matrix: row-major N x M matrix.    auto const global_memory_layout_dst{cute::make_layout(        tensor_shape_transposed, cute::GenRowMajor{})}; // (N, M) : (M, 1)    // Same output matrix, but different view: column-major M x N matrix.    auto const global_memory_layout_dst_transposed{cute::make_layout(        tensor_shape, cute::GenColMajor{})}; // (M, N) : (1, M)    auto const tensor_src{cute::make_tensor(cute::make_gmem_ptr(input_matrix),                                            global_memory_layout_src)};    auto const tensor_dst{cute::make_tensor(cute::make_gmem_ptr(output_matrix),                                            global_memory_layout_dst)};    auto const tensor_dst_transposed{        cute::make_tensor(cute::make_gmem_ptr(output_matrix),                          global_memory_layout_dst_transposed)};    using TileSizeX = cute::Int<128>; // bN    // Such padding is necessary for the byte alignment of the vectorized    // access. However, the shared memory bank conflict mitigation can be    // compromised.    using TILE_SIZE_X_PADDED =        cute::Int<128 + NUM_VECTOR_ELEMENTS>; // bN + NUM_VECTOR_ELEMENTS    using TileSizeY = cute::Int<32>;          // bM    constexpr auto block_shape{cute::make_shape(TileSizeY{}, TileSizeX{})};    constexpr auto block_shape_transposed{        cute::make_shape(TileSizeX{}, TileSizeY{})};    auto const shared_memory_layout_src{cute::make_layout(        block_shape, cute::GenRowMajor{})}; // (bM, bN) : (bN, 1)    auto const shared_memory_layout_src_padded{cute::make_layout(        block_shape,        cute::make_stride(TILE_SIZE_X_PADDED{},                          cute::Int<1>{}))}; // (bM, bN) : (bN + 1, 1)    auto const shared_memory_layout_dst{cute::make_layout(        block_shape_transposed, cute::GenRowMajor{})}; // (bN, bM) : (bM, 1)    auto const shared_memory_layout_dst_transposed{cute::make_layout(        block_shape, cute::GenColMajor{})}; // (bM, bN) : (1, bM)    auto const tiled_tensor_src{cute::tiled_divide(        tensor_src, block_shape)}; // ((TileSizeY, TileSizeX), M /                                   // TileSizeY, N / TileSizeX)    auto const tiled_tensor_dst{cute::tiled_divide(        tensor_dst, block_shape_transposed)}; // ((TileSizeX, TileSizeY), N                                              // / TileSizeX, M / TileSizeY)    auto const tiled_tensor_dst_transposed{cute::tiled_divide(        tensor_dst_transposed, block_shape)}; // ((TileSizeY, TileSizeX), M                                              // / TileSizeY, N / TileSizeX)    using ThreadBlockSizeX = cute::Int<32>; // tN    using ThreadBlockSizeY = cute::Int<8>;  // tM    CUTE_STATIC_ASSERT_V(TileSizeX{} % ThreadBlockSizeX{} == cute::Int<0>{},                         "TileSizeX must be divisible by ThreadBlockSizeX");    CUTE_STATIC_ASSERT_V(TileSizeY{} % ThreadBlockSizeY{} == cute::Int<0>{},                         "TileSizeY must be divisible by ThreadBlockSizeY");    constexpr auto thread_block_shape{        cute::make_shape(ThreadBlockSizeY{}, ThreadBlockSizeX{})};    constexpr auto thread_block_shape_transposed{        cute::make_shape(ThreadBlockSizeX{}, ThreadBlockSizeY{})};    constexpr auto thread_layout{        cute::make_layout(thread_block_shape, cute::GenRowMajor{})};    constexpr auto thread_layout_transposed{        cute::make_layout(thread_block_shape_transposed, cute::GenColMajor{})};    dim3 const grid_dim{cute::size<2>(tiled_tensor_src),                        cute::size<1>(tiled_tensor_src)};    dim3 const thread_dim{ThreadBlockSizeX::value * ThreadBlockSizeY::value};    using VECTOR_SIZE_X = cute::Int<NUM_VECTOR_ELEMENTS>;    constexpr auto vector_shape{        cute::make_shape(cute::Int<1>{}, VECTOR_SIZE_X{})};    // Copy atom vector layout.    constexpr auto vector_layout{        cute::make_layout(vector_shape, cute::GenRowMajor{})};    matrix_transpose_shared_memory_vectorized<<<grid_dim, thread_dim, 0,                                                stream>>>(        tiled_tensor_src, tiled_tensor_dst_transposed,        shared_memory_layout_src_padded, shared_memory_layout_src_padded,        thread_layout, thread_layout_transposed, vector_layout);    return cudaGetLastError();}template <class SHARED_MEMORY_LAYOUT>static voidprint_shared_memory_bank_ids(SHARED_MEMORY_LAYOUT shared_memory_layout){    // Print the shared memory bank ids.    for (unsigned int i{0}; i < cute::size<0>(shared_memory_layout); ++i)    {        for (unsigned int j{0}; j < cute::size<1>(shared_memory_layout); ++j)        {            std::cout << std::setw(2) << shared_memory_layout(i, j) % 32 << " ";        }        std::cout << std::endl;    }}constexpr int constexpr_log2(int n){    return ((n < 2) ? 0 : 1 + constexpr_log2(n / 2));}template <typename T>static cudaError_t launch_matrix_transpose_shared_memory_swizzled(    T const* input_matrix, T* output_matrix, unsigned int M, unsigned int N,    cudaStream_t stream){    auto const tensor_shape{cute::make_shape(M, N)};    auto const tensor_shape_transposed{cute::make_shape(N, M)};    // Input matrix: row-major M x N matrix.    auto const global_memory_layout_src{cute::make_layout(        tensor_shape, cute::GenRowMajor{})}; // (M, N) : (N, 1)    // Output matrix: row-major N x M matrix.    auto const global_memory_layout_dst{cute::make_layout(        tensor_shape_transposed, cute::GenRowMajor{})}; // (N, M) : (M, 1)    // Same output matrix, but different view: column-major M x N matrix.    auto const global_memory_layout_dst_transposed{cute::make_layout(        tensor_shape, cute::GenColMajor{})}; // (M, N) : (1, M)    auto const tensor_src{cute::make_tensor(cute::make_gmem_ptr(input_matrix),                                            global_memory_layout_src)};    auto const tensor_dst{cute::make_tensor(cute::make_gmem_ptr(output_matrix),                                            global_memory_layout_dst)};    auto const tensor_dst_transposed{        cute::make_tensor(cute::make_gmem_ptr(output_matrix),                          global_memory_layout_dst_transposed)};    using TileSizeX = cute::Int<64>; // bN    using TileSizeY = cute::Int<32>; // bM    constexpr int NUM_BASE_BITS{constexpr_log2(1)};    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};    constexpr auto block_shape{cute::make_shape(TileSizeY{}, TileSizeX{})};    constexpr auto block_shape_transposed{        cute::make_shape(TileSizeX{}, TileSizeY{})};    auto const shared_memory_layout_src{cute::make_layout(        block_shape, cute::GenRowMajor{})}; // (bM, bN) : (bN, 1)    auto const shared_memory_layout_dst{cute::make_layout(        block_shape_transposed, cute::GenRowMajor{})}; // (bN, bM) : (bM, 1)    auto const shared_memory_layout_dst_transposed{cute::make_layout(        block_shape, cute::GenColMajor{})}; // (bM, bN) : (1, bM)    auto const swizzle_src{        cute::Swizzle<NUM_MASK_BITS, NUM_BASE_BITS, NUM_SHIFT_BITS>{}};    auto const shared_memory_layout_swizzled_src{        cute::composition(swizzle_src, shared_memory_layout_src)};    // Inspect if the swizzling reduces the shared memory bank conflicts.    // print_shared_memory_bank_ids(shared_memory_layout_swizzled_src);    auto const tiled_tensor_src{cute::tiled_divide(        tensor_src, block_shape)}; // ((TileSizeY, TileSizeX), M /                                   // TileSizeY, N / TileSizeX)    auto const tiled_tensor_dst{cute::tiled_divide(        tensor_dst, block_shape_transposed)}; // ((TileSizeX, TileSizeY), N                                              // / TileSizeX, M / TileSizeY)    auto const tiled_tensor_dst_transposed{cute::tiled_divide(        tensor_dst_transposed, block_shape)}; // ((TileSizeY, TileSizeX), M                                              // / TileSizeY, N / TileSizeX)    using ThreadBlockSizeX = cute::Int<32>; // tN    using ThreadBlockSizeY = cute::Int<8>;  // tM    CUTE_STATIC_ASSERT_V(TileSizeX{} % ThreadBlockSizeX{} == cute::Int<0>{},                         "TileSizeX must be divisible by ThreadBlockSizeX");    CUTE_STATIC_ASSERT_V(TileSizeY{} % ThreadBlockSizeY{} == cute::Int<0>{},                         "TileSizeY must be divisible by ThreadBlockSizeY");    constexpr auto thread_block_shape{        cute::make_shape(ThreadBlockSizeY{}, ThreadBlockSizeX{})};    constexpr auto thread_block_shape_transposed{        cute::make_shape(ThreadBlockSizeX{}, ThreadBlockSizeY{})};    constexpr auto thread_layout{        cute::make_layout(thread_block_shape, cute::GenRowMajor{})};    constexpr auto thread_layout_transposed{        cute::make_layout(thread_block_shape_transposed, cute::GenColMajor{})};    dim3 const grid_dim{cute::size<2>(tiled_tensor_src),                        cute::size<1>(tiled_tensor_src)};    dim3 const thread_dim{ThreadBlockSizeX::value * ThreadBlockSizeY::value};    matrix_transpose_shared_memory<<<grid_dim, thread_dim, 0, stream>>>(        tiled_tensor_src, tiled_tensor_dst_transposed,        shared_memory_layout_swizzled_src, shared_memory_layout_swizzled_src,        thread_layout, thread_layout_transposed);    return cudaGetLastError();}template <typename T>static cudaError_t launch_matrix_transpose_shared_memory_vectorized_swizzled(    T const* input_matrix, T* output_matrix, unsigned int M, unsigned int N,    cudaStream_t stream){    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)};    if (N % NUM_VECTOR_ELEMENTS != 0)    {        return cudaErrorInvalidValue;    }    auto const tensor_shape{cute::make_shape(M, N)};    auto const tensor_shape_transposed{cute::make_shape(N, M)};    // Input matrix: row-major M x N matrix.    auto const global_memory_layout_src{cute::make_layout(        tensor_shape, cute::GenRowMajor{})}; // (M, N) : (N, 1)    // Output matrix: row-major N x M matrix.    auto const global_memory_layout_dst{cute::make_layout(        tensor_shape_transposed, cute::GenRowMajor{})}; // (N, M) : (M, 1)    // Same output matrix, but different view: column-major M x N matrix.    auto const global_memory_layout_dst_transposed{cute::make_layout(        tensor_shape, cute::GenColMajor{})}; // (M, N) : (1, M)    auto const tensor_src{cute::make_tensor(cute::make_gmem_ptr(input_matrix),                                            global_memory_layout_src)};    auto const tensor_dst{cute::make_tensor(cute::make_gmem_ptr(output_matrix),                                            global_memory_layout_dst)};    auto const tensor_dst_transposed{        cute::make_tensor(cute::make_gmem_ptr(output_matrix),                          global_memory_layout_dst_transposed)};    using TileSizeX = cute::Int<128>; // bN    using TileSizeY = cute::Int<32>;  // bM    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};    constexpr auto block_shape{cute::make_shape(TileSizeY{}, TileSizeX{})};    constexpr auto block_shape_transposed{        cute::make_shape(TileSizeX{}, TileSizeY{})};    auto const shared_memory_layout_src{cute::make_layout(        block_shape, cute::GenRowMajor{})}; // (bM, bN) : (bN, 1)    auto const shared_memory_layout_dst{cute::make_layout(        block_shape_transposed, cute::GenRowMajor{})}; // (bN, bM) : (bM, 1)    auto const shared_memory_layout_dst_transposed{cute::make_layout(        block_shape, cute::GenColMajor{})}; // (bM, bN) : (1, bM)    // Because of the vectorized access, NUM_BASE_BITS cannot be zero.    // The shared memory bank conflict mitigation can be compromised.    // Print the shared memory bank ids to see the details.    auto const swizzle_src{        cute::Swizzle<NUM_MASK_BITS, NUM_BASE_BITS, NUM_SHIFT_BITS>{}};    auto const shared_memory_layout_swizzled_src{        cute::composition(swizzle_src, shared_memory_layout_src)};    // Inspect if the swizzling reduces the shared memory bank conflicts.    // print_shared_memory_bank_ids(shared_memory_layout_swizzled_src);    auto const tiled_tensor_src{cute::tiled_divide(        tensor_src, block_shape)}; // ((TileSizeY, TileSizeX), M /                                   // TileSizeY, N / TileSizeX)    auto const tiled_tensor_dst{cute::tiled_divide(        tensor_dst, block_shape_transposed)}; // ((TileSizeX, TileSizeY), N                                              // / TileSizeX, M / TileSizeY)    auto const tiled_tensor_dst_transposed{cute::tiled_divide(        tensor_dst_transposed, block_shape)}; // ((TileSizeY, TileSizeX), M                                              // / TileSizeY, N / TileSizeX)    using ThreadBlockSizeX = cute::Int<32>; // tN    using ThreadBlockSizeY = cute::Int<8>;  // tM    CUTE_STATIC_ASSERT_V(TileSizeX{} % ThreadBlockSizeX{} == cute::Int<0>{},                         "TileSizeX must be divisible by ThreadBlockSizeX");    CUTE_STATIC_ASSERT_V(TileSizeY{} % ThreadBlockSizeY{} == cute::Int<0>{},                         "TileSizeY must be divisible by ThreadBlockSizeY");    constexpr auto thread_block_shape{        cute::make_shape(ThreadBlockSizeY{}, ThreadBlockSizeX{})};    constexpr auto thread_block_shape_transposed{        cute::make_shape(ThreadBlockSizeX{}, ThreadBlockSizeY{})};    constexpr auto thread_layout{        cute::make_layout(thread_block_shape, cute::GenRowMajor{})};    constexpr auto thread_layout_transposed{        cute::make_layout(thread_block_shape_transposed, cute::GenColMajor{})};    using VECTOR_SIZE_X = cute::Int<NUM_VECTOR_ELEMENTS>;    constexpr auto vector_shape{        cute::make_shape(cute::Int<1>{}, VECTOR_SIZE_X{})};    // Copy atom vector layout.    constexpr auto vector_layout{        cute::make_layout(vector_shape, cute::GenRowMajor{})};    dim3 const grid_dim{cute::size<2>(tiled_tensor_src),                        cute::size<1>(tiled_tensor_src)};    dim3 const thread_dim{ThreadBlockSizeX::value * ThreadBlockSizeY::value};    matrix_transpose_shared_memory_vectorized<<<grid_dim, thread_dim, 0,                                                stream>>>(        tiled_tensor_src, tiled_tensor_dst_transposed,        shared_memory_layout_swizzled_src, shared_memory_layout_swizzled_src,        thread_layout, thread_layout_transposed, vector_layout);    return cudaGetLastError();}// Explicit instantiation.template cudaError_tlaunch_matrix_transpose_shared_memory_bank_conflict_read<float>(    float const* input_matrix, float* output_matrix, unsigned int M,    unsigned int N, cudaStream_t stream);template cudaError_tlaunch_matrix_transpose_shared_memory_bank_conflict_read<double>(    double const* input_matrix, double* output_matrix, unsigned int M,    unsigned int N, cudaStream_t stream);template cudaError_tlaunch_matrix_transpose_shared_memory_vectorized_bank_conflict_read<float>(    float const* input_matrix, float* output_matrix, unsigned int M,    unsigned int N, cudaStream_t stream);template cudaError_tlaunch_matrix_transpose_shared_memory_vectorized_bank_conflict_read<double>(    double const* input_matrix, double* output_matrix, unsigned int M,    unsigned int N, cudaStream_t stream);template cudaError_tlaunch_matrix_transpose_shared_memory_bank_conflict_write<float>(    float const* input_matrix, float* output_matrix, unsigned int M,    unsigned int N, cudaStream_t stream);template cudaError_tlaunch_matrix_transpose_shared_memory_bank_conflict_write<double>(    double const* input_matrix, double* output_matrix, unsigned int M,    unsigned int N, cudaStream_t stream);template cudaError_t launch_matrix_transpose_shared_memory_padded<float>(    float const* input_matrix, float* output_matrix, unsigned int M,    unsigned int N, cudaStream_t stream);template cudaError_t launch_matrix_transpose_shared_memory_padded<double>(    double const* input_matrix, double* output_matrix, unsigned int M,    unsigned int N, cudaStream_t stream);template cudaError_tlaunch_matrix_transpose_shared_memory_vectorized_padded<float>(    float const* input_matrix, float* output_matrix, unsigned int M,    unsigned int N, cudaStream_t stream);template cudaError_tlaunch_matrix_transpose_shared_memory_vectorized_padded<double>(    double const* input_matrix, double* output_matrix, unsigned int M,    unsigned int N, cudaStream_t stream);template cudaError_t launch_matrix_transpose_shared_memory_swizzled<float>(    float const* input_matrix, float* output_matrix, unsigned int M,    unsigned int N, cudaStream_t stream);template cudaError_t launch_matrix_transpose_shared_memory_swizzled<double>(    double const* input_matrix, double* output_matrix, unsigned int M,    unsigned int N, cudaStream_t stream);template cudaError_tlaunch_matrix_transpose_shared_memory_vectorized_swizzled<float>(    float const* input_matrix, float* output_matrix, unsigned int M,    unsigned int N, cudaStream_t stream);template cudaError_tlaunch_matrix_transpose_shared_memory_vectorized_swizzled<double>(    double const* input_matrix, double* output_matrix, unsigned int M,    unsigned int N, cudaStream_t stream);

Shared Memory Layout and CUDA Thread Block Layout

Because the strided memory access will be performed on shared memory, the global memory reads and writes can be fully coalesced. Then we have two options about how to perform the strided memory reads or writes on shared memory. The first option is to perform matrix transpose when reading from global memory to shared memory, and then perform matrix copy from shared memory to global memory, resulting in strided memory writes on shared memory. The second option is to perform matrix copy when reading from global memory to shared memory, and then perform matrix transpose when writing from shared memory to global memory, resulting in strided memory reads on shared memory.

To implement the first option, the shared memory layout has to be column-major if the input matrix layout is row-major. Two different CUDA thread block layouts are used for reading from global memory to shared memory and writing from shared memory to global memory. The first CUDA thread block layout is row-major if the input matrix layout is row-major, resulting in coalesced memory reads from global memory and strided memory writes to shared memory. The second CUDA thread block layout is column-major if the input matrix layout is row-major, which is the same as the output matrix layout, resulting in coalesced memory reads from shared memory and coalesced memory writes to global memory.

To implement the second option, the shared memory layout has to be row-major if the input matrix layout is row-major. Two different CUDA thread block layouts are used for reading from global memory to shared memory and writing from shared memory to global memory. The first CUDA thread block layout is row-major if the input matrix layout is row-major, resulting in coalesced memory reads from global memory and coalesced memory writes to shared memory. The second CUDA thread block layout is column-major if the input matrix layout is row-major, which is the same as the output matrix layout, resulting in strided memory reads from shared memory and coalesced memory writes to global memory.

The strided reads and writes on shared memory will result in as severe as 32-way shared memory bank conflicts. On certain platforms, this will significantly reduce the performance.

123456789101112131415161718192021222324252627282930313233343536

using ThreadBlockSizeX = cute::Int<32>; // tNusing ThreadBlockSizeY = cute::Int<8>;  // tMCUTE_STATIC_ASSERT_V(TileSizeX{} % ThreadBlockSizeX{} == cute::Int<0>{},                     "TileSizeX must be divisible by ThreadBlockSizeX");CUTE_STATIC_ASSERT_V(TileSizeY{} % ThreadBlockSizeY{} == cute::Int<0>{},                     "TileSizeY must be divisible by ThreadBlockSizeY");constexpr auto thread_block_shape{    cute::make_shape(ThreadBlockSizeY{}, ThreadBlockSizeX{})};constexpr auto thread_block_shape_transposed{    cute::make_shape(ThreadBlockSizeX{}, ThreadBlockSizeY{})};constexpr auto thread_layout{    cute::make_layout(thread_block_shape, cute::GenRowMajor{})};constexpr auto thread_layout_transposed{    cute::make_layout(thread_block_shape_transposed, cute::GenColMajor{})};dim3 const grid_dim{cute::size<2>(tiled_tensor_src),                    cute::size<1>(tiled_tensor_src)};dim3 const thread_dim{ThreadBlockSizeX::value * ThreadBlockSizeY::value};if (bank_conflict_access_mode == SharedMemoryBankConflictAccessMode::Read){    matrix_transpose_shared_memory<<<grid_dim, thread_dim, 0, stream>>>(        tiled_tensor_src, tiled_tensor_dst_transposed,        shared_memory_layout_src, shared_memory_layout_src, thread_layout,        thread_layout_transposed);}else{    matrix_transpose_shared_memory<<<grid_dim, thread_dim, 0, stream>>>(        tiled_tensor_src, tiled_tensor_dst_transposed,        shared_memory_layout_dst_transposed,        shared_memory_layout_dst_transposed, thread_layout,        thread_layout_transposed);}

Predicates and Boundary Checking

One typical mistake I would make in the implementation is to use the same predicate for both reading from global memory to shared memory and writing from shared memory to global memory because the shapes of the global memory input matrix tile layout, the global memory output matrix tile layout, the shared memory layout are the same. We could not reuse the predicate because the thread layouts for reading from global memory to shared memory and writing from shared memory to global memory are different. Therefore, even for the same thread, different identity tuples are assigned for reading from global memory to shared memory and writing from shared memory to global memory, and we have to use two sets of predicates.

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455

// A 2D array of tuples that maps (x, y) to (x, y).auto const identity_tensor_src{cute::make_identity_tensor(cute::make_shape(    cute::size<0>(global_tile_src), cute::size<1>(global_tile_src)))};auto const thread_identity_tensor_src{cute::local_partition(    identity_tensor_src, ThreadLayoutSrc{}, threadIdx.x)};auto predicator_src{cute::make_tensor<bool>(    cute::make_shape(cute::size<0>(thread_global_tile_src),                     cute::size<1>(thread_global_tile_src)))};auto const identity_tensor_dst{cute::make_identity_tensor(cute::make_shape(    cute::size<0>(global_tile_dst), cute::size<1>(global_tile_dst)))};auto const thread_identity_tensor_dst{cute::local_partition(    identity_tensor_dst, ThreadLayoutDst{}, threadIdx.x)};auto predicator_dst{cute::make_tensor<bool>(    cute::make_shape(cute::size<0>(thread_global_tile_dst),                     cute::size<1>(thread_global_tile_dst)))};auto const num_max_columns{cute::stride<0>(global_tile_src)};auto const num_max_rows{cute::stride<1>(global_tile_dst)};constexpr auto global_tile_columns{cute::size<1>(global_tile_src)};constexpr auto global_tile_rows{cute::size<0>(global_tile_src)};CUTE_UNROLLfor (unsigned int i{0}; i < cute::size<0>(predicator_src); ++i){    CUTE_UNROLL    for (unsigned int j{0}; j < cute::size<1>(predicator_src); ++j)    {        auto const thread_identity{thread_identity_tensor_src(i, j)};        bool const is_row_in_bound{cute::get<0>(thread_identity) +                                       blockIdx.y * global_tile_rows <                                   num_max_rows};        bool const is_column_in_bound{cute::get<1>(thread_identity) +                                          blockIdx.x * global_tile_columns <                                      num_max_columns};        predicator_src(i, j) = is_row_in_bound && is_column_in_bound;    }}CUTE_UNROLLfor (unsigned int i{0}; i < cute::size<0>(predicator_dst); ++i){    CUTE_UNROLL    for (unsigned int j{0}; j < cute::size<1>(predicator_dst); ++j)    {        auto const thread_identity{thread_identity_tensor_dst(i, j)};        bool const is_row_in_bound{cute::get<0>(thread_identity) +                                       blockIdx.y * global_tile_rows <                                   num_max_rows};        bool const is_column_in_bound{cute::get<1>(thread_identity) +                                          blockIdx.x * global_tile_columns <                                      num_max_columns};        predicator_dst(i, j) = is_row_in_bound && is_column_in_bound;    }}

Thread Block Synchronization

Because shared memory is used as a cache to store the intermediate matrix tile for transpose and all the threads in the same thread block are synergistically reading the matrix tile from global memory to shared memory, before the writing the matrix tile from shared memory to global memory, we have to make sure all the threads in the same thread block have finished reading the matrix tile from global memory to shared memory. In addition to the commonly used __syncthreads(), cute::cp_async_fence() and cute::cp_async_wait<0>() are also used in CuTe for thread block synchronization. This is because cute::copy_if and cute::copy can be asynchronous operations on SM80 and above platforms. cute::cp_async_fence() and cute::cp_async_wait<0>() are no-ops on platforms lower than SM80.

1234567

cute::copy_if(predicator_src, thread_global_tile_src,              thread_shared_tile_src);cute::cp_async_fence();cute::cp_async_wait<0>();__syncthreads();cute::copy_if(predicator_dst, thread_shared_tile_dst,              thread_global_tile_dst);

Shared Memory Padding

The shared memory padding is a common trick to avoid shared memory bank conflicts when a warp of threads is accessing shared memory.

In our case, assuming we have the strided memory read on shared memory. Then instead of using the shared memory layout of $(bM, bN) : (bN, 1)$, the padded shared memory layout should be $(bM, bN) : (bN + 1, 1)$. Notice that the shared memory shape remains unchanged, but the stride of the shared memory layout gets changed, resulting in the shared memory cosize, i.e. the shared memory that needs to be allocated, is also changed. Using the shared memory layout of $(bM, bN + 1) : (bN + 1, 1)$ is incorrect.

12345678910111213141516171819

using TILE_SIZE_X = cute::Int<64>;        // bNusing TILE_SIZE_X_PADDED = cute::Int<65>; // bN + 1using TILE_SIZE_Y = cute::Int<32>;        // bMconstexpr auto block_shape{cute::make_shape(TILE_SIZE_Y{}, TILE_SIZE_X{})};constexpr auto block_shape_transposed{    cute::make_shape(TILE_SIZE_X{}, TILE_SIZE_Y{})};auto const shared_memory_layout_src{cute::make_layout(    block_shape, cute::GenRowMajor{})}; // (bM, bN) : (bN, 1)auto const shared_memory_layout_src_padded{cute::make_layout(    block_shape,    cute::make_stride(TILE_SIZE_X_PADDED{},                      cute::Int<1>{}))}; // (bM, bN) : (bN + 1, 1)transpose_shared_memory<<<grid_dim, thread_dim, 0, stream>>>(    tiled_tensor_src, tiled_tensor_dst_transposed,    shared_memory_layout_src_padded, shared_memory_layout_src_padded,    thread_layout, thread_layout_transposed);

Because the shared memory shape remains the same, the CUDA kernel previously implemented can be just reused.

Shared Memory Swizzling

The shared memory swizzling is another common trick to avoid shared memory bank conflicts when a warp of threads is accessing shared memory. Comparing to the shared memory padding, the shared memory swizzling will not allocate extract shared memory that is not used, and is therefore a more favorable approach. However, the formula of shared memory swizzling is very brain-twisting and the implementation can be very error-prone. In CuTe, fortunately, the shared memory swizzling is implemented as a simple template class, and the shared memory swizzling can be easily applied to the shared memory layout via CuTe layout composition. After verifying the shared memory swizzled bank ids are meeting our requirement, we could just reuse the CUDA kernel previously implemented for the shared memory swizzled layout.

1234567891011121314151617181920212223

using TileSizeX = cute::Int<64>; // bNusing TileSizeY = cute::Int<32>; // bMconstexpr int NUM_BASE_BITS{constexpr_log2(1)};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};constexpr auto block_shape{cute::make_shape(TileSizeY{}, TileSizeX{})};constexpr auto block_shape_transposed{    cute::make_shape(TileSizeX{}, TileSizeY{})};auto const shared_memory_layout_src{cute::make_layout(    block_shape, cute::GenRowMajor{})}; // (bM, bN) : (bN, 1)auto const shared_memory_layout_dst{cute::make_layout(    block_shape_transposed, cute::GenRowMajor{})}; // (bN, bM) : (bM, 1)auto const shared_memory_layout_dst_transposed{cute::make_layout(    block_shape, cute::GenColMajor{})}; // (bM, bN) : (1, bM)auto const swizzle_src{    cute::Swizzle<NUM_MASK_BITS, NUM_BASE_BITS, NUM_SHIFT_BITS>{}};auto const shared_memory_layout_swizzled_src{    cute::composition(swizzle_src, shared_memory_layout_src)};

In our case, given the shared memory of shape $(bM, bN) : (bN, 1) = (32, 64) : (64, 1)$, the shared memory bank id before and after applying CuTe swizzling are as follows, respectively.

1234567891011121314151617181920212223242526272829303132

0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 310  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31

1234567891011121314151617181920212223242526272829303132

0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 1  0  3  2  5  4  7  6  9  8 11 10 13 12 15 14 17 16 19 18 21 20 23 22 25 24 27 26 29 28 31 30  1  0  3  2  5  4  7  6  9  8 11 10 13 12 15 14 17 16 19 18 21 20 23 22 25 24 27 26 29 28 31 30 2  3  0  1  6  7  4  5 10 11  8  9 14 15 12 13 18 19 16 17 22 23 20 21 26 27 24 25 30 31 28 29  2  3  0  1  6  7  4  5 10 11  8  9 14 15 12 13 18 19 16 17 22 23 20 21 26 27 24 25 30 31 28 29 3  2  1  0  7  6  5  4 11 10  9  8 15 14 13 12 19 18 17 16 23 22 21 20 27 26 25 24 31 30 29 28  3  2  1  0  7  6  5  4 11 10  9  8 15 14 13 12 19 18 17 16 23 22 21 20 27 26 25 24 31 30 29 28 4  5  6  7  0  1  2  3 12 13 14 15  8  9 10 11 20 21 22 23 16 17 18 19 28 29 30 31 24 25 26 27  4  5  6  7  0  1  2  3 12 13 14 15  8  9 10 11 20 21 22 23 16 17 18 19 28 29 30 31 24 25 26 27 5  4  7  6  1  0  3  2 13 12 15 14  9  8 11 10 21 20 23 22 17 16 19 18 29 28 31 30 25 24 27 26  5  4  7  6  1  0  3  2 13 12 15 14  9  8 11 10 21 20 23 22 17 16 19 18 29 28 31 30 25 24 27 26 6  7  4  5  2  3  0  1 14 15 12 13 10 11  8  9 22 23 20 21 18 19 16 17 30 31 28 29 26 27 24 25  6  7  4  5  2  3  0  1 14 15 12 13 10 11  8  9 22 23 20 21 18 19 16 17 30 31 28 29 26 27 24 25 7  6  5  4  3  2  1  0 15 14 13 12 11 10  9  8 23 22 21 20 19 18 17 16 31 30 29 28 27 26 25 24  7  6  5  4  3  2  1  0 15 14 13 12 11 10  9  8 23 22 21 20 19 18 17 16 31 30 29 28 27 26 25 24 8  9 10 11 12 13 14 15  0  1  2  3  4  5  6  7 24 25 26 27 28 29 30 31 16 17 18 19 20 21 22 23  8  9 10 11 12 13 14 15  0  1  2  3  4  5  6  7 24 25 26 27 28 29 30 31 16 17 18 19 20 21 22 23 9  8 11 10 13 12 15 14  1  0  3  2  5  4  7  6 25 24 27 26 29 28 31 30 17 16 19 18 21 20 23 22  9  8 11 10 13 12 15 14  1  0  3  2  5  4  7  6 25 24 27 26 29 28 31 30 17 16 19 18 21 20 23 2210 11  8  9 14 15 12 13  2  3  0  1  6  7  4  5 26 27 24 25 30 31 28 29 18 19 16 17 22 23 20 21 10 11  8  9 14 15 12 13  2  3  0  1  6  7  4  5 26 27 24 25 30 31 28 29 18 19 16 17 22 23 20 2111 10  9  8 15 14 13 12  3  2  1  0  7  6  5  4 27 26 25 24 31 30 29 28 19 18 17 16 23 22 21 20 11 10  9  8 15 14 13 12  3  2  1  0  7  6  5  4 27 26 25 24 31 30 29 28 19 18 17 16 23 22 21 2012 13 14 15  8  9 10 11  4  5  6  7  0  1  2  3 28 29 30 31 24 25 26 27 20 21 22 23 16 17 18 19 12 13 14 15  8  9 10 11  4  5  6  7  0  1  2  3 28 29 30 31 24 25 26 27 20 21 22 23 16 17 18 1913 12 15 14  9  8 11 10  5  4  7  6  1  0  3  2 29 28 31 30 25 24 27 26 21 20 23 22 17 16 19 18 13 12 15 14  9  8 11 10  5  4  7  6  1  0  3  2 29 28 31 30 25 24 27 26 21 20 23 22 17 16 19 1814 15 12 13 10 11  8  9  6  7  4  5  2  3  0  1 30 31 28 29 26 27 24 25 22 23 20 21 18 19 16 17 14 15 12 13 10 11  8  9  6  7  4  5  2  3  0  1 30 31 28 29 26 27 24 25 22 23 20 21 18 19 16 1715 14 13 12 11 10  9  8  7  6  5  4  3  2  1  0 31 30 29 28 27 26 25 24 23 22 21 20 19 18 17 16 15 14 13 12 11 10  9  8  7  6  5  4  3  2  1  0 31 30 29 28 27 26 25 24 23 22 21 20 19 18 17 1616 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 1517 16 19 18 21 20 23 22 25 24 27 26 29 28 31 30  1  0  3  2  5  4  7  6  9  8 11 10 13 12 15 14 17 16 19 18 21 20 23 22 25 24 27 26 29 28 31 30  1  0  3  2  5  4  7  6  9  8 11 10 13 12 15 1418 19 16 17 22 23 20 21 26 27 24 25 30 31 28 29  2  3  0  1  6  7  4  5 10 11  8  9 14 15 12 13 18 19 16 17 22 23 20 21 26 27 24 25 30 31 28 29  2  3  0  1  6  7  4  5 10 11  8  9 14 15 12 1319 18 17 16 23 22 21 20 27 26 25 24 31 30 29 28  3  2  1  0  7  6  5  4 11 10  9  8 15 14 13 12 19 18 17 16 23 22 21 20 27 26 25 24 31 30 29 28  3  2  1  0  7  6  5  4 11 10  9  8 15 14 13 1220 21 22 23 16 17 18 19 28 29 30 31 24 25 26 27  4  5  6  7  0  1  2  3 12 13 14 15  8  9 10 11 20 21 22 23 16 17 18 19 28 29 30 31 24 25 26 27  4  5  6  7  0  1  2  3 12 13 14 15  8  9 10 1121 20 23 22 17 16 19 18 29 28 31 30 25 24 27 26  5  4  7  6  1  0  3  2 13 12 15 14  9  8 11 10 21 20 23 22 17 16 19 18 29 28 31 30 25 24 27 26  5  4  7  6  1  0  3  2 13 12 15 14  9  8 11 1022 23 20 21 18 19 16 17 30 31 28 29 26 27 24 25  6  7  4  5  2  3  0  1 14 15 12 13 10 11  8  9 22 23 20 21 18 19 16 17 30 31 28 29 26 27 24 25  6  7  4  5  2  3  0  1 14 15 12 13 10 11  8  923 22 21 20 19 18 17 16 31 30 29 28 27 26 25 24  7  6  5  4  3  2  1  0 15 14 13 12 11 10  9  8 23 22 21 20 19 18 17 16 31 30 29 28 27 26 25 24  7  6  5  4  3  2  1  0 15 14 13 12 11 10  9  824 25 26 27 28 29 30 31 16 17 18 19 20 21 22 23  8  9 10 11 12 13 14 15  0  1  2  3  4  5  6  7 24 25 26 27 28 29 30 31 16 17 18 19 20 21 22 23  8  9 10 11 12 13 14 15  0  1  2  3  4  5  6  725 24 27 26 29 28 31 30 17 16 19 18 21 20 23 22  9  8 11 10 13 12 15 14  1  0  3  2  5  4  7  6 25 24 27 26 29 28 31 30 17 16 19 18 21 20 23 22  9  8 11 10 13 12 15 14  1  0  3  2  5  4  7  626 27 24 25 30 31 28 29 18 19 16 17 22 23 20 21 10 11  8  9 14 15 12 13  2  3  0  1  6  7  4  5 26 27 24 25 30 31 28 29 18 19 16 17 22 23 20 21 10 11  8  9 14 15 12 13  2  3  0  1  6  7  4  527 26 25 24 31 30 29 28 19 18 17 16 23 22 21 20 11 10  9  8 15 14 13 12  3  2  1  0  7  6  5  4 27 26 25 24 31 30 29 28 19 18 17 16 23 22 21 20 11 10  9  8 15 14 13 12  3  2  1  0  7  6  5  428 29 30 31 24 25 26 27 20 21 22 23 16 17 18 19 12 13 14 15  8  9 10 11  4  5  6  7  0  1  2  3 28 29 30 31 24 25 26 27 20 21 22 23 16 17 18 19 12 13 14 15  8  9 10 11  4  5  6  7  0  1  2  329 28 31 30 25 24 27 26 21 20 23 22 17 16 19 18 13 12 15 14  9  8 11 10  5  4  7  6  1  0  3  2 29 28 31 30 25 24 27 26 21 20 23 22 17 16 19 18 13 12 15 14  9  8 11 10  5  4  7  6  1  0  3  230 31 28 29 26 27 24 25 22 23 20 21 18 19 16 17 14 15 12 13 10 11  8  9  6  7  4  5  2  3  0  1 30 31 28 29 26 27 24 25 22 23 20 21 18 19 16 17 14 15 12 13 10 11  8  9  6  7  4  5  2  3  0  131 30 29 28 27 26 25 24 23 22 21 20 19 18 17 16 15 14 13 12 11 10  9  8  7  6  5  4  3  2  1  0 31 30 29 28 27 26 25 24 23 22 21 20 19 18 17 16 15 14 13 12 11 10  9  8  7  6  5  4  3  2  1  0

We could see that when a warp of threads read or write on the column of the shared memory of row-major, it’s shared memory bank conflict free.

Performances

The following tables show the performance measurements of the matrix transpose CUDA kernels on NVIDIA GeForce RTX 3090.

It’s somewhat surprising that except the native coalesced read CUDA kernel, all the other CUDA kernels have similar effective bandwidth and the bandwidth is very close to the ones that can be achieved in practice. Whether having shared memory bank conflicts in this CUDA kernel does not affect the performance significantly, because the performance bottleneck is in the global memory access.

After profiling using NVIDIA Nsight Compute, we could confirm that the global memory access is not fully coalesced for the native coalesced read and the native coalesced write CUDA kernels, shared memory bank load conflicts present in the shared memory bank conflict read CUDA kernel, shared memory bank store conflicts present in the shared memory bank conflict write CUDA kernel, and shared memory bank conflicts are free in the shared memory padded and shared memory swizzled CUDA kernels.

References

CuTe Matrix Transpose

https://leimao.github.io/article/CuTe-Matrix-Transpose/

Author

Lei Mao

Posted on

11-20-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: