CUDA Shared Memory Swizzling

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367#include <algorithm>#include <cassert>#include <chrono>#include <cstdio>#include <functional>#include <iomanip>#include <iostream>#include <random>#include <vector>#include <cuda_runtime.h>#define CHECK_CUDA_ERROR(val) check((val), #val, __FILE__, __LINE__)void check(cudaError_t err, char const* func, char const* file, 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); }}#define CHECK_LAST_CUDA_ERROR() check_last(__FILE__, __LINE__)void check_last(char const* file, 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); }}template <class T>float measure_performance(std::function<T(cudaStream_t)> bound_function, cudaStream_t stream, size_t num_repeats = 10, size_t num_warmups = 10){ 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;}constexpr size_t div_up(size_t a, size_t b) { return (a + b - 1) / b; }template <typename T, size_t BLOCK_TILE_SIZE_X = 32, size_t BLOCK_TILE_SIZE_Y = 32, size_t BLOCK_TILE_SKEW_SIZE_X = 0>__global__ void transpose(T* output_matrix, T const* input_matrix, size_t M, size_t N){ __shared__ T shm[BLOCK_TILE_SIZE_Y][BLOCK_TILE_SIZE_X + BLOCK_TILE_SKEW_SIZE_X]; size_t const input_matrix_from_idx_x{threadIdx.x + blockIdx.x * blockDim.x}; size_t const input_matrix_from_idx_y{threadIdx.y + blockIdx.y * blockDim.y}; size_t const input_matrix_from_idx{input_matrix_from_idx_x + input_matrix_from_idx_y * N}; size_t const shm_to_idx_x{threadIdx.x}; size_t const shm_to_idx_y{threadIdx.y}; if ((input_matrix_from_idx_y < M) && (input_matrix_from_idx_x < N)) { shm[shm_to_idx_y][shm_to_idx_x] = input_matrix[input_matrix_from_idx]; } __syncthreads(); size_t const block_thread_idx{threadIdx.x + threadIdx.y * blockDim.x}; size_t const shm_from_idx_x{block_thread_idx / BLOCK_TILE_SIZE_Y}; size_t const shm_from_idx_y{block_thread_idx % BLOCK_TILE_SIZE_Y}; size_t const output_matrix_to_idx_x{shm_from_idx_y + blockIdx.y * blockDim.y}; size_t const output_matrix_to_idx_y{shm_from_idx_x + blockIdx.x * blockDim.x}; size_t const output_matrix_to_idx{output_matrix_to_idx_x + output_matrix_to_idx_y * M}; if ((output_matrix_to_idx_y < N) && (output_matrix_to_idx_x < M)) { output_matrix[output_matrix_to_idx] = shm[shm_from_idx_y][shm_from_idx_x]; }}template <typename T, size_t BLOCK_TILE_SIZE_X = 32, size_t BLOCK_TILE_SIZE_Y = 32>__global__ void transpose_swizzling(T* output_matrix, T const* input_matrix, size_t M, size_t N){ __shared__ T shm[BLOCK_TILE_SIZE_Y][BLOCK_TILE_SIZE_X]; size_t const input_matrix_from_idx_x{threadIdx.x + blockIdx.x * blockDim.x}; size_t const input_matrix_from_idx_y{threadIdx.y + blockIdx.y * blockDim.y}; size_t const input_matrix_from_idx{input_matrix_from_idx_x + input_matrix_from_idx_y * N}; size_t const shm_to_idx_x{threadIdx.x}; size_t const shm_to_idx_y{threadIdx.y}; size_t const shm_to_idx_x_swizzled{(shm_to_idx_x ^ shm_to_idx_y) % BLOCK_TILE_SIZE_X}; if ((input_matrix_from_idx_y < M) && (input_matrix_from_idx_x < N)) { shm[shm_to_idx_y][shm_to_idx_x_swizzled] = input_matrix[input_matrix_from_idx]; } __syncthreads(); size_t const block_thread_idx{threadIdx.x + threadIdx.y * blockDim.x}; size_t const shm_from_idx_x{block_thread_idx / BLOCK_TILE_SIZE_Y}; size_t const shm_from_idx_y{block_thread_idx % BLOCK_TILE_SIZE_Y}; size_t const shm_from_idx_x_swizzled{(shm_from_idx_x ^ shm_from_idx_y) % BLOCK_TILE_SIZE_X}; size_t const output_matrix_to_idx_x{shm_from_idx_y + blockIdx.y * blockDim.y}; size_t const output_matrix_to_idx_y{shm_from_idx_x + blockIdx.x * blockDim.x}; size_t const output_matrix_to_idx{output_matrix_to_idx_x + output_matrix_to_idx_y * M}; if ((output_matrix_to_idx_y < N) && (output_matrix_to_idx_x < M)) { output_matrix[output_matrix_to_idx] = shm[shm_from_idx_y][shm_from_idx_x_swizzled]; }}template <typename T>void launch_transpose_with_shm_bank_conflict(T* d_output_matrix, T const* d_input_matrix, size_t M, size_t N, cudaStream_t stream){ constexpr size_t BLOCK_TILE_SIZE_X{32}; constexpr size_t BLOCK_TILE_SIZE_Y{32}; constexpr size_t BLOCK_TILE_SKEW_SIZE_X{0}; dim3 const block_size{BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y}; dim3 const grid_size{static_cast<unsigned int>(div_up(N, block_size.x)), static_cast<unsigned int>(div_up(M, block_size.y))}; transpose<T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SKEW_SIZE_X> <<<grid_size, block_size, 0, stream>>>(d_output_matrix, d_input_matrix, M, N); CHECK_LAST_CUDA_ERROR();}template <typename T>void launch_transpose_without_shm_bank_conflict_via_padding( T* d_output_matrix, T const* d_input_matrix, size_t M, size_t N, cudaStream_t stream){ constexpr size_t BLOCK_TILE_SIZE_X{32}; constexpr size_t BLOCK_TILE_SIZE_Y{32}; constexpr size_t BLOCK_TILE_SKEW_SIZE_X{1}; dim3 const block_size{BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y}; dim3 const grid_size{static_cast<unsigned int>(div_up(N, block_size.x)), static_cast<unsigned int>(div_up(M, block_size.y))}; transpose<T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SKEW_SIZE_X> <<<grid_size, block_size, 0, stream>>>(d_output_matrix, d_input_matrix, M, N); CHECK_LAST_CUDA_ERROR();}template <typename T>void launch_transpose_without_shm_bank_conflict_via_swizzling( T* d_output_matrix, T const* d_input_matrix, size_t M, size_t N, cudaStream_t stream){ constexpr size_t BLOCK_TILE_SIZE_X{32}; constexpr size_t BLOCK_TILE_SIZE_Y{32}; dim3 const block_size{BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y}; dim3 const grid_size{static_cast<unsigned int>(div_up(N, block_size.x)), static_cast<unsigned int>(div_up(M, block_size.y))}; transpose_swizzling<T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y><<<grid_size, block_size, 0, stream>>>( d_output_matrix, d_input_matrix, M, N); CHECK_LAST_CUDA_ERROR();}template <typename T>bool is_equal(T const* data_1, T const* data_2, size_t size){ for (size_t i{0}; i < size; ++i) { if (data_1[i] != data_2[i]) { return false; } } return true;}template <typename T>bool verify_transpose_implementation( std::function<void(T*, T const*, size_t, size_t, cudaStream_t)> transpose_function, size_t M, size_t N){ std::mt19937 gen{0}; cudaStream_t stream; size_t const matrix_size{M * N}; std::vector<T> matrix(matrix_size, 0.0f); std::vector<T> matrix_transposed(matrix_size, 1.0f); std::vector<T> matrix_transposed_reference(matrix_size, 2.0f); std::uniform_real_distribution<T> uniform_dist(-256, 256); for (size_t i{0}; i < matrix_size; ++i) { matrix[i] = uniform_dist(gen); } for (size_t i{0}; i < M; ++i) { for (size_t j{0}; j < N; ++j) { size_t const from_idx{i * N + j}; size_t const to_idx{j * M + i}; matrix_transposed_reference[to_idx] = matrix[from_idx]; } } T* d_matrix; T* d_matrix_transposed;