Naive GEMM vs Shared Memory Tiling GEMM Test

Naive GEMM

__global__ void gemm_naive_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int M, int N, int K)
{
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    if (row >= M || col >= N) return;

    float acc = 0.0f;
    for (int k = 0; k < K; ++k) {
        float a = A[row * K + k];     // row 고정, k만 변함
        float b = B[k * N + col];     // col 고정, k만 변함
        acc += a * b;
    }
    C[row * N + col] = acc;
}

• thread 당 C[row, col] 하나 계산
• K 루프 동안
  • A[row, k], B[k, col] 를 매번 global memory 에서 직접 load
• 같은 block 안의 thread 들이 같이 쓰는 A/B 데이터도 공유 없이 각자 다시 읽음

Shared Memory Tiling GEMM 

__global__ void gemm_shared_tiling_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int M, int N, int K)
{
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    __shared__ float As[BLOCK_SIZE][BLOCK_SIZE];
    __shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE];

    float acc = 0.0f;
    int numTiles = (K + BLOCK_SIZE - 1) / BLOCK_SIZE;

    for (int t = 0; t < numTiles; ++t) {
        int kA = t * BLOCK_SIZE + threadIdx.x; // A: column index
        int kB = t * BLOCK_SIZE + threadIdx.y; // B: row index

        if (row < M && kA < K)
            As[threadIdx.y][threadIdx.x] = A[row * K + kA];
        else
            As[threadIdx.y][threadIdx.x] = 0.0f;

        if (kB < K && col < N)
            Bs[threadIdx.y][threadIdx.x] = B[kB * N + col];
        else
            Bs[threadIdx.y][threadIdx.x] = 0.0f;

        __syncthreads();

        #pragma unroll
        for (int k = 0; k < BLOCK_SIZE; ++k) {
            acc += As[threadIdx.y][k] * Bs[k][threadIdx.x];
        }

        __syncthreads();
    }

    if (row < M && col < N)
        C[row * N + col] = acc;
}

• 각 block 이 C 의 32 x 32 tile 담당
• K 방향을 BLOCK_SIZE 단위로 쪼개서 numTiles 만큼 루프
• 루프마다
  • As : A 의 (row, kA) 구간을 tile 로 로드
  • Bs : B 의 (kA, col) 구간을 tile 로 로드
  • 한 번 로르된 tile 을 block 내 32 x 32 threads 가 모두 재사용
• tile 당 한 번만 global load, 

 

GEMM config: C[1024 x 1024] = A[1024 x 1024] * B[1024 x 1024]
BLOCK_SIZE = 32
Computing host reference...

[naive] warm-up + timing
[naive] kernel time: 2.163 ms
[naive] max |C_ref - C_naive| = 7.450581e-08

[shared_tiling] warm-up + timing
[shared_tiling] kernel time: 1.758 ms
[shared_tiling] max |C_ref - C_tiled| = 7.450581e-08

 

성능 향상 관측 확