Register/Shared Memory Tiling

개요

Register/Shared Memory Tiling은 GPU에서 GEMM(Generic Matrix Multiply) 연산을 극한 성능으로 수행하기 위한 메모리 계층 기반 타일링 기법이다. Square Tiling이 캐시 수준에서 타일을 분할하는 개념이라면, Register/Shared Memory Tiling은 글로벌 메모리 → 공유 메모리(SMEM) → 레지스터까지 세 계층에 걸쳐 데이터 흐름을 체계적으로 관리하여 연산 재사용(Operand Reuse)을 극대화한다.

이 기법의 핵심은 각 스레드가 단일 출력 요소가 아닌 TM×TN 크기의 작은 부분 행렬(sub-matrix)을 계산하는 "아웃터 프로덕트(Outer Product)" 패턴이다. 공유 메모리에서 타일을 로드한 후, 각 스레드가 자신의 레지스터에 A의 행 조각과 B의 열 조각을 저장하고, TM×TN 누산기에서 외적 연산을 수행한다. 이를 통해 스레드당 연산 밀도(Arithmetic Intensity)를 극적으로 높이고, 메모리 대역폭 병목을 해소한다.

cuBLAS, CUTLASS, Triton 등 모든 고성능 GEMM 라이브러리가 이 타일링 전략을 기본으로 채택하고 있으며, NVIDIA A100/H100에서 FP16 텐서 코어 기준 90% 이상의 피크 성능을 달성할 수 있게 하는 핵심 최적화 기법이다.

핵심 개념

메모리 계층과 타일링

GPU 메모리 계층은 세 가지 주요 수준으로 구성된다:

각 계층별 타일 크기 파라미터:

• BM, BN: 공유 메모리에 로드할 A와 B의 타일 크기 (높이 × 너비)
• BK: K 차원 방향으로 한 번에 처리할 깊이
• TM, TN: 각 스레드가 레지스터에서 계산하는 출력 서브 매트릭스 크기

아웃터 프로덕트 GEMM

기존 Square Tiling이 한 스레드에 대해 한 번에 하나의 요소를 계산하는 반면, Register Tiling은 아웃터 프로덕트 방식으로 작동한다:

for k = 0 to K step BK:
    A_tile[BM×BK]를 공유 메모리로 로드
    B_tile[BK×BN]를 공유 메모리로 로드

    for t = 0 to BK:
        a = A_tile[:, t]    → BM개 요소를 레지스터로 로드
        b = B_tile[t, :]    → BN개 요소를 레지스버로 로드

        C_reg[i][j] += a[i] × b[j]  (외적 연산)

각 k 반복에서 BM개의 A 요소와 BN개의 B 요소를 로드하여 TM×TN 크기의 누산기를 업데이트한다. 이를 통해 연산 대 메모리 비율(Compute-to-Memory Ratio)이 BM×BN/BK로 증가한다.

공유 메모리 타일 로딩

블록의 모든 스레드가 협력적으로 글로벌 메모리에서 공유 메모리로 타일을 로드한다:

// 타일 로딩 코디네이션
int thread_id = threadIdx.x + threadIdx.y * blockDim.x;
int total_threads = blockDim.x * blockDim.y;

// A[BM×BK] 타일을 total_threads로 균등 분할 로딩
for (int i = thread_id; i < BM * BK; i += total_threads) {
    int row = i / BK;
    int col = i % BK;
    smem_A[row][col] = global_A[block_row * BM + row][k_tile * BK + col];
}

// B[BK×BN] 타일도 동일하게 로딩
for (int i = thread_id; i < BK * BN; i += total_threads) {
    int row = i / BN;
    int col = i % BN;
    smem_B[row][col] = global_B[k_tile * BK + row][block_col * BN + col];
}
__syncthreads();

더블 버퍼링

공유 메모리 타일 로딩과 연산을 겹쳐 실행하여 지연 시간을 숨기는 기법이다:

1. 버퍼 0: 현재 타일 연산 수행
2. 버퍼 1: 다음 타일을 글로벌 메모리에서 로딩
3. 다음 반복에서 버퍼 교체

이를 통해 메모리 로딩 지연을 연산으로 숨길 수 있다.

비교/분석

타일링 전략 비교

데이터플로우 패턴 비교

파라미터 튜닝 영향

점유율(Occupancy) 분석

레지스터/스레드 = TM + TN + TM×TN  (입력 조각 + 누산기)
블록당 스레드 = (BM/TM) × (BN/TN)
SMEM/블록 = BM×BK×4 + BK×BN×4  (FP32 기준)

