diff --git a/upstream_ref/cuda_gemm_optimization/00_non_coalesced_global_memory_access.cu b/upstream_ref/cuda_gemm_optimization/00_non_coalesced_global_memory_access.cu new file mode 100644 index 00000000..65d2acfa --- /dev/null +++ b/upstream_ref/cuda_gemm_optimization/00_non_coalesced_global_memory_access.cu @@ -0,0 +1,65 @@ +#include + +#include "cuda_gemm.hpp" +#include "cuda_gemm_utils.hpp" + +// GEMM kernel v00. +// Non-coalesced read and write from global memory. +template +__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(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 +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(m) + block_dim.x - 1U) / block_dim.x, + (static_cast(n) + block_dim.y - 1U) / block_dim.y, 1U}; + gemm_v00<<>>(m, n, k, *alpha, A, lda, B, + ldb, *beta, C, ldc); + CHECK_LAST_CUDA_ERROR(); +} + +// Explicit instantiation. +template void launch_gemm_kernel_v00(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(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); \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/01_coalesced_global_memory_access.cu b/upstream_ref/cuda_gemm_optimization/01_coalesced_global_memory_access.cu new file mode 100644 index 00000000..e9b8bc69 --- /dev/null +++ b/upstream_ref/cuda_gemm_optimization/01_coalesced_global_memory_access.cu @@ -0,0 +1,65 @@ +#include + +#include "cuda_gemm.hpp" +#include "cuda_gemm_utils.hpp" + +// GEMM kernel v01. +// Coalesced read and write from global memory. +template +__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(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 +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(n) + block_dim.x - 1U) / block_dim.x, + (static_cast(m) + block_dim.y - 1U) / block_dim.y, 1U}; + gemm_v01<<>>(m, n, k, *alpha, A, lda, B, + ldb, *beta, C, ldc); + CHECK_LAST_CUDA_ERROR(); +} + +// Explicit instantiation. +template void launch_gemm_kernel_v01(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(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); \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/02_2d_block_tiling.cu b/upstream_ref/cuda_gemm_optimization/02_2d_block_tiling.cu new file mode 100644 index 00000000..11063d00 --- /dev/null +++ b/upstream_ref/cuda_gemm_optimization/02_2d_block_tiling.cu @@ -0,0 +1,112 @@ +#include + +#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 +__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(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 +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(n) + block_dim.x - 1U) / block_dim.x, + (static_cast(m) + block_dim.y - 1U) / block_dim.y, 1U}; + gemm_v02 + <<>>(m, n, k, *alpha, A, lda, B, ldb, + *beta, C, ldc); + CHECK_LAST_CUDA_ERROR(); +} + +// Explicit instantiation. +template void launch_gemm_kernel_v02(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(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); \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/02_2d_block_tiling_vectorized_memory_access.cu b/upstream_ref/cuda_gemm_optimization/02_2d_block_tiling_vectorized_memory_access.cu new file mode 100644 index 00000000..0dcb9bcb --- /dev/null +++ b/upstream_ref/cuda_gemm_optimization/02_2d_block_tiling_vectorized_memory_access.cu @@ -0,0 +1,108 @@ +#include + +#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 +__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(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 +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(n) + block_dim.x - 1U) / block_dim.x, + (static_cast(m) + block_dim.y - 1U) / block_dim.y, 1U}; + gemm_v02_vectorized<<>>( + m, n, k, *alpha, A, lda, B, ldb, *beta, C, ldc); + CHECK_LAST_CUDA_ERROR(); +} + +// Explicit instantiation. +template void launch_gemm_kernel_v02_vectorized( + 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( + 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); \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/03_2d_block_tiling_1d_thread_tiling.cu b/upstream_ref/cuda_gemm_optimization/03_2d_block_tiling_1d_thread_tiling.cu new file mode 100644 index 00000000..3782ab15 --- /dev/null +++ b/upstream_ref/cuda_gemm_optimization/03_2d_block_tiling_1d_thread_tiling.cu @@ -0,0 +1,144 @@ +#include + +#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 +__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(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 +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(n) + BLOCK_TILE_SIZE_X - 1U) / + BLOCK_TILE_SIZE_X, + (static_cast(m) + BLOCK_TILE_SIZE_Y - 1U) / + BLOCK_TILE_SIZE_Y, + 1U}; + gemm_v03<<>>( + m, n, k, *alpha, A, lda, B, ldb, *beta, C, ldc); + CHECK_LAST_CUDA_ERROR(); +} + +// Explicit instantiation. +template void launch_gemm_kernel_v03(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(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); \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/03_2d_block_tiling_1d_thread_tiling_vectorized_memory_access.cu b/upstream_ref/cuda_gemm_optimization/03_2d_block_tiling_1d_thread_tiling_vectorized_memory_access.cu new file mode 100644 index 00000000..91935030 --- /dev/null +++ b/upstream_ref/cuda_gemm_optimization/03_2d_block_tiling_1d_thread_tiling_vectorized_memory_access.cu @@ -0,0 +1,141 @@ +#include + +#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 +__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(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 +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(n) + BLOCK_TILE_SIZE_X - 1U) / + BLOCK_TILE_SIZE_X, + (static_cast(m) + BLOCK_TILE_SIZE_Y - 1U) / + BLOCK_TILE_SIZE_Y, + 1U}; + gemm_v03_vectorized + <<>>(m, n, k, *alpha, A, lda, B, ldb, + *beta, C, ldc); + CHECK_LAST_CUDA_ERROR(); +} + +// Explicit instantiation. +template void launch_gemm_kernel_v03_vectorized( + 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( + 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); \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/04_2d_block_tiling_2d_thread_tiling.cu b/upstream_ref/cuda_gemm_optimization/04_2d_block_tiling_2d_thread_tiling.cu new file mode 100644 index 00000000..45c149fc --- /dev/null +++ b/upstream_ref/cuda_gemm_optimization/04_2d_block_tiling_2d_thread_tiling.cu @@ -0,0 +1,199 @@ +#include + +#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 +__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(0)}; + // A_vals is cached in the register. + T A_vals[THREAD_TILE_SIZE_Y] = {static_cast(0)}; + // B_vals is cached in the register. + T B_vals[THREAD_TILE_SIZE_X] = {static_cast(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 +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(n) + BLOCK_TILE_SIZE_X - 1U) / + BLOCK_TILE_SIZE_X, + (static_cast(m) + BLOCK_TILE_SIZE_Y - 1U) / + BLOCK_TILE_SIZE_Y, + 1U}; + gemm_v04 + <<>>(m, n, k, *alpha, A, lda, B, ldb, + *beta, C, ldc); + CHECK_LAST_CUDA_ERROR(); +} + +// Explicit instantiation. +template void launch_gemm_kernel_v04(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(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); \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/04_2d_block_tiling_2d_thread_tiling_vectorized_memory_access.cu b/upstream_ref/cuda_gemm_optimization/04_2d_block_tiling_2d_thread_tiling_vectorized_memory_access.cu new file mode 100644 index 00000000..ad355a13 --- /dev/null +++ b/upstream_ref/cuda_gemm_optimization/04_2d_block_tiling_2d_thread_tiling_vectorized_memory_access.cu @@ -0,0 +1,225 @@ +#include + +#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 +__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(0)}; + // A_vals is cached in the register. + T A_vals[THREAD_TILE_SIZE_Y] = {static_cast(0)}; + // B_vals is cached in the register. + T B_vals[THREAD_TILE_SIZE_X] = {static_cast(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( + &B_vals[thread_tile_col_vector_idx * NUM_VECTOR_UNITS]) = + *reinterpret_cast( + &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( + &C[C_row_idx * ldc + C_col_idx])}; + // Vectorized read from C_thread_results. + int4 const C_thread_results_row_vector_vals{ + *reinterpret_cast( + &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(&C_row_vector_vals)[i] = + alpha * reinterpret_cast( + &C_thread_results_row_vector_vals)[i] + + beta * reinterpret_cast(&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(&C[C_row_idx * ldc + C_col_idx]) = + C_row_vector_vals; + } + } + } +} + +template +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(n) + BLOCK_TILE_SIZE_X - 1U) / + BLOCK_TILE_SIZE_X, + (static_cast(m) + BLOCK_TILE_SIZE_Y - 1U) / + BLOCK_TILE_SIZE_Y, + 1U}; + gemm_v04_vectorized + <<>>(m, n, k, *alpha, A, lda, B, ldb, + *beta, C, ldc); + CHECK_LAST_CUDA_ERROR(); +} + +// Explicit instantiation. +template void launch_gemm_kernel_v04_vectorized( + 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( + 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); \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/05_2d_block_tiling_2d_thread_tiling_matrix_transpose.cu b/upstream_ref/cuda_gemm_optimization/05_2d_block_tiling_2d_thread_tiling_matrix_transpose.cu new file mode 100644 index 00000000..64cdc76d --- /dev/null +++ b/upstream_ref/cuda_gemm_optimization/05_2d_block_tiling_2d_thread_tiling_matrix_transpose.cu @@ -0,0 +1,196 @@ +#include + +#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 +__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(0)}; + // A_vals is cached in the register. + T A_vals[THREAD_TILE_SIZE_Y] = {static_cast(0)}; + // B_vals is cached in the register. + T B_vals[THREAD_TILE_SIZE_X] = {static_cast(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 +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(n) + BLOCK_TILE_SIZE_X - 1U) / + BLOCK_TILE_SIZE_X, + (static_cast(m) + BLOCK_TILE_SIZE_Y - 1U) / + BLOCK_TILE_SIZE_Y, + 1U}; + gemm_v05 + <<>>(m, n, k, *alpha, A, lda, B, ldb, + *beta, C, ldc); + CHECK_LAST_CUDA_ERROR(); +} + +// Explicit instantiation. +template void launch_gemm_kernel_v05(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(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); \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/05_2d_block_tiling_2d_thread_tiling_matrix_transpose_vectorized_memory_access.cu b/upstream_ref/cuda_gemm_optimization/05_2d_block_tiling_2d_thread_tiling_matrix_transpose_vectorized_memory_access.cu new file mode 100644 index 00000000..f3cf919f --- /dev/null +++ b/upstream_ref/cuda_gemm_optimization/05_2d_block_tiling_2d_thread_tiling_matrix_transpose_vectorized_memory_access.cu @@ -0,0 +1,222 @@ +#include + +#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 +__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(0)}; + // A_vals is cached in the register. + T A_vals[THREAD_TILE_SIZE_Y] = {static_cast(0)}; + // B_vals is cached in the register. + T B_vals[THREAD_TILE_SIZE_X] = {static_cast(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( + &B_vals[thread_tile_col_vector_idx * NUM_VECTOR_UNITS]) = + *reinterpret_cast( + &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( + &C[C_row_idx * ldc + C_col_idx])}; + // Vectorized read from C_thread_results. + int4 const C_thread_results_row_vector_vals{ + *reinterpret_cast( + &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(&C_row_vector_vals)[i] = + alpha * reinterpret_cast( + &C_thread_results_row_vector_vals)[i] + + beta * reinterpret_cast(&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(&C[C_row_idx * ldc + C_col_idx]) = + C_row_vector_vals; + } + } + } +} + +template +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(n) + BLOCK_TILE_SIZE_X - 1U) / + BLOCK_TILE_SIZE_X, + (static_cast(m) + BLOCK_TILE_SIZE_Y - 1U) / + BLOCK_TILE_SIZE_Y, + 1U}; + gemm_v05_vectorized + <<>>(m, n, k, *alpha, A, lda, B, ldb, + *beta, C, ldc); + CHECK_LAST_CUDA_ERROR(); +} + +// Explicit instantiation. +template void launch_gemm_kernel_v05_vectorized( + 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( + 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); \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/06_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose.cu b/upstream_ref/cuda_gemm_optimization/06_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose.cu new file mode 100644 index 00000000..db98fe25 --- /dev/null +++ b/upstream_ref/cuda_gemm_optimization/06_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose.cu @@ -0,0 +1,334 @@ +#include + +#include "cuda_gemm.hpp" +#include "cuda_gemm_utils.cuh" +#include "cuda_gemm_utils.hpp" + +template +__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 +__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 +__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 +__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(0)}; + // B_vals is cached in the register. + T B_vals[NUM_THREAD_TILES_PER_WARP_X][THREAD_TILE_SIZE_X] = { + static_cast(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(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( + 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 +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(n) + BLOCK_TILE_SIZE_X - 1U) / + BLOCK_TILE_SIZE_X, + (static_cast(m) + BLOCK_TILE_SIZE_Y - 1U) / + BLOCK_TILE_SIZE_Y, + 1U}; + gemm_v06 + <<>>(m, n, k, *alpha, A, lda, B, ldb, + *beta, C, ldc); + CHECK_LAST_CUDA_ERROR(); +} + +// Explicit instantiation. +template void launch_gemm_kernel_v06(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(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); \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/06_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose_vectorized_memory_access.cu b/upstream_ref/cuda_gemm_optimization/06_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose_vectorized_memory_access.cu new file mode 100644 index 00000000..8cd73c18 --- /dev/null +++ b/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 + +#include "cuda_gemm.hpp" +#include "cuda_gemm_utils.cuh" +#include "cuda_gemm_utils.hpp" + +template +__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( + ®ister_values[thread_tile_repeat_row_idx] + [thread_tile_vector_idx * NUM_VECTOR_UNITS]) = + *reinterpret_cast( + &thread_block_tile[thread_block_tile_row_idx + + thread_tile_vector_idx * + NUM_VECTOR_UNITS]); + } + } +} + +template +__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 +__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( + &C[C_row_idx * ldc + C_col_idx])}; +#pragma unroll + for (size_t i{0U}; i < NUM_VECTOR_UNITS; ++i) + { + reinterpret_cast(&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(&C_vals)[i]; + } + *reinterpret_cast( + &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 +__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(0)}; + // B_vals is cached in the register. + T B_vals[NUM_THREAD_TILES_PER_WARP_X][THREAD_TILE_SIZE_X] = { + static_cast(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(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( + 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 +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(n) + BLOCK_TILE_SIZE_X - 1U) / + BLOCK_TILE_SIZE_X, + (static_cast(m) + BLOCK_TILE_SIZE_Y - 1U) / + BLOCK_TILE_SIZE_Y, + 1U}; + gemm_v06_vectorized + <<>>(m, n, k, *alpha, A, lda, B, ldb, + *beta, C, ldc); + CHECK_LAST_CUDA_ERROR(); +} + +// Explicit instantiation. +template void launch_gemm_kernel_v06_vectorized( + 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( + 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); \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/06_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose_vectorized_memory_access_double_buffered.cu b/upstream_ref/cuda_gemm_optimization/06_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose_vectorized_memory_access_double_buffered.cu new file mode 100644 index 00000000..3c4952e7 --- /dev/null +++ b/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 + +#include "cuda_gemm.hpp" +#include "cuda_gemm_utils.cuh" +#include "cuda_gemm_utils.hpp" + +template +__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( + ®ister_values[thread_tile_repeat_row_idx] + [thread_tile_vector_idx * NUM_VECTOR_UNITS]) = + *reinterpret_cast( + &thread_block_tile[thread_block_tile_row_idx + + thread_tile_vector_idx * + NUM_VECTOR_UNITS]); + } + } +} + +template +__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 +__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( + A_vals, B_vals, C_thread_results); + } +} + +template +__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( + &C[C_row_idx * ldc + C_col_idx])}; +#pragma unroll + for (size_t i{0U}; i < NUM_VECTOR_UNITS; ++i) + { + reinterpret_cast(&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(&C_vals)[i]; + } + *reinterpret_cast( + &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 +__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(0)}; + // B_vals is cached in the register. + T B_vals[NUM_THREAD_TILES_PER_WARP_X][THREAD_TILE_SIZE_X] = { + static_cast(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(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 +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(n) + BLOCK_TILE_SIZE_X - 1U) / + BLOCK_TILE_SIZE_X, + (static_cast(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> + <<>>(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( + 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); \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/07_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose_wmma.cu b/upstream_ref/cuda_gemm_optimization/07_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose_wmma.cu new file mode 100644 index 00000000..fff93618 --- /dev/null +++ b/upstream_ref/cuda_gemm_optimization/07_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose_wmma.cu @@ -0,0 +1,247 @@ +#include +#include + +#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 +__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 + a_frags[NUM_WMMA_TILES_Y]; + nvcuda::wmma::fragment + b_frags[NUM_WMMA_TILES_X]; + nvcuda::wmma::fragment + acc_frags[NUM_WMMA_TILES_Y][NUM_WMMA_TILES_X]; + nvcuda::wmma::fragment + 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(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 +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(n) + BLOCK_TILE_SIZE_X - 1U) / + BLOCK_TILE_SIZE_X, + (static_cast(m) + BLOCK_TILE_SIZE_Y - 1U) / + BLOCK_TILE_SIZE_Y, + 1U}; + gemm_v07 + <<>>(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); \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/07_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose_wmma_vectorized_memory_access.cu b/upstream_ref/cuda_gemm_optimization/07_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose_wmma_vectorized_memory_access.cu new file mode 100644 index 00000000..1fb2ef8f --- /dev/null +++ b/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 +#include + +#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 +__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 + a_frags[NUM_WMMA_TILES_Y]; + nvcuda::wmma::fragment + b_frags[NUM_WMMA_TILES_X]; + nvcuda::wmma::fragment + acc_frags[NUM_WMMA_TILES_Y][NUM_WMMA_TILES_X]; + nvcuda::wmma::fragment + 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(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 +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(n) + BLOCK_TILE_SIZE_X - 1U) / + BLOCK_TILE_SIZE_X, + (static_cast(m) + BLOCK_TILE_SIZE_Y - 1U) / + BLOCK_TILE_SIZE_Y, + 1U}; + gemm_v07_vectorized + <<>>(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); \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/07_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose_wmma_vectorized_memory_access_double_buffered.cu b/upstream_ref/cuda_gemm_optimization/07_2d_block_tiling_2d_warp_tiling_2d_thread_tiling_matrix_transpose_wmma_vectorized_memory_access_double_buffered.cu new file mode 100644 index 00000000..90a82c7a --- /dev/null +++ b/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 +#include + +#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 + a_frags[NUM_WMMA_TILES_Y], + nvcuda::wmma::fragment + b_frags[NUM_WMMA_TILES_X], + nvcuda::wmma::fragment + 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 +__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 + a_frags[NUM_WMMA_TILES_Y]; + nvcuda::wmma::fragment + b_frags[NUM_WMMA_TILES_X]; + nvcuda::wmma::fragment + acc_frags[NUM_WMMA_TILES_Y][NUM_WMMA_TILES_X]; + nvcuda::wmma::fragment + 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(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 +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(n) + BLOCK_TILE_SIZE_X - 1U) / + BLOCK_TILE_SIZE_X, + (static_cast(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><<>>( + 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); \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/cuda_gemm.hpp b/upstream_ref/cuda_gemm_optimization/cuda_gemm.hpp new file mode 100644 index 00000000..1681e5e0 --- /dev/null +++ b/upstream_ref/cuda_gemm_optimization/cuda_gemm.hpp @@ -0,0 +1,98 @@ +#ifndef CUDA_GEMM_HPP +#define CUDA_GEMM_HPP + +#include + +template +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 +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 +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 +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 +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 +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 +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 +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 +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 +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 +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 +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 +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 +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 +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 +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 \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/cuda_gemm_utils.cu b/upstream_ref/cuda_gemm_optimization/cuda_gemm_utils.cu new file mode 100644 index 00000000..42d91290 --- /dev/null +++ b/upstream_ref/cuda_gemm_optimization/cuda_gemm_utils.cu @@ -0,0 +1,29 @@ +#include + +#include + +#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); + } +} diff --git a/upstream_ref/cuda_gemm_optimization/cuda_gemm_utils.cuh b/upstream_ref/cuda_gemm_optimization/cuda_gemm_utils.cuh new file mode 100644 index 00000000..4dd95bd2 --- /dev/null +++ b/upstream_ref/cuda_gemm_optimization/cuda_gemm_utils.cuh @@ -0,0 +1,486 @@ +#ifndef CUDA_GEMM_UTILS_CUH +#define CUDA_GEMM_UTILS_CUH + +#include + +#include "cuda_gemm_utils.hpp" + +template +__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(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(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 +__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(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(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 +__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( + &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(&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(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( + &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( + &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(&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(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( + &B_thread_block_tile[B_thread_block_tile_row_idx] + [B_thread_block_tile_col_idx]) = + B_row_vector_vals; + } + } +} + +template +__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(&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(&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(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(&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(&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(&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(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( + &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 \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/cuda_gemm_utils.hpp b/upstream_ref/cuda_gemm_optimization/cuda_gemm_utils.hpp new file mode 100644 index 00000000..a2cf345b --- /dev/null +++ b/upstream_ref/cuda_gemm_optimization/cuda_gemm_utils.hpp @@ -0,0 +1,13 @@ +#ifndef CUDA_GEMM_UTILS_HPP +#define CUDA_GEMM_UTILS_HPP + +#include + +#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 \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/profile_cuda_gemm_fp16.cu b/upstream_ref/cuda_gemm_optimization/profile_cuda_gemm_fp16.cu new file mode 100644 index 00000000..d0b4268b --- /dev/null +++ b/upstream_ref/cuda_gemm_optimization/profile_cuda_gemm_fp16.cu @@ -0,0 +1,109 @@ +#include +#include + +#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>> 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>> 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; +} \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/profile_cuda_gemm_fp32.cu b/upstream_ref/cuda_gemm_optimization/profile_cuda_gemm_fp32.cu new file mode 100644 index 00000000..a985dbe4 --- /dev/null +++ b/upstream_ref/cuda_gemm_optimization/profile_cuda_gemm_fp32.cu @@ -0,0 +1,78 @@ +#include +#include + +#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>> const + gemm_kernel_launch_functions{ + {"Custom GEMM Kernel V00", launch_gemm_kernel_v00}, + {"Custom GEMM Kernel V01", launch_gemm_kernel_v01}, + {"Custom GEMM Kernel V02", launch_gemm_kernel_v02}, + {"Custom GEMM Kernel V02 Vectorized", + launch_gemm_kernel_v02_vectorized}, + {"Custom GEMM Kernel V03", launch_gemm_kernel_v03}, + {"Custom GEMM Kernel V03 Vectorized", + launch_gemm_kernel_v03_vectorized}, + {"Custom GEMM Kernel V04", launch_gemm_kernel_v04}, + {"Custom GEMM Kernel V04 Vectorized", + launch_gemm_kernel_v04_vectorized}, + {"Custom GEMM Kernel V05", launch_gemm_kernel_v05}, + {"Custom GEMM Kernel V05 Vectorized", + launch_gemm_kernel_v05_vectorized}, + {"Custom GEMM Kernel V06", launch_gemm_kernel_v06}, + {"Custom GEMM Kernel V06 Vectorized", + launch_gemm_kernel_v06_vectorized}, + {"Custom GEMM Kernel V06 Vectorized Double Buffered", + launch_gemm_kernel_v06_vectorized_double_buffered}, + }; + + for (auto const& gemm_kernel_launch_function : gemm_kernel_launch_functions) + { + std::cout << gemm_kernel_launch_function.first << std::endl; + std::pair const gemm_kernel_profile_result{ + profile_gemm( + 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; +} \ No newline at end of file diff --git a/upstream_ref/cuda_gemm_optimization/profile_utils.cuh b/upstream_ref/cuda_gemm_optimization/profile_utils.cuh new file mode 100644 index 00000000..f62f8097 --- /dev/null +++ b/upstream_ref/cuda_gemm_optimization/profile_utils.cuh @@ -0,0 +1,400 @@ +#ifndef PROFILE_UTILS_CUH +#define PROFILE_UTILS_CUH + +#include +#include +#include +#include +#include + +#include "cuda_gemm.hpp" +#include "cuda_gemm_utils.cuh" + +#include +#include +#include + +template +float measure_performance(std::function 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 ::value || + std::is_same::value || + std::is_same::value, + bool>::type = true> +constexpr cudaDataType_t cuda_data_type_trait() +{ + if (std::is_same::value) + { + return CUDA_R_32F; + } + else if (std::is_same::value) + { + return CUDA_R_64F; + } + else if (std::is_same::value) + { + return CUDA_R_16F; + } + else + { + throw std::runtime_error("Unsupported data type."); + } +} + +template ::value || + std::is_same::value || + std::is_same::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()}; + // 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 ::value || + std::is_same::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(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 ::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 +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(C[i * ldc + j])}; + double const C_ref_val{static_cast(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(abs_tol), + static_cast(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(abs_tol) + << " Rel->Abs Diff Threshold: " + << static_cast( + static_cast(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(device_prop.totalGlobalMem) / + (1 << 30)}; + std::cout << "Memory Size: " << memory_size << " GB" << std::endl; + float const peak_bandwidth{ + static_cast(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 +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 ::value || + std::is_same::value || + std::is_same::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 dis(0, 5); + // std::uniform_real_distribution 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(rand()); + } + } +} + +void print_performance_result(size_t m, size_t n, size_t k, float latency) +{ + float const effective_bandwidth{ + compute_effective_bandwidth(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 ::value || + std::is_same::value || + std::is_same::value, + bool>::type = true> +std::pair profile_gemm( + size_t m, size_t n, size_t k, size_t lda, size_t ldb, size_t ldc, + std::function + 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(1.0)}; + T const beta{static_cast(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(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(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(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(C_host_from_device, C_host_ref, m, n, ldc, abs_tol, + rel_tol)); + + // Launch cuBLAS GEMM. + float const latency_cublas{measure_performance( + [&](cudaStream_t stream) + { + launch_gemm_cublas(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( + [&](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{latency_cublas, latency_cuda_gemm}; +} + +#endif // PROFILE_UTILS_CUH \ No newline at end of file diff --git a/upstream_ref/nvidia_sgemm_practice/kernel_1.cuh b/upstream_ref/nvidia_sgemm_practice/kernel_1.cuh new file mode 100644 index 00000000..94ce9a77 --- /dev/null +++ b/upstream_ref/nvidia_sgemm_practice/kernel_1.cuh @@ -0,0 +1,19 @@ +#pragma once + +#include +#include +#include +#include + +__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]; +} \ No newline at end of file diff --git a/upstream_ref/nvidia_sgemm_practice/kernel_2.cuh b/upstream_ref/nvidia_sgemm_practice/kernel_2.cuh new file mode 100644 index 00000000..73bde1a1 --- /dev/null +++ b/upstream_ref/nvidia_sgemm_practice/kernel_2.cuh @@ -0,0 +1,45 @@ +#pragma once + +#include +#include +#include +#include + +template +__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]; +} \ No newline at end of file diff --git a/upstream_ref/nvidia_sgemm_practice/kernel_3.cuh b/upstream_ref/nvidia_sgemm_practice/kernel_3.cuh new file mode 100644 index 00000000..f55a32ce --- /dev/null +++ b/upstream_ref/nvidia_sgemm_practice/kernel_3.cuh @@ -0,0 +1,71 @@ +#pragma once + +#include +#include +#include +#include + +template +__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]; + } +} \ No newline at end of file diff --git a/upstream_ref/nvidia_sgemm_practice/kernel_4.cuh b/upstream_ref/nvidia_sgemm_practice/kernel_4.cuh new file mode 100644 index 00000000..bb2214fa --- /dev/null +++ b/upstream_ref/nvidia_sgemm_practice/kernel_4.cuh @@ -0,0 +1,76 @@ +#pragma once + +#include +#include +#include +#include + +template +__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]; + } +} \ No newline at end of file diff --git a/upstream_ref/nvidia_sgemm_practice/kernel_5.cuh b/upstream_ref/nvidia_sgemm_practice/kernel_5.cuh new file mode 100644 index 00000000..2fcd0587 --- /dev/null +++ b/upstream_ref/nvidia_sgemm_practice/kernel_5.cuh @@ -0,0 +1,88 @@ +#pragma once + +#include +#include +#include +#include + +template +__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]; + } +} \ No newline at end of file diff --git a/upstream_ref/nvidia_sgemm_practice/kernel_6.cuh b/upstream_ref/nvidia_sgemm_practice/kernel_6.cuh new file mode 100644 index 00000000..f40c2501 --- /dev/null +++ b/upstream_ref/nvidia_sgemm_practice/kernel_6.cuh @@ -0,0 +1,110 @@ +#pragma once + +#include +#include +#include +#include + +#define OFFSET(row, col, ld) ((row)*(ld)+(col)) +#define FETCH_FLOAT4(pointer) (reinterpret_cast(&(pointer))[0]) + +template +__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; + } + } +} \ No newline at end of file diff --git a/upstream_ref/nvidia_sgemm_practice/kernel_7.cuh b/upstream_ref/nvidia_sgemm_practice/kernel_7.cuh new file mode 100644 index 00000000..53639644 --- /dev/null +++ b/upstream_ref/nvidia_sgemm_practice/kernel_7.cuh @@ -0,0 +1,180 @@ +#pragma once + +#include +#include +#include +#include + +#define OFFSET(row, col, ld) ((row)*(ld)+(col)) +#define FETCH_FLOAT4(pointer) (reinterpret_cast(&(pointer))[0]) + +template +__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; + } + } +} \ No newline at end of file diff --git a/upstream_ref/nvidia_sgemm_practice/sgemm.cu b/upstream_ref/nvidia_sgemm_practice/sgemm.cu new file mode 100644 index 00000000..b93c59d3 --- /dev/null +++ b/upstream_ref/nvidia_sgemm_practice/sgemm.cu @@ -0,0 +1,119 @@ +#include +#include +#include +#include + +#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; +}; diff --git a/upstream_ref/nvidia_sgemm_practice/utils.cuh b/upstream_ref/nvidia_sgemm_practice/utils.cuh new file mode 100644 index 00000000..fa423199 --- /dev/null +++ b/upstream_ref/nvidia_sgemm_practice/utils.cuh @@ -0,0 +1,42 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +/* +===================================== +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); \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/10_kernel_warptiling.cuh b/upstream_ref/sgemm_cuda/10_kernel_warptiling.cuh new file mode 100644 index 00000000..2cc66f36 --- /dev/null +++ b/upstream_ref/sgemm_cuda/10_kernel_warptiling.cuh @@ -0,0 +1,187 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) +const int WARPSIZE = 32; // warpSize is not constexpr + +namespace wt { +template +__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( + &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( + &Bs[(innerRowB + offset) * BN + innerColB * 4])[0] = + reinterpret_cast( + &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 +__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 +__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( + N, K, A, B, As, Bs, innerRowA, innerColA, innerRowB, innerColB); + __syncthreads(); + wt::processFromSmem(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( + &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( + &C_interim[(threadRowInWarp * TM + resIdxM) * N + + threadColInWarp * TN + resIdxN])[0] = tmp; + } + } + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/11_kernel_double_buffering.cuh b/upstream_ref/sgemm_cuda/11_kernel_double_buffering.cuh new file mode 100644 index 00000000..f54b61cf --- /dev/null +++ b/upstream_ref/sgemm_cuda/11_kernel_double_buffering.cuh @@ -0,0 +1,220 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +namespace db { + +template +__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( + &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( + &Bs[(innerRowB + offset) * BN + innerColB * 4])[0] = + reinterpret_cast( + &B[(innerRowB + offset) * N + innerColB * 4])[0]; + } +} + +template +__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 +__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( + 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(regM, regN, threadResults, As, Bs, warpRow, + warpCol, threadRowInWarp, threadColInWarp); + __syncthreads(); + + // process current+1 (B1) + if (bkIdx + BK < K) { + db::processFromSmem(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( + 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( + N, K, A + BK, B + BK * N, As + (BM * BK), Bs + (BK * BN), innerRowA, + innerColA, innerRowB, innerColB); + } + __syncthreads(); + + // process current (B0) + db::processFromSmem(regM, regN, threadResults, As, Bs, warpRow, + warpCol, threadRowInWarp, threadColInWarp); + __syncthreads(); + + // process current+1 (B1) + if (bkIdx + BK < K) { + db::processFromSmem(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( + &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( + &C_interim[(threadRowInWarp * TM + resIdxM) * N + + threadColInWarp * TN + resIdxN])[0] = tmp; + } + } + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/12_kernel_double_buffering.cuh b/upstream_ref/sgemm_cuda/12_kernel_double_buffering.cuh new file mode 100644 index 00000000..c0b0a40f --- /dev/null +++ b/upstream_ref/sgemm_cuda/12_kernel_double_buffering.cuh @@ -0,0 +1,229 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +namespace { +template +__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)), + barrier); + cuda::memcpy_async(&As[(innerColA * 4 + 1) * BM + innerRowA + offset], + &A[(innerRowA + offset) * K + innerColA * 4 + 1], + cuda::aligned_size_t(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)), + barrier); + cuda::memcpy_async(&As[(innerColA * 4 + 3) * BM + innerRowA + offset], + &A[(innerRowA + offset) * K + innerColA * 4 + 3], + cuda::aligned_size_t(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)), + barrier); + } +} + +template +__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 +__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 frontBarrier; + __shared__ cuda::barrier 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( + 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( + 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( + 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( + 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( + &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( + &C_interim[(threadRowInWarp * TM + resIdxM) * N + + threadColInWarp * TN + resIdxN])[0] = tmp; + } + } + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/1_naive.cuh b/upstream_ref/sgemm_cuda/1_naive.cuh new file mode 100644 index 00000000..47038a37 --- /dev/null +++ b/upstream_ref/sgemm_cuda/1_naive.cuh @@ -0,0 +1,29 @@ +#pragma once + +#include +#include +#include +#include + +/* + +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]; + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/2_kernel_global_mem_coalesce.cuh b/upstream_ref/sgemm_cuda/2_kernel_global_mem_coalesce.cuh new file mode 100644 index 00000000..ef0a8b94 --- /dev/null +++ b/upstream_ref/sgemm_cuda/2_kernel_global_mem_coalesce.cuh @@ -0,0 +1,24 @@ +#pragma once + +#include +#include +#include +#include +#include + +template +__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]; + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/3_kernel_shared_mem_blocking.cuh b/upstream_ref/sgemm_cuda/3_kernel_shared_mem_blocking.cuh new file mode 100644 index 00000000..806b6982 --- /dev/null +++ b/upstream_ref/sgemm_cuda/3_kernel_shared_mem_blocking.cuh @@ -0,0 +1,57 @@ +#pragma once + +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +template +__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]; +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/4_kernel_1D_blocktiling.cuh b/upstream_ref/sgemm_cuda/4_kernel_1D_blocktiling.cuh new file mode 100644 index 00000000..12c9c5d7 --- /dev/null +++ b/upstream_ref/sgemm_cuda/4_kernel_1D_blocktiling.cuh @@ -0,0 +1,80 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +template +__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]; + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/5_kernel_2D_blocktiling.cuh b/upstream_ref/sgemm_cuda/5_kernel_2D_blocktiling.cuh new file mode 100644 index 00000000..0b361f7c --- /dev/null +++ b/upstream_ref/sgemm_cuda/5_kernel_2D_blocktiling.cuh @@ -0,0 +1,102 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +template +__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]; + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/6_kernel_vectorize.cuh b/upstream_ref/sgemm_cuda/6_kernel_vectorize.cuh new file mode 100644 index 00000000..665d4e50 --- /dev/null +++ b/upstream_ref/sgemm_cuda/6_kernel_vectorize.cuh @@ -0,0 +1,98 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +template +__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(&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(&Bs[innerRowB * BN + innerColB * 4])[0] = + reinterpret_cast(&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( + &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( + &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0] = + tmp; + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/7_kernel_resolve_bank_conflicts.cuh b/upstream_ref/sgemm_cuda/7_kernel_resolve_bank_conflicts.cuh new file mode 100644 index 00000000..d571008e --- /dev/null +++ b/upstream_ref/sgemm_cuda/7_kernel_resolve_bank_conflicts.cuh @@ -0,0 +1,103 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +template +__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(&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(&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( + &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( + &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0] = + tmp; + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/8_kernel_bank_extra_col.cuh b/upstream_ref/sgemm_cuda/8_kernel_bank_extra_col.cuh new file mode 100644 index 00000000..c362cba8 --- /dev/null +++ b/upstream_ref/sgemm_cuda/8_kernel_bank_extra_col.cuh @@ -0,0 +1,103 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +template +__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(&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(&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( + &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( + &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0] = + tmp; + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_cuda/9_kernel_autotuned.cuh b/upstream_ref/sgemm_cuda/9_kernel_autotuned.cuh new file mode 100644 index 00000000..6841952b --- /dev/null +++ b/upstream_ref/sgemm_cuda/9_kernel_autotuned.cuh @@ -0,0 +1,127 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) +const int K9_NUM_THREADS = 256; + +template +__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( + &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( + &Bs[(innerRowB + offset) * BN + innerColB * 4])[0] = + reinterpret_cast( + &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( + &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(&C_interim[(threadRow * TM + resIdxM) * N + + threadCol * TN + resIdxN])[0] = + tmp; + } + } + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_edtallison/01_naive.cuh b/upstream_ref/sgemm_edtallison/01_naive.cuh new file mode 100644 index 00000000..04ed8e8c --- /dev/null +++ b/upstream_ref/sgemm_edtallison/01_naive.cuh @@ -0,0 +1,37 @@ +# pragma once + +#include +#include +#include +#include + +/* + +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]; + } +} diff --git a/upstream_ref/sgemm_edtallison/02_kernel_global_mem_coalesce.cuh b/upstream_ref/sgemm_edtallison/02_kernel_global_mem_coalesce.cuh new file mode 100644 index 00000000..2f6f4504 --- /dev/null +++ b/upstream_ref/sgemm_edtallison/02_kernel_global_mem_coalesce.cuh @@ -0,0 +1,25 @@ +#pragma once + +#include +#include +#include +#include +#include + +template +// __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]; + } +} diff --git a/upstream_ref/sgemm_edtallison/03_kernel_shared_mem_blocking.cuh b/upstream_ref/sgemm_edtallison/03_kernel_shared_mem_blocking.cuh new file mode 100644 index 00000000..975cc47f --- /dev/null +++ b/upstream_ref/sgemm_edtallison/03_kernel_shared_mem_blocking.cuh @@ -0,0 +1,54 @@ +#pragma once + +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +template +__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]; +} \ No newline at end of file diff --git a/upstream_ref/sgemm_edtallison/04_kernel_1D_blocktiling.cuh b/upstream_ref/sgemm_edtallison/04_kernel_1D_blocktiling.cuh new file mode 100644 index 00000000..12c9c5d7 --- /dev/null +++ b/upstream_ref/sgemm_edtallison/04_kernel_1D_blocktiling.cuh @@ -0,0 +1,80 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +template +__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]; + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_edtallison/05_kernel_2D_blocktiling.cuh b/upstream_ref/sgemm_edtallison/05_kernel_2D_blocktiling.cuh new file mode 100644 index 00000000..0b361f7c --- /dev/null +++ b/upstream_ref/sgemm_edtallison/05_kernel_2D_blocktiling.cuh @@ -0,0 +1,102 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +template +__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]; + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_edtallison/06_kernel_vectorize.cuh b/upstream_ref/sgemm_edtallison/06_kernel_vectorize.cuh new file mode 100644 index 00000000..665d4e50 --- /dev/null +++ b/upstream_ref/sgemm_edtallison/06_kernel_vectorize.cuh @@ -0,0 +1,98 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +template +__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(&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(&Bs[innerRowB * BN + innerColB * 4])[0] = + reinterpret_cast(&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( + &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( + &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0] = + tmp; + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_edtallison/07_kernel_resolve_bank_conflicts.cuh b/upstream_ref/sgemm_edtallison/07_kernel_resolve_bank_conflicts.cuh new file mode 100644 index 00000000..d571008e --- /dev/null +++ b/upstream_ref/sgemm_edtallison/07_kernel_resolve_bank_conflicts.cuh @@ -0,0 +1,103 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +template +__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(&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(&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( + &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( + &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0] = + tmp; + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_edtallison/08_kernel_bank_extra_col.cuh b/upstream_ref/sgemm_edtallison/08_kernel_bank_extra_col.cuh new file mode 100644 index 00000000..c362cba8 --- /dev/null +++ b/upstream_ref/sgemm_edtallison/08_kernel_bank_extra_col.cuh @@ -0,0 +1,103 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +template +__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(&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(&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( + &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( + &C[(threadRow * TM + resIdxM) * N + threadCol * TN + resIdxN])[0] = + tmp; + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_edtallison/09_kernel_autotuned.cuh b/upstream_ref/sgemm_edtallison/09_kernel_autotuned.cuh new file mode 100644 index 00000000..6841952b --- /dev/null +++ b/upstream_ref/sgemm_edtallison/09_kernel_autotuned.cuh @@ -0,0 +1,127 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) +const int K9_NUM_THREADS = 256; + +template +__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( + &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( + &Bs[(innerRowB + offset) * BN + innerColB * 4])[0] = + reinterpret_cast( + &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( + &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(&C_interim[(threadRow * TM + resIdxM) * N + + threadCol * TN + resIdxN])[0] = + tmp; + } + } + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_edtallison/10_kernel_warptiling.cuh b/upstream_ref/sgemm_edtallison/10_kernel_warptiling.cuh new file mode 100644 index 00000000..2cc66f36 --- /dev/null +++ b/upstream_ref/sgemm_edtallison/10_kernel_warptiling.cuh @@ -0,0 +1,187 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) +const int WARPSIZE = 32; // warpSize is not constexpr + +namespace wt { +template +__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( + &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( + &Bs[(innerRowB + offset) * BN + innerColB * 4])[0] = + reinterpret_cast( + &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 +__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 +__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( + N, K, A, B, As, Bs, innerRowA, innerColA, innerRowB, innerColB); + __syncthreads(); + wt::processFromSmem(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( + &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( + &C_interim[(threadRowInWarp * TM + resIdxM) * N + + threadColInWarp * TN + resIdxN])[0] = tmp; + } + } + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_edtallison/11_kernel_double_buffering.cuh b/upstream_ref/sgemm_edtallison/11_kernel_double_buffering.cuh new file mode 100644 index 00000000..f54b61cf --- /dev/null +++ b/upstream_ref/sgemm_edtallison/11_kernel_double_buffering.cuh @@ -0,0 +1,220 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +namespace db { + +template +__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( + &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( + &Bs[(innerRowB + offset) * BN + innerColB * 4])[0] = + reinterpret_cast( + &B[(innerRowB + offset) * N + innerColB * 4])[0]; + } +} + +template +__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 +__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( + 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(regM, regN, threadResults, As, Bs, warpRow, + warpCol, threadRowInWarp, threadColInWarp); + __syncthreads(); + + // process current+1 (B1) + if (bkIdx + BK < K) { + db::processFromSmem(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( + 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( + N, K, A + BK, B + BK * N, As + (BM * BK), Bs + (BK * BN), innerRowA, + innerColA, innerRowB, innerColB); + } + __syncthreads(); + + // process current (B0) + db::processFromSmem(regM, regN, threadResults, As, Bs, warpRow, + warpCol, threadRowInWarp, threadColInWarp); + __syncthreads(); + + // process current+1 (B1) + if (bkIdx + BK < K) { + db::processFromSmem(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( + &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( + &C_interim[(threadRowInWarp * TM + resIdxM) * N + + threadColInWarp * TN + resIdxN])[0] = tmp; + } + } + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_edtallison/12_kernel_double_buffering.cuh b/upstream_ref/sgemm_edtallison/12_kernel_double_buffering.cuh new file mode 100644 index 00000000..c0b0a40f --- /dev/null +++ b/upstream_ref/sgemm_edtallison/12_kernel_double_buffering.cuh @@ -0,0 +1,229 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +namespace { +template +__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)), + barrier); + cuda::memcpy_async(&As[(innerColA * 4 + 1) * BM + innerRowA + offset], + &A[(innerRowA + offset) * K + innerColA * 4 + 1], + cuda::aligned_size_t(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)), + barrier); + cuda::memcpy_async(&As[(innerColA * 4 + 3) * BM + innerRowA + offset], + &A[(innerRowA + offset) * K + innerColA * 4 + 3], + cuda::aligned_size_t(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)), + barrier); + } +} + +template +__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 +__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 frontBarrier; + __shared__ cuda::barrier 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( + 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( + 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( + 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( + 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( + &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( + &C_interim[(threadRowInWarp * TM + resIdxM) * N + + threadColInWarp * TN + resIdxN])[0] = tmp; + } + } + } + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_edtallison/kernels.cuh b/upstream_ref/sgemm_edtallison/kernels.cuh new file mode 100644 index 00000000..b027b017 --- /dev/null +++ b/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" \ No newline at end of file diff --git a/upstream_ref/sgemm_edtallison/runner.cu b/upstream_ref/sgemm_edtallison/runner.cu new file mode 100644 index 00000000..4c0142dc --- /dev/null +++ b/upstream_ref/sgemm_edtallison/runner.cu @@ -0,0 +1,549 @@ +#include "kernels.cuh" +#include "runner.cuh" +#include +#include +#include +#include + +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<<>>(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> + <<>>(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> + <<>>(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 + <<>>(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 + <<>>(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 + <<>>(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 + <<>>(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 + <<>>(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 + <<>>(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 + <<>>(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 + <<>>(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 + <<>>(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 + <<>>(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 + <<>>(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 + <<>>(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 + <<>>(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"); + } +} \ No newline at end of file diff --git a/upstream_ref/sgemm_edtallison/runner.cuh b/upstream_ref/sgemm_edtallison/runner.cuh new file mode 100644 index 00000000..3ac967e6 --- /dev/null +++ b/upstream_ref/sgemm_edtallison/runner.cuh @@ -0,0 +1,26 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include + +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); \ No newline at end of file