Matrix Multiplication Acceleration

행렬 곱셈은 단순하지만 상당히 비용이 큰 연산이다.

C = A * B 일때 C의 원소 하나를 구하기 위해선 A의 행벡터와 B의 열벡터를 곱해야 한다. 이를 모든 원소에 대하여 수행해야하기 때문에 행렬의 크기가 클 경우 컴퓨터로 연산을 하더라도 상당한 시간이 걸릴 수 있다.


다만 다행인 것은 행렬의 연산은 모든 원소에 대해 같은 방식, 다른 데이터로 구하는 것이라는 거다. 즉 SIMD task 라는 것이다. 따라서 GPU로 가속한다면 성능향상을 기대할 수 있다.

1. Naive implementation
가장 단순하게 구현하는 방법을 생각해보자. CPU에서의 연산을 그대로 사용하되 이를 GPU를 이용해 병렬적으로 계산하는 것이다.

__kernel
void matrix_multiplication1(const int M, const int N, const int K,
 const __global float* A,const __global float* B, __global float* C ) {
    const int globalRow = get_global_id(1);
    const int globalCol = get_global_id(0);

    float localsum = 0.0f;
    for (int k = 0; k < K; k++) {
        localsum += A[globalRow*K+k]*B[N*k + globalCol];
    }
    C[K*globalRow + globalCol] = localsum;
}

위는 내가 수치 컴퓨팅을 수강하자 마자 집가서 작성해본 코드이다. 간단하지만 그만큼 가속 성능이 좋지 않다. 그 이유론 잦은 Global Memory access를 생각해 볼 수 있다.
GPU에서 Global memory를 access 하는것은 상당히 오래 걸리는 작업이다. 물론 I/O 동안 context switching, caching 이 일어나지만 이로썬 부족하다.

register usage : 6(parameters) + 5(local variables)
memory access : (2K+1)*size^2[GM] for WG

그렇다면 이를 최적화 해서 성능을 끌어올릴 방법은 어떤 것이 있을까?
첫째로 Local Memory 이용을 생각해 볼 수 있다.

2.Implementation using Local Memory
위와 같은 방식이지만 Global memory를 줄이고 Local Memory 를 이용하는 것이다.

#define TS 16

__kernel 
void matrix_multiplication2(const int M, const int N, const int K,
    const __global float* A, const __global float* B, __global float* C) {
    const int row = get_local_id(1);
    const int col = get_local_id(0);
    const int globalRow = TS * get_group_id(1) + row;
    const int globalCol = TS * get_group_id(0) + col; 
 
    __local float Asub[TS][TS];
    __local float Bsub[TS][TS];

    float localsum = 0.0f;
    const int numTiles = K / TS;
    for (int i = 0; i < numTiles; i++) {
        const int tiledRow = TS * i + row;
        const int tiledCol = TS * i + col;
        Asub[row][col] = A[globalRow*K + tiledCol];
        Bsub[row][col] = B[tiledRow*N + globalCol];
  
        barrier(CLK_LOCAL_MEM_FENCE);
        for (int k = 0; k < TS; k++) {
            localsum += Asub[row][k] * Bsub[k][col];
        }
        barrier(CLK_LOCAL_MEM_FENCE);
    }

    C[globalRow*N + globalCol] = localsum;
}

Local Memory를 이용하기 위해 Global Memory에서 데이터를 복사하는 과정이 추가되었다. 또한 Local Memory를 할당받는 것은 하드웨어적 제약이 따르므로 Tiling 을 한다.
그 외 나머지 작업은 동일하다.

register usage : 6(parameters) + 9(local variables)
memory access : (2K+size)*size[GM] + 2K*size^2[LM] for WG

이로써 어느정도 최적화를 했다고 생각할 수 있다.
하지만 우린 파워 공대생이기 때문에 최적화를 멈출 수 없다.

위는 Local Memory를 이용한 행렬곱셈 코드를 PTX assembly code 로 dump 시킨 것이다. C의 원소 계산을 위한 fma instruction당 두번의 Local Memory Access가 필요하단 것이다.
그렇다면 하나의 CU에 더 효율적으로, 더 많은 일을 시킬 수 있을까?

3.More Work per Thread
나의 실력이 부족해 이 방식을 이해하고 직접 구현하는데 정말 많은 고민을 했다.
설명할 자신은 없고 다만 커널에서 원소를 한개가 아니고 여러개 계산한다는 것이 포인트다. 친절한 설명따윈 없으니 고민해 보길 바란다.

#define TS 64
#define WPT 8
#define RTS 8

__kernel 
void matrix_multiplication3(const int M, const int N, const int K,
    const __global float* A, const __global float* B, __global float* C) {
    const int row = get_local_id(1);
    const int col = get_local_id(0);
    const int globalRow = TS * get_group_id(1) + row;
    const int globalCol = TS * get_group_id(0) + col; 

    __local float Asub[TS][TS];
    __local float Bsub[TS][TS];

    float acc[8];
    for (int i = 0; i < WPT; i++)
        acc[i] = 0.0f;

    const int numTiles = K / TS;
    for (int t = 0; t < numTiles; t++) {           
        for (int w = 0; w < WPT; w++) {            
            const int tiledRow = TS * t + row;
            const int tiledCol = TS * t + col;
            Asub[row][col + w * RTS] = A[globalRow * K + (tiledCol + w * RTS)];
            Bsub[row][col + w * RTS] = B[tiledRow * N + (globalCol + w * RTS)];
        }
        barrier(CLK_LOCAL_MEM_FENCE);
  
        for (int w = 0; w < WPT; w++) {
            for (int k = 0; k < TS; k++) {
                acc[w] += Asub[row][k] * Bsub[k][col + w * RTS];
            }
        }
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    for (int w = 0; w < WPT; w++) {
        C[globalRow * N + globalCol + w * RTS] = acc[w];
    }
}

이로써 우리는 clBlas와 비슷한 성능의 코드를 작성할 수 있었다. 하지만 파워 공대인은 멈추지 않는다. 최적화 방법은 상당히 많이 남아있고 아직 그 부분은 구현을 하지 않았다.
내가 생각하고 있는 것은 Transpose를 이용한 memory access pattern을 조금이나마 효율적으로 바꾸는 것과 벡터를 이용한 계산등을 시도해 볼 예정이다.

관심이 있는 사람은 https://cnugteren.github.io/tutorial/pages/page1.html 에서 많은 도움을 얻을 수 있을 것이다.