예시 (BM=128, BN=128, BK=16, TM=8, TN=8):
- 레지스터/스레드: 8 + 8 + 64 = 80개
- 블록당 스레드: 16 × 16 = 256개
- SMEM/블록: 128×16×4 + 16×128×4 = 16KB
- SM당 블록 수: 제한적 (레지스터 256KB 기준 최대 40블록)

동작 원리

전체 알고리즘 구조

// BM=128, BN=128, BK=16, TM=8, TN=8 기준 GEMM 커널
__global__ void gemm_tiled(float* A, float* B, float* C, int M, int N, int K) {
    __shared__ float smem_A[BM][BK];  // 128×16 = 2KB
    __shared__ float smem_B[BK][BN];  // 16×128 = 2KB

    // 각 스레드의 TM×TN 출력 타일
    float C_reg[TM][TN] = {0.0f};
    float A_reg[TM];
    float B_reg[TN];

    int bx = blockIdx.x;
    int by = blockIdx.y;

    // 1. 타일 루프: K 차원을 BK 단위로 순회
    for (int k_tile = 0; k_tile < K; k_tile += BK) {

        // 2. 협력적 타일 로딩 (글로벌 → 공유 메모리)
        load_tile_to_smem(A, smem_A, by, k_tile);
        load_tile_to_smem(B, smem_B, k_tile, bx);
        __syncthreads();

        // 3. 내부 루프: BK 깊이만큼 외적 연산
        for (int k = 0; k < BK; k++) {
            // A 조각 로드: smem_A에서 TM개 요소
            for (int i = 0; i < TM; i++) {
                A_reg[i] = smem_A[threadRow * TM + i][k];
            }
            // B 조각 로드: smem_B에서 TN개 요소
            for (int j = 0; j < TN; j++) {
                B_reg[j] = smem_B[k][threadCol * TN + j];
            }
            // 외적 연산: TM×TN 누산기 업데이트
            for (int i = 0; i < TM; i++) {
                for (int j = 0; j < TN; j++) {
                    C_reg[i][j] += A_reg[i] * B_reg[j];
                }
            }
        }
        __syncthreads();
    }

    // 4. 결과 저장 (레지스터 → 글로벌 메모리)
    for (int i = 0; i < TM; i++) {
        for (int j = 0; j < TN; j++) {
            C[(by * BM + threadRow * TM + i) * N + bx * BN + threadCol * TN + j]
                = C_reg[i][j];
        }
    }
}

메모리 접근 패턴 분석

블록당 글로벌 메모리 로딩량:
- A: BM×K = 128×K 요소
- B: K×BN = K×128 요소
- 총: 256K 요소

블록당 연산량:
- BM×BN×K = 128×128×K = 16384K FLOP

연산 대 메모리 비율:

FLOP/Byte = 16384K / (256K × 4) = 16 FLOP/Byte

Square Tiling 대비 16배 향상된 연산 밀도를 달성한다.

레지스터 레벨 외적 연산 상세

k=0:  A_reg[0..7] × B_reg[0..7] → C_reg[0..7][0..7]  (64 FLOP)
k=1:  A_reg[0..7] × B_reg[0..7] → C_reg[0..7][0..7]  (64 FLOP)
...
k=15: A_reg[0..7] × B_reg[0..7] → C_reg[0..7][0..7]  (64 FLOP)
총: 16 × 64 = 1024 FLOP per thread

각 스레드가 BK=16번의 외적 연산으로 TM×TN=64개 출력 요소를 누산한다.

컨티어 버퍼링 구현

__shared__ float smem_A[2][BM][BK];  // 두 개의 A 타일 버퍼
__shared__ float smem_B[2][BK][BN];  // 두 개의 B 타일 버퍼

// 첫 번째 타일 미리 로딩
load_tile_to_smem(A, smem_A[0], by, 0);
load_tile_to_smem(B, smem_B[0], 0, bx);
__syncthreads();

for (int k_tile = 0; k_tile < K; k_tile += BK) {
    int curr = (k_tile / BK) % 2;
    int next = 1 - curr;

    // 현재 버퍼에서 연산
    compute_outer_product(smem_A[curr], smem_B[curr], C_reg);

    // 다음 타일을 다른 버퍼로 미리 로딩
    if (k_tile + BK < K) {
        load_tile_to_smem(A, smem_A[next], by, k_tile + BK);
        load_tile_to_smem(B, smem_B[next], k_tile + BK, bx);
    }
    __syncthreads();
}

장단점

장점

