Micro Kernel 입문하기 - 레지스터 기반 1 스레드용 소형 GEMM 엔진 구현

1. 마이크로 커널 실험의 핵심 목표

초기 입문에서는 스레드 하나가 작은 C 블록을 레지스터만으로 계산하도록 하는 것을 목표로 한다.

• 1 thread 가 C(4x4) 블록을 직접 계산
• A 는 4 x K, B 는 K x 4 형태로 레지스터에 배치
• K 루프마다 A_frag, B_frag 불러와서 FMA 수행
• 공유메모리 / 전역메모리 없이 오직 레지스터 기반으로 C 누적

이걸 구현하면

• 레지스터 타일 layout
• inner FMA pipeline
• unroll factor
• micro-kernel 이라는 개념

위 감각을 잡을 수 있음

 

2. 실행 가능한 Micro-GEMM 4x4 실험 코드

CUDA 디바이스에서 완벽히 동작하고, 작동 방식이 단순한 코드

// micro_gemm_4x4.cu
#include <cstdio>

__global__ void micro_gemm_4x4(const float* __restrict__ A,
                               const float* __restrict__ B,
                               float* __restrict__ C,
                               int K)
{
    // C 4x4 block in registers
    float c00=0, c01=0, c02=0, c03=0;
    float c10=0, c11=0, c12=0, c13=0;
    float c20=0, c21=0, c22=0, c23=0;
    float c30=0, c31=0, c32=0, c33=0;

    for (int k = 0; k < K; k++) {
        // load A(4) and B(4) into registers
        float a0 = A[0*K + k];
        float a1 = A[1*K + k];
        float a2 = A[2*K + k];
        float a3 = A[3*K + k];

        float b0 = B[k*4 + 0];
        float b1 = B[k*4 + 1];
        float b2 = B[k*4 + 2];
        float b3 = B[k*4 + 3];

        // Micro-kernel FMAs (4×4)
        c00 += a0*b0; c01 += a0*b1; c02 += a0*b2; c03 += a0*b3;
        c10 += a1*b0; c11 += a1*b1; c12 += a1*b2; c13 += a1*b3;
        c20 += a2*b0; c21 += a2*b1; c22 += a2*b2; c23 += a2*b3;
        c30 += a3*b0; c31 += a3*b1; c32 += a3*b2; c33 += a3*b3;
    }

    // single thread writes C
    C[0] = c00; C[1] = c01; C[2] = c02; C[3] = c03;
    C[4] = c10; C[5] = c11; C[6] = c12; C[7] = c13;
    C[8] = c20; C[9] = c21; C[10]= c22; C[11]= c23;
    C[12]= c30; C[13]= c31; C[14]= c32; C[15]= c33;
}

int main() {
    int K = 1024;

    float *A, *B, *C;
    cudaMallocManaged(&A, sizeof(float)*4*K);
    cudaMallocManaged(&B, sizeof(float)*K*4);
    cudaMallocManaged(&C, sizeof(float)*16);

    // init
    for (int i=0;i<4*K;i++) A[i] = (i%7)*0.1f;
    for (int i=0;i<K*4;i++) B[i] = (i%5)*0.2f;

    micro_gemm_4x4<<<1,1>>>(A,B,C,K);
    cudaDeviceSynchronize();

    printf("C[0..3]: %f %f %f %f\n", C[0],C[1],C[2],C[3]);

    cudaFree(A); cudaFree(B); cudaFree(C);
}

중요한 부분

• A, B, C 모두 레지스터 기반
• K 루프를 완전히 thread 단위 micro-kernel 로 구현
• 공유 메모리, 타일링 없음 - 본질만 남김
• 실제 GEMM 커널의 inner kernel 의 1:1 축소판

이 코드 자체가 micro-kernel 개념의 가장 원초적 형태

 

3. 이 실험으로 얻는 감각

1) 마이크로 커널의 기본 구조

• 작은 C 타일 유지
• A/B fragment 로딩 - FMA 반복 - 레지스터 누적
• GEMM 내부 루프의 본질

2) 레지스터 타일 설계의 의미

• 왜 c00, c01, c02 ... c33 가 필요한지
  • 마이크로 커널은 곧 accumulator layout 설계

3) unroll의 가치

• 나중에 unroll 4/8 을 적용하면 속도 체감

4) 스레드 1개 vs warp 1개 vs block 1개의 차이

