도입

지난 시간을 요약해보면, 근거없는 자신감으로 NumPy한테 도전하기 위해 pybind11로 코드를 포팅해서 덤볐다가, 처참하게 졌다. 고 할 수 있겠다.

그런 의미에서 이번 시간에는 다시 지하실로 돌아가서, 다시 최적화를 진행해보려고 한다.

현재 알고리즘

지금 내 커널은 어떻게 동작할까. 우선 그것부터 명확하게 하고 가자.

M행 K열의 행렬 A와, K행 N열의 행렬 B의 행렬곱에 대해서 결과적으로 M×N개의 내적을 계산해야 하고, 각 내적에 대해서는 A행렬의 Row Vector 및 B행렬의 Column Vector만 있으면 된다.

따라서? 16×16의 타일을 만들고, A행렬에서는 Column 방향(오른쪽)으로, B행렬에서는 Row 방향(아래쪽)으로 움직이며 내적을 수행해도 아무런 문제가 없으므로, 16×16의 가상 작업자(이하 스레드) 팀에게 각 타일의 1칸씩을 담당하게 한 뒤, 메모리에서 한번에 한 타일을 로컬 메모리로 옮겨오고, 각 스레드는 로컬 메모리만 왔다갔다 하며 16번의 곱셈 및 덧셈(sum에 곱한 결과를 더해야 하므로)을 수행한다. 이것이 타일링이었고, 이걸 통해 큰 폭으로 시간을 줄이는 데 성공했다(#2 참고).

그리하여 현재의 최적화 상태에서, 커널의 메모리 접근 횟수를 분석해보면 다음과 같다.

1. 시스템 메모리 -> 로컬 메모리로 데이터를 옮기는 횟수:

   {행렬의 수} × {팀의 수} × {타일의 수} × {타일의 크기}
   = {2} × {(M/16) × (N/16)} × {(K/16)} × {(16 × 16)}
   = (1/8) × MNK (회)

2. 로컬 메모리 -> 레지스터로 데이터를 옮기는 횟수:

   {전체 내적의 횟수} × {내적할 벡터의 길이} × {피연산자 수}
   = {M × N} × {K} × {2}
   = 2MNK (회)

여기서 시스템 메모리 -> 로컬 메모리 구간이 로컬 메모리 -> 레지스터 구간보다 1회당 시간 비용이 더 크기 때문에, 현재의 알짜 시간 비용이 과거 시스템 메모리 -> 레지스터로 직접 2MNK의 데이터를 옮기던 때보다는 훨씬 빠르지만, 아직 부족하다.

이전의 성능 극복의 본질은, 시스템 메모리에서 레지스터로 직접 가져온 데이터를 여러번 활용할 수 있음에도 불구하고 바로 버리는 문제를, 캐싱을 통해 해결하는 것이었다.

레지스터를 써보자

자. 자연스럽게 레지스터라는 말을 언급했는데, 이번에 최적화를 시도하기 전까지, 나는 이 녀석을 그냥 연산장치의 입력 포트 정도로 생각하고 있었다. 피연산자를 복사해서 레지스터에 담고, 연산장치에 명령을 내리면 레지스터를 읽고 명령을 실행한다. 정도의 개념으로.

그런데, 잘 생각해보면 이 녀석, 데이터를 "저장"하고 있다. 즉, 메모리와 똑같은 일을 하고 있다. 그런 주제에, 접근 시간은 진짜 압도적으로 짧다. 이건 안 써먹을 수가 없지 않은가?

그러고 좀 조사해 보니, 세상에, sycl에서 지역 변수를 선언하면 기본적으로 그것이 레지스터에 할당된다(조금이라도 크면 시스템 메모리로 쫓겨나지만)는 사실을 알았다.

즉, 각 스레드별로 엄청나게 빠른 메모리가 하나씩 있다고 생각하면 될 것 같은데, 지금까지대로의 알고리즘이라면 이것이 하등 쓸모가 없다.
왜냐하면? 한 스레드 당 entry 하나를 채우기 때문이다. 로컬 메모리에서 값 하나를 가져와서 굳이 가지고 있을 필요가 없다. 한번 쓰고 버릴 것이기 때문이다.

즉, 한 스레드가 복수의 entry를 채우기 시작할 때 이 조그마한 유사 메모리가 빛을 발할 것이라는 얘기다.

이 원리를 스레드 작업량 증가(Thread Coarsening)라고 하는 것 같다. 바로 적용해보...기 전에 한가지 수학적 사실을 짚고 넘어가야 한다.

행렬곱: by 벡터의 외적(Outer Product)

지금까지는 한 entry를 일일이 계산했기 때문에, 전혀 신경 쓸 여지가 없었지만, M×K 행렬 A와 K×N 행렬 B의 행렬곱을 구하는 방법은 사실 두가지이다.

1. M행 N열인, (AB)의 각 entry를 내적을 통해 하나하나 구한다.
2. A의 모든 열벡터와 B의 모든 행벡터를 외적한 후 합한다.

사실 정답 행렬의 각 entry에 스레드가 하나씩 대응될 때는 2번 방법이 별 의미를 갖지 못했지만, 이렇게 한 스레드가 여러 entry를 다룰 이 방법은 의미를 갖는걸 넘어 훨씬 효율적인 방법으로 거듭나는데, 이유는 레지스터의 높은 재사용성 때문이다.

1번을 고를 경우, 레지스터를 100% 활용하기 위해서는 각 칸의 지금까지의 합을 저장할 배열(크기: batch_size × batch_size), 로컬메모리에서 필요한 정보를 각각 A와 B로부터 긁어와서 저장할 배열(크기: 2 × batch_size × tile_size)이 필요한데, 너무 많다! 레지스터 용량이 터져버리고 말 것이다.

따라서 2번을 고르게 된다. 2번 같은 경우, 루프를 돌면서 각 열벡터와 행벡터를 한번씩만 방문하고 외적한 다음 합 변수에 더하기만 하면 되므로, 합을 저장할 배열(크기: batch_size × batch_size), 열벡터를 저장할 배열(크기: batch_size), 행벡터를 저장할 배열(크기: batch_size) 이렇게가 끝이다! 훨씬 메모리 집약적이다.

구현!

우선 너무 복잡하게 생각하지 않기 위해, 행렬의 크기는 이전처럼 항상 2의 n제곱 꼴(n은 n>4인 자연수)이라고 생각하고 구현해보자.

아래는 파이썬으로 포팅하는 부분을 생략하고, sycl 세팅 및 커널 부분만 가져온 코드이다.

 // coarsening을 위해 batch_size를 지정
            const int BATCH = 4;
            // q에 행렬 연산 작업을 제출해준 후,
            q.submit([&](sycl::handler& h) {
                // Accessor 세팅
                sycl::accessor accA(bufA, h, sycl::read_only);
                sycl::accessor accB(bufB, h, sycl::read_only);
                sycl::accessor accR(bufR, h, sycl::write_only, sycl::no_init);

                // Tiling을 위한 local memory 할당
                // A는 row vector, B는 column vector만 필요하므로,
                // coarsening을 할 시, 한 쪽 차원으로만 늘려도 ok.
                sycl::local_accessor<float, 2> tileA(sycl::range<2>(TILE*BATCH, TILE), h);
                sycl::local_accessor<float, 2> tileB(sycl::range<2>(TILE, TILE*BATCH), h);

                // coarsening으로 인해 전체 쓰레드의 수가 1 / (batch_size ^ 2) 로 감소함.
                // 또한 TILE의 크기도 자연스럽게 (batch_size ^ 2)배가 된다.
                // 한 스레드는 이제 (batch_size × batch_size) 영역을 커버하게 된다.
                h.parallel_for(sycl::nd_range<2>(sycl::range<2>(M/BATCH, N/BATCH), sycl::range<2>(TILE, TILE)), [=](sycl::nd_item<2> item) {
                    // 우선 전체를 batch_size로 쪼개서 global id를 먼저 부여한 후,
                    const int r = item.get_global_id(0), c = item.get_global_id(1);
                    // 생성된 녀석들을 tile_size 단위로 묶어서 local id를 추가로 부여하는 방식.
                    const int local_r = item.get_local_id(0), local_c = item.get_local_id(1);
                    // 따라서 global_id에 batch_size를 곱하면 자기가 담당하는 구역의 시작 인덱스가 구해짐.
                    // local_id는 좀 더 복잡한데, local_r은 A입장에서는 채워야 하는 entry들의 행에 해당하고,
                    // B 입장에서는 batch_size를 곱해야 채워야할 entry들의 시작 행에 해당한다.
                    // local_c는 그 반대이다.

                    // NUM_TILES는 동일하다.
                    const int NUM_TILES = K / TILE;
                    // 외적의 합을 저장할 배열을 생성 및 초기화
                    float sum[BATCH][BATCH];
                    for(int i=0; i<BATCH*BATCH; ++i) sum[i/BATCH][i%BATCH] = 0;
                    // 
                    for(int t=0; t<NUM_TILES; ++t)
                    {
                        // 먼저 로컬 메모리부터 채운다. 
                        // 이전과 달리 담당구역 전체를 채워야 하기 때문에 루프.
                        // 헷갈리면 안되는 부분이, TILE 크기 자체는 동일하기 때문에,
                        // K축(A의 열축, B의 행축. 이하 설명 생략.)으로 진행하는 속도는 똑같음.
                        // 즉 이 경우, 로컬 메모리의 형태가 정사각형이 아니다.
                        for(int i=0; i<BATCH; ++i)
                        {
                            tileA[local_r*BATCH+i][local_c] = accA[(r*BATCH+i)*K + (local_c+TILE*t)];
                            tileB[local_r][local_c*BATCH+i] = accB[(local_r+TILE*t)*N + (c*BATCH+i)];
                        }
                        // 동료들의 메모리 채우기를 기다린다.
                        item.barrier(sycl::access::fence_space::local_space);
                        float regA[BATCH], regB[BATCH];
                        // K축으로 루프 시작
                        for(int k=0; k<TILE; ++k)
                        {
                            // 로컬메모리 -> 레지스터로 ( ,k) / (k, ) 벡터를 옮긴다.
                            for (int i=0; i<BATCH; ++i)
                            {
                                regA[i] = tileA[local_r*BATCH+i][k];
                                regB[i] = tileB[k][local_c*BATCH+i];
                            }
                            // 이제 오직 레지스터에만 접근하면서 외적을 계산할 수 있음!
                            for (int i=0; i<BATCH; ++i)
                            {
                                for(int j=0; j<BATCH; ++j)
                                {
                                    sum[i][j] += regA[i] * regB[j];
                                }
                            }
                        }
                        // 동료들의 연산 종료를 기다린다.
                        item.barrier(sycl::access::fence_space::local_space);
                    }
                    // 레지스터 -> 시스템 메모리로 정답을 돌려주면 끝.
                    for (int i=0; i<BATCH; ++i)
                    {
                        for(int j=0; j<BATCH; ++j)
                        {
                            accR[(r*BATCH+i)*N + (c*BATCH+j)] = sum[i][j];
                        }
                    }
                    });
                });
            // 결과를 기다리고, res를 return하면 끝.
            q.wait();
            return res;

그래서 결과는??

오? 생각보다 많이 따라잡았다. 역시 로컬 메모리도 메모리는 메모리라는 것일까.

마무리

정말 하면 할수록 느끼는 것이지만, 항상 메모리가 문제다! 느린 메모리를 들르는 횟수를 줄이는 것. 즉, 가능한 한 빠른 메모리에 들르는 것. 이것이 핵심인 것 같다.

지금까지는 느린 메모리에 접근하는 횟수를 줄이는 것을, 빠른 메모리로 데이터를 옮겨놓고 여러번 우려먹는 방법으로 해결해왔다.

그런데, GPU에는 그것 말고도 메모리 접근 횟수를 줄일 수 있는 방법이 하나 더 있다. 그것은 한번에 왕창 가져오는 것이다!

그런 의미에서 다음 시간에는 sycl::float4를 활용해서 데이터를 뭉텅이로 가져오는 최적화를 적용해보고자 한다.