GEMM 커널 테스트 코드 작성 및 최적화 (1024^3 크기 실험)

실험 목적

이 실험의 목적은 다음 두 가지

1. 행렬 곱 GEMM 연산의 실제 동작 방식을 코드 수준에서 이해
2. naive 구현 vs shared memory tiled 구현의 성능 차이를 수치로 확인한다.

다음의 설정 사용

• 행렬 크기
  • M = 1024, N = 1024, K = 1024
  • A (M, K), B (K, N), C (M, N)
• 데이터 타입 : float (FP32)
• 반복 횟수 : iters = 50
• 타일 크기 : TILE = 32

측정 결과

[naive]  avg time: 2.9437 ms
[tiled]  avg time: 2.5054 ms
tiled / naive = 0.8511x
max |diff| = 0.000000e+00

 

수학적 정의 

C = A × B

각 원소 C_ij 의 계산

연산량(FLOPs) = 2.147 * 10^9 FLOPs

곱셈 : M * N * K

덧셈 : M * N * k

보통 GEMM FLOPs ~ = 2 * MNK

 

메모리 레이아웃

코드는 row-major 사용

• A: 크기 M×K, 인덱싱: A[row * K + k]
  • “row번째 행의 k번째 열”에 해당
• B: 크기 K×N, 인덱싱: B[k * N + col]
  • “k번째 행의 col번째 열”
• C: 크기 M×N, 인덱싱: C[row * N + col]
  • “row번째 행의 col번째 열”

이 구조에 맞춰 커널에서 row/col/k를 1D 인덱스로 변환해서 접근한다.

 

Naive GEMM 커널 동작 방식

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

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

    T acc = T(0);
    int a_row_offset = row * K;
    int b_col_offset = col;
    for (int k = 0; k < K; ++k) {
        acc += A[a_row_offset + k] * B[k * N + b_col_offset];
    }
    C[row * N + col] = acc;
}

쓰레드 매핑

• 블록 크기 : (blockDim.x, blockDim.y) = (16, 16)
• 그리드 크기

dim3 block(16, 16);
dim3 grid((N + block.x - 1) / block.x,
          (M + block.y - 1) / block.y);

• 한 스레드가 C 의 한 원소 C[row, col] 을 계산
  • row = blockIdx.y * 16 + threadIdx.y
  • col = blockIdx.x * 16 + threadIdx.x

즉, 16*16 블록 하나가 C 의 16*16 서브타일 하나를 계산한다.

연산 흐름

acc = 0 초기화

k 루프를 돌면서

for (int k = 0; k < K; ++k) {
    acc += A[row, k] * B[k, col];
}

루프가 끝나면 acc 가 해당 C[row, col] 의 최종 값이 되고 메모리에 기록

 

메모리 접근 특성

A [row * K + k]

같은 row 에서 k 가 0 -> K-1 로 증가 -> row 안에서 연속 접근

워프 내 neighboring threads 는 다른 row 에 있으므로, A 는 워프 기준으로는 크게 coalesced 하지는 않지만, L1/L2캐시가 도와줄 수 있음

 

B[k * N + col]

같은 k 에 대해 col 이 0 -> N-1 로 변함 -> warp 가 열 방향으로 나뉘어 있으면 coalesced 접근이 잘 나온다.

 

모든 스레드가 매번 global memory 에서 A, B 를 직접 읽는다.

데이터 재사용은 캐시가 알아서 해주는 수준

shared memory 를 사용하지 않는 가장 단순한 GEMM 구조

 

 

Shared Memory Tiled GEMM 연산 방식

template<typename T, int TILE>
__global__ void gemm_tiled_kernel(
    int M, int N, int K,
    const T* __restrict__ A,
    const T* __restrict__ B,
    T* __restrict__ C)
{
    __shared__ T As[TILE][TILE];
    __shared__ T Bs[TILE][TILE];

    int row = blockIdx.y * TILE + threadIdx.y;
    int col = blockIdx.x * TILE + threadIdx.x;

    T acc = T(0);

    int num_tiles = (K + TILE - 1) / TILE;

    for (int t = 0; t < num_tiles; ++t) {
        int kA = t * TILE + threadIdx.x; // A: (row, kA)
        int kB = t * TILE + threadIdx.y; // B: (kB, col)

        if (row < M && kA < K) {
            As[threadIdx.y][threadIdx.x] = A[row * K + kA];
        } else {
            As[threadIdx.y][threadIdx.x] = T(0);
        }

        if (kB < K && col < N) {
            Bs[threadIdx.y][threadIdx.x] = B[kB * N + col];
        } else {
            Bs[threadIdx.y][threadIdx.x] = T(0);
        }

        __syncthreads();

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

        __syncthreads();
    }

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

타일 개념

블록 크기 : ( TILE, TILE )  (TILE = 32)

한 블록은 A, B 의 TILE * TILE 서브 행렬 을 shared memory 에 가져와서, C 의 TILE * TILE 블록을 계산

 

K 축을 따라 타일 단위로 잘라서 반복하는 구조

 

타일 로딩 단계

int kA = t * TILE + threadIdx.x; // A(row, kA)
int kB = t * TILE + threadIdx.y; // B(kB, col)

각 스레드는 

A 에서 한 원소를 읽어 As[ty][tx] 에 저장

B 에서 한 원소를 읽어 Bs[ty][tx] 에 저장

패딩된 영역은 0으로 채운다.

이후 스레드가 타일 로딩을 끝낼 때까지 동기화

 

한 타일이 shared 에 준비되면

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

As 의 한 행 , Bs 의 한 열을 따라 dot product 수행

타일별로 이 루프를 반복하면서 acc 를 쌓는다.

 

핵심 아이디어 : shared memory 를 통한 재사용

Naive 는 C 의 한 원소마다 A, B 원소를 global 에서 직접 읽음

tiled 는 한 타일에서 A/B 의 tile * tile 원소를 한 번 global 에서 읽어 shared 에 저장,

이후 같은 타일 내에서 tile 번씩 재사용하여 곱/합 수행

 

shared 용량 사용 및 .__syncthreads() 오버헤드 존재

 

 

Host 코드와 타이밍 방식

데이터 초기화

for (int i = 0; i < M * K; ++i)
    hA[i] = (float)((i % 13) - 6);  // -6 ~ 6

for (int i = 0; i < K * N; ++i)
    hB[i] = (float)((i % 7) - 3);   // -3 ~ 3

재현 가능한 패턴으로 초기화

 

cudaEvent_t 를 사용해 커널 실행 시간을 측정

 

결과 해석

[naive]  avg time: 2.9437 ms
[tiled]  avg time: 2.5054 ms
tiled / naive = 0.8511x

전체 FLOPs ≈ 2.147e9.

• naive:
  • 시간: 2.9437 ms = 0.0029437 s
  • 성능 ≈ 2.147×109/0.0029437≈7302.147 \times 10^9 / 0.0029437 ≈ 7302.147×109/0.0029437≈730 GFLOP/s
• tiled:
  • 시간: 2.5054 ms = 0.0025054 s
  • 성능 ≈ 2.147×109/0.0025054≈8572.147 \times 10^9 / 0.0025054 ≈ 8572.147×109/0.0025054≈857 GFLOP/s

 

공간 재사용(shared memory 타일링) 덕분에 약 17% 성능 향상 (0.8511×)