• 마이크로 커널은 스레드 하나의 수학 엔진, warp 타일링보다 더 미세한 설계

 

4. 다음 단계 진행 계획 설계

1. 단일 thraed 4 x 4 micro-kernel (현재 코드)
2. 단일 warp 가 16 x 16 블록을 micro-kernel 조합으로 계산, 기존 커널에 연결 가능
3. cp.async 기반 pipeline 을 micro-kernel 에 넣기
4. micro-kernel 크기 튜닝, MR, NR = (4x4),(8x4),(16x8) 등
5. regN 커널에 마이크로 커널 삽입 - 성능 20~40% 향상 기대 가능??!!

 

출력 결과

C[0..3]: 122.720276 122.460274 122.700279 122.740273

1. 레지스터 기반 4x4 마이크로 GEMM 계산이 정상 수행됨
2. 전역 메모리 - 레지스터 - FMA 반복 - 레지스터 누적 - 출력 
   1. 위 흐름이 제대로 돌아감
3. CUDA 커널 런타임, launch config, 메모리 접근이 올바름

지금 환경에서 마이크로 커널을 실험할 수 있는 준비가 된 것

 

현재는 스레드 1 개가 C(4x4) 만 계산한 형

 

1. 커널 시그니처

__global__ void micro_gemm_4x4(const float* __restrict__ A,
                               const float* __restrict__ B,
                               float* __restrict__ C,
                               int K)

• __global__ : 디바이스에서 실행, 호스트에서 호출하는 CUDA 커널
• 인자
  • A : 크기 4 x K 행렬
  • B : 크기 K x 4 행렬
  • C : 크기 4 x 4 결과 행렬
  • K : 내부 공통 차원
• __register__ : 포인터 alias 없다고 힌트 - 컴파일러 최적화에 도움

 

2. C(4x4) 를 레지스터에 잡는 부분

    // C 4x4 block in registers
    float c00=0, c01=0, c02=0, c03=0;
    float c10=0, c11=0, c12=0, c13=0;
    float c20=0, c21=0, c22=0, c23=0;
    float c30=0, c31=0, c32=0, c33=0;

• cij = C(i,j) 원소 하나에 해당하는 accumulator
• 전부 float 지역 변수 -> GPU 레지스터에 올라감
• micro-kernel 핵심 포인트
  • C 타일 (4x4) 을 전부 레지스터에 유지하면서 K 루프 동안 누적

GEMM 수식으로 보면

• C[i, j] = sum_{k=0,..K-1} A[i,k] * B[k,j]
• 그 sum 을 담는 그릇이 각각 cij

 

3. K 루프 - 마이크로 커널의 내부 엔진

