CUDA 기반 행렬 곱셈

행렬 연산은 CUDA 연산에 가장 어울리는 문제이다. 따라서 행렬 곱셈을 GPU에서 수행하는 방법을 이해한다.

스레드 레이아웃 설정

대규모 행렬 곱을 위한 CUDA 프로그램을 위해 스레드 레이아웃을 먼저 결정해야 한다. 어떻게 레이아웃 기준을 잡아야 할까? 두 가지 경우를 생각할 수 있다.

  • 데이터를 읽는 행렬 A, B 기준
  • 결과가 저장되는 행렬 C 기준

C = AB

이 행렬 연산에서 C의 (row, col) 값 계산을 위해서는 A(row, k) B (k, col) 원소들을 불러와야 한다.

// 호스트 코드
for (int r = 0; r < m; r++) 
    for (int c = 0; c < n; c++)
        for (int i = 0; i < k; i++)
            C[r][c] += A[r][i] * B[i][c]

메모리 접근 측면에서 보면 행렬 A, B에서 read가 반복된다. 그리고 행렬 C의 (row, col) 원소에는 write가 반복된다. 이러한 메모리 접근을 기억하고, 2가지 레이아웃 방법을 비교한다.

A, B 기준 스레드 레이아웃

행렬 A의 각 행과 행렬 B의 각 열을 하나의 스레드가 처리하면, 스레드가 하는 일이 완전히 독립적이며 이상적일 것 같으나, 스레드의 수가 행 개수 또는 열 개수로 제한된다. CUDA 코어보다도 수십 배의 스레드를 사용해야 하는 것 특성 상, 이러한 스레드 레이아웃은 GPU 유틸라지션에서 적절하지 않다.

두 행렬 입력을 동시에 고려해볼 수 있다. 모든 for 무느이 각 반복을 하나의 스레드가 처리한다. 행, 열 그리고 차원의 수 만큼 스레드가 할당되어 대규모 병렬 처리가 가능하다. 하지만 C의 입장에서 C(row, col)에 여러 개의 스레드가 동시에 접근하여 값을 갱신해야 한다.

읽기 연산은 동일한 데이터를 여러 스레드가 읽어도 문제되지 않는다. 하지만 동일한 메모리 영역에 여러 스레드가 쓰기 연산을 하면 잘못된 값이 생성될 수 있다. 이를 해결하기 위해 동기화 기법을 사용할 수 있지만, 스레드들이 직렬화되어 병렬처리 성능을 크게 떨어트린다.

C 기준 스레드 레이아웃

동기화 문제가 발생하는 연산은 write이다. 행렬 C의 한 원소처럼 동일한 메모리에 여러 스레드가 동시에 write를 하지 않도록 하는 방법 중 하나는 write 영역을 분할하여 스레드에게 분배하는 것이다. 행렬 곱셈 연산은 C의 각 원소를 스레드에 분배함으로써 실현할 수 있다.

C(row, col)을 담당하는 스레드는 행렬 A의 row 행과 행렬 B의 col 열의 원소들을 순차적으로 읽으며, C의 값을 계산한다. C(row, col + n)을 담당하는 스레드도 행렬 A의 row 행에 동시 접근 가능하지만, read 연산은 동기화 문제가 발생하지 않는다. 따라서 스레드 수가 행렬 C의 원소 수와 동일하다. 대규모 행렬의 경우, GPU에서 충분한 병렬 처리 연산을 수행할 수 있다.

행렬은 2차원 데이터로, 레이아웃 또한 2차원으로 잡는 것이 일반적이다. 따라서 각 차원의 길이를 행렬 C의 행과 열의 길이와 같게 하는 것이 스레드 인덱싱에서 직관적이다.

스레드 인덱싱

스레드 레이아웃을 결정하면, 커널에서 각 스레드가 접근할 데이터 위치를 결정하는 스레드 인덱싱 작업을 해야 한다. 결과 행렬 C를 기준으로 스레드 레이아웃을 설정함에 따라, 각 스레드가 접근하는 데이터는 아래와 같다.

  • 행렬 C 중 자신이 담당하는 원소
  • C의 원소를 계산하기 위한 행렬 A의 한 행과 B의 한 열

C의 크기가 블록 최대 크기(1,024)보다 작은 경우

이 설명은 블록 내 최대 스레드 수인 1,024개 인 경우, 블록 1개를 사용하는 행렬 연산이다. C가 2차원 스레드 레이아웃을 사용하기 때문에, 각 스레드의 C(row, col) 인덱스는 스레드 ID를 사용하여 지정할 수 있다.

row = threadIdx.x
col = threadIdx.y

메모리상의 C(row, col) 위치는 row, col 값과 C의 행 길이 n을 이용하여 계산한다. 또한 이처럼 단일 블록의 행렬 연산에서는 행 길이는 스레드 블록의 y-차원 길이와 같다.