1. 극대화된 연산 밀도: TM×TN 크기의 외적 연산을 통해 스레드당 FLOP을 대폭 증가시켜 메모리 대역폭 병목을 최소화한다
2. 계층적 데이터 재사용: 글로벌 → 공유 메모리 → 레지스터 세 계층에서 각각 BM×BN/BK, BM/BK, 1배의 데이터 재사용을 달성한다
3. 텐서 코어와의 자연스러운 결합: NVIDIA Tensor Core의 WMMA/MMA 명령과 타일 구조가 호환되어 혼합 정밀도 연산에서 추가적인 성능 향상을 제공한다
4. 유연한 파라미터 조정: BM/BN/BK/TM/TN을 하드웨어 특성에 맞게 조정하여 다양한 아키텍처에 최적화할 수 있다
5. 높은 실용성: cuBLAS, CUTLASS, Triton 등 검증된 구현체가 존재하여 프로덕션 환경에서 즉시 활용 가능하다

단점

1. 구현 복잡도: 메모리 계층 간 동기화, 레지스터 할당, 타일 로딩 순서 등 고려할 요소가 많아 디버깅이 어렵다
2. 레지스터 압박(Register Pressure): TM×TN 누산기가 많은 레지스터를 소모하여, 점유율(Occupancy)이 낮아질 수 있다. 스레드 수 감소는 병렬성을 저하시킨다
3. 공유 메모리 제약: SMEM 크기가 제한적이므로 타일 크기를 과도하게 키울 수 없으며, 뱅크 충돌(Bank Conflict) 회피를 위한 패딩이 필요하다
4. 비정형 행렬 처리: 행렬 차원이 타일 크기로 나누어떨어지지 않으면 패딩이 발생하여 연산 낭비와 복잡도 증가를 초래한다
5. 하드웨어 의존성: 최적 타일 크기는 GPU 아키텍처(Ampere, Hopper 등)에 따라 달라지므로 이식성이 떨어진다

관련 기술

프레임워크 및 라이브러리

• NVIDIA cuBLAS: 아키텍처별 최적화된 GEMM 커널 집합으로, Register/Shared Memory Tiling을 자동 선택하여 사용
• NVIDIA CUTLASS: 템플릿 기반 GEMM 라이브러리로, BM/BN/BK/TM/TN 파라미터를 컴파일 시점에 설정하여 다양한 타일링 전략 구현
• OpenAI Triton: 고수준 언어로 타일링을 기술하면 컴파일러가 자동으로 SMEM/레지스터 할당 및 동기화를 생성
• FlashAttention: 타일링을 어텐션 연산에 적용하여 메모리 O(N²) → O(N)으로 감소
• XGemm/CK (Composable Kernel): AMD GPU용 타일링 GEMM 라이브러리

관련 개념

• Bank Conflict: 공유 메모리의 32개 뱅크에 대한 동시 접근 충돌. 패딩(+4) 또는 스위즐(Swizzle)로 해결
• Warp-level GEMM: 워프 내 32개 스레드가 협력적으로 큰 타일을 처리하는 기법
• Tensor Core: NVIDIA GPU의 매트릭스 곱셈 전용 유닛. WMMA/MMA 인터페이스로 접근
• Pipeline Prefetching: 다음 타일을 비동기로 로딩하여 현재 연산과 겹치게 하는 기법
• Occupancy: SM당 활성 스레드 수. 레지스터/SMEM 사용량에 반비례

참고 문헌

• "Anatomy of High-Performance Many-Threaded Matrix Multiplication" - Soliman & Rodriguez (2019)
• "CUTLASS: CUDA Templates for Linear Algebra Subroutines" - NVIDIA
• "Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations" - Tillet et al. (2019)
• "Optimizing Sgemm on NVIDIA GPUs" - NVIDIA Developer Blog
• "FlashAttention: Fast and Memory-Efficient Exact Attention" - Dao et al. (2022)
• "Efficient Memory Management for Large Language Model Serving with PagedAttention" - Kwon et al. (2023)

핵심 정리

1. Register/Shared Memory Tiling은 글로벌 메모리 → 공유 메모리 → 레지스터 세 계층에서 타일을 체계적으로 관리하는 GPU GEMM 최적화 기법이다
2. 각 스레드가 TM×TN 크기의 출력 서브 매트릭스를 외적(Outer Product) 방식으로 계산하여, 스레드당 연산 밀도를 극대화한다
3. 공유 메모리 타일은 블록 내 모든 스레드가 협력적으로 로딩하여 데이터 재사용을 달성하고, 레지스터 타일은 스레드 내부에서 연산을 수행한다
4. 더블 버퍼링을 통해 타일 로딩과 연산을 겹쳐 실행하여 메모리 지연 시간을 숨길 수 있다
5. cuBLAS, CUTLASS, Triton 등 모든 고성능 GEMM 라이브러리가 이 전략을 기본으로 채택하며, NVIDIA 텐서 코어와 결합하여 피크 성능의 90% 이상을 달성할 수 있다