Optimizing Matrix Multiplication in CUDA
Introduction
This post explores several techniques for optimizing matrix multiplication in CUDA.
Baseline CPU Implementation
First, we implement $C = A \times B$ on the CPU.
fn matmul(a: &[f32], b: &[f32], c: &mut [f32], m: usize, k: usize, n: usize) {
for i in 0..m {
for j in 0..n {
c[i * n + j] = 0.0;
for p in 0..k {
c[i * n + j] += a[i * k + p] * b[p * n + j];
}
}
}
}
Naive GPU Implementation
// Naive GPU implementation
__global__ void matmul_kernel(const float *a, const float *b, float *c, int m, int k, int n) {
const int col = blockIdx.x * blockDim.x + threadIdx.x;
const int row = blockIdx.y * blockDim.y + threadIdx.y;
if (row < m && col < n) {
float sum = 0.0f;
for (int i = 0; i < k; i++) {
sum += a[row * k + i] * b[i * n + col];
}
c[row * n + col] = sum;
}
}
Each thread computes one element of the output matrix $C$.
Global-Memory Coalescing
Before discussing global-memory coalescing, we first need to understand how GPU threads are organized into warps. Threads in a block are linearized and divided into warps, typically containing 32 threads. Warps are scheduled for execution by the warp schedulers in a streaming multiprocessor (SM), and the threads in a warp execute the same instruction in SIMT fashion. The previous implementation does not take this execution model into account. To use memory bandwidth efficiently, neighboring threads in a warp should access neighboring memory addresses whenever possible. For a three-dimensional block, the linear thread index is computed as follows: \(threadId = threadIdx.x + blockDim.x * threadIdx.y + blockDim.x * blockDim.y * threadIdx.z\)
Threads are divided into groups of 32, so if tid1 / 32 == tid2 / 32, the two threads belong to the same warp. When the threads in a warp access consecutive float values, the hardware can combine their requests into a small number of memory transactions. Strided accesses may require additional transactions and therefore waste memory bandwidth.
In the previous kernel, we assigned each thread an element of $C$ as follows:
const uint x = blockDim.x * blockIdx.x + threadIdx.x;
const uint y = blockDim.y * blockIdx.y + threadIdx.y;
Depending on how the block is configured, this mapping may cause threads in the same warp—those with consecutive threadIdx.x values—to access different rows of $A$ with a large stride.
We can change the mapping with a small modification:
const uint x = blockDim.x * blockIdx.x + (threadIdx.x / BLOCKSIZE);
const uint y = blockDim.y * blockIdx.y + (threadIdx.y % BLOCKSIZE);
if (x < M && y < N) {
float temp = 0.0f;
for (int i = 0; i < K; i++) {
temp += A[x * K + i] * B[i * N + y];
}
A[x * N + y] = alpha * temp + beta * C[x * N + y];
}
In the code above, threads in the same warp share the same x coordinate. As a result, they read the same element of $A$ during each loop iteration, which can be served as a broadcast. Their accesses to $B$—B[i * N + y], B[i * N + (y + 1)], …, B[i * N + (y + 31)]—refer to consecutive memory locations and can therefore be coalesced. Writes to $C$ follow the same consecutive access pattern.
We launch the kernel as follows:
dim3 gridDim(CELL_DIV(M, 32), CELL_DIV(N, 32));
dim3 blockDim(32 * 32);
sgemm_coalescing<<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);
Shared-Memory Tiling
A GPU provides a small, on-chip memory space called shared memory (SMEM). Shared memory is allocated per thread block, allowing threads in the same block to share data and cooperate. It offers much lower latency and higher bandwidth than global memory. In the next kernel, we load a tile of $A$ and a tile of $B$ from global memory into shared memory, then perform most of the computation using the cached data. Each thread still computes one element of $C$.
template <int BLOCKSIZE>
__global__ void matmul_kernel_3(const float *__restrict__ A, const float *__restrict__ B,
float *__restrict__ C, int M, int K, int N, float alpha, float beta) {
__shared__ float As[BLOCKSIZE * BLOCKSIZE];
__shared__ float Bs[BLOCKSIZE * BLOCKSIZE];
// row/col in block
const int threadRow = threadIdx.x;
const int threadCol = threadIdx.y;
// top-left of the C submatrix
const int blockRow = blockIdx.x * BLOCKSIZE;
const int blockCol = blockIdx.y * BLOCKSIZE;
// Position of the element of C computed by the current thread
const int row = blockRow + threadRow;
const int col = blockCol + threadCol;
float temp = 0.0f;
for (int blockId = 0; blockId < K; blockId += BLOCKSIZE) {
const int aCol = blockId + threadCol;
const int bRow = blockId + threadRow;
if (row < M && aCol < K) {
As[threadRow * BLOCKSIZE + threadCol] = A[row * K + aCol];
} else {
As[threadRow * BLOCKSIZE + threadCol] = 0.0f;
}
if (bRow < K && col < N) {
Bs[threadRow * BLOCKSIZE + threadCol] = B[bRow * N + col];
} else {
Bs[threadRow * BLOCKSIZE + threadCol] = 0.0f;
}
__syncthreads();
#pragma unroll
for (int dotIdx = 0; dotIdx < BLOCKSIZE; dotIdx++) {
temp += As[threadRow * BLOCKSIZE + dotIdx] * Bs[dotIdx * BLOCKSIZE + threadCol];
}
__syncthreads();
}
if (row < M && col < N) {
const int index = row * N + col;
C[index] = alpha * temp + beta * C[index];
}
}
1D Thread Tiling: Computing Multiple Results per Thread
If we inspect the executed instructions of the previous implementation in a profiler, we can see that most of them are shared-memory loads.
The inner loop produces PTX instructions similar to the following:
ld.shared.f32 %f91, [%r8+3456];
ld.shared.f32 %f92, [%r7+108];
fma.rn.f32 %f93, %f92, %f91, %f90;
If we inspect the profiler’s warp-state samples, we find that warps spend much of their time in the Stall MIO Throttle state. The profiler describes this state as follows:
A warp was stalled while waiting for space in the MIO (memory input/output) instruction queue. This stall reason is common when the MIO pipelines—which handle shared-memory instructions, special math instructions, and dynamic branches—are heavily utilized.
How can we reduce the number of memory instructions issued by the kernel? One approach is to have each thread compute more than one output element. This allows more of the computation to take place in registers and reduces the pressure on shared memory.
Each thread now computes TM output elements and stores the intermediate results in registers. We add an inner loop that updates all TM results.
The process is illustrated below:
template <int BM, int BK, int BN, int TM>
__global__ void matmul_kernel_4(const float *__restrict__ A, const float *__restrict__ B,
float *__restrict__ C, int M, int K, int N, float alpha, float beta) {
static_assert(BM > 0);
static_assert(BK > 0);
static_assert(BN > 0);
static_assert(TM > 0);
static_assert(BM % TM == 0);
constexpr int NUM_THREADS = BM * BN / TM;
static_assert(NUM_THREADS <= 1024);
__shared__ float As[BM * BK];
__shared__ float Bs[BK * BN];
// The block is one-dimensional
const int tid = threadIdx.x;
// thread position
const int threadRow = tid / BN;
const int threadCol = tid % BN;
// block position
const int blockRow = blockIdx.x * BM;
const int blockCol = blockIdx.x * BN;
float threadResults[TM] = {0.0f};
for (int bkIdx = 0; bkIdx < BK; bkIdx++) {
// load block A into shared memory
for (int index = tid; index < BM * BK; index += NUM_THREADS) {
const int tileRow = tid / BK;
const int tileCol = tid % BK;
const int globalRow = blockRow + tileRow;
const int globalCol = blockCol + tileCol;
if (globalRow < M && globalCol < K) {
As[tileRow * BK + tileCol] = A[globalRow * K + globalCol];
} else {
As[tileRow * BK + tileCol] = 0.0f;
}
}
// load block B into shared memory
for (int index = tid; index < BK * BN; index += NUM_THREADS) {
const int tileRow = tid / BN;
const int tileCol = tid % BN;
const int globalRow = blockRow + tileRow;
const int globalCol = blockCol + tileCol;
if (globalRow < K && globalCol < N) {
Bs[tileRow * BN + tileCol] = B[globalRow * N + globalCol];
} else {
Bs[tileRow * BN + tileCol] = 0.0f;
}
}
__syncthreads();
#pragma unroll
for (int dotIdx = 0; dotIdx < BK; dotIdx++) {
float bTemp = Bs[dotIdx * BN + threadCol];
for (int resIdx = 0; resIdx < TM; resIdx++) {
const int localRow = threadRow * TM + resIdx;
threadResults[resIdx] += A[localRow * BK + dotIdx] * bTemp;
}
}
__syncthreads();
}
#pragma unroll
for (int resIdx = 0; resIdx < TM; resIdx++) {
const int localRow = threadRow * TM + resIdx;
const int globalRow = blockRow + localRow;
const int globalCol = blockCol + threadCol;
if (globalRow < M && globalCol < N) {
int index = globalRow * N + globalCol;
if (beta == 0.0f) {
C[index] = alpha * threadResults[resIdx];
} else {
C[index] = alpha * threadResults[resIdx] + beta * C[index];
}
}
}
}
Increasing Arithmetic Intensity with 2D Thread Tiling
The basic idea of this kernel is to compute TM * TN results per thread. First, all threads cooperate to populate the shared-memory tiles. Each thread then loads TM values from $A$ and TN values from $B$ into registers, computes TM * TN results, and stores them in the output matrix $C$. The process is illustrated below:
// Increasing arithmetic intensity with 2D thread tiling
template <int BM, int BK, int BN, int TM, int TN>
__global__ void matmul_kernel_5(
const float *__restrict__ A,
const float *__restrict__ B,
float *__restrict__ C,
int M,
int K,
int N,
float alpha,
float beta) {
static_assert(BM > 0);
static_assert(BK > 0);
static_assert(BN > 0);
static_assert(TM > 0);
static_assert(TN > 0);
static_assert(BM % TM == 0);
static_assert(BN % TN == 0);
constexpr int THREAD_ROWS = BM / TM;
constexpr int THREAD_COLS = BN / TN;
constexpr int NUM_THREADS = THREAD_ROWS * THREAD_COLS;
static_assert(NUM_THREADS <= 1024);
__shared__ float As[BM * BK];
__shared__ float Bs[BK * BN];
float threadResults[TM * TN] = {0.0f};
float regM[TM] = {0.0f};
float regN[TN] = {0.0f};
const int tid = threadIdx.x;
const int blockRow = blockIdx.x * BM;
const int blockCol = blockIdx.y * BN;
const int threadRow = tid / THREAD_COLS;
const int threadCol = tid % THREAD_COLS;
#pragma unroll
for (int blkId = 0; blkId < K; blkId += BK) {
#pragma unroll
for (int index = tid; index < BM * BK; index += NUM_THREADS) {
int localRow = index / BK;
int localCol = index % BK;
int globalRow = blockRow + localRow;
int globalCol = blkId + localCol;
if (globalRow < M && globalCol < K) {
As[localRow * BK + localCol] = A[globalRow * K + globalCol];
} else {
As[localRow * BK + localCol] = 0.0f;
}
}
#pragma unroll
for (int index = tid; index < BK * BN; index += NUM_THREADS) {
int localRow = index / BN;
int localCol = index % BN;
int globalRow = blkId + localRow;
int globalCol = blockCol + localCol;
if (globalRow < K && globalCol < N) {
Bs[localRow * BN + localCol] = B[globalRow * N + globalCol];
} else {
Bs[localRow * BN + localCol] = 0.0f;
}
}
__syncthreads();
#pragma unroll
for (int dotIdx = 0; dotIdx < BK; dotIdx++) {
#pragma unroll
for (int i = 0; i < TM; i++) {
int localRow = threadRow * TM + i;
int localCol = dotIdx;
regM[i] = As[localRow * BK + localCol];
}
#pragma unroll
for (int i = 0; i < TN; i++) {
int localRow = dotIdx;
int localCol = threadCol * TN + i;
regN[i] = Bs[localRow * BN + localCol];
}
#pragma unroll
for (int i = 0; i < TM; i++) {
#pragma unroll
for (int j = 0; j < TN; j++) {
threadResults[i * TN + j] += regM[i] * regN[j];
}
}
}
__syncthreads();
}
#pragma unroll
for (int i = 0; i < TM; i++) {
#pragma unroll
for (int j = 0; j < TN; j++) {
int localRow = threadRow * TM + i;
int localCol = threadCol * TN + j;
int globalRow = blockRow + localRow;
int globalCol = blockCol + localCol;
if (globalRow < M && globalCol < N) {
int index = globalRow * N + globalCol;
const float result = threadResults[i * TN + j];
C[index] = alpha * result + beta * C[index];
}
}
}
}
Vectorizing Shared- and Global-Memory Access
We can transpose the As tile when loading it into shared memory. In the previous implementation, values such as As[row + 0][col], …, As[row + m][col] had to be loaded individually because they were not stored consecutively. After transposition, the same values are arranged as As[row][col + 0], …, As[row][col + m]. Because they are now consecutive in memory, we can load four elements at a time using vectorized accesses.
// Vectorize SMEM and GMEM Access
template <int BM, int BK, int BN, int TM, int TN>
__global__ void matmul_kernel_6(
const float *__restrict__ A,
const float *__restrict__ B,
float *__restrict__ C,
int M,
int K,
int N,
float alpha,
float beta) {
static_assert(BM > 0);
static_assert(BK > 0);
static_assert(BN > 0);
static_assert(TM > 0);
static_assert(TN > 0);
static_assert(BM % TM == 0);
static_assert(BN % TN == 0);
constexpr int VECTOR_SIZE = 4;
constexpr int A_VECTOR_COLS = BK / VECTOR_SIZE;
constexpr int B_VECTOR_COLS = BN / VECTOR_SIZE;
constexpr int THREAD_ROWS = BM / TM;
constexpr int THREAD_COLS = BN / TN;
constexpr int NUM_THREADS = THREAD_ROWS * THREAD_COLS;
static_assert(BK % VECTOR_SIZE == 0);
static_assert(BN % VECTOR_SIZE == 0);
static_assert(NUM_THREADS <= 1024);
__shared__ float As[BK * BM];
__shared__ float Bs[BK * BN];
float threadResults[TM * TN] = {0.0f};
float regM[TM] = {0.0f};
float regN[TN] = {0.0f};
const int tid = threadIdx.x;
const int blockRow = blockIdx.x * BM;
const int blockCol = blockIdx.y * BN;
const int threadRow = tid / THREAD_COLS;
const int threadCol = tid % THREAD_COLS;
for (int blkId = 0; blkId < K; blkId += BK) {
#pragma unroll
for (int index = tid; index < BM * A_VECTOR_COLS; index += NUM_THREADS) {
const int localRow = index / A_VECTOR_COLS;
const int localVecCol = index % A_VECTOR_COLS;
const int localCol = localVecCol * VECTOR_SIZE;
const int globalRow = blockRow + localRow;
const int globalCol = blkId + localCol;
float4 tmp = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
if (globalRow < M && globalCol + VECTOR_SIZE <= K) {
tmp = *reinterpret_cast<const float4 *>(&A[globalRow * K + globalCol]);
} else if (globalRow < M) {
if (globalCol + 0 < K) {
tmp.x = A[globalRow * K + globalCol + 0];
}
if (globalCol + 1 < K) {
tmp.y = A[globalRow * K + globalCol + 1];
}
if (globalCol + 2 < K) {
tmp.z = A[globalRow * K + globalCol + 2];
}
if (globalCol + 3 < K) {
tmp.w = A[globalRow * K + globalCol + 3];
}
}
As[(localCol + 0) * BM + localRow] = tmp.x;
As[(localCol + 1) * BM + localRow] = tmp.y;
As[(localCol + 2) * BM + localRow] = tmp.z;
As[(localCol + 3) * BM + localRow] = tmp.w;
}
#pragma unroll
for (int index = tid; index < BK * B_VECTOR_COLS; index += NUM_THREADS) {
const int localRow = index / B_VECTOR_COLS;
const int localVecCol = index % B_VECTOR_COLS;
const int localCol = localVecCol * VECTOR_SIZE;
const int globalRow = blkId + localRow;
const int globalCol = blockCol + localCol;
float4 tmp = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
if (globalRow < K && globalCol + VECTOR_SIZE <= N) {
tmp = *reinterpret_cast<const float4 *>(&B[globalRow * N + globalCol]);
} else if (globalRow < K) {
if (globalCol + 0 < N) {
tmp.x = B[globalRow * N + globalCol + 0];
}
if (globalCol + 1 < N) {
tmp.y = B[globalRow * N + globalCol + 1];
}
if (globalCol + 2 < N) {
tmp.z = B[globalRow * N + globalCol + 2];
}
if (globalCol + 3 < N) {
tmp.w = B[globalRow * N + globalCol + 3];
}
}
*reinterpret_cast<float4 *>(&Bs[localRow * BN + localCol]) = tmp;
}
__syncthreads();
#pragma unroll
for (int dotIdx = 0; dotIdx < BK; dotIdx++) {
#pragma unroll
for (int i = 0; i < TM; i++) {
int localRow = threadRow * TM + i;
int localCol = dotIdx;
regM[i] = As[localCol * BM + localRow];
}
#pragma unroll
for (int i = 0; i < TN; i++) {
int localRow = dotIdx;
int localCol = threadCol * TN + i;
regN[i] = Bs[localRow * BN + localCol];
}
#pragma unroll
for (int i = 0; i < TM; i++) {
#pragma unroll
for (int j = 0; j < TN; j++) {
threadResults[i * TN + j] += regM[i] * regN[j];
}
}
}
__syncthreads();
}
#pragma unroll
for (int i = 0; i < TM; i++) {
#pragma unroll
for (int j = 0; j < TN; j++) {
int localRow = threadRow * TM + i;
int localCol = threadCol * TN + j;
int globalRow = blockRow + localRow;
int globalCol = blockCol + localCol;
if (globalRow < M && globalCol < N) {
int index = globalRow * N + globalCol;
const float result = threadResults[i * TN + j];
C[index] = alpha * result + beta * C[index];
}
}
}
}
Warp Tiling
We now introduce another level of tiling between block tiling and thread tiling: warp tiling. There are three main reasons for adding it:
- Warps are the scheduling units assigned to the warp schedulers within an SM.
- Shared-memory bank conflicts occur among threads in the same warp.
- Recent GPUs include a register-file cache, and tighter thread tiling can improve register-cache locality.
// Warp tiling
template <int BM, int BK, int BN, int WM, int WN, int WMITER, int WNITER, int TM, int TN>
__global__ void matmul_kernel_7(
const float *__restrict__ A,
const float *__restrict__ B,
float *__restrict__ C,
int M,
int K,
int N,
float alpha,
float beta) {
static_assert(WM > 0);
static_assert(WN > 0);
static_assert(WMITER > 0);
static_assert(WNITER > 0);
static_assert(BM > 0);
static_assert(BK > 0);
static_assert(BN > 0);
static_assert(TM > 0);
static_assert(TN > 0);
static_assert(BM % TM == 0);
static_assert(BN % TN == 0);
static_assert(WM % WMITER == 0);
static_assert(WN % WNITER == 0);
static_assert(BM % WM == 0);
static_assert(BN % WN == 0);
static_assert(TM % 4 == 0);
static_assert(TN % 4 == 0);
static_assert(BM % 4 == 0);
static_assert(WM % 4 == 0);
static_assert(BN % 4 == 0);
static_assert(WN % 4 == 0);
constexpr int WARP_SIZE = 32;
constexpr int VECTOR_SIZE = 4;
constexpr int A_VECTOR_COLS = BK / VECTOR_SIZE;
constexpr int B_VECTOR_COLS = BN / VECTOR_SIZE;
constexpr int BLOCK_WARP_ROWS = BM / WM;
constexpr int BLOCK_WARP_COLS = BN / WN;
constexpr int BLOCK_NUM_WARPS = BLOCK_WARP_ROWS * BLOCK_WARP_COLS;
constexpr int BLOCK_NUM_THREADS = BLOCK_NUM_WARPS * WARP_SIZE;
constexpr int WSUBM = WM / WMITER;
constexpr int WSUBN = WN / WNITER;
constexpr int WARP_THREAD_ROWS = WSUBM / TM;
constexpr int WARP_THREAD_COLS = WSUBN / TN;
static_assert(BK % VECTOR_SIZE == 0);
static_assert(BN % VECTOR_SIZE == 0);
static_assert(BLOCK_NUM_THREADS <= 1024);
static_assert(WSUBM % TM == 0);
static_assert(WSUBN % TN == 0);
static_assert(WARP_THREAD_ROWS * WARP_THREAD_COLS == WARP_SIZE);
static_assert(WSUBM % 4 == 0);
static_assert(WSUBN % 4 == 0);
__shared__ __align__(16) float As[BK * BM];
__shared__ __align__(16) float Bs[BK * BN];
float threadResults[WMITER * WNITER * TM * TN] = {0.0f};
float regM[TM * WMITER] = {0.0f};
float regN[TN * WNITER] = {0.0f};
const int tid = threadIdx.x;
const int warpId = tid / WARP_SIZE;
const int laneId = tid % WARP_SIZE;
const int blockRow = blockIdx.x * BM;
const int blockCol = blockIdx.y * BN;
const int warpRowInBlock = warpId / BLOCK_WARP_COLS;
const int warpColInBlock = warpId % BLOCK_WARP_COLS;
const int threadRowInWarp = laneId / WARP_THREAD_COLS;
const int threadColInWarp = laneId % WARP_THREAD_COLS;
for (int blkId = 0; blkId < K; blkId += BK) {
#pragma unroll
for (int index = tid; index < BM * A_VECTOR_COLS; index += BLOCK_NUM_THREADS) {
const int localRow = index / A_VECTOR_COLS;
const int localVecCol = index % A_VECTOR_COLS;
const int localCol = localVecCol * VECTOR_SIZE;
const int globalRow = blockRow + localRow;
const int globalCol = blkId + localCol;
float4 tmp = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
if (globalRow < M && globalCol + VECTOR_SIZE <= K) {
tmp = *reinterpret_cast<const float4 *>(&A[globalRow * K + globalCol]);
} else if (globalRow < M) {
if (globalCol + 0 < K) {
tmp.x = A[globalRow * K + globalCol + 0];
}
if (globalCol + 1 < K) {
tmp.y = A[globalRow * K + globalCol + 1];
}
if (globalCol + 2 < K) {
tmp.z = A[globalRow * K + globalCol + 2];
}
if (globalCol + 3 < K) {
tmp.w = A[globalRow * K + globalCol + 3];
}
}
As[(localCol + 0) * BM + localRow] = tmp.x;
As[(localCol + 1) * BM + localRow] = tmp.y;
As[(localCol + 2) * BM + localRow] = tmp.z;
As[(localCol + 3) * BM + localRow] = tmp.w;
}
#pragma unroll
for (int index = tid; index < BK * B_VECTOR_COLS; index += BLOCK_NUM_THREADS) {
const int localRow = index / B_VECTOR_COLS;
const int localVecCol = index % B_VECTOR_COLS;
const int localCol = localVecCol * VECTOR_SIZE;
const int globalRow = blkId + localRow;
const int globalCol = blockCol + localCol;
float4 tmp = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
if (globalRow < K && globalCol + VECTOR_SIZE <= N) {
tmp = *reinterpret_cast<const float4 *>(&B[globalRow * N + globalCol]);
} else if (globalRow < K) {
if (globalCol + 0 < N) {
tmp.x = B[globalRow * N + globalCol + 0];
}
if (globalCol + 1 < N) {
tmp.y = B[globalRow * N + globalCol + 1];
}
if (globalCol + 2 < N) {
tmp.z = B[globalRow * N + globalCol + 2];
}
if (globalCol + 3 < N) {
tmp.w = B[globalRow * N + globalCol + 3];
}
}
*reinterpret_cast<float4 *>(&Bs[localRow * BN + localCol]) = tmp;
}
__syncthreads();
#pragma unroll
for (int dotIdx = 0; dotIdx < BK; dotIdx++) {
#pragma unroll
for (int wSubRow = 0; wSubRow < WMITER; wSubRow++) {
const int baseRow = warpRowInBlock * WM + wSubRow * WSUBM + threadRowInWarp * TM;
#pragma unroll
for (int i = 0; i < TM; i++) {
const int sharedIndex = dotIdx * BM + baseRow + i;
const int registerIndex = wSubRow * TM + i;
const float4 tmp = *reinterpret_cast<const float4 *>(&As[sharedIndex]);
regM[registerIndex + 0] = tmp.x;
regM[registerIndex + 1] = tmp.y;
regM[registerIndex + 2] = tmp.z;
regM[registerIndex + 3] = tmp.w;
}
}
#pragma unroll
for (int wSubCol = 0; wSubCol < WNITER; wSubCol++) {
const int baseCol = warpColInBlock * WN + wSubCol * WSUBN + threadColInWarp * TN;
#pragma unroll
for (int i = 0; i < TN; i++) {
const int sharedIndex = dotIdx * BN + baseCol + i;
const int registerIndex = wSubCol * TN + i;
const float4 tmp = *reinterpret_cast<const float4 *>(&Bs[sharedIndex]);
regN[registerIndex + 0] = tmp.x;
regN[registerIndex + 1] = tmp.y;
regN[registerIndex + 2] = tmp.z;
regN[registerIndex + 3] = tmp.w;
}
}
#pragma unroll
for (int wSubRow = 0; wSubRow < WMITER; wSubRow++) {
#pragma unroll
for (int wSubCol = 0; wSubCol < WNITER; wSubCol++) {
#pragma unroll
for (int i = 0; i < TM; i++) {
#pragma unroll
for (int j = 0; j < TN; j++) {
const int resultRow = wSubRow * TM + i;
const int resultCol = wSubCol * TN + j;
const int resultIndex = resultRow * (WNITER * TN) + resultCol;
threadResults[resultIndex] +=
regM[wSubRow * TM + i] * regN[wSubCol * TN + j];
}
}
}
}
}
__syncthreads();
}
#pragma unroll
for (int wSubRow = 0; wSubRow < WMITER; wSubRow++) {
for (int wSubCol = 0; wSubCol < WNITER; wSubCol++) {
for (int i = 0; i < TM; i++) {
for (int j = 0; j < TN; j++) {
const int localRow =
warpRowInBlock * WM + wSubRow * WSUBM + threadRowInWarp * TM + i;
const int localCol =
warpColInBlock * WN + wSubCol * WSUBN + threadColInWarp * TN + j;
const int globalRow = blockRow + localRow;
const int globalCol = blockCol + localCol;
if (globalRow < M && globalCol < N) {
int index = globalRow * N + globalCol;
const int resultRow = wSubRow * TM + i;
const int resultCol = wSubCol * TN + j;
const int resultIndex = resultRow * (WNITER * TN) + resultCol;
const float result = threadResults[resultIndex];
C[index] = alpha * result + beta * C[index];
}
}
}
}
}
}
Comments
Loading comments...