index = row * n + col
index = row * blockDim.y + col

C(row, col) 값 연산은 A의 row 번째 행 원소와 B col 번째 원소들을 차례대로 방문해야 한다. offset = 0, 1, … k일때 A(row, offset)과 B(offset, col)에 접근한다. 이를 각 행렬의 목표 원소라 할 때, 행렬에서 목표 원소에 접근하기 위한 인덱싱은 다음과 같다.

A(row, offset) = row * k + offset
B(offset, col) = col + (offset * n)

행렬 A는 행 길이가 k이며, rok * k는 row 행의 첫 번째 원소를 가리킨다. 접근하고자 하는 원소는 row 행의 offset 원소이므로, row * k + offset이다. 반면에 행렬 B는 열이 고정되며 행을 따라 원소를 읽는다. 행렬 B의 행 길이는 n이며, offset 원소 행의 첫 번째 원소는 n * offset이다. 해당 행의 col 번째 원소를 읽으면 된다. 즉 n * offet + col로서 메모리 상 위치를 계산할 수 있다.

곱셈 커널

아래 코드는 스레드 수가 1024개 이하인 경우에서의 행렬 곱셈 커널이다. 호스트 코드에서 첫 번째, 두 번째 반복문이 병렬화되어 스레드에 분배됨에 따라, 사실상 offset 레벨의 반복문만 남아있다. 스레드 레이아웃에 의해 C의 인덱스를 지정하여, 디바이스 메모리 영역의 실제 메모리에 접근한다.

// 단일 블록에서의 CUDA 기반 행렬 곱셉 커널
__global__ void matMul_krenl(int* A, int* B, int* C, int m, int n, int k) {
    int row = threadIdx.x;
    int col = threadIdx.y;
    int index = row * n + col;

    C[index] = 0
    for (int offset = 0; offset < k; offset++) {
        C[index] += A[row*k + offset] * B[col + offset*n];
    }
}

C의 크기가 블록 최대 크기(1,024)보다 큰 경우

이 경우에는 여러 개의 블록을 사용해야 한다. 다만 스레드 레이아웃이 행렬 C와 같은 크기를 가지는 점은 동일하다. 따라서 스레드의 x-차원과 y-차원을 이용해서 행렬 C의 각 원소를 지정할 수 있다. 단일 블록에서는 블록 내 스레드 번호만 사용하면 되었지만, 여러 개의 블록은 스레드의 글로벌 인덱스를 사용해야 한다. 즉, (2차원 그리드, 1차원 블록)이 아닌, (2차원 그리드, 2차원 블록)이다.

이전에서 사용한 글로벌 인덱스를 구하는 방법을 활용할 수 있다. 목표 원소가 속한 블록의 첫 번째 원소까지 이동한 후, 블록 내부에서 자신의 인덱스까지 이동하면 된다.

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

곱셈 커널

입력 행렬 A, B의 인덱싱은 C(row, col)을 기준으로 계산한다. 커널 역시 row, col을 계산하는 부분만 수정하면 된다.

// 1024개 스레드보다 큰 여러 블록에서의 행렬 곱셉 커널
__global__ void matMul_kernel (int* A, int* B, int* C, int m, int n, int k) {
    int row = (blockDim.x * blockIdx.x) + threadIdx.x;
    int col = (blockDim.y * blockIdx.y) + threadIdx.y;
    int index = row * n + col;

    C[index] = 0;
    for (int offset = 0; offset < k; offset++) {
        C[index] += A[row * k + offset] * B[col + offset * n];
    }
}

구현 및 성능 평가

전체 프로그램 코드

#define BLOCK_SIZE 16

int main(int argc, char* argv[]) {
    int m, n, k;
    m = atoi(argv[1]); n = atoi(argv[2]); k = atoi(argv[3]);

    int sizeA = m * k;
    int sizeB = k * n;
    int sizeC = m * n;

    int* dA, * dB, * dC;

    // 1. dA, dB, dC를 위한 디바이스 메모리 할당
    cudaMalloc(&dA, sizeA * sizeof(int));
    cudaMalloc(&dB, sizeB * sizeof(int));
    cudaMalloc(&dC, sizeC * sizeof(int));

    cudaMemset(&dA, sizeA * sizeof(int));
    cudaMemset(&dB, sizeB * sizeof(int));
    cudaMemset(&dC, sizeC * sizeof(int));

    // 2. 행렬을 GPU로 데이터 복사
    cudaMemcpy(dA, A, sizeA * sizeof(int), cudaMemcpyHostToDevice);
    cudaMemcpy(dB, B, sizeB * sizeof(int), cudaMemcpyHostToDevice);
    cudaMemcpy(dC, C, sizeC * sizeof(int), cudaMemcpyHostToDevice);

    // 3. 스레드 레이아웃 설정
    dim3 gridDim(ceil((float) m  / BLOCK_SIZE), ceil((float) n / BLOCK_SIZE));
    dim3 blockDim(BLOCK_SIZE, BLOCK_SIZE);
    
    // 4. 커널 호출
    MatMul_kernel <<<gridDim. blockDim>>> (dA, dB, dC, m, n, k);
    
    // 5. 디바이스 메모리에서 호스트 메모리로 데이터 복사
    cudaMemcpy(Cgpu, dC, sizeC * sizeof(int), cudaMemcpyDeviceToHost);

    // 6. 디바이스 메모리 해제
    cudaFree(dA);
    cudaFree(dB);
    cudaFree(dC);

    return 0;
}

