스칼라 FMA, WMMA 방식의 GEMM 연산 분해 방식의 차이 이해

1. 스칼라 FMA 관점 (1 thread <-> C 셀 1개)

기존에 구현한 타일 GEMM 의 느낌

// thread 하나가 C[row, col] 한 칸 담당
float acc = 0;
for (int k = 0; k < K; ++k) {
    acc += A[row, k] * B[k, col];
}
C[row, col] = acc;

• A[row, k] 와 B[k, col] 셀 1쌍 - > C[row, col] 셀에 1번 기여
• 스레드는 C 의 셀 1개를 들고, k 루프를 돌면서 1:1로 누적
• 타일 알고리즘을 써도 구조는 동일

C 의 특정 위치 1칸 기준으로 연산이 돌아가는 느낌

 

2. WMMA 방식 : 타일/행렬 연산 관점 (1 warp <-> C 타일, 1 lane <-> 여러 셀)

WMMA/Tensor Core 관점에선 바뀜

// warp 전체가 16x16 C 타일 하나 담당
wmma::fragment<matrix_a, 16,16,16, half>   a_frag;
wmma::fragment<matrix_b, 16,16,16, half>   b_frag;
wmma::fragment<accumulator, 16,16,16,float> c_frag;

wmma::fill_fragment(c_frag, 0.0f);

for (kk = 0; kk < K; kk += 16) {
    // A, B 에서 16x16 타일씩 로드
    wmma::load_matrix_sync(a_frag, A_tile_ptr, lda);
    wmma::load_matrix_sync(b_frag, B_tile_ptr, ldb);

    // 여기서 한 번에 16x16x16 "블록" 곱셈/누적
    wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
}

wmma::store_matrix_sync(C_tile_ptr, c_frag, ldc, mem_row_major);

• warp 전체가 C 의 16, 16 타일 하나를 통째로 맡는다.
• mma_sync 한 번 호출할 때
  • C 타일의 256 개 셀에 대한 부분합을 동시에 업데이트
• lane 관점에서 보면
  • lane 하나가 C 타일 안의 8 개 셀을 들고 있고
  • mma_sync 가 호출될 때마다 그 8개 셀의 값이 한 번에 업데이트 됨
  • 그 안에는 4의 곱의 부분합이 한 번에 들어오는 느낌

즉, WMMA 에선 

더 이상 A, B 한 쌍 -> C 1칸을 스레드 수준에서 직접 돌리지 않고, 16, 16, 16 블록 단위로 처리하게 Tensor Core 에 맡긴다고 보면 된다.