for (int k = 0; k < K; k++) {

• k = 0...K-1 까지 반복
• 일반 GEMM 에서 k 축을 도는 inner loop 와 동일
• 여기서 매번 A 의 열, B 의행 일부를 레지스터에 올리고 FMA

 

3-1. A의 4개 원소를 레지스터로 로드

        // load A(4) and B(4) into registers
        float a0 = A[0*K + k];
        float a1 = A[1*K + k];
        float a2 = A[2*K + k];
        float a3 = A[3*K + k];

A 의 메모리 레이아웃 가정

• A 크기 : 4 x K
• row-major, leading dimension = K
• 인덱싱 : A[row * K + col]

그러면

• A[0*K + k] -> A(0, k)
• A[1*K + k] -> A(1, k)
• A[2*K + k] -> A(2, k)
• A[3*K + k] -> A(3, k)

즉, k 고정, row 0~3 로우 하나씩 - A 의 열 k 를 통째로 4개 원소 읽는 것

결과

• a0 = A(0, k)
• a1 = A(1, k)
• a2 = A(2, k)
• a3 = A(3, k)

 

3-2 B 의 4개 원소를 레지스터로 로드

        float b0 = B[k*4 + 0];
        float b1 = B[k*4 + 1];
        float b2 = B[k*4 + 2];
        float b3 = B[k*4 + 3];

B 의 메모리 레이아웃

• B 크기 : K x 4
• row-major, leading demension = 4
• 인덱싱: B[row * 4 + col]

여기서 row = k, col = 0..3:

• B[k*4 + 0] -> B(k, 0)
• B[k*4 + 1] -> B(k, 1)
• B[k*4 + 2] -> B(k, 2)
• B[k*4 + 3] -> B(k, 3)

즉, k 번째 행의 4개 열을 한 번에 레지스터로 로드

결과

• b0 = B(k, 0)
• b1 = B(k, 1)
• b2 = B(k, 2)
• b3 = B(k, 3)

 

3-3. 4x4 마이크로 커널 FMA

        // Micro-kernel FMAs (4×4)
        c00 += a0*b0; c01 += a0*b1; c02 += a0*b2; c03 += a0*b3;
        c10 += a1*b0; c11 += a1*b1; c12 += a1*b2; c13 += a1*b3;
        c20 += a2*b0; c21 += a2*b1; c22 += a2*b2; c23 += a2*b3;
        c30 += a3*b0; c31 += a3*b1; c32 += a3*b2; c33 += a3*b3;

 

여기서 실제로 GEMM 의 수식이 전개된다.

c00 에서

• 루프 시작 시 c00 = 0
• 각 k 마다 c00 += a0 * b0
• a0 = A(0, k), b0 = B(k, 0)

즉

• c00 = sum_k A(0, k) * B(k, 0) = C(0,0)

다른 것도 동일

전체적으로 보면

• 위 4 줄이 4x4 GEMM micro-kernel
• A 의 4x1 벡터, B 의 1x4 벡터를 outer-product 형태로 C(4x4) 에 누적하는 구조

그걸 K 번 반복하는 구조 = GEMM

 

4. 레지스터에 들고 있던 C 를 메모리로 저장

    // single thread writes C
    C[0] = c00; C[1] = c01; C[2] = c02; C[3] = c03;
    C[4] = c10; C[5] = c11; C[6] = c12; C[7] = c13;
    C[8] = c20; C[9] = c21; C[10]= c22; C[11]= c23;
    C[12]= c30; C[13]= c31; C[14]= c32; C[15]= c33;

• C 도 row-major 로 저장된다고 가정

즉, cij 들을 적절한 위치에 다시 써주는 단계

중요한 포인트로 이 단계 이전까진 C 는 전적으로 레지스터에만 존재했음, 이것이 바로 micro-kernel 의 핵심 감각

 

5. main - K, 메모리 할당

int mina() {
	int K = 1024;
    
	float *A, *B, *C;
	cudaMallocMAnaged(&A, sizeof(float)*4*K);
	cudaMallocMAnaged(&B, sizeof(float)*K*4);
	cudaMallocMAnaged(&C, sizeof(float)*16);

• GEMM 의 공통 차원 K = 1024 -> A(4x1024), B(1024x4)
• cudaMallocManaged : Unified Memory 할당
• Unified Memory 는 호스트/디바이스 공용

 

6. A, B 초기화

    // init
    for (int i=0;i<4*K;i++) A[i] = (i%7)*0.1f;
    for (int i=0;i<K*4;i++) B[i] = (i%5)*0.2f;

값 자체는 중요한 것이 아닌, 비 0 값 / 다양한 패턴으로 곱이 잘되는지 확인

 

7. 커널 런치 및 동기화

    micro_gemm_4x4<<<1,1>>>(A,B,C,K);
    cudaDeviceSynchronize();

• <<<1.1>>>
  • girdDim = 1 block
  • blockDim = 1 thread
    • 단일 스레드 실행
• 이 스레드가 위에서 본 micro-kernel 전체를 수행
  • A (4 x K), B (K x 4) 읽어서
  • C (4 x 4) 를 계산
• cudaDeviceSynchronize()
  • 커널 종료 대기

 

8. 결과 출력 & 메모리 해제

 

 

이 코드가 마이크로 커널인 이유

• C 의 작은 블록 전부 레지스터에 유지
  • c00 .. c33 가 바로 micro-tile
• K 루프마다 A 의 한 열 (4 x 1) + B 의 한 행 (1 x 4)
  • outer product 형태로 C 레지스터에 누적
• 메모리 계층
  • 글로벌 메모리에서 A[k열], B[k행] aks dlfrdma
  • C 는 끝가지 DRAM/Shared 에 안 놓이고 레지스터에서만 누적됨

이 구조를

• 스레드 1개 -> warp 32개 스레드
• 4x4 -> 16x16, 32x64 등으로 키우고
• shared memory 에 A/B tile 을 올려서 여러 스레드가 같이 쓰게 만들면, 

tiled GEMM 의 inner core 가 된다

거기에  FMA 패턴 최적화 + unroll + pipeline 넣으면 cuBLAS 류 마이크로 커널