upstream: add GEMM kernel references from 4 repos for BI-V100 porting

Sources (all CUDA 10.2 compatible, no CUTLASS/Triton dependency):
- leimao/CUDA-GEMM-Optimization: v00-v07, fp16 WMMA variant, double buffered
- siboehm/SGEMM_CUDA: kernel 1-12, warp tiling + double buffering
- wangzyon/NVIDIA_SGEMM_PRACTICE: kernel 1-7
- edtallison/sgemm-cuda: kernel 1-12 (reimplementation with notes)

Key porting issue: ALL kernels hardcode WARPSIZE=32.
BI-V100 has warp_size=64. Need to:
1. Replace all 32U / WARPSIZE constants with 64
2. Adjust warp subtile decomposition (WMITER, WNITER, WSUBM, WSUBN)
3. Adjust shared memory bank conflict avoidance (may have different bank count)
4. Test __shfl_down_sync with mask=0xFFFFFFFFFFFFFFFF (64-bit)

upstream_ref/cuda_gemm_optimization/00_non_coalesced_global_memory_access.cu

@@ -0,0 +1,65 @@
+#include <cuda_fp16.h>
+
+#include "cuda_gemm.hpp"
+#include "cuda_gemm_utils.hpp"
+
+// GEMM kernel v00.
+// Non-coalesced read and write from global memory.
+template <typename T>
+__global__ void gemm_v00(size_t m, size_t n, size_t k, T alpha, T const* A,
+                         size_t lda, T const* B, size_t ldb, T beta, T* C,
+                         size_t ldc)
+{
+    // Compute the row and column of C that this thread is responsible for.
+    size_t const C_row_idx{blockIdx.x * blockDim.x + threadIdx.x};
+    size_t const C_col_idx{blockIdx.y * blockDim.y + threadIdx.y};
+
+    // Each thread compute
+    // C[C_row_idx, C_col_idx] = alpha * A[C_row_idx, :] * B[:, C_col_idx] +
+    // beta * C[C_row_idx, C_col_idx].
+    if (C_row_idx < m && C_col_idx < n)
+    {
+        T sum{static_cast<T>(0)};
+        for (size_t k_idx{0U}; k_idx < k; ++k_idx)
+        {
+            sum += A[C_row_idx * lda + k_idx] * B[k_idx * ldb + C_col_idx];
+        }
+        C[C_row_idx * ldc + C_col_idx] =
+            alpha * sum + beta * C[C_row_idx * ldc + C_col_idx];
+    }
+}
+
+template <typename T>
+void launch_gemm_kernel_v00(size_t m, size_t n, size_t k, T const* alpha,
+                            T const* A, size_t lda, T const* B, size_t ldb,
+                            T const* beta, T* C, size_t ldc,
+                            cudaStream_t stream)
+{
+    dim3 const block_dim{32U, 32U, 1U};
+    dim3 const grid_dim{
+        (static_cast<unsigned int>(m) + block_dim.x - 1U) / block_dim.x,
+        (static_cast<unsigned int>(n) + block_dim.y - 1U) / block_dim.y, 1U};
+    gemm_v00<T><<<grid_dim, block_dim, 0U, stream>>>(m, n, k, *alpha, A, lda, B,
+                                                     ldb, *beta, C, ldc);
+    CHECK_LAST_CUDA_ERROR();
+}
+
+// Explicit instantiation.
+template void launch_gemm_kernel_v00<float>(size_t m, size_t n, size_t k,
+                                            float const* alpha, float const* A,
+                                            size_t lda, float const* B,
+                                            size_t ldb, float const* beta,
+                                            float* C, size_t ldc,
+                                            cudaStream_t stream);
+template void launch_gemm_kernel_v00<double>(size_t m, size_t n, size_t k,
+                                             double const* alpha,
+                                             double const* A, size_t lda,
+                                             double const* B, size_t ldb,
+                                             double const* beta, double* C,
+                                             size_t ldc, cudaStream_t stream);
+template void launch_gemm_kernel_v00<__half>(size_t m, size_t n, size_t k,
+                                             __half const* alpha,
+                                             __half const* A, size_t lda,
+                                             __half const* B, size_t ldb,
+                                             __half const* beta, __half* C,
+                                             size_t ldc, cudaStream_t stream);

upstream_ref/cuda_gemm_optimization/01_coalesced_global_memory_access.cu

@@ -0,0 +1,65 @@
+#include <cuda_fp16.h>
+
+#include "cuda_gemm.hpp"
+#include "cuda_gemm_utils.hpp"
+
+// GEMM kernel v01.
+// Coalesced read and write from global memory.
+template <typename T>
+__global__ void gemm_v01(size_t m, size_t n, size_t k, T alpha, T const* A,
+                         size_t lda, T const* B, size_t ldb, T beta, T* C,
+                         size_t ldc)
+{
+    // Compute the row and column of C that this thread is responsible for.
+    size_t const C_col_idx{blockIdx.x * blockDim.x + threadIdx.x};
+    size_t const C_row_idx{blockIdx.y * blockDim.y + threadIdx.y};
+
+    // Each thread compute
+    // C[C_row_idx, C_col_idx] = alpha * A[C_row_idx, :] * B[:, C_col_idx] +
+    // beta * C[C_row_idx, C_col_idx].
+    if (C_row_idx < m && C_col_idx < n)
+    {
+        T sum{static_cast<T>(0)};
+        for (size_t k_idx{0U}; k_idx < k; ++k_idx)
+        {
+            sum += A[C_row_idx * lda + k_idx] * B[k_idx * ldb + C_col_idx];
+        }
+        C[C_row_idx * ldc + C_col_idx] =
+            alpha * sum + beta * C[C_row_idx * ldc + C_col_idx];
+    }
+}
+
+template <typename T>
+void launch_gemm_kernel_v01(size_t m, size_t n, size_t k, T const* alpha,
+                            T const* A, size_t lda, T const* B, size_t ldb,
+                            T const* beta, T* C, size_t ldc,
+                            cudaStream_t stream)
+{
+    dim3 const block_dim{32U, 32U, 1U};
+    dim3 const grid_dim{
+        (static_cast<unsigned int>(n) + block_dim.x - 1U) / block_dim.x,
+        (static_cast<unsigned int>(m) + block_dim.y - 1U) / block_dim.y, 1U};
+    gemm_v01<T><<<grid_dim, block_dim, 0U, stream>>>(m, n, k, *alpha, A, lda, B,
+                                                     ldb, *beta, C, ldc);
+    CHECK_LAST_CUDA_ERROR();
+}
+
+// Explicit instantiation.
+template void launch_gemm_kernel_v01<float>(size_t m, size_t n, size_t k,
+                                            float const* alpha, float const* A,
+                                            size_t lda, float const* B,
+                                            size_t ldb, float const* beta,
+                                            float* C, size_t ldc,
+                                            cudaStream_t stream);
+template void launch_gemm_kernel_v01<double>(size_t m, size_t n, size_t k,
+                                             double const* alpha,
+                                             double const* A, size_t lda,
+                                             double const* B, size_t ldb,
+                                             double const* beta, double* C,
+                                             size_t ldc, cudaStream_t stream);
+template void launch_gemm_kernel_v01<__half>(size_t m, size_t n, size_t k,
+                                             __half const* alpha,
+                                             __half const* A, size_t lda,
+                                             __half const* B, size_t ldb,
+                                             __half const* beta, __half* C,
+                                             size_t ldc, cudaStream_t stream);

upstream_ref/cuda_gemm_optimization/02_2d_block_tiling.cu

@@ -0,0 +1,112 @@
+#include <cuda_fp16.h>
+
+#include "cuda_gemm.hpp"
+#include "cuda_gemm_utils.cuh"
+#include "cuda_gemm_utils.hpp"
+
+// GEMM kernel v02.
+// Coalesced read and write from global memory.
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t BLOCK_TILE_SIZE_K>
+__global__ void gemm_v02(size_t m, size_t n, size_t k, T alpha, T const* A,
+                         size_t lda, T const* B, size_t ldb, T beta, T* C,
+                         size_t ldc)
+{
+    // Avoid using blockDim.x * blockDim.y as the number of threads per block.
+    // Because it is a runtime constant and the compiler cannot optimize the
+    // loop unrolling based on that.
+    // Use a compile time constant instead.
+    constexpr size_t NUM_THREADS{BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y};
+    size_t const thread_linear_idx{threadIdx.y * blockDim.x + threadIdx.x};
+
+    // Compute the row and column of C that this thread is responsible for.
+    size_t const C_col_idx{blockIdx.x * blockDim.x + threadIdx.x};
+    size_t const C_row_idx{blockIdx.y * blockDim.y + threadIdx.y};
+
+    // Cache a tile of A and B in shared memory for data reuse.
+    __shared__ T A_thread_block_tile[BLOCK_TILE_SIZE_Y][BLOCK_TILE_SIZE_K];
+    __shared__ T B_thread_block_tile[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_X];
+
+    size_t const num_thread_block_tiles{(k + BLOCK_TILE_SIZE_K - 1) /
+                                        BLOCK_TILE_SIZE_K};
+
+    T sum{static_cast<T>(0)};
+    for (size_t thread_block_tile_idx{0U};
+         thread_block_tile_idx < num_thread_block_tiles;
+         ++thread_block_tile_idx)
+    {
+        load_data_from_global_memory_to_shared_memory<
+            T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+            NUM_THREADS>(A, lda, B, ldb, A_thread_block_tile,
+                         B_thread_block_tile, thread_block_tile_idx,
+                         thread_linear_idx, m, n, k);
+        __syncthreads();
+
+#pragma unroll
+        for (size_t k_i{0U}; k_i < BLOCK_TILE_SIZE_K; ++k_i)
+        {
+            // Doing this results in 2 TOPS.
+            // Suppose blockDim.x = blockDim.y = 32.
+            // Effectively, for a warp, in one iteration, we read the value from
+            // A_thread_block_tile at the same location on the shared memory
+            // resulting in a broadcast, we also read 32 values that have no
+            // bank conflicts from B_thread_block_tile. Even with that, all the
+            // values have to be read from the shared memory and consequence is
+            // the shared memory instruction runs very intensively just to
+            // compute a small number of values using simple arithmetic
+            // instructions, which is not efficient.
+            sum += A_thread_block_tile[threadIdx.y][k_i] *
+                   B_thread_block_tile[k_i][threadIdx.x];
+        }
+        __syncthreads();
+    }
+    if (C_row_idx < m && C_col_idx < n)
+    {
+        C[C_row_idx * ldc + C_col_idx] =
+            alpha * sum + beta * C[C_row_idx * ldc + C_col_idx];
+    }
+}
+
+template <typename T>
+void launch_gemm_kernel_v02(size_t m, size_t n, size_t k, T const* alpha,
+                            T const* A, size_t lda, T const* B, size_t ldb,
+                            T const* beta, T* C, size_t ldc,
+                            cudaStream_t stream)
+{
+    // Feel free to play with the block tile sizes.
+    // The algorithm correctness should always be guaranteed.
+    constexpr unsigned int BLOCK_TILE_SIZE_X{32U};
+    constexpr unsigned int BLOCK_TILE_SIZE_Y{32U};
+    constexpr unsigned int BLOCK_TILE_SIZE_K{32U};
+    constexpr unsigned int NUM_THREADS{BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y};
+    static_assert(BLOCK_TILE_SIZE_K * BLOCK_TILE_SIZE_Y % NUM_THREADS == 0U);
+    static_assert(BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_K % NUM_THREADS == 0U);
+    dim3 const block_dim{BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, 1U};
+    dim3 const grid_dim{
+        (static_cast<unsigned int>(n) + block_dim.x - 1U) / block_dim.x,
+        (static_cast<unsigned int>(m) + block_dim.y - 1U) / block_dim.y, 1U};
+    gemm_v02<T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K>
+        <<<grid_dim, block_dim, 0U, stream>>>(m, n, k, *alpha, A, lda, B, ldb,
+                                              *beta, C, ldc);
+    CHECK_LAST_CUDA_ERROR();
+}
+
+// Explicit instantiation.
+template void launch_gemm_kernel_v02<float>(size_t m, size_t n, size_t k,
+                                            float const* alpha, float const* A,
+                                            size_t lda, float const* B,
+                                            size_t ldb, float const* beta,
+                                            float* C, size_t ldc,
+                                            cudaStream_t stream);
+template void launch_gemm_kernel_v02<double>(size_t m, size_t n, size_t k,
+                                             double const* alpha,
+                                             double const* A, size_t lda,
+                                             double const* B, size_t ldb,
+                                             double const* beta, double* C,
+                                             size_t ldc, cudaStream_t stream);
+template void launch_gemm_kernel_v02<__half>(size_t m, size_t n, size_t k,
+                                             __half const* alpha,
+                                             __half const* A, size_t lda,
+                                             __half const* B, size_t ldb,
+                                             __half const* beta, __half* C,
+                                             size_t ldc, cudaStream_t stream);

upstream_ref/cuda_gemm_optimization/02_2d_block_tiling_vectorized_memory_access.cu

@@ -0,0 +1,108 @@
+#include <cuda_fp16.h>
+
+#include "cuda_gemm.hpp"
+#include "cuda_gemm_utils.cuh"
+#include "cuda_gemm_utils.hpp"
+
+// GEMM kernel v02.
+// Coalesced read and write from global memory.
+// We guarantee that matrix A, B, and C are 32 byte aligned.
+// This implementation is slower because we waste a lot of threads.
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t BLOCK_TILE_SIZE_K>
+__global__ void gemm_v02_vectorized(size_t m, size_t n, size_t k, T alpha,
+                                    T const* A, size_t lda, T const* B,
+                                    size_t ldb, T beta, T* C, size_t ldc)
+{
+    // Avoid using blockDim.x * blockDim.y as the number of threads per block.
+    // Because it is a runtime constant and the compiler cannot optimize the
+    // loop unrolling based on that.
+    // Use a compile time constant instead.
+    constexpr size_t NUM_THREADS{BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y};
+    size_t const thread_linear_idx{threadIdx.y * blockDim.x + threadIdx.x};
+
+    // Compute the row and column of C that this thread is responsible for.
+    size_t const C_col_idx{blockIdx.x * blockDim.x + threadIdx.x};
+    size_t const C_row_idx{blockIdx.y * blockDim.y + threadIdx.y};
+
+    // Cache a tile of A and B in shared memory for data reuse.
+    __shared__ T A_thread_block_tile[BLOCK_TILE_SIZE_Y][BLOCK_TILE_SIZE_K];
+    __shared__ T B_thread_block_tile[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_X];
+
+    size_t const num_thread_block_tiles{(k + BLOCK_TILE_SIZE_K - 1) /
+                                        BLOCK_TILE_SIZE_K};
+
+    T sum{static_cast<T>(0)};
+    for (size_t thread_block_tile_idx{0U};
+         thread_block_tile_idx < num_thread_block_tiles;
+         ++thread_block_tile_idx)
+    {
+        load_data_from_global_memory_to_shared_memory_vectorized<
+            T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+            NUM_THREADS>(A, lda, B, ldb, A_thread_block_tile,
+                         B_thread_block_tile, thread_block_tile_idx,
+                         thread_linear_idx, m, n, k);
+        __syncthreads();
+
+#pragma unroll
+        for (size_t k_i{0U}; k_i < BLOCK_TILE_SIZE_K; ++k_i)
+        {
+            // Doing this results in 2 TOPS.
+            // Suppose blockDim.x = blockDim.y = 32.
+            // Effectively, for a warp, in one iteration, we read the value from
+            // A_thread_block_tile at the same location on the shared memory
+            // resulting in a broadcast, we also read 32 values that have no
+            // bank conflicts from B_thread_block_tile. Even with that, all the
+            // values have to be read from the shared memory and consequence is
+            // the shared memory instruction runs very intensively just to
+            // compute a small number of values using simple arithmetic
+            // instructions, which is not efficient.
+            sum += A_thread_block_tile[threadIdx.y][k_i] *
+                   B_thread_block_tile[k_i][threadIdx.x];
+        }
+        __syncthreads();
+    }
+    if (C_row_idx < m && C_col_idx < n)
+    {
+        C[C_row_idx * ldc + C_col_idx] =
+            alpha * sum + beta * C[C_row_idx * ldc + C_col_idx];
+    }
+}
+
+template <typename T>
+void launch_gemm_kernel_v02_vectorized(size_t m, size_t n, size_t k,
+                                       T const* alpha, T const* A, size_t lda,
+                                       T const* B, size_t ldb, T const* beta,
+                                       T* C, size_t ldc, cudaStream_t stream)
+{
+    // Feel free to play with the block tile sizes.
+    // The algorithm correctness should always be guaranteed.
+    constexpr unsigned int BLOCK_TILE_SIZE_X{32U};
+    constexpr unsigned int BLOCK_TILE_SIZE_Y{32U};
+    constexpr unsigned int BLOCK_TILE_SIZE_K{32U};
+    constexpr unsigned int NUM_THREADS{BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y};
+    static_assert(BLOCK_TILE_SIZE_K * BLOCK_TILE_SIZE_Y % NUM_THREADS == 0U);
+    static_assert(BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_K % NUM_THREADS == 0U);
+    dim3 const block_dim{BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, 1U};
+    dim3 const grid_dim{
+        (static_cast<unsigned int>(n) + block_dim.x - 1U) / block_dim.x,
+        (static_cast<unsigned int>(m) + block_dim.y - 1U) / block_dim.y, 1U};
+    gemm_v02_vectorized<T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y,
+                        BLOCK_TILE_SIZE_K><<<grid_dim, block_dim, 0U, stream>>>(
+        m, n, k, *alpha, A, lda, B, ldb, *beta, C, ldc);
+    CHECK_LAST_CUDA_ERROR();
+}
+
+// Explicit instantiation.
+template void launch_gemm_kernel_v02_vectorized<float>(
+    size_t m, size_t n, size_t k, float const* alpha, float const* A,
+    size_t lda, float const* B, size_t ldb, float const* beta, float* C,
+    size_t ldc, cudaStream_t stream);
+template void launch_gemm_kernel_v02_vectorized<double>(
+    size_t m, size_t n, size_t k, double const* alpha, double const* A,
+    size_t lda, double const* B, size_t ldb, double const* beta, double* C,
+    size_t ldc, cudaStream_t stream);
+template void launch_gemm_kernel_v02_vectorized<__half>(
+    size_t m, size_t n, size_t k, __half const* alpha, __half const* A,
+    size_t lda, __half const* B, size_t ldb, __half const* beta, __half* C,
+    size_t ldc, cudaStream_t stream);

upstream_ref/cuda_gemm_optimization/03_2d_block_tiling_1d_thread_tiling.cu

@@ -0,0 +1,144 @@
+#include <cuda_fp16.h>
+
+#include "cuda_gemm.hpp"
+#include "cuda_gemm_utils.cuh"
+#include "cuda_gemm_utils.hpp"
+
+// GEMM kernel v03.
+// Coalesced read and write from global memory.
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t BLOCK_TILE_SIZE_K, size_t THREAD_TILE_SIZE_Y>
+__global__ void gemm_v03(size_t m, size_t n, size_t k, T alpha, T const* A,
+                         size_t lda, T const* B, size_t ldb, T beta, T* C,
+                         size_t ldc)
+{
+    // Avoid using blockDim.x * blockDim.y as the number of threads per block.
+    // Because it is a runtime constant and the compiler cannot optimize the
+    // loop unrolling based on that.
+    // Use a compile time constant instead.
+    constexpr size_t NUM_THREADS{BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y /
+                                 THREAD_TILE_SIZE_Y};
+    size_t const thread_linear_idx{threadIdx.y * blockDim.x + threadIdx.x};
+
+    // Cache a tile of A and B in shared memory for data reuse.
+    __shared__ T A_thread_block_tile[BLOCK_TILE_SIZE_Y][BLOCK_TILE_SIZE_K];
+    __shared__ T B_thread_block_tile[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_X];
+
+    size_t const num_thread_block_tiles{(k + BLOCK_TILE_SIZE_K - 1) /
+                                        BLOCK_TILE_SIZE_K};
+
+    // Each thread in the block processes BLOCK_TILE_SIZE_Y output values.
+    // Specifically, these values corresponds to
+    // C[blockIdx.y * BLOCK_TILE_SIZE_Y + threadIdx.x / BLOCK_TILE_SIZE_X *
+    // THREAD_TILE_SIZE_Y : blockIdx.y * BLOCK_TILE_SIZE_Y + (threadIdx.x /
+    // BLOCK_TILE_SIZE_X + 1) * THREAD_TILE_SIZE_Y][blockIdx.x *
+    // BLOCK_TILE_SIZE_X + threadIdx.x % BLOCK_TILE_SIZE_X]
+    T C_thread_results[THREAD_TILE_SIZE_Y] = {static_cast<T>(0)};
+
+    for (size_t thread_block_tile_idx{0U};
+         thread_block_tile_idx < num_thread_block_tiles;
+         ++thread_block_tile_idx)
+    {
+        load_data_from_global_memory_to_shared_memory<
+            T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+            NUM_THREADS>(A, lda, B, ldb, A_thread_block_tile,
+                         B_thread_block_tile, thread_block_tile_idx,
+                         thread_linear_idx, m, n, k);
+        __syncthreads();
+
+#pragma unroll
+        for (size_t k_i{0U}; k_i < BLOCK_TILE_SIZE_K; ++k_i)
+        {
+            size_t const B_thread_block_tile_row_idx{k_i};
+            // B_val is cached in the register to alleviate the pressure on the
+            // shared memory access.
+            T const B_val{
+                B_thread_block_tile[B_thread_block_tile_row_idx]
+                                   [thread_linear_idx % BLOCK_TILE_SIZE_X]};
+#pragma unroll
+            for (size_t thread_tile_row_idx{0U};
+                 thread_tile_row_idx < THREAD_TILE_SIZE_Y;
+                 ++thread_tile_row_idx)
+            {
+                size_t const A_thread_block_tile_row_idx{
+                    thread_linear_idx / BLOCK_TILE_SIZE_X * THREAD_TILE_SIZE_Y +
+                    thread_tile_row_idx};
+                size_t const A_thread_block_tile_col_idx{k_i};
+                T const A_val{A_thread_block_tile[A_thread_block_tile_row_idx]
+                                                 [A_thread_block_tile_col_idx]};
+                C_thread_results[thread_tile_row_idx] += A_val * B_val;
+            }
+        }
+        __syncthreads();
+    }
+
+// Write the results to DRAM.
+#pragma unroll
+    for (size_t thread_tile_row_idx{0U};
+         thread_tile_row_idx < THREAD_TILE_SIZE_Y; ++thread_tile_row_idx)
+    {
+        size_t const C_row_idx{blockIdx.y * BLOCK_TILE_SIZE_Y +
+                               thread_linear_idx / BLOCK_TILE_SIZE_X *
+                                   THREAD_TILE_SIZE_Y +
+                               thread_tile_row_idx};
+        size_t const C_col_idx{blockIdx.x * BLOCK_TILE_SIZE_X +
+                               thread_linear_idx % BLOCK_TILE_SIZE_X};
+        if (C_row_idx < m && C_col_idx < n)
+        {
+            C[C_row_idx * ldc + C_col_idx] =
+                alpha * C_thread_results[thread_tile_row_idx] +
+                beta * C[C_row_idx * ldc + C_col_idx];
+        }
+    }
+}
+
+template <typename T>
+void launch_gemm_kernel_v03(size_t m, size_t n, size_t k, T const* alpha,
+                            T const* A, size_t lda, T const* B, size_t ldb,
+                            T const* beta, T* C, size_t ldc,
+                            cudaStream_t stream)
+{
+    // Feel free to play with the block tile sizes.
+    // The algorithm correctness should always be guaranteed.
+    constexpr unsigned int BLOCK_TILE_SIZE_X{64U};
+    constexpr unsigned int BLOCK_TILE_SIZE_Y{64U};
+    constexpr unsigned int BLOCK_TILE_SIZE_K{8U};
+    // Each thread computes THREAD_TILE_SIZE_Y values of C.
+    constexpr unsigned int THREAD_TILE_SIZE_Y{8U};
+    constexpr unsigned int NUM_THREADS_PER_BLOCK{
+        BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y / THREAD_TILE_SIZE_Y};
+    static_assert(BLOCK_TILE_SIZE_Y % THREAD_TILE_SIZE_Y == 0U);
+    static_assert(NUM_THREADS_PER_BLOCK % BLOCK_TILE_SIZE_K == 0U);
+    static_assert(NUM_THREADS_PER_BLOCK % BLOCK_TILE_SIZE_X == 0U);
+    dim3 const block_dim{NUM_THREADS_PER_BLOCK, 1U, 1U};
+    dim3 const grid_dim{
+        (static_cast<unsigned int>(n) + BLOCK_TILE_SIZE_X - 1U) /
+            BLOCK_TILE_SIZE_X,
+        (static_cast<unsigned int>(m) + BLOCK_TILE_SIZE_Y - 1U) /
+            BLOCK_TILE_SIZE_Y,
+        1U};
+    gemm_v03<T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+             THREAD_TILE_SIZE_Y><<<grid_dim, block_dim, 0U, stream>>>(
+        m, n, k, *alpha, A, lda, B, ldb, *beta, C, ldc);
+    CHECK_LAST_CUDA_ERROR();
+}
+
+// Explicit instantiation.
+template void launch_gemm_kernel_v03<float>(size_t m, size_t n, size_t k,
+                                            float const* alpha, float const* A,
+                                            size_t lda, float const* B,
+                                            size_t ldb, float const* beta,
+                                            float* C, size_t ldc,
+                                            cudaStream_t stream);
+template void launch_gemm_kernel_v03<double>(size_t m, size_t n, size_t k,
+                                             double const* alpha,
+                                             double const* A, size_t lda,
+                                             double const* B, size_t ldb,
+                                             double const* beta, double* C,
+                                             size_t ldc, cudaStream_t stream);
+template void launch_gemm_kernel_v03<__half>(size_t m, size_t n, size_t k,
+                                             __half const* alpha,
+                                             __half const* A, size_t lda,
+                                             __half const* B, size_t ldb,
+                                             __half const* beta, __half* C,
+                                             size_t ldc, cudaStream_t stream);

upstream_ref/cuda_gemm_optimization/03_2d_block_tiling_1d_thread_tiling_vectorized_memory_access.cu

@@ -0,0 +1,141 @@
+#include <cuda_fp16.h>
+
+#include "cuda_gemm.hpp"
+#include "cuda_gemm_utils.cuh"
+#include "cuda_gemm_utils.hpp"
+
+// GEMM kernel v03.
+// Coalesced read and write from global memory.
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t BLOCK_TILE_SIZE_K, size_t THREAD_TILE_SIZE_Y>
+__global__ void gemm_v03_vectorized(size_t m, size_t n, size_t k, T alpha,
+                                    T const* A, size_t lda, T const* B,
+                                    size_t ldb, T beta, T* C, size_t ldc)
+{
+    // Avoid using blockDim.x * blockDim.y as the number of threads per block.
+    // Because it is a runtime constant and the compiler cannot optimize the
+    // loop unrolling based on that.
+    // Use a compile time constant instead.
+    constexpr size_t NUM_THREADS{BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y /
+                                 THREAD_TILE_SIZE_Y};
+    size_t const thread_linear_idx{threadIdx.y * blockDim.x + threadIdx.x};
+
+    // Cache a tile of A and B in shared memory for data reuse.
+    __shared__ T A_thread_block_tile[BLOCK_TILE_SIZE_Y][BLOCK_TILE_SIZE_K];
+    __shared__ T B_thread_block_tile[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_X];
+
+    size_t const num_thread_block_tiles{(k + BLOCK_TILE_SIZE_K - 1) /
+                                        BLOCK_TILE_SIZE_K};
+
+    // Each thread in the block processes BLOCK_TILE_SIZE_Y output values.
+    // Specifically, these values corresponds to
+    // C[blockIdx.y * BLOCK_TILE_SIZE_Y + threadIdx.x / BLOCK_TILE_SIZE_X *
+    // THREAD_TILE_SIZE_Y : blockIdx.y * BLOCK_TILE_SIZE_Y + (threadIdx.x /
+    // BLOCK_TILE_SIZE_X + 1) * THREAD_TILE_SIZE_Y][blockIdx.x *
+    // BLOCK_TILE_SIZE_X + threadIdx.x % BLOCK_TILE_SIZE_X]
+    T C_thread_results[THREAD_TILE_SIZE_Y] = {static_cast<T>(0)};
+
+    for (size_t thread_block_tile_idx{0U};
+         thread_block_tile_idx < num_thread_block_tiles;
+         ++thread_block_tile_idx)
+    {
+        load_data_from_global_memory_to_shared_memory_vectorized<
+            T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+            NUM_THREADS>(A, lda, B, ldb, A_thread_block_tile,
+                         B_thread_block_tile, thread_block_tile_idx,
+                         thread_linear_idx, m, n, k);
+        __syncthreads();
+
+#pragma unroll
+        for (size_t k_i{0U}; k_i < BLOCK_TILE_SIZE_K; ++k_i)
+        {
+            size_t const B_thread_block_tile_row_idx{k_i};
+            // B_val is cached in the register to alleviate the pressure on the
+            // shared memory access.
+            T const B_val{
+                B_thread_block_tile[B_thread_block_tile_row_idx]
+                                   [thread_linear_idx % BLOCK_TILE_SIZE_X]};
+#pragma unroll
+            for (size_t thread_tile_row_idx{0U};
+                 thread_tile_row_idx < THREAD_TILE_SIZE_Y;
+                 ++thread_tile_row_idx)
+            {
+                size_t const A_thread_block_tile_row_idx{
+                    thread_linear_idx / BLOCK_TILE_SIZE_X * THREAD_TILE_SIZE_Y +
+                    thread_tile_row_idx};
+                size_t const A_thread_block_tile_col_idx{k_i};
+                T const A_val{A_thread_block_tile[A_thread_block_tile_row_idx]
+                                                 [A_thread_block_tile_col_idx]};
+                C_thread_results[thread_tile_row_idx] += A_val * B_val;
+            }
+        }
+        __syncthreads();
+    }
+
+// Write the results to DRAM.
+// Cannot vectorized the write to DRAM because we are writting to a column
+// instead of a row in C.
+#pragma unroll
+    for (size_t thread_tile_row_idx{0U};
+         thread_tile_row_idx < THREAD_TILE_SIZE_Y; ++thread_tile_row_idx)
+    {
+        size_t const C_row_idx{blockIdx.y * BLOCK_TILE_SIZE_Y +
+                               thread_linear_idx / BLOCK_TILE_SIZE_X *
+                                   THREAD_TILE_SIZE_Y +
+                               thread_tile_row_idx};
+        size_t const C_col_idx{blockIdx.x * BLOCK_TILE_SIZE_X +
+                               thread_linear_idx % BLOCK_TILE_SIZE_X};
+        if (C_row_idx < m && C_col_idx < n)
+        {
+            C[C_row_idx * ldc + C_col_idx] =
+                alpha * C_thread_results[thread_tile_row_idx] +
+                beta * C[C_row_idx * ldc + C_col_idx];
+        }
+    }
+}
+
+template <typename T>
+void launch_gemm_kernel_v03_vectorized(size_t m, size_t n, size_t k,
+                                       T const* alpha, T const* A, size_t lda,
+                                       T const* B, size_t ldb, T const* beta,
+                                       T* C, size_t ldc, cudaStream_t stream)
+{
+    // Feel free to play with the block tile sizes.
+    // The algorithm correctness should always be guaranteed.
+    constexpr unsigned int BLOCK_TILE_SIZE_X{64U};
+    constexpr unsigned int BLOCK_TILE_SIZE_Y{64U};
+    constexpr unsigned int BLOCK_TILE_SIZE_K{8U};
+    // Each thread computes THREAD_TILE_SIZE_Y values of C.
+    constexpr unsigned int THREAD_TILE_SIZE_Y{8U};
+    constexpr unsigned int NUM_THREADS_PER_BLOCK{
+        BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y / THREAD_TILE_SIZE_Y};
+    static_assert(BLOCK_TILE_SIZE_Y % THREAD_TILE_SIZE_Y == 0U);
+    static_assert(NUM_THREADS_PER_BLOCK % BLOCK_TILE_SIZE_K == 0U);
+    static_assert(NUM_THREADS_PER_BLOCK % BLOCK_TILE_SIZE_X == 0U);
+    dim3 const block_dim{NUM_THREADS_PER_BLOCK, 1U, 1U};
+    dim3 const grid_dim{
+        (static_cast<unsigned int>(n) + BLOCK_TILE_SIZE_X - 1U) /
+            BLOCK_TILE_SIZE_X,
+        (static_cast<unsigned int>(m) + BLOCK_TILE_SIZE_Y - 1U) /
+            BLOCK_TILE_SIZE_Y,
+        1U};
+    gemm_v03_vectorized<T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y,
+                        BLOCK_TILE_SIZE_K, THREAD_TILE_SIZE_Y>
+        <<<grid_dim, block_dim, 0U, stream>>>(m, n, k, *alpha, A, lda, B, ldb,
+                                              *beta, C, ldc);
+    CHECK_LAST_CUDA_ERROR();
+}
+
+// Explicit instantiation.
+template void launch_gemm_kernel_v03_vectorized<float>(
+    size_t m, size_t n, size_t k, float const* alpha, float const* A,
+    size_t lda, float const* B, size_t ldb, float const* beta, float* C,
+    size_t ldc, cudaStream_t stream);
+template void launch_gemm_kernel_v03_vectorized<double>(
+    size_t m, size_t n, size_t k, double const* alpha, double const* A,
+    size_t lda, double const* B, size_t ldb, double const* beta, double* C,
+    size_t ldc, cudaStream_t stream);
+template void launch_gemm_kernel_v03_vectorized<__half>(
+    size_t m, size_t n, size_t k, __half const* alpha, __half const* A,
+    size_t lda, __half const* B, size_t ldb, __half const* beta, __half* C,
+    size_t ldc, cudaStream_t stream);

upstream_ref/cuda_gemm_optimization/04_2d_block_tiling_2d_thread_tiling.cu

@@ -0,0 +1,199 @@
+#include <cuda_fp16.h>
+
+#include "cuda_gemm.hpp"
+#include "cuda_gemm_utils.cuh"
+#include "cuda_gemm_utils.hpp"
+
+// GEMM kernel v04.
+// Coalesced read and write from global memory.
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t BLOCK_TILE_SIZE_K, size_t THREAD_TILE_SIZE_X,
+          size_t THREAD_TILE_SIZE_Y>
+__global__ void gemm_v04(size_t m, size_t n, size_t k, T alpha, T const* A,
+                         size_t lda, T const* B, size_t ldb, T beta, T* C,
+                         size_t ldc)
+{
+    // Avoid using blockDim.x * blockDim.y as the number of threads per block.
+    // Because it is a runtime constant and the compiler cannot optimize the
+    // loop unrolling based on that.
+    // Use a compile time constant instead.
+    constexpr size_t NUM_THREADS{BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y /
+                                 (THREAD_TILE_SIZE_X * THREAD_TILE_SIZE_Y)};
+    size_t const thread_linear_idx{threadIdx.y * blockDim.x + threadIdx.x};
+
+    // Cache a tile of A and B in shared memory for data reuse.
+    __shared__ T A_thread_block_tile[BLOCK_TILE_SIZE_Y][BLOCK_TILE_SIZE_K];
+    __shared__ T B_thread_block_tile[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_X];
+
+    size_t const num_thread_block_tiles{(k + BLOCK_TILE_SIZE_K - 1) /
+                                        BLOCK_TILE_SIZE_K};
+
+    // Each thread in the block processes BLOCK_TILE_SIZE_Y output values.
+    // Specifically, these values corresponds to
+    // C[blockIdx.y * BLOCK_TILE_SIZE_Y + threadIdx.x / BLOCK_TILE_SIZE_X *
+    // THREAD_TILE_SIZE_Y : blockIdx.y * BLOCK_TILE_SIZE_Y + (threadIdx.x /
+    // BLOCK_TILE_SIZE_X + 1) * THREAD_TILE_SIZE_Y][blockIdx.x *
+    // BLOCK_TILE_SIZE_X + threadIdx.x % BLOCK_TILE_SIZE_X *
+    // THREAD_TILE_SIZE_X : blockIdx.x * BLOCK_TILE_SIZE_X + (threadIdx.x %
+    // BLOCK_TILE_SIZE_X + 1) * THREAD_TILE_SIZE_X]
+    T C_thread_results[THREAD_TILE_SIZE_Y][THREAD_TILE_SIZE_X] = {
+        static_cast<T>(0)};
+    // A_vals is cached in the register.
+    T A_vals[THREAD_TILE_SIZE_Y] = {static_cast<T>(0)};
+    // B_vals is cached in the register.
+    T B_vals[THREAD_TILE_SIZE_X] = {static_cast<T>(0)};
+
+    for (size_t thread_block_tile_idx{0U};
+         thread_block_tile_idx < num_thread_block_tiles;
+         ++thread_block_tile_idx)
+    {
+
+        load_data_from_global_memory_to_shared_memory<
+            T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+            NUM_THREADS>(A, lda, B, ldb, A_thread_block_tile,
+                         B_thread_block_tile, thread_block_tile_idx,
+                         thread_linear_idx, m, n, k);
+        __syncthreads();
+
+#pragma unroll
+        for (size_t k_i{0U}; k_i < BLOCK_TILE_SIZE_K; ++k_i)
+        {
+            size_t const A_thread_block_tile_row_idx{
+                thread_linear_idx / (BLOCK_TILE_SIZE_X / THREAD_TILE_SIZE_X) *
+                THREAD_TILE_SIZE_Y};
+            size_t const A_thread_block_tile_col_idx{k_i};
+
+#pragma unroll
+            for (size_t thread_tile_row_idx{0U};
+                 thread_tile_row_idx < THREAD_TILE_SIZE_Y;
+                 ++thread_tile_row_idx)
+            {
+                // There will be shared memory bank conflicts accessing the
+                // values from A_thread_block_tile. We can do it better by
+                // transposing the A_thread_block_tile when we load the data
+                // from DRAM.
+                A_vals[thread_tile_row_idx] =
+                    A_thread_block_tile[A_thread_block_tile_row_idx +
+                                        thread_tile_row_idx]
+                                       [A_thread_block_tile_col_idx];
+            }
+
+            size_t const B_thread_block_tile_row_idx{k_i};
+            size_t const B_thread_block_tile_col_idx{
+                thread_linear_idx % (BLOCK_TILE_SIZE_X / THREAD_TILE_SIZE_X) *
+                THREAD_TILE_SIZE_X};
+#pragma unroll
+            for (size_t thread_tile_col_idx{0U};
+                 thread_tile_col_idx < THREAD_TILE_SIZE_X;
+                 ++thread_tile_col_idx)
+            {
+                B_vals[thread_tile_col_idx] =
+                    B_thread_block_tile[B_thread_block_tile_row_idx]
+                                       [B_thread_block_tile_col_idx +
+                                        thread_tile_col_idx];
+            }
+
+            for (size_t thread_tile_row_idx{0U};
+                 thread_tile_row_idx < THREAD_TILE_SIZE_Y;
+                 ++thread_tile_row_idx)
+            {
+                for (size_t thread_tile_col_idx{0U};
+                     thread_tile_col_idx < THREAD_TILE_SIZE_X;
+                     ++thread_tile_col_idx)
+                {
+                    C_thread_results[thread_tile_row_idx]
+                                    [thread_tile_col_idx] +=
+                        A_vals[thread_tile_row_idx] *
+                        B_vals[thread_tile_col_idx];
+                }
+            }
+        }
+        __syncthreads();
+    }
+
+    // Write the results to DRAM.
+    for (size_t thread_tile_row_idx{0U};
+         thread_tile_row_idx < THREAD_TILE_SIZE_Y; ++thread_tile_row_idx)
+    {
+        for (size_t thread_tile_col_idx{0U};
+             thread_tile_col_idx < THREAD_TILE_SIZE_X; ++thread_tile_col_idx)
+        {
+            size_t const C_row_idx{
+                blockIdx.y * BLOCK_TILE_SIZE_Y +
+                threadIdx.x / (BLOCK_TILE_SIZE_X / THREAD_TILE_SIZE_X) *
+                    THREAD_TILE_SIZE_Y +
+                thread_tile_row_idx};
+            size_t const C_col_idx{
+                blockIdx.x * BLOCK_TILE_SIZE_X +
+                threadIdx.x % (BLOCK_TILE_SIZE_X / THREAD_TILE_SIZE_X) *
+                    THREAD_TILE_SIZE_X +
+                thread_tile_col_idx};
+            if (C_row_idx < m && C_col_idx < n)
+            {
+                C[C_row_idx * ldc + C_col_idx] =
+                    alpha * C_thread_results[thread_tile_row_idx]
+                                            [thread_tile_col_idx] +
+                    beta * C[C_row_idx * ldc + C_col_idx];
+            }
+        }
+    }
+}
+
+template <typename T>
+void launch_gemm_kernel_v04(size_t m, size_t n, size_t k, T const* alpha,
+                            T const* A, size_t lda, T const* B, size_t ldb,
+                            T const* beta, T* C, size_t ldc,
+                            cudaStream_t stream)
+{
+    // Feel free to play with the block tile sizes.
+    // The algorithm correctness should always be guaranteed.
+    constexpr unsigned int BLOCK_TILE_SIZE_X{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_Y{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_K{16U};
+    // Each thread computes THREAD_TILE_SIZE_X * THREAD_TILE_SIZE_Y values of C.
+    constexpr unsigned int THREAD_TILE_SIZE_X{8U};
+    constexpr unsigned int THREAD_TILE_SIZE_Y{8U};
+    constexpr unsigned int NUM_THREADS_PER_BLOCK{
+        BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y /
+        (THREAD_TILE_SIZE_X * THREAD_TILE_SIZE_Y)};
+    static_assert(BLOCK_TILE_SIZE_X % THREAD_TILE_SIZE_X == 0U);
+    static_assert(BLOCK_TILE_SIZE_Y % THREAD_TILE_SIZE_Y == 0U);
+    static_assert(NUM_THREADS_PER_BLOCK % BLOCK_TILE_SIZE_K == 0U);
+    static_assert(NUM_THREADS_PER_BLOCK % BLOCK_TILE_SIZE_X == 0U);
+    static_assert(
+        BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_K % NUM_THREADS_PER_BLOCK == 0U);
+    static_assert(
+        BLOCK_TILE_SIZE_K * BLOCK_TILE_SIZE_Y % NUM_THREADS_PER_BLOCK == 0U);
+    dim3 const block_dim{NUM_THREADS_PER_BLOCK, 1U, 1U};
+    dim3 const grid_dim{
+        (static_cast<unsigned int>(n) + BLOCK_TILE_SIZE_X - 1U) /
+            BLOCK_TILE_SIZE_X,
+        (static_cast<unsigned int>(m) + BLOCK_TILE_SIZE_Y - 1U) /
+            BLOCK_TILE_SIZE_Y,
+        1U};
+    gemm_v04<T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+             THREAD_TILE_SIZE_X, THREAD_TILE_SIZE_Y>
+        <<<grid_dim, block_dim, 0U, stream>>>(m, n, k, *alpha, A, lda, B, ldb,
+                                              *beta, C, ldc);
+    CHECK_LAST_CUDA_ERROR();
+}
+
+// Explicit instantiation.
+template void launch_gemm_kernel_v04<float>(size_t m, size_t n, size_t k,
+                                            float const* alpha, float const* A,
+                                            size_t lda, float const* B,
+                                            size_t ldb, float const* beta,
+                                            float* C, size_t ldc,
+                                            cudaStream_t stream);
+template void launch_gemm_kernel_v04<double>(size_t m, size_t n, size_t k,
+                                             double const* alpha,
+                                             double const* A, size_t lda,
+                                             double const* B, size_t ldb,
+                                             double const* beta, double* C,
+                                             size_t ldc, cudaStream_t stream);
+template void launch_gemm_kernel_v04<__half>(size_t m, size_t n, size_t k,
+                                             __half const* alpha,
+                                             __half const* A, size_t lda,
+                                             __half const* B, size_t ldb,
+                                             __half const* beta, __half* C,
+                                             size_t ldc, cudaStream_t stream);

upstream_ref/cuda_gemm_optimization/04_2d_block_tiling_2d_thread_tiling_vectorized_memory_access.cu

@@ -0,0 +1,225 @@
+#include <cuda_fp16.h>
+
+#include "cuda_gemm.hpp"
+#include "cuda_gemm_utils.cuh"
+#include "cuda_gemm_utils.hpp"
+
+// GEMM kernel v04.
+// Coalesced read and write from global memory.
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t BLOCK_TILE_SIZE_K, size_t THREAD_TILE_SIZE_X,
+          size_t THREAD_TILE_SIZE_Y>
+__global__ void gemm_v04_vectorized(size_t m, size_t n, size_t k, T alpha,
+                                    T const* A, size_t lda, T const* B,
+                                    size_t ldb, T beta, T* C, size_t ldc)
+{
+    // Avoid using blockDim.x * blockDim.y as the number of threads per block.
+    // Because it is a runtime constant and the compiler cannot optimize the
+    // loop unrolling based on that.
+    // Use a compile time constant instead.
+    constexpr size_t NUM_THREADS{BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y /
+                                 (THREAD_TILE_SIZE_X * THREAD_TILE_SIZE_Y)};
+    size_t const thread_linear_idx{threadIdx.y * blockDim.x + threadIdx.x};
+
+    // Cache a tile of A and B in shared memory for data reuse.
+    __shared__ T A_thread_block_tile[BLOCK_TILE_SIZE_Y][BLOCK_TILE_SIZE_K];
+    __shared__ T B_thread_block_tile[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_X];
+
+    size_t const num_thread_block_tiles{(k + BLOCK_TILE_SIZE_K - 1) /
+                                        BLOCK_TILE_SIZE_K};
+
+    // Each thread in the block processes BLOCK_TILE_SIZE_Y output values.
+    // Specifically, these values corresponds to
+    // C[blockIdx.y * BLOCK_TILE_SIZE_Y + threadIdx.x / BLOCK_TILE_SIZE_X *
+    // THREAD_TILE_SIZE_Y : blockIdx.y * BLOCK_TILE_SIZE_Y + (threadIdx.x /
+    // BLOCK_TILE_SIZE_X + 1) * THREAD_TILE_SIZE_Y][blockIdx.x *
+    // BLOCK_TILE_SIZE_X + threadIdx.x % BLOCK_TILE_SIZE_X *
+    // THREAD_TILE_SIZE_X : blockIdx.x * BLOCK_TILE_SIZE_X + (threadIdx.x %
+    // BLOCK_TILE_SIZE_X + 1) * THREAD_TILE_SIZE_X]
+    T C_thread_results[THREAD_TILE_SIZE_Y][THREAD_TILE_SIZE_X] = {
+        static_cast<T>(0)};
+    // A_vals is cached in the register.
+    T A_vals[THREAD_TILE_SIZE_Y] = {static_cast<T>(0)};
+    // B_vals is cached in the register.
+    T B_vals[THREAD_TILE_SIZE_X] = {static_cast<T>(0)};
+
+    constexpr size_t NUM_VECTOR_UNITS{sizeof(int4) / sizeof(T)};
+    static_assert(sizeof(int4) % sizeof(T) == 0U);
+    static_assert(BLOCK_TILE_SIZE_K % NUM_VECTOR_UNITS == 0U);
+    static_assert(BLOCK_TILE_SIZE_X % NUM_VECTOR_UNITS == 0U);
+    constexpr size_t VECTORIZED_THREAD_TILE_SIZE_X{THREAD_TILE_SIZE_X /
+                                                   NUM_VECTOR_UNITS};
+    static_assert(THREAD_TILE_SIZE_X % NUM_VECTOR_UNITS == 0U);
+
+    for (size_t thread_block_tile_idx{0U};
+         thread_block_tile_idx < num_thread_block_tiles;
+         ++thread_block_tile_idx)
+    {
+        load_data_from_global_memory_to_shared_memory<
+            T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+            NUM_THREADS>(A, lda, B, ldb, A_thread_block_tile,
+                         B_thread_block_tile, thread_block_tile_idx,
+                         thread_linear_idx, m, n, k);
+        __syncthreads();
+
+#pragma unroll
+        for (size_t k_i{0U}; k_i < BLOCK_TILE_SIZE_K; ++k_i)
+        {
+            size_t const A_thread_block_tile_row_idx{
+                thread_linear_idx / (BLOCK_TILE_SIZE_X / THREAD_TILE_SIZE_X) *
+                THREAD_TILE_SIZE_Y};
+            size_t const A_thread_block_tile_col_idx{k_i};
+
+#pragma unroll
+            for (size_t thread_tile_row_idx{0U};
+                 thread_tile_row_idx < THREAD_TILE_SIZE_Y;
+                 ++thread_tile_row_idx)
+            {
+                // There will be shared memory bank conflicts accessing the
+                // values from A_thread_block_tile. We can do it better by
+                // transposing the A_thread_block_tile when we load the data
+                // from DRAM.
+                A_vals[thread_tile_row_idx] =
+                    A_thread_block_tile[A_thread_block_tile_row_idx +
+                                        thread_tile_row_idx]
+                                       [A_thread_block_tile_col_idx];
+            }
+
+            size_t const B_thread_block_tile_row_idx{k_i};
+            size_t const B_thread_block_tile_col_idx{
+                thread_linear_idx % (BLOCK_TILE_SIZE_X / THREAD_TILE_SIZE_X) *
+                THREAD_TILE_SIZE_X};
+// Although the read from A_thread_block_tile cannot be vectorized, the read
+// from B_thread_block_tile can be vectorized.
+#pragma unroll
+            for (size_t thread_tile_col_vector_idx{0U};
+                 thread_tile_col_vector_idx < VECTORIZED_THREAD_TILE_SIZE_X;
+                 ++thread_tile_col_vector_idx)
+            {
+                *reinterpret_cast<int4*>(
+                    &B_vals[thread_tile_col_vector_idx * NUM_VECTOR_UNITS]) =
+                    *reinterpret_cast<int4 const*>(
+                        &B_thread_block_tile[B_thread_block_tile_row_idx]
+                                            [B_thread_block_tile_col_idx +
+                                             thread_tile_col_vector_idx *
+                                                 NUM_VECTOR_UNITS]);
+            }
+
+            for (size_t thread_tile_row_idx{0U};
+                 thread_tile_row_idx < THREAD_TILE_SIZE_Y;
+                 ++thread_tile_row_idx)
+            {
+                for (size_t thread_tile_col_idx{0U};
+                     thread_tile_col_idx < THREAD_TILE_SIZE_X;
+                     ++thread_tile_col_idx)
+                {
+                    C_thread_results[thread_tile_row_idx]
+                                    [thread_tile_col_idx] +=
+                        A_vals[thread_tile_row_idx] *
+                        B_vals[thread_tile_col_idx];
+                }
+            }
+        }
+        __syncthreads();
+    }
+
+    // Vectorized writing the results to DRAM.
+    for (size_t thread_tile_row_idx{0U};
+         thread_tile_row_idx < THREAD_TILE_SIZE_Y; ++thread_tile_row_idx)
+    {
+        for (size_t thread_tile_col_vector_idx{0U};
+             thread_tile_col_vector_idx < VECTORIZED_THREAD_TILE_SIZE_X;
+             ++thread_tile_col_vector_idx)
+        {
+            size_t const C_row_idx{
+                blockIdx.y * BLOCK_TILE_SIZE_Y +
+                thread_linear_idx / (BLOCK_TILE_SIZE_X / THREAD_TILE_SIZE_X) *
+                    THREAD_TILE_SIZE_Y +
+                thread_tile_row_idx};
+            size_t const C_col_idx{
+                blockIdx.x * BLOCK_TILE_SIZE_X +
+                thread_linear_idx % (BLOCK_TILE_SIZE_X / THREAD_TILE_SIZE_X) *
+                    THREAD_TILE_SIZE_X +
+                thread_tile_col_vector_idx * NUM_VECTOR_UNITS};
+            // Vectorized read from C.
+            int4 C_row_vector_vals{*reinterpret_cast<int4 const*>(
+                &C[C_row_idx * ldc + C_col_idx])};
+            // Vectorized read from C_thread_results.
+            int4 const C_thread_results_row_vector_vals{
+                *reinterpret_cast<int4 const*>(
+                    &C_thread_results[thread_tile_row_idx]
+                                     [thread_tile_col_vector_idx *
+                                      NUM_VECTOR_UNITS])};
+            // Update the values in C_row_vector_vals
+            for (size_t i{0U}; i < NUM_VECTOR_UNITS; ++i)
+            {
+                reinterpret_cast<T*>(&C_row_vector_vals)[i] =
+                    alpha * reinterpret_cast<T const*>(
+                                &C_thread_results_row_vector_vals)[i] +
+                    beta * reinterpret_cast<T const*>(&C_row_vector_vals)[i];
+            }
+            // Vectorized write to C.
+            if (C_row_idx < m && C_col_idx < n)
+            {
+                // No need to mask out the out-of-bound invalid elements,
+                // because the row of C matrix is 32-byte aligned.
+                *reinterpret_cast<int4*>(&C[C_row_idx * ldc + C_col_idx]) =
+                    C_row_vector_vals;
+            }
+        }
+    }
+}
+
+template <typename T>
+void launch_gemm_kernel_v04_vectorized(size_t m, size_t n, size_t k,
+                                       T const* alpha, T const* A, size_t lda,
+                                       T const* B, size_t ldb, T const* beta,
+                                       T* C, size_t ldc, cudaStream_t stream)
+{
+    // Feel free to play with the block tile sizes.
+    // The algorithm correctness should always be guaranteed.
+    constexpr unsigned int BLOCK_TILE_SIZE_X{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_Y{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_K{16U};
+    // Each thread computes THREAD_TILE_SIZE_X * THREAD_TILE_SIZE_Y values of C.
+    constexpr unsigned int THREAD_TILE_SIZE_X{8U};
+    constexpr unsigned int THREAD_TILE_SIZE_Y{8U};
+    constexpr unsigned int NUM_THREADS_PER_BLOCK{
+        BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y /
+        (THREAD_TILE_SIZE_X * THREAD_TILE_SIZE_Y)};
+    static_assert(BLOCK_TILE_SIZE_X % THREAD_TILE_SIZE_X == 0U);
+    static_assert(BLOCK_TILE_SIZE_Y % THREAD_TILE_SIZE_Y == 0U);
+    static_assert(NUM_THREADS_PER_BLOCK % BLOCK_TILE_SIZE_K == 0U);
+    static_assert(NUM_THREADS_PER_BLOCK % BLOCK_TILE_SIZE_X == 0U);
+    static_assert(
+        BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_K % NUM_THREADS_PER_BLOCK == 0U);
+    static_assert(
+        BLOCK_TILE_SIZE_K * BLOCK_TILE_SIZE_Y % NUM_THREADS_PER_BLOCK == 0U);
+    dim3 const block_dim{NUM_THREADS_PER_BLOCK, 1U, 1U};
+    dim3 const grid_dim{
+        (static_cast<unsigned int>(n) + BLOCK_TILE_SIZE_X - 1U) /
+            BLOCK_TILE_SIZE_X,
+        (static_cast<unsigned int>(m) + BLOCK_TILE_SIZE_Y - 1U) /
+            BLOCK_TILE_SIZE_Y,
+        1U};
+    gemm_v04_vectorized<T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y,
+                        BLOCK_TILE_SIZE_K, THREAD_TILE_SIZE_X,
+                        THREAD_TILE_SIZE_Y>
+        <<<grid_dim, block_dim, 0U, stream>>>(m, n, k, *alpha, A, lda, B, ldb,
+                                              *beta, C, ldc);
+    CHECK_LAST_CUDA_ERROR();
+}
+
+// Explicit instantiation.
+template void launch_gemm_kernel_v04_vectorized<float>(
+    size_t m, size_t n, size_t k, float const* alpha, float const* A,
+    size_t lda, float const* B, size_t ldb, float const* beta, float* C,
+    size_t ldc, cudaStream_t stream);
+template void launch_gemm_kernel_v04_vectorized<double>(
+    size_t m, size_t n, size_t k, double const* alpha, double const* A,
+    size_t lda, double const* B, size_t ldb, double const* beta, double* C,
+    size_t ldc, cudaStream_t stream);
+template void launch_gemm_kernel_v04_vectorized<__half>(
+    size_t m, size_t n, size_t k, __half const* alpha, __half const* A,
+    size_t lda, __half const* B, size_t ldb, __half const* beta, __half* C,
+    size_t ldc, cudaStream_t stream);

upstream_ref/cuda_gemm_optimization/05_2d_block_tiling_2d_thread_tiling_matrix_transpose.cu

@@ -0,0 +1,196 @@
+#include <cuda_fp16.h>
+
+#include "cuda_gemm.hpp"
+#include "cuda_gemm_utils.cuh"
+#include "cuda_gemm_utils.hpp"
+
+// GEMM kernel v05.
+// Coalesced read and write from global memory.
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t BLOCK_TILE_SIZE_K, size_t THREAD_TILE_SIZE_X,
+          size_t THREAD_TILE_SIZE_Y>
+__global__ void gemm_v05(size_t m, size_t n, size_t k, T alpha, T const* A,
+                         size_t lda, T const* B, size_t ldb, T beta, T* C,
+                         size_t ldc)
+{
+    // Avoid using blockDim.x * blockDim.y as the number of threads per block.
+    // Because it is a runtime constant and the compiler cannot optimize the
+    // loop unrolling based on that.
+    // Use a compile time constant instead.
+    constexpr size_t NUM_THREADS{BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y /
+                                 (THREAD_TILE_SIZE_X * THREAD_TILE_SIZE_Y)};
+    size_t const thread_linear_idx{threadIdx.y * blockDim.x + threadIdx.x};
+
+    // Cache a tile of A and B in shared memory for data reuse.
+    __shared__ T
+        A_thread_block_tile_transposed[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_Y];
+    __shared__ T B_thread_block_tile[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_X];
+
+    size_t const num_thread_block_tiles{(k + BLOCK_TILE_SIZE_K - 1) /
+                                        BLOCK_TILE_SIZE_K};
+
+    // Each thread in the block processes BLOCK_TILE_SIZE_Y output values.
+    // Specifically, these values corresponds to
+    // C[blockIdx.y * BLOCK_TILE_SIZE_Y + threadIdx.x / BLOCK_TILE_SIZE_X *
+    // THREAD_TILE_SIZE_Y : blockIdx.y * BLOCK_TILE_SIZE_Y + (threadIdx.x /
+    // BLOCK_TILE_SIZE_X + 1) * THREAD_TILE_SIZE_Y][blockIdx.x *
+    // BLOCK_TILE_SIZE_X + threadIdx.x % BLOCK_TILE_SIZE_X *
+    // THREAD_TILE_SIZE_X : blockIdx.x * BLOCK_TILE_SIZE_X + (threadIdx.x %
+    // BLOCK_TILE_SIZE_X + 1) * THREAD_TILE_SIZE_X]
+    T C_thread_results[THREAD_TILE_SIZE_Y][THREAD_TILE_SIZE_X] = {
+        static_cast<T>(0)};
+    // A_vals is cached in the register.
+    T A_vals[THREAD_TILE_SIZE_Y] = {static_cast<T>(0)};
+    // B_vals is cached in the register.
+    T B_vals[THREAD_TILE_SIZE_X] = {static_cast<T>(0)};
+
+    for (size_t thread_block_tile_idx{0U};
+         thread_block_tile_idx < num_thread_block_tiles;
+         ++thread_block_tile_idx)
+    {
+
+        load_data_from_global_memory_to_shared_memory_transposed<
+            T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+            NUM_THREADS>(A, lda, B, ldb, A_thread_block_tile_transposed,
+                         B_thread_block_tile, thread_block_tile_idx,
+                         thread_linear_idx, m, n, k);
+        __syncthreads();
+
+#pragma unroll
+        for (size_t k_i{0U}; k_i < BLOCK_TILE_SIZE_K; ++k_i)
+        {
+            size_t const A_thread_block_tile_row_idx{
+                thread_linear_idx / (BLOCK_TILE_SIZE_X / THREAD_TILE_SIZE_X) *
+                THREAD_TILE_SIZE_Y};
+            size_t const A_thread_block_tile_col_idx{k_i};
+
+#pragma unroll
+            for (size_t thread_tile_row_idx{0U};
+                 thread_tile_row_idx < THREAD_TILE_SIZE_Y;
+                 ++thread_tile_row_idx)
+            {
+                A_vals[thread_tile_row_idx] =
+                    A_thread_block_tile_transposed[A_thread_block_tile_col_idx]
+                                                  [A_thread_block_tile_row_idx +
+                                                   thread_tile_row_idx];
+            }
+
+            size_t const B_thread_block_tile_row_idx{k_i};
+            size_t const B_thread_block_tile_col_idx{
+                thread_linear_idx % (BLOCK_TILE_SIZE_X / THREAD_TILE_SIZE_X) *
+                THREAD_TILE_SIZE_X};
+#pragma unroll
+            for (size_t thread_tile_col_idx{0U};
+                 thread_tile_col_idx < THREAD_TILE_SIZE_X;
+                 ++thread_tile_col_idx)
+            {
+                B_vals[thread_tile_col_idx] =
+                    B_thread_block_tile[B_thread_block_tile_row_idx]
+                                       [B_thread_block_tile_col_idx +
+                                        thread_tile_col_idx];
+            }
+
+            for (size_t thread_tile_row_idx{0U};
+                 thread_tile_row_idx < THREAD_TILE_SIZE_Y;
+                 ++thread_tile_row_idx)
+            {
+                for (size_t thread_tile_col_idx{0U};
+                     thread_tile_col_idx < THREAD_TILE_SIZE_X;
+                     ++thread_tile_col_idx)
+                {
+                    C_thread_results[thread_tile_row_idx]
+                                    [thread_tile_col_idx] +=
+                        A_vals[thread_tile_row_idx] *
+                        B_vals[thread_tile_col_idx];
+                }
+            }
+        }
+        __syncthreads();
+    }
+
+    // Write the results to DRAM.
+    for (size_t thread_tile_row_idx{0U};
+         thread_tile_row_idx < THREAD_TILE_SIZE_Y; ++thread_tile_row_idx)
+    {
+        for (size_t thread_tile_col_idx{0U};
+             thread_tile_col_idx < THREAD_TILE_SIZE_X; ++thread_tile_col_idx)
+        {
+            size_t const C_row_idx{
+                blockIdx.y * BLOCK_TILE_SIZE_Y +
+                threadIdx.x / (BLOCK_TILE_SIZE_X / THREAD_TILE_SIZE_X) *
+                    THREAD_TILE_SIZE_Y +
+                thread_tile_row_idx};
+            size_t const C_col_idx{
+                blockIdx.x * BLOCK_TILE_SIZE_X +
+                threadIdx.x % (BLOCK_TILE_SIZE_X / THREAD_TILE_SIZE_X) *
+                    THREAD_TILE_SIZE_X +
+                thread_tile_col_idx};
+            if (C_row_idx < m && C_col_idx < n)
+            {
+                C[C_row_idx * ldc + C_col_idx] =
+                    alpha * C_thread_results[thread_tile_row_idx]
+                                            [thread_tile_col_idx] +
+                    beta * C[C_row_idx * ldc + C_col_idx];
+            }
+        }
+    }
+}
+
+template <typename T>
+void launch_gemm_kernel_v05(size_t m, size_t n, size_t k, T const* alpha,
+                            T const* A, size_t lda, T const* B, size_t ldb,
+                            T const* beta, T* C, size_t ldc,
+                            cudaStream_t stream)
+{
+    // Feel free to play with the block tile sizes.
+    // The algorithm correctness should always be guaranteed.
+    constexpr unsigned int BLOCK_TILE_SIZE_X{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_Y{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_K{16U};
+    // Each thread computes THREAD_TILE_SIZE_X * THREAD_TILE_SIZE_Y values of C.
+    constexpr unsigned int THREAD_TILE_SIZE_X{8U};
+    constexpr unsigned int THREAD_TILE_SIZE_Y{8U};
+    constexpr unsigned int NUM_THREADS_PER_BLOCK{
+        BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y /
+        (THREAD_TILE_SIZE_X * THREAD_TILE_SIZE_Y)};
+    static_assert(BLOCK_TILE_SIZE_X % THREAD_TILE_SIZE_X == 0U);
+    static_assert(BLOCK_TILE_SIZE_Y % THREAD_TILE_SIZE_Y == 0U);
+    static_assert(NUM_THREADS_PER_BLOCK % BLOCK_TILE_SIZE_K == 0U);
+    static_assert(NUM_THREADS_PER_BLOCK % BLOCK_TILE_SIZE_X == 0U);
+    static_assert(
+        BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_K % NUM_THREADS_PER_BLOCK == 0U);
+    static_assert(
+        BLOCK_TILE_SIZE_K * BLOCK_TILE_SIZE_Y % NUM_THREADS_PER_BLOCK == 0U);
+    dim3 const block_dim{NUM_THREADS_PER_BLOCK, 1U, 1U};
+    dim3 const grid_dim{
+        (static_cast<unsigned int>(n) + BLOCK_TILE_SIZE_X - 1U) /
+            BLOCK_TILE_SIZE_X,
+        (static_cast<unsigned int>(m) + BLOCK_TILE_SIZE_Y - 1U) /
+            BLOCK_TILE_SIZE_Y,
+        1U};
+    gemm_v05<T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+             THREAD_TILE_SIZE_X, THREAD_TILE_SIZE_Y>
+        <<<grid_dim, block_dim, 0U, stream>>>(m, n, k, *alpha, A, lda, B, ldb,
+                                              *beta, C, ldc);
+    CHECK_LAST_CUDA_ERROR();
+}
+
+// Explicit instantiation.
+template void launch_gemm_kernel_v05<float>(size_t m, size_t n, size_t k,
+                                            float const* alpha, float const* A,
+                                            size_t lda, float const* B,
+                                            size_t ldb, float const* beta,
+                                            float* C, size_t ldc,
+                                            cudaStream_t stream);
+template void launch_gemm_kernel_v05<double>(size_t m, size_t n, size_t k,
+                                             double const* alpha,
+                                             double const* A, size_t lda,
+                                             double const* B, size_t ldb,
+                                             double const* beta, double* C,
+                                             size_t ldc, cudaStream_t stream);
+template void launch_gemm_kernel_v05<__half>(size_t m, size_t n, size_t k,
+                                             __half const* alpha,
+                                             __half const* A, size_t lda,
+                                             __half const* B, size_t ldb,
+                                             __half const* beta, __half* C,
+                                             size_t ldc, cudaStream_t stream);

upstream_ref/cuda_gemm_optimization/05_2d_block_tiling_2d_thread_tiling_matrix_transpose_vectorized_memory_access.cu

@@ -0,0 +1,222 @@
+#include <cuda_fp16.h>
+
+#include "cuda_gemm.hpp"
+#include "cuda_gemm_utils.cuh"
+#include "cuda_gemm_utils.hpp"
+
+// GEMM kernel v05.
+// Coalesced read and write from global memory.
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t BLOCK_TILE_SIZE_K, size_t THREAD_TILE_SIZE_X,
+          size_t THREAD_TILE_SIZE_Y>
+__global__ void gemm_v05_vectorized(size_t m, size_t n, size_t k, T alpha,
+                                    T const* A, size_t lda, T const* B,
+                                    size_t ldb, T beta, T* C, size_t ldc)
+{
+    // Avoid using blockDim.x * blockDim.y as the number of threads per block.
+    // Because it is a runtime constant and the compiler cannot optimize the
+    // loop unrolling based on that.
+    // Use a compile time constant instead.
+    constexpr size_t NUM_THREADS{BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y /
+                                 (THREAD_TILE_SIZE_X * THREAD_TILE_SIZE_Y)};
+    size_t const thread_linear_idx{threadIdx.y * blockDim.x + threadIdx.x};
+
+    // Cache a tile of A and B in shared memory for data reuse.
+    __shared__ T
+        A_thread_block_tile_transposed[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_Y];
+    __shared__ T B_thread_block_tile[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_X];
+
+    size_t const num_thread_block_tiles{(k + BLOCK_TILE_SIZE_K - 1) /
+                                        BLOCK_TILE_SIZE_K};
+
+    // Each thread in the block processes BLOCK_TILE_SIZE_Y output values.
+    // Specifically, these values corresponds to
+    // C[blockIdx.y * BLOCK_TILE_SIZE_Y + threadIdx.x / BLOCK_TILE_SIZE_X *
+    // THREAD_TILE_SIZE_Y : blockIdx.y * BLOCK_TILE_SIZE_Y + (threadIdx.x /
+    // BLOCK_TILE_SIZE_X + 1) * THREAD_TILE_SIZE_Y][blockIdx.x *
+    // BLOCK_TILE_SIZE_X + threadIdx.x % BLOCK_TILE_SIZE_X *
+    // THREAD_TILE_SIZE_X : blockIdx.x * BLOCK_TILE_SIZE_X + (threadIdx.x %
+    // BLOCK_TILE_SIZE_X + 1) * THREAD_TILE_SIZE_X]
+    T C_thread_results[THREAD_TILE_SIZE_Y][THREAD_TILE_SIZE_X] = {
+        static_cast<T>(0)};
+    // A_vals is cached in the register.
+    T A_vals[THREAD_TILE_SIZE_Y] = {static_cast<T>(0)};
+    // B_vals is cached in the register.
+    T B_vals[THREAD_TILE_SIZE_X] = {static_cast<T>(0)};
+
+    constexpr size_t NUM_VECTOR_UNITS{sizeof(int4) / sizeof(T)};
+    static_assert(sizeof(int4) % sizeof(T) == 0U);
+    static_assert(BLOCK_TILE_SIZE_K % NUM_VECTOR_UNITS == 0U);
+    static_assert(BLOCK_TILE_SIZE_X % NUM_VECTOR_UNITS == 0U);
+    constexpr size_t VECTORIZED_THREAD_TILE_SIZE_X{THREAD_TILE_SIZE_X /
+                                                   NUM_VECTOR_UNITS};
+    static_assert(THREAD_TILE_SIZE_X % NUM_VECTOR_UNITS == 0U);
+
+    for (size_t thread_block_tile_idx{0U};
+         thread_block_tile_idx < num_thread_block_tiles;
+         ++thread_block_tile_idx)
+    {
+        load_data_from_global_memory_to_shared_memory_transposed_vectorized<
+            T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+            NUM_THREADS>(A, lda, B, ldb, A_thread_block_tile_transposed,
+                         B_thread_block_tile, thread_block_tile_idx,
+                         thread_linear_idx, m, n, k);
+        __syncthreads();
+
+#pragma unroll
+        for (size_t k_i{0U}; k_i < BLOCK_TILE_SIZE_K; ++k_i)
+        {
+            size_t const A_thread_block_tile_row_idx{
+                thread_linear_idx / (BLOCK_TILE_SIZE_X / THREAD_TILE_SIZE_X) *
+                THREAD_TILE_SIZE_Y};
+            size_t const A_thread_block_tile_col_idx{k_i};
+
+#pragma unroll
+            for (size_t thread_tile_row_idx{0U};
+                 thread_tile_row_idx < THREAD_TILE_SIZE_Y;
+                 ++thread_tile_row_idx)
+            {
+                A_vals[thread_tile_row_idx] =
+                    A_thread_block_tile_transposed[A_thread_block_tile_col_idx]
+                                                  [A_thread_block_tile_row_idx +
+                                                   thread_tile_row_idx];
+            }
+
+            size_t const B_thread_block_tile_row_idx{k_i};
+            size_t const B_thread_block_tile_col_idx{
+                thread_linear_idx % (BLOCK_TILE_SIZE_X / THREAD_TILE_SIZE_X) *
+                THREAD_TILE_SIZE_X};
+// Although the read from A_thread_block_tile cannot be vectorized, the read
+// from B_thread_block_tile can be vectorized.
+#pragma unroll
+            for (size_t thread_tile_col_vector_idx{0U};
+                 thread_tile_col_vector_idx < VECTORIZED_THREAD_TILE_SIZE_X;
+                 ++thread_tile_col_vector_idx)
+            {
+                *reinterpret_cast<int4*>(
+                    &B_vals[thread_tile_col_vector_idx * NUM_VECTOR_UNITS]) =
+                    *reinterpret_cast<int4 const*>(
+                        &B_thread_block_tile[B_thread_block_tile_row_idx]
+                                            [B_thread_block_tile_col_idx +
+                                             thread_tile_col_vector_idx *
+                                                 NUM_VECTOR_UNITS]);
+            }
+
+            for (size_t thread_tile_row_idx{0U};
+                 thread_tile_row_idx < THREAD_TILE_SIZE_Y;
+                 ++thread_tile_row_idx)
+            {
+                for (size_t thread_tile_col_idx{0U};
+                     thread_tile_col_idx < THREAD_TILE_SIZE_X;
+                     ++thread_tile_col_idx)
+                {
+                    C_thread_results[thread_tile_row_idx]
+                                    [thread_tile_col_idx] +=
+                        A_vals[thread_tile_row_idx] *
+                        B_vals[thread_tile_col_idx];
+                }
+            }
+        }
+        __syncthreads();
+    }
+
+    // Vectorized writing the results to DRAM.
+    for (size_t thread_tile_row_idx{0U};
+         thread_tile_row_idx < THREAD_TILE_SIZE_Y; ++thread_tile_row_idx)
+    {
+        for (size_t thread_tile_col_vector_idx{0U};
+             thread_tile_col_vector_idx < VECTORIZED_THREAD_TILE_SIZE_X;
+             ++thread_tile_col_vector_idx)
+        {
+            size_t const C_row_idx{
+                blockIdx.y * BLOCK_TILE_SIZE_Y +
+                thread_linear_idx / (BLOCK_TILE_SIZE_X / THREAD_TILE_SIZE_X) *
+                    THREAD_TILE_SIZE_Y +
+                thread_tile_row_idx};
+            size_t const C_col_idx{
+                blockIdx.x * BLOCK_TILE_SIZE_X +
+                thread_linear_idx % (BLOCK_TILE_SIZE_X / THREAD_TILE_SIZE_X) *
+                    THREAD_TILE_SIZE_X +
+                thread_tile_col_vector_idx * NUM_VECTOR_UNITS};
+            // Vectorized read from C.
+            int4 C_row_vector_vals{*reinterpret_cast<int4 const*>(
+                &C[C_row_idx * ldc + C_col_idx])};
+            // Vectorized read from C_thread_results.
+            int4 const C_thread_results_row_vector_vals{
+                *reinterpret_cast<int4 const*>(
+                    &C_thread_results[thread_tile_row_idx]
+                                     [thread_tile_col_vector_idx *
+                                      NUM_VECTOR_UNITS])};
+            // Update the values in C_row_vector_vals
+            for (size_t i{0U}; i < NUM_VECTOR_UNITS; ++i)
+            {
+                reinterpret_cast<T*>(&C_row_vector_vals)[i] =
+                    alpha * reinterpret_cast<T const*>(
+                                &C_thread_results_row_vector_vals)[i] +
+                    beta * reinterpret_cast<T const*>(&C_row_vector_vals)[i];
+            }
+            // Vectorized write to C.
+            if (C_row_idx < m && C_col_idx < n)
+            {
+                // No need to mask out the out-of-bound invalid elements,
+                // because the row of C matrix is 32-byte aligned.
+                *reinterpret_cast<int4*>(&C[C_row_idx * ldc + C_col_idx]) =
+                    C_row_vector_vals;
+            }
+        }
+    }
+}
+
+template <typename T>
+void launch_gemm_kernel_v05_vectorized(size_t m, size_t n, size_t k,
+                                       T const* alpha, T const* A, size_t lda,
+                                       T const* B, size_t ldb, T const* beta,
+                                       T* C, size_t ldc, cudaStream_t stream)
+{
+    // Feel free to play with the block tile sizes.
+    // The algorithm correctness should always be guaranteed.
+    constexpr unsigned int BLOCK_TILE_SIZE_X{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_Y{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_K{16U};
+    // Each thread computes THREAD_TILE_SIZE_X * THREAD_TILE_SIZE_Y values of C.
+    constexpr unsigned int THREAD_TILE_SIZE_X{8U};
+    constexpr unsigned int THREAD_TILE_SIZE_Y{8U};
+    constexpr unsigned int NUM_THREADS_PER_BLOCK{
+        BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y /
+        (THREAD_TILE_SIZE_X * THREAD_TILE_SIZE_Y)};
+    static_assert(BLOCK_TILE_SIZE_X % THREAD_TILE_SIZE_X == 0U);
+    static_assert(BLOCK_TILE_SIZE_Y % THREAD_TILE_SIZE_Y == 0U);
+    static_assert(NUM_THREADS_PER_BLOCK % BLOCK_TILE_SIZE_K == 0U);
+    static_assert(NUM_THREADS_PER_BLOCK % BLOCK_TILE_SIZE_X == 0U);
+    static_assert(
+        BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_K % NUM_THREADS_PER_BLOCK == 0U);
+    static_assert(
+        BLOCK_TILE_SIZE_K * BLOCK_TILE_SIZE_Y % NUM_THREADS_PER_BLOCK == 0U);
+    dim3 const block_dim{NUM_THREADS_PER_BLOCK, 1U, 1U};
+    dim3 const grid_dim{
+        (static_cast<unsigned int>(n) + BLOCK_TILE_SIZE_X - 1U) /
+            BLOCK_TILE_SIZE_X,
+        (static_cast<unsigned int>(m) + BLOCK_TILE_SIZE_Y - 1U) /
+            BLOCK_TILE_SIZE_Y,
+        1U};
+    gemm_v05_vectorized<T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y,
+                        BLOCK_TILE_SIZE_K, THREAD_TILE_SIZE_X,
+                        THREAD_TILE_SIZE_Y>
+        <<<grid_dim, block_dim, 0U, stream>>>(m, n, k, *alpha, A, lda, B, ldb,
+                                              *beta, C, ldc);
+    CHECK_LAST_CUDA_ERROR();
+}
+
+// Explicit instantiation.
+template void launch_gemm_kernel_v05_vectorized<float>(
+    size_t m, size_t n, size_t k, float const* alpha, float const* A,
+    size_t lda, float const* B, size_t ldb, float const* beta, float* C,
+    size_t ldc, cudaStream_t stream);
+template void launch_gemm_kernel_v05_vectorized<double>(
+    size_t m, size_t n, size_t k, double const* alpha, double const* A,
+    size_t lda, double const* B, size_t ldb, double const* beta, double* C,
+    size_t ldc, cudaStream_t stream);
+template void launch_gemm_kernel_v05_vectorized<__half>(
+    size_t m, size_t n, size_t k, __half const* alpha, __half const* A,
+    size_t lda, __half const* B, size_t ldb, __half const* beta, __half* C,
+    size_t ldc, cudaStream_t stream);

upstream_ref/cuda_gemm_optimization/06_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose.cu

@@ -0,0 +1,334 @@
+#include <cuda_fp16.h>
+
+#include "cuda_gemm.hpp"
+#include "cuda_gemm_utils.cuh"
+#include "cuda_gemm_utils.hpp"
+
+template <typename T, size_t BLOCK_TILE_SIZE, size_t WARP_TILE_SIZE,
+          size_t NUM_THREAD_TILES_PER_WARP, size_t THREAD_TILE_SIZE>
+__device__ void load_data_from_shared_memory_to_register_file(
+    T const thread_block_tile[BLOCK_TILE_SIZE],
+    T register_values[NUM_THREAD_TILES_PER_WARP][THREAD_TILE_SIZE],
+    size_t warp_idx, size_t thread_idx)
+{
+    static_assert(BLOCK_TILE_SIZE % THREAD_TILE_SIZE == 0U);
+#pragma unroll
+    for (size_t thread_tile_repeat_idx{0U};
+         thread_tile_repeat_idx < NUM_THREAD_TILES_PER_WARP;
+         ++thread_tile_repeat_idx)
+    {
+        size_t const thread_block_tile_idx{
+            warp_idx * WARP_TILE_SIZE +
+            thread_tile_repeat_idx *
+                (WARP_TILE_SIZE / NUM_THREAD_TILES_PER_WARP) +
+            thread_idx * THREAD_TILE_SIZE};
+#pragma unroll
+        for (size_t thread_tile_idx{0U}; thread_tile_idx < THREAD_TILE_SIZE;
+             ++thread_tile_idx)
+        {
+            register_values[thread_tile_repeat_idx][thread_tile_idx] =
+                thread_block_tile[thread_block_tile_idx + thread_tile_idx];
+        }
+    }
+}
+
+template <typename T, size_t NUM_THREAD_TILES_PER_WARP_X,
+          size_t NUM_THREAD_TILES_PER_WARP_Y, size_t THREAD_TILE_SIZE_X,
+          size_t THREAD_TILE_SIZE_Y>
+__device__ void compute_thread_tile_results(
+    T const A_vals[NUM_THREAD_TILES_PER_WARP_Y][THREAD_TILE_SIZE_Y],
+    T const B_vals[NUM_THREAD_TILES_PER_WARP_X][THREAD_TILE_SIZE_X],
+    T C_thread_results[NUM_THREAD_TILES_PER_WARP_Y][NUM_THREAD_TILES_PER_WARP_X]
+                      [THREAD_TILE_SIZE_Y][THREAD_TILE_SIZE_X])
+{
+// Compute NUM_THREAD_TILES_PER_WARP_Y * NUM_THREAD_TILES_PER_WARP_X outer
+// products.
+#pragma unroll
+    for (size_t thread_tile_repeat_row_idx{0U};
+         thread_tile_repeat_row_idx < NUM_THREAD_TILES_PER_WARP_Y;
+         ++thread_tile_repeat_row_idx)
+    {
+#pragma unroll
+        for (size_t thread_tile_repeat_col_idx{0U};
+             thread_tile_repeat_col_idx < NUM_THREAD_TILES_PER_WARP_X;
+             ++thread_tile_repeat_col_idx)
+        {
+#pragma unroll
+            for (size_t thread_tile_y_idx{0U};
+                 thread_tile_y_idx < THREAD_TILE_SIZE_Y; ++thread_tile_y_idx)
+            {
+#pragma unroll
+                for (size_t thread_tile_x_idx{0U};
+                     thread_tile_x_idx < THREAD_TILE_SIZE_X;
+                     ++thread_tile_x_idx)
+                {
+                    C_thread_results[thread_tile_repeat_row_idx]
+                                    [thread_tile_repeat_col_idx]
+                                    [thread_tile_y_idx][thread_tile_x_idx] +=
+                        A_vals[thread_tile_repeat_row_idx][thread_tile_y_idx] *
+                        B_vals[thread_tile_repeat_col_idx][thread_tile_x_idx];
+                }
+            }
+        }
+    }
+}
+
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t WARP_TILE_SIZE_X, size_t WARP_TILE_SIZE_Y,
+          size_t THREAD_TILE_SIZE_X, size_t THREAD_TILE_SIZE_Y,
+          size_t NUM_THREAD_TILES_PER_WARP_X,
+          size_t NUM_THREAD_TILES_PER_WARP_Y>
+__device__ void write_results_from_register_file_to_global_memory(
+    T const C_thread_results[NUM_THREAD_TILES_PER_WARP_Y]
+                            [NUM_THREAD_TILES_PER_WARP_X][THREAD_TILE_SIZE_Y]
+                            [THREAD_TILE_SIZE_X],
+    T alpha, T beta, T* C, size_t ldc, size_t m, size_t n, size_t block_row_idx,
+    size_t block_col_idx, size_t warp_row_idx, size_t warp_col_idx,
+    size_t thread_row_idx_in_warp, size_t thread_col_idx_in_warp)
+{
+// Write the results to DRAM.
+#pragma unroll
+    for (size_t thread_tile_repeat_row_idx{0U};
+         thread_tile_repeat_row_idx < NUM_THREAD_TILES_PER_WARP_Y;
+         ++thread_tile_repeat_row_idx)
+    {
+#pragma unroll
+        for (size_t thread_tile_repeat_col_idx{0U};
+             thread_tile_repeat_col_idx < NUM_THREAD_TILES_PER_WARP_X;
+             ++thread_tile_repeat_col_idx)
+        {
+#pragma unroll
+            for (size_t thread_tile_y_idx{0U};
+                 thread_tile_y_idx < THREAD_TILE_SIZE_Y; ++thread_tile_y_idx)
+            {
+#pragma unroll
+                for (size_t thread_tile_x_idx{0U};
+                     thread_tile_x_idx < THREAD_TILE_SIZE_X;
+                     ++thread_tile_x_idx)
+                {
+                    size_t const C_row_idx{
+                        block_row_idx * BLOCK_TILE_SIZE_Y +
+                        warp_row_idx * WARP_TILE_SIZE_Y +
+                        thread_tile_repeat_row_idx *
+                            (WARP_TILE_SIZE_Y / NUM_THREAD_TILES_PER_WARP_Y) +
+                        thread_row_idx_in_warp * THREAD_TILE_SIZE_Y +
+                        thread_tile_y_idx};
+                    size_t const C_col_idx{
+                        block_col_idx * BLOCK_TILE_SIZE_X +
+                        warp_col_idx * WARP_TILE_SIZE_X +
+                        thread_tile_repeat_col_idx *
+                            (WARP_TILE_SIZE_X / NUM_THREAD_TILES_PER_WARP_X) +
+                        thread_col_idx_in_warp * THREAD_TILE_SIZE_X +
+                        thread_tile_x_idx};
+                    if (C_row_idx < m && C_col_idx < n)
+                    {
+                        C[C_row_idx * ldc + C_col_idx] =
+                            alpha * C_thread_results[thread_tile_repeat_row_idx]
+                                                    [thread_tile_repeat_col_idx]
+                                                    [thread_tile_y_idx]
+                                                    [thread_tile_x_idx] +
+                            beta * C[C_row_idx * ldc + C_col_idx];
+                    }
+                }
+            }
+        }
+    }
+}
+
+// GEMM kernel v06.
+// Each thread in the block processes THREAD_TILE_SIZE_Y *
+// THREAD_TILE_SIZE_X output values. Number of threads BLOCK_TILE_SIZE_Y *
+// BLOCK_TILE_SIZE_X / (THREAD_TILE_SIZE_Y * THREAD_TILE_SIZE_X)
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t BLOCK_TILE_SIZE_K, size_t WARP_TILE_SIZE_X,
+          size_t WARP_TILE_SIZE_Y, size_t THREAD_TILE_SIZE_X,
+          size_t THREAD_TILE_SIZE_Y, size_t NUM_THREADS_PER_WARP_X,
+          size_t NUM_THREADS_PER_WARP_Y>
+__global__ void gemm_v06(size_t m, size_t n, size_t k, T alpha, T const* A,
+                         size_t lda, T const* B, size_t ldb, T beta, T* C,
+                         size_t ldc)
+{
+    static_assert(NUM_THREADS_PER_WARP_X * NUM_THREADS_PER_WARP_Y == 32U);
+    constexpr size_t NUM_WARPS_X{BLOCK_TILE_SIZE_X / WARP_TILE_SIZE_X};
+    static_assert(BLOCK_TILE_SIZE_X % WARP_TILE_SIZE_X == 0U);
+    constexpr size_t NUM_WARPS_Y{BLOCK_TILE_SIZE_Y / WARP_TILE_SIZE_Y};
+    static_assert(BLOCK_TILE_SIZE_Y % WARP_TILE_SIZE_Y == 0U);
+    constexpr unsigned int NUM_THREAD_TILES_PER_WARP_X{
+        WARP_TILE_SIZE_X / (THREAD_TILE_SIZE_X * NUM_THREADS_PER_WARP_X)};
+    constexpr unsigned int NUM_THREAD_TILES_PER_WARP_Y{
+        WARP_TILE_SIZE_Y / (THREAD_TILE_SIZE_Y * NUM_THREADS_PER_WARP_Y)};
+    static_assert(
+        WARP_TILE_SIZE_X % (THREAD_TILE_SIZE_X * NUM_THREADS_PER_WARP_X) == 0U);
+    static_assert(
+        WARP_TILE_SIZE_Y % (THREAD_TILE_SIZE_Y * NUM_THREADS_PER_WARP_Y) == 0U);
+
+    constexpr unsigned int NUM_THREADS_X{NUM_WARPS_X * NUM_THREADS_PER_WARP_X};
+    constexpr unsigned int NUM_THREADS_Y{NUM_WARPS_Y * NUM_THREADS_PER_WARP_Y};
+    // Avoid using blockDim.x * blockDim.y as the number of threads per block.
+    // Because it is a runtime constant and the compiler cannot optimize the
+    // loop unrolling based on that.
+    // Use a compile time constant instead.
+    constexpr size_t NUM_THREADS{NUM_THREADS_X * NUM_THREADS_Y};
+
+    // Cache a tile of A and B in shared memory for data reuse.
+    __shared__ T
+        A_thread_block_tile_transposed[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_Y];
+    __shared__ T B_thread_block_tile[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_X];
+
+    // A_vals is cached in the register.
+    T A_vals[NUM_THREAD_TILES_PER_WARP_Y][THREAD_TILE_SIZE_Y] = {
+        static_cast<T>(0)};
+    // B_vals is cached in the register.
+    T B_vals[NUM_THREAD_TILES_PER_WARP_X][THREAD_TILE_SIZE_X] = {
+        static_cast<T>(0)};
+
+    size_t const thread_linear_idx{threadIdx.y * blockDim.x + threadIdx.x};
+    size_t const warp_linear_idx{thread_linear_idx / 32U};
+    size_t const warp_row_idx{warp_linear_idx / NUM_WARPS_X};
+    size_t const warp_col_idx{warp_linear_idx % NUM_WARPS_X};
+    size_t const thread_linear_idx_in_warp{thread_linear_idx % 32U};
+    size_t const thread_linear_row_idx_in_warp{thread_linear_idx_in_warp /
+                                               NUM_THREADS_PER_WARP_X};
+    size_t const thread_linear_col_idx_in_warp{thread_linear_idx_in_warp %
+                                               NUM_THREADS_PER_WARP_X};
+
+    // Number of outer loops to perform the sum of inner products.
+    // C_thread_block_tile =
+    // \sigma_{thread_block_tile_idx=0}^{num_thread_block_tiles-1} A[:,
+    // thread_block_tile_idx:BLOCK_TILE_SIZE_K] *
+    // B[thread_block_tile_idx:BLOCK_TILE_SIZE_K, :]
+    size_t const num_thread_block_tiles{(k + BLOCK_TILE_SIZE_K - 1) /
+                                        BLOCK_TILE_SIZE_K};
+    // Each thread in the block processes NUM_THREAD_TILES_PER_WARP_Y *
+    // NUM_THREAD_TILES_PER_WARP_X * THREAD_TILE_SIZE_Y *
+    // THREAD_TILE_SIZE_X output values.
+    T C_thread_results[NUM_THREAD_TILES_PER_WARP_Y][NUM_THREAD_TILES_PER_WARP_X]
+                      [THREAD_TILE_SIZE_Y][THREAD_TILE_SIZE_X] = {
+                          static_cast<T>(0)};
+
+    for (size_t thread_block_tile_idx{0U};
+         thread_block_tile_idx < num_thread_block_tiles;
+         ++thread_block_tile_idx)
+    {
+        load_data_from_global_memory_to_shared_memory_transposed<
+            T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+            NUM_THREADS>(A, lda, B, ldb, A_thread_block_tile_transposed,
+                         B_thread_block_tile, thread_block_tile_idx,
+                         thread_linear_idx, m, n, k);
+        __syncthreads();
+
+// Perform A[:, thread_block_tile_idx:BLOCK_TILE_SIZE_K] *
+// B[thread_block_tile_idx:BLOCK_TILE_SIZE_K, :] where A[:,
+// thread_block_tile_idx:BLOCK_TILE_SIZE_K] and
+// B[thread_block_tile_idx:BLOCK_TILE_SIZE_K, :] are cached in the
+// shared memory as A_thread_block_tile and B_thread_block_tile,
+// respectively. This inner product is further decomposed to
+// BLOCK_TILE_SIZE_K outer products. A_thread_block_tile *
+// B_thread_block_tile = \sigma_{k_i=0}^{BLOCK_TILE_SIZE_K-1}
+// A_thread_block_tile[:, k_i] @ B_thread_block_tile[k_i, :] Note that
+// both A_thread_block_tile and B_thread_block_tile can be cached in the
+// register.
+#pragma unroll
+        for (size_t k_i{0U}; k_i < BLOCK_TILE_SIZE_K; ++k_i)
+        {
+            // Load data from shared memory to register file for A.
+            load_data_from_shared_memory_to_register_file<
+                T, BLOCK_TILE_SIZE_Y, WARP_TILE_SIZE_Y, NUM_THREADS_PER_WARP_Y,
+                THREAD_TILE_SIZE_Y>(A_thread_block_tile_transposed[k_i], A_vals,
+                                    warp_row_idx,
+                                    thread_linear_row_idx_in_warp);
+            // Load data from shared memory to register file for B.
+            load_data_from_shared_memory_to_register_file<
+                T, BLOCK_TILE_SIZE_X, WARP_TILE_SIZE_X, NUM_THREADS_PER_WARP_X,
+                THREAD_TILE_SIZE_X>(B_thread_block_tile[k_i], B_vals,
+                                    warp_col_idx,
+                                    thread_linear_col_idx_in_warp);
+            // Compute NUM_THREAD_TILES_PER_WARP_Y * NUM_THREAD_TILES_PER_WARP_X
+            // outer products.
+            compute_thread_tile_results<T, NUM_THREAD_TILES_PER_WARP_X,
+                                        NUM_THREAD_TILES_PER_WARP_Y,
+                                        THREAD_TILE_SIZE_X, THREAD_TILE_SIZE_Y>(
+                A_vals, B_vals, C_thread_results);
+        }
+        __syncthreads();
+    }
+
+    // Write the results to DRAM.
+    write_results_from_register_file_to_global_memory<
+        T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, WARP_TILE_SIZE_X,
+        WARP_TILE_SIZE_Y, THREAD_TILE_SIZE_X, THREAD_TILE_SIZE_Y,
+        NUM_THREAD_TILES_PER_WARP_X, NUM_THREAD_TILES_PER_WARP_Y>(
+        C_thread_results, alpha, beta, C, ldc, m, n, blockIdx.y, blockIdx.x,
+        warp_row_idx, warp_col_idx, thread_linear_row_idx_in_warp,
+        thread_linear_col_idx_in_warp);
+}
+
+template <typename T>
+void launch_gemm_kernel_v06(size_t m, size_t n, size_t k, T const* alpha,
+                            T const* A, size_t lda, T const* B, size_t ldb,
+                            T const* beta, T* C, size_t ldc,
+                            cudaStream_t stream)
+{
+    // Feel free to play with the block tile sizes.
+    // The algorithm correctness should always be guaranteed.
+    constexpr unsigned int BLOCK_TILE_SIZE_X{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_Y{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_K{16U};
+
+    constexpr unsigned int WARP_TILE_SIZE_X{32U};
+    constexpr unsigned int WARP_TILE_SIZE_Y{64U};
+    constexpr unsigned int NUM_WARPS_X{BLOCK_TILE_SIZE_X / WARP_TILE_SIZE_X};
+    constexpr unsigned int NUM_WARPS_Y{BLOCK_TILE_SIZE_Y / WARP_TILE_SIZE_Y};
+    static_assert(BLOCK_TILE_SIZE_X % WARP_TILE_SIZE_X == 0U);
+    static_assert(BLOCK_TILE_SIZE_Y % WARP_TILE_SIZE_Y == 0U);
+
+    constexpr unsigned int THREAD_TILE_SIZE_X{8U};
+    constexpr unsigned int THREAD_TILE_SIZE_Y{8U};
+
+    constexpr unsigned int NUM_THREADS_PER_WARP_X{4U};
+    constexpr unsigned int NUM_THREADS_PER_WARP_Y{8U};
+    static_assert(NUM_THREADS_PER_WARP_X * NUM_THREADS_PER_WARP_Y == 32U);
+    static_assert(
+        WARP_TILE_SIZE_X % (THREAD_TILE_SIZE_X * NUM_THREADS_PER_WARP_X) == 0U);
+    static_assert(
+        WARP_TILE_SIZE_Y % (THREAD_TILE_SIZE_Y * NUM_THREADS_PER_WARP_Y) == 0U);
+
+    constexpr unsigned int NUM_THREADS_X{NUM_WARPS_X * NUM_THREADS_PER_WARP_X};
+    constexpr unsigned int NUM_THREADS_Y{NUM_WARPS_Y * NUM_THREADS_PER_WARP_Y};
+
+    constexpr unsigned int NUM_THREADS_PER_BLOCK{NUM_THREADS_X * NUM_THREADS_Y};
+
+    dim3 const block_dim{NUM_THREADS_PER_BLOCK, 1U, 1U};
+    dim3 const grid_dim{
+        (static_cast<unsigned int>(n) + BLOCK_TILE_SIZE_X - 1U) /
+            BLOCK_TILE_SIZE_X,
+        (static_cast<unsigned int>(m) + BLOCK_TILE_SIZE_Y - 1U) /
+            BLOCK_TILE_SIZE_Y,
+        1U};
+    gemm_v06<T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+             WARP_TILE_SIZE_X, WARP_TILE_SIZE_Y, THREAD_TILE_SIZE_X,
+             THREAD_TILE_SIZE_Y, NUM_THREADS_PER_WARP_X, NUM_THREADS_PER_WARP_Y>
+        <<<grid_dim, block_dim, 0U, stream>>>(m, n, k, *alpha, A, lda, B, ldb,
+                                              *beta, C, ldc);
+    CHECK_LAST_CUDA_ERROR();
+}
+
+// Explicit instantiation.
+template void launch_gemm_kernel_v06<float>(size_t m, size_t n, size_t k,
+                                            float const* alpha, float const* A,
+                                            size_t lda, float const* B,
+                                            size_t ldb, float const* beta,
+                                            float* C, size_t ldc,
+                                            cudaStream_t stream);
+template void launch_gemm_kernel_v06<double>(size_t m, size_t n, size_t k,
+                                             double const* alpha,
+                                             double const* A, size_t lda,
+                                             double const* B, size_t ldb,
+                                             double const* beta, double* C,
+                                             size_t ldc, cudaStream_t stream);
+template void launch_gemm_kernel_v06<__half>(size_t m, size_t n, size_t k,
+                                             __half const* alpha,
+                                             __half const* A, size_t lda,
+                                             __half const* B, size_t ldb,
+                                             __half const* beta, __half* C,
+                                             size_t ldc, cudaStream_t stream);

upstream_ref/cuda_gemm_optimization/06_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose_vectorized_memory_access.cu

@@ -0,0 +1,361 @@
+#include <cuda_fp16.h>
+
+#include "cuda_gemm.hpp"
+#include "cuda_gemm_utils.cuh"
+#include "cuda_gemm_utils.hpp"
+
+template <typename T, size_t BLOCK_TILE_SIZE, size_t WARP_TILE_SIZE,
+          size_t NUM_THREAD_TILES_PER_WARP, size_t THREAD_TILE_SIZE>
+__device__ void load_data_from_shared_memory_to_register_file_vectorized(
+    T const thread_block_tile[BLOCK_TILE_SIZE],
+    T register_values[NUM_THREAD_TILES_PER_WARP][THREAD_TILE_SIZE],
+    size_t warp_idx, size_t thread_idx)
+{
+    static_assert(BLOCK_TILE_SIZE % THREAD_TILE_SIZE == 0U);
+    constexpr size_t NUM_VECTOR_UNITS{sizeof(int4) / sizeof(T)};
+    static_assert(sizeof(int4) % sizeof(T) == 0U);
+    constexpr size_t VECTORIZED_THREAD_TILE_SIZE{THREAD_TILE_SIZE /
+                                                 NUM_VECTOR_UNITS};
+    static_assert(THREAD_TILE_SIZE % NUM_VECTOR_UNITS == 0U);
+
+#pragma unroll
+    for (size_t thread_tile_repeat_row_idx{0U};
+         thread_tile_repeat_row_idx < NUM_THREAD_TILES_PER_WARP;
+         ++thread_tile_repeat_row_idx)
+    {
+        size_t const thread_block_tile_row_idx{
+            warp_idx * WARP_TILE_SIZE +
+            thread_tile_repeat_row_idx *
+                (WARP_TILE_SIZE / NUM_THREAD_TILES_PER_WARP) +
+            thread_idx * THREAD_TILE_SIZE};
+#pragma unroll
+        for (size_t thread_tile_vector_idx{0U};
+             thread_tile_vector_idx < VECTORIZED_THREAD_TILE_SIZE;
+             ++thread_tile_vector_idx)
+        {
+            *reinterpret_cast<int4*>(
+                &register_values[thread_tile_repeat_row_idx]
+                                [thread_tile_vector_idx * NUM_VECTOR_UNITS]) =
+                *reinterpret_cast<int4 const*>(
+                    &thread_block_tile[thread_block_tile_row_idx +
+                                       thread_tile_vector_idx *
+                                           NUM_VECTOR_UNITS]);
+        }
+    }
+}
+
+template <typename T, size_t NUM_THREAD_TILES_PER_WARP_X,
+          size_t NUM_THREAD_TILES_PER_WARP_Y, size_t THREAD_TILE_SIZE_X,
+          size_t THREAD_TILE_SIZE_Y>
+__device__ void compute_thread_tile_results(
+    T const A_vals[NUM_THREAD_TILES_PER_WARP_Y][THREAD_TILE_SIZE_Y],
+    T const B_vals[NUM_THREAD_TILES_PER_WARP_X][THREAD_TILE_SIZE_X],
+    T C_thread_results[NUM_THREAD_TILES_PER_WARP_Y][NUM_THREAD_TILES_PER_WARP_X]
+                      [THREAD_TILE_SIZE_Y][THREAD_TILE_SIZE_X])
+{
+// Compute NUM_THREAD_TILES_PER_WARP_Y * NUM_THREAD_TILES_PER_WARP_X outer
+// products.
+#pragma unroll
+    for (size_t thread_tile_repeat_row_idx{0U};
+         thread_tile_repeat_row_idx < NUM_THREAD_TILES_PER_WARP_Y;
+         ++thread_tile_repeat_row_idx)
+    {
+#pragma unroll
+        for (size_t thread_tile_repeat_col_idx{0U};
+             thread_tile_repeat_col_idx < NUM_THREAD_TILES_PER_WARP_X;
+             ++thread_tile_repeat_col_idx)
+        {
+#pragma unroll
+            for (size_t thread_tile_y_idx{0U};
+                 thread_tile_y_idx < THREAD_TILE_SIZE_Y; ++thread_tile_y_idx)
+            {
+#pragma unroll
+                for (size_t thread_tile_x_idx{0U};
+                     thread_tile_x_idx < THREAD_TILE_SIZE_X;
+                     ++thread_tile_x_idx)
+                {
+                    C_thread_results[thread_tile_repeat_row_idx]
+                                    [thread_tile_repeat_col_idx]
+                                    [thread_tile_y_idx][thread_tile_x_idx] +=
+                        A_vals[thread_tile_repeat_row_idx][thread_tile_y_idx] *
+                        B_vals[thread_tile_repeat_col_idx][thread_tile_x_idx];
+                }
+            }
+        }
+    }
+}
+
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t WARP_TILE_SIZE_X, size_t WARP_TILE_SIZE_Y,
+          size_t THREAD_TILE_SIZE_X, size_t THREAD_TILE_SIZE_Y,
+          size_t NUM_THREAD_TILES_PER_WARP_X,
+          size_t NUM_THREAD_TILES_PER_WARP_Y>
+__device__ void write_results_from_register_file_to_global_memory_vectorized(
+    T const C_thread_results[NUM_THREAD_TILES_PER_WARP_Y]
+                            [NUM_THREAD_TILES_PER_WARP_X][THREAD_TILE_SIZE_Y]
+                            [THREAD_TILE_SIZE_X],
+    T alpha, T beta, T* C, size_t ldc, size_t m, size_t n, size_t block_row_idx,
+    size_t block_col_idx, size_t warp_row_idx, size_t warp_col_idx,
+    size_t thread_row_idx_in_warp, size_t thread_col_idx_in_warp)
+{
+    constexpr size_t NUM_VECTOR_UNITS{sizeof(int4) / sizeof(T)};
+    static_assert(sizeof(int4) % sizeof(T) == 0U);
+    static_assert(BLOCK_TILE_SIZE_X % NUM_VECTOR_UNITS == 0U);
+    constexpr size_t VECTORIZED_THREAD_TILE_SIZE_X{THREAD_TILE_SIZE_X /
+                                                   NUM_VECTOR_UNITS};
+    static_assert(THREAD_TILE_SIZE_X % NUM_VECTOR_UNITS == 0U);
+
+// Write the results to DRAM.
+#pragma unroll
+    for (size_t thread_tile_repeat_row_idx{0U};
+         thread_tile_repeat_row_idx < NUM_THREAD_TILES_PER_WARP_Y;
+         ++thread_tile_repeat_row_idx)
+    {
+#pragma unroll
+        for (size_t thread_tile_repeat_col_idx{0U};
+             thread_tile_repeat_col_idx < NUM_THREAD_TILES_PER_WARP_X;
+             ++thread_tile_repeat_col_idx)
+        {
+#pragma unroll
+            for (size_t thread_tile_y_idx{0U};
+                 thread_tile_y_idx < THREAD_TILE_SIZE_Y; ++thread_tile_y_idx)
+            {
+#pragma unroll
+                for (size_t thread_tile_x_vector_idx{0U};
+                     thread_tile_x_vector_idx < VECTORIZED_THREAD_TILE_SIZE_X;
+                     ++thread_tile_x_vector_idx)
+                {
+                    size_t const C_row_idx{
+                        blockIdx.y * BLOCK_TILE_SIZE_Y +
+                        warp_row_idx * WARP_TILE_SIZE_Y +
+                        thread_tile_repeat_row_idx *
+                            (WARP_TILE_SIZE_Y / NUM_THREAD_TILES_PER_WARP_Y) +
+                        thread_row_idx_in_warp * THREAD_TILE_SIZE_Y +
+                        thread_tile_y_idx};
+                    size_t const C_col_idx{
+                        blockIdx.x * BLOCK_TILE_SIZE_X +
+                        warp_col_idx * WARP_TILE_SIZE_X +
+                        thread_tile_repeat_col_idx *
+                            (WARP_TILE_SIZE_X / NUM_THREAD_TILES_PER_WARP_X) +
+                        thread_col_idx_in_warp * THREAD_TILE_SIZE_X +
+                        thread_tile_x_vector_idx * NUM_VECTOR_UNITS};
+
+                    if (C_row_idx < m && C_col_idx < n)
+                    {
+                        int4 C_vals{*reinterpret_cast<int4 const*>(
+                            &C[C_row_idx * ldc + C_col_idx])};
+#pragma unroll
+                        for (size_t i{0U}; i < NUM_VECTOR_UNITS; ++i)
+                        {
+                            reinterpret_cast<T*>(&C_vals)[i] =
+                                alpha *
+                                    C_thread_results[thread_tile_repeat_row_idx]
+                                                    [thread_tile_repeat_col_idx]
+                                                    [thread_tile_y_idx]
+                                                    [thread_tile_x_vector_idx *
+                                                         NUM_VECTOR_UNITS +
+                                                     i] +
+                                beta * reinterpret_cast<T const*>(&C_vals)[i];
+                        }
+                        *reinterpret_cast<int4*>(
+                            &C[C_row_idx * ldc + C_col_idx]) = C_vals;
+                    }
+                }
+            }
+        }
+    }
+}
+
+// GEMM kernel v06.
+// Each thread in the block processes THREAD_TILE_SIZE_Y *
+// THREAD_TILE_SIZE_X output values. Number of threads BLOCK_TILE_SIZE_Y *
+// BLOCK_TILE_SIZE_X / (THREAD_TILE_SIZE_Y * THREAD_TILE_SIZE_X)
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t BLOCK_TILE_SIZE_K, size_t WARP_TILE_SIZE_X,
+          size_t WARP_TILE_SIZE_Y, size_t THREAD_TILE_SIZE_X,
+          size_t THREAD_TILE_SIZE_Y, size_t NUM_THREADS_PER_WARP_X,
+          size_t NUM_THREADS_PER_WARP_Y>
+__global__ void gemm_v06_vectorized(size_t m, size_t n, size_t k, T alpha,
+                                    T const* A, size_t lda, T const* B,
+                                    size_t ldb, T beta, T* C, size_t ldc)
+{
+    static_assert(NUM_THREADS_PER_WARP_X * NUM_THREADS_PER_WARP_Y == 32U);
+    constexpr size_t NUM_WARPS_X{BLOCK_TILE_SIZE_X / WARP_TILE_SIZE_X};
+    static_assert(BLOCK_TILE_SIZE_X % WARP_TILE_SIZE_X == 0U);
+    constexpr size_t NUM_WARPS_Y{BLOCK_TILE_SIZE_Y / WARP_TILE_SIZE_Y};
+    static_assert(BLOCK_TILE_SIZE_Y % WARP_TILE_SIZE_Y == 0U);
+    constexpr unsigned int NUM_THREAD_TILES_PER_WARP_X{
+        WARP_TILE_SIZE_X / (THREAD_TILE_SIZE_X * NUM_THREADS_PER_WARP_X)};
+    constexpr unsigned int NUM_THREAD_TILES_PER_WARP_Y{
+        WARP_TILE_SIZE_Y / (THREAD_TILE_SIZE_Y * NUM_THREADS_PER_WARP_Y)};
+    static_assert(
+        WARP_TILE_SIZE_X % (THREAD_TILE_SIZE_X * NUM_THREADS_PER_WARP_X) == 0U);
+    static_assert(
+        WARP_TILE_SIZE_Y % (THREAD_TILE_SIZE_Y * NUM_THREADS_PER_WARP_Y) == 0U);
+
+    constexpr unsigned int NUM_THREADS_X{NUM_WARPS_X * NUM_THREADS_PER_WARP_X};
+    constexpr unsigned int NUM_THREADS_Y{NUM_WARPS_Y * NUM_THREADS_PER_WARP_Y};
+    // Avoid using blockDim.x * blockDim.y as the number of threads per block.
+    // Because it is a runtime constant and the compiler cannot optimize the
+    // loop unrolling based on that.
+    // Use a compile time constant instead.
+    constexpr size_t NUM_THREADS{NUM_THREADS_X * NUM_THREADS_Y};
+
+    // Cache a tile of A and B in shared memory for data reuse.
+    __shared__ T
+        A_thread_block_tile_transposed[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_Y];
+    __shared__ T B_thread_block_tile[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_X];
+
+    // A_vals is cached in the register.
+    T A_vals[NUM_THREAD_TILES_PER_WARP_Y][THREAD_TILE_SIZE_Y] = {
+        static_cast<T>(0)};
+    // B_vals is cached in the register.
+    T B_vals[NUM_THREAD_TILES_PER_WARP_X][THREAD_TILE_SIZE_X] = {
+        static_cast<T>(0)};
+
+    size_t const thread_linear_idx{threadIdx.y * blockDim.x + threadIdx.x};
+    size_t const warp_linear_idx{thread_linear_idx / 32U};
+    size_t const warp_row_idx{warp_linear_idx / NUM_WARPS_X};
+    size_t const warp_col_idx{warp_linear_idx % NUM_WARPS_X};
+    size_t const thread_linear_idx_in_warp{thread_linear_idx % 32U};
+    size_t const thread_linear_row_idx_in_warp{thread_linear_idx_in_warp /
+                                               NUM_THREADS_PER_WARP_X};
+    size_t const thread_linear_col_idx_in_warp{thread_linear_idx_in_warp %
+                                               NUM_THREADS_PER_WARP_X};
+
+    // Number of outer loops to perform the sum of inner products.
+    // C_thread_block_tile =
+    // \sigma_{thread_block_tile_idx=0}^{num_thread_block_tiles-1} A[:,
+    // thread_block_tile_idx:BLOCK_TILE_SIZE_K] *
+    // B[thread_block_tile_idx:BLOCK_TILE_SIZE_K, :]
+    size_t const num_thread_block_tiles{(k + BLOCK_TILE_SIZE_K - 1) /
+                                        BLOCK_TILE_SIZE_K};
+    // Each thread in the block processes NUM_THREAD_TILES_PER_WARP_Y *
+    // NUM_THREAD_TILES_PER_WARP_X * THREAD_TILE_SIZE_Y *
+    // THREAD_TILE_SIZE_X output values.
+    T C_thread_results[NUM_THREAD_TILES_PER_WARP_Y][NUM_THREAD_TILES_PER_WARP_X]
+                      [THREAD_TILE_SIZE_Y][THREAD_TILE_SIZE_X] = {
+                          static_cast<T>(0)};
+
+    for (size_t thread_block_tile_idx{0U};
+         thread_block_tile_idx < num_thread_block_tiles;
+         ++thread_block_tile_idx)
+    {
+        load_data_from_global_memory_to_shared_memory_transposed_vectorized<
+            T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+            NUM_THREADS>(A, lda, B, ldb, A_thread_block_tile_transposed,
+                         B_thread_block_tile, thread_block_tile_idx,
+                         thread_linear_idx, m, n, k);
+        __syncthreads();
+
+// Perform A[:, thread_block_tile_idx:BLOCK_TILE_SIZE_K] *
+// B[thread_block_tile_idx:BLOCK_TILE_SIZE_K, :] where A[:,
+// thread_block_tile_idx:BLOCK_TILE_SIZE_K] and
+// B[thread_block_tile_idx:BLOCK_TILE_SIZE_K, :] are cached in the
+// shared memory as A_thread_block_tile and B_thread_block_tile,
+// respectively. This inner product is further decomposed to
+// BLOCK_TILE_SIZE_K outer products. A_thread_block_tile *
+// B_thread_block_tile = \sigma_{k_i=0}^{BLOCK_TILE_SIZE_K-1}
+// A_thread_block_tile[:, k_i] @ B_thread_block_tile[k_i, :] Note that
+// both A_thread_block_tile and B_thread_block_tile can be cached in the
+// register.
+#pragma unroll
+        for (size_t k_i{0U}; k_i < BLOCK_TILE_SIZE_K; ++k_i)
+        {
+            // Load data from shared memory to register file for A.
+            load_data_from_shared_memory_to_register_file_vectorized<
+                T, BLOCK_TILE_SIZE_Y, WARP_TILE_SIZE_Y, NUM_THREADS_PER_WARP_Y,
+                THREAD_TILE_SIZE_Y>(A_thread_block_tile_transposed[k_i], A_vals,
+                                    warp_row_idx,
+                                    thread_linear_row_idx_in_warp);
+            // Load data from shared memory to register file for B.
+            load_data_from_shared_memory_to_register_file_vectorized<
+                T, BLOCK_TILE_SIZE_X, WARP_TILE_SIZE_X, NUM_THREADS_PER_WARP_X,
+                THREAD_TILE_SIZE_X>(B_thread_block_tile[k_i], B_vals,
+                                    warp_col_idx,
+                                    thread_linear_col_idx_in_warp);
+
+            // Compute NUM_THREAD_TILES_PER_WARP_Y * NUM_THREAD_TILES_PER_WARP_X
+            // outer products.
+            compute_thread_tile_results<T, NUM_THREAD_TILES_PER_WARP_X,
+                                        NUM_THREAD_TILES_PER_WARP_Y,
+                                        THREAD_TILE_SIZE_X, THREAD_TILE_SIZE_Y>(
+                A_vals, B_vals, C_thread_results);
+        }
+        __syncthreads();
+    }
+
+    // Write the results to DRAM.
+    write_results_from_register_file_to_global_memory_vectorized<
+        T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, WARP_TILE_SIZE_X,
+        WARP_TILE_SIZE_Y, THREAD_TILE_SIZE_X, THREAD_TILE_SIZE_Y,
+        NUM_THREAD_TILES_PER_WARP_X, NUM_THREAD_TILES_PER_WARP_Y>(
+        C_thread_results, alpha, beta, C, ldc, m, n, blockIdx.y, blockIdx.x,
+        warp_row_idx, warp_col_idx, thread_linear_row_idx_in_warp,
+        thread_linear_col_idx_in_warp);
+}
+
+template <typename T>
+void launch_gemm_kernel_v06_vectorized(size_t m, size_t n, size_t k,
+                                       T const* alpha, T const* A, size_t lda,
+                                       T const* B, size_t ldb, T const* beta,
+                                       T* C, size_t ldc, cudaStream_t stream)
+{
+    // Feel free to play with the block tile sizes.
+    // The algorithm correctness should always be guaranteed.
+    constexpr unsigned int BLOCK_TILE_SIZE_X{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_Y{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_K{16U};
+
+    constexpr unsigned int WARP_TILE_SIZE_X{32U};
+    constexpr unsigned int WARP_TILE_SIZE_Y{64U};
+    constexpr unsigned int NUM_WARPS_X{BLOCK_TILE_SIZE_X / WARP_TILE_SIZE_X};
+    constexpr unsigned int NUM_WARPS_Y{BLOCK_TILE_SIZE_Y / WARP_TILE_SIZE_Y};
+    static_assert(BLOCK_TILE_SIZE_X % WARP_TILE_SIZE_X == 0U);
+    static_assert(BLOCK_TILE_SIZE_Y % WARP_TILE_SIZE_Y == 0U);
+
+    constexpr unsigned int THREAD_TILE_SIZE_X{8U};
+    constexpr unsigned int THREAD_TILE_SIZE_Y{8U};
+
+    constexpr unsigned int NUM_THREADS_PER_WARP_X{4U};
+    constexpr unsigned int NUM_THREADS_PER_WARP_Y{8U};
+    static_assert(NUM_THREADS_PER_WARP_X * NUM_THREADS_PER_WARP_Y == 32U);
+    static_assert(
+        WARP_TILE_SIZE_X % (THREAD_TILE_SIZE_X * NUM_THREADS_PER_WARP_X) == 0U);
+    static_assert(
+        WARP_TILE_SIZE_Y % (THREAD_TILE_SIZE_Y * NUM_THREADS_PER_WARP_Y) == 0U);
+
+    constexpr unsigned int NUM_THREADS_X{NUM_WARPS_X * NUM_THREADS_PER_WARP_X};
+    constexpr unsigned int NUM_THREADS_Y{NUM_WARPS_Y * NUM_THREADS_PER_WARP_Y};
+
+    constexpr unsigned int NUM_THREADS_PER_BLOCK{NUM_THREADS_X * NUM_THREADS_Y};
+
+    dim3 const block_dim{NUM_THREADS_PER_BLOCK, 1U, 1U};
+    dim3 const grid_dim{
+        (static_cast<unsigned int>(n) + BLOCK_TILE_SIZE_X - 1U) /
+            BLOCK_TILE_SIZE_X,
+        (static_cast<unsigned int>(m) + BLOCK_TILE_SIZE_Y - 1U) /
+            BLOCK_TILE_SIZE_Y,
+        1U};
+    gemm_v06_vectorized<T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y,
+                        BLOCK_TILE_SIZE_K, WARP_TILE_SIZE_X, WARP_TILE_SIZE_Y,
+                        THREAD_TILE_SIZE_X, THREAD_TILE_SIZE_Y,
+                        NUM_THREADS_PER_WARP_X, NUM_THREADS_PER_WARP_Y>
+        <<<grid_dim, block_dim, 0U, stream>>>(m, n, k, *alpha, A, lda, B, ldb,
+                                              *beta, C, ldc);
+    CHECK_LAST_CUDA_ERROR();
+}
+
+// Explicit instantiation.
+template void launch_gemm_kernel_v06_vectorized<float>(
+    size_t m, size_t n, size_t k, float const* alpha, float const* A,
+    size_t lda, float const* B, size_t ldb, float const* beta, float* C,
+    size_t ldc, cudaStream_t stream);
+template void launch_gemm_kernel_v06_vectorized<double>(
+    size_t m, size_t n, size_t k, double const* alpha, double const* A,
+    size_t lda, double const* B, size_t ldb, double const* beta, double* C,
+    size_t ldc, cudaStream_t stream);
+template void launch_gemm_kernel_v06_vectorized<__half>(
+    size_t m, size_t n, size_t k, __half const* alpha, __half const* A,
+    size_t lda, __half const* B, size_t ldb, __half const* beta, __half* C,
+    size_t ldc, cudaStream_t stream);

upstream_ref/cuda_gemm_optimization/06_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose_vectorized_memory_access_double_buffered.cu

@@ -0,0 +1,481 @@
+#include <cuda_fp16.h>
+
+#include "cuda_gemm.hpp"
+#include "cuda_gemm_utils.cuh"
+#include "cuda_gemm_utils.hpp"
+
+template <typename T, size_t BLOCK_TILE_SIZE, size_t WARP_TILE_SIZE,
+          size_t NUM_THREAD_TILES_PER_WARP, size_t THREAD_TILE_SIZE>
+__device__ void load_data_from_shared_memory_to_register_file_vectorized(
+    T const thread_block_tile[BLOCK_TILE_SIZE],
+    T register_values[NUM_THREAD_TILES_PER_WARP][THREAD_TILE_SIZE],
+    size_t warp_idx, size_t thread_idx)
+{
+    static_assert(BLOCK_TILE_SIZE % THREAD_TILE_SIZE == 0U);
+    constexpr size_t NUM_VECTOR_UNITS{sizeof(int4) / sizeof(T)};
+    static_assert(sizeof(int4) % sizeof(T) == 0U);
+    constexpr size_t VECTORIZED_THREAD_TILE_SIZE{THREAD_TILE_SIZE /
+                                                 NUM_VECTOR_UNITS};
+    static_assert(THREAD_TILE_SIZE % NUM_VECTOR_UNITS == 0U);
+
+#pragma unroll
+    for (size_t thread_tile_repeat_row_idx{0U};
+         thread_tile_repeat_row_idx < NUM_THREAD_TILES_PER_WARP;
+         ++thread_tile_repeat_row_idx)
+    {
+        size_t const thread_block_tile_row_idx{
+            warp_idx * WARP_TILE_SIZE +
+            thread_tile_repeat_row_idx *
+                (WARP_TILE_SIZE / NUM_THREAD_TILES_PER_WARP) +
+            thread_idx * THREAD_TILE_SIZE};
+#pragma unroll
+        for (size_t thread_tile_vector_idx{0U};
+             thread_tile_vector_idx < VECTORIZED_THREAD_TILE_SIZE;
+             ++thread_tile_vector_idx)
+        {
+            *reinterpret_cast<int4*>(
+                &register_values[thread_tile_repeat_row_idx]
+                                [thread_tile_vector_idx * NUM_VECTOR_UNITS]) =
+                *reinterpret_cast<int4 const*>(
+                    &thread_block_tile[thread_block_tile_row_idx +
+                                       thread_tile_vector_idx *
+                                           NUM_VECTOR_UNITS]);
+        }
+    }
+}
+
+template <typename T, size_t NUM_THREAD_TILES_PER_WARP_X,
+          size_t NUM_THREAD_TILES_PER_WARP_Y, size_t THREAD_TILE_SIZE_X,
+          size_t THREAD_TILE_SIZE_Y>
+__device__ void compute_thread_tile_results(
+    T const A_vals[NUM_THREAD_TILES_PER_WARP_Y][THREAD_TILE_SIZE_Y],
+    T const B_vals[NUM_THREAD_TILES_PER_WARP_X][THREAD_TILE_SIZE_X],
+    T C_thread_results[NUM_THREAD_TILES_PER_WARP_Y][NUM_THREAD_TILES_PER_WARP_X]
+                      [THREAD_TILE_SIZE_Y][THREAD_TILE_SIZE_X])
+{
+// Compute NUM_THREAD_TILES_PER_WARP_Y * NUM_THREAD_TILES_PER_WARP_X outer
+// products.
+#pragma unroll
+    for (size_t thread_tile_repeat_row_idx{0U};
+         thread_tile_repeat_row_idx < NUM_THREAD_TILES_PER_WARP_Y;
+         ++thread_tile_repeat_row_idx)
+    {
+#pragma unroll
+        for (size_t thread_tile_repeat_col_idx{0U};
+             thread_tile_repeat_col_idx < NUM_THREAD_TILES_PER_WARP_X;
+             ++thread_tile_repeat_col_idx)
+        {
+#pragma unroll
+            for (size_t thread_tile_y_idx{0U};
+                 thread_tile_y_idx < THREAD_TILE_SIZE_Y; ++thread_tile_y_idx)
+            {
+#pragma unroll
+                for (size_t thread_tile_x_idx{0U};
+                     thread_tile_x_idx < THREAD_TILE_SIZE_X;
+                     ++thread_tile_x_idx)
+                {
+                    C_thread_results[thread_tile_repeat_row_idx]
+                                    [thread_tile_repeat_col_idx]
+                                    [thread_tile_y_idx][thread_tile_x_idx] +=
+                        A_vals[thread_tile_repeat_row_idx][thread_tile_y_idx] *
+                        B_vals[thread_tile_repeat_col_idx][thread_tile_x_idx];
+                }
+            }
+        }
+    }
+}
+
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t BLOCK_TILE_SIZE_K, size_t WARP_TILE_SIZE_X,
+          size_t WARP_TILE_SIZE_Y, size_t THREAD_TILE_SIZE_X,
+          size_t THREAD_TILE_SIZE_Y, size_t NUM_THREADS_PER_WARP_X,
+          size_t NUM_THREADS_PER_WARP_Y, size_t NUM_THREAD_TILES_PER_WARP_X,
+          size_t NUM_THREAD_TILES_PER_WARP_Y>
+__device__ void process_data_from_shared_memory_using_register_file_vectorized(
+    T A_vals[NUM_THREAD_TILES_PER_WARP_Y][THREAD_TILE_SIZE_Y],
+    T B_vals[NUM_THREAD_TILES_PER_WARP_X][THREAD_TILE_SIZE_X],
+    T C_thread_results[NUM_THREAD_TILES_PER_WARP_Y][NUM_THREAD_TILES_PER_WARP_X]
+                      [THREAD_TILE_SIZE_Y][THREAD_TILE_SIZE_X],
+    T const A_thread_block_tile_transposed[BLOCK_TILE_SIZE_K]
+                                          [BLOCK_TILE_SIZE_Y],
+    T const B_thread_block_tile[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_X],
+    size_t warp_row_idx, size_t warp_col_idx, size_t thread_row_idx_in_warp,
+    size_t thread_col_idx_in_warp)
+{
+#pragma unroll
+    for (size_t k_i{0U}; k_i < BLOCK_TILE_SIZE_K; ++k_i)
+    {
+        // Load data from shared memory to register file for A.
+        load_data_from_shared_memory_to_register_file_vectorized<
+            T, BLOCK_TILE_SIZE_Y, WARP_TILE_SIZE_Y, NUM_THREADS_PER_WARP_Y,
+            THREAD_TILE_SIZE_Y>(A_thread_block_tile_transposed[k_i], A_vals,
+                                warp_row_idx, thread_row_idx_in_warp);
+        // Load data from shared memory to register file for B.
+        load_data_from_shared_memory_to_register_file_vectorized<
+            T, BLOCK_TILE_SIZE_X, WARP_TILE_SIZE_X, NUM_THREADS_PER_WARP_X,
+            THREAD_TILE_SIZE_X>(B_thread_block_tile[k_i], B_vals, warp_col_idx,
+                                thread_col_idx_in_warp);
+
+        // Compute NUM_THREAD_TILES_PER_WARP_Y *
+        // NUM_THREAD_TILES_PER_WARP_X outer products.
+        compute_thread_tile_results<T, NUM_THREAD_TILES_PER_WARP_X,
+                                    NUM_THREAD_TILES_PER_WARP_Y,
+                                    THREAD_TILE_SIZE_X, THREAD_TILE_SIZE_Y>(
+            A_vals, B_vals, C_thread_results);
+    }
+}
+
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t WARP_TILE_SIZE_X, size_t WARP_TILE_SIZE_Y,
+          size_t THREAD_TILE_SIZE_X, size_t THREAD_TILE_SIZE_Y,
+          size_t NUM_THREAD_TILES_PER_WARP_X,
+          size_t NUM_THREAD_TILES_PER_WARP_Y>
+__device__ void write_results_from_register_file_to_global_memory_vectorized(
+    T const C_thread_results[NUM_THREAD_TILES_PER_WARP_Y]
+                            [NUM_THREAD_TILES_PER_WARP_X][THREAD_TILE_SIZE_Y]
+                            [THREAD_TILE_SIZE_X],
+    T alpha, T beta, T* C, size_t ldc, size_t m, size_t n, size_t block_row_idx,
+    size_t block_col_idx, size_t warp_row_idx, size_t warp_col_idx,
+    size_t thread_row_idx_in_warp, size_t thread_col_idx_in_warp)
+{
+    constexpr size_t NUM_VECTOR_UNITS{sizeof(int4) / sizeof(T)};
+    static_assert(sizeof(int4) % sizeof(T) == 0U);
+    static_assert(BLOCK_TILE_SIZE_X % NUM_VECTOR_UNITS == 0U);
+    constexpr size_t VECTORIZED_THREAD_TILE_SIZE_X{THREAD_TILE_SIZE_X /
+                                                   NUM_VECTOR_UNITS};
+    static_assert(THREAD_TILE_SIZE_X % NUM_VECTOR_UNITS == 0U);
+
+// Write the results to DRAM.
+#pragma unroll
+    for (size_t thread_tile_repeat_row_idx{0U};
+         thread_tile_repeat_row_idx < NUM_THREAD_TILES_PER_WARP_Y;
+         ++thread_tile_repeat_row_idx)
+    {
+#pragma unroll
+        for (size_t thread_tile_repeat_col_idx{0U};
+             thread_tile_repeat_col_idx < NUM_THREAD_TILES_PER_WARP_X;
+             ++thread_tile_repeat_col_idx)
+        {
+#pragma unroll
+            for (size_t thread_tile_y_idx{0U};
+                 thread_tile_y_idx < THREAD_TILE_SIZE_Y; ++thread_tile_y_idx)
+            {
+#pragma unroll
+                for (size_t thread_tile_x_vector_idx{0U};
+                     thread_tile_x_vector_idx < VECTORIZED_THREAD_TILE_SIZE_X;
+                     ++thread_tile_x_vector_idx)
+                {
+                    size_t const C_row_idx{
+                        blockIdx.y * BLOCK_TILE_SIZE_Y +
+                        warp_row_idx * WARP_TILE_SIZE_Y +
+                        thread_tile_repeat_row_idx *
+                            (WARP_TILE_SIZE_Y / NUM_THREAD_TILES_PER_WARP_Y) +
+                        thread_row_idx_in_warp * THREAD_TILE_SIZE_Y +
+                        thread_tile_y_idx};
+                    size_t const C_col_idx{
+                        blockIdx.x * BLOCK_TILE_SIZE_X +
+                        warp_col_idx * WARP_TILE_SIZE_X +
+                        thread_tile_repeat_col_idx *
+                            (WARP_TILE_SIZE_X / NUM_THREAD_TILES_PER_WARP_X) +
+                        thread_col_idx_in_warp * THREAD_TILE_SIZE_X +
+                        thread_tile_x_vector_idx * NUM_VECTOR_UNITS};
+
+                    if (C_row_idx < m && C_col_idx < n)
+                    {
+                        int4 C_vals{*reinterpret_cast<int4 const*>(
+                            &C[C_row_idx * ldc + C_col_idx])};
+#pragma unroll
+                        for (size_t i{0U}; i < NUM_VECTOR_UNITS; ++i)
+                        {
+                            reinterpret_cast<T*>(&C_vals)[i] =
+                                alpha *
+                                    C_thread_results[thread_tile_repeat_row_idx]
+                                                    [thread_tile_repeat_col_idx]
+                                                    [thread_tile_y_idx]
+                                                    [thread_tile_x_vector_idx *
+                                                         NUM_VECTOR_UNITS +
+                                                     i] +
+                                beta * reinterpret_cast<T const*>(&C_vals)[i];
+                        }
+                        *reinterpret_cast<int4*>(
+                            &C[C_row_idx * ldc + C_col_idx]) = C_vals;
+                    }
+                }
+            }
+        }
+    }
+}
+
+// GEMM kernel v06.
+// Each thread in the block processes THREAD_TILE_SIZE_Y *
+// THREAD_TILE_SIZE_X output values. Number of threads BLOCK_TILE_SIZE_Y *
+// BLOCK_TILE_SIZE_X / (THREAD_TILE_SIZE_Y * THREAD_TILE_SIZE_X)
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t BLOCK_TILE_SIZE_K, size_t WARP_TILE_SIZE_X,
+          size_t WARP_TILE_SIZE_Y, size_t THREAD_TILE_SIZE_X,
+          size_t THREAD_TILE_SIZE_Y, size_t NUM_THREADS_PER_WARP_X,
+          size_t NUM_THREADS_PER_WARP_Y>
+__global__ void
+gemm_v06_vectorized_double_buffered(size_t m, size_t n, size_t k, T alpha,
+                                    T const* A, size_t lda, T const* B,
+                                    size_t ldb, T beta, T* C, size_t ldc)
+{
+    static_assert(NUM_THREADS_PER_WARP_X * NUM_THREADS_PER_WARP_Y == 32U);
+    constexpr size_t NUM_WARPS_X{BLOCK_TILE_SIZE_X / WARP_TILE_SIZE_X};
+    static_assert(BLOCK_TILE_SIZE_X % WARP_TILE_SIZE_X == 0U);
+    constexpr size_t NUM_WARPS_Y{BLOCK_TILE_SIZE_Y / WARP_TILE_SIZE_Y};
+    static_assert(BLOCK_TILE_SIZE_Y % WARP_TILE_SIZE_Y == 0U);
+    constexpr unsigned int NUM_THREAD_TILES_PER_WARP_X{
+        WARP_TILE_SIZE_X / (THREAD_TILE_SIZE_X * NUM_THREADS_PER_WARP_X)};
+    constexpr unsigned int NUM_THREAD_TILES_PER_WARP_Y{
+        WARP_TILE_SIZE_Y / (THREAD_TILE_SIZE_Y * NUM_THREADS_PER_WARP_Y)};
+    static_assert(
+        WARP_TILE_SIZE_X % (THREAD_TILE_SIZE_X * NUM_THREADS_PER_WARP_X) == 0U);
+    static_assert(
+        WARP_TILE_SIZE_Y % (THREAD_TILE_SIZE_Y * NUM_THREADS_PER_WARP_Y) == 0U);
+
+    constexpr unsigned int NUM_THREADS_X{NUM_WARPS_X * NUM_THREADS_PER_WARP_X};
+    constexpr unsigned int NUM_THREADS_Y{NUM_WARPS_Y * NUM_THREADS_PER_WARP_Y};
+    // Avoid using blockDim.x * blockDim.y as the number of threads per block.
+    // Because it is a runtime constant and the compiler cannot optimize the
+    // loop unrolling based on that.
+    // Use a compile time constant instead.
+    constexpr size_t NUM_THREADS{NUM_THREADS_X * NUM_THREADS_Y};
+
+    constexpr size_t NUM_PIPELINES{2U};
+    // Only double buffer is supported in the implementation.
+    // But even more number of pipelines can be supported if the implementation
+    // is modified.
+    static_assert(NUM_PIPELINES == 2U);
+    static_assert((NUM_WARPS_X * NUM_WARPS_Y) % NUM_PIPELINES == 0U);
+    static_assert(NUM_THREADS % NUM_PIPELINES == 0U);
+    constexpr size_t NUM_THREADS_PER_PIPELINE{NUM_THREADS / NUM_PIPELINES};
+    constexpr size_t NUM_WARPS_PER_PIPELINE{(NUM_WARPS_X * NUM_WARPS_Y) /
+                                            NUM_PIPELINES};
+
+    // Cache a tile of A and B in shared memory for data reuse.
+    __shared__ T
+        A_thread_block_tile_transposed[NUM_PIPELINES][BLOCK_TILE_SIZE_K]
+                                      [BLOCK_TILE_SIZE_Y];
+    __shared__ T B_thread_block_tile[NUM_PIPELINES][BLOCK_TILE_SIZE_K]
+                                    [BLOCK_TILE_SIZE_X];
+
+    // A_vals is cached in the register.
+    T A_vals[NUM_THREAD_TILES_PER_WARP_Y][THREAD_TILE_SIZE_Y] = {
+        static_cast<T>(0)};
+    // B_vals is cached in the register.
+    T B_vals[NUM_THREAD_TILES_PER_WARP_X][THREAD_TILE_SIZE_X] = {
+        static_cast<T>(0)};
+
+    size_t const thread_linear_idx{threadIdx.y * blockDim.x + threadIdx.x};
+    size_t const warp_linear_idx{thread_linear_idx / 32U};
+    size_t const warp_row_idx{warp_linear_idx / NUM_WARPS_X};
+    size_t const warp_col_idx{warp_linear_idx % NUM_WARPS_X};
+    size_t const thread_linear_idx_in_warp{thread_linear_idx % 32U};
+    size_t const thread_linear_row_idx_in_warp{thread_linear_idx_in_warp /
+                                               NUM_THREADS_PER_WARP_X};
+    size_t const thread_linear_col_idx_in_warp{thread_linear_idx_in_warp %
+                                               NUM_THREADS_PER_WARP_X};
+    // Separate the warps to different pipelines.
+    size_t const pipeline_index{warp_linear_idx / NUM_WARPS_PER_PIPELINE};
+
+    // Number of outer loops to perform the sum of inner products.
+    // C_thread_block_tile =
+    // \sigma_{thread_block_tile_idx=0}^{num_thread_block_tiles-1} A[:,
+    // thread_block_tile_idx:BLOCK_TILE_SIZE_K] *
+    // B[thread_block_tile_idx:BLOCK_TILE_SIZE_K, :]
+    size_t const num_thread_block_tiles{(k + BLOCK_TILE_SIZE_K - 1) /
+                                        BLOCK_TILE_SIZE_K};
+    // Each thread in the block processes NUM_THREAD_TILES_PER_WARP_Y *
+    // NUM_THREAD_TILES_PER_WARP_X * THREAD_TILE_SIZE_Y *
+    // THREAD_TILE_SIZE_X output values.
+    T C_thread_results[NUM_THREAD_TILES_PER_WARP_Y][NUM_THREAD_TILES_PER_WARP_X]
+                      [THREAD_TILE_SIZE_Y][THREAD_TILE_SIZE_X] = {
+                          static_cast<T>(0)};
+
+    if (pipeline_index == 0U)
+    {
+        // Pipeline 0 warps load buffer 0.
+        load_data_from_global_memory_to_shared_memory_transposed_vectorized<
+            T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+            NUM_THREADS_PER_PIPELINE>(
+            A, lda, B, ldb, A_thread_block_tile_transposed[pipeline_index],
+            B_thread_block_tile[pipeline_index], 0U,
+            thread_linear_idx - pipeline_index * NUM_THREADS_PER_PIPELINE, m, n,
+            k);
+    }
+    __syncthreads();
+
+    for (size_t thread_block_tile_idx{0U};
+         thread_block_tile_idx < num_thread_block_tiles;
+         thread_block_tile_idx += NUM_PIPELINES)
+    {
+        if (pipeline_index == 0U)
+        {
+            // Pipeline 0 warps process buffer 0.
+            process_data_from_shared_memory_using_register_file_vectorized<
+                T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+                WARP_TILE_SIZE_X, WARP_TILE_SIZE_Y, THREAD_TILE_SIZE_X,
+                THREAD_TILE_SIZE_Y, NUM_THREADS_PER_WARP_X,
+                NUM_THREADS_PER_WARP_Y, NUM_THREAD_TILES_PER_WARP_X,
+                NUM_THREAD_TILES_PER_WARP_Y>(
+                A_vals, B_vals, C_thread_results,
+                A_thread_block_tile_transposed[pipeline_index],
+                B_thread_block_tile[pipeline_index], warp_row_idx, warp_col_idx,
+                thread_linear_row_idx_in_warp, thread_linear_col_idx_in_warp);
+            __syncthreads();
+
+            // Pipeline 0 warps process buffer 1.
+            if (thread_block_tile_idx + 1U < num_thread_block_tiles)
+            {
+                process_data_from_shared_memory_using_register_file_vectorized<
+                    T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+                    WARP_TILE_SIZE_X, WARP_TILE_SIZE_Y, THREAD_TILE_SIZE_X,
+                    THREAD_TILE_SIZE_Y, NUM_THREADS_PER_WARP_X,
+                    NUM_THREADS_PER_WARP_Y, NUM_THREAD_TILES_PER_WARP_X,
+                    NUM_THREAD_TILES_PER_WARP_Y>(
+                    A_vals, B_vals, C_thread_results,
+                    A_thread_block_tile_transposed[pipeline_index + 1],
+                    B_thread_block_tile[pipeline_index + 1], warp_row_idx,
+                    warp_col_idx, thread_linear_row_idx_in_warp,
+                    thread_linear_col_idx_in_warp);
+            }
+            __syncthreads();
+
+            // Pipeline 0 warps load buffer 0.
+            if (thread_block_tile_idx + 2U < num_thread_block_tiles)
+            {
+                load_data_from_global_memory_to_shared_memory_transposed_vectorized<
+                    T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+                    NUM_THREADS_PER_PIPELINE>(
+                    A, lda, B, ldb,
+                    A_thread_block_tile_transposed[pipeline_index],
+                    B_thread_block_tile[pipeline_index],
+                    thread_block_tile_idx + 2,
+                    thread_linear_idx -
+                        pipeline_index * NUM_THREADS_PER_PIPELINE,
+                    m, n, k);
+            }
+            __syncthreads();
+        }
+        else
+        {
+            // Pipeline 1 warps load buffer 1.
+            if (thread_block_tile_idx + 1U < num_thread_block_tiles)
+            {
+                load_data_from_global_memory_to_shared_memory_transposed_vectorized<
+                    T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+                    NUM_THREADS_PER_PIPELINE>(
+                    A, lda, B, ldb,
+                    A_thread_block_tile_transposed[pipeline_index],
+                    B_thread_block_tile[pipeline_index],
+                    thread_block_tile_idx + 1,
+                    thread_linear_idx -
+                        pipeline_index * NUM_THREADS_PER_PIPELINE,
+                    m, n, k);
+            }
+            __syncthreads();
+
+            // Pipeline 1 warps process buffer 0.
+            process_data_from_shared_memory_using_register_file_vectorized<
+                T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+                WARP_TILE_SIZE_X, WARP_TILE_SIZE_Y, THREAD_TILE_SIZE_X,
+                THREAD_TILE_SIZE_Y, NUM_THREADS_PER_WARP_X,
+                NUM_THREADS_PER_WARP_Y, NUM_THREAD_TILES_PER_WARP_X,
+                NUM_THREAD_TILES_PER_WARP_Y>(
+                A_vals, B_vals, C_thread_results,
+                A_thread_block_tile_transposed[pipeline_index - 1],
+                B_thread_block_tile[pipeline_index - 1], warp_row_idx,
+                warp_col_idx, thread_linear_row_idx_in_warp,
+                thread_linear_col_idx_in_warp);
+            __syncthreads();
+
+            // Pipeline 1 warps process buffer 1.
+            if (thread_block_tile_idx + 1U < num_thread_block_tiles)
+            {
+                process_data_from_shared_memory_using_register_file_vectorized<
+                    T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+                    WARP_TILE_SIZE_X, WARP_TILE_SIZE_Y, THREAD_TILE_SIZE_X,
+                    THREAD_TILE_SIZE_Y, NUM_THREADS_PER_WARP_X,
+                    NUM_THREADS_PER_WARP_Y, NUM_THREAD_TILES_PER_WARP_X,
+                    NUM_THREAD_TILES_PER_WARP_Y>(
+                    A_vals, B_vals, C_thread_results,
+                    A_thread_block_tile_transposed[pipeline_index],
+                    B_thread_block_tile[pipeline_index], warp_row_idx,
+                    warp_col_idx, thread_linear_row_idx_in_warp,
+                    thread_linear_col_idx_in_warp);
+            }
+            __syncthreads();
+        }
+    }
+
+    // Write the results to DRAM.
+    write_results_from_register_file_to_global_memory_vectorized<
+        T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, WARP_TILE_SIZE_X,
+        WARP_TILE_SIZE_Y, THREAD_TILE_SIZE_X, THREAD_TILE_SIZE_Y,
+        NUM_THREAD_TILES_PER_WARP_X, NUM_THREAD_TILES_PER_WARP_Y>(
+        C_thread_results, alpha, beta, C, ldc, m, n, blockIdx.y, blockIdx.x,
+        warp_row_idx, warp_col_idx, thread_linear_row_idx_in_warp,
+        thread_linear_col_idx_in_warp);
+}
+
+template <typename T>
+void launch_gemm_kernel_v06_vectorized_double_buffered(
+    size_t m, size_t n, size_t k, T const* alpha, T const* A, size_t lda,
+    T const* B, size_t ldb, T const* beta, T* C, size_t ldc,
+    cudaStream_t stream)
+{
+    // Feel free to play with the block tile sizes.
+    // The algorithm correctness should always be guaranteed.
+    constexpr unsigned int BLOCK_TILE_SIZE_X{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_Y{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_K{16U};
+
+    constexpr unsigned int WARP_TILE_SIZE_X{32U};
+    constexpr unsigned int WARP_TILE_SIZE_Y{64U};
+    constexpr unsigned int NUM_WARPS_X{BLOCK_TILE_SIZE_X / WARP_TILE_SIZE_X};
+    constexpr unsigned int NUM_WARPS_Y{BLOCK_TILE_SIZE_Y / WARP_TILE_SIZE_Y};
+    static_assert(BLOCK_TILE_SIZE_X % WARP_TILE_SIZE_X == 0U);
+    static_assert(BLOCK_TILE_SIZE_Y % WARP_TILE_SIZE_Y == 0U);
+
+    constexpr unsigned int THREAD_TILE_SIZE_X{8U};
+    constexpr unsigned int THREAD_TILE_SIZE_Y{8U};
+
+    constexpr unsigned int NUM_THREADS_PER_WARP_X{4U};
+    constexpr unsigned int NUM_THREADS_PER_WARP_Y{8U};
+    static_assert(NUM_THREADS_PER_WARP_X * NUM_THREADS_PER_WARP_Y == 32U);
+    static_assert(
+        WARP_TILE_SIZE_X % (THREAD_TILE_SIZE_X * NUM_THREADS_PER_WARP_X) == 0U);
+    static_assert(
+        WARP_TILE_SIZE_Y % (THREAD_TILE_SIZE_Y * NUM_THREADS_PER_WARP_Y) == 0U);
+
+    constexpr unsigned int NUM_THREADS_X{NUM_WARPS_X * NUM_THREADS_PER_WARP_X};
+    constexpr unsigned int NUM_THREADS_Y{NUM_WARPS_Y * NUM_THREADS_PER_WARP_Y};
+
+    constexpr unsigned int NUM_THREADS_PER_BLOCK{NUM_THREADS_X * NUM_THREADS_Y};
+
+    dim3 const block_dim{NUM_THREADS_PER_BLOCK, 1U, 1U};
+    dim3 const grid_dim{
+        (static_cast<unsigned int>(n) + BLOCK_TILE_SIZE_X - 1U) /
+            BLOCK_TILE_SIZE_X,
+        (static_cast<unsigned int>(m) + BLOCK_TILE_SIZE_Y - 1U) /
+            BLOCK_TILE_SIZE_Y,
+        1U};
+    gemm_v06_vectorized_double_buffered<
+        T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+        WARP_TILE_SIZE_X, WARP_TILE_SIZE_Y, THREAD_TILE_SIZE_X,
+        THREAD_TILE_SIZE_Y, NUM_THREADS_PER_WARP_X, NUM_THREADS_PER_WARP_Y>
+        <<<grid_dim, block_dim, 0U, stream>>>(m, n, k, *alpha, A, lda, B, ldb,
+                                              *beta, C, ldc);
+    CHECK_LAST_CUDA_ERROR();
+}
+
+// Explicit instantiation.
+template void launch_gemm_kernel_v06_vectorized_double_buffered<float>(
+    size_t m, size_t n, size_t k, float const* alpha, float const* A,
+    size_t lda, float const* B, size_t ldb, float const* beta, float* C,
+    size_t ldc, cudaStream_t stream);
+template void launch_gemm_kernel_v06_vectorized_double_buffered<__half>(
+    size_t m, size_t n, size_t k, __half const* alpha, __half const* A,
+    size_t lda, __half const* B, size_t ldb, __half const* beta, __half* C,
+    size_t ldc, cudaStream_t stream);

upstream_ref/cuda_gemm_optimization/07_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose_wmma.cu

@@ -0,0 +1,247 @@
+#include <cuda_fp16.h>
+#include <mma.h>
+
+#include "cuda_gemm.hpp"
+#include "cuda_gemm_utils.cuh"
+#include "cuda_gemm_utils.hpp"
+
+// https://developer.nvidia.com/blog/cutlass-linear-algebra-cuda/
+// https://github.com/NVIDIA/cutlass/blob/b7508e337938137a699e486d8997646980acfc58/media/docs/programming_guidelines.md
+
+// GEMM kernel v07.
+// Each thread in the block processes THREAD_TILE_SIZE_Y *
+// THREAD_TILE_SIZE_X output values. Number of threads BLOCK_TILE_SIZE_Y *
+// BLOCK_TILE_SIZE_X / (THREAD_TILE_SIZE_Y * THREAD_TILE_SIZE_X)
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t BLOCK_TILE_SIZE_K, size_t BLOCK_TILE_SKEW_SIZE_X,
+          size_t BLOCK_TILE_SKEW_SIZE_Y, size_t WARP_TILE_SIZE_X,
+          size_t WARP_TILE_SIZE_Y, size_t WMMA_TILE_SIZE_X,
+          size_t WMMA_TILE_SIZE_Y, size_t WMMA_TILE_SIZE_K, size_t NUM_THREADS>
+__global__ void gemm_v07(size_t m, size_t n, size_t k, T alpha, T const* A,
+                         size_t lda, T const* B, size_t ldb, T beta, T* C,
+                         size_t ldc)
+{
+    constexpr size_t NUM_WARPS_X{BLOCK_TILE_SIZE_X / WARP_TILE_SIZE_X};
+    static_assert(BLOCK_TILE_SIZE_X % WARP_TILE_SIZE_X == 0U);
+    static_assert(BLOCK_TILE_SIZE_Y % WARP_TILE_SIZE_Y == 0U);
+
+    // Cache a tile of A and B in shared memory for data reuse.
+    __shared__ T A_thread_block_tile_transposed[BLOCK_TILE_SIZE_K]
+                                               [BLOCK_TILE_SIZE_Y +
+                                                BLOCK_TILE_SKEW_SIZE_Y];
+    __shared__ T B_thread_block_tile[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_X +
+                                                        BLOCK_TILE_SKEW_SIZE_X];
+
+    constexpr size_t NUM_WMMA_TILES_X{WARP_TILE_SIZE_X / WMMA_TILE_SIZE_X};
+    static_assert(WARP_TILE_SIZE_X % WMMA_TILE_SIZE_X == 0U);
+    constexpr size_t NUM_WMMA_TILES_Y{WARP_TILE_SIZE_Y / WMMA_TILE_SIZE_Y};
+    static_assert(WARP_TILE_SIZE_Y % WMMA_TILE_SIZE_Y == 0U);
+    constexpr size_t NUM_WMMA_TILES_K{BLOCK_TILE_SIZE_K / WMMA_TILE_SIZE_K};
+    static_assert(BLOCK_TILE_SIZE_K % WMMA_TILE_SIZE_K == 0U);
+
+    // Declare the fragments.
+    nvcuda::wmma::fragment<nvcuda::wmma::matrix_a, WMMA_TILE_SIZE_Y,
+                           WMMA_TILE_SIZE_X, WMMA_TILE_SIZE_K, T,
+                           nvcuda::wmma::col_major>
+        a_frags[NUM_WMMA_TILES_Y];
+    nvcuda::wmma::fragment<nvcuda::wmma::matrix_b, WMMA_TILE_SIZE_Y,
+                           WMMA_TILE_SIZE_X, WMMA_TILE_SIZE_K, T,
+                           nvcuda::wmma::row_major>
+        b_frags[NUM_WMMA_TILES_X];
+    nvcuda::wmma::fragment<nvcuda::wmma::accumulator, WMMA_TILE_SIZE_Y,
+                           WMMA_TILE_SIZE_X, WMMA_TILE_SIZE_K, T>
+        acc_frags[NUM_WMMA_TILES_Y][NUM_WMMA_TILES_X];
+    nvcuda::wmma::fragment<nvcuda::wmma::accumulator, WMMA_TILE_SIZE_Y,
+                           WMMA_TILE_SIZE_X, WMMA_TILE_SIZE_K, T>
+        c_frag;
+
+// Make sure the accumulator starts from 0.
+#pragma unroll
+    for (size_t wmma_tile_row_idx{0U}; wmma_tile_row_idx < NUM_WMMA_TILES_Y;
+         ++wmma_tile_row_idx)
+    {
+        for (size_t wmma_tile_col_idx{0U}; wmma_tile_col_idx < NUM_WMMA_TILES_X;
+             ++wmma_tile_col_idx)
+        {
+            nvcuda::wmma::fill_fragment(
+                acc_frags[wmma_tile_row_idx][wmma_tile_col_idx],
+                static_cast<T>(0));
+        }
+    }
+
+    size_t const thread_linear_idx{threadIdx.y * blockDim.x + threadIdx.x};
+    size_t const warp_linear_idx{thread_linear_idx / 32U};
+    size_t const warp_row_idx{warp_linear_idx / NUM_WARPS_X};
+    size_t const warp_col_idx{warp_linear_idx % NUM_WARPS_X};
+
+    // Number of outer loops to perform the sum of inner products.
+    // C_thread_block_tile =
+    // \sigma_{thread_block_tile_idx=0}^{num_thread_block_tiles-1} A[:,
+    // thread_block_tile_idx:BLOCK_TILE_SIZE_K] *
+    // B[thread_block_tile_idx:BLOCK_TILE_SIZE_K, :]
+    size_t const num_thread_block_tiles{(k + BLOCK_TILE_SIZE_K - 1) /
+                                        BLOCK_TILE_SIZE_K};
+
+    for (size_t thread_block_tile_idx{0U};
+         thread_block_tile_idx < num_thread_block_tiles;
+         ++thread_block_tile_idx)
+    {
+        load_data_from_global_memory_to_shared_memory_transposed<
+            T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+            NUM_THREADS, BLOCK_TILE_SKEW_SIZE_X, BLOCK_TILE_SKEW_SIZE_Y>(
+            A, lda, B, ldb, A_thread_block_tile_transposed, B_thread_block_tile,
+            thread_block_tile_idx, thread_linear_idx, m, n, k);
+        __syncthreads();
+
+// Perform A[:, thread_block_tile_idx:BLOCK_TILE_SIZE_K] *
+// B[thread_block_tile_idx:BLOCK_TILE_SIZE_K, :] where A[:,
+// thread_block_tile_idx:BLOCK_TILE_SIZE_K] and
+// B[thread_block_tile_idx:BLOCK_TILE_SIZE_K, :] are cached in the
+// shared memory as A_thread_block_tile and B_thread_block_tile,
+// respectively. This inner product is further decomposed to
+// BLOCK_TILE_SIZE_K outer products. A_thread_block_tile *
+// B_thread_block_tile = \sigma_{k_i=0}^{BLOCK_TILE_SIZE_K-1}
+// A_thread_block_tile[:, k_i] @ B_thread_block_tile[k_i, :] Note that
+// both A_thread_block_tile and B_thread_block_tile can be cached in the
+// register.
+#pragma unroll
+        for (size_t k_i{0U}; k_i < NUM_WMMA_TILES_K; ++k_i)
+        {
+#pragma unroll
+            for (size_t wmma_tile_row_idx{0U};
+                 wmma_tile_row_idx < NUM_WMMA_TILES_Y; ++wmma_tile_row_idx)
+            {
+                nvcuda::wmma::load_matrix_sync(
+                    a_frags[wmma_tile_row_idx],
+                    &A_thread_block_tile_transposed[k_i * WMMA_TILE_SIZE_K]
+                                                   [warp_row_idx *
+                                                        WARP_TILE_SIZE_Y +
+                                                    wmma_tile_row_idx *
+                                                        WMMA_TILE_SIZE_Y],
+                    BLOCK_TILE_SIZE_Y + BLOCK_TILE_SKEW_SIZE_Y);
+            }
+#pragma unroll
+            for (size_t wmma_tile_col_idx{0U};
+                 wmma_tile_col_idx < NUM_WMMA_TILES_X; ++wmma_tile_col_idx)
+            {
+                nvcuda::wmma::load_matrix_sync(
+                    b_frags[wmma_tile_col_idx],
+                    &B_thread_block_tile[k_i * WMMA_TILE_SIZE_K]
+                                        [warp_col_idx * WARP_TILE_SIZE_X +
+                                         wmma_tile_col_idx * WMMA_TILE_SIZE_X],
+                    BLOCK_TILE_SIZE_X + BLOCK_TILE_SKEW_SIZE_X);
+            }
+#pragma unroll
+            for (size_t wmma_tile_row_idx{0U};
+                 wmma_tile_row_idx < NUM_WMMA_TILES_Y; ++wmma_tile_row_idx)
+            {
+#pragma unroll
+                for (size_t wmma_tile_col_idx{0U};
+                     wmma_tile_col_idx < NUM_WMMA_TILES_X; ++wmma_tile_col_idx)
+                {
+                    // Perform the matrix multiplication.
+                    nvcuda::wmma::mma_sync(
+                        acc_frags[wmma_tile_row_idx][wmma_tile_col_idx],
+                        a_frags[wmma_tile_row_idx], b_frags[wmma_tile_col_idx],
+                        acc_frags[wmma_tile_row_idx][wmma_tile_col_idx]);
+                }
+            }
+        }
+        __syncthreads();
+    }
+
+// Write the results to DRAM.
+#pragma unroll
+    for (size_t wmma_tile_row_idx{0U}; wmma_tile_row_idx < NUM_WMMA_TILES_Y;
+         ++wmma_tile_row_idx)
+    {
+#pragma unroll
+        for (size_t wmma_tile_col_idx{0U}; wmma_tile_col_idx < NUM_WMMA_TILES_X;
+             ++wmma_tile_col_idx)
+        {
+            // Load the fragment from global memory.
+            nvcuda::wmma::load_matrix_sync(
+                c_frag,
+                &C[(blockIdx.y * BLOCK_TILE_SIZE_Y +
+                    warp_row_idx * WARP_TILE_SIZE_Y +
+                    wmma_tile_row_idx * WMMA_TILE_SIZE_Y) *
+                       n +
+                   blockIdx.x * BLOCK_TILE_SIZE_X +
+                   warp_col_idx * WARP_TILE_SIZE_X +
+                   wmma_tile_col_idx * WMMA_TILE_SIZE_X],
+                n, nvcuda::wmma::mem_row_major);
+            // Perform scaling and addition.
+            for (size_t i{0}; i < c_frag.num_elements; ++i)
+            {
+                c_frag.x[i] =
+                    alpha *
+                        acc_frags[wmma_tile_row_idx][wmma_tile_col_idx].x[i] +
+                    beta * c_frag.x[i];
+            }
+            // Store the fragment back to global memory.
+            nvcuda::wmma::store_matrix_sync(
+                &C[(blockIdx.y * BLOCK_TILE_SIZE_Y +
+                    warp_row_idx * WARP_TILE_SIZE_Y +
+                    wmma_tile_row_idx * WMMA_TILE_SIZE_Y) *
+                       n +
+                   blockIdx.x * BLOCK_TILE_SIZE_X +
+                   warp_col_idx * WARP_TILE_SIZE_X +
+                   wmma_tile_col_idx * WMMA_TILE_SIZE_X],
+                c_frag, n, nvcuda::wmma::mem_row_major);
+        }
+    }
+}
+
+template <typename T>
+void launch_gemm_kernel_v07(size_t m, size_t n, size_t k, T const* alpha,
+                            T const* A, size_t lda, T const* B, size_t ldb,
+                            T const* beta, T* C, size_t ldc,
+                            cudaStream_t stream)
+{
+    // Feel free to play with the block tile sizes.
+    // The algorithm correctness should always be guaranteed.
+    constexpr unsigned int BLOCK_TILE_SIZE_X{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_Y{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_K{16U};
+
+    constexpr unsigned int WARP_TILE_SIZE_X{32U};
+    constexpr unsigned int WARP_TILE_SIZE_Y{64U};
+    constexpr unsigned int NUM_WARPS_X{BLOCK_TILE_SIZE_X / WARP_TILE_SIZE_X};
+    constexpr unsigned int NUM_WARPS_Y{BLOCK_TILE_SIZE_Y / WARP_TILE_SIZE_Y};
+    static_assert(BLOCK_TILE_SIZE_X % WARP_TILE_SIZE_X == 0U);
+    static_assert(BLOCK_TILE_SIZE_Y % WARP_TILE_SIZE_Y == 0U);
+
+    // The skew size is used to avoid bank conflicts in shared memory.
+    constexpr size_t BLOCK_TILE_SKEW_SIZE_X{16U};
+    constexpr size_t BLOCK_TILE_SKEW_SIZE_Y{16U};
+
+    constexpr unsigned int WMMA_TILE_SIZE_X{16U};
+    constexpr unsigned int WMMA_TILE_SIZE_Y{16U};
+    constexpr unsigned int WMMA_TILE_SIZE_K{16U};
+
+    constexpr unsigned int NUM_THREADS_PER_BLOCK{NUM_WARPS_X * NUM_WARPS_Y *
+                                                 32U};
+
+    dim3 const block_dim{NUM_THREADS_PER_BLOCK, 1U, 1U};
+    dim3 const grid_dim{
+        (static_cast<unsigned int>(n) + BLOCK_TILE_SIZE_X - 1U) /
+            BLOCK_TILE_SIZE_X,
+        (static_cast<unsigned int>(m) + BLOCK_TILE_SIZE_Y - 1U) /
+            BLOCK_TILE_SIZE_Y,
+        1U};
+    gemm_v07<T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+             BLOCK_TILE_SKEW_SIZE_X, BLOCK_TILE_SKEW_SIZE_Y, WARP_TILE_SIZE_X,
+             WARP_TILE_SIZE_Y, WMMA_TILE_SIZE_X, WMMA_TILE_SIZE_Y,
+             WMMA_TILE_SIZE_K, NUM_THREADS_PER_BLOCK>
+        <<<grid_dim, block_dim, 0U, stream>>>(m, n, k, *alpha, A, lda, B, ldb,
+                                              *beta, C, ldc);
+    CHECK_LAST_CUDA_ERROR();
+}
+
+// Explicit instantiation.
+template void launch_gemm_kernel_v07<__half>(size_t m, size_t n, size_t k,
+                                             __half const* alpha,
+                                             __half const* A, size_t lda,
+                                             __half const* B, size_t ldb,
+                                             __half const* beta, __half* C,
+                                             size_t ldc, cudaStream_t stream);

upstream_ref/cuda_gemm_optimization/07_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose_wmma_vectorized_memory_access.cu

@@ -0,0 +1,246 @@
+#include <cuda_fp16.h>
+#include <mma.h>
+
+#include "cuda_gemm.hpp"
+#include "cuda_gemm_utils.cuh"
+#include "cuda_gemm_utils.hpp"
+
+// https://developer.nvidia.com/blog/cutlass-linear-algebra-cuda/
+// https://github.com/NVIDIA/cutlass/blob/b7508e337938137a699e486d8997646980acfc58/media/docs/programming_guidelines.md
+
+// GEMM kernel v07.
+// Each thread in the block processes THREAD_TILE_SIZE_Y *
+// THREAD_TILE_SIZE_X output values. Number of threads BLOCK_TILE_SIZE_Y *
+// BLOCK_TILE_SIZE_X / (THREAD_TILE_SIZE_Y * THREAD_TILE_SIZE_X)
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t BLOCK_TILE_SIZE_K, size_t BLOCK_TILE_SKEW_SIZE_X,
+          size_t BLOCK_TILE_SKEW_SIZE_Y, size_t WARP_TILE_SIZE_X,
+          size_t WARP_TILE_SIZE_Y, size_t WMMA_TILE_SIZE_X,
+          size_t WMMA_TILE_SIZE_Y, size_t WMMA_TILE_SIZE_K, size_t NUM_THREADS>
+__global__ void gemm_v07_vectorized(size_t m, size_t n, size_t k, T alpha,
+                                    T const* A, size_t lda, T const* B,
+                                    size_t ldb, T beta, T* C, size_t ldc)
+{
+    constexpr size_t NUM_WARPS_X{BLOCK_TILE_SIZE_X / WARP_TILE_SIZE_X};
+    static_assert(BLOCK_TILE_SIZE_X % WARP_TILE_SIZE_X == 0U);
+    static_assert(BLOCK_TILE_SIZE_Y % WARP_TILE_SIZE_Y == 0U);
+
+    // Cache a tile of A and B in shared memory for data reuse.
+    __shared__ T A_thread_block_tile_transposed[BLOCK_TILE_SIZE_K]
+                                               [BLOCK_TILE_SIZE_Y +
+                                                BLOCK_TILE_SKEW_SIZE_Y];
+    __shared__ T B_thread_block_tile[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_X +
+                                                        BLOCK_TILE_SKEW_SIZE_X];
+
+    constexpr size_t NUM_WMMA_TILES_X{WARP_TILE_SIZE_X / WMMA_TILE_SIZE_X};
+    static_assert(WARP_TILE_SIZE_X % WMMA_TILE_SIZE_X == 0U);
+    constexpr size_t NUM_WMMA_TILES_Y{WARP_TILE_SIZE_Y / WMMA_TILE_SIZE_Y};
+    static_assert(WARP_TILE_SIZE_Y % WMMA_TILE_SIZE_Y == 0U);
+    constexpr size_t NUM_WMMA_TILES_K{BLOCK_TILE_SIZE_K / WMMA_TILE_SIZE_K};
+    static_assert(BLOCK_TILE_SIZE_K % WMMA_TILE_SIZE_K == 0U);
+
+    // Declare the fragments.
+    nvcuda::wmma::fragment<nvcuda::wmma::matrix_a, WMMA_TILE_SIZE_Y,
+                           WMMA_TILE_SIZE_X, WMMA_TILE_SIZE_K, T,
+                           nvcuda::wmma::col_major>
+        a_frags[NUM_WMMA_TILES_Y];
+    nvcuda::wmma::fragment<nvcuda::wmma::matrix_b, WMMA_TILE_SIZE_Y,
+                           WMMA_TILE_SIZE_X, WMMA_TILE_SIZE_K, T,
+                           nvcuda::wmma::row_major>
+        b_frags[NUM_WMMA_TILES_X];
+    nvcuda::wmma::fragment<nvcuda::wmma::accumulator, WMMA_TILE_SIZE_Y,
+                           WMMA_TILE_SIZE_X, WMMA_TILE_SIZE_K, T>
+        acc_frags[NUM_WMMA_TILES_Y][NUM_WMMA_TILES_X];
+    nvcuda::wmma::fragment<nvcuda::wmma::accumulator, WMMA_TILE_SIZE_Y,
+                           WMMA_TILE_SIZE_X, WMMA_TILE_SIZE_K, T>
+        c_frag;
+
+// Make sure the accumulator starts from 0.
+#pragma unroll
+    for (size_t wmma_tile_row_idx{0U}; wmma_tile_row_idx < NUM_WMMA_TILES_Y;
+         ++wmma_tile_row_idx)
+    {
+        for (size_t wmma_tile_col_idx{0U}; wmma_tile_col_idx < NUM_WMMA_TILES_X;
+             ++wmma_tile_col_idx)
+        {
+            nvcuda::wmma::fill_fragment(
+                acc_frags[wmma_tile_row_idx][wmma_tile_col_idx],
+                static_cast<T>(0));
+        }
+    }
+
+    size_t const thread_linear_idx{threadIdx.y * blockDim.x + threadIdx.x};
+    size_t const warp_linear_idx{thread_linear_idx / 32U};
+    size_t const warp_row_idx{warp_linear_idx / NUM_WARPS_X};
+    size_t const warp_col_idx{warp_linear_idx % NUM_WARPS_X};
+
+    // Number of outer loops to perform the sum of inner products.
+    // C_thread_block_tile =
+    // \sigma_{thread_block_tile_idx=0}^{num_thread_block_tiles-1} A[:,
+    // thread_block_tile_idx:BLOCK_TILE_SIZE_K] *
+    // B[thread_block_tile_idx:BLOCK_TILE_SIZE_K, :]
+    size_t const num_thread_block_tiles{(k + BLOCK_TILE_SIZE_K - 1) /
+                                        BLOCK_TILE_SIZE_K};
+
+    for (size_t thread_block_tile_idx{0U};
+         thread_block_tile_idx < num_thread_block_tiles;
+         ++thread_block_tile_idx)
+    {
+        load_data_from_global_memory_to_shared_memory_transposed_vectorized<
+            T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+            NUM_THREADS, BLOCK_TILE_SKEW_SIZE_X, BLOCK_TILE_SKEW_SIZE_Y>(
+            A, lda, B, ldb, A_thread_block_tile_transposed, B_thread_block_tile,
+            thread_block_tile_idx, thread_linear_idx, m, n, k);
+        __syncthreads();
+
+// Perform A[:, thread_block_tile_idx:BLOCK_TILE_SIZE_K] *
+// B[thread_block_tile_idx:BLOCK_TILE_SIZE_K, :] where A[:,
+// thread_block_tile_idx:BLOCK_TILE_SIZE_K] and
+// B[thread_block_tile_idx:BLOCK_TILE_SIZE_K, :] are cached in the
+// shared memory as A_thread_block_tile and B_thread_block_tile,
+// respectively. This inner product is further decomposed to
+// BLOCK_TILE_SIZE_K outer products. A_thread_block_tile *
+// B_thread_block_tile = \sigma_{k_i=0}^{BLOCK_TILE_SIZE_K-1}
+// A_thread_block_tile[:, k_i] @ B_thread_block_tile[k_i, :] Note that
+// both A_thread_block_tile and B_thread_block_tile can be cached in the
+// register.
+#pragma unroll
+        for (size_t k_i{0U}; k_i < NUM_WMMA_TILES_K; ++k_i)
+        {
+#pragma unroll
+            for (size_t wmma_tile_row_idx{0U};
+                 wmma_tile_row_idx < NUM_WMMA_TILES_Y; ++wmma_tile_row_idx)
+            {
+                nvcuda::wmma::load_matrix_sync(
+                    a_frags[wmma_tile_row_idx],
+                    &A_thread_block_tile_transposed[k_i * WMMA_TILE_SIZE_K]
+                                                   [warp_row_idx *
+                                                        WARP_TILE_SIZE_Y +
+                                                    wmma_tile_row_idx *
+                                                        WMMA_TILE_SIZE_Y],
+                    BLOCK_TILE_SIZE_Y + BLOCK_TILE_SKEW_SIZE_Y);
+            }
+#pragma unroll
+            for (size_t wmma_tile_col_idx{0U};
+                 wmma_tile_col_idx < NUM_WMMA_TILES_X; ++wmma_tile_col_idx)
+            {
+                nvcuda::wmma::load_matrix_sync(
+                    b_frags[wmma_tile_col_idx],
+                    &B_thread_block_tile[k_i * WMMA_TILE_SIZE_K]
+                                        [warp_col_idx * WARP_TILE_SIZE_X +
+                                         wmma_tile_col_idx * WMMA_TILE_SIZE_X],
+                    BLOCK_TILE_SIZE_X + BLOCK_TILE_SKEW_SIZE_X);
+            }
+#pragma unroll
+            for (size_t wmma_tile_row_idx{0U};
+                 wmma_tile_row_idx < NUM_WMMA_TILES_Y; ++wmma_tile_row_idx)
+            {
+#pragma unroll
+                for (size_t wmma_tile_col_idx{0U};
+                     wmma_tile_col_idx < NUM_WMMA_TILES_X; ++wmma_tile_col_idx)
+                {
+                    // Perform the matrix multiplication.
+                    nvcuda::wmma::mma_sync(
+                        acc_frags[wmma_tile_row_idx][wmma_tile_col_idx],
+                        a_frags[wmma_tile_row_idx], b_frags[wmma_tile_col_idx],
+                        acc_frags[wmma_tile_row_idx][wmma_tile_col_idx]);
+                }
+            }
+        }
+        __syncthreads();
+    }
+
+// Write the results to DRAM.
+#pragma unroll
+    for (size_t wmma_tile_row_idx{0U}; wmma_tile_row_idx < NUM_WMMA_TILES_Y;
+         ++wmma_tile_row_idx)
+    {
+#pragma unroll
+        for (size_t wmma_tile_col_idx{0U}; wmma_tile_col_idx < NUM_WMMA_TILES_X;
+             ++wmma_tile_col_idx)
+        {
+            // Load the fragment from global memory.
+            nvcuda::wmma::load_matrix_sync(
+                c_frag,
+                &C[(blockIdx.y * BLOCK_TILE_SIZE_Y +
+                    warp_row_idx * WARP_TILE_SIZE_Y +
+                    wmma_tile_row_idx * WMMA_TILE_SIZE_Y) *
+                       n +
+                   blockIdx.x * BLOCK_TILE_SIZE_X +
+                   warp_col_idx * WARP_TILE_SIZE_X +
+                   wmma_tile_col_idx * WMMA_TILE_SIZE_X],
+                n, nvcuda::wmma::mem_row_major);
+            // Perform scaling and addition.
+            for (size_t i{0}; i < c_frag.num_elements; ++i)
+            {
+                c_frag.x[i] =
+                    alpha *
+                        acc_frags[wmma_tile_row_idx][wmma_tile_col_idx].x[i] +
+                    beta * c_frag.x[i];
+            }
+            // Store the fragment back to global memory.
+            nvcuda::wmma::store_matrix_sync(
+                &C[(blockIdx.y * BLOCK_TILE_SIZE_Y +
+                    warp_row_idx * WARP_TILE_SIZE_Y +
+                    wmma_tile_row_idx * WMMA_TILE_SIZE_Y) *
+                       n +
+                   blockIdx.x * BLOCK_TILE_SIZE_X +
+                   warp_col_idx * WARP_TILE_SIZE_X +
+                   wmma_tile_col_idx * WMMA_TILE_SIZE_X],
+                c_frag, n, nvcuda::wmma::mem_row_major);
+        }
+    }
+}
+
+template <typename T>
+void launch_gemm_kernel_v07_vectorized(size_t m, size_t n, size_t k,
+                                       T const* alpha, T const* A, size_t lda,
+                                       T const* B, size_t ldb, T const* beta,
+                                       T* C, size_t ldc, cudaStream_t stream)
+{
+    // Feel free to play with the block tile sizes.
+    // The algorithm correctness should always be guaranteed.
+    constexpr unsigned int BLOCK_TILE_SIZE_X{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_Y{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_K{16U};
+
+    // The skew size is used to avoid bank conflicts in shared memory.
+    constexpr size_t BLOCK_TILE_SKEW_SIZE_X{16U};
+    constexpr size_t BLOCK_TILE_SKEW_SIZE_Y{16U};
+
+    constexpr unsigned int WARP_TILE_SIZE_X{32U};
+    constexpr unsigned int WARP_TILE_SIZE_Y{64U};
+    constexpr unsigned int NUM_WARPS_X{BLOCK_TILE_SIZE_X / WARP_TILE_SIZE_X};
+    constexpr unsigned int NUM_WARPS_Y{BLOCK_TILE_SIZE_Y / WARP_TILE_SIZE_Y};
+    static_assert(BLOCK_TILE_SIZE_X % WARP_TILE_SIZE_X == 0U);
+    static_assert(BLOCK_TILE_SIZE_Y % WARP_TILE_SIZE_Y == 0U);
+
+    constexpr unsigned int WMMA_TILE_SIZE_X{16U};
+    constexpr unsigned int WMMA_TILE_SIZE_Y{16U};
+    constexpr unsigned int WMMA_TILE_SIZE_K{16U};
+
+    constexpr unsigned int NUM_THREADS_PER_BLOCK{NUM_WARPS_X * NUM_WARPS_Y *
+                                                 32U};
+
+    dim3 const block_dim{NUM_THREADS_PER_BLOCK, 1U, 1U};
+    dim3 const grid_dim{
+        (static_cast<unsigned int>(n) + BLOCK_TILE_SIZE_X - 1U) /
+            BLOCK_TILE_SIZE_X,
+        (static_cast<unsigned int>(m) + BLOCK_TILE_SIZE_Y - 1U) /
+            BLOCK_TILE_SIZE_Y,
+        1U};
+    gemm_v07_vectorized<T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y,
+                        BLOCK_TILE_SIZE_K, BLOCK_TILE_SKEW_SIZE_X,
+                        BLOCK_TILE_SKEW_SIZE_Y, WARP_TILE_SIZE_X,
+                        WARP_TILE_SIZE_Y, WMMA_TILE_SIZE_X, WMMA_TILE_SIZE_Y,
+                        WMMA_TILE_SIZE_K, NUM_THREADS_PER_BLOCK>
+        <<<grid_dim, block_dim, 0U, stream>>>(m, n, k, *alpha, A, lda, B, ldb,
+                                              *beta, C, ldc);
+    CHECK_LAST_CUDA_ERROR();
+}
+
+// Explicit instantiation.
+template void launch_gemm_kernel_v07_vectorized<__half>(
+    size_t m, size_t n, size_t k, __half const* alpha, __half const* A,
+    size_t lda, __half const* B, size_t ldb, __half const* beta, __half* C,
+    size_t ldc, cudaStream_t stream);

upstream_ref/cuda_gemm_optimization/07_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose_wmma_vectorized_memory_access_double_buffered.cu

@@ -0,0 +1,380 @@
+#include <cuda_fp16.h>
+#include <mma.h>
+
+#include "cuda_gemm.hpp"
+#include "cuda_gemm_utils.cuh"
+#include "cuda_gemm_utils.hpp"
+
+// https://developer.nvidia.com/blog/cutlass-linear-algebra-cuda/
+// https://github.com/NVIDIA/cutlass/blob/b7508e337938137a699e486d8997646980acfc58/media/docs/programming_guidelines.md
+
+template <
+    typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+    size_t BLOCK_TILE_SIZE_K, size_t WARP_TILE_SIZE_X, size_t WARP_TILE_SIZE_Y,
+    size_t WMMA_TILE_SIZE_X, size_t WMMA_TILE_SIZE_Y, size_t WMMA_TILE_SIZE_K,
+    size_t NUM_WMMA_TILES_X, size_t NUM_WMMA_TILES_Y, size_t NUM_WMMA_TILES_K,
+    size_t BLOCK_TILE_SKEW_SIZE_X, size_t BLOCK_TILE_SKEW_SIZE_Y>
+__device__ void process_data_from_shared_memory_using_wmma(
+    nvcuda::wmma::fragment<nvcuda::wmma::matrix_a, WMMA_TILE_SIZE_Y,
+                           WMMA_TILE_SIZE_X, WMMA_TILE_SIZE_K, T,
+                           nvcuda::wmma::col_major>
+        a_frags[NUM_WMMA_TILES_Y],
+    nvcuda::wmma::fragment<nvcuda::wmma::matrix_b, WMMA_TILE_SIZE_Y,
+                           WMMA_TILE_SIZE_X, WMMA_TILE_SIZE_K, T,
+                           nvcuda::wmma::row_major>
+        b_frags[NUM_WMMA_TILES_X],
+    nvcuda::wmma::fragment<nvcuda::wmma::accumulator, WMMA_TILE_SIZE_Y,
+                           WMMA_TILE_SIZE_X, WMMA_TILE_SIZE_K, T>
+        acc_frags[NUM_WMMA_TILES_Y][NUM_WMMA_TILES_X],
+    T const A_thread_block_tile_transposed[BLOCK_TILE_SIZE_K]
+                                          [BLOCK_TILE_SIZE_Y +
+                                           BLOCK_TILE_SKEW_SIZE_Y],
+    T const B_thread_block_tile[BLOCK_TILE_SIZE_K]
+                               [BLOCK_TILE_SIZE_X + BLOCK_TILE_SKEW_SIZE_X],
+    size_t warp_row_idx, size_t warp_col_idx)
+{
+#pragma unroll
+    for (size_t k_i{0U}; k_i < NUM_WMMA_TILES_K; ++k_i)
+    {
+#pragma unroll
+        for (size_t wmma_tile_row_idx{0U}; wmma_tile_row_idx < NUM_WMMA_TILES_Y;
+             ++wmma_tile_row_idx)
+        {
+            nvcuda::wmma::load_matrix_sync(
+                a_frags[wmma_tile_row_idx],
+                &A_thread_block_tile_transposed[k_i * WMMA_TILE_SIZE_K]
+                                               [warp_row_idx *
+                                                    WARP_TILE_SIZE_Y +
+                                                wmma_tile_row_idx *
+                                                    WMMA_TILE_SIZE_Y],
+                BLOCK_TILE_SIZE_Y + BLOCK_TILE_SKEW_SIZE_Y);
+        }
+#pragma unroll
+        for (size_t wmma_tile_col_idx{0U}; wmma_tile_col_idx < NUM_WMMA_TILES_X;
+             ++wmma_tile_col_idx)
+        {
+            nvcuda::wmma::load_matrix_sync(
+                b_frags[wmma_tile_col_idx],
+                &B_thread_block_tile[k_i * WMMA_TILE_SIZE_K]
+                                    [warp_col_idx * WARP_TILE_SIZE_X +
+                                     wmma_tile_col_idx * WMMA_TILE_SIZE_X],
+                BLOCK_TILE_SIZE_X + BLOCK_TILE_SKEW_SIZE_X);
+        }
+#pragma unroll
+        for (size_t wmma_tile_row_idx{0U}; wmma_tile_row_idx < NUM_WMMA_TILES_Y;
+             ++wmma_tile_row_idx)
+        {
+#pragma unroll
+            for (size_t wmma_tile_col_idx{0U};
+                 wmma_tile_col_idx < NUM_WMMA_TILES_X; ++wmma_tile_col_idx)
+            {
+                // Perform the matrix multiplication.
+                nvcuda::wmma::mma_sync(
+                    acc_frags[wmma_tile_row_idx][wmma_tile_col_idx],
+                    a_frags[wmma_tile_row_idx], b_frags[wmma_tile_col_idx],
+                    acc_frags[wmma_tile_row_idx][wmma_tile_col_idx]);
+            }
+        }
+    }
+}
+
+// GEMM kernel v07.
+// Each thread in the block processes THREAD_TILE_SIZE_Y *
+// THREAD_TILE_SIZE_X output values. Number of threads BLOCK_TILE_SIZE_Y *
+// BLOCK_TILE_SIZE_X / (THREAD_TILE_SIZE_Y * THREAD_TILE_SIZE_X)
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t BLOCK_TILE_SIZE_K, size_t BLOCK_TILE_SKEW_SIZE_X,
+          size_t BLOCK_TILE_SKEW_SIZE_Y, size_t WARP_TILE_SIZE_X,
+          size_t WARP_TILE_SIZE_Y, size_t WMMA_TILE_SIZE_X,
+          size_t WMMA_TILE_SIZE_Y, size_t WMMA_TILE_SIZE_K, size_t NUM_THREADS>
+__global__ void
+gemm_v07_vectorized_double_buffered(size_t m, size_t n, size_t k, T alpha,
+                                    T const* A, size_t lda, T const* B,
+                                    size_t ldb, T beta, T* C, size_t ldc)
+{
+    constexpr size_t NUM_WARPS_X{BLOCK_TILE_SIZE_X / WARP_TILE_SIZE_X};
+    constexpr size_t NUM_WARPS_Y{BLOCK_TILE_SIZE_Y / WARP_TILE_SIZE_Y};
+    static_assert(BLOCK_TILE_SIZE_X % WARP_TILE_SIZE_X == 0U);
+    static_assert(BLOCK_TILE_SIZE_Y % WARP_TILE_SIZE_Y == 0U);
+
+    constexpr size_t NUM_WMMA_TILES_X{WARP_TILE_SIZE_X / WMMA_TILE_SIZE_X};
+    static_assert(WARP_TILE_SIZE_X % WMMA_TILE_SIZE_X == 0U);
+    constexpr size_t NUM_WMMA_TILES_Y{WARP_TILE_SIZE_Y / WMMA_TILE_SIZE_Y};
+    static_assert(WARP_TILE_SIZE_Y % WMMA_TILE_SIZE_Y == 0U);
+    constexpr size_t NUM_WMMA_TILES_K{BLOCK_TILE_SIZE_K / WMMA_TILE_SIZE_K};
+    static_assert(BLOCK_TILE_SIZE_K % WMMA_TILE_SIZE_K == 0U);
+
+    constexpr size_t NUM_PIPELINES{2U};
+    // Only double buffer is supported in the implementation.
+    // But even more number of pipelines can be supported if the implementation
+    // is modified.
+    static_assert(NUM_PIPELINES == 2U);
+    static_assert((NUM_WARPS_X * NUM_WARPS_Y) % NUM_PIPELINES == 0U);
+    static_assert(NUM_THREADS % NUM_PIPELINES == 0U);
+    constexpr size_t NUM_THREADS_PER_PIPELINE{NUM_THREADS / NUM_PIPELINES};
+    constexpr size_t NUM_WARPS_PER_PIPELINE{(NUM_WARPS_X * NUM_WARPS_Y) /
+                                            NUM_PIPELINES};
+
+    // Cache a tile of A and B in shared memory for data reuse.
+    __shared__ T
+        A_thread_block_tile_transposed[NUM_PIPELINES][BLOCK_TILE_SIZE_K]
+                                      [BLOCK_TILE_SIZE_Y +
+                                       BLOCK_TILE_SKEW_SIZE_Y];
+    __shared__ T
+        B_thread_block_tile[NUM_PIPELINES][BLOCK_TILE_SIZE_K]
+                           [BLOCK_TILE_SIZE_X + BLOCK_TILE_SKEW_SIZE_X];
+
+    // Declare the fragments.
+    nvcuda::wmma::fragment<nvcuda::wmma::matrix_a, WMMA_TILE_SIZE_Y,
+                           WMMA_TILE_SIZE_X, WMMA_TILE_SIZE_K, T,
+                           nvcuda::wmma::col_major>
+        a_frags[NUM_WMMA_TILES_Y];
+    nvcuda::wmma::fragment<nvcuda::wmma::matrix_b, WMMA_TILE_SIZE_Y,
+                           WMMA_TILE_SIZE_X, WMMA_TILE_SIZE_K, T,
+                           nvcuda::wmma::row_major>
+        b_frags[NUM_WMMA_TILES_X];
+    nvcuda::wmma::fragment<nvcuda::wmma::accumulator, WMMA_TILE_SIZE_Y,
+                           WMMA_TILE_SIZE_X, WMMA_TILE_SIZE_K, T>
+        acc_frags[NUM_WMMA_TILES_Y][NUM_WMMA_TILES_X];
+    nvcuda::wmma::fragment<nvcuda::wmma::accumulator, WMMA_TILE_SIZE_Y,
+                           WMMA_TILE_SIZE_X, WMMA_TILE_SIZE_K, T>
+        c_frag;
+
+// Make sure the accumulator starts from 0.
+#pragma unroll
+    for (size_t wmma_tile_row_idx{0U}; wmma_tile_row_idx < NUM_WMMA_TILES_Y;
+         ++wmma_tile_row_idx)
+    {
+        for (size_t wmma_tile_col_idx{0U}; wmma_tile_col_idx < NUM_WMMA_TILES_X;
+             ++wmma_tile_col_idx)
+        {
+            nvcuda::wmma::fill_fragment(
+                acc_frags[wmma_tile_row_idx][wmma_tile_col_idx],
+                static_cast<T>(0));
+        }
+    }
+
+    size_t const thread_linear_idx{threadIdx.y * blockDim.x + threadIdx.x};
+    size_t const warp_linear_idx{thread_linear_idx / 32U};
+    size_t const warp_row_idx{warp_linear_idx / NUM_WARPS_X};
+    size_t const warp_col_idx{warp_linear_idx % NUM_WARPS_X};
+    // Separate the warps to different pipelines.
+    size_t const pipeline_index{warp_linear_idx / NUM_WARPS_PER_PIPELINE};
+
+    // Number of outer loops to perform the sum of inner products.
+    // C_thread_block_tile =
+    // \sigma_{thread_block_tile_idx=0}^{num_thread_block_tiles-1} A[:,
+    // thread_block_tile_idx:BLOCK_TILE_SIZE_K] *
+    // B[thread_block_tile_idx:BLOCK_TILE_SIZE_K, :]
+    size_t const num_thread_block_tiles{(k + BLOCK_TILE_SIZE_K - 1) /
+                                        BLOCK_TILE_SIZE_K};
+
+    if (pipeline_index == 0U)
+    {
+        // Pipeline 0 warps load buffer 0.
+        load_data_from_global_memory_to_shared_memory_transposed_vectorized<
+            T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+            NUM_THREADS_PER_PIPELINE, BLOCK_TILE_SKEW_SIZE_X,
+            BLOCK_TILE_SKEW_SIZE_Y>(
+            A, lda, B, ldb, A_thread_block_tile_transposed[pipeline_index],
+            B_thread_block_tile[pipeline_index], 0U,
+            thread_linear_idx - pipeline_index * NUM_THREADS_PER_PIPELINE, m, n,
+            k);
+    }
+    __syncthreads();
+
+    for (size_t thread_block_tile_idx{0U};
+         thread_block_tile_idx < num_thread_block_tiles;
+         thread_block_tile_idx += NUM_PIPELINES)
+    {
+        if (pipeline_index == 0U)
+        {
+            // Pipeline 0 warps process buffer 0.
+            process_data_from_shared_memory_using_wmma<
+                T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+                WARP_TILE_SIZE_X, WARP_TILE_SIZE_Y, WMMA_TILE_SIZE_X,
+                WMMA_TILE_SIZE_Y, WMMA_TILE_SIZE_K, NUM_WMMA_TILES_X,
+                NUM_WMMA_TILES_Y, NUM_WMMA_TILES_K, BLOCK_TILE_SKEW_SIZE_X,
+                BLOCK_TILE_SKEW_SIZE_Y>(
+                a_frags, b_frags, acc_frags,
+                A_thread_block_tile_transposed[pipeline_index],
+                B_thread_block_tile[pipeline_index], warp_row_idx,
+                warp_col_idx);
+            __syncthreads();
+
+            // Pipeline 0 warps process buffer 1.
+            if (thread_block_tile_idx + 1U < num_thread_block_tiles)
+            {
+                process_data_from_shared_memory_using_wmma<
+                    T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+                    WARP_TILE_SIZE_X, WARP_TILE_SIZE_Y, WMMA_TILE_SIZE_X,
+                    WMMA_TILE_SIZE_Y, WMMA_TILE_SIZE_K, NUM_WMMA_TILES_X,
+                    NUM_WMMA_TILES_Y, NUM_WMMA_TILES_K, BLOCK_TILE_SKEW_SIZE_X,
+                    BLOCK_TILE_SKEW_SIZE_Y>(
+                    a_frags, b_frags, acc_frags,
+                    A_thread_block_tile_transposed[pipeline_index + 1],
+                    B_thread_block_tile[pipeline_index + 1], warp_row_idx,
+                    warp_col_idx);
+            }
+            __syncthreads();
+
+            // Pipeline 0 warps load buffer 0.
+            if (thread_block_tile_idx + 2U < num_thread_block_tiles)
+            {
+                load_data_from_global_memory_to_shared_memory_transposed_vectorized<
+                    T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+                    NUM_THREADS_PER_PIPELINE, BLOCK_TILE_SKEW_SIZE_X,
+                    BLOCK_TILE_SKEW_SIZE_Y>(
+                    A, lda, B, ldb,
+                    A_thread_block_tile_transposed[pipeline_index],
+                    B_thread_block_tile[pipeline_index],
+                    thread_block_tile_idx + 2,
+                    thread_linear_idx -
+                        pipeline_index * NUM_THREADS_PER_PIPELINE,
+                    m, n, k);
+            }
+            __syncthreads();
+        }
+        else
+        {
+            // Pipeline 1 warps load buffer 1.
+            if (thread_block_tile_idx + 1U < num_thread_block_tiles)
+            {
+                load_data_from_global_memory_to_shared_memory_transposed_vectorized<
+                    T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+                    NUM_THREADS_PER_PIPELINE, BLOCK_TILE_SKEW_SIZE_X,
+                    BLOCK_TILE_SKEW_SIZE_Y>(
+                    A, lda, B, ldb,
+                    A_thread_block_tile_transposed[pipeline_index],
+                    B_thread_block_tile[pipeline_index],
+                    thread_block_tile_idx + 1,
+                    thread_linear_idx -
+                        pipeline_index * NUM_THREADS_PER_PIPELINE,
+                    m, n, k);
+            }
+            __syncthreads();
+
+            // Pipeline 1 warps process buffer 0.
+            process_data_from_shared_memory_using_wmma<
+                T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+                WARP_TILE_SIZE_X, WARP_TILE_SIZE_Y, WMMA_TILE_SIZE_X,
+                WMMA_TILE_SIZE_Y, WMMA_TILE_SIZE_K, NUM_WMMA_TILES_X,
+                NUM_WMMA_TILES_Y, NUM_WMMA_TILES_K, BLOCK_TILE_SKEW_SIZE_X,
+                BLOCK_TILE_SKEW_SIZE_Y>(
+                a_frags, b_frags, acc_frags,
+                A_thread_block_tile_transposed[pipeline_index - 1],
+                B_thread_block_tile[pipeline_index - 1], warp_row_idx,
+                warp_col_idx);
+            __syncthreads();
+
+            // Pipeline 1 warps process buffer 1.
+            if (thread_block_tile_idx + 1U < num_thread_block_tiles)
+            {
+                process_data_from_shared_memory_using_wmma<
+                    T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+                    WARP_TILE_SIZE_X, WARP_TILE_SIZE_Y, WMMA_TILE_SIZE_X,
+                    WMMA_TILE_SIZE_Y, WMMA_TILE_SIZE_K, NUM_WMMA_TILES_X,
+                    NUM_WMMA_TILES_Y, NUM_WMMA_TILES_K, BLOCK_TILE_SKEW_SIZE_X,
+                    BLOCK_TILE_SKEW_SIZE_Y>(
+                    a_frags, b_frags, acc_frags,
+                    A_thread_block_tile_transposed[pipeline_index],
+                    B_thread_block_tile[pipeline_index], warp_row_idx,
+                    warp_col_idx);
+            }
+            __syncthreads();
+        }
+    }
+
+// Write the results to DRAM.
+#pragma unroll
+    for (size_t wmma_tile_row_idx{0U}; wmma_tile_row_idx < NUM_WMMA_TILES_Y;
+         ++wmma_tile_row_idx)
+    {
+#pragma unroll
+        for (size_t wmma_tile_col_idx{0U}; wmma_tile_col_idx < NUM_WMMA_TILES_X;
+             ++wmma_tile_col_idx)
+        {
+            // Load the fragment from global memory.
+            nvcuda::wmma::load_matrix_sync(
+                c_frag,
+                &C[(blockIdx.y * BLOCK_TILE_SIZE_Y +
+                    warp_row_idx * WARP_TILE_SIZE_Y +
+                    wmma_tile_row_idx * WMMA_TILE_SIZE_Y) *
+                       n +
+                   blockIdx.x * BLOCK_TILE_SIZE_X +
+                   warp_col_idx * WARP_TILE_SIZE_X +
+                   wmma_tile_col_idx * WMMA_TILE_SIZE_X],
+                n, nvcuda::wmma::mem_row_major);
+            // Perform scaling and addition.
+            for (size_t i{0}; i < c_frag.num_elements; ++i)
+            {
+                c_frag.x[i] =
+                    alpha *
+                        acc_frags[wmma_tile_row_idx][wmma_tile_col_idx].x[i] +
+                    beta * c_frag.x[i];
+            }
+            // Store the fragment back to global memory.
+            nvcuda::wmma::store_matrix_sync(
+                &C[(blockIdx.y * BLOCK_TILE_SIZE_Y +
+                    warp_row_idx * WARP_TILE_SIZE_Y +
+                    wmma_tile_row_idx * WMMA_TILE_SIZE_Y) *
+                       n +
+                   blockIdx.x * BLOCK_TILE_SIZE_X +
+                   warp_col_idx * WARP_TILE_SIZE_X +
+                   wmma_tile_col_idx * WMMA_TILE_SIZE_X],
+                c_frag, n, nvcuda::wmma::mem_row_major);
+        }
+    }
+}
+
+template <typename T>
+void launch_gemm_kernel_v07_vectorized_double_buffered(
+    size_t m, size_t n, size_t k, T const* alpha, T const* A, size_t lda,
+    T const* B, size_t ldb, T const* beta, T* C, size_t ldc,
+    cudaStream_t stream)
+{
+    // Feel free to play with the block tile sizes.
+    // The algorithm correctness should always be guaranteed.
+    constexpr unsigned int BLOCK_TILE_SIZE_X{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_Y{128U};
+    constexpr unsigned int BLOCK_TILE_SIZE_K{16U};
+
+    // The skew size is used to avoid bank conflicts in shared memory.
+    constexpr size_t BLOCK_TILE_SKEW_SIZE_X{16U};
+    constexpr size_t BLOCK_TILE_SKEW_SIZE_Y{16U};
+
+    constexpr unsigned int WARP_TILE_SIZE_X{32U};
+    constexpr unsigned int WARP_TILE_SIZE_Y{64U};
+    constexpr unsigned int NUM_WARPS_X{BLOCK_TILE_SIZE_X / WARP_TILE_SIZE_X};
+    constexpr unsigned int NUM_WARPS_Y{BLOCK_TILE_SIZE_Y / WARP_TILE_SIZE_Y};
+    static_assert(BLOCK_TILE_SIZE_X % WARP_TILE_SIZE_X == 0U);
+    static_assert(BLOCK_TILE_SIZE_Y % WARP_TILE_SIZE_Y == 0U);
+
+    constexpr unsigned int WMMA_TILE_SIZE_X{16U};
+    constexpr unsigned int WMMA_TILE_SIZE_Y{16U};
+    constexpr unsigned int WMMA_TILE_SIZE_K{16U};
+
+    constexpr unsigned int NUM_THREADS_PER_BLOCK{NUM_WARPS_X * NUM_WARPS_Y *
+                                                 32U};
+
+    dim3 const block_dim{NUM_THREADS_PER_BLOCK, 1U, 1U};
+    dim3 const grid_dim{
+        (static_cast<unsigned int>(n) + BLOCK_TILE_SIZE_X - 1U) /
+            BLOCK_TILE_SIZE_X,
+        (static_cast<unsigned int>(m) + BLOCK_TILE_SIZE_Y - 1U) /
+            BLOCK_TILE_SIZE_Y,
+        1U};
+    gemm_v07_vectorized_double_buffered<
+        T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
+        BLOCK_TILE_SKEW_SIZE_X, BLOCK_TILE_SKEW_SIZE_Y, WARP_TILE_SIZE_X,
+        WARP_TILE_SIZE_Y, WMMA_TILE_SIZE_X, WMMA_TILE_SIZE_Y, WMMA_TILE_SIZE_K,
+        NUM_THREADS_PER_BLOCK><<<grid_dim, block_dim, 0U, stream>>>(
+        m, n, k, *alpha, A, lda, B, ldb, *beta, C, ldc);
+    CHECK_LAST_CUDA_ERROR();
+}
+
+// Explicit instantiation.
+template void launch_gemm_kernel_v07_vectorized_double_buffered<__half>(
+    size_t m, size_t n, size_t k, __half const* alpha, __half const* A,
+    size_t lda, __half const* B, size_t ldb, __half const* beta, __half* C,
+    size_t ldc, cudaStream_t stream);

upstream_ref/cuda_gemm_optimization/cuda_gemm.hpp

@@ -0,0 +1,98 @@
+#ifndef CUDA_GEMM_HPP
+#define CUDA_GEMM_HPP
+
+#include <cuda_runtime.h>
+
+template <typename T>
+void launch_gemm_kernel_v00(size_t m, size_t n, size_t k, T const* alpha,
+                            T const* A, size_t lda, T const* B, size_t ldb,
+                            T const* beta, T* C, size_t ldc,
+                            cudaStream_t stream);
+
+template <typename T>
+void launch_gemm_kernel_v01(size_t m, size_t n, size_t k, T const* alpha,
+                            T const* A, size_t lda, T const* B, size_t ldb,
+                            T const* beta, T* C, size_t ldc,
+                            cudaStream_t stream);
+
+template <typename T>
+void launch_gemm_kernel_v02(size_t m, size_t n, size_t k, T const* alpha,
+                            T const* A, size_t lda, T const* B, size_t ldb,
+                            T const* beta, T* C, size_t ldc,
+                            cudaStream_t stream);
+
+template <typename T>
+void launch_gemm_kernel_v02_vectorized(size_t m, size_t n, size_t k,
+                                       T const* alpha, T const* A, size_t lda,
+                                       T const* B, size_t ldb, T const* beta,
+                                       T* C, size_t ldc, cudaStream_t stream);
+
+template <typename T>
+void launch_gemm_kernel_v03(size_t m, size_t n, size_t k, T const* alpha,
+                            T const* A, size_t lda, T const* B, size_t ldb,
+                            T const* beta, T* C, size_t ldc,
+                            cudaStream_t stream);
+
+template <typename T>
+void launch_gemm_kernel_v03_vectorized(size_t m, size_t n, size_t k,
+                                       T const* alpha, T const* A, size_t lda,
+                                       T const* B, size_t ldb, T const* beta,
+                                       T* C, size_t ldc, cudaStream_t stream);
+template <typename T>
+void launch_gemm_kernel_v04(size_t m, size_t n, size_t k, T const* alpha,
+                            T const* A, size_t lda, T const* B, size_t ldb,
+                            T const* beta, T* C, size_t ldc,
+                            cudaStream_t stream);
+
+template <typename T>
+void launch_gemm_kernel_v04_vectorized(size_t m, size_t n, size_t k,
+                                       T const* alpha, T const* A, size_t lda,
+                                       T const* B, size_t ldb, T const* beta,
+                                       T* C, size_t ldc, cudaStream_t stream);
+
+template <typename T>
+void launch_gemm_kernel_v05(size_t m, size_t n, size_t k, T const* alpha,
+                            T const* A, size_t lda, T const* B, size_t ldb,
+                            T const* beta, T* C, size_t ldc,
+                            cudaStream_t stream);
+template <typename T>
+void launch_gemm_kernel_v05_vectorized(size_t m, size_t n, size_t k,
+                                       T const* alpha, T const* A, size_t lda,
+                                       T const* B, size_t ldb, T const* beta,
+                                       T* C, size_t ldc, cudaStream_t stream);
+
+template <typename T>
+void launch_gemm_kernel_v06(size_t m, size_t n, size_t k, T const* alpha,
+                            T const* A, size_t lda, T const* B, size_t ldb,
+                            T const* beta, T* C, size_t ldc,
+                            cudaStream_t stream);
+template <typename T>
+void launch_gemm_kernel_v06_vectorized(size_t m, size_t n, size_t k,
+                                       T const* alpha, T const* A, size_t lda,
+                                       T const* B, size_t ldb, T const* beta,
+                                       T* C, size_t ldc, cudaStream_t stream);
+
+template <typename T>
+void launch_gemm_kernel_v06_vectorized_double_buffered(
+    size_t m, size_t n, size_t k, T const* alpha, T const* A, size_t lda,
+    T const* B, size_t ldb, T const* beta, T* C, size_t ldc,
+    cudaStream_t stream);
+
+template <typename T>
+void launch_gemm_kernel_v07(size_t m, size_t n, size_t k, T const* alpha,
+                            T const* A, size_t lda, T const* B, size_t ldb,
+                            T const* beta, T* C, size_t ldc,
+                            cudaStream_t stream);
+
+template <typename T>
+void launch_gemm_kernel_v07_vectorized(size_t m, size_t n, size_t k,
+                                       T const* alpha, T const* A, size_t lda,
+                                       T const* B, size_t ldb, T const* beta,
+                                       T* C, size_t ldc, cudaStream_t stream);
+
+template <typename T>
+void launch_gemm_kernel_v07_vectorized_double_buffered(
+    size_t m, size_t n, size_t k, T const* alpha, T const* A, size_t lda,
+    T const* B, size_t ldb, T const* beta, T* C, size_t ldc,
+    cudaStream_t stream);
+#endif

upstream_ref/cuda_gemm_optimization/cuda_gemm_utils.cu

@@ -0,0 +1,29 @@
+#include <iostream>
+
+#include <cuda_runtime.h>
+
+#include "cuda_gemm_utils.hpp"
+
+void check_cuda(cudaError_t err, const char* const func, const char* const file,
+                const int line)
+{
+    if (err != cudaSuccess)
+    {
+        std::cerr << "CUDA Runtime Error at: " << file << ":" << line
+                  << std::endl;
+        std::cerr << cudaGetErrorString(err) << " " << func << std::endl;
+        std::exit(EXIT_FAILURE);
+    }
+}
+
+void check_cuda_last(const char* const file, const int line)
+{
+    cudaError_t const err{cudaGetLastError()};
+    if (err != cudaSuccess)
+    {
+        std::cerr << "CUDA Runtime Error at: " << file << ":" << line
+                  << std::endl;
+        std::cerr << cudaGetErrorString(err) << std::endl;
+        std::exit(EXIT_FAILURE);
+    }
+}

upstream_ref/cuda_gemm_optimization/cuda_gemm_utils.cuh

@@ -0,0 +1,486 @@
+#ifndef CUDA_GEMM_UTILS_CUH
+#define CUDA_GEMM_UTILS_CUH
+
+#include <cuda_runtime.h>
+
+#include "cuda_gemm_utils.hpp"
+
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t BLOCK_TILE_SIZE_K, size_t NUM_THREADS,
+          size_t BLOCK_TILE_SKEW_SIZE_X = 0U,
+          size_t BLOCK_TILE_SKEW_SIZE_K = 0U>
+__device__ void load_data_from_global_memory_to_shared_memory(
+    T const* A, size_t lda, T const* B, size_t ldb,
+    T A_thread_block_tile[BLOCK_TILE_SIZE_Y]
+                         [BLOCK_TILE_SIZE_K + BLOCK_TILE_SKEW_SIZE_K],
+    T B_thread_block_tile[BLOCK_TILE_SIZE_K]
+                         [BLOCK_TILE_SIZE_X + BLOCK_TILE_SKEW_SIZE_X],
+    size_t thread_block_tile_idx, size_t thread_linear_idx, size_t m, size_t n,
+    size_t k)
+{
+    // Load data from A on DRAM to A_thread_block_tile on shared memory.
+#pragma unroll
+    for (size_t load_idx{0U};
+         load_idx < (BLOCK_TILE_SIZE_Y * BLOCK_TILE_SIZE_K + NUM_THREADS - 1U) /
+                        NUM_THREADS;
+         ++load_idx)
+    {
+        size_t const A_thread_block_tile_row_idx{
+            (thread_linear_idx + load_idx * NUM_THREADS) / BLOCK_TILE_SIZE_K};
+        size_t const A_thread_block_tile_col_idx{
+            (thread_linear_idx + load_idx * NUM_THREADS) % BLOCK_TILE_SIZE_K};
+        size_t const A_row_idx{blockIdx.y * BLOCK_TILE_SIZE_Y +
+                               A_thread_block_tile_row_idx};
+        size_t const A_col_idx{thread_block_tile_idx * BLOCK_TILE_SIZE_K +
+                               A_thread_block_tile_col_idx};
+
+        // These boundary checks might slow down the kernel to some extent.
+        // But they guarantee the correctness of the kernel for all
+        // different GEMM configurations.
+        T val{static_cast<T>(0)};
+        if (A_row_idx < m && A_col_idx < k)
+        {
+            val = A[A_row_idx * lda + A_col_idx];
+        }
+        // This if will slow down the kernel.
+        // Add static asserts from the host code to guarantee this if is
+        // always true.
+        static_assert(BLOCK_TILE_SIZE_K * BLOCK_TILE_SIZE_Y % NUM_THREADS ==
+                      0U);
+        // if (A_thread_block_tile_row_idx < BLOCK_TILE_SIZE_Y &&
+        //     A_thread_block_tile_col_idx < BLOCK_TILE_SIZE_K)
+        // {
+        //     A_thread_block_tile[A_thread_block_tile_row_idx]
+        //                        [A_thread_block_tile_col_idx] = val;
+        // }
+        A_thread_block_tile[A_thread_block_tile_row_idx]
+                           [A_thread_block_tile_col_idx] = val;
+    }
+// Load data from B on DRAM to B_thread_block_tile on shared memory.
+#pragma unroll
+    for (size_t load_idx{0U};
+         load_idx < (BLOCK_TILE_SIZE_K * BLOCK_TILE_SIZE_X + NUM_THREADS - 1U) /
+                        NUM_THREADS;
+         ++load_idx)
+    {
+        size_t const B_thread_block_tile_row_idx{
+            (thread_linear_idx + load_idx * NUM_THREADS) / BLOCK_TILE_SIZE_X};
+        size_t const B_thread_block_tile_col_idx{
+            (thread_linear_idx + load_idx * NUM_THREADS) % BLOCK_TILE_SIZE_X};
+        size_t const B_row_idx{thread_block_tile_idx * BLOCK_TILE_SIZE_K +
+                               B_thread_block_tile_row_idx};
+        size_t const B_col_idx{blockIdx.x * BLOCK_TILE_SIZE_X +
+                               B_thread_block_tile_col_idx};
+
+        // These boundary checks might slow down the kernel to some extent.
+        // But they guarantee the correctness of the kernel for all
+        // different GEMM configurations.
+        T val{static_cast<T>(0)};
+        if (B_row_idx < k && B_col_idx < n)
+        {
+            val = B[B_row_idx * ldb + B_col_idx];
+        }
+        // This if will slow down the kernel.
+        // Add static asserts from the host code to guarantee this if is
+        // always true.
+        static_assert(BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_K % NUM_THREADS ==
+                      0U);
+        // if (B_thread_block_tile_row_idx < BLOCK_TILE_SIZE_K &&
+        //     B_thread_block_tile_col_idx < BLOCK_TILE_SIZE_X)
+        // {
+        //     B_thread_block_tile[B_thread_block_tile_row_idx]
+        //                        [B_thread_block_tile_col_idx] = val;
+        // }
+        B_thread_block_tile[B_thread_block_tile_row_idx]
+                           [B_thread_block_tile_col_idx] = val;
+    }
+}
+
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t BLOCK_TILE_SIZE_K, size_t NUM_THREADS,
+          size_t BLOCK_TILE_SKEW_SIZE_X = 0U,
+          size_t BLOCK_TILE_SKEW_SIZE_Y = 0U>
+__device__ void load_data_from_global_memory_to_shared_memory_transposed(
+    T const* A, size_t lda, T const* B, size_t ldb,
+    T A_thread_block_tile_transposed[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_Y +
+                                                        BLOCK_TILE_SKEW_SIZE_Y],
+    T B_thread_block_tile[BLOCK_TILE_SIZE_K]
+                         [BLOCK_TILE_SIZE_X + BLOCK_TILE_SKEW_SIZE_X],
+    size_t thread_block_tile_idx, size_t thread_linear_idx, size_t m, size_t n,
+    size_t k)
+{
+// Load data from A on DRAM to A_thread_block_tile on shared memory.
+#pragma unroll
+    for (size_t load_idx{0U};
+         load_idx < (BLOCK_TILE_SIZE_Y * BLOCK_TILE_SIZE_K + NUM_THREADS - 1U) /
+                        NUM_THREADS;
+         ++load_idx)
+    {
+        size_t const A_thread_block_tile_row_idx{
+            (thread_linear_idx + load_idx * NUM_THREADS) / BLOCK_TILE_SIZE_K};
+        size_t const A_thread_block_tile_col_idx{
+            (thread_linear_idx + load_idx * NUM_THREADS) % BLOCK_TILE_SIZE_K};
+        size_t const A_row_idx{blockIdx.y * BLOCK_TILE_SIZE_Y +
+                               A_thread_block_tile_row_idx};
+        size_t const A_col_idx{thread_block_tile_idx * BLOCK_TILE_SIZE_K +
+                               A_thread_block_tile_col_idx};
+
+        // These boundary checks might slow down the kernel to some extent.
+        // But they guarantee the correctness of the kernel for all
+        // different GEMM configurations.
+        T val{static_cast<T>(0)};
+        if (A_row_idx < m && A_col_idx < k)
+        {
+            val = A[A_row_idx * lda + A_col_idx];
+        }
+        // Removing the if will give another ~2 FLOPs performance on RTX
+        // 3090. But it will make the kernel incorrect for some GEMM
+        // configurations. T val{A[A_row_idx * lda + A_col_idx]}; This if
+        // will slow down the kernel. Add static asserts from the host code
+        // to guarantee this if is always true.
+        static_assert(BLOCK_TILE_SIZE_K * BLOCK_TILE_SIZE_Y % NUM_THREADS ==
+                      0U);
+        // if (A_thread_block_tile_row_idx < BLOCK_TILE_SIZE_Y &&
+        //     A_thread_block_tile_col_idx < BLOCK_TILE_SIZE_K)
+        // {
+        //     A_thread_block_tile[A_thread_block_tile_row_idx]
+        //                        [A_thread_block_tile_col_idx] = val;
+        // }
+        A_thread_block_tile_transposed[A_thread_block_tile_col_idx]
+                                      [A_thread_block_tile_row_idx] = val;
+    }
+// Load data from B on DRAM to B_thread_block_tile on shared memory.
+#pragma unroll
+    for (size_t load_idx{0U};
+         load_idx < (BLOCK_TILE_SIZE_K * BLOCK_TILE_SIZE_X + NUM_THREADS - 1U) /
+                        NUM_THREADS;
+         ++load_idx)
+    {
+        size_t const B_thread_block_tile_row_idx{
+            (thread_linear_idx + load_idx * NUM_THREADS) / BLOCK_TILE_SIZE_X};
+        size_t const B_thread_block_tile_col_idx{
+            (thread_linear_idx + load_idx * NUM_THREADS) % BLOCK_TILE_SIZE_X};
+        size_t const B_row_idx{thread_block_tile_idx * BLOCK_TILE_SIZE_K +
+                               B_thread_block_tile_row_idx};
+        size_t const B_col_idx{blockIdx.x * BLOCK_TILE_SIZE_X +
+                               B_thread_block_tile_col_idx};
+
+        // These boundary checks might slow down the kernel to some extent.
+        // But they guarantee the correctness of the kernel for all
+        // different GEMM configurations.
+        T val{static_cast<T>(0)};
+        if (B_row_idx < k && B_col_idx < n)
+        {
+            val = B[B_row_idx * ldb + B_col_idx];
+        }
+        // Removing the if will give another ~2 FLOPs performance on RTX
+        // 3090. But it will make the kernel incorrect for some GEMM
+        // configurations. T val{B[B_row_idx * ldb + B_col_idx]}; This if
+        // will slow down the kernel. Add static asserts from the host code
+        // to guarantee this if is always true.
+        static_assert(BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_K % NUM_THREADS ==
+                      0U);
+        // if (B_thread_block_tile_row_idx < BLOCK_TILE_SIZE_K &&
+        //     B_thread_block_tile_col_idx < BLOCK_TILE_SIZE_X)
+        // {
+        //     B_thread_block_tile[B_thread_block_tile_row_idx]
+        //                        [B_thread_block_tile_col_idx] = val;
+        // }
+        B_thread_block_tile[B_thread_block_tile_row_idx]
+                           [B_thread_block_tile_col_idx] = val;
+    }
+}
+
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t BLOCK_TILE_SIZE_K, size_t NUM_THREADS,
+          size_t BLOCK_TILE_SKEW_SIZE_X = 0U,
+          size_t BLOCK_TILE_SKEW_SIZE_K = 0U, typename VECTOR_TYPE = int4>
+__device__ void load_data_from_global_memory_to_shared_memory_vectorized(
+    T const* A, size_t lda, T const* B, size_t ldb,
+    T A_thread_block_tile[BLOCK_TILE_SIZE_Y]
+                         [BLOCK_TILE_SIZE_K + BLOCK_TILE_SKEW_SIZE_K],
+    T B_thread_block_tile[BLOCK_TILE_SIZE_K]
+                         [BLOCK_TILE_SIZE_X + BLOCK_TILE_SKEW_SIZE_X],
+    size_t thread_block_tile_idx, size_t thread_linear_idx, size_t m, size_t n,
+    size_t k)
+{
+    constexpr size_t NUM_VECTOR_UNITS{sizeof(VECTOR_TYPE) / sizeof(T)};
+    static_assert(sizeof(VECTOR_TYPE) % sizeof(T) == 0U);
+    static_assert(BLOCK_TILE_SIZE_K % NUM_VECTOR_UNITS == 0U);
+    static_assert(BLOCK_TILE_SIZE_X % NUM_VECTOR_UNITS == 0U);
+    constexpr size_t VECTORIZED_BLOCK_TILE_SIZE_K{BLOCK_TILE_SIZE_K /
+                                                  NUM_VECTOR_UNITS};
+    static_assert(BLOCK_TILE_SIZE_K % NUM_VECTOR_UNITS == 0U);
+    constexpr size_t VECTORIZED_BLOCK_TILE_SIZE_X{BLOCK_TILE_SIZE_X /
+                                                  NUM_VECTOR_UNITS};
+    static_assert(BLOCK_TILE_SIZE_X % NUM_VECTOR_UNITS == 0U);
+
+    // The skew size could affect the data alignment in shared memory when we
+    // use vectorized load. We need to make sure the data alignment is correct.
+    static_assert((BLOCK_TILE_SIZE_K) * sizeof(T) % sizeof(VECTOR_TYPE) == 0U);
+    static_assert((BLOCK_TILE_SIZE_X) * sizeof(T) % sizeof(VECTOR_TYPE) == 0U);
+    static_assert((BLOCK_TILE_SIZE_K + BLOCK_TILE_SKEW_SIZE_K) * sizeof(T) %
+                      sizeof(VECTOR_TYPE) ==
+                  0U);
+    static_assert((BLOCK_TILE_SIZE_X + BLOCK_TILE_SKEW_SIZE_X) * sizeof(T) %
+                      sizeof(VECTOR_TYPE) ==
+                  0U);
+
+// Load data from A on DRAM to A_thread_block_tile on shared memory.
+#pragma unroll
+    for (size_t load_idx{0U};
+         load_idx <
+         (BLOCK_TILE_SIZE_Y * VECTORIZED_BLOCK_TILE_SIZE_K + NUM_THREADS - 1U) /
+             NUM_THREADS;
+         ++load_idx)
+    {
+        size_t const A_thread_block_tile_row_idx{
+            (thread_linear_idx + load_idx * NUM_THREADS) /
+            VECTORIZED_BLOCK_TILE_SIZE_K};
+        size_t const A_thread_block_tile_col_idx{
+            (thread_linear_idx + load_idx * NUM_THREADS) %
+            VECTORIZED_BLOCK_TILE_SIZE_K * NUM_VECTOR_UNITS};
+        size_t const A_row_idx{blockIdx.y * BLOCK_TILE_SIZE_Y +
+                               A_thread_block_tile_row_idx};
+        size_t const A_col_idx{thread_block_tile_idx * BLOCK_TILE_SIZE_K +
+                               A_thread_block_tile_col_idx};
+
+        // These boundary checks might slow down the kernel to some extent.
+        // But they guarantee the correctness of the kernel for all
+        // different GEMM configurations.
+        VECTOR_TYPE A_row_vector_vals{0, 0, 0, 0};
+        if (A_row_idx < m && A_col_idx < k)
+        {
+            A_row_vector_vals = *reinterpret_cast<VECTOR_TYPE const*>(
+                &A[A_row_idx * lda + A_col_idx]);
+        }
+        if (A_col_idx + NUM_VECTOR_UNITS > k)
+        {
+            // Number of invalid elements in the last vector.
+            size_t const num_invalid_elements{A_col_idx + NUM_VECTOR_UNITS - k};
+            // Mask out the invalid elements.
+            T* const A_row_vector_vals_ptr{
+                reinterpret_cast<T*>(&A_row_vector_vals)};
+            for (size_t i{0U}; i < num_invalid_elements; ++i)
+            {
+                A_row_vector_vals_ptr[NUM_VECTOR_UNITS - 1U - i] =
+                    static_cast<T>(0);
+            }
+        }
+        // If this is true, the following if can be removed.
+        // static_assert(VECTORIZED_BLOCK_TILE_SIZE_K * BLOCK_TILE_SIZE_Y %
+        // NUM_THREADS == 0U);
+        if (A_thread_block_tile_row_idx < BLOCK_TILE_SIZE_Y &&
+            A_thread_block_tile_col_idx < BLOCK_TILE_SIZE_K)
+        {
+            *reinterpret_cast<int4*>(
+                &A_thread_block_tile[A_thread_block_tile_row_idx]
+                                    [A_thread_block_tile_col_idx]) =
+                A_row_vector_vals;
+        }
+    }
+// Load data from B on DRAM to B_thread_block_tile on shared memory.
+#pragma unroll
+    for (size_t load_idx{0U};
+         load_idx <
+         (BLOCK_TILE_SIZE_K * VECTORIZED_BLOCK_TILE_SIZE_X + NUM_THREADS - 1U) /
+             NUM_THREADS;
+         ++load_idx)
+    {
+        size_t const B_thread_block_tile_row_idx{
+            (thread_linear_idx + load_idx * NUM_THREADS) /
+            VECTORIZED_BLOCK_TILE_SIZE_X};
+        size_t const B_thread_block_tile_col_idx{
+            (thread_linear_idx + load_idx * NUM_THREADS) %
+            VECTORIZED_BLOCK_TILE_SIZE_X * NUM_VECTOR_UNITS};
+        size_t const B_row_idx{thread_block_tile_idx * BLOCK_TILE_SIZE_K +
+                               B_thread_block_tile_row_idx};
+        size_t const B_col_idx{blockIdx.x * BLOCK_TILE_SIZE_X +
+                               B_thread_block_tile_col_idx};
+
+        // These boundary checks might slow down the kernel to some extent.
+        // But they guarantee the correctness of the kernel for all
+        // different GEMM configurations.
+        VECTOR_TYPE B_row_vector_vals{0, 0, 0, 0};
+        if (B_row_idx < k && B_col_idx < n)
+        {
+            B_row_vector_vals = *reinterpret_cast<VECTOR_TYPE const*>(
+                &B[B_row_idx * ldb + B_col_idx]);
+        }
+        if (B_col_idx + NUM_VECTOR_UNITS > n)
+        {
+            // Number of invalid elements in the last vector.
+            size_t const num_invalid_elements{B_col_idx + NUM_VECTOR_UNITS - n};
+            // Mask out the invalid elements.
+            T* const B_row_vector_vals_ptr{
+                reinterpret_cast<T*>(&B_row_vector_vals)};
+            for (size_t i{0U}; i < num_invalid_elements; ++i)
+            {
+                B_row_vector_vals_ptr[NUM_VECTOR_UNITS - 1U - i] =
+                    static_cast<T>(0);
+            }
+        }
+        // If this is true, the following if can be removed.
+        // static_assert(VECTORIZED_BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_K %
+        // NUM_THREADS ==
+        //               0U);
+        if (B_thread_block_tile_row_idx < BLOCK_TILE_SIZE_K &&
+            B_thread_block_tile_col_idx < BLOCK_TILE_SIZE_X)
+        {
+            *reinterpret_cast<int4*>(
+                &B_thread_block_tile[B_thread_block_tile_row_idx]
+                                    [B_thread_block_tile_col_idx]) =
+                B_row_vector_vals;
+        }
+    }
+}
+
+template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
+          size_t BLOCK_TILE_SIZE_K, size_t NUM_THREADS,
+          size_t BLOCK_TILE_SKEW_SIZE_X = 0U,
+          size_t BLOCK_TILE_SKEW_SIZE_Y = 0U, typename VECTOR_TYPE = int4>
+__device__ void
+load_data_from_global_memory_to_shared_memory_transposed_vectorized(
+    T const* A, size_t lda, T const* B, size_t ldb,
+    T A_thread_block_tile_transposed[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_Y +
+                                                        BLOCK_TILE_SKEW_SIZE_Y],
+    T B_thread_block_tile[BLOCK_TILE_SIZE_K]
+                         [BLOCK_TILE_SIZE_X + BLOCK_TILE_SKEW_SIZE_X],
+    size_t thread_block_tile_idx, size_t thread_linear_idx, size_t m, size_t n,
+    size_t k)
+{
+    constexpr size_t NUM_VECTOR_UNITS{sizeof(VECTOR_TYPE) / sizeof(T)};
+    static_assert(sizeof(VECTOR_TYPE) % sizeof(T) == 0U);
+    static_assert(BLOCK_TILE_SIZE_K % NUM_VECTOR_UNITS == 0U);
+    static_assert(BLOCK_TILE_SIZE_X % NUM_VECTOR_UNITS == 0U);
+    constexpr size_t VECTORIZED_BLOCK_TILE_SIZE_K{BLOCK_TILE_SIZE_K /
+                                                  NUM_VECTOR_UNITS};
+    static_assert(BLOCK_TILE_SIZE_K % NUM_VECTOR_UNITS == 0U);
+    constexpr size_t VECTORIZED_BLOCK_TILE_SIZE_X{BLOCK_TILE_SIZE_X /
+                                                  NUM_VECTOR_UNITS};
+    static_assert(BLOCK_TILE_SIZE_X % NUM_VECTOR_UNITS == 0U);
+
+    // The skew size could affect the data alignment in shared memory when we
+    // use vectorized load. We need to make sure the data alignment is correct.
+    static_assert((BLOCK_TILE_SIZE_Y) * sizeof(T) % sizeof(VECTOR_TYPE) == 0U);
+    static_assert((BLOCK_TILE_SIZE_X) * sizeof(T) % sizeof(VECTOR_TYPE) == 0U);
+    static_assert((BLOCK_TILE_SIZE_Y + BLOCK_TILE_SKEW_SIZE_Y) * sizeof(T) %
+                      sizeof(VECTOR_TYPE) ==
+                  0U);
+    static_assert((BLOCK_TILE_SIZE_X + BLOCK_TILE_SKEW_SIZE_X) * sizeof(T) %
+                      sizeof(VECTOR_TYPE) ==
+                  0U);
+
+// Load data from A on DRAM to A_thread_block_tile on shared memory.
+#pragma unroll
+    for (size_t load_idx{0U};
+         load_idx <
+         (BLOCK_TILE_SIZE_Y * VECTORIZED_BLOCK_TILE_SIZE_K + NUM_THREADS - 1U) /
+             NUM_THREADS;
+         ++load_idx)
+    {
+        size_t const A_thread_block_tile_row_idx{
+            (thread_linear_idx + load_idx * NUM_THREADS) /
+            VECTORIZED_BLOCK_TILE_SIZE_K};
+        size_t const A_thread_block_tile_col_idx{
+            (thread_linear_idx + load_idx * NUM_THREADS) %
+            VECTORIZED_BLOCK_TILE_SIZE_K * NUM_VECTOR_UNITS};
+        size_t const A_row_idx{blockIdx.y * BLOCK_TILE_SIZE_Y +
+                               A_thread_block_tile_row_idx};
+        size_t const A_col_idx{thread_block_tile_idx * BLOCK_TILE_SIZE_K +
+                               A_thread_block_tile_col_idx};
+
+        // These boundary checks might slow down the kernel to some extent.
+        // But they guarantee the correctness of the kernel for all
+        // different GEMM configurations.
+        int4 A_row_vector_vals{0, 0, 0, 0};
+        if (A_row_idx < m && A_col_idx < k)
+        {
+            A_row_vector_vals =
+                *reinterpret_cast<int4 const*>(&A[A_row_idx * lda + A_col_idx]);
+        }
+        if (A_col_idx + NUM_VECTOR_UNITS > k)
+        {
+            // Number of invalid elements in the last vector.
+            size_t const num_invalid_elements{A_col_idx + NUM_VECTOR_UNITS - k};
+            // Mask out the invalid elements.
+            T* const A_row_vector_vals_ptr{
+                reinterpret_cast<T*>(&A_row_vector_vals)};
+            for (size_t i{0U}; i < num_invalid_elements; ++i)
+            {
+                A_row_vector_vals_ptr[NUM_VECTOR_UNITS - 1U - i] =
+                    static_cast<T>(0);
+            }
+        }
+        // If this is true, the following if can be removed.
+        // static_assert(VECTORIZED_BLOCK_TILE_SIZE_K * BLOCK_TILE_SIZE_Y %
+        // NUM_THREADS ==
+        //               0U);
+        if (A_thread_block_tile_row_idx < BLOCK_TILE_SIZE_Y &&
+            A_thread_block_tile_col_idx < BLOCK_TILE_SIZE_K)
+        {
+            for (size_t i{0U}; i < NUM_VECTOR_UNITS; ++i)
+            {
+                A_thread_block_tile_transposed[A_thread_block_tile_col_idx +
+                                               i][A_thread_block_tile_row_idx] =
+                    reinterpret_cast<T const*>(&A_row_vector_vals)[i];
+            }
+        }
+    }
+// Load data from B on DRAM to B_thread_block_tile on shared memory.
+#pragma unroll
+    for (size_t load_idx{0U};
+         load_idx <
+         (BLOCK_TILE_SIZE_K * VECTORIZED_BLOCK_TILE_SIZE_X + NUM_THREADS - 1U) /
+             NUM_THREADS;
+         ++load_idx)
+    {
+        size_t const B_thread_block_tile_row_idx{
+            (thread_linear_idx + load_idx * NUM_THREADS) /
+            VECTORIZED_BLOCK_TILE_SIZE_X};
+        size_t const B_thread_block_tile_col_idx{
+            (thread_linear_idx + load_idx * NUM_THREADS) %
+            VECTORIZED_BLOCK_TILE_SIZE_X * NUM_VECTOR_UNITS};
+        size_t const B_row_idx{thread_block_tile_idx * BLOCK_TILE_SIZE_K +
+                               B_thread_block_tile_row_idx};
+        size_t const B_col_idx{blockIdx.x * BLOCK_TILE_SIZE_X +
+                               B_thread_block_tile_col_idx};
+
+        // These boundary checks might slow down the kernel to some extent.
+        // But they guarantee the correctness of the kernel for all
+        // different GEMM configurations.
+        int4 B_row_vector_vals{0, 0, 0, 0};
+        if (B_row_idx < k && B_col_idx < n)
+        {
+            B_row_vector_vals =
+                *reinterpret_cast<int4 const*>(&B[B_row_idx * ldb + B_col_idx]);
+        }
+        if (B_col_idx + NUM_VECTOR_UNITS > n)
+        {
+            // Number of invalid elements in the last vector.
+            size_t const num_invalid_elements{B_col_idx + NUM_VECTOR_UNITS - n};
+            // Mask out the invalid elements.
+            T* const B_row_vector_vals_ptr{
+                reinterpret_cast<T*>(&B_row_vector_vals)};
+            for (size_t i{0U}; i < num_invalid_elements; ++i)
+            {
+                B_row_vector_vals_ptr[NUM_VECTOR_UNITS - 1U - i] =
+                    static_cast<T>(0);
+            }
+        }
+        // If this is true, the following if can be removed.
+        // static_assert(VECTORIZED_BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_K %
+        // NUM_THREADS ==
+        //               0U);
+        if (B_thread_block_tile_row_idx < BLOCK_TILE_SIZE_K &&
+            B_thread_block_tile_col_idx < BLOCK_TILE_SIZE_X)
+        {
+            *reinterpret_cast<int4*>(
+                &B_thread_block_tile[B_thread_block_tile_row_idx]
+                                    [B_thread_block_tile_col_idx]) =
+                B_row_vector_vals;
+        }
+    }
+}
+
+#endif // CUDA_GEMM_UTILS_CUH

upstream_ref/cuda_gemm_optimization/cuda_gemm_utils.hpp

@@ -0,0 +1,13 @@
+#ifndef CUDA_GEMM_UTILS_HPP
+#define CUDA_GEMM_UTILS_HPP
+
+#include <cuda_runtime.h>
+
+#define CHECK_CUDA_ERROR(val) check_cuda((val), #val, __FILE__, __LINE__)
+void check_cuda(cudaError_t err, const char* const func, const char* const file,
+                const int line);
+
+#define CHECK_LAST_CUDA_ERROR() check_cuda_last(__FILE__, __LINE__)
+void check_cuda_last(const char* const file, const int line);
+
+#endif // CUDA_GEMM_UTILS_HPP

upstream_ref/cuda_gemm_optimization/profile_cuda_gemm_fp16.cu

@@ -0,0 +1,109 @@
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#include "cuda_gemm.hpp"
+#include "profile_utils.cuh"
+
+int main()
+{
+    print_device_info();
+
+    constexpr size_t num_repeats{1U};
+    constexpr size_t num_warmups{1U};
+
+    __half const fp16_abs_tol{__float2half(5.0e-2f)};
+    double const fp16_rel_tol{1.0e-1f};
+
+    __half const fp16_tensor_core_abs_tol{__float2half(5.0e-2f)};
+    double const fp16_tensor_core_rel_tol{1.0e-2f};
+
+    constexpr size_t m{4096U};
+    constexpr size_t k{4096U};
+    constexpr size_t n{4096U};
+
+    constexpr size_t lda{(k + 16U - 1U) / 16U * 16U};
+    constexpr size_t ldb{(n + 16U - 1U) / 16U * 16U};
+    constexpr size_t ldc{(n + 16U - 1U) / 16U * 16U};
+
+    static_assert(lda >= k);
+    static_assert(ldb >= n);
+    static_assert(ldc >= n);
+
+    std::cout << "Matrix Size: " << "M = " << m << " N = " << n << " K = " << k
+              << std::endl;
+    std::cout << "Matrix A: " << m << " x " << k
+              << " Leading Dimension Size = " << lda << std::endl;
+    std::cout << "Matrix B: " << k << " x " << n
+              << " Leading Dimension Size = " << ldb << std::endl;
+    std::cout << "Matrix C: " << m << " x " << n
+              << " Leading Dimension Size = " << ldc << std::endl;
+    std::cout << std::endl;
+
+    // Define all the GEMM kernel launch functions to be profiled.
+    std::vector<std::pair<
+        std::string,
+        std::function<void(size_t, size_t, size_t, __half const*, __half const*,
+                           size_t, __half const*, size_t, __half const*,
+                           __half*, size_t, cudaStream_t)>>> const
+        gemm_fp16_kernel_launch_functions{
+            {"Custom GEMM Kernel V00", launch_gemm_kernel_v00<__half>},
+            {"Custom GEMM Kernel V01", launch_gemm_kernel_v01<__half>},
+            {"Custom GEMM Kernel V02", launch_gemm_kernel_v02<__half>},
+            {"Custom GEMM Kernel V02 Vectorized",
+             launch_gemm_kernel_v02_vectorized<__half>},
+            {"Custom GEMM Kernel V03", launch_gemm_kernel_v03<__half>},
+            {"Custom GEMM Kernel V03 Vectorized",
+             launch_gemm_kernel_v03_vectorized<__half>},
+            {"Custom GEMM Kernel V04", launch_gemm_kernel_v04<__half>},
+            {"Custom GEMM Kernel V04 Vectorized",
+             launch_gemm_kernel_v04_vectorized<__half>},
+            {"Custom GEMM Kernel V05", launch_gemm_kernel_v05<__half>},
+            {"Custom GEMM Kernel V05 Vectorized",
+             launch_gemm_kernel_v05_vectorized<__half>},
+            {"Custom GEMM Kernel V06", launch_gemm_kernel_v06<__half>},
+            {"Custom GEMM Kernel V06 Vectorized",
+             launch_gemm_kernel_v06_vectorized<__half>},
+            {"Custom GEMM Kernel V06 Vectorized Double Buffered",
+             launch_gemm_kernel_v06_vectorized_double_buffered<__half>},
+        };
+
+    for (auto const& gemm_fp16_kernel_launch_function :
+         gemm_fp16_kernel_launch_functions)
+    {
+        std::cout << gemm_fp16_kernel_launch_function.first << std::endl;
+        std::pair<__half, __half> const gemm_kernel_profile_result{
+            profile_gemm<__half>(
+                m, n, k, lda, ldb, ldc, gemm_fp16_kernel_launch_function.second,
+                fp16_abs_tol, fp16_rel_tol, num_repeats, num_warmups)};
+        std::cout << std::endl;
+    }
+
+    std::vector<std::pair<
+        std::string,
+        std::function<void(size_t, size_t, size_t, __half const*, __half const*,
+                           size_t, __half const*, size_t, __half const*,
+                           __half*, size_t, cudaStream_t)>>> const
+        gemm_fp16_tensor_core_kernel_launch_functions{
+            {"Custom GEMM Kernel V07", launch_gemm_kernel_v07<__half>},
+            {"Custom GEMM Kernel V07 Vectorized",
+             launch_gemm_kernel_v07_vectorized<__half>},
+            {"Custom GEMM Kernel V07 Vectorized Double Buffered",
+             launch_gemm_kernel_v07_vectorized_double_buffered<__half>},
+        };
+
+    for (auto const& gemm_fp16_tensor_core_kernel_launch_function :
+         gemm_fp16_tensor_core_kernel_launch_functions)
+    {
+        std::cout << gemm_fp16_tensor_core_kernel_launch_function.first
+                  << std::endl;
+        std::pair<__half, __half> const gemm_kernel_profile_result{
+            profile_gemm<__half>(
+                m, n, k, lda, ldb, ldc,
+                gemm_fp16_tensor_core_kernel_launch_function.second,
+                fp16_tensor_core_abs_tol, fp16_tensor_core_rel_tol, num_repeats,
+                num_warmups)};
+        std::cout << std::endl;
+    }
+
+    return 0;
+}

upstream_ref/cuda_gemm_optimization/profile_cuda_gemm_fp32.cu

@@ -0,0 +1,78 @@
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#include "cuda_gemm.hpp"
+#include "profile_utils.cuh"
+
+int main()
+{
+    print_device_info();
+
+    constexpr size_t num_repeats{1U};
+    constexpr size_t num_warmups{1U};
+
+    float const fp32_abs_tol{1.0e-3f};
+    double const fp32_rel_tol{0.0e-4f};
+
+    constexpr size_t m{4096U};
+    constexpr size_t k{4096U};
+    constexpr size_t n{4096U};
+
+    constexpr size_t lda{(k + 16U - 1U) / 16U * 16U};
+    constexpr size_t ldb{(n + 16U - 1U) / 16U * 16U};
+    constexpr size_t ldc{(n + 16U - 1U) / 16U * 16U};
+
+    static_assert(lda >= k);
+    static_assert(ldb >= n);
+    static_assert(ldc >= n);
+
+    std::cout << "Matrix Size: " << "M = " << m << " N = " << n << " K = " << k
+              << std::endl;
+    std::cout << "Matrix A: " << m << " x " << k
+              << " Leading Dimension Size = " << lda << std::endl;
+    std::cout << "Matrix B: " << k << " x " << n
+              << " Leading Dimension Size = " << ldb << std::endl;
+    std::cout << "Matrix C: " << m << " x " << n
+              << " Leading Dimension Size = " << ldc << std::endl;
+    std::cout << std::endl;
+
+    // Define all the GEMM kernel launch functions to be profiled.
+    std::vector<std::pair<
+        std::string,
+        std::function<void(size_t, size_t, size_t, float const*, float const*,
+                           size_t, float const*, size_t, float const*, float*,
+                           size_t, cudaStream_t)>>> const
+        gemm_kernel_launch_functions{
+            {"Custom GEMM Kernel V00", launch_gemm_kernel_v00<float>},
+            {"Custom GEMM Kernel V01", launch_gemm_kernel_v01<float>},
+            {"Custom GEMM Kernel V02", launch_gemm_kernel_v02<float>},
+            {"Custom GEMM Kernel V02 Vectorized",
+             launch_gemm_kernel_v02_vectorized<float>},
+            {"Custom GEMM Kernel V03", launch_gemm_kernel_v03<float>},
+            {"Custom GEMM Kernel V03 Vectorized",
+             launch_gemm_kernel_v03_vectorized<float>},
+            {"Custom GEMM Kernel V04", launch_gemm_kernel_v04<float>},
+            {"Custom GEMM Kernel V04 Vectorized",
+             launch_gemm_kernel_v04_vectorized<float>},
+            {"Custom GEMM Kernel V05", launch_gemm_kernel_v05<float>},
+            {"Custom GEMM Kernel V05 Vectorized",
+             launch_gemm_kernel_v05_vectorized<float>},
+            {"Custom GEMM Kernel V06", launch_gemm_kernel_v06<float>},
+            {"Custom GEMM Kernel V06 Vectorized",
+             launch_gemm_kernel_v06_vectorized<float>},
+            {"Custom GEMM Kernel V06 Vectorized Double Buffered",
+             launch_gemm_kernel_v06_vectorized_double_buffered<float>},
+        };
+
+    for (auto const& gemm_kernel_launch_function : gemm_kernel_launch_functions)
+    {
+        std::cout << gemm_kernel_launch_function.first << std::endl;
+        std::pair<float, float> const gemm_kernel_profile_result{
+            profile_gemm<float>(
+                m, n, k, lda, ldb, ldc, gemm_kernel_launch_function.second,
+                fp32_abs_tol, fp32_rel_tol, num_repeats, num_warmups)};
+        std::cout << std::endl;
+    }
+
+    return 0;
+}

upstream_ref/cuda_gemm_optimization/profile_utils.cuh

@@ -0,0 +1,400 @@
+#ifndef PROFILE_UTILS_CUH
+#define PROFILE_UTILS_CUH
+
+#include <cassert>
+#include <cmath>
+#include <functional>
+#include <iostream>
+#include <random>
+
+#include "cuda_gemm.hpp"
+#include "cuda_gemm_utils.cuh"
+
+#include <cublas_v2.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+template <typename T>
+float measure_performance(std::function<T(cudaStream_t)> bound_function,
+                          cudaStream_t stream, size_t num_repeats = 100,
+                          size_t num_warmups = 100)
+{
+    cudaEvent_t start, stop;
+    float time;
+
+    CHECK_CUDA_ERROR(cudaEventCreate(&start));
+    CHECK_CUDA_ERROR(cudaEventCreate(&stop));
+
+    for (size_t i{0}; i < num_warmups; ++i)
+    {
+        bound_function(stream);
+    }
+
+    CHECK_CUDA_ERROR(cudaStreamSynchronize(stream));
+
+    CHECK_CUDA_ERROR(cudaEventRecord(start, stream));
+    for (size_t i{0}; i < num_repeats; ++i)
+    {
+        bound_function(stream);
+    }
+    CHECK_CUDA_ERROR(cudaEventRecord(stop, stream));
+    CHECK_CUDA_ERROR(cudaEventSynchronize(stop));
+    CHECK_LAST_CUDA_ERROR();
+    CHECK_CUDA_ERROR(cudaEventElapsedTime(&time, start, stop));
+    CHECK_CUDA_ERROR(cudaEventDestroy(start));
+    CHECK_CUDA_ERROR(cudaEventDestroy(stop));
+
+    float const latency{time / num_repeats};
+
+    return latency;
+}
+
+#define CHECK_CUBLASS_ERROR(val) check_cublass((val), #val, __FILE__, __LINE__)
+void check_cublass(cublasStatus_t err, const char* const func,
+                   const char* const file, const int line)
+{
+    if (err != CUBLAS_STATUS_SUCCESS)
+    {
+        std::cerr << "cuBLAS Error at: " << file << ":" << line << std::endl;
+        std::cerr << cublasGetStatusString(err) << std::endl;
+        std::exit(EXIT_FAILURE);
+    }
+}
+
+// Determine CUDA data type from type.
+template <typename T,
+          typename std::enable_if<std::is_same<T, float>::value ||
+                                      std::is_same<T, double>::value ||
+                                      std::is_same<T, __half>::value,
+                                  bool>::type = true>
+constexpr cudaDataType_t cuda_data_type_trait()
+{
+    if (std::is_same<T, float>::value)
+    {
+        return CUDA_R_32F;
+    }
+    else if (std::is_same<T, double>::value)
+    {
+        return CUDA_R_64F;
+    }
+    else if (std::is_same<T, __half>::value)
+    {
+        return CUDA_R_16F;
+    }
+    else
+    {
+        throw std::runtime_error("Unsupported data type.");
+    }
+}
+
+template <typename T,
+          typename std::enable_if<std::is_same<T, float>::value ||
+                                      std::is_same<T, double>::value ||
+                                      std::is_same<T, __half>::value,
+                                  bool>::type = true>
+void launch_gemm_cublas(size_t m, size_t n, size_t k, T const* alpha,
+                        T const* A, size_t lda, T const* B, size_t ldb,
+                        T const* beta, T* C, size_t ldc, cublasHandle_t handle)
+{
+    // Non-TensorCore algorithm?
+    constexpr cublasGemmAlgo_t algo{CUBLAS_GEMM_DEFAULT};
+    constexpr cudaDataType_t data_type{cuda_data_type_trait<T>()};
+    // All the matrix are in row-major order.
+    // https://docs.nvidia.com/cuda/cublas/#cublasgemmex
+    // A: m x k row-major -> A: k x m column-major non-transposed
+    // B: k x n row-major -> B: n x k column-major non-transposed
+    // C: m x n row-major -> C: n x m column-major non-transposed
+    // Thus, without padding, the leading dimension of the matrix in row-major
+    // order is the number of columns, i.e., k for A, n for B, and n for C.
+    // Row-major order: C = AB + C
+    // Column-major order: C = BA + C
+    // The cuBLAS API requires the leading dimension of the matrix in
+    // column-major order. This API call looks non-intuitive, but it is correct.
+    CHECK_CUBLASS_ERROR(cublasGemmEx(
+        handle, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, alpha, B, data_type, ldb, A,
+        data_type, lda, beta, C, data_type, ldc, data_type, algo));
+}
+
+template <typename T,
+          typename std::enable_if<std::is_same<T, float>::value ||
+                                      std::is_same<T, double>::value,
+                                  bool>::type = true>
+void launch_gemm_cpu(size_t m, size_t n, size_t k, T const* alpha, T const* A,
+                     size_t lda, T const* B, size_t ldb, T const* beta, T* C,
+                     size_t ldc)
+{
+    // Compute GEMM using CPU.
+    for (size_t i{0U}; i < m; ++i)
+    {
+        for (size_t j{0U}; j < n; ++j)
+        {
+            T sum{static_cast<T>(0)};
+            for (size_t l{0U}; l < k; ++l)
+            {
+                sum += A[i * lda + l] * B[l * ldb + j];
+            }
+            C[i * ldc + j] = (*alpha) * sum + (*beta) * C[i * ldc + j];
+        }
+    }
+}
+
+// Many different implementations have been tried for FP16 GEMM on CPU.
+// There is always a discrepancy between the results from CPU and GPU (cuBLAS or
+// custom kernel).
+template <typename T, typename std::enable_if<std::is_same<T, __half>::value,
+                                              bool>::type = true>
+void launch_gemm_cpu(size_t m, size_t n, size_t k, T const* alpha, T const* A,
+                     size_t lda, T const* B, size_t ldb, T const* beta, T* C,
+                     size_t ldc)
+{
+    // Compute GEMM using CPU.
+    for (size_t i{0U}; i < m; ++i)
+    {
+        for (size_t j{0U}; j < n; ++j)
+        {
+            float sum{0.0f};
+            for (size_t l{0U}; l < k; ++l)
+            {
+                sum += __half2float(__hmul(A[i * lda + l], B[l * ldb + j]));
+            }
+            C[i * ldc + j] = __float2half(__half2float(*alpha) * sum +
+                                          __half2float(*beta) *
+                                              __half2float(C[i * ldc + j]));
+        }
+    }
+}
+
+template <typename T>
+bool all_close(T const* C, T const* C_ref, size_t m, size_t n, size_t ldc,
+               T abs_tol, double rel_tol)
+{
+    bool status{true};
+    for (size_t i{0U}; i < m; ++i)
+    {
+        for (size_t j{0U}; j < n; ++j)
+        {
+            double const C_val{static_cast<double>(C[i * ldc + j])};
+            double const C_ref_val{static_cast<double>(C_ref[i * ldc + j])};
+            double const diff{C_val - C_ref_val};
+            double const diff_val{std::abs(diff)};
+            if (diff_val >
+                std::max(static_cast<double>(abs_tol),
+                         static_cast<double>(std::abs(C_ref_val)) * rel_tol))
+            {
+                std::cout << "C[" << i << ", " << j << "] = " << C_val
+                          << " C_ref[" << i << ", " << j << "] = " << C_ref_val
+                          << " Abs Diff: " << diff_val
+                          << " Abs Diff Threshold: "
+                          << static_cast<double>(abs_tol)
+                          << " Rel->Abs Diff Threshold: "
+                          << static_cast<double>(
+                                 static_cast<double>(std::abs(C_ref_val)) *
+                                 rel_tol)
+                          << std::endl;
+                status = false;
+                return status;
+            }
+        }
+    }
+    return status;
+}
+
+void print_device_info()
+{
+    int device_id{0};
+    cudaGetDevice(&device_id);
+    cudaDeviceProp device_prop;
+    cudaGetDeviceProperties(&device_prop, device_id);
+    std::cout << "Device Name: " << device_prop.name << std::endl;
+    float const memory_size{static_cast<float>(device_prop.totalGlobalMem) /
+                            (1 << 30)};
+    std::cout << "Memory Size: " << memory_size << " GB" << std::endl;
+    float const peak_bandwidth{
+        static_cast<float>(2.0f * device_prop.memoryClockRate *
+                           (device_prop.memoryBusWidth / 8) / 1.0e6)};
+    std::cout << "Peak Bandwitdh: " << peak_bandwidth << " GB/s" << std::endl;
+    std::cout << std::endl;
+}
+
+template <typename T>
+float compute_effective_bandwidth(size_t m, size_t n, size_t k, float latency)
+{
+    return ((m * k + k * n + m * n) * sizeof(T)) / (latency * 1e-3) / 1e9;
+}
+
+float compute_effective_tflops(size_t m, size_t n, size_t k, float latency)
+{
+    return (2.0 * m * k * n) / (latency * 1e-3) / 1e12;
+}
+
+template <typename T,
+          typename std::enable_if<std::is_same<T, float>::value ||
+                                      std::is_same<T, double>::value ||
+                                      std::is_same<T, __half>::value,
+                                  bool>::type = true>
+void random_initialize_matrix(T* A, size_t m, size_t n, size_t lda,
+                              unsigned int seed = 0U)
+{
+    std::default_random_engine eng(seed);
+    // The best way to verify is to use integer values.
+    std::uniform_int_distribution<int> dis(0, 5);
+    // std::uniform_real_distribution<float> dis(-1.0f, 1.0f);
+    auto const rand = [&dis, &eng]() { return dis(eng); };
+    for (size_t i{0U}; i < m; ++i)
+    {
+        for (size_t j{0U}; j < n; ++j)
+        {
+            A[i * lda + j] = static_cast<T>(rand());
+        }
+    }
+}
+
+void print_performance_result(size_t m, size_t n, size_t k, float latency)
+{
+    float const effective_bandwidth{
+        compute_effective_bandwidth<float>(m, n, k, latency)};
+    float const effective_tflops{compute_effective_tflops(m, n, k, latency)};
+
+    std::cout << "Latency: " << latency << " ms" << std::endl;
+    std::cout << "Effective Bandwidth: " << effective_bandwidth << " GB/s"
+              << std::endl;
+    std::cout << "Effective TFLOPS: " << effective_tflops << " TFLOPS"
+              << std::endl;
+}
+
+template <typename T,
+          typename std::enable_if<std::is_same<T, float>::value ||
+                                      std::is_same<T, double>::value ||
+                                      std::is_same<T, __half>::value,
+                                  bool>::type = true>
+std::pair<float, float> profile_gemm(
+    size_t m, size_t n, size_t k, size_t lda, size_t ldb, size_t ldc,
+    std::function<void(size_t, size_t, size_t, T const*, T const*, size_t,
+                       T const*, size_t, T const*, T*, size_t, cudaStream_t)>
+        gemm_kernel_launch_function,
+    T abs_tol, double rel_tol, size_t num_repeats = 10, size_t num_warmups = 10,
+    unsigned int seed = 0U)
+{
+    T const alpha{static_cast<T>(1.0)};
+    T const beta{static_cast<T>(0.0)};
+
+    // Create CUDA stream.
+    cudaStream_t stream;
+    CHECK_CUDA_ERROR(cudaStreamCreate(&stream));
+
+    // Allocate memory on host.
+    T* A_host{nullptr};
+    T* B_host{nullptr};
+    T* C_host{nullptr};
+    T* C_host_ref{nullptr};
+    T* C_host_from_device{nullptr};
+    CHECK_CUDA_ERROR(cudaMallocHost(&A_host, m * lda * sizeof(T)));
+    CHECK_CUDA_ERROR(cudaMallocHost(&B_host, k * ldb * sizeof(T)));
+    CHECK_CUDA_ERROR(cudaMallocHost(&C_host, m * ldc * sizeof(T)));
+    CHECK_CUDA_ERROR(cudaMallocHost(&C_host_ref, m * ldc * sizeof(T)));
+    CHECK_CUDA_ERROR(cudaMallocHost(&C_host_from_device, m * ldc * sizeof(T)));
+
+    // Initialize matrix A and B.
+    random_initialize_matrix(A_host, m, k, lda);
+    random_initialize_matrix(B_host, k, n, ldb);
+    random_initialize_matrix(C_host, m, n, ldc);
+
+    // Allocate memory on device.
+    T* A_device{nullptr};
+    T* B_device{nullptr};
+    T* C_device{nullptr};
+    CHECK_CUDA_ERROR(cudaMalloc(&A_device, m * lda * sizeof(T)));
+    CHECK_CUDA_ERROR(cudaMalloc(&B_device, k * ldb * sizeof(T)));
+    CHECK_CUDA_ERROR(cudaMalloc(&C_device, m * ldc * sizeof(T)));
+
+    // Copy matrix A and B from host to device.
+    CHECK_CUDA_ERROR(cudaMemcpy(A_device, A_host, m * lda * sizeof(T),
+                                cudaMemcpyHostToDevice));
+    CHECK_CUDA_ERROR(cudaMemcpy(B_device, B_host, k * ldb * sizeof(T),
+                                cudaMemcpyHostToDevice));
+    CHECK_CUDA_ERROR(cudaMemcpy(C_device, C_host, m * ldc * sizeof(T),
+                                cudaMemcpyHostToDevice));
+    CHECK_CUDA_ERROR(cudaMemcpy(C_host_ref, C_host, m * ldc * sizeof(T),
+                                cudaMemcpyHostToHost));
+
+    // Create cuBLAS handle.
+    cublasHandle_t handle;
+    CHECK_CUBLASS_ERROR(cublasCreate(&handle));
+    CHECK_CUBLASS_ERROR(cublasSetStream(handle, stream));
+
+    // Compute reference output using cuBLAS.
+    launch_gemm_cublas<T>(m, n, k, &alpha, A_device, lda, B_device, ldb, &beta,
+                          C_device, ldc, handle);
+    CHECK_CUDA_ERROR(cudaStreamSynchronize(stream));
+
+    // Copy matrix C from device to host.
+    CHECK_CUDA_ERROR(cudaMemcpy(C_host_ref, C_device, m * ldc * sizeof(T),
+                                cudaMemcpyDeviceToHost));
+
+    // // Compute reference output using CPU.
+    // std::cout << "Computing reference output using CPU..." << std::endl;
+    // launch_gemm_cpu<T>(m, n, k, &alpha, A_host, lda, B_host, ldb, &beta,
+    //                    C_host_ref, ldc);
+    // std::cout << "Done." << std::endl;
+
+    // Launch CUDA GEMM.
+    CHECK_CUDA_ERROR(cudaMemcpy(C_device, C_host, m * ldc * sizeof(T),
+                                cudaMemcpyHostToDevice));
+    // Verify the correctness of CUDA GEMM.
+    gemm_kernel_launch_function(m, n, k, &alpha, A_device, lda, B_device, ldb,
+                                &beta, C_device, ldc, stream);
+
+    // launch_gemm_cublas<T>(m, n, k, &alpha, A_device, lda, B_device, ldb,
+    // &beta,
+    //                       C_device, ldc, handle);
+
+    CHECK_CUDA_ERROR(cudaStreamSynchronize(stream));
+    CHECK_CUDA_ERROR(cudaMemcpy(C_host_from_device, C_device,
+                                m * ldc * sizeof(T), cudaMemcpyDeviceToHost));
+    assert(all_close<T>(C_host_from_device, C_host_ref, m, n, ldc, abs_tol,
+                        rel_tol));
+
+    // Launch cuBLAS GEMM.
+    float const latency_cublas{measure_performance<void>(
+        [&](cudaStream_t stream)
+        {
+            launch_gemm_cublas<T>(m, n, k, &alpha, A_device, lda, B_device, ldb,
+                                  &beta, C_device, ldc, handle);
+            return;
+        },
+        stream, num_repeats, num_warmups)};
+
+    float const latency_cuda_gemm{measure_performance<void>(
+        [&](cudaStream_t stream)
+        {
+            gemm_kernel_launch_function(m, n, k, &alpha, A_device, lda,
+                                        B_device, ldb, &beta, C_device, ldc,
+                                        stream);
+            return;
+        },
+        stream, num_repeats, num_warmups)};
+
+    // Release resources.
+    CHECK_CUDA_ERROR(cudaFree(A_device));
+    CHECK_CUDA_ERROR(cudaFree(B_device));
+    CHECK_CUDA_ERROR(cudaFree(C_device));
+    CHECK_CUDA_ERROR(cudaFreeHost(A_host));
+    CHECK_CUDA_ERROR(cudaFreeHost(B_host));
+    CHECK_CUDA_ERROR(cudaFreeHost(C_host));
+    CHECK_CUDA_ERROR(cudaFreeHost(C_host_ref));
+    CHECK_CUDA_ERROR(cudaFreeHost(C_host_from_device));
+    CHECK_CUBLASS_ERROR(cublasDestroy(handle));
+    CHECK_CUDA_ERROR(cudaStreamDestroy(stream));
+
+    std::cout << "cuBLAS GEMM Kernel Performance" << std::endl;
+    print_performance_result(m, n, k, latency_cublas);
+    std::cout << "Custom GEMM Kernel Performance" << std::endl;
+    print_performance_result(m, n, k, latency_cuda_gemm);
+    std::cout << "Custom GEMM VS cuBLAS GEMM Performance: "
+              << latency_cublas / latency_cuda_gemm * 100.0f << "%"
+              << std::endl;
+
+    return std::pair<float, float>{latency_cublas, latency_cuda_gemm};
+}
+
+#endif // PROFILE_UTILS_CUH

upstream_ref/nvidia_sgemm_practice/kernel_1.cuh

@@ -0,0 +1,19 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+__global__ __launch_bounds__(1024) void
+mysgemm_v1(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) {
+
+    int gx = blockIdx.x * blockDim.x + threadIdx.x; // 全局x
+    int gy = blockIdx.y * blockDim.y + threadIdx.y; // 全局y
+
+    float tmp = 0.;
+    for (int i = 0; i < K; i++) {
+        tmp += A[gy * K + i] * B[i * N + gx]; // 两次全局内存访问和一次FMA（累加乘）
+    }
+    C[gy * N + gx] = alpha * tmp + beta * C[gy * N + gx];
+}

upstream_ref/nvidia_sgemm_practice/kernel_2.cuh

@@ -0,0 +1,45 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+template<const int BLOCK_SIZE>
+__global__ void mysgemm_v2(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) {
+    int bx = blockIdx.x;
+    int by = blockIdx.y;
+
+    const int BM = BLOCK_SIZE;
+    const int BN = BLOCK_SIZE;
+    const int BK = BLOCK_SIZE;
+
+    int tx = threadIdx.x % BN;
+    int ty = threadIdx.x / BN;
+
+    // 申请共享内存空间
+    __shared__ float As[BM * BK];
+    __shared__ float Bs[BK * BN];
+
+    // 移动到当前block
+    A = &A[by * BM * K];
+    B = &B[bx * BN];
+    C = &C[by * BM * N + bx * BN];
+
+    float tmp = 0.;
+    for (int k = 0; k < K; k += BK) {
+        // 缓存A_tile和B_tile
+        As[ty * BK + tx] = A[ty * K + tx];
+        Bs[ty * BN + tx] = B[ty * N + tx];
+        // 同步所有线程缓存完成
+        __syncthreads();
+        A += BK;
+        B += BK * N;
+        for (int i = 0; i < BK; i++) {
+            tmp += As[ty * BK + i] * Bs[i * BN + tx];
+        }
+        // FMA计算需要读取缓存数据，在新一轮写入缓存前进行同步，确保所有线程计算完成
+        __syncthreads();
+    }
+    C[ty * N + tx] = alpha * tmp + beta * C[ty * N + tx];
+}

upstream_ref/nvidia_sgemm_practice/kernel_3.cuh

@@ -0,0 +1,71 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+template<const int BM,
+        const int BN,
+        const int BK,
+        const int TM>
+__global__ void mysgemm_v3(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) {
+    int bx = blockIdx.x;
+    int by = blockIdx.y;
+    int thread_num = BM * BN / TM; // 一个线程负责block中计算TM个元素
+
+    int tx = threadIdx.x % BN;
+    int ty = threadIdx.x / BN * TM;
+
+    __shared__ float As[BM * BK];
+    __shared__ float Bs[BK * BN];
+
+    // 移动到当前block
+    A = &A[by * BM * K];
+    B = &B[bx * BN];
+    C = &C[by * BM * N + bx * BN];
+
+    /*
+    当前线程负责搬运全局内存中第a_tile_row行，第a_tile_col列元素至共享内存第a_tile_row行，第a_tile_col列
+    a_tile_stride表示block中线程可搬运a_tile_stride行至共享内存；
+
+    若BM=64,BK=8,thread_num=512,则a_tile_stride=64,a_tile_stride=BM，表示每个线程搬运一轮即可完成所需元素的搬运;
+    若BM=128,BK=8,thread_num=512,则a_tile_stride=64,表示每个线程搬运两轮即可完成所需元素的搬运;
+    */
+    int a_tile_row = threadIdx.x / BK;
+    int a_tile_col = threadIdx.x % BK;
+    int a_tile_stride = thread_num / BK;
+
+    int b_tile_row = threadIdx.x / BN;
+    int b_tile_col = threadIdx.x % BN;
+    int b_tile_stride = thread_num / BN;
+
+    float tmp[TM + 1] = {0.}; // 每个线程负责TM个元素，则需要申请TM个寄存器保存累加值，额外的一个寄存器用于缓存；
+    #pragma unroll
+    for (int k = 0; k < K; k += BK) {
+        #pragma unroll
+        for (int i = 0; i < BM; i += a_tile_stride) {
+            As[(a_tile_row + i) * BK + a_tile_col] = A[(a_tile_row + i) * K + a_tile_col];
+        }
+        #pragma unroll
+        for (int i = 0; i < BK; i += b_tile_stride) {
+            Bs[(b_tile_row + i) * BN + b_tile_col] = B[(b_tile_row + i) * N + b_tile_col];
+        }
+        __syncthreads();
+        A += BK;
+        B += BK * N;
+        #pragma unroll
+        for (int i = 0; i < BK; i++) {
+            tmp[TM] = Bs[tx + i * BN]; // 额外的一个寄存器，避免反复从共享内存中读取Bs[tx + i * BN]
+            #pragma unroll  // 循环展开，增加指令并行度
+            for (int j = 0; j < TM; j++) {
+                tmp[j] += As[(ty + j) * BK + i] * tmp[TM];
+            }
+        }
+        __syncthreads();
+    }
+    #pragma unroll
+    for (int j = 0; j < TM; j++) {
+        C[(ty + j) * N + tx] = alpha * tmp[j] + beta * C[(ty + j) * N + tx];
+    }
+}

upstream_ref/nvidia_sgemm_practice/kernel_4.cuh

@@ -0,0 +1,76 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+template<const int BM,
+        const int BN,
+        const int BK,
+        const int TM,
+        const int TN>
+__global__ void mysgemm_v4(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) {
+    int bx = blockIdx.x;
+    int by = blockIdx.y;
+
+    int block_row_thread = BN / TN;
+    int block_col_thread = BM / TM;
+    int thread_num = block_row_thread * block_col_thread; // 一个线程负责计算block中TM*TN个元素
+
+    int tx = (threadIdx.x % block_row_thread) * TN;
+    int ty = (threadIdx.x / block_row_thread) * TM;
+
+    __shared__ float As[BM * BK];
+    __shared__ float Bs[BK * BN];
+
+    // 移动到当前block
+    A = &A[by * BM * K];
+    B = &B[bx * BN];
+    C = &C[by * BM * N + bx * BN];
+
+    /*
+    当前线程负责搬运全局内存中第a_tile_row行，第a_tile_col列元素至共享内存第a_tile_row行，第a_tile_col列
+    a_tile_stride表示block中线程可搬运a_tile_stride行至共享内存；
+
+    若BM=64,BK=8,thread_num=512,则a_tile_stride=64,a_tile_stride=BM，表示每个线程搬运一轮即可完成所需元素的搬运;
+    若BM=128,BK=8,thread_num=512,则a_tile_stride=64,表示每个线程搬运两轮即可完成所需元素的搬运;
+    */
+    int a_tile_row = threadIdx.x / BK;
+    int a_tile_col = threadIdx.x % BK;
+    int a_tile_stride = thread_num / BK;
+
+    int b_tile_row = threadIdx.x / BN;
+    int b_tile_col = threadIdx.x % BN;
+    int b_tile_stride = thread_num / BN;
+
+    float tmp[TM][TN] = {0.}; // 每个线程负责TM*TN个元素，则需要申请TM*TN个寄存器保存累加值，额外的一个寄存器用于缓存；
+    #pragma unroll
+    for (int k = 0; k < K; k += BK) {
+        #pragma unroll
+        for (int i = 0; i < BM; i += a_tile_stride) {
+            As[(a_tile_row + i) * BK + a_tile_col] = A[(a_tile_row + i) * K + a_tile_col];
+        }
+        #pragma unroll
+        for (int i = 0; i < BK; i += b_tile_stride) {
+            Bs[(b_tile_row + i) * BN + b_tile_col] = B[(b_tile_row + i) * N + b_tile_col];
+        }
+        __syncthreads();
+        A += BK;
+        B += BK * N;
+        #pragma unroll
+        for (int i = 0; i < BK; i++) {
+            #pragma unroll  // 循环展开，增加指令并行度
+            for (int j = 0; j < TM; j++) {
+                for (int l = 0; l < TN; l++)
+                    tmp[j][l] += As[(ty + j) * BK + i] * Bs[tx + l + i * BN];
+            }
+        }
+        __syncthreads();
+    }
+    #pragma unroll
+    for (int j = 0; j < TM; j++) {
+        for (int l = 0; l < TN; l++)
+            C[(ty + j) * N + tx + l] = alpha * tmp[j][l] + beta * C[(ty + j) * N + tx + l];
+    }
+}

upstream_ref/nvidia_sgemm_practice/kernel_5.cuh

@@ -0,0 +1,88 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+template<const int BM,
+        const int BN,
+        const int BK,
+        const int TM,
+        const int TN>
+__global__ void mysgemm_v5(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) {
+    int bx = blockIdx.x;
+    int by = blockIdx.y;
+
+    int block_row_thread = BN / TN;
+    int block_col_thread = BM / TM;
+    int thread_num = block_row_thread * block_col_thread; // 一个线程负责计算block中TM*TN个元素
+
+    int tx = (threadIdx.x % block_row_thread) * TN;
+    int ty = (threadIdx.x / block_row_thread) * TM;
+
+    __shared__ float As[BM * BK];
+    __shared__ float Bs[BK * BN];
+
+    // 移动到当前block
+    A = &A[by * BM * K];
+    B = &B[bx * BN];
+    C = &C[by * BM * N + bx * BN];
+
+    /*
+    当前线程负责搬运全局内存中第a_tile_row行，第a_tile_col列元素至共享内存第a_tile_row行，第a_tile_col列
+    a_tile_stride表示block中线程可搬运a_tile_stride行至共享内存；
+
+    若BM=64,BK=8,thread_num=512,则a_tile_stride=64,a_tile_stride=BM，表示每个线程搬运一轮即可完成所需元素的搬运;
+    若BM=128,BK=8,thread_num=512,则a_tile_stride=64,表示每个线程搬运两轮即可完成所需元素的搬运;
+    */
+    int a_tile_row = threadIdx.x / BK;
+    int a_tile_col = threadIdx.x % BK;
+    int a_tile_stride = thread_num / BK;
+
+    int b_tile_row = threadIdx.x / BN;
+    int b_tile_col = threadIdx.x % BN;
+    int b_tile_stride = thread_num / BN;
+
+    float tmp[TM][TN] = {0.}; // 每个线程负责TM*TN个元素，则需要申请TM*TN个寄存器保存累加值，额外的一个寄存器用于缓存；
+    float a_frag[TM] = {0.};
+    float b_frag[TN] = {0.};
+
+    #pragma unroll
+    for (int k = 0; k < K; k += BK) {
+        #pragma unroll
+        for (int i = 0; i < BM; i += a_tile_stride) {
+            As[(a_tile_row + i) * BK + a_tile_col] = A[(a_tile_row + i) * K + a_tile_col];
+        }
+        #pragma unroll
+        for (int i = 0; i < BK; i += b_tile_stride) {
+            Bs[(b_tile_row + i) * BN + b_tile_col] = B[(b_tile_row + i) * N + b_tile_col];
+        }
+        __syncthreads();
+        A += BK;
+        B += BK * N;
+        #pragma unroll
+        for (int i = 0; i < BK; i++) {
+            #pragma unroll
+            for (int j = 0; j < TM; j++) {
+                a_frag[j] = As[(ty + j) * BK + i];
+            }
+            #pragma unroll
+            for (int l = 0; l < TN; l++) {
+                b_frag[l] = Bs[tx + l + i * BN];
+            }
+            #pragma unroll
+            for (int j = 0; j < TM; j++) {
+                #pragma unroll
+                for (int l = 0; l < TN; l++)
+                    tmp[j][l] += a_frag[j] * b_frag[l];
+            }
+        }
+        __syncthreads();
+    }
+    #pragma unroll
+    for (int j = 0; j < TM; j++) {
+        for (int l = 0; l < TN; l++)
+            C[(ty + j) * N + tx + l] = alpha * tmp[j][l] + beta * C[(ty + j) * N + tx + l];
+    }
+}

upstream_ref/nvidia_sgemm_practice/kernel_6.cuh

@@ -0,0 +1,110 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#define OFFSET(row, col, ld) ((row)*(ld)+(col))
+#define FETCH_FLOAT4(pointer) (reinterpret_cast<float4*>(&(pointer))[0])
+
+template<const int BM,
+        const int BN,
+        const int BK,
+        const int TM,
+        const int TN>
+__global__ void mysgemm_v6(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) {
+    int bx = blockIdx.x;
+    int by = blockIdx.y;
+
+    const int block_row_thread = BN / TN;
+    const int block_col_thread = BM / TM;
+    const int thread_num = block_row_thread * block_col_thread; // 一个线程负责计算block中TM*TN个元素
+
+    // 当前线程对应thread tile的左上角元素在block中的位置
+    int tx = (threadIdx.x % block_row_thread) * TN;
+    int ty = (threadIdx.x / block_row_thread) * TM;
+
+    __shared__ float As[BK * BM];
+    __shared__ float Bs[BK * BN];
+
+
+    const int ldg_a_num = BK * BM / thread_num / 4; // 每个线程搬运4个浮点数，完成搬运至As需要所有线程搬运ldg_a_num轮
+    const int ldg_b_num = BK * BN / thread_num / 4; // 每个线程搬运4个浮点数，完成搬运至Bs需要所有线程搬运ldg_b_num轮
+
+    int a_tile_row = threadIdx.x / (BK / 4); // 每行4个字节作为一个内存块，当前线程负责第a_tile_row行的第a_tile_col个内存块的搬运
+    int a_tile_col = threadIdx.x % (BK / 4) * 4;
+    int a_tile_stride = BM / ldg_a_num; // 一共BM行，搬运ldg_a_num轮，每论搬运a_tile_stride行
+
+    int b_tile_row = threadIdx.x / (BN / 4); // 每行4个字节作为一个内存块，当前线程负责第b_tile_row行的第b_tile_col个内存块的搬运
+    int b_tile_col = threadIdx.x % (BN / 4) * 4;
+    int b_tile_stride = BK / ldg_b_num; // 一共BK行，搬运ldg_b_num轮，每论搬运b_tile_stride行
+
+    float accum[TM][TN] = {0.}; // 每个线程负责TM*TN个元素，则需要申请TM*TN个寄存器保存累加值，额外的一个寄存器用于缓存；
+
+    // 计算ldg_a_num的所有参数必须全部是const，否则不能用来申明数组大小
+    float ldg_a_reg[4 * ldg_a_num] = {0.}; // 每个线程搬运ldg_a_num轮，寄存器缓存ldg_a_num个float4元素，用于转置As矩阵
+
+    float a_frag[TM];  // 缓存As共享内存
+    float b_frag[TN];  // 缓存Bs共享内存
+
+    // 移动到当前block
+    A = &A[by * BM * K];
+    B = &B[bx * BN];
+    C = &C[by * BM * N + bx * BN];
+
+#pragma unroll
+    for (int k = 0; k < K; k += BK) {
+#pragma unroll
+        for (int i = 0; i < BM; i += a_tile_stride) {
+            int ldg_index = i / a_tile_stride * 4;  // 第ldg_index轮
+            FETCH_FLOAT4(ldg_a_reg[ldg_index]) =
+                    FETCH_FLOAT4(A[OFFSET(a_tile_row + i, a_tile_col, K)]);
+            // As转置存，其中ldg_a_reg做中间缓存，目的是读取时可以按FLOAT4读取
+            As[OFFSET(a_tile_col, i + a_tile_row, BM)] = ldg_a_reg[ldg_index];
+            As[OFFSET(a_tile_col + 1, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 1];
+            As[OFFSET(a_tile_col + 2, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 2];
+            As[OFFSET(a_tile_col + 3, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 3];
+        }
+#pragma unroll
+        for (int i = 0; i < BK; i += b_tile_stride) {
+            FETCH_FLOAT4(Bs[OFFSET(b_tile_row + i, b_tile_col, BN)]) =
+                    FETCH_FLOAT4(B[OFFSET(b_tile_row + i, b_tile_col, N)]); // 不需要转置
+        }
+        __syncthreads();
+        A += BK;
+        B += BK * N;
+#pragma unroll
+        for (int i = 0; i < BK; i++) {
+#pragma unroll
+            for (int m = 0; m < TM; m += 4) {
+                FETCH_FLOAT4(a_frag[m]) = FETCH_FLOAT4(As[OFFSET(i, ty + m, BM)]); // 偏移到当前thread tile
+            }
+#pragma unroll
+            for (int n = 0; n < TN; n += 4) {
+                FETCH_FLOAT4(b_frag[n]) = FETCH_FLOAT4(Bs[OFFSET(i, tx + n, BN)]); // 偏移到当前thread tile
+            }
+#pragma unroll
+            for (int m = 0; m < TM; m++) {
+#pragma unroll
+                for (int n = 0; n < TN; n++) {
+                    accum[m][n] += a_frag[m] * b_frag[n];
+                }
+            }
+        }
+        __syncthreads();
+    }
+#pragma unroll
+    for (int m = 0; m < TM; m++) {
+#pragma unroll
+        for (int n = 0; n < TN; n += 4) {
+            float4 ctmp = FETCH_FLOAT4(C[OFFSET(ty + m, tx + n, N)]);
+            //float4 atmp = FETCH_FLOAT4(accum[m][n]);
+            ctmp.x = alpha * accum[m][n] + beta * ctmp.x;
+            ctmp.y = alpha * accum[m][n + 1] + beta * ctmp.y;
+            ctmp.z = alpha * accum[m][n + 2] + beta * ctmp.z;
+            ctmp.w = alpha * accum[m][n + 3] + beta * ctmp.w;
+            FETCH_FLOAT4(C[OFFSET(ty + m, tx + n, N)]) = ctmp;
+        }
+    }
+}

upstream_ref/nvidia_sgemm_practice/kernel_7.cuh

@@ -0,0 +1,180 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#define OFFSET(row, col, ld) ((row)*(ld)+(col))
+#define FETCH_FLOAT4(pointer) (reinterpret_cast<float4*>(&(pointer))[0])
+
+template<const int BM,
+        const int BN,
+        const int BK,
+        const int TM,
+        const int TN>
+__global__ void mysgemm_v7(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) {
+    int bx = blockIdx.x;
+    int by = blockIdx.y;
+
+    const int block_row_thread = BN / TN;
+    const int block_col_thread = BM / TM;
+    const int thread_num = block_row_thread * block_col_thread; // 一个线程负责计算block中TM*TN个元素
+
+    // 当前线程对应thread tile的左上角元素在block中的位置
+    int tx = (threadIdx.x % block_row_thread) * TN;
+    int ty = (threadIdx.x / block_row_thread) * TM;
+
+    __shared__ float As[2][BK * BM]; // 增加一倍共享内存大小用于缓存
+    __shared__ float Bs[2][BK * BN];
+
+
+    const int ldg_a_num = BK * BM / thread_num / 4; // 每个线程搬运4个浮点数，完成搬运至As需要所有线程搬运ldg_a_num轮
+    const int ldg_b_num = BK * BN / thread_num / 4; // 每个线程搬运4个浮点数，完成搬运至Bs需要所有线程搬运ldg_b_num轮
+
+    int a_tile_row = threadIdx.x / (BK / 4); // 每行4个字节作为一个内存块，当前线程负责第a_tile_row行的第a_tile_col个内存块的搬运
+    int a_tile_col = threadIdx.x % (BK / 4) * 4;
+    int a_tile_stride = BM / ldg_a_num; // 一共BM行，搬运ldg_a_num轮，每论搬运a_tile_stride行
+
+    int b_tile_row = threadIdx.x / (BN / 4); // 每行4个字节作为一个内存块，当前线程负责第b_tile_row行的第b_tile_col个内存块的搬运
+    int b_tile_col = threadIdx.x % (BN / 4) * 4;
+    int b_tile_stride = BK / ldg_b_num; // 一共BK行，搬运ldg_b_num轮，每论搬运b_tile_stride行
+
+    float accum[TM][TN] = {0.}; // 每个线程负责TM*TN个元素，则需要申请TM*TN个寄存器保存累加值，额外的一个寄存器用于缓存；
+
+    // 计算ldg_a_num的所有参数必须全部是const，否则不能用来申明数组大小
+    float ldg_a_reg[4 * ldg_a_num] = {0.}; // 每个线程搬运ldg_a_num轮，寄存器缓存ldg_a_num个float4元素，用于转置As矩阵
+    float ldg_b_reg[4 * ldg_b_num] = {0.}; // 每个线程搬运ldg_a_num轮，寄存器缓存ldg_a_num个float4元素，用于转置As矩阵
+
+    float a_frag[2][TM];  // 缓存As共享内存,增加一倍寄存器大小用于缓存
+    float b_frag[2][TN];  // 缓存Bs共享内存,增加一倍寄存器大小用于缓存
+
+    // 移动到当前block
+    A = &A[by * BM * K];
+    B = &B[bx * BN];
+    C = &C[by * BM * N + bx * BN];
+
+    // first global to shared
+#pragma unroll
+    for (int i = 0; i < BM; i += a_tile_stride) {
+        int ldg_index = i / a_tile_stride * 4;  // 第ldg_index轮
+        FETCH_FLOAT4(ldg_a_reg[ldg_index]) =
+                FETCH_FLOAT4(A[OFFSET(a_tile_row + i, a_tile_col, K)]);
+        // As转置存，其中ldg_a_reg做中间缓存，目的是读取时可以按FLOAT4读取
+        As[0][OFFSET(a_tile_col, i + a_tile_row, BM)] = ldg_a_reg[ldg_index];
+        As[0][OFFSET(a_tile_col + 1, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 1];
+        As[0][OFFSET(a_tile_col + 2, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 2];
+        As[0][OFFSET(a_tile_col + 3, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 3];
+    }
+#pragma unroll
+    for (int i = 0; i < BK; i += b_tile_stride) {
+        FETCH_FLOAT4(Bs[0][OFFSET(b_tile_row + i, b_tile_col, BN)]) =
+                FETCH_FLOAT4(B[OFFSET(b_tile_row + i, b_tile_col, N)]); // 不需要转置
+    }
+    __syncthreads();
+
+    // first shared to frag
+#pragma unroll
+    for (int m = 0; m < TM; m += 4) {
+        FETCH_FLOAT4(a_frag[0][m]) = FETCH_FLOAT4(As[0][OFFSET(0, ty + m, BM)]); // 偏移到当前thread tile
+    }
+#pragma unroll
+    for (int n = 0; n < TN; n += 4) {
+        FETCH_FLOAT4(b_frag[0][n]) = FETCH_FLOAT4(Bs[0][OFFSET(0, tx + n, BN)]); // 偏移到当前thread tile
+    }
+
+
+    int write_index = 1;
+    int load_index;
+    int k = 0;
+    do {
+        k += BK;
+        // load global to reg
+        if (k < K) {
+#pragma unroll
+            for (int i = 0; i < BM; i += a_tile_stride) {
+                int ldg_index = i / a_tile_stride * 4;  // 第ldg_index轮
+                FETCH_FLOAT4(ldg_a_reg[ldg_index]) =
+                        FETCH_FLOAT4(A[OFFSET(a_tile_row + i, k + a_tile_col, K)]);
+            }
+#pragma unroll
+            for (int i = 0; i < BK; i += b_tile_stride) {
+                int ldg_index = i / b_tile_stride * 4;  // 第ldg_index轮
+                FETCH_FLOAT4(ldg_b_reg[ldg_index]) =
+                        FETCH_FLOAT4(B[OFFSET(k + b_tile_row + i, b_tile_col, N)]);
+            }
+        }
+
+        load_index = write_index ^ 1;
+#pragma unroll
+        for (int bk = 0; bk < BK - 1; bk++) {
+            for (int m = 0; m < TM; m += 4) {
+                FETCH_FLOAT4(a_frag[(bk + 1) % 2][m]) = FETCH_FLOAT4(
+                        As[load_index][OFFSET(bk + 1, ty + m, BM)]); // 偏移到当前thread tile
+            }
+#pragma unroll
+            for (int n = 0; n < TN; n += 4) {
+                FETCH_FLOAT4(b_frag[(bk + 1) % 2][n]) = FETCH_FLOAT4(
+                        Bs[load_index][OFFSET(bk + 1, tx + n, BN)]); // 偏移到当前thread tile
+            }
+#pragma unroll
+            for (int m = 0; m < TM; m++) {
+                for (int n = 0; n < TN; n++) {
+                    accum[m][n] += a_frag[bk % 2][m] * b_frag[bk % 2][n];
+                }
+            }
+        }
+        if (k < K) {
+#pragma unroll
+            for (int i = 0; i < BM; i += a_tile_stride) {
+                int ldg_index = i / a_tile_stride * 4;
+                As[write_index][OFFSET(a_tile_col, i + a_tile_row, BM)] = ldg_a_reg[ldg_index];
+                As[write_index][OFFSET(a_tile_col + 1, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 1];
+                As[write_index][OFFSET(a_tile_col + 2, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 2];
+                As[write_index][OFFSET(a_tile_col + 3, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 3];
+            }
+#pragma unroll
+            for (int i = 0; i < BK; i += b_tile_stride) {
+                int ldg_index = i / b_tile_stride * 4;
+                FETCH_FLOAT4(Bs[write_index][OFFSET(b_tile_row + i, b_tile_col, BN)]) =
+                        FETCH_FLOAT4(ldg_b_reg[ldg_index]);
+            }
+            __syncthreads();
+#pragma unroll
+            for (int m = 0; m < TM; m += 4) {
+                FETCH_FLOAT4(a_frag[0][m]) = FETCH_FLOAT4(
+                        As[write_index][OFFSET(0, ty + m, BM)]); // 偏移到当前thread tile
+            }
+#pragma unroll
+            for (int n = 0; n < TN; n += 4) {
+                FETCH_FLOAT4(b_frag[0][n]) = FETCH_FLOAT4(
+                        Bs[write_index][OFFSET(0, tx + n, BN)]); // 偏移到当前thread tile
+            }
+
+            write_index ^= 1;
+        }
+#pragma unroll
+        for (int m = 0; m < TM; m++) {
+#pragma unroll
+            for (int n = 0; n < TN; n++) {
+                accum[m][n] += a_frag[(BK - 1) % 2][m] * b_frag[(BK - 1) % 2][n];
+            }
+        }
+
+
+    } while (k < K);
+    
+    // C = alpha*AB+C
+#pragma unroll
+    for (int m = 0; m < TM; m++) {
+#pragma unroll
+        for (int n = 0; n < TN; n += 4) {
+            float4 ctmp = FETCH_FLOAT4(C[OFFSET(ty + m, tx + n, N)]);
+            ctmp.x = alpha * accum[m][n] + beta * ctmp.x;
+            ctmp.y = alpha * accum[m][n + 1] + beta * ctmp.y;
+            ctmp.z = alpha * accum[m][n + 2] + beta * ctmp.z;
+            ctmp.w = alpha * accum[m][n + 3] + beta * ctmp.w;
+            FETCH_FLOAT4(C[OFFSET(ty + m, tx + n, N)]) = ctmp;
+        }
+    }
+}

upstream_ref/nvidia_sgemm_practice/sgemm.cu

@@ -0,0 +1,119 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <sys/time.h>
+#include <utils.cuh>
+
+#define cudaCheck(err) (cudaCheck(err, __FILE__, __LINE__))
+
+int main(int argc, char **argv) {
+    if (argc != 2) {
+        printf("Please select a kernel (range 0 - 11, here 0 is for NVIDIA cuBLAS).\n");
+        exit(EXIT_FAILURE);
+    }
+
+    // cuda kernel num
+    int kernel_num = atoi(argv[1]);
+    if (kernel_num < 0 || kernel_num > 11) {
+        printf("Please enter a valid kernel number (0-11).\n");
+        exit(EXIT_FAILURE);
+    } else {
+        printf("Select kernel %d.\n", kernel_num);
+    };
+
+    // 申明句柄，创建句柄, cublasCreate会返回一个cublasStatus_t类型的值，用来判断句柄是否创建成功(值为0)
+    cublasHandle_t handle;
+    if (cublasCreate(&handle)) {
+        printf("Create cublas handle error.\n");
+        exit(EXIT_FAILURE);
+    };
+
+    // 采用cudaEvent进行gpu流计时，cudaEvent相当于在目标流中发布事件任务
+    float elapsed_time;
+    cudaEvent_t beg, end;
+    cudaEventCreate(&beg);
+    cudaEventCreate(&end);
+
+    // matrix size
+    int size_len = 24;
+    int SIZE[size_len];
+    for (int i = 0; i < size_len; i++)
+        SIZE[i] = 256 * (i + 1);
+
+    int m, n, k, max_size;
+    max_size = SIZE[size_len - 1];
+    printf("max_size=%d\n", max_size);
+
+    float alpha = 1.0, beta = 0.; //two arbitary input parameters，C=α*AB+β*C
+
+    float *A = NULL, *B = NULL, *C = NULL, *C_ref = NULL;     //host matrices
+    float *dA = NULL, *dB = NULL, *dC = NULL, *dC_ref = NULL; //device matrices
+
+    A = (float *) malloc(sizeof(float) * max_size * max_size);
+    B = (float *) malloc(sizeof(float) * max_size * max_size);
+    C = (float *) malloc(sizeof(float) * max_size * max_size);
+    C_ref = (float *) malloc(sizeof(float) * max_size * max_size);
+
+    randomize_matrix(A, max_size * max_size);
+    randomize_matrix(B, max_size * max_size);
+    randomize_matrix(C, max_size * max_size);
+    copy_matrix(C, C_ref, max_size * max_size);
+
+    cudaCheck(cudaMalloc((void **) &dA, sizeof(float) * max_size * max_size));
+    cudaCheck(cudaMalloc((void **) &dB, sizeof(float) * max_size * max_size));
+    cudaCheck(cudaMalloc((void **) &dC, sizeof(float) * max_size * max_size));
+    cudaCheck(cudaMalloc((void **) &dC_ref, sizeof(float) * max_size * max_size));
+
+    cudaCheck(cudaMemcpy(dA, A, sizeof(float) * max_size * max_size, cudaMemcpyHostToDevice));
+    cudaCheck(cudaMemcpy(dB, B, sizeof(float) * max_size * max_size, cudaMemcpyHostToDevice));
+    cudaCheck(cudaMemcpy(dC, C, sizeof(float) * max_size * max_size, cudaMemcpyHostToDevice));
+    cudaCheck(cudaMemcpy(dC_ref, C_ref, sizeof(float) * max_size * max_size, cudaMemcpyHostToDevice));
+
+    int repeat_times = 10;
+    for (int i = 0; i < size_len; i++) {
+        m = n = k = SIZE[i];
+
+        printf("m=n=k=%d\n", m);
+        // 验证计算正确性，同时在核函数计时前预先执行一次，避免冷启动误差
+        if (kernel_num != 0) {
+            test_kernel(0, m, n, k, alpha, dA, dB, beta, dC_ref, handle);      // cuBLAS
+            test_kernel(kernel_num, m, n, k, alpha, dA, dB, beta, dC, handle); // user define
+            cudaDeviceSynchronize();
+            cudaMemcpy(C, dC, sizeof(float) * m * n, cudaMemcpyDeviceToHost);
+            cudaMemcpy(C_ref, dC_ref, sizeof(float) * m * n, cudaMemcpyDeviceToHost);
+            cudaDeviceSynchronize();
+
+            if (!verify_matrix(C_ref, C, m * n)) {
+                printf("Failed to pass the correctness verification against NVIDIA cuBLAS. Exited.\n");
+                exit(EXIT_FAILURE);
+            }
+        }
+        cudaDeviceSynchronize();
+
+        cudaEventRecord(beg);
+        for (int j = 0; j < repeat_times; j++) {
+            test_kernel(kernel_num, m, n, k, alpha, dA, dB, beta, dC, handle);
+        }
+        cudaEventRecord(end);
+        cudaEventSynchronize(beg);
+        cudaEventSynchronize(end);
+        cudaEventElapsedTime(&elapsed_time, beg, end);
+        elapsed_time /= 1000.; //换算成秒
+
+        printf("Average elasped time: (%f) second, performance: (%f) GFLOPS. size: (%d).\n",
+               elapsed_time / repeat_times, 2. * 1e-9 * repeat_times * m * n * k / elapsed_time, m);
+        fflush(stdout);
+        copy_matrix(C_ref, C, m * n); //sync C with cuBLAS to prepare for the next run
+    }
+
+    // 释放CPU和GPU空间
+    free(A);
+    free(B);
+    free(C);
+    free(C_ref);
+    cudaFree(dA);
+    cudaFree(dB);
+    cudaFree(dC);
+    cudaFree(dC_ref);
+
+    return 0;
+};

upstream_ref/nvidia_sgemm_practice/utils.cuh

@@ -0,0 +1,42 @@
+#pragma once
+#include <stdio.h>
+#include <stdlib.h>
+#include <time.h>
+#include <unistd.h>
+#include <sys/time.h>
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+
+/*
+=====================================
+CUDA操作
+=====================================
+*/
+void cudaCheck(cudaError_t error, const char *file, int line); //CUDA错误检查
+void CudaDeviceInfo();                                         // 打印CUDA信息
+
+/*
+=====================================
+矩阵操作
+=====================================
+*/
+void randomize_matrix(float *mat, int N);            // 随机初始化矩阵
+void copy_matrix(float *src, float *dest, int N);    // 复制矩阵
+void print_matrix(const float *A, int M, int N);     // 打印矩阵
+bool verify_matrix(float *mat1, float *mat2, int N); // 验证矩阵
+
+/*
+=====================================
+计时操作
+=====================================
+*/
+float get_current_sec();                        // 获取当前时刻
+float cpu_elapsed_time(float &beg, float &end); // 计算时间差
+
+/*
+=====================================
+kernel操作
+=====================================
+*/
+//调用指定核函数计算矩阵乘法
+void test_kernel(int kernel_num, int m, int n, int k, float alpha, float *A, float *B, float beta, float *C, cublasHandle_t handle);

upstream_ref/sgemm_cuda/10_kernel_warptiling.cuh

@@ -0,0 +1,187 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+const int WARPSIZE = 32; // warpSize is not constexpr
+
+namespace wt {
+template <const int BM, const int BN, const int BK, const int rowStrideA,
+          const int rowStrideB>
+__device__ void loadFromGmem(int N, int K, const float *A, const float *B,
+                             float *As, float *Bs, int innerRowA, int innerColA,
+                             int innerRowB, int innerColB) {
+  for (uint offset = 0; offset + rowStrideA <= BM; offset += rowStrideA) {
+    const float4 tmp = reinterpret_cast<const float4 *>(
+        &A[(innerRowA + offset) * K + innerColA * 4])[0];
+    // float4 tmp;
+    // asm("ld.global.nc.v4.f32 {%0, %1, %2, %3}, [%4];"
+    //     : "=f"(tmp.x), "=f"(tmp.y), "=f"(tmp.z), "=f"(tmp.w)
+    //     : "l"(&A[(innerRowA + offset) * K + innerColA * 4]));
+    As[(innerColA * 4 + 0) * BM + innerRowA + offset] = tmp.x;
+    As[(innerColA * 4 + 1) * BM + innerRowA + offset] = tmp.y;
+    As[(innerColA * 4 + 2) * BM + innerRowA + offset] = tmp.z;
+    As[(innerColA * 4 + 3) * BM + innerRowA + offset] = tmp.w;
+  }
+
+  for (uint offset = 0; offset + rowStrideB <= BK; offset += rowStrideB) {
+    reinterpret_cast<float4 *>(
+        &Bs[(innerRowB + offset) * BN + innerColB * 4])[0] =
+        reinterpret_cast<const float4 *>(
+            &B[(innerRowB + offset) * N + innerColB * 4])[0];
+    // asm("ld.global.v4.f32 {%0, %1, %2, %3}, [%4];"
+    //     : "=f"(Bs[(innerRowB + offset) * BN + innerColB * 4 + 0]),
+    //       "=f"(Bs[(innerRowB + offset) * BN + innerColB * 4 + 1]),
+    //       "=f"(Bs[(innerRowB + offset) * BN + innerColB * 4 + 2]),
+    //       "=f"(Bs[(innerRowB + offset) * BN + innerColB * 4 + 3])
+    //     : "l"(&B[(innerRowB + offset) * N + innerColB * 4]));
+  }
+}
+
+template <const int BM, const int BN, const int BK, const int WM, const int WN,
+          const int WMITER, const int WNITER, const int WSUBM, const int WSUBN,
+          const int TM, const int TN>
+__device__ void
+processFromSmem(float *regM, float *regN, float *threadResults, const float *As,
+                const float *Bs, const uint warpRow, const uint warpCol,
+                const uint threadRowInWarp, const uint threadColInWarp) {
+  for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
+    // populate registers for whole warptile
+    for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
+      for (uint i = 0; i < TM; ++i) {
+        regM[wSubRowIdx * TM + i] =
+            As[(dotIdx * BM) + warpRow * WM + wSubRowIdx * WSUBM +
+               threadRowInWarp * TM + i];
+      }
+    }
+    for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
+      for (uint i = 0; i < TN; ++i) {
+        regN[wSubColIdx * TN + i] =
+            Bs[(dotIdx * BN) + warpCol * WN + wSubColIdx * WSUBN +
+               threadColInWarp * TN + i];
+      }
+    }
+
+    // execute warptile matmul
+    for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
+      for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
+        // calculate per-thread results
+        for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
+          for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
+            threadResults[(wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
+                          (wSubColIdx * TN) + resIdxN] +=
+                regM[wSubRowIdx * TM + resIdxM] *
+                regN[wSubColIdx * TN + resIdxN];
+          }
+        }
+      }
+    }
+  }
+}
+
+} // namespace wt
+
+/*
+ * @tparam BM The threadblock size for M dimension SMEM caching.
+ * @tparam BN The threadblock size for N dimension SMEM caching.
+ * @tparam BK The threadblock size for K dimension SMEM caching.
+ * @tparam WM M dim of continuous tile computed by each warp
+ * @tparam WN N dim of continuous tile computed by each warp
+ * @tparam WMITER The number of subwarp tiling steps in M dimension.
+ * @tparam WNITER The number of subwarp tiling steps in N dimension.
+ * @tparam TM The per-thread tile size for M dimension.
+ * @tparam TN The per-thread tile size for N dimension.
+ */
+template <const int BM, const int BN, const int BK, const int WM, const int WN,
+          const int WNITER, const int TM, const int TN, const int NUM_THREADS>
+__global__ void __launch_bounds__(NUM_THREADS)
+    sgemmWarptiling(int M, int N, int K, float alpha, float *A, float *B,
+                    float beta, float *C) {
+  const uint cRow = blockIdx.y;
+  const uint cCol = blockIdx.x;
+
+  // Placement of the warp in the threadblock tile
+  const uint warpIdx = threadIdx.x / WARPSIZE; // the warp this thread is in
+  const uint warpCol = warpIdx % (BN / WN);
+  const uint warpRow = warpIdx / (BN / WN);
+
+  // size of the warp subtile
+  constexpr uint WMITER = (WM * WN) / (WARPSIZE * TM * TN * WNITER);
+  constexpr uint WSUBM = WM / WMITER; // 64/2=32
+  constexpr uint WSUBN = WN / WNITER; // 32/2=16
+
+  // Placement of the thread in the warp subtile
+  const uint threadIdxInWarp = threadIdx.x % WARPSIZE;         // [0, 31]
+  const uint threadColInWarp = threadIdxInWarp % (WSUBN / TN); // i%(16/4)
+  const uint threadRowInWarp = threadIdxInWarp / (WSUBN / TN); // i/4
+
+  // allocate space for the current blocktile in SMEM
+  __shared__ float As[BM * BK];
+  __shared__ float Bs[BK * BN];
+
+  // Move blocktile to beginning of A's row and B's column
+  A += cRow * BM * K;
+  B += cCol * BN;
+  // Move C_ptr to warp's output tile
+  C += (cRow * BM + warpRow * WM) * N + cCol * BN + warpCol * WN;
+
+  // calculating the indices that this thread will load into SMEM
+  // we'll load 128bit / 32bit = 4 elements per thread at each step
+  const uint innerRowA = threadIdx.x / (BK / 4);
+  const uint innerColA = threadIdx.x % (BK / 4);
+  constexpr uint rowStrideA = (NUM_THREADS * 4) / BK;
+  const uint innerRowB = threadIdx.x / (BN / 4);
+  const uint innerColB = threadIdx.x % (BN / 4);
+  constexpr uint rowStrideB = NUM_THREADS / (BN / 4);
+
+  // allocate thread-local cache for results in registerfile
+  float threadResults[WMITER * TM * WNITER * TN] = {0.0};
+  // we cache into registers on the warptile level
+  float regM[WMITER * TM] = {0.0};
+  float regN[WNITER * TN] = {0.0};
+
+  // outer-most loop over block tiles
+  for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
+    wt::loadFromGmem<BM, BN, BK, rowStrideA, rowStrideB>(
+        N, K, A, B, As, Bs, innerRowA, innerColA, innerRowB, innerColB);
+    __syncthreads();
+    wt::processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN, TM,
+                        TN>(regM, regN, threadResults, As, Bs, warpRow, warpCol,
+                            threadRowInWarp, threadColInWarp);
+    A += BK;     // move BK columns to right
+    B += BK * N; // move BK rows down
+    __syncthreads();
+  }
+
+  // write out the results
+  for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
+    for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
+      // move C pointer to current warp subtile
+      float *C_interim = C + (wSubRowIdx * WSUBM) * N + wSubColIdx * WSUBN;
+      for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
+        for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) {
+          // load C vector into registers
+          float4 tmp = reinterpret_cast<float4 *>(
+              &C_interim[(threadRowInWarp * TM + resIdxM) * N +
+                         threadColInWarp * TN + resIdxN])[0];
+          // perform GEMM update in reg
+          const int i = (wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
+                        wSubColIdx * TN + resIdxN;
+          tmp.x = alpha * threadResults[i + 0] + beta * tmp.x;
+          tmp.y = alpha * threadResults[i + 1] + beta * tmp.y;
+          tmp.z = alpha * threadResults[i + 2] + beta * tmp.z;
+          tmp.w = alpha * threadResults[i + 3] + beta * tmp.w;
+          // write back
+          reinterpret_cast<float4 *>(
+              &C_interim[(threadRowInWarp * TM + resIdxM) * N +
+                         threadColInWarp * TN + resIdxN])[0] = tmp;
+        }
+      }
+    }
+  }
+}

upstream_ref/sgemm_cuda/11_kernel_double_buffering.cuh

@@ -0,0 +1,220 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+namespace db {
+
+template <const int BM, const int BN, const int BK, const int rowStrideA,
+          const int rowStrideB>
+__device__ void loadFromGmem(const int N, const int K, float *A, float *B,
+                             float *As, float *Bs, const int innerRowA,
+                             const int innerColA, const int innerRowB,
+                             const int innerColB) {
+  for (uint offset = 0; offset + rowStrideA <= BM; offset += rowStrideA) {
+    float4 tmp = reinterpret_cast<float4 *>(
+        &A[(innerRowA + offset) * K + innerColA * 4])[0];
+    // transpose A while storing it
+    As[(innerColA * 4 + 0) * BM + innerRowA + offset] = tmp.x;
+    As[(innerColA * 4 + 1) * BM + innerRowA + offset] = tmp.y;
+    As[(innerColA * 4 + 2) * BM + innerRowA + offset] = tmp.z;
+    As[(innerColA * 4 + 3) * BM + innerRowA + offset] = tmp.w;
+  }
+
+  for (uint offset = 0; offset + rowStrideB <= BK; offset += rowStrideB) {
+    reinterpret_cast<float4 *>(
+        &Bs[(innerRowB + offset) * BN + innerColB * 4])[0] =
+        reinterpret_cast<float4 *>(
+            &B[(innerRowB + offset) * N + innerColB * 4])[0];
+  }
+}
+
+template <const int BM, const int BN, const int BK, const int WM, const int WN,
+          const int WMITER, const int WNITER, const int WSUBM, const int WSUBN,
+          const int TM, const int TN>
+__device__ void
+processFromSmem(float *regM, float *regN, float *threadResults, const float *As,
+                const float *Bs, const uint warpRow, const uint warpCol,
+                const uint threadRowInWarp, const uint threadColInWarp) {
+  for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
+    // populate registers for whole warptile
+    for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
+      for (uint i = 0; i < TM; ++i) {
+        regM[wSubRowIdx * TM + i] =
+            As[(dotIdx * BM) + warpRow * WM + wSubRowIdx * WSUBM +
+               threadRowInWarp * TM + i];
+      }
+    }
+    for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
+      for (uint i = 0; i < TN; ++i) {
+        regN[wSubColIdx * TN + i] =
+            Bs[(dotIdx * BN) + warpCol * WN + wSubColIdx * WSUBN +
+               threadColInWarp * TN + i];
+      }
+    }
+
+    // execute warptile matmul
+    for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
+      for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
+        // calculate per-thread results
+        for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
+          for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
+            threadResults[(wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
+                          (wSubColIdx * TN) + resIdxN] +=
+                regM[wSubRowIdx * TM + resIdxM] *
+                regN[wSubColIdx * TN + resIdxN];
+          }
+        }
+      }
+    }
+  }
+}
+
+} // namespace db
+
+template <const int BM, const int BN, const int BK, const int WM, const int WN,
+          const int WNITER, const int TM, const int TN, const int NUM_THREADS>
+__global__ void __launch_bounds__(NUM_THREADS)
+    sgemmDoubleBuffering(const int M, const int N, const int K,
+                         const float alpha, float *A, float *B, float beta,
+                         float *C) {
+  const uint cRow = blockIdx.y;
+  const uint cCol = blockIdx.x;
+
+  // Placement of the warp in the threadblock tile
+  const uint warpIdx = threadIdx.x / WARPSIZE; // the warp this thread is in
+  const uint warpCol = warpIdx % (BN / WN);
+  const uint warpRow = warpIdx / (BN / WN);
+
+  // size of the warp subtile
+  constexpr uint WMITER = (WM * WN) / (WARPSIZE * TM * TN * WNITER);
+  constexpr uint WSUBM = WM / WMITER; // 64/2=32
+  constexpr uint WSUBN = WN / WNITER; // 32/2=16
+
+  // Placement of the thread in the warp subtile
+  const uint threadIdxInWarp = threadIdx.x % WARPSIZE;         // [0, 31]
+  const uint threadColInWarp = threadIdxInWarp % (WSUBN / TN); // i%(16/4)
+  const uint threadRowInWarp = threadIdxInWarp / (WSUBN / TN); // i/4
+
+  // allocate space for the current blocktile in SMEM
+  __shared__ float As[2 * BM * BK];
+  __shared__ float Bs[2 * BK * BN];
+
+  // setup double buffering split
+  bool doubleBufferIdx = threadIdx.x >= (NUM_THREADS / 2);
+
+  // Move blocktile to beginning of A's row and B's column
+  A += cRow * BM * K;
+  B += cCol * BN;
+  // Move C_ptr to warp's output tile
+  C += (cRow * BM + warpRow * WM) * N + cCol * BN + warpCol * WN;
+
+  // calculating the indices that this thread will load into SMEM
+  // for the loading, we're pretending like there's half as many threads
+  // as there actually are
+  const uint innerRowA = (threadIdx.x % (NUM_THREADS / 2)) / (BK / 4);
+  const uint innerColA = (threadIdx.x % (NUM_THREADS / 2)) % (BK / 4);
+  constexpr uint rowStrideA = ((NUM_THREADS / 2) * 4) / BK;
+  const uint innerRowB = (threadIdx.x % (NUM_THREADS / 2)) / (BN / 4);
+  const uint innerColB = (threadIdx.x % (NUM_THREADS / 2)) % (BN / 4);
+  constexpr uint rowStrideB = (NUM_THREADS / 2) / (BN / 4);
+
+  // allocate thread-local cache for results in registerfile
+  float threadResults[WMITER * TM * WNITER * TN] = {0.0};
+  // we cache into registers on the warptile level
+  float regM[WMITER * TM] = {0.0};
+  float regN[WNITER * TN] = {0.0};
+
+  if (doubleBufferIdx == 0) {
+    // load first (B0)
+    db::loadFromGmem<BM, BN, BK, rowStrideA, rowStrideB>(
+        N, K, A, B, As, Bs, innerRowA, innerColA, innerRowB, innerColB);
+  }
+  __syncthreads();
+
+  // outer-most loop over block tiles
+  for (uint bkIdx = 0; bkIdx < K; bkIdx += 2 * BK) {
+    if (doubleBufferIdx == 0) {
+      // process current (B0)
+      db::processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN, TM,
+                          TN>(regM, regN, threadResults, As, Bs, warpRow,
+                              warpCol, threadRowInWarp, threadColInWarp);
+      __syncthreads();
+
+      // process current+1 (B1)
+      if (bkIdx + BK < K) {
+        db::processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN,
+                            TM, TN>(regM, regN, threadResults, As + (BM * BK),
+                                    Bs + (BK * BN), warpRow, warpCol,
+                                    threadRowInWarp, threadColInWarp);
+      }
+      __syncthreads();
+
+      // load current + 2 (B0)
+      if (bkIdx + 2 * BK < K) {
+        db::loadFromGmem<BM, BN, BK, rowStrideA, rowStrideB>(
+            N, K, A + 2 * BK, B + 2 * BK * N, As, Bs, innerRowA, innerColA,
+            innerRowB, innerColB);
+      }
+    } else {
+      // load current + 1 (B1)
+      if (bkIdx + BK < K) {
+        db::loadFromGmem<BM, BN, BK, rowStrideA, rowStrideB>(
+            N, K, A + BK, B + BK * N, As + (BM * BK), Bs + (BK * BN), innerRowA,
+            innerColA, innerRowB, innerColB);
+      }
+      __syncthreads();
+
+      // process current (B0)
+      db::processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN, TM,
+                          TN>(regM, regN, threadResults, As, Bs, warpRow,
+                              warpCol, threadRowInWarp, threadColInWarp);
+      __syncthreads();
+
+      // process current+1 (B1)
+      if (bkIdx + BK < K) {
+        db::processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN,
+                            TM, TN>(regM, regN, threadResults, As + (BM * BK),
+                                    Bs + (BK * BN), warpRow, warpCol,
+                                    threadRowInWarp, threadColInWarp);
+      }
+    }
+
+    A += 2 * BK;     // move BK columns to right
+    B += 2 * BK * N; // move BK rows down
+    __syncthreads();
+  }
+
+  // write out the results
+  for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
+    for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
+      // move C pointer to current warp subtile
+      float *C_interim = C + (wSubRowIdx * WSUBM) * N + wSubColIdx * WSUBN;
+      for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
+        for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) {
+          // load C vector into registers
+          float4 tmp = reinterpret_cast<float4 *>(
+              &C_interim[(threadRowInWarp * TM + resIdxM) * N +
+                         threadColInWarp * TN + resIdxN])[0];
+          // perform GEMM update in reg
+          const int i = (wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
+                        wSubColIdx * TN + resIdxN;
+          tmp.x = alpha * threadResults[i + 0] + beta * tmp.x;
+          tmp.y = alpha * threadResults[i + 1] + beta * tmp.y;
+          tmp.z = alpha * threadResults[i + 2] + beta * tmp.z;
+          tmp.w = alpha * threadResults[i + 3] + beta * tmp.w;
+          // write back
+          reinterpret_cast<float4 *>(
+              &C_interim[(threadRowInWarp * TM + resIdxM) * N +
+                         threadColInWarp * TN + resIdxN])[0] = tmp;
+        }
+      }
+    }
+  }
+}

upstream_ref/sgemm_cuda/12_kernel_double_buffering.cuh

@@ -0,0 +1,229 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cooperative_groups.h>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda/barrier>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+namespace {
+template <const int BM, const int BN, const int BK, const int rowStrideA,
+          const int rowStrideB, typename T>
+__device__ void loadFromGmem(int N, int K, float *A, float *B, float *As,
+                             float *Bs, int innerRowA, int innerColA,
+                             int innerRowB, int innerColB, T &barrier) {
+
+  for (uint offset = 0; offset + rowStrideA <= BM; offset += rowStrideA) {
+    cuda::memcpy_async(&As[(innerColA * 4 + 0) * BM + innerRowA + offset],
+                       &A[(innerRowA + offset) * K + innerColA * 4],
+                       cuda::aligned_size_t<sizeof(float)>(sizeof(float)),
+                       barrier);
+    cuda::memcpy_async(&As[(innerColA * 4 + 1) * BM + innerRowA + offset],
+                       &A[(innerRowA + offset) * K + innerColA * 4 + 1],
+                       cuda::aligned_size_t<sizeof(float)>(sizeof(float)),
+                       barrier);
+    cuda::memcpy_async(&As[(innerColA * 4 + 2) * BM + innerRowA + offset],
+                       &A[(innerRowA + offset) * K + innerColA * 4 + 2],
+                       cuda::aligned_size_t<sizeof(float)>(sizeof(float)),
+                       barrier);
+    cuda::memcpy_async(&As[(innerColA * 4 + 3) * BM + innerRowA + offset],
+                       &A[(innerRowA + offset) * K + innerColA * 4 + 3],
+                       cuda::aligned_size_t<sizeof(float)>(sizeof(float)),
+                       barrier);
+  }
+
+  for (uint offset = 0; offset + rowStrideB <= BK; offset += rowStrideB) {
+    cuda::memcpy_async(&Bs[(innerRowB + offset) * BN + innerColB * 4],
+                       &B[(innerRowB + offset) * N + innerColB * 4],
+                       cuda::aligned_size_t<sizeof(float4)>(sizeof(float4)),
+                       barrier);
+  }
+}
+
+template <const int BM, const int BN, const int BK, const int WM, const int WN,
+          const int WMITER, const int WNITER, const int WSUBM, const int WSUBN,
+          const int TM, const int TN>
+__device__ void
+processFromSmem(float *regM, float *regN, float *threadResults, const float *As,
+                const float *Bs, const uint warpRow, const uint warpCol,
+                const uint threadRowInWarp, const uint threadColInWarp) {
+  for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
+    // populate registers for whole warptile
+    for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
+      for (uint i = 0; i < TM; ++i) {
+        regM[wSubRowIdx * TM + i] =
+            As[(dotIdx * BM) + warpRow * WM + wSubRowIdx * WSUBM +
+               threadRowInWarp * TM + i];
+      }
+    }
+    for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
+      for (uint i = 0; i < TN; ++i) {
+        regN[wSubColIdx * TN + i] =
+            Bs[(dotIdx * BN) + warpCol * WN + wSubColIdx * WSUBN +
+               threadColInWarp * TN + i];
+      }
+    }
+
+    // execute warptile matmul
+    for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
+      for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
+        // calculate per-thread results
+        for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
+          for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
+            threadResults[(wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
+                          (wSubColIdx * TN) + resIdxN] +=
+                regM[wSubRowIdx * TM + resIdxM] *
+                regN[wSubColIdx * TN + resIdxN];
+          }
+        }
+      }
+    }
+  }
+}
+
+} // namespace
+
+/*
+ * @tparam BM The threadblock size for M dimension SMEM caching.
+ * @tparam BN The threadblock size for N dimension SMEM caching.
+ * @tparam BK The threadblock size for K dimension SMEM caching.
+ * @tparam WM M dim of continuous tile computed by each warp
+ * @tparam WN N dim of continuous tile computed by each warp
+ * @tparam WMITER The number of subwarp tiling steps in M dimension.
+ * @tparam WNITER The number of subwarp tiling steps in N dimension.
+ * @tparam TM The per-thread tile size for M dimension.
+ * @tparam TN The per-thread tile size for N dimension.
+ */
+template <const int BM, const int BN, const int BK, const int WM, const int WN,
+          const int WNITER, const int TM, const int TN, const int NUM_THREADS>
+__global__ void __launch_bounds__(NUM_THREADS)
+    runSgemmDoubleBuffering2(int M, int N, int K, float alpha, float *A,
+                             float *B, float beta, float *C) {
+  auto block = cooperative_groups::this_thread_block();
+  __shared__ cuda::barrier<cuda::thread_scope::thread_scope_block> frontBarrier;
+  __shared__ cuda::barrier<cuda::thread_scope::thread_scope_block> backBarrier;
+  auto frontBarrierPtr = &frontBarrier;
+  auto backBarrierPtr = &backBarrier;
+  if (block.thread_rank() == 0) {
+    init(&frontBarrier, block.size());
+    init(&backBarrier, block.size());
+  }
+  __syncthreads();
+
+  const uint cRow = blockIdx.y;
+  const uint cCol = blockIdx.x;
+
+  // Placement of the warp in the threadblock tile
+  const uint warpIdx = threadIdx.x / WARPSIZE; // the warp this thread is in
+  const uint warpCol = warpIdx % (BN / WN);
+  const uint warpRow = warpIdx / (BN / WN);
+
+  // size of the warp subtile
+  constexpr uint WMITER = (WM * WN) / (WARPSIZE * TM * TN * WNITER);
+  constexpr uint WSUBM = WM / WMITER; // 64/2=32
+  constexpr uint WSUBN = WN / WNITER; // 32/2=16
+
+  // Placement of the thread in the warp subtile
+  const uint threadIdxInWarp = threadIdx.x % WARPSIZE;         // [0, 31]
+  const uint threadColInWarp = threadIdxInWarp % (WSUBN / TN); // i%(16/4)
+  const uint threadRowInWarp = threadIdxInWarp / (WSUBN / TN); // i/4
+
+  // allocate space for the current blocktile in SMEM
+  __shared__ float As[2 * BM * BK];
+  __shared__ float Bs[2 * BK * BN];
+
+  // Move blocktile to beginning of A's row and B's column
+  A += cRow * BM * K;
+  B += cCol * BN;
+  // Move C_ptr to warp's output tile
+  C += (cRow * BM + warpRow * WM) * N + cCol * BN + warpCol * WN;
+
+  // calculating the indices that this thread will load into SMEM
+  // we'll load 128bit / 32bit = 4 elements per thread at each step
+  const uint innerRowA = threadIdx.x / (BK / 4);
+  const uint innerColA = threadIdx.x % (BK / 4);
+  constexpr uint rowStrideA = (NUM_THREADS * 4) / BK;
+  const uint innerRowB = threadIdx.x / (BN / 4);
+  const uint innerColB = threadIdx.x % (BN / 4);
+  constexpr uint rowStrideB = NUM_THREADS / (BN / 4);
+
+  // allocate thread-local cache for results in registerfile
+  float threadResults[WMITER * TM * WNITER * TN] = {0.0};
+  // we cache into registers on the warptile level
+  float regM[WMITER * TM] = {0.0};
+  float regN[WNITER * TN] = {0.0};
+
+  int As_offset = 0;
+  int Bs_offset = 0;
+
+  // double-buffering: load first blocktile into SMEM
+  loadFromGmem<BM, BN, BK, rowStrideA, rowStrideB>(
+      N, K, A, B, As + As_offset * BM * BK, Bs + Bs_offset * BK * BN, innerRowA,
+      innerColA, innerRowB, innerColB, (*frontBarrierPtr));
+
+  // outer-most loop over block tiles
+  for (uint bkIdx = 0; bkIdx < K - BK; bkIdx += BK) {
+    // double-buffering: load next blocktile into SMEM
+    loadFromGmem<BM, BN, BK, rowStrideA, rowStrideB>(
+        N, K, A + BK, B + BK * N, As + (1 - As_offset) * BM * BK,
+        Bs + (1 - Bs_offset) * BK * BN, innerRowA, innerColA, innerRowB,
+        innerColB, (*backBarrierPtr));
+
+    // compute the current blocktile
+    (*frontBarrierPtr).arrive_and_wait();
+    processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN, TM, TN>(
+        regM, regN, threadResults, As + As_offset * BM * BK,
+        Bs + Bs_offset * BK * BN, warpRow, warpCol, threadRowInWarp,
+        threadColInWarp);
+    A += BK;     // move BK columns to right
+    B += BK * N; // move BK rows down
+
+    As_offset = 1 - As_offset;
+    Bs_offset = 1 - Bs_offset;
+    // swap the front and back barriers
+    auto tmp = frontBarrierPtr;
+    frontBarrierPtr = backBarrierPtr;
+    backBarrierPtr = tmp;
+
+    __syncthreads();
+  }
+
+  // compute the last blocktile
+  (*frontBarrierPtr).arrive_and_wait();
+  processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN, TM, TN>(
+      regM, regN, threadResults, As + As_offset * BM * BK,
+      Bs + Bs_offset * BK * BN, warpRow, warpCol, threadRowInWarp,
+      threadColInWarp);
+
+  // write out the results
+  for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
+    for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
+      // move C pointer to current warp subtile
+      float *C_interim = C + (wSubRowIdx * WSUBM) * N + wSubColIdx * WSUBN;
+      for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
+        for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) {
+          // load C vector into registers
+          float4 tmp = reinterpret_cast<float4 *>(
+              &C_interim[(threadRowInWarp * TM + resIdxM) * N +
+                         threadColInWarp * TN + resIdxN])[0];
+          // perform GEMM update in reg
+          const int i = (wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
+                        wSubColIdx * TN + resIdxN;
+          tmp.x = alpha * threadResults[i + 0] + beta * tmp.x;
+          tmp.y = alpha * threadResults[i + 1] + beta * tmp.y;
+          tmp.z = alpha * threadResults[i + 2] + beta * tmp.z;
+          tmp.w = alpha * threadResults[i + 3] + beta * tmp.w;
+          // write back
+          reinterpret_cast<float4 *>(
+              &C_interim[(threadRowInWarp * TM + resIdxM) * N +
+                         threadColInWarp * TN + resIdxN])[0] = tmp;
+        }
+      }
+    }
+  }
+}

upstream_ref/sgemm_cuda/1_naive.cuh

@@ -0,0 +1,29 @@
+#pragma once
+
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+/*
+
+Matrix sizes:
+MxK * KxN = MxN
+
+*/
+
+__global__ void sgemm_naive(int M, int N, int K, float alpha, const float *A,
+                            const float *B, float beta, float *C) {
+  const uint x = blockIdx.x * blockDim.x + threadIdx.x;
+  const uint y = blockIdx.y * blockDim.y + threadIdx.y;
+
+  // if statement is necessary to make things work under tile quantization
+  if (x < M && y < N) {
+    float tmp = 0.0;
+    for (int i = 0; i < K; ++i) {
+      tmp += A[x * K + i] * B[i * N + y];
+    }
+    // C = α*(A@B)+β*C
+    C[x * N + y] = alpha * tmp + beta * C[x * N + y];
+  }
+}

upstream_ref/sgemm_cuda/2_kernel_global_mem_coalesce.cuh

@@ -0,0 +1,24 @@
+#pragma once
+
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+template <const uint BLOCKSIZE>
+__global__ void sgemm_global_mem_coalesce(int M, int N, int K, float alpha,
+                                          const float *A, const float *B,
+                                          float beta, float *C) {
+  const int cRow = blockIdx.x * BLOCKSIZE + (threadIdx.x / BLOCKSIZE);
+  const int cCol = blockIdx.y * BLOCKSIZE + (threadIdx.x % BLOCKSIZE);
+
+  // if statement is necessary to make things work under tile quantization
+  if (cRow < M && cCol < N) {
+    float tmp = 0.0;
+    for (int i = 0; i < K; ++i) {
+      tmp += A[cRow * K + i] * B[i * N + cCol];
+    }
+    C[cRow * N + cCol] = alpha * tmp + beta * C[cRow * N + cCol];
+  }
+}

upstream_ref/sgemm_cuda/3_kernel_shared_mem_blocking.cuh

@@ -0,0 +1,57 @@
+#pragma once
+
+#include <algorithm>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+template <const int BLOCKSIZE>
+__global__ void sgemm_shared_mem_block(int M, int N, int K, float alpha,
+                                       const float *A, const float *B,
+                                       float beta, float *C) {
+  // the output block that we want to compute in this threadblock
+  const uint cRow = blockIdx.x;
+  const uint cCol = blockIdx.y;
+
+  // allocate buffer for current block in fast shared mem
+  // shared mem is shared between all threads in a block
+  __shared__ float As[BLOCKSIZE * BLOCKSIZE];
+  __shared__ float Bs[BLOCKSIZE * BLOCKSIZE];
+
+  // the inner row & col that we're accessing in this thread
+  const uint threadCol = threadIdx.x % BLOCKSIZE;
+  const uint threadRow = threadIdx.x / BLOCKSIZE;
+
+  // advance pointers to the starting positions
+  A += cRow * BLOCKSIZE * K;                    // row=cRow, col=0
+  B += cCol * BLOCKSIZE;                        // row=0, col=cCol
+  C += cRow * BLOCKSIZE * N + cCol * BLOCKSIZE; // row=cRow, col=cCol
+
+  float tmp = 0.0;
+  for (int bkIdx = 0; bkIdx < K; bkIdx += BLOCKSIZE) {
+    // Have each thread load one of the elements in A & B
+    // Make the threadCol (=threadIdx.x) the consecutive index
+    // to allow global memory access coalescing
+    As[threadRow * BLOCKSIZE + threadCol] = A[threadRow * K + threadCol];
+    Bs[threadRow * BLOCKSIZE + threadCol] = B[threadRow * N + threadCol];
+
+    // block threads in this block until cache is fully populated
+    __syncthreads();
+    A += BLOCKSIZE;
+    B += BLOCKSIZE * N;
+
+    // execute the dotproduct on the currently cached block
+    for (int dotIdx = 0; dotIdx < BLOCKSIZE; ++dotIdx) {
+      tmp += As[threadRow * BLOCKSIZE + dotIdx] *
+             Bs[dotIdx * BLOCKSIZE + threadCol];
+    }
+    // need to sync again at the end, to avoid faster threads
+    // fetching the next block into the cache before slower threads are done
+    __syncthreads();
+  }
+  C[threadRow * N + threadCol] =
+      alpha * tmp + beta * C[threadRow * N + threadCol];
+}

upstream_ref/sgemm_cuda/4_kernel_1D_blocktiling.cuh

@@ -0,0 +1,80 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+template <const int BM, const int BN, const int BK, const int TM>
+__global__ void sgemm1DBlocktiling(int M, int N, int K, float alpha,
+                                   const float *A, const float *B, float beta,
+                                   float *C) {
+  // If we flip x and y here we get ~30% less performance for large matrices.
+  // The current, 30% faster configuration ensures that blocks with sequential
+  // blockIDs access columns of B sequentially, while sharing the same row of A.
+  // The slower configuration would share columns of A, but access into B would
+  // be non-sequential. So the faster configuration has better spatial locality
+  // and hence a greater L2 hit rate.
+  const uint cRow = blockIdx.y;
+  const uint cCol = blockIdx.x;
+
+  // each warp will calculate 32*TM elements, with 32 being the columnar dim.
+  const int threadCol = threadIdx.x % BN;
+  const int threadRow = threadIdx.x / BN;
+
+  // allocate space for the current blocktile in SMEM
+  __shared__ float As[BM * BK];
+  __shared__ float Bs[BK * BN];
+
+  // Move blocktile to beginning of A's row and B's column
+  A += cRow * BM * K;
+  B += cCol * BN;
+  C += cRow * BM * N + cCol * BN;
+
+  // todo: adjust this to each thread to load multiple entries and
+  // better exploit the cache sizes
+  assert(BM * BK == blockDim.x);
+  assert(BN * BK == blockDim.x);
+  const uint innerColA = threadIdx.x % BK; // warp-level GMEM coalescing
+  const uint innerRowA = threadIdx.x / BK;
+  const uint innerColB = threadIdx.x % BN; // warp-level GMEM coalescing
+  const uint innerRowB = threadIdx.x / BN;
+
+  // allocate thread-local cache for results in registerfile
+  float threadResults[TM] = {0.0};
+
+  // outer loop over block tiles
+  for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
+    // populate the SMEM caches
+    As[innerRowA * BK + innerColA] = A[innerRowA * K + innerColA];
+    Bs[innerRowB * BN + innerColB] = B[innerRowB * N + innerColB];
+    __syncthreads();
+
+    // advance blocktile
+    A += BK;
+    B += BK * N;
+
+    // calculate per-thread results
+    for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
+      // we make the dotproduct loop the outside loop, which facilitates
+      // reuse of the Bs entry, which we can cache in a tmp var.
+      float tmpB = Bs[dotIdx * BN + threadCol];
+      for (uint resIdx = 0; resIdx < TM; ++resIdx) {
+        threadResults[resIdx] +=
+            As[(threadRow * TM + resIdx) * BK + dotIdx] * tmpB;
+      }
+    }
+    __syncthreads();
+  }
+
+  // write out the results
+  for (uint resIdx = 0; resIdx < TM; ++resIdx) {
+    C[(threadRow * TM + resIdx) * N + threadCol] =
+        alpha * threadResults[resIdx] +
+        beta * C[(threadRow * TM + resIdx) * N + threadCol];
+  }
+}

upstream_ref/sgemm_cuda/5_kernel_2D_blocktiling.cuh

@@ -0,0 +1,102 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+template <const int BM, const int BN, const int BK, const int TM, const int TN>
+__global__ void __launch_bounds__((BM * BN) / (TM * TN), 1)
+    sgemm2DBlocktiling(int M, int N, int K, float alpha, const float *A,
+                       const float *B, float beta, float *C) {
+  const uint cRow = blockIdx.y;
+  const uint cCol = blockIdx.x;
+
+  const uint totalResultsBlocktile = BM * BN;
+  // A thread is responsible for calculating TM*TN elements in the blocktile
+  const uint numThreadsBlocktile = totalResultsBlocktile / (TM * TN);
+
+  // ResultsPerBlock / ResultsPerThread == ThreadsPerBlock
+  assert(numThreadsBlocktile == blockDim.x);
+
+  // BN/TN are the number of threads to span a column
+  const int threadCol = threadIdx.x % (BN / TN);
+  const int threadRow = threadIdx.x / (BN / TN);
+
+  // allocate space for the current blocktile in smem
+  __shared__ float As[BM * BK];
+  __shared__ float Bs[BK * BN];
+
+  // Move blocktile to beginning of A's row and B's column
+  A += cRow * BM * K;
+  B += cCol * BN;
+  C += cRow * BM * N + cCol * BN;
+
+  // calculating the indices that this thread will load into SMEM
+  const uint innerRowA = threadIdx.x / BK;
+  const uint innerColA = threadIdx.x % BK;
+  // calculates the number of rows of As that are being loaded in a single step
+  // by a single block
+  const uint strideA = numThreadsBlocktile / BK;
+  const uint innerRowB = threadIdx.x / BN;
+  const uint innerColB = threadIdx.x % BN;
+  // for both As and Bs we want each load to span the full column-width, for
+  // better GMEM coalescing (as opposed to spanning full row-width and iterating
+  // across columns)
+  const uint strideB = numThreadsBlocktile / BN;
+
+  // allocate thread-local cache for results in registerfile
+  float threadResults[TM * TN] = {0.0};
+  // register caches for As and Bs
+  float regM[TM] = {0.0};
+  float regN[TN] = {0.0};
+
+  // outer-most loop over block tiles
+  for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
+    // populate the SMEM caches
+    for (uint loadOffset = 0; loadOffset < BM; loadOffset += strideA) {
+      As[(innerRowA + loadOffset) * BK + innerColA] =
+          A[(innerRowA + loadOffset) * K + innerColA];
+    }
+    for (uint loadOffset = 0; loadOffset < BK; loadOffset += strideB) {
+      Bs[(innerRowB + loadOffset) * BN + innerColB] =
+          B[(innerRowB + loadOffset) * N + innerColB];
+    }
+    __syncthreads();
+
+    // advance blocktile
+    A += BK;     // move BK columns to right
+    B += BK * N; // move BK rows down
+
+    // calculate per-thread results
+    for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
+      // block into registers
+      for (uint i = 0; i < TM; ++i) {
+        regM[i] = As[(threadRow * TM + i) * BK + dotIdx];
+      }
+      for (uint i = 0; i < TN; ++i) {
+        regN[i] = Bs[dotIdx * BN + threadCol * TN + i];
+      }
+      for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
+        for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
+          threadResults[resIdxM * TN + resIdxN] +=
+              regM[resIdxM] * regN[resIdxN];
+        }
+      }
+    }
+    __syncthreads();
+  }
+
+  // write out the results
+  for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
+    for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
+      C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN] =
+          alpha * threadResults[resIdxM * TN + resIdxN] +
+          beta * C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN];
+    }
+  }
+}

upstream_ref/sgemm_cuda/6_kernel_vectorize.cuh

@@ -0,0 +1,98 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+template <const int BM, const int BN, const int BK, const int TM, const int TN>
+__global__ void sgemmVectorize(int M, int N, int K, float alpha, float *A,
+                               float *B, float beta, float *C) {
+  const uint cRow = blockIdx.y;
+  const uint cCol = blockIdx.x;
+
+  // BN/TN are the number of threads to span a column
+  const int threadCol = threadIdx.x % (BN / TN);
+  const int threadRow = threadIdx.x / (BN / TN);
+
+  // allocate space for the current blocktile in smem
+  __shared__ float As[BM * BK];
+  __shared__ float Bs[BK * BN];
+
+  // Move blocktile to beginning of A's row and B's column
+  A += cRow * BM * K;
+  B += cCol * BN;
+  C += cRow * BM * N + cCol * BN;
+
+  // calculating the indices that this thread will load into SMEM
+  // we'll load 128bit / 32bit = 4 elements per thread at each step
+  const uint innerRowA = threadIdx.x / (BK / 4);
+  const uint innerColA = threadIdx.x % (BK / 4);
+  const uint innerRowB = threadIdx.x / (BN / 4);
+  const uint innerColB = threadIdx.x % (BN / 4);
+
+  // allocate thread-local cache for results in registerfile
+  float threadResults[TM * TN] = {0.0};
+  float regM[TM] = {0.0};
+  float regN[TN] = {0.0};
+
+  // outer-most loop over block tiles
+  for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
+    // populate the SMEM caches
+    // transpose A while loading it
+    float4 tmp =
+        reinterpret_cast<float4 *>(&A[innerRowA * K + innerColA * 4])[0];
+    As[(innerColA * 4 + 0) * BM + innerRowA] = tmp.x;
+    As[(innerColA * 4 + 1) * BM + innerRowA] = tmp.y;
+    As[(innerColA * 4 + 2) * BM + innerRowA] = tmp.z;
+    As[(innerColA * 4 + 3) * BM + innerRowA] = tmp.w;
+
+    reinterpret_cast<float4 *>(&Bs[innerRowB * BN + innerColB * 4])[0] =
+        reinterpret_cast<float4 *>(&B[innerRowB * N + innerColB * 4])[0];
+    __syncthreads();
+
+    // advance blocktile
+    A += BK;     // move BK columns to right
+    B += BK * N; // move BK rows down
+
+    // calculate per-thread results
+    for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
+      // block into registers
+      for (uint i = 0; i < TM; ++i) {
+        regM[i] = As[dotIdx * BM + threadRow * TM + i];
+      }
+      for (uint i = 0; i < TN; ++i) {
+        regN[i] = Bs[dotIdx * BN + threadCol * TN + i];
+      }
+      for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
+        for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
+          threadResults[resIdxM * TN + resIdxN] +=
+              regM[resIdxM] * regN[resIdxN];
+        }
+      }
+    }
+    __syncthreads();
+  }
+
+  // write out the results
+  for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
+    for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) {
+      // load C vector into registers
+      float4 tmp = reinterpret_cast<float4 *>(
+          &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0];
+      // perform GEMM update in reg
+      tmp.x = alpha * threadResults[resIdxM * TN + resIdxN] + beta * tmp.x;
+      tmp.y = alpha * threadResults[resIdxM * TN + resIdxN + 1] + beta * tmp.y;
+      tmp.z = alpha * threadResults[resIdxM * TN + resIdxN + 2] + beta * tmp.z;
+      tmp.w = alpha * threadResults[resIdxM * TN + resIdxN + 3] + beta * tmp.w;
+      // write back
+      reinterpret_cast<float4 *>(
+          &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0] =
+          tmp;
+    }
+  }
+}

upstream_ref/sgemm_cuda/7_kernel_resolve_bank_conflicts.cuh

@@ -0,0 +1,103 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+template <const int BM, const int BN, const int BK, const int TM, const int TN>
+__global__ void sgemmResolveBankConflicts(int M, int N, int K, float alpha,
+                                          float *A, float *B, float beta,
+                                          float *C) {
+  const uint cRow = blockIdx.y;
+  const uint cCol = blockIdx.x;
+
+  // BN/TN are the number of threads to span a column
+  const int threadCol = threadIdx.x % (BN / TN);
+  const int threadRow = threadIdx.x / (BN / TN);
+
+  // allocate space for the current blocktile in smem
+  __shared__ float As[BM * BK];
+  __shared__ float Bs[BK * BN];
+
+  // Move blocktile to beginning of A's row and B's column
+  A += cRow * BM * K;
+  B += cCol * BN;
+  C += cRow * BM * N + cCol * BN;
+
+  // calculating the indices that this thread will load into SMEM
+  // we'll load 128bit / 32bit = 4 elements per thread at each step
+  const uint innerRowA = threadIdx.x / (BK / 4);
+  const uint innerColA = threadIdx.x % (BK / 4);
+  const uint innerRowB = threadIdx.x / (BN / 4);
+  const uint innerColB = threadIdx.x % (BN / 4);
+
+  // allocate thread-local cache for results in registerfile
+  float threadResults[TM * TN] = {0.0};
+  float regM[TM] = {0.0};
+  float regN[TN] = {0.0};
+
+  // outer-most loop over block tiles
+  for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
+    // populate the SMEM caches
+    // transpose A while loading it
+    float4 tmp =
+        reinterpret_cast<float4 *>(&A[innerRowA * K + innerColA * 4])[0];
+    As[(innerColA * 4 + 0) * BM + innerRowA] = tmp.x;
+    As[(innerColA * 4 + 1) * BM + innerRowA] = tmp.y;
+    As[(innerColA * 4 + 2) * BM + innerRowA] = tmp.z;
+    As[(innerColA * 4 + 3) * BM + innerRowA] = tmp.w;
+
+    // "linearize" Bs while storing it
+    tmp = reinterpret_cast<float4 *>(&B[innerRowB * N + innerColB * 4])[0];
+    Bs[((innerColB % 2) * 4 + innerRowB * 8 + 0) * 16 + innerColB / 2] = tmp.x;
+    Bs[((innerColB % 2) * 4 + innerRowB * 8 + 1) * 16 + innerColB / 2] = tmp.y;
+    Bs[((innerColB % 2) * 4 + innerRowB * 8 + 2) * 16 + innerColB / 2] = tmp.z;
+    Bs[((innerColB % 2) * 4 + innerRowB * 8 + 3) * 16 + innerColB / 2] = tmp.w;
+    __syncthreads();
+
+    // advance blocktile
+    A += BK;     // move BK columns to right
+    B += BK * N; // move BK rows down
+
+    // calculate per-thread results
+    for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
+      // block into registers
+      for (uint i = 0; i < TM; ++i) {
+        regM[i] = As[dotIdx * BM + threadRow * TM + i];
+      }
+      for (uint i = 0; i < TN; ++i) {
+        regN[i] = Bs[(dotIdx * 8 + i) * 16 + threadCol];
+      }
+      for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
+        for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
+          threadResults[resIdxM * TN + resIdxN] +=
+              regM[resIdxM] * regN[resIdxN];
+        }
+      }
+    }
+    __syncthreads();
+  }
+
+  // write out the results
+  for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
+    for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) {
+      // load C vector into registers
+      float4 tmp = reinterpret_cast<float4 *>(
+          &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0];
+      // perform GEMM update in reg
+      tmp.x = alpha * threadResults[resIdxM * TN + resIdxN] + beta * tmp.x;
+      tmp.y = alpha * threadResults[resIdxM * TN + resIdxN + 1] + beta * tmp.y;
+      tmp.z = alpha * threadResults[resIdxM * TN + resIdxN + 2] + beta * tmp.z;
+      tmp.w = alpha * threadResults[resIdxM * TN + resIdxN + 3] + beta * tmp.w;
+      // write back
+      reinterpret_cast<float4 *>(
+          &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0] =
+          tmp;
+    }
+  }
+}

upstream_ref/sgemm_cuda/8_kernel_bank_extra_col.cuh

@@ -0,0 +1,103 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+template <const int BM, const int BN, const int BK, const int TM, const int TN>
+__global__ void sgemmResolveBankExtraCol(int M, int N, int K, float alpha,
+                                         float *A, float *B, float beta,
+                                         float *C) {
+  const uint cRow = blockIdx.y;
+  const uint cCol = blockIdx.x;
+
+  // BN/TN are the number of threads to span a column
+  const int threadCol = threadIdx.x % (BN / TN);
+  const int threadRow = threadIdx.x / (BN / TN);
+
+  // allocate space for the current blocktile in smem
+  __shared__ float As[BM * BK];
+  const int extraCols = 5;
+  __shared__ float Bs[BK * (BN + extraCols)];
+
+  // Move blocktile to beginning of A's row and B's column
+  A += cRow * BM * K;
+  B += cCol * BN;
+  C += cRow * BM * N + cCol * BN;
+
+  // calculating the indices that this thread will load into SMEM
+  // we'll load 128bit / 32bit = 4 elements per thread at each step
+  const uint innerRowA = threadIdx.x / (BK / 4);
+  const uint innerColA = threadIdx.x % (BK / 4);
+  const uint innerRowB = threadIdx.x / (BN / 4);
+  const uint innerColB = threadIdx.x % (BN / 4);
+
+  // allocate thread-local cache for results in registerfile
+  float threadResults[TM * TN] = {0.0};
+  float regM[TM] = {0.0};
+  float regN[TN] = {0.0};
+
+  // outer-most loop over block tiles
+  for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
+    // populate the SMEM caches
+    // transpose A while loading it
+    float4 tmp =
+        reinterpret_cast<float4 *>(&A[innerRowA * K + innerColA * 4])[0];
+    As[(innerColA * 4 + 0) * BM + innerRowA] = tmp.x;
+    As[(innerColA * 4 + 1) * BM + innerRowA] = tmp.y;
+    As[(innerColA * 4 + 2) * BM + innerRowA] = tmp.z;
+    As[(innerColA * 4 + 3) * BM + innerRowA] = tmp.w;
+
+    tmp = reinterpret_cast<float4 *>(&B[innerRowB * N + innerColB * 4])[0];
+    Bs[innerRowB * (BN + extraCols) + innerColB * 4 + 0] = tmp.x;
+    Bs[innerRowB * (BN + extraCols) + innerColB * 4 + 1] = tmp.y;
+    Bs[innerRowB * (BN + extraCols) + innerColB * 4 + 2] = tmp.z;
+    Bs[innerRowB * (BN + extraCols) + innerColB * 4 + 3] = tmp.w;
+    __syncthreads();
+
+    // advance blocktile
+    A += BK;     // move BK columns to right
+    B += BK * N; // move BK rows down
+
+    // calculate per-thread results
+    for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
+      // block into registers
+      for (uint i = 0; i < TM; ++i) {
+        regM[i] = As[dotIdx * BM + threadRow * TM + i];
+      }
+      for (uint i = 0; i < TN; ++i) {
+        regN[i] = Bs[dotIdx * (BN + extraCols) + threadCol * TN + i];
+      }
+      for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
+        for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
+          threadResults[resIdxM * TN + resIdxN] +=
+              regM[resIdxM] * regN[resIdxN];
+        }
+      }
+    }
+    __syncthreads();
+  }
+
+  // write out the results
+  for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
+    for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) {
+      // load C vector into registers
+      float4 tmp = reinterpret_cast<float4 *>(
+          &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0];
+      // perform GEMM update in reg
+      tmp.x = alpha * threadResults[resIdxM * TN + resIdxN] + beta * tmp.x;
+      tmp.y = alpha * threadResults[resIdxM * TN + resIdxN + 1] + beta * tmp.y;
+      tmp.z = alpha * threadResults[resIdxM * TN + resIdxN + 2] + beta * tmp.z;
+      tmp.w = alpha * threadResults[resIdxM * TN + resIdxN + 3] + beta * tmp.w;
+      // write back
+      reinterpret_cast<float4 *>(
+          &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0] =
+          tmp;
+    }
+  }
+}

upstream_ref/sgemm_cuda/9_kernel_autotuned.cuh

@@ -0,0 +1,127 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+const int K9_NUM_THREADS = 256;
+
+template <const int BM, const int BN, const int BK, const int TM, const int TN>
+__global__ void __launch_bounds__(K9_NUM_THREADS)
+    sgemmAutotuned(int M, int N, int K, float alpha, float *A, float *B,
+                   float beta, float *C) {
+  const uint cRow = blockIdx.y;
+  const uint cCol = blockIdx.x;
+
+  // size of warptile
+  constexpr int WM = TM * 16;
+  constexpr int WN = TN * 16;
+  // iterations of warptile
+  constexpr int WMITER = CEIL_DIV(BM, WM);
+  constexpr int WNITER = CEIL_DIV(BN, WN);
+
+  // Placement of the thread in the warptile
+  const int threadCol = threadIdx.x % (WN / TN);
+  const int threadRow = threadIdx.x / (WN / TN);
+
+  // allocate space for the current blocktile in smem
+  __shared__ float As[BM * BK];
+  __shared__ float Bs[BK * BN];
+
+  // Move blocktile to beginning of A's row and B's column
+  A += cRow * BM * K;
+  B += cCol * BN;
+  C += cRow * BM * N + cCol * BN;
+
+  // calculating the indices that this thread will load into SMEM
+  // we'll load 128bit / 32bit = 4 elements per thread at each step
+  const uint innerRowA = threadIdx.x / (BK / 4);
+  const uint innerColA = threadIdx.x % (BK / 4);
+  constexpr uint rowStrideA = (K9_NUM_THREADS * 4) / BK;
+  const uint innerRowB = threadIdx.x / (BN / 4);
+  const uint innerColB = threadIdx.x % (BN / 4);
+  constexpr uint rowStrideB = K9_NUM_THREADS / (BN / 4);
+
+  // allocate thread-local cache for results in registerfile
+  float threadResults[WMITER * WNITER * TM * TN] = {0.0};
+  float regM[TM] = {0.0};
+  float regN[TN] = {0.0};
+
+  // outer-most loop over block tiles
+  for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
+    // populate the SMEM caches
+    for (uint offset = 0; offset + rowStrideA <= BM; offset += rowStrideA) {
+      float4 tmp = reinterpret_cast<float4 *>(
+          &A[(innerRowA + offset) * K + innerColA * 4])[0];
+      // transpose A while storing it
+      As[(innerColA * 4 + 0) * BM + innerRowA + offset] = tmp.x;
+      As[(innerColA * 4 + 1) * BM + innerRowA + offset] = tmp.y;
+      As[(innerColA * 4 + 2) * BM + innerRowA + offset] = tmp.z;
+      As[(innerColA * 4 + 3) * BM + innerRowA + offset] = tmp.w;
+    }
+
+    for (uint offset = 0; offset + rowStrideB <= BK; offset += rowStrideB) {
+      reinterpret_cast<float4 *>(
+          &Bs[(innerRowB + offset) * BN + innerColB * 4])[0] =
+          reinterpret_cast<float4 *>(
+              &B[(innerRowB + offset) * N + innerColB * 4])[0];
+    }
+    __syncthreads();
+
+    for (uint wmIdx = 0; wmIdx < WMITER; ++wmIdx) {
+      for (uint wnIdx = 0; wnIdx < WNITER; ++wnIdx) {
+        // calculate per-thread results
+        for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
+          // block into registers
+          for (uint i = 0; i < TM; ++i) {
+            regM[i] = As[dotIdx * BM + (wmIdx * WM) + threadRow * TM + i];
+          }
+          for (uint i = 0; i < TN; ++i) {
+            regN[i] = Bs[dotIdx * BN + (wnIdx * WN) + threadCol * TN + i];
+          }
+          for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
+            for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
+              threadResults[(wmIdx * TM + resIdxM) * (WNITER * TN) +
+                            wnIdx * TN + resIdxN] +=
+                  regM[resIdxM] * regN[resIdxN];
+            }
+          }
+        }
+      }
+    }
+    __syncthreads();
+    // advance blocktile
+    A += BK;     // move BK columns to right
+    B += BK * N; // move BK rows down
+  }
+
+  // write out the results
+  for (uint wmIdx = 0; wmIdx < WMITER; ++wmIdx) {
+    for (uint wnIdx = 0; wnIdx < WNITER; ++wnIdx) {
+      float *C_interim = C + (wmIdx * WM * N) + (wnIdx * WN);
+      for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
+        for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) {
+          // load C vector into registers
+          float4 tmp = reinterpret_cast<float4 *>(
+              &C_interim[(threadRow * TM + resIdxM) * N + threadCol * TN +
+                         resIdxN])[0];
+          // perform GEMM update in reg
+          const int i =
+              (wmIdx * TM + resIdxM) * (WNITER * TN) + wnIdx * TN + resIdxN;
+          tmp.x = alpha * threadResults[i + 0] + beta * tmp.x;
+          tmp.y = alpha * threadResults[i + 1] + beta * tmp.y;
+          tmp.z = alpha * threadResults[i + 2] + beta * tmp.z;
+          tmp.w = alpha * threadResults[i + 3] + beta * tmp.w;
+          // write back
+          reinterpret_cast<float4 *>(&C_interim[(threadRow * TM + resIdxM) * N +
+                                                threadCol * TN + resIdxN])[0] =
+              tmp;
+        }
+      }
+    }
+  }
+}

upstream_ref/sgemm_edtallison/01_naive.cuh

@@ -0,0 +1,37 @@
+# pragma once
+
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+/*
+
+Matrix sizes:
+MxK * KxN = MxN
+
+*/
+
+__global__ void sgemm_naive(
+    int M, int N, int K, // sizes
+    float alpha, const float *A, const float *B, float beta, float *C // pointers used to point to matrices
+) {
+    // compute position in C that this thread is responsible for
+    // "which block" * "width of block" to get to start of block + "which thread"
+    const uint x = blockIdx.x * blockDim.x + threadIdx.x; // "which row?" (inverted from graphical intuition, confusingly)
+    const uint y = blockIdx.y * blockDim.y + threadIdx.y; // "which column?"
+
+    // if M or N are not multiples of 32, there will be "extra"/"remainder" threads on the last block in x/y.
+    // we don't want those leftover threads to do anything (tile quantisation)
+    if (x < M && y < N) {
+        float tmp = 0.0;
+        for (int i = 0; i < K; ++i) { // K is the size of the row in A, col in B i.e. the dot product
+            // A: x * K gives the start of relevant row, i enumerates across the row (col by col)
+            // B: y gives the relevant column, i * N enumerates down the column, (row by row)
+            tmp += A[x * K + i] * B[i * N + y];
+        }
+        // C = alpha*(A@B) + beta*C
+        // x * N takes to start of relevant row, y moves across to the relevant column
+        C[x * N + y] = alpha * tmp + beta * C[x * N + y];
+    }
+}

upstream_ref/sgemm_edtallison/02_kernel_global_mem_coalesce.cuh

@@ -0,0 +1,25 @@
+#pragma once
+
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+template <const uint BLOCKSIZE>
+// __global__ is used to specify that the function is run on GPU, called by host (CPU)
+__global__ void sgemm_global_mem_coalesce(int M, int N, int K, float alpha,
+                                          const float *A, const float *B,
+                                          float beta, float *C) {
+    const int cRow = blockIdx.x * BLOCKSIZE + (threadIdx.x / BLOCKSIZE); // note that blockDim is now 1-dimensional
+    const int cCol = blockIdx.y * BLOCKSIZE + (threadIdx.x % BLOCKSIZE);
+
+    // if statement is necessary to make things work under tile quantization
+    if (cRow < M && cCol < N) {
+        float tmp = 0.0;
+        for (int i = 0; i < K; ++i) {
+            tmp += A[cRow * K + i] * B[i * N + cCol];
+        }
+        C[cRow * N + cCol] = alpha * tmp + beta * C[cRow * N + cCol];
+    }
+}

upstream_ref/sgemm_edtallison/03_kernel_shared_mem_blocking.cuh

@@ -0,0 +1,54 @@
+#pragma once
+
+#include <algorithm>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+template <const int BLOCKSIZE>
+__global__ void sgemm_shared_mem_block(int M, int N, int K, float alpha,
+                                       const float *A, const float *B,
+                                       float beta, float *C) {
+    // output C block we want to compute with this threadBlock
+    const uint cRow = blockIdx.x;
+    const uint cCol = blockIdx.y;
+
+    // allocate buffer for current block in fast SMEM (shared between all threads in block)
+    __shared__ float As[BLOCKSIZE * BLOCKSIZE];
+    __shared__ float Bs[BLOCKSIZE * BLOCKSIZE];
+
+    // the inner row and col that we are accessing in this specific thread
+    const uint threadRow = threadIdx.x / BLOCKSIZE; // note similarity to previous kernel
+    const uint threadCol = threadIdx.x % BLOCKSIZE;
+
+    // advance pointers to the starting positions (they are input as pointers to first elements in the matrices)
+    A += cRow * BLOCKSIZE * K;                    // row=cRow, col=0 (the start of the relevant row)
+    B += cCol * BLOCKSIZE;                        // row=0, col=cCol (top of relevant col)
+    C += cRow * BLOCKSIZE * N + cCol * BLOCKSIZE; // row=cRow, col=cCol
+
+    float tmp = 0.0;
+    for (int bkIdx=0; bkIdx < K; bkIdx+=BLOCKSIZE) { // shifting the whole block along the row of A and col of B
+        // have each thread load one of the elements in A and B
+        // make the threadCol (=threadIdx.x) the consecutive index
+        // to allow GMEM access coalescing
+        As[threadRow * BLOCKSIZE + threadCol] = A[threadRow * K + threadCol];
+        Bs[threadRow * BLOCKSIZE + threadCol] = B[threadRow * N + threadCol];
+
+        // ensure cache is fully populated
+        __syncthreads();
+        A += BLOCKSIZE; // for next iteration
+        B += BLOCKSIZE * N;
+
+        // execute the dotproduct on the currently cached block
+        for (int dotIdx = 0; dotIdx < BLOCKSIZE; ++dotIdx) {
+            tmp += As[threadRow * BLOCKSIZE + dotIdx] *
+                   Bs[dotIdx * BLOCKSIZE + threadCol];
+        }
+        // sync so faster threads don't fetch the next block into cache
+        _syncthreads();
+    }
+    C[threadRow * N + threadCol] = alpha * tmp + beta * C[threadRow * N + threadCol];
+}

upstream_ref/sgemm_edtallison/04_kernel_1D_blocktiling.cuh

@@ -0,0 +1,80 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+template <const int BM, const int BN, const int BK, const int TM>
+__global__ void sgemm1DBlocktiling(int M, int N, int K, float alpha,
+                                   const float *A, const float *B, float beta,
+                                   float *C) {
+  // If we flip x and y here we get ~30% less performance for large matrices.
+  // The current, 30% faster configuration ensures that blocks with sequential
+  // blockIDs access columns of B sequentially, while sharing the same row of A.
+  // The slower configuration would share columns of A, but access into B would
+  // be non-sequential. So the faster configuration has better spatial locality
+  // and hence a greater L2 hit rate.
+  const uint cRow = blockIdx.y;
+  const uint cCol = blockIdx.x;
+
+  // each warp will calculate 32*TM elements, with 32 being the columnar dim.
+  const int threadCol = threadIdx.x % BN;
+  const int threadRow = threadIdx.x / BN;
+
+  // allocate space for the current blocktile in SMEM
+  __shared__ float As[BM * BK];
+  __shared__ float Bs[BK * BN];
+
+  // Move blocktile to beginning of A's row and B's column
+  A += cRow * BM * K;
+  B += cCol * BN;
+  C += cRow * BM * N + cCol * BN;
+
+  // todo: adjust this to each thread to load multiple entries and
+  // better exploit the cache sizes
+  assert(BM * BK == blockDim.x);
+  assert(BN * BK == blockDim.x);
+  const uint innerColA = threadIdx.x % BK; // warp-level GMEM coalescing
+  const uint innerRowA = threadIdx.x / BK;
+  const uint innerColB = threadIdx.x % BN; // warp-level GMEM coalescing
+  const uint innerRowB = threadIdx.x / BN;
+
+  // allocate thread-local cache for results in registerfile
+  float threadResults[TM] = {0.0};
+
+  // outer loop over block tiles
+  for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
+    // populate the SMEM caches
+    As[innerRowA * BK + innerColA] = A[innerRowA * K + innerColA];
+    Bs[innerRowB * BN + innerColB] = B[innerRowB * N + innerColB];
+    __syncthreads();
+
+    // advance blocktile
+    A += BK;
+    B += BK * N;
+
+    // calculate per-thread results
+    for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
+      // we make the dotproduct loop the outside loop, which facilitates
+      // reuse of the Bs entry, which we can cache in a tmp var.
+      float tmpB = Bs[dotIdx * BN + threadCol];
+      for (uint resIdx = 0; resIdx < TM; ++resIdx) {
+        threadResults[resIdx] +=
+            As[(threadRow * TM + resIdx) * BK + dotIdx] * tmpB;
+      }
+    }
+    __syncthreads();
+  }
+
+  // write out the results
+  for (uint resIdx = 0; resIdx < TM; ++resIdx) {
+    C[(threadRow * TM + resIdx) * N + threadCol] =
+        alpha * threadResults[resIdx] +
+        beta * C[(threadRow * TM + resIdx) * N + threadCol];
+  }
+}

upstream_ref/sgemm_edtallison/05_kernel_2D_blocktiling.cuh

@@ -0,0 +1,102 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+template <const int BM, const int BN, const int BK, const int TM, const int TN>
+__global__ void __launch_bounds__((BM * BN) / (TM * TN), 1)
+    sgemm2DBlocktiling(int M, int N, int K, float alpha, const float *A,
+                       const float *B, float beta, float *C) {
+  const uint cRow = blockIdx.y;
+  const uint cCol = blockIdx.x;
+
+  const uint totalResultsBlocktile = BM * BN;
+  // A thread is responsible for calculating TM*TN elements in the blocktile
+  const uint numThreadsBlocktile = totalResultsBlocktile / (TM * TN);
+
+  // ResultsPerBlock / ResultsPerThread == ThreadsPerBlock
+  assert(numThreadsBlocktile == blockDim.x);
+
+  // BN/TN are the number of threads to span a column
+  const int threadCol = threadIdx.x % (BN / TN);
+  const int threadRow = threadIdx.x / (BN / TN);
+
+  // allocate space for the current blocktile in smem
+  __shared__ float As[BM * BK];
+  __shared__ float Bs[BK * BN];
+
+  // Move blocktile to beginning of A's row and B's column
+  A += cRow * BM * K;
+  B += cCol * BN;
+  C += cRow * BM * N + cCol * BN;
+
+  // calculating the indices that this thread will load into SMEM
+  const uint innerRowA = threadIdx.x / BK;
+  const uint innerColA = threadIdx.x % BK;
+  // calculates the number of rows of As that are being loaded in a single step
+  // by a single block
+  const uint strideA = numThreadsBlocktile / BK;
+  const uint innerRowB = threadIdx.x / BN;
+  const uint innerColB = threadIdx.x % BN;
+  // for both As and Bs we want each load to span the full column-width, for
+  // better GMEM coalescing (as opposed to spanning full row-width and iterating
+  // across columns)
+  const uint strideB = numThreadsBlocktile / BN;
+
+  // allocate thread-local cache for results in registerfile
+  float threadResults[TM * TN] = {0.0};
+  // register caches for As and Bs
+  float regM[TM] = {0.0};
+  float regN[TN] = {0.0};
+
+  // outer-most loop over block tiles
+  for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
+    // populate the SMEM caches
+    for (uint loadOffset = 0; loadOffset < BM; loadOffset += strideA) {
+      As[(innerRowA + loadOffset) * BK + innerColA] =
+          A[(innerRowA + loadOffset) * K + innerColA];
+    }
+    for (uint loadOffset = 0; loadOffset < BK; loadOffset += strideB) {
+      Bs[(innerRowB + loadOffset) * BN + innerColB] =
+          B[(innerRowB + loadOffset) * N + innerColB];
+    }
+    __syncthreads();
+
+    // advance blocktile
+    A += BK;     // move BK columns to right
+    B += BK * N; // move BK rows down
+
+    // calculate per-thread results
+    for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
+      // block into registers
+      for (uint i = 0; i < TM; ++i) {
+        regM[i] = As[(threadRow * TM + i) * BK + dotIdx];
+      }
+      for (uint i = 0; i < TN; ++i) {
+        regN[i] = Bs[dotIdx * BN + threadCol * TN + i];
+      }
+      for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
+        for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
+          threadResults[resIdxM * TN + resIdxN] +=
+              regM[resIdxM] * regN[resIdxN];
+        }
+      }
+    }
+    __syncthreads();
+  }
+
+  // write out the results
+  for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
+    for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
+      C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN] =
+          alpha * threadResults[resIdxM * TN + resIdxN] +
+          beta * C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN];
+    }
+  }
+}

upstream_ref/sgemm_edtallison/06_kernel_vectorize.cuh

@@ -0,0 +1,98 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+template <const int BM, const int BN, const int BK, const int TM, const int TN>
+__global__ void sgemmVectorize(int M, int N, int K, float alpha, float *A,
+                               float *B, float beta, float *C) {
+  const uint cRow = blockIdx.y;
+  const uint cCol = blockIdx.x;
+
+  // BN/TN are the number of threads to span a column
+  const int threadCol = threadIdx.x % (BN / TN);
+  const int threadRow = threadIdx.x / (BN / TN);
+
+  // allocate space for the current blocktile in smem
+  __shared__ float As[BM * BK];
+  __shared__ float Bs[BK * BN];
+
+  // Move blocktile to beginning of A's row and B's column
+  A += cRow * BM * K;
+  B += cCol * BN;
+  C += cRow * BM * N + cCol * BN;
+
+  // calculating the indices that this thread will load into SMEM
+  // we'll load 128bit / 32bit = 4 elements per thread at each step
+  const uint innerRowA = threadIdx.x / (BK / 4);
+  const uint innerColA = threadIdx.x % (BK / 4);
+  const uint innerRowB = threadIdx.x / (BN / 4);
+  const uint innerColB = threadIdx.x % (BN / 4);
+
+  // allocate thread-local cache for results in registerfile
+  float threadResults[TM * TN] = {0.0};
+  float regM[TM] = {0.0};
+  float regN[TN] = {0.0};
+
+  // outer-most loop over block tiles
+  for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
+    // populate the SMEM caches
+    // transpose A while loading it
+    float4 tmp =
+        reinterpret_cast<float4 *>(&A[innerRowA * K + innerColA * 4])[0];
+    As[(innerColA * 4 + 0) * BM + innerRowA] = tmp.x;
+    As[(innerColA * 4 + 1) * BM + innerRowA] = tmp.y;
+    As[(innerColA * 4 + 2) * BM + innerRowA] = tmp.z;
+    As[(innerColA * 4 + 3) * BM + innerRowA] = tmp.w;
+
+    reinterpret_cast<float4 *>(&Bs[innerRowB * BN + innerColB * 4])[0] =
+        reinterpret_cast<float4 *>(&B[innerRowB * N + innerColB * 4])[0];
+    __syncthreads();
+
+    // advance blocktile
+    A += BK;     // move BK columns to right
+    B += BK * N; // move BK rows down
+
+    // calculate per-thread results
+    for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
+      // block into registers
+      for (uint i = 0; i < TM; ++i) {
+        regM[i] = As[dotIdx * BM + threadRow * TM + i];
+      }
+      for (uint i = 0; i < TN; ++i) {
+        regN[i] = Bs[dotIdx * BN + threadCol * TN + i];
+      }
+      for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
+        for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
+          threadResults[resIdxM * TN + resIdxN] +=
+              regM[resIdxM] * regN[resIdxN];
+        }
+      }
+    }
+    __syncthreads();
+  }
+
+  // write out the results
+  for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
+    for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) {
+      // load C vector into registers
+      float4 tmp = reinterpret_cast<float4 *>(
+          &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0];
+      // perform GEMM update in reg
+      tmp.x = alpha * threadResults[resIdxM * TN + resIdxN] + beta * tmp.x;
+      tmp.y = alpha * threadResults[resIdxM * TN + resIdxN + 1] + beta * tmp.y;
+      tmp.z = alpha * threadResults[resIdxM * TN + resIdxN + 2] + beta * tmp.z;
+      tmp.w = alpha * threadResults[resIdxM * TN + resIdxN + 3] + beta * tmp.w;
+      // write back
+      reinterpret_cast<float4 *>(
+          &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0] =
+          tmp;
+    }
+  }
+}

upstream_ref/sgemm_edtallison/07_kernel_resolve_bank_conflicts.cuh

@@ -0,0 +1,103 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+template <const int BM, const int BN, const int BK, const int TM, const int TN>
+__global__ void sgemmResolveBankConflicts(int M, int N, int K, float alpha,
+                                          float *A, float *B, float beta,
+                                          float *C) {
+  const uint cRow = blockIdx.y;
+  const uint cCol = blockIdx.x;
+
+  // BN/TN are the number of threads to span a column
+  const int threadCol = threadIdx.x % (BN / TN);
+  const int threadRow = threadIdx.x / (BN / TN);
+
+  // allocate space for the current blocktile in smem
+  __shared__ float As[BM * BK];
+  __shared__ float Bs[BK * BN];
+
+  // Move blocktile to beginning of A's row and B's column
+  A += cRow * BM * K;
+  B += cCol * BN;
+  C += cRow * BM * N + cCol * BN;
+
+  // calculating the indices that this thread will load into SMEM
+  // we'll load 128bit / 32bit = 4 elements per thread at each step
+  const uint innerRowA = threadIdx.x / (BK / 4);
+  const uint innerColA = threadIdx.x % (BK / 4);
+  const uint innerRowB = threadIdx.x / (BN / 4);
+  const uint innerColB = threadIdx.x % (BN / 4);
+
+  // allocate thread-local cache for results in registerfile
+  float threadResults[TM * TN] = {0.0};
+  float regM[TM] = {0.0};
+  float regN[TN] = {0.0};
+
+  // outer-most loop over block tiles
+  for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
+    // populate the SMEM caches
+    // transpose A while loading it
+    float4 tmp =
+        reinterpret_cast<float4 *>(&A[innerRowA * K + innerColA * 4])[0];
+    As[(innerColA * 4 + 0) * BM + innerRowA] = tmp.x;
+    As[(innerColA * 4 + 1) * BM + innerRowA] = tmp.y;
+    As[(innerColA * 4 + 2) * BM + innerRowA] = tmp.z;
+    As[(innerColA * 4 + 3) * BM + innerRowA] = tmp.w;
+
+    // "linearize" Bs while storing it
+    tmp = reinterpret_cast<float4 *>(&B[innerRowB * N + innerColB * 4])[0];
+    Bs[((innerColB % 2) * 4 + innerRowB * 8 + 0) * 16 + innerColB / 2] = tmp.x;
+    Bs[((innerColB % 2) * 4 + innerRowB * 8 + 1) * 16 + innerColB / 2] = tmp.y;
+    Bs[((innerColB % 2) * 4 + innerRowB * 8 + 2) * 16 + innerColB / 2] = tmp.z;
+    Bs[((innerColB % 2) * 4 + innerRowB * 8 + 3) * 16 + innerColB / 2] = tmp.w;
+    __syncthreads();
+
+    // advance blocktile
+    A += BK;     // move BK columns to right
+    B += BK * N; // move BK rows down
+
+    // calculate per-thread results
+    for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
+      // block into registers
+      for (uint i = 0; i < TM; ++i) {
+        regM[i] = As[dotIdx * BM + threadRow * TM + i];
+      }
+      for (uint i = 0; i < TN; ++i) {
+        regN[i] = Bs[(dotIdx * 8 + i) * 16 + threadCol];
+      }
+      for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
+        for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
+          threadResults[resIdxM * TN + resIdxN] +=
+              regM[resIdxM] * regN[resIdxN];
+        }
+      }
+    }
+    __syncthreads();
+  }
+
+  // write out the results
+  for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
+    for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) {
+      // load C vector into registers
+      float4 tmp = reinterpret_cast<float4 *>(
+          &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0];
+      // perform GEMM update in reg
+      tmp.x = alpha * threadResults[resIdxM * TN + resIdxN] + beta * tmp.x;
+      tmp.y = alpha * threadResults[resIdxM * TN + resIdxN + 1] + beta * tmp.y;
+      tmp.z = alpha * threadResults[resIdxM * TN + resIdxN + 2] + beta * tmp.z;
+      tmp.w = alpha * threadResults[resIdxM * TN + resIdxN + 3] + beta * tmp.w;
+      // write back
+      reinterpret_cast<float4 *>(
+          &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0] =
+          tmp;
+    }
+  }
+}

upstream_ref/sgemm_edtallison/08_kernel_bank_extra_col.cuh

@@ -0,0 +1,103 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+template <const int BM, const int BN, const int BK, const int TM, const int TN>
+__global__ void sgemmResolveBankExtraCol(int M, int N, int K, float alpha,
+                                         float *A, float *B, float beta,
+                                         float *C) {
+  const uint cRow = blockIdx.y;
+  const uint cCol = blockIdx.x;
+
+  // BN/TN are the number of threads to span a column
+  const int threadCol = threadIdx.x % (BN / TN);
+  const int threadRow = threadIdx.x / (BN / TN);
+
+  // allocate space for the current blocktile in smem
+  __shared__ float As[BM * BK];
+  const int extraCols = 5;
+  __shared__ float Bs[BK * (BN + extraCols)];
+
+  // Move blocktile to beginning of A's row and B's column
+  A += cRow * BM * K;
+  B += cCol * BN;
+  C += cRow * BM * N + cCol * BN;
+
+  // calculating the indices that this thread will load into SMEM
+  // we'll load 128bit / 32bit = 4 elements per thread at each step
+  const uint innerRowA = threadIdx.x / (BK / 4);
+  const uint innerColA = threadIdx.x % (BK / 4);
+  const uint innerRowB = threadIdx.x / (BN / 4);
+  const uint innerColB = threadIdx.x % (BN / 4);
+
+  // allocate thread-local cache for results in registerfile
+  float threadResults[TM * TN] = {0.0};
+  float regM[TM] = {0.0};
+  float regN[TN] = {0.0};
+
+  // outer-most loop over block tiles
+  for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
+    // populate the SMEM caches
+    // transpose A while loading it
+    float4 tmp =
+        reinterpret_cast<float4 *>(&A[innerRowA * K + innerColA * 4])[0];
+    As[(innerColA * 4 + 0) * BM + innerRowA] = tmp.x;
+    As[(innerColA * 4 + 1) * BM + innerRowA] = tmp.y;
+    As[(innerColA * 4 + 2) * BM + innerRowA] = tmp.z;
+    As[(innerColA * 4 + 3) * BM + innerRowA] = tmp.w;
+
+    tmp = reinterpret_cast<float4 *>(&B[innerRowB * N + innerColB * 4])[0];
+    Bs[innerRowB * (BN + extraCols) + innerColB * 4 + 0] = tmp.x;
+    Bs[innerRowB * (BN + extraCols) + innerColB * 4 + 1] = tmp.y;
+    Bs[innerRowB * (BN + extraCols) + innerColB * 4 + 2] = tmp.z;
+    Bs[innerRowB * (BN + extraCols) + innerColB * 4 + 3] = tmp.w;
+    __syncthreads();
+
+    // advance blocktile
+    A += BK;     // move BK columns to right
+    B += BK * N; // move BK rows down
+
+    // calculate per-thread results
+    for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
+      // block into registers
+      for (uint i = 0; i < TM; ++i) {
+        regM[i] = As[dotIdx * BM + threadRow * TM + i];
+      }
+      for (uint i = 0; i < TN; ++i) {
+        regN[i] = Bs[dotIdx * (BN + extraCols) + threadCol * TN + i];
+      }
+      for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
+        for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
+          threadResults[resIdxM * TN + resIdxN] +=
+              regM[resIdxM] * regN[resIdxN];
+        }
+      }
+    }
+    __syncthreads();
+  }
+
+  // write out the results
+  for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
+    for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) {
+      // load C vector into registers
+      float4 tmp = reinterpret_cast<float4 *>(
+          &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0];
+      // perform GEMM update in reg
+      tmp.x = alpha * threadResults[resIdxM * TN + resIdxN] + beta * tmp.x;
+      tmp.y = alpha * threadResults[resIdxM * TN + resIdxN + 1] + beta * tmp.y;
+      tmp.z = alpha * threadResults[resIdxM * TN + resIdxN + 2] + beta * tmp.z;
+      tmp.w = alpha * threadResults[resIdxM * TN + resIdxN + 3] + beta * tmp.w;
+      // write back
+      reinterpret_cast<float4 *>(
+          &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0] =
+          tmp;
+    }
+  }
+}

upstream_ref/sgemm_edtallison/09_kernel_autotuned.cuh

@@ -0,0 +1,127 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+const int K9_NUM_THREADS = 256;
+
+template <const int BM, const int BN, const int BK, const int TM, const int TN>
+__global__ void __launch_bounds__(K9_NUM_THREADS)
+    sgemmAutotuned(int M, int N, int K, float alpha, float *A, float *B,
+                   float beta, float *C) {
+  const uint cRow = blockIdx.y;
+  const uint cCol = blockIdx.x;
+
+  // size of warptile
+  constexpr int WM = TM * 16;
+  constexpr int WN = TN * 16;
+  // iterations of warptile
+  constexpr int WMITER = CEIL_DIV(BM, WM);
+  constexpr int WNITER = CEIL_DIV(BN, WN);
+
+  // Placement of the thread in the warptile
+  const int threadCol = threadIdx.x % (WN / TN);
+  const int threadRow = threadIdx.x / (WN / TN);
+
+  // allocate space for the current blocktile in smem
+  __shared__ float As[BM * BK];
+  __shared__ float Bs[BK * BN];
+
+  // Move blocktile to beginning of A's row and B's column
+  A += cRow * BM * K;
+  B += cCol * BN;
+  C += cRow * BM * N + cCol * BN;
+
+  // calculating the indices that this thread will load into SMEM
+  // we'll load 128bit / 32bit = 4 elements per thread at each step
+  const uint innerRowA = threadIdx.x / (BK / 4);
+  const uint innerColA = threadIdx.x % (BK / 4);
+  constexpr uint rowStrideA = (K9_NUM_THREADS * 4) / BK;
+  const uint innerRowB = threadIdx.x / (BN / 4);
+  const uint innerColB = threadIdx.x % (BN / 4);
+  constexpr uint rowStrideB = K9_NUM_THREADS / (BN / 4);
+
+  // allocate thread-local cache for results in registerfile
+  float threadResults[WMITER * WNITER * TM * TN] = {0.0};
+  float regM[TM] = {0.0};
+  float regN[TN] = {0.0};
+
+  // outer-most loop over block tiles
+  for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
+    // populate the SMEM caches
+    for (uint offset = 0; offset + rowStrideA <= BM; offset += rowStrideA) {
+      float4 tmp = reinterpret_cast<float4 *>(
+          &A[(innerRowA + offset) * K + innerColA * 4])[0];
+      // transpose A while storing it
+      As[(innerColA * 4 + 0) * BM + innerRowA + offset] = tmp.x;
+      As[(innerColA * 4 + 1) * BM + innerRowA + offset] = tmp.y;
+      As[(innerColA * 4 + 2) * BM + innerRowA + offset] = tmp.z;
+      As[(innerColA * 4 + 3) * BM + innerRowA + offset] = tmp.w;
+    }
+
+    for (uint offset = 0; offset + rowStrideB <= BK; offset += rowStrideB) {
+      reinterpret_cast<float4 *>(
+          &Bs[(innerRowB + offset) * BN + innerColB * 4])[0] =
+          reinterpret_cast<float4 *>(
+              &B[(innerRowB + offset) * N + innerColB * 4])[0];
+    }
+    __syncthreads();
+
+    for (uint wmIdx = 0; wmIdx < WMITER; ++wmIdx) {
+      for (uint wnIdx = 0; wnIdx < WNITER; ++wnIdx) {
+        // calculate per-thread results
+        for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
+          // block into registers
+          for (uint i = 0; i < TM; ++i) {
+            regM[i] = As[dotIdx * BM + (wmIdx * WM) + threadRow * TM + i];
+          }
+          for (uint i = 0; i < TN; ++i) {
+            regN[i] = Bs[dotIdx * BN + (wnIdx * WN) + threadCol * TN + i];
+          }
+          for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
+            for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
+              threadResults[(wmIdx * TM + resIdxM) * (WNITER * TN) +
+                            wnIdx * TN + resIdxN] +=
+                  regM[resIdxM] * regN[resIdxN];
+            }
+          }
+        }
+      }
+    }
+    __syncthreads();
+    // advance blocktile
+    A += BK;     // move BK columns to right
+    B += BK * N; // move BK rows down
+  }
+
+  // write out the results
+  for (uint wmIdx = 0; wmIdx < WMITER; ++wmIdx) {
+    for (uint wnIdx = 0; wnIdx < WNITER; ++wnIdx) {
+      float *C_interim = C + (wmIdx * WM * N) + (wnIdx * WN);
+      for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
+        for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) {
+          // load C vector into registers
+          float4 tmp = reinterpret_cast<float4 *>(
+              &C_interim[(threadRow * TM + resIdxM) * N + threadCol * TN +
+                         resIdxN])[0];
+          // perform GEMM update in reg
+          const int i =
+              (wmIdx * TM + resIdxM) * (WNITER * TN) + wnIdx * TN + resIdxN;
+          tmp.x = alpha * threadResults[i + 0] + beta * tmp.x;
+          tmp.y = alpha * threadResults[i + 1] + beta * tmp.y;
+          tmp.z = alpha * threadResults[i + 2] + beta * tmp.z;
+          tmp.w = alpha * threadResults[i + 3] + beta * tmp.w;
+          // write back
+          reinterpret_cast<float4 *>(&C_interim[(threadRow * TM + resIdxM) * N +
+                                                threadCol * TN + resIdxN])[0] =
+              tmp;
+        }
+      }
+    }
+  }
+}

upstream_ref/sgemm_edtallison/10_kernel_warptiling.cuh

@@ -0,0 +1,187 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+const int WARPSIZE = 32; // warpSize is not constexpr
+
+namespace wt {
+template <const int BM, const int BN, const int BK, const int rowStrideA,
+          const int rowStrideB>
+__device__ void loadFromGmem(int N, int K, const float *A, const float *B,
+                             float *As, float *Bs, int innerRowA, int innerColA,
+                             int innerRowB, int innerColB) {
+  for (uint offset = 0; offset + rowStrideA <= BM; offset += rowStrideA) {
+    const float4 tmp = reinterpret_cast<const float4 *>(
+        &A[(innerRowA + offset) * K + innerColA * 4])[0];
+    // float4 tmp;
+    // asm("ld.global.nc.v4.f32 {%0, %1, %2, %3}, [%4];"
+    //     : "=f"(tmp.x), "=f"(tmp.y), "=f"(tmp.z), "=f"(tmp.w)
+    //     : "l"(&A[(innerRowA + offset) * K + innerColA * 4]));
+    As[(innerColA * 4 + 0) * BM + innerRowA + offset] = tmp.x;
+    As[(innerColA * 4 + 1) * BM + innerRowA + offset] = tmp.y;
+    As[(innerColA * 4 + 2) * BM + innerRowA + offset] = tmp.z;
+    As[(innerColA * 4 + 3) * BM + innerRowA + offset] = tmp.w;
+  }
+
+  for (uint offset = 0; offset + rowStrideB <= BK; offset += rowStrideB) {
+    reinterpret_cast<float4 *>(
+        &Bs[(innerRowB + offset) * BN + innerColB * 4])[0] =
+        reinterpret_cast<const float4 *>(
+            &B[(innerRowB + offset) * N + innerColB * 4])[0];
+    // asm("ld.global.v4.f32 {%0, %1, %2, %3}, [%4];"
+    //     : "=f"(Bs[(innerRowB + offset) * BN + innerColB * 4 + 0]),
+    //       "=f"(Bs[(innerRowB + offset) * BN + innerColB * 4 + 1]),
+    //       "=f"(Bs[(innerRowB + offset) * BN + innerColB * 4 + 2]),
+    //       "=f"(Bs[(innerRowB + offset) * BN + innerColB * 4 + 3])
+    //     : "l"(&B[(innerRowB + offset) * N + innerColB * 4]));
+  }
+}
+
+template <const int BM, const int BN, const int BK, const int WM, const int WN,
+          const int WMITER, const int WNITER, const int WSUBM, const int WSUBN,
+          const int TM, const int TN>
+__device__ void
+processFromSmem(float *regM, float *regN, float *threadResults, const float *As,
+                const float *Bs, const uint warpRow, const uint warpCol,
+                const uint threadRowInWarp, const uint threadColInWarp) {
+  for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
+    // populate registers for whole warptile
+    for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
+      for (uint i = 0; i < TM; ++i) {
+        regM[wSubRowIdx * TM + i] =
+            As[(dotIdx * BM) + warpRow * WM + wSubRowIdx * WSUBM +
+               threadRowInWarp * TM + i];
+      }
+    }
+    for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
+      for (uint i = 0; i < TN; ++i) {
+        regN[wSubColIdx * TN + i] =
+            Bs[(dotIdx * BN) + warpCol * WN + wSubColIdx * WSUBN +
+               threadColInWarp * TN + i];
+      }
+    }
+
+    // execute warptile matmul
+    for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
+      for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
+        // calculate per-thread results
+        for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
+          for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
+            threadResults[(wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
+                          (wSubColIdx * TN) + resIdxN] +=
+                regM[wSubRowIdx * TM + resIdxM] *
+                regN[wSubColIdx * TN + resIdxN];
+          }
+        }
+      }
+    }
+  }
+}
+
+} // namespace wt
+
+/*
+ * @tparam BM The threadblock size for M dimension SMEM caching.
+ * @tparam BN The threadblock size for N dimension SMEM caching.
+ * @tparam BK The threadblock size for K dimension SMEM caching.
+ * @tparam WM M dim of continuous tile computed by each warp
+ * @tparam WN N dim of continuous tile computed by each warp
+ * @tparam WMITER The number of subwarp tiling steps in M dimension.
+ * @tparam WNITER The number of subwarp tiling steps in N dimension.
+ * @tparam TM The per-thread tile size for M dimension.
+ * @tparam TN The per-thread tile size for N dimension.
+ */
+template <const int BM, const int BN, const int BK, const int WM, const int WN,
+          const int WNITER, const int TM, const int TN, const int NUM_THREADS>
+__global__ void __launch_bounds__(NUM_THREADS)
+    sgemmWarptiling(int M, int N, int K, float alpha, float *A, float *B,
+                    float beta, float *C) {
+  const uint cRow = blockIdx.y;
+  const uint cCol = blockIdx.x;
+
+  // Placement of the warp in the threadblock tile
+  const uint warpIdx = threadIdx.x / WARPSIZE; // the warp this thread is in
+  const uint warpCol = warpIdx % (BN / WN);
+  const uint warpRow = warpIdx / (BN / WN);
+
+  // size of the warp subtile
+  constexpr uint WMITER = (WM * WN) / (WARPSIZE * TM * TN * WNITER);
+  constexpr uint WSUBM = WM / WMITER; // 64/2=32
+  constexpr uint WSUBN = WN / WNITER; // 32/2=16
+
+  // Placement of the thread in the warp subtile
+  const uint threadIdxInWarp = threadIdx.x % WARPSIZE;         // [0, 31]
+  const uint threadColInWarp = threadIdxInWarp % (WSUBN / TN); // i%(16/4)
+  const uint threadRowInWarp = threadIdxInWarp / (WSUBN / TN); // i/4
+
+  // allocate space for the current blocktile in SMEM
+  __shared__ float As[BM * BK];
+  __shared__ float Bs[BK * BN];
+
+  // Move blocktile to beginning of A's row and B's column
+  A += cRow * BM * K;
+  B += cCol * BN;
+  // Move C_ptr to warp's output tile
+  C += (cRow * BM + warpRow * WM) * N + cCol * BN + warpCol * WN;
+
+  // calculating the indices that this thread will load into SMEM
+  // we'll load 128bit / 32bit = 4 elements per thread at each step
+  const uint innerRowA = threadIdx.x / (BK / 4);
+  const uint innerColA = threadIdx.x % (BK / 4);
+  constexpr uint rowStrideA = (NUM_THREADS * 4) / BK;
+  const uint innerRowB = threadIdx.x / (BN / 4);
+  const uint innerColB = threadIdx.x % (BN / 4);
+  constexpr uint rowStrideB = NUM_THREADS / (BN / 4);
+
+  // allocate thread-local cache for results in registerfile
+  float threadResults[WMITER * TM * WNITER * TN] = {0.0};
+  // we cache into registers on the warptile level
+  float regM[WMITER * TM] = {0.0};
+  float regN[WNITER * TN] = {0.0};
+
+  // outer-most loop over block tiles
+  for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
+    wt::loadFromGmem<BM, BN, BK, rowStrideA, rowStrideB>(
+        N, K, A, B, As, Bs, innerRowA, innerColA, innerRowB, innerColB);
+    __syncthreads();
+    wt::processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN, TM,
+                        TN>(regM, regN, threadResults, As, Bs, warpRow, warpCol,
+                            threadRowInWarp, threadColInWarp);
+    A += BK;     // move BK columns to right
+    B += BK * N; // move BK rows down
+    __syncthreads();
+  }
+
+  // write out the results
+  for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
+    for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
+      // move C pointer to current warp subtile
+      float *C_interim = C + (wSubRowIdx * WSUBM) * N + wSubColIdx * WSUBN;
+      for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
+        for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) {
+          // load C vector into registers
+          float4 tmp = reinterpret_cast<float4 *>(
+              &C_interim[(threadRowInWarp * TM + resIdxM) * N +
+                         threadColInWarp * TN + resIdxN])[0];
+          // perform GEMM update in reg
+          const int i = (wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
+                        wSubColIdx * TN + resIdxN;
+          tmp.x = alpha * threadResults[i + 0] + beta * tmp.x;
+          tmp.y = alpha * threadResults[i + 1] + beta * tmp.y;
+          tmp.z = alpha * threadResults[i + 2] + beta * tmp.z;
+          tmp.w = alpha * threadResults[i + 3] + beta * tmp.w;
+          // write back
+          reinterpret_cast<float4 *>(
+              &C_interim[(threadRowInWarp * TM + resIdxM) * N +
+                         threadColInWarp * TN + resIdxN])[0] = tmp;
+        }
+      }
+    }
+  }
+}

upstream_ref/sgemm_edtallison/11_kernel_double_buffering.cuh

@@ -0,0 +1,220 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+namespace db {
+
+template <const int BM, const int BN, const int BK, const int rowStrideA,
+          const int rowStrideB>
+__device__ void loadFromGmem(const int N, const int K, float *A, float *B,
+                             float *As, float *Bs, const int innerRowA,
+                             const int innerColA, const int innerRowB,
+                             const int innerColB) {
+  for (uint offset = 0; offset + rowStrideA <= BM; offset += rowStrideA) {
+    float4 tmp = reinterpret_cast<float4 *>(
+        &A[(innerRowA + offset) * K + innerColA * 4])[0];
+    // transpose A while storing it
+    As[(innerColA * 4 + 0) * BM + innerRowA + offset] = tmp.x;
+    As[(innerColA * 4 + 1) * BM + innerRowA + offset] = tmp.y;
+    As[(innerColA * 4 + 2) * BM + innerRowA + offset] = tmp.z;
+    As[(innerColA * 4 + 3) * BM + innerRowA + offset] = tmp.w;
+  }
+
+  for (uint offset = 0; offset + rowStrideB <= BK; offset += rowStrideB) {
+    reinterpret_cast<float4 *>(
+        &Bs[(innerRowB + offset) * BN + innerColB * 4])[0] =
+        reinterpret_cast<float4 *>(
+            &B[(innerRowB + offset) * N + innerColB * 4])[0];
+  }
+}
+
+template <const int BM, const int BN, const int BK, const int WM, const int WN,
+          const int WMITER, const int WNITER, const int WSUBM, const int WSUBN,
+          const int TM, const int TN>
+__device__ void
+processFromSmem(float *regM, float *regN, float *threadResults, const float *As,
+                const float *Bs, const uint warpRow, const uint warpCol,
+                const uint threadRowInWarp, const uint threadColInWarp) {
+  for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
+    // populate registers for whole warptile
+    for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
+      for (uint i = 0; i < TM; ++i) {
+        regM[wSubRowIdx * TM + i] =
+            As[(dotIdx * BM) + warpRow * WM + wSubRowIdx * WSUBM +
+               threadRowInWarp * TM + i];
+      }
+    }
+    for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
+      for (uint i = 0; i < TN; ++i) {
+        regN[wSubColIdx * TN + i] =
+            Bs[(dotIdx * BN) + warpCol * WN + wSubColIdx * WSUBN +
+               threadColInWarp * TN + i];
+      }
+    }
+
+    // execute warptile matmul
+    for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
+      for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
+        // calculate per-thread results
+        for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
+          for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
+            threadResults[(wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
+                          (wSubColIdx * TN) + resIdxN] +=
+                regM[wSubRowIdx * TM + resIdxM] *
+                regN[wSubColIdx * TN + resIdxN];
+          }
+        }
+      }
+    }
+  }
+}
+
+} // namespace db
+
+template <const int BM, const int BN, const int BK, const int WM, const int WN,
+          const int WNITER, const int TM, const int TN, const int NUM_THREADS>
+__global__ void __launch_bounds__(NUM_THREADS)
+    sgemmDoubleBuffering(const int M, const int N, const int K,
+                         const float alpha, float *A, float *B, float beta,
+                         float *C) {
+  const uint cRow = blockIdx.y;
+  const uint cCol = blockIdx.x;
+
+  // Placement of the warp in the threadblock tile
+  const uint warpIdx = threadIdx.x / WARPSIZE; // the warp this thread is in
+  const uint warpCol = warpIdx % (BN / WN);
+  const uint warpRow = warpIdx / (BN / WN);
+
+  // size of the warp subtile
+  constexpr uint WMITER = (WM * WN) / (WARPSIZE * TM * TN * WNITER);
+  constexpr uint WSUBM = WM / WMITER; // 64/2=32
+  constexpr uint WSUBN = WN / WNITER; // 32/2=16
+
+  // Placement of the thread in the warp subtile
+  const uint threadIdxInWarp = threadIdx.x % WARPSIZE;         // [0, 31]
+  const uint threadColInWarp = threadIdxInWarp % (WSUBN / TN); // i%(16/4)
+  const uint threadRowInWarp = threadIdxInWarp / (WSUBN / TN); // i/4
+
+  // allocate space for the current blocktile in SMEM
+  __shared__ float As[2 * BM * BK];
+  __shared__ float Bs[2 * BK * BN];
+
+  // setup double buffering split
+  bool doubleBufferIdx = threadIdx.x >= (NUM_THREADS / 2);
+
+  // Move blocktile to beginning of A's row and B's column
+  A += cRow * BM * K;
+  B += cCol * BN;
+  // Move C_ptr to warp's output tile
+  C += (cRow * BM + warpRow * WM) * N + cCol * BN + warpCol * WN;
+
+  // calculating the indices that this thread will load into SMEM
+  // for the loading, we're pretending like there's half as many threads
+  // as there actually are
+  const uint innerRowA = (threadIdx.x % (NUM_THREADS / 2)) / (BK / 4);
+  const uint innerColA = (threadIdx.x % (NUM_THREADS / 2)) % (BK / 4);
+  constexpr uint rowStrideA = ((NUM_THREADS / 2) * 4) / BK;
+  const uint innerRowB = (threadIdx.x % (NUM_THREADS / 2)) / (BN / 4);
+  const uint innerColB = (threadIdx.x % (NUM_THREADS / 2)) % (BN / 4);
+  constexpr uint rowStrideB = (NUM_THREADS / 2) / (BN / 4);
+
+  // allocate thread-local cache for results in registerfile
+  float threadResults[WMITER * TM * WNITER * TN] = {0.0};
+  // we cache into registers on the warptile level
+  float regM[WMITER * TM] = {0.0};
+  float regN[WNITER * TN] = {0.0};
+
+  if (doubleBufferIdx == 0) {
+    // load first (B0)
+    db::loadFromGmem<BM, BN, BK, rowStrideA, rowStrideB>(
+        N, K, A, B, As, Bs, innerRowA, innerColA, innerRowB, innerColB);
+  }
+  __syncthreads();
+
+  // outer-most loop over block tiles
+  for (uint bkIdx = 0; bkIdx < K; bkIdx += 2 * BK) {
+    if (doubleBufferIdx == 0) {
+      // process current (B0)
+      db::processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN, TM,
+                          TN>(regM, regN, threadResults, As, Bs, warpRow,
+                              warpCol, threadRowInWarp, threadColInWarp);
+      __syncthreads();
+
+      // process current+1 (B1)
+      if (bkIdx + BK < K) {
+        db::processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN,
+                            TM, TN>(regM, regN, threadResults, As + (BM * BK),
+                                    Bs + (BK * BN), warpRow, warpCol,
+                                    threadRowInWarp, threadColInWarp);
+      }
+      __syncthreads();
+
+      // load current + 2 (B0)
+      if (bkIdx + 2 * BK < K) {
+        db::loadFromGmem<BM, BN, BK, rowStrideA, rowStrideB>(
+            N, K, A + 2 * BK, B + 2 * BK * N, As, Bs, innerRowA, innerColA,
+            innerRowB, innerColB);
+      }
+    } else {
+      // load current + 1 (B1)
+      if (bkIdx + BK < K) {
+        db::loadFromGmem<BM, BN, BK, rowStrideA, rowStrideB>(
+            N, K, A + BK, B + BK * N, As + (BM * BK), Bs + (BK * BN), innerRowA,
+            innerColA, innerRowB, innerColB);
+      }
+      __syncthreads();
+
+      // process current (B0)
+      db::processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN, TM,
+                          TN>(regM, regN, threadResults, As, Bs, warpRow,
+                              warpCol, threadRowInWarp, threadColInWarp);
+      __syncthreads();
+
+      // process current+1 (B1)
+      if (bkIdx + BK < K) {
+        db::processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN,
+                            TM, TN>(regM, regN, threadResults, As + (BM * BK),
+                                    Bs + (BK * BN), warpRow, warpCol,
+                                    threadRowInWarp, threadColInWarp);
+      }
+    }
+
+    A += 2 * BK;     // move BK columns to right
+    B += 2 * BK * N; // move BK rows down
+    __syncthreads();
+  }
+
+  // write out the results
+  for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
+    for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
+      // move C pointer to current warp subtile
+      float *C_interim = C + (wSubRowIdx * WSUBM) * N + wSubColIdx * WSUBN;
+      for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
+        for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) {
+          // load C vector into registers
+          float4 tmp = reinterpret_cast<float4 *>(
+              &C_interim[(threadRowInWarp * TM + resIdxM) * N +
+                         threadColInWarp * TN + resIdxN])[0];
+          // perform GEMM update in reg
+          const int i = (wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
+                        wSubColIdx * TN + resIdxN;
+          tmp.x = alpha * threadResults[i + 0] + beta * tmp.x;
+          tmp.y = alpha * threadResults[i + 1] + beta * tmp.y;
+          tmp.z = alpha * threadResults[i + 2] + beta * tmp.z;
+          tmp.w = alpha * threadResults[i + 3] + beta * tmp.w;
+          // write back
+          reinterpret_cast<float4 *>(
+              &C_interim[(threadRowInWarp * TM + resIdxM) * N +
+                         threadColInWarp * TN + resIdxN])[0] = tmp;
+        }
+      }
+    }
+  }
+}

upstream_ref/sgemm_edtallison/12_kernel_double_buffering.cuh

@@ -0,0 +1,229 @@
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cooperative_groups.h>
+#include <cstdio>
+#include <cstdlib>
+#include <cublas_v2.h>
+#include <cuda/barrier>
+#include <cuda_runtime.h>
+
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+namespace {
+template <const int BM, const int BN, const int BK, const int rowStrideA,
+          const int rowStrideB, typename T>
+__device__ void loadFromGmem(int N, int K, float *A, float *B, float *As,
+                             float *Bs, int innerRowA, int innerColA,
+                             int innerRowB, int innerColB, T &barrier) {
+
+  for (uint offset = 0; offset + rowStrideA <= BM; offset += rowStrideA) {
+    cuda::memcpy_async(&As[(innerColA * 4 + 0) * BM + innerRowA + offset],
+                       &A[(innerRowA + offset) * K + innerColA * 4],
+                       cuda::aligned_size_t<sizeof(float)>(sizeof(float)),
+                       barrier);
+    cuda::memcpy_async(&As[(innerColA * 4 + 1) * BM + innerRowA + offset],
+                       &A[(innerRowA + offset) * K + innerColA * 4 + 1],
+                       cuda::aligned_size_t<sizeof(float)>(sizeof(float)),
+                       barrier);
+    cuda::memcpy_async(&As[(innerColA * 4 + 2) * BM + innerRowA + offset],
+                       &A[(innerRowA + offset) * K + innerColA * 4 + 2],
+                       cuda::aligned_size_t<sizeof(float)>(sizeof(float)),
+                       barrier);
+    cuda::memcpy_async(&As[(innerColA * 4 + 3) * BM + innerRowA + offset],
+                       &A[(innerRowA + offset) * K + innerColA * 4 + 3],
+                       cuda::aligned_size_t<sizeof(float)>(sizeof(float)),
+                       barrier);
+  }
+
+  for (uint offset = 0; offset + rowStrideB <= BK; offset += rowStrideB) {
+    cuda::memcpy_async(&Bs[(innerRowB + offset) * BN + innerColB * 4],
+                       &B[(innerRowB + offset) * N + innerColB * 4],
+                       cuda::aligned_size_t<sizeof(float4)>(sizeof(float4)),
+                       barrier);
+  }
+}
+
+template <const int BM, const int BN, const int BK, const int WM, const int WN,
+          const int WMITER, const int WNITER, const int WSUBM, const int WSUBN,
+          const int TM, const int TN>
+__device__ void
+processFromSmem(float *regM, float *regN, float *threadResults, const float *As,
+                const float *Bs, const uint warpRow, const uint warpCol,
+                const uint threadRowInWarp, const uint threadColInWarp) {
+  for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
+    // populate registers for whole warptile
+    for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
+      for (uint i = 0; i < TM; ++i) {
+        regM[wSubRowIdx * TM + i] =
+            As[(dotIdx * BM) + warpRow * WM + wSubRowIdx * WSUBM +
+               threadRowInWarp * TM + i];
+      }
+    }
+    for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
+      for (uint i = 0; i < TN; ++i) {
+        regN[wSubColIdx * TN + i] =
+            Bs[(dotIdx * BN) + warpCol * WN + wSubColIdx * WSUBN +
+               threadColInWarp * TN + i];
+      }
+    }
+
+    // execute warptile matmul
+    for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
+      for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
+        // calculate per-thread results
+        for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
+          for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
+            threadResults[(wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
+                          (wSubColIdx * TN) + resIdxN] +=
+                regM[wSubRowIdx * TM + resIdxM] *
+                regN[wSubColIdx * TN + resIdxN];
+          }
+        }
+      }
+    }
+  }
+}
+
+} // namespace
+
+/*
+ * @tparam BM The threadblock size for M dimension SMEM caching.
+ * @tparam BN The threadblock size for N dimension SMEM caching.
+ * @tparam BK The threadblock size for K dimension SMEM caching.
+ * @tparam WM M dim of continuous tile computed by each warp
+ * @tparam WN N dim of continuous tile computed by each warp
+ * @tparam WMITER The number of subwarp tiling steps in M dimension.
+ * @tparam WNITER The number of subwarp tiling steps in N dimension.
+ * @tparam TM The per-thread tile size for M dimension.
+ * @tparam TN The per-thread tile size for N dimension.
+ */
+template <const int BM, const int BN, const int BK, const int WM, const int WN,
+          const int WNITER, const int TM, const int TN, const int NUM_THREADS>
+__global__ void __launch_bounds__(NUM_THREADS)
+    runSgemmDoubleBuffering2(int M, int N, int K, float alpha, float *A,
+                             float *B, float beta, float *C) {
+  auto block = cooperative_groups::this_thread_block();
+  __shared__ cuda::barrier<cuda::thread_scope::thread_scope_block> frontBarrier;
+  __shared__ cuda::barrier<cuda::thread_scope::thread_scope_block> backBarrier;
+  auto frontBarrierPtr = &frontBarrier;
+  auto backBarrierPtr = &backBarrier;
+  if (block.thread_rank() == 0) {
+    init(&frontBarrier, block.size());
+    init(&backBarrier, block.size());
+  }
+  __syncthreads();
+
+  const uint cRow = blockIdx.y;
+  const uint cCol = blockIdx.x;
+
+  // Placement of the warp in the threadblock tile
+  const uint warpIdx = threadIdx.x / WARPSIZE; // the warp this thread is in
+  const uint warpCol = warpIdx % (BN / WN);
+  const uint warpRow = warpIdx / (BN / WN);
+
+  // size of the warp subtile
+  constexpr uint WMITER = (WM * WN) / (WARPSIZE * TM * TN * WNITER);
+  constexpr uint WSUBM = WM / WMITER; // 64/2=32
+  constexpr uint WSUBN = WN / WNITER; // 32/2=16
+
+  // Placement of the thread in the warp subtile
+  const uint threadIdxInWarp = threadIdx.x % WARPSIZE;         // [0, 31]
+  const uint threadColInWarp = threadIdxInWarp % (WSUBN / TN); // i%(16/4)
+  const uint threadRowInWarp = threadIdxInWarp / (WSUBN / TN); // i/4
+
+  // allocate space for the current blocktile in SMEM
+  __shared__ float As[2 * BM * BK];
+  __shared__ float Bs[2 * BK * BN];
+
+  // Move blocktile to beginning of A's row and B's column
+  A += cRow * BM * K;
+  B += cCol * BN;
+  // Move C_ptr to warp's output tile
+  C += (cRow * BM + warpRow * WM) * N + cCol * BN + warpCol * WN;
+
+  // calculating the indices that this thread will load into SMEM
+  // we'll load 128bit / 32bit = 4 elements per thread at each step
+  const uint innerRowA = threadIdx.x / (BK / 4);
+  const uint innerColA = threadIdx.x % (BK / 4);
+  constexpr uint rowStrideA = (NUM_THREADS * 4) / BK;
+  const uint innerRowB = threadIdx.x / (BN / 4);
+  const uint innerColB = threadIdx.x % (BN / 4);
+  constexpr uint rowStrideB = NUM_THREADS / (BN / 4);
+
+  // allocate thread-local cache for results in registerfile
+  float threadResults[WMITER * TM * WNITER * TN] = {0.0};
+  // we cache into registers on the warptile level
+  float regM[WMITER * TM] = {0.0};
+  float regN[WNITER * TN] = {0.0};
+
+  int As_offset = 0;
+  int Bs_offset = 0;
+
+  // double-buffering: load first blocktile into SMEM
+  loadFromGmem<BM, BN, BK, rowStrideA, rowStrideB>(
+      N, K, A, B, As + As_offset * BM * BK, Bs + Bs_offset * BK * BN, innerRowA,
+      innerColA, innerRowB, innerColB, (*frontBarrierPtr));
+
+  // outer-most loop over block tiles
+  for (uint bkIdx = 0; bkIdx < K - BK; bkIdx += BK) {
+    // double-buffering: load next blocktile into SMEM
+    loadFromGmem<BM, BN, BK, rowStrideA, rowStrideB>(
+        N, K, A + BK, B + BK * N, As + (1 - As_offset) * BM * BK,
+        Bs + (1 - Bs_offset) * BK * BN, innerRowA, innerColA, innerRowB,
+        innerColB, (*backBarrierPtr));
+
+    // compute the current blocktile
+    (*frontBarrierPtr).arrive_and_wait();
+    processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN, TM, TN>(
+        regM, regN, threadResults, As + As_offset * BM * BK,
+        Bs + Bs_offset * BK * BN, warpRow, warpCol, threadRowInWarp,
+        threadColInWarp);
+    A += BK;     // move BK columns to right
+    B += BK * N; // move BK rows down
+
+    As_offset = 1 - As_offset;
+    Bs_offset = 1 - Bs_offset;
+    // swap the front and back barriers
+    auto tmp = frontBarrierPtr;
+    frontBarrierPtr = backBarrierPtr;
+    backBarrierPtr = tmp;
+
+    __syncthreads();
+  }
+
+  // compute the last blocktile
+  (*frontBarrierPtr).arrive_and_wait();
+  processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN, TM, TN>(
+      regM, regN, threadResults, As + As_offset * BM * BK,
+      Bs + Bs_offset * BK * BN, warpRow, warpCol, threadRowInWarp,
+      threadColInWarp);
+
+  // write out the results
+  for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
+    for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
+      // move C pointer to current warp subtile
+      float *C_interim = C + (wSubRowIdx * WSUBM) * N + wSubColIdx * WSUBN;
+      for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
+        for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) {
+          // load C vector into registers
+          float4 tmp = reinterpret_cast<float4 *>(
+              &C_interim[(threadRowInWarp * TM + resIdxM) * N +
+                         threadColInWarp * TN + resIdxN])[0];
+          // perform GEMM update in reg
+          const int i = (wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
+                        wSubColIdx * TN + resIdxN;
+          tmp.x = alpha * threadResults[i + 0] + beta * tmp.x;
+          tmp.y = alpha * threadResults[i + 1] + beta * tmp.y;
+          tmp.z = alpha * threadResults[i + 2] + beta * tmp.z;
+          tmp.w = alpha * threadResults[i + 3] + beta * tmp.w;
+          // write back
+          reinterpret_cast<float4 *>(
+              &C_interim[(threadRowInWarp * TM + resIdxM) * N +
+                         threadColInWarp * TN + resIdxN])[0] = tmp;
+        }
+      }
+    }
+  }
+}

upstream_ref/sgemm_edtallison/kernels.cuh

@@ -0,0 +1,14 @@
+#pragma once
+
+#include "kernels/01_naive.cuh"
+#include "kernels/02_kernel_global_mem_coalesce.cuh"
+#include "kernels/03_kernel_shared_mem_blocking.cuh"
+#include "kernels/04_kernel_1D_blocktiling.cuh"
+#include "kernels/05_kernel_2D_blocktiling.cuh"
+#include "kernels/06_kernel_vectorize.cuh"
+#include "kernels/07_kernel_resolve_bank_conflicts.cuh"
+#include "kernels/08_kernel_bank_extra_col.cuh"
+#include "kernels/09_kernel_autotuned.cuh"
+#include "kernels/10_kernel_warptiling.cuh"
+#include "kernels/11_kernel_double_buffering.cuh"
+#include "kernels/12_kernel_double_buffering.cuh"

upstream_ref/sgemm_edtallison/runner.cu

@@ -0,0 +1,549 @@
+#include "kernels.cuh"
+#include "runner.cuh"
+#include <cmath>
+#include <cstdio>
+#include <fstream>
+#include <iomanip>
+
+float get_sec() {
+  struct timeval time;
+  gettimeofday(&time, NULL);
+  return (1e6 * time.tv_sec + time.tv_usec);
+}
+
+float cpu_elapsed_time(float &beg, float &end) { return 1.0e-6 * (end - beg); }
+
+void cudaCheck(cudaError_t error, const char *file, int line) {
+  if (error != cudaSuccess) {
+    printf("[CUDA ERROR] at file %s:%d:\n%s\n", file, line,
+           cudaGetErrorString(error));
+    exit(EXIT_FAILURE);
+  }
+};
+
+void CudaDeviceInfo() {
+  int deviceId;
+
+  cudaGetDevice(&deviceId);
+
+  cudaDeviceProp props{};
+  cudaGetDeviceProperties(&props, deviceId);
+
+  printf("Device ID: %d\n\
+    Name: %s\n\
+    Compute Capability: %d.%d\n\
+    memoryBusWidth: %d\n\
+    maxThreadsPerBlock: %d\n\
+    maxThreadsPerMultiProcessor: %d\n\
+    maxRegsPerBlock: %d\n\
+    maxRegsPerMultiProcessor: %d\n\
+    totalGlobalMem: %zuMB\n\
+    sharedMemPerBlock: %zuKB\n\
+    sharedMemPerMultiprocessor: %zuKB\n\
+    totalConstMem: %zuKB\n\
+    multiProcessorCount: %d\n\
+    Warp Size: %d\n",
+         deviceId, props.name, props.major, props.minor, props.memoryBusWidth,
+         props.maxThreadsPerBlock, props.maxThreadsPerMultiProcessor,
+         props.regsPerBlock, props.regsPerMultiprocessor,
+         props.totalGlobalMem / 1024 / 1024, props.sharedMemPerBlock / 1024,
+         props.sharedMemPerMultiprocessor / 1024, props.totalConstMem / 1024,
+         props.multiProcessorCount, props.warpSize);
+};
+
+void randomize_matrix(float *mat, int N) {
+  // NOTICE: Use gettimeofday instead of srand((unsigned)time(NULL)); the time
+  // precision is too low and the same random number is generated.
+  struct timeval time {};
+  gettimeofday(&time, nullptr);
+  srand(time.tv_usec);
+  for (int i = 0; i < N; i++) {
+    float tmp = (float)(rand() % 5) + 0.01 * (rand() % 5);
+    tmp = (rand() % 2 == 0) ? tmp : tmp * (-1.);
+    mat[i] = tmp;
+  }
+}
+
+void range_init_matrix(float *mat, int N) {
+  for (int i = 0; i < N; i++) {
+    mat[i] = i;
+  }
+}
+
+void zero_init_matrix(float *mat, int N) {
+  for (int i = 0; i < N; i++) {
+    mat[i] = 0.0;
+  }
+}
+
+void copy_matrix(const float *src, float *dest, int N) {
+  int i;
+  for (i = 0; src + i && dest + i && i < N; i++)
+    *(dest + i) = *(src + i);
+  if (i != N)
+    printf("copy failed at %d while there are %d elements in total.\n", i, N);
+}
+
+void print_matrix(const float *A, int M, int N, std::ofstream &fs) {
+  int i;
+  fs << std::setprecision(2)
+     << std::fixed; // Set floating-point precision and fixed notation
+  fs << "[";
+  for (i = 0; i < M * N; i++) {
+    if ((i + 1) % N == 0)
+      fs << std::setw(5) << A[i]; // Set field width and write the value
+    else
+      fs << std::setw(5) << A[i] << ", ";
+    if ((i + 1) % N == 0) {
+      if (i + 1 < M * N)
+        fs << ";\n";
+    }
+  }
+  fs << "]\n";
+}
+
+bool verify_matrix(float *matRef, float *matOut, int N) {
+  double diff = 0.0;
+  int i;
+  for (i = 0; i < N; i++) {
+    diff = std::fabs(matRef[i] - matOut[i]);
+    if (diff > 0.01) {
+      printf("Divergence! Should %5.2f, Is %5.2f (Diff %5.2f) at %d\n",
+             matRef[i], matOut[i], diff, i);
+      return false;
+    }
+  }
+  return true;
+}
+
+int div_ceil(int numerator, int denominator) {
+  std::div_t res = std::div(numerator, denominator);
+  return res.rem ? (res.quot + 1) : res.quot;
+}
+
+void runCublasFP32(cublasHandle_t handle, int M, int N, int K, float alpha,
+                   float *A, float *B, float beta, float *C) {
+  // cuBLAS uses column-major order. So we change the order of our row-major A &
+  // B, since (B^T*A^T)^T = (A*B)
+  // This runs cuBLAS in full fp32 mode
+  cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, B, CUDA_R_32F,
+               N, A, CUDA_R_32F, K, &beta, C, CUDA_R_32F, N, CUBLAS_COMPUTE_32F,
+               CUBLAS_GEMM_DEFAULT_TENSOR_OP);
+}
+
+void runCublasBF16(cublasHandle_t handle, int M, int N, int K, float alpha,
+                   float *A, float *B, float beta, float *C) {
+  // This runs cuBLAS with mixed precision (performing the mul with operands
+  // downcast to bf16), which is ~4x faster
+  cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, B, CUDA_R_32F,
+               N, A, CUDA_R_32F, K, &beta, C, CUDA_R_32F, N,
+               CUBLAS_COMPUTE_32F_FAST_16BF, CUBLAS_GEMM_DEFAULT_TENSOR_OP);
+}
+
+void runCublasTF32(cublasHandle_t handle, int M, int N, int K, float alpha,
+                   float *A, float *B, float beta, float *C) {
+  // This runs cuBLAS with mixed precision (performing the mul with operands
+  // downcast to bf16), which is ~4x faster
+  cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, B, CUDA_R_32F,
+               N, A, CUDA_R_32F, K, &beta, C, CUDA_R_32F, N,
+               CUBLAS_COMPUTE_32F_FAST_TF32, CUBLAS_GEMM_DEFAULT_TENSOR_OP);
+}
+
+void run_sgemm_naive(int M, int N, int K, float alpha, float *A, float *B,
+                     float beta, float *C) {
+  dim3 gridDim(CEIL_DIV(M, 32), CEIL_DIV(N, 32));
+  dim3 blockDim(32, 32);
+  sgemm_naive<<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);
+}
+
+void run_sgemm_coalesce(int M, int N, int K, float alpha, float *A, float *B,
+                        float beta, float *C) {
+  dim3 gridDim(CEIL_DIV(M, 32), CEIL_DIV(N, 32));
+  dim3 blockDim(32 * 32);
+  sgemm_global_mem_coalesce<32>
+      <<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);
+}
+
+void run_sgemm_shared_mem_block(int M, int N, int K, float alpha, float *A,
+                                float *B, float beta, float *C) {
+  dim3 gridDim(CEIL_DIV(M, 32), CEIL_DIV(N, 32));
+  dim3 blockDim(32 * 32);
+  // L1 cache becomes useless, since we access GMEM only via SMEM, so we carve
+  // out all of L1 to SMEM. This doesn't currently make a difference, since
+  // occupancy is limited by reg and thread count, but it's good to do anyway.
+  cudaFuncSetAttribute(sgemm_shared_mem_block<32>,
+                       cudaFuncAttributePreferredSharedMemoryCarveout,
+                       cudaSharedmemCarveoutMaxShared);
+  sgemm_shared_mem_block<32>
+      <<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);
+}
+
+void runSgemm1DBlocktiling(int M, int N, int K, float alpha, float *A, float *B,
+                           float beta, float *C) {
+  const uint BM = 64;
+  const uint BN = 64;
+  const uint BK = 8;
+  const uint TM = 8;
+  dim3 gridDim(CEIL_DIV(N, BN), CEIL_DIV(M, BM));
+  dim3 blockDim((BM * BN) / TM);
+  sgemm1DBlocktiling<BM, BN, BK, TM>
+      <<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);
+}
+
+void runSgemm2DBlocktiling(int M, int N, int K, float alpha, float *A, float *B,
+                           float beta, float *C) {
+  const uint BK = 8;
+  const uint TM = 8;
+  const uint TN = 8;
+  if (M >= 128 and N >= 128) {
+    const uint BM = 128;
+    const uint BN = 128;
+    dim3 gridDim(CEIL_DIV(N, BN), CEIL_DIV(M, BM));
+    dim3 blockDim((BM * BN) / (TM * TN));
+    sgemm2DBlocktiling<BM, BN, BK, TM, TN>
+        <<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);
+  } else {
+    // this is a hacky solution to the underlying problem
+    // of not having proper bounds checking in the kernel
+    const uint BM = 64;
+    const uint BN = 64;
+    dim3 gridDim(CEIL_DIV(N, BN), CEIL_DIV(M, BM));
+    dim3 blockDim((BM * BN) / (TM * TN));
+    sgemm2DBlocktiling<BM, BN, BK, TM, TN>
+        <<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);
+  }
+}
+
+void runSgemmVectorize(int M, int N, int K, float alpha, float *A, float *B,
+                       float beta, float *C) {
+  const uint BK = 8;
+  const uint TM = 8;
+  const uint TN = 8;
+  if (M >= 128 and N >= 128) {
+    const uint BM = 128;
+    const uint BN = 128;
+    dim3 gridDim(CEIL_DIV(N, BN), CEIL_DIV(M, BM));
+    dim3 blockDim((BM * BN) / (TM * TN));
+    sgemmVectorize<BM, BN, BK, TM, TN>
+        <<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);
+  } else {
+    // this is a hacky solution to the underlying problem
+    // of not having proper bounds checking in the kernel
+    const uint BM = 64;
+    const uint BN = 64;
+    dim3 gridDim(CEIL_DIV(N, BN), CEIL_DIV(M, BM));
+    dim3 blockDim((BM * BN) / (TM * TN));
+    sgemmVectorize<BM, BN, BK, TM, TN>
+        <<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);
+  }
+}
+
+void runSgemmResolveBankConflicts(int M, int N, int K, float alpha, float *A,
+                                  float *B, float beta, float *C) {
+  const uint BK = 8;
+  const uint TM = 8;
+  const uint TN = 8;
+  if (M >= 128 and N >= 128) {
+    const uint BM = 128;
+    const uint BN = 128;
+    dim3 gridDim(CEIL_DIV(N, BN), CEIL_DIV(M, BM));
+    dim3 blockDim((BM * BN) / (TM * TN));
+    sgemmResolveBankConflicts<BM, BN, BK, TM, TN>
+        <<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);
+  } else {
+    // this is a hacky solution to the underlying problem
+    // of not having proper bounds checking in the kernel
+    const uint BM = 64;
+    const uint BN = 64;
+    dim3 gridDim(CEIL_DIV(N, BN), CEIL_DIV(M, BM));
+    dim3 blockDim((BM * BN) / (TM * TN));
+    sgemmResolveBankConflicts<BM, BN, BK, TM, TN>
+        <<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);
+  }
+}
+
+void runSgemmResolveBankExtraCol(int M, int N, int K, float alpha, float *A,
+                                 float *B, float beta, float *C) {
+  const uint BK = 8;
+  const uint TM = 8;
+  const uint TN = 8;
+  if (M >= 128 and N >= 128) {
+    const uint BM = 128;
+    const uint BN = 128;
+    dim3 gridDim(CEIL_DIV(N, BN), CEIL_DIV(M, BM));
+    dim3 blockDim((BM * BN) / (TM * TN));
+    sgemmResolveBankExtraCol<BM, BN, BK, TM, TN>
+        <<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);
+  } else {
+    // this is a hacky solution to the underlying problem
+    // of not having proper bounds checking in the kernel
+    const uint BM = 64;
+    const uint BN = 64;
+    dim3 gridDim(CEIL_DIV(N, BN), CEIL_DIV(M, BM));
+    dim3 blockDim((BM * BN) / (TM * TN));
+    sgemmResolveBankExtraCol<BM, BN, BK, TM, TN>
+        <<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);
+  }
+}
+
+void runSgemmAutotuned(int M, int N, int K, float alpha, float *A, float *B,
+                       float beta, float *C) {
+  // A100
+  // const uint K9_BK = 16;
+  // const uint K9_TM = 4;
+  // const uint K9_TN = 4;
+  // const uint K9_BM = 64;
+  // const uint K9_BN = 64;
+  // A6000
+  const uint K9_BK = 16;
+  const uint K9_TM = 8;
+  const uint K9_TN = 8;
+  const uint K9_BM = 128;
+  const uint K9_BN = 128;
+  dim3 blockDim(K9_NUM_THREADS);
+
+  static_assert(
+      (K9_NUM_THREADS * 4) % K9_BK == 0,
+      "NUM_THREADS*4 must be multiple of K9_BK to avoid quantization issues "
+      "during GMEM->SMEM tiling (loading only parts of the final row of Bs "
+      "during each iteraion)");
+  static_assert(
+      (K9_NUM_THREADS * 4) % K9_BN == 0,
+      "NUM_THREADS*4 must be multiple of K9_BN to avoid quantization issues "
+      "during GMEM->SMEM tiling (loading only parts of the final row of As "
+      "during each iteration)");
+  static_assert(
+      K9_BN % (16 * K9_TN) == 0,
+      "K9_BN must be a multiple of 16*K9_TN to avoid quantization effects");
+  static_assert(
+      K9_BM % (16 * K9_TM) == 0,
+      "K9_BM must be a multiple of 16*K9_TM to avoid quantization effects");
+  static_assert((K9_BM * K9_BK) % (4 * K9_NUM_THREADS) == 0,
+                "K9_BM*K9_BK must be a multiple of 4*256 to vectorize loads");
+  static_assert((K9_BN * K9_BK) % (4 * K9_NUM_THREADS) == 0,
+                "K9_BN*K9_BK must be a multiple of 4*256 to vectorize loads");
+
+  dim3 gridDim(CEIL_DIV(N, K9_BN), CEIL_DIV(M, K9_BM));
+  sgemmAutotuned<K9_BM, K9_BN, K9_BK, K9_TM, K9_TN>
+      <<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);
+}
+
+void runSgemmWarptiling(int M, int N, int K, float alpha, float *A, float *B,
+                        float beta, float *C) {
+  // Settings for A100
+  // const uint K10_NUM_THREADS = 128;
+  // const uint K10_BN = 128;
+  // const uint K10_BM = 64;
+  // const uint K10_BK = 16;
+  // const uint K10_WN = 64;
+  // const uint K10_WM = 32;
+  // const uint K10_WNITER = 1;
+  // const uint K10_TN = 4;
+  // const uint K10_TM = 4;
+  // Settings for A6000
+  const uint K10_NUM_THREADS = 128;
+  const uint K10_BN = 128;
+  const uint K10_BM = 128;
+  const uint K10_BK = 16;
+  const uint K10_WN = 64;
+  const uint K10_WM = 64;
+  const uint K10_WNITER = 4;
+  const uint K10_TN = 4;
+  const uint K10_TM = 8;
+  dim3 blockDim(K10_NUM_THREADS);
+
+  constexpr uint NUM_WARPS = K10_NUM_THREADS / 32;
+
+  // warptile in threadblocktile
+  static_assert((K10_BN % K10_WN == 0) and (K10_BM % K10_WM == 0));
+  static_assert((K10_BN / K10_WN) * (K10_BM / K10_WM) == NUM_WARPS);
+
+  // threads in warpsubtile
+  static_assert((K10_WM * K10_WN) % (WARPSIZE * K10_TM * K10_TN * K10_WNITER) ==
+                0);
+  constexpr uint K10_WMITER =
+      (K10_WM * K10_WN) / (32 * K10_TM * K10_TN * K10_WNITER);
+  // warpsubtile in warptile
+  static_assert((K10_WM % K10_WMITER == 0) and (K10_WN % K10_WNITER == 0));
+
+  static_assert((K10_NUM_THREADS * 4) % K10_BK == 0,
+                "NUM_THREADS*4 must be multiple of K9_BK to avoid quantization "
+                "issues during GMEM->SMEM tiling (loading only parts of the "
+                "final row of Bs during each iteraion)");
+  static_assert((K10_NUM_THREADS * 4) % K10_BN == 0,
+                "NUM_THREADS*4 must be multiple of K9_BN to avoid quantization "
+                "issues during GMEM->SMEM tiling (loading only parts of the "
+                "final row of As during each iteration)");
+  static_assert(K10_BN % (16 * K10_TN) == 0,
+                "BN must be a multiple of 16*TN to avoid quantization effects");
+  static_assert(K10_BM % (16 * K10_TM) == 0,
+                "BM must be a multiple of 16*TM to avoid quantization effects");
+  static_assert((K10_BM * K10_BK) % (4 * K10_NUM_THREADS) == 0,
+                "BM*BK must be a multiple of 4*256 to vectorize loads");
+  static_assert((K10_BN * K10_BK) % (4 * K10_NUM_THREADS) == 0,
+                "BN*BK must be a multiple of 4*256 to vectorize loads");
+
+  dim3 gridDim(CEIL_DIV(N, K10_BN), CEIL_DIV(M, K10_BM));
+  sgemmWarptiling<K10_BM, K10_BN, K10_BK, K10_WM, K10_WN, K10_WNITER, K10_TM,
+                  K10_TN, K10_NUM_THREADS>
+      <<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);
+}
+
+void runSgemmDoubleBuffering(int M, int N, int K, float alpha, float *A,
+                             float *B, float beta, float *C) {
+  // Settings for A100
+  // const uint K11_NUM_THREADS = 256;
+  // const uint K11_BN = 128;
+  // const uint K11_BM = 64;
+  // const uint K11_BK = 16;
+  // const uint K11_WN = 32;
+  // const uint K11_WM = 32;
+  // const uint K11_WNITER = 2;
+  // const uint K11_TN = 4;
+  // const uint K11_TM = 4;
+  // Settings for A6000
+  const uint K11_NUM_THREADS = 256;
+  const uint K11_BN = 256;
+  const uint K11_BM = 128;
+  const uint K11_BK = 16;
+  const uint K11_WN = 32;
+  const uint K11_WM = 128;
+  const uint K11_WNITER = 1;
+  const uint K11_TN = 8;
+  const uint K11_TM = 8;
+  dim3 blockDim(K11_NUM_THREADS);
+
+  constexpr uint NUM_WARPS = K11_NUM_THREADS / 32;
+
+  // warptile in threadblocktile
+  static_assert((K11_BN % K11_WN == 0) and (K11_BM % K11_WM == 0));
+  static_assert((K11_BN / K11_WN) * (K11_BM / K11_WM) == NUM_WARPS);
+
+  // threads in warpsubtile
+  static_assert((K11_WM * K11_WN) % (WARPSIZE * K11_TM * K11_TN * K11_WNITER) ==
+                0);
+  constexpr uint K11_WMITER =
+      (K11_WM * K11_WN) / (32 * K11_TM * K11_TN * K11_WNITER);
+  // warpsubtile in warptile
+  static_assert((K11_WM % K11_WMITER == 0) and (K11_WN % K11_WNITER == 0));
+
+  static_assert((K11_NUM_THREADS / 2 * 4) % K11_BK == 0,
+                "NUM_THREADS*4 must be multiple of BK to avoid quantization "
+                "issues during GMEM->SMEM tiling (loading only parts of the "
+                "final row of Bs during each iteraion)");
+  static_assert((K11_NUM_THREADS / 2 * 4) % K11_BN == 0,
+                "NUM_THREADS*4 must be multiple of BN to avoid quantization "
+                "issues during GMEM->SMEM tiling (loading only parts of the "
+                "final row of As during each iteration)");
+  static_assert(K11_BN % (16 * K11_TN) == 0,
+                "BN must be a multiple of 16*TN to avoid quantization effects");
+  static_assert(K11_BM % (16 * K11_TM) == 0,
+                "BM must be a multiple of 16*TM to avoid quantization effects");
+  static_assert((K11_BM * K11_BK) % (4 * K11_NUM_THREADS / 2) == 0,
+                "BM*BK must be a multiple of 4*256 to vectorize loads");
+  static_assert((K11_BN * K11_BK) % (4 * K11_NUM_THREADS / 2) == 0,
+                "BN*BK must be a multiple of 4*256 to vectorize loads");
+
+  dim3 gridDim(CEIL_DIV(N, K11_BN), CEIL_DIV(M, K11_BM));
+  sgemmDoubleBuffering<K11_BM, K11_BN, K11_BK, K11_WM, K11_WN, K11_WNITER,
+                       K11_TM, K11_TN, K11_NUM_THREADS>
+      <<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);
+}
+
+void runSgemmDoubleBuffering2(int M, int N, int K, float alpha, float *A,
+                              float *B, float beta, float *C) {
+  // Settings for A6000
+  const uint K12_NUM_THREADS = 128;
+  const uint K12_BN = 128;
+  const uint K12_BM = 128;
+  const uint K12_BK = 16;
+  const uint K12_WN = 64;
+  const uint K12_WM = 64;
+  const uint K12_WNITER = 4;
+  const uint K12_TN = 4;
+  const uint K12_TM = 8;
+  dim3 blockDim(K12_NUM_THREADS);
+
+  constexpr uint NUM_WARPS = K12_NUM_THREADS / 32;
+
+  // warptile in threadblocktile
+  static_assert((K12_BN % K12_WN == 0) and (K12_BM % K12_WM == 0));
+  static_assert((K12_BN / K12_WN) * (K12_BM / K12_WM) == NUM_WARPS);
+
+  // threads in warpsubtile
+  static_assert((K12_WM * K12_WN) % (WARPSIZE * K12_TM * K12_TN * K12_WNITER) ==
+                0);
+  constexpr uint K12_WMITER =
+      (K12_WM * K12_WN) / (32 * K12_TM * K12_TN * K12_WNITER);
+  // warpsubtile in warptile
+  static_assert((K12_WM % K12_WMITER == 0) and (K12_WN % K12_WNITER == 0));
+
+  static_assert((K12_NUM_THREADS * 4) % K12_BK == 0,
+                "NUM_THREADS*4 must be multiple of K9_BK to avoid quantization "
+                "issues during GMEM->SMEM tiling (loading only parts of the "
+                "final row of Bs during each iteraion)");
+  static_assert((K12_NUM_THREADS * 4) % K12_BN == 0,
+                "NUM_THREADS*4 must be multiple of K9_BN to avoid quantization "
+                "issues during GMEM->SMEM tiling (loading only parts of the "
+                "final row of As during each iteration)");
+  static_assert(K12_BN % (16 * K12_TN) == 0,
+                "BN must be a multiple of 16*TN to avoid quantization effects");
+  static_assert(K12_BM % (16 * K12_TM) == 0,
+                "BM must be a multiple of 16*TM to avoid quantization effects");
+  static_assert((K12_BM * K12_BK) % (4 * K12_NUM_THREADS) == 0,
+                "BM*BK must be a multiple of 4*256 to vectorize loads");
+  static_assert((K12_BN * K12_BK) % (4 * K12_NUM_THREADS) == 0,
+                "BN*BK must be a multiple of 4*256 to vectorize loads");
+
+  dim3 gridDim(CEIL_DIV(N, K12_BN), CEIL_DIV(M, K12_BM));
+  runSgemmDoubleBuffering2<K12_BM, K12_BN, K12_BK, K12_WM, K12_WN, K12_WNITER,
+                           K12_TM, K12_TN, K12_NUM_THREADS>
+      <<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);
+}
+
+void run_kernel(int kernel_num, int M, int N, int K, float alpha, float *A,
+                float *B, float beta, float *C, cublasHandle_t handle) {
+  switch (kernel_num) {
+  case 0:
+    runCublasFP32(handle, M, N, K, alpha, A, B, beta, C);
+    break;
+  case 1:
+    run_sgemm_naive(M, N, K, alpha, A, B, beta, C);
+    break;
+  case 2:
+    run_sgemm_coalesce(M, N, K, alpha, A, B, beta, C);
+    break;
+  case 3:
+    run_sgemm_shared_mem_block(M, N, K, alpha, A, B, beta, C);
+    break;
+  case 4:
+    runSgemm1DBlocktiling(M, N, K, alpha, A, B, beta, C);
+    break;
+  case 5:
+    runSgemm2DBlocktiling(M, N, K, alpha, A, B, beta, C);
+    break;
+  case 6:
+    runSgemmVectorize(M, N, K, alpha, A, B, beta, C);
+    break;
+  case 7:
+    runSgemmResolveBankConflicts(M, N, K, alpha, A, B, beta, C);
+    break;
+  case 8:
+    runSgemmResolveBankExtraCol(M, N, K, alpha, A, B, beta, C);
+    break;
+  case 9:
+    runSgemmAutotuned(M, N, K, alpha, A, B, beta, C);
+    break;
+  case 10:
+    runSgemmWarptiling(M, N, K, alpha, A, B, beta, C);
+    break;
+  case 11:
+    runSgemmDoubleBuffering(M, N, K, alpha, A, B, beta, C);
+    break;
+  case 12:
+    runSgemmDoubleBuffering2(M, N, K, alpha, A, B, beta, C);
+    break;
+  default:
+    throw std::invalid_argument("Unknown kernel number");
+  }
+}

upstream_ref/sgemm_edtallison/runner.cuh

@@ -0,0 +1,26 @@
+#pragma once
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+#include <fstream>
+#include <stdio.h>
+#include <stdlib.h>
+#include <sys/time.h>
+#include <time.h>
+#include <unistd.h>
+
+void cudaCheck(cudaError_t error, const char *file,
+               int line); // CUDA error check
+void CudaDeviceInfo();    // print CUDA information
+
+void range_init_matrix(float *mat, int N);
+void randomize_matrix(float *mat, int N);
+void zero_init_matrix(float *mat, int N);
+void copy_matrix(const float *src, float *dest, int N);
+void print_matrix(const float *A, int M, int N, std::ofstream &fs);
+bool verify_matrix(float *mat1, float *mat2, int N);
+
+float get_current_sec();                        // Get the current moment
+float cpu_elapsed_time(float &beg, float &end); // Calculate time difference
+
+void run_kernel(int kernel_num, int m, int n, int k, float alpha, float *A,
+                float *B, float beta, float *C, cublasHandle_t handle);