커널을 실제로 호출하기 위해서는 대상 행렬의 크기에 따라 스레드 레이아웃을 결정해야 한다. 또한 커널 수행 시, 사용할 스레드 블록의 크기를 결정해야 한다. 스레드 블록의 크기를 결정했다면, 그리드 크기를 계산해야 한다. 위 구현에서 사용하는 스레드 레이아웃은 연산 결과인 행렬 C를 기준으로 한다. 따라서 C의 원소 수만큼 스레드를 사용해야 한다.

// 행이 m이고 열이 n인 행렬 C의 필요한 블록 수
x-차원: m / BLOCK_SIZE
y-차원: n / BLOCK_SIZE

여기서 문제는 BLOCK_SIZE로 나누어 떨어지지 않는 경우이다. 그런데 그리드는 반드시 정수여야 한다. 만약 내림 연산을 한다면, 오른쪽 원소들을 담당할 스레드가 없어서 해당 원소들을 연산할 수 없다. 따라서 실수 값을 올림 처리하는 것이 유리하다. 이때 ceil() 함수를 통해 누락되는 부분 없이 모든 원소를 연산할 수 있다.

ceil() 연산으로 올림 처리하여, 모든 행렬의 원소를 연산할 수 있지만 행렬을 벗어나는 영역까지 연산을 위해 메모리에 접근하려는 문제가 생긴다. 따라서 메모리 할당을 받지 않는 영역을 접근하는 것은 메모리 접근 위반으로 프로그램에 오류를 발생시킬 수 있다.

__global__ void matMul_kernel (int* A, int* B, int* C, int m, int n, int k) {
    int row = (blockDim.x * blockIdx.x) + threadIdx.x;
    int col = (blockDim.y * blockIdx.y) + threadIdx.y;
    int index = row * n + col;

    // 메모리 접근 위반을 방지하는 예외처리 코드
    if (row >= m || col >= n) {
        return;
    }

    C[index] = 0
    for (int offset = 0; offset < k; offset++) {
        C[index] += A[row * k + offset] * B[col + offset * n];
    }
}

스레드가 담당할 원소 (row, col)이 행렬 C의 범위를 벗어나는지 체크한다. 만약 범위를 벗어나면, 연산할 원소가 아니기 때문에 스레드를 종료하면 된다. (이를 return 처리 한다.) 마지막으로 GPU로 연산을 위해서는, 디바이스 메모리 공간을 확보해야 한다. 메모리를 직접 할당하고 초기화하기 위해, 연산이 종료된 데이터의 메모리는 항상 메모리 해제를 하는 습관을 가진다.

부동소수점 정밀도

CUDA 연산 중, 쉬운 예제는 데이터가 정수인 경우이다. 하지만 딥러닝, 수치해석 같이 다양한 분야에서 부동소수점 연산을 한다. 따라서 float 데이터형을 가지는 경우를 고려해야 한다. float 연산에서는 CPU와 GPU의 동일한 로직 연산이 다른 경우가 발생할 수 있다.

이는 부동소수점 표현 방식에서 가지는 정밀도 차이로 발생하는 문제이다. 부동소수점은 다음과 같이 이루어진다.

  • 부호(sign)
  • 지수부(exponenet)
  • 가수부(fraction)

부호는 += 여부를 결정한다. 지수부는 숫자 표현의 범위를 결정한다. 그리고 가수부는 숫자 표현의 정밀도를 결정한다. 예를 들어서 딥러닝에서 쓰이는 부동소수점 표현 중 같은 16비트 데이터 형이여도, bfloat16, float16은 서로 표현이 다르다.

bfloat16: 지수부가 8, 가수부가 7이다.
float16:  지수부가 5, 가수부가 10이다.

bf16은 fp16에 비해 수의 범위가 넓지만, 표현 정밀도가 떨어져 0에 가까운 값은 모두 0으로 표현된다.

bf16이 fp16과 다른 특징 때문에 딥러닝 구조와 연산에 따라서 bf16이 더 유리할 수 있다. 예를 들어서 LLM같이 매우 디코더 레이어가 많은 연산이라면,