Naive GEMM Kernel

GEMM은 dense layer와 Transformer projection의 중심 연산이다.

C[M, N] = A[M, K] x B[K, N]

naive CUDA GEMM에서는 thread 하나가 C[row, col] 하나를 계산한다.

float acc = 0.0f;
for (int kk = 0; kk < K; kk++) {
    acc += A[row * K + kk] * B[kk * N + col];
}
C[row * N + col] = acc;

2D grid

output C가 2D이므로 2D grid를 쓴다.

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

row는 C의 행, col은 C의 열이다.

Naive kernel

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

    if (row < M && col < N) {
        float acc = 0.0f;
        for (int kk = 0; kk < K; kk++) {
            acc += A[row * K + kk] * B[kk * N + col];
        }
        C[row * N + col] = acc;
    }
}

아직 최적화하지 않는다

이 kernel은 느리다. 같은 A, B 원소를 global memory에서 반복해서 읽는다.

하지만 지금 목표는 빠른 GEMM이 아니다.

목표:
  행렬곱 수식을 CUDA row/col/index로 옮길 수 있다.

나중 목표:
  shared memory tiling, block matmul, tensor core

확인

• C[row, col] 하나를 만들 때 어떤 A row와 B column이 필요한가?
• A[row * K + kk]에서 K는 왜 stride인가?
• block matmul을 이 카드에서 다루지 않는 이유는 무엇인가?