OpenCL에서 효율적인 행렬곱(matrix multiplication) 수행

2020. 5. 10. 21:58Study

행렬 A: M x K

행렬 B: K x N

행렬 C: M x N

 

이 글에서는 OpenCL에서 $C = A \times B$를 효율적으로 계산하는 방법에 대한 내용을 다룬다. 총 4가지 버전의 구현이 있으며 설명 순서대로 incrementally 성능이 개선된다. 이 글은 https://cnugteren.github.io/tutorial/pages/page1.html의 많은 내용을 참고해서 작성 되었다. (이 글에서 사용한 행렬곱을 시각화한 그림은 모두 위의 링크에서 가져왔다.)

 

OpenCL 없는 구현

1. naive version

가장 일반적으로 생각 해낼 수 있는 행렬곱 구현이다.

for (int i=0; i<M; m++) {
  for (int j=0; j<N; n++) {
    float acc = 0.0f;
    for (int k=0; k<K; k++) {
      acc += A[i*K + k] * B[k*N + j];
    }
    C[i*N + j] = acc;
  }
}

 

2. cache friendly version

배열의 모든 원소 접근이 row-major로 이루어지도록 하여 cache miss를 크게 줄인 구현 방법이다.

for (i = 0; i < M; i++) {
  for (k = 0; k < K; k++) {
    r = A[i * K + k];
    for (j = 0; j < N; j++) {
      acc += r * B[k * N + j];
    }
    C[i * N + j] = acc;
}

 

1. Naive Implementation

 

첫 번째 버전의 커널은 가장 간단하지만 성능은 안 좋은 방법이다. 여기서는 앞에서 제시한 OpenCL을 사용하지 않은 코드에서 가장 바깥의 2개의 loop(각각 M번, N번 iterations)를 2차원 kernel index space를 통해서 병렬처리 한다. M * N 만큼의 연산이 GPU threads에서 병렬처리 되므로 총 M * N 개의 work-item들이 만들어진다.

 

// 커널 버전 1

__kernel void sgemm(__global float *A, __global float *B, __global float *C, int M, int N, int K) {
  // Thread identifiers
  int row = get_global_id(0); // row index of C
  int col = get_global_id(1); // column index of C

  float intermediate_val = 0.0f;
  for (int k = 0; k < K; k++) {	   
    intermediate_val += A[row * K + k] * B[k * N + col];
  }
 
  C[row * N + col] = intermediate_val;
}

 

위의 코드를 보면 i에 대한 loop(0~(M-1))는 TID(thread identifier) row로, j에 대한 loop(0~(N-1))는 TID col로 할당되어 병렬처리 되는 것을 알 수 있다. 위의 커널이 제대로 동작하기 위해서는 CPU host code에서 M * N번 만큼 kernel launch를 하여 M * N개의 커널 인스턴스를 생성하여 연산을 수행하여야 한다. OpenCL 프로그래밍의 기본에 대해서는 다른 포스팅에서 다루도록 하고 여기에서는 host code의 일부만 보이도록 하겠다.

 

# define TS    32

err = clSetKernelArg(kernel, 0, sizeof(cl_mem), &a_d);
err = clSetKernelArg(kernel, 1, sizeof(cl_mem), &b_d);
err = clSetKernelArg(kernel, 2, sizeof(cl_mem), &c_d);
err = clSetKernelArg(kernel, 3, sizeof(int), &M);
err = clSetKernelArg(kernel, 4, sizeof(int), &N);
err = clSetKernelArg(kernel, 5, sizeof(int), &K);

// Setup global work size and local work size
size_t gws[2] = {M, N}, lws[2] = {TS, TS};
for (int i = 0; i < 2; ++i) {
  // By OpenCL spec, global work size should be MULTIPLE of local work size
  // Formula below achieve it
  // e.g., gws = 25, lws = 16, then (25 + 16 - 1) / 16 * 16 = 40 / 16 * 16 = 2 * 16 = 32
  gws[i] = (gws[i] + lws[i] - 1) / lws[i] * lws[i];
}

// Launch kernel
err = clEnqueueNDRangeKernel(queue, kernel, 2, NULL, gws, lws, 0, NULL, NULL);

// Synchronize
err = clFinish(queue);

 

성능 (FLOPS)

이 글에서 성능 측정에 사용한 GPU는 NVIDIA GeForce GTX 1080이며, M = K = N = 4096 옵션에 대해서 성능을 측정했다. 

 

GLOPS: 44

 

2. Local Memory Tiling

커널 버전 1에서 성능이 좋지 못한 이유는 GPU의 off-chip memory에 너무 많이 접근하기 때문이다.

 

출처: 자일링스(Xilinx) SDAccel Development Environment Help

위의 그림은 GPU 기반의 이종 시스템(heterogeneous system)의 전체적인 메모리 및 연산 유닛 구조를 보여준다. 커널 버전 1에서는 for loop에서 FMA 연산을 수행할 때마다 global memory 접근이 발생할 수 있기 때문에 매우 비효율적이다. 이를 해결하기 위해서는 global memory 접근을 최소화 시켜야하고, 가장 먼저 떠올릴 수 있는 방법은 global memory 대신 on-chip local memory에 값을 저장/접근하도록 캐싱하는 것이다.

 

 

위의 그림을 캐싱의 이점을 활용한 tiling 방법을 보여준다. C의 sub block인 $T_{C1,3}$를 계산하기 위해서 그에 해당하는 A의 sub block(row major)과 B의 sub block(column major)이 필요하다. 여기에서는 A의 sub block인 $T_{A1,1}$, $T_{A1,2}$와 B의 sub block인 $T_{B1,3}$, $T_{B2,3}$의 곱을 누적합하는 방식으로 그에 해당하는 Csub를 계산해낼 수 있다. 즉, 각 타일을 일반적인 행렬곱의 원소와 같이 취급하되, 각 타일 내에서는 일반적인 행렬곱을 수행한다. 이와 같은  방법이 성능 개선에 도움이 되는 이유는 tile 내에서 데이터의 재사용이 여러 번 이루어지기 때문이다.

 

이를 OpenCL 코드로 구현해보자. 주의할 점은 Asub와 Bsub 타일의 데이터를 하나의 work-group에 할당하도록 하는 것이다. (참고: work-group은 kernel의 scheduling 단위이며, 32 배수의 work-item들을 포함한다.) 같은 work-group이면 같은 local memory를 공유하며, 같은 local memory를 공유하면 global memory에 접근을 하지 않을 수 있다. 타일의 크기는 work-group이 수용가능하다면 최대한 크게 유지하는 것이 데이터 재사용의 이점을 최대화할 수 있기 때문에 성능 개선에 유리하다.

 

// 커널 버전 2

__kernel void sgemm(__global float *A, __global float *B, __global float *C, int M, int N, int K) {
  // Thread identifiers
  const int row = get_local_id(0);   // local row ID (0~31) 
  const int col = get_local_id(1);   // local col ID (0~31)
  const int global_row = 32 * get_group_id(0) + row;   // row ID of C (0~M)
  const int global_col = 32 * get_group_id(1) + col;   // col ID of C (0~N)

  __local float Asub[32][32];
  __local float Bsub[32][32];

  float intermediate_val = 0.0f;
  const int num_tiles = K / 32;
  for (int t = 0; t < num_tiles; t++) {
    const int t_row = 32 * t + row;
    const int t_col = 32 * t + col;
    Asub[row][col] = A[global_row * K + t_col];
    Bsub[row][col] = B[t_row * N + global_col];

    barrier(CLK_LOCAL_MEM_FENCE);

    for (int k = 0; k < 32; k++) {
      intermediate_val += Asub[row][k] * Bsub[k][col];
    }

    barrier(CLK_LOCAL_MEM_FENCE);
  }

  C[global_row*N + global_col] = intermediate_val;
}

 

코드에서 눈 여겨볼 점은 다음과 같다.

 

  • 커널 버전 1에서 K 번 반복되던 loop는 2개의 nested loop로 분리되어 구현되었다. 바깥 loop는 K/32번(work-group 혹은 타일 개수만큼) 반복되며 그 안의 loop는 하나의 work-group 내에 있는 32개의 work-item들의 수만큼 반복된다.
  • 첫 번째 synchronization barrier 앞에서는 off-chip global memory의 데이터를 local memory로 로드한다.
  • Inner loop에서는 local memory로 로드된 데이터에 대해서 계산을 수행한다.

 

이와 같은 방법을 통해서 타일 당 오직 2번의 global memory 접근이 이루어지고, 타일 내부에서 계산을 할 때는 local memory 접근만이 이루어진다.

 

성능 (FLOPS)

이 글에서 성능 측정에 사용한 GPU는 NVIDIA GeForce GTX 1080이며, M = K = N = 4096 옵션에 대해서 성능을 측정했다. 

 

GLOPS: 197

 

3. 쓰레드마다 더 많은 workload 부여하기

2번 버전에서 성능을 개선할 수 있었던 또 다른 이유는 각 GPU thread 마다 할당되는 workload의 양이 증가했기 때문이다. (따라서 사용하는 전체 thread의 수가 감소함) 이것이 왜 성능 개선에 도움이 되었을까? 2번 버전 커널의 PTX assembly 코드를 살펴보면 그 이유를 알 수 있다. 다음은 커널 버전 2에서 32번의 loop를 돌면서 intermediate_val += Asub[row][k] * Bsub[k][col];를 계산하는 부분을 loop unrolling한 것이다. (제시된 코드는 2번의 iteration에 대한 코드)

 

ld.shared.f32   %f50, [%r18+56];
ld.shared.f32   %f51, [%r17+1792];
fma.rn.f32  %f52, %f51, %f50, %f49;
ld.shared.f32   %f53, [%r18+60];
ld.shared.f32   %f54, [%r17+1920];
fma.rn.f32  %f55, %f54, %f53, %f52;

 

각 iteration마다 다음과 같은 과정이 이루어진다.

 

  1. Asub의 원소들을 local memory에서 register로 로드
  2. Bsub의 원소들을 local memory에서 register로 로드
  3. FMA 연산 수행

위의 과정을 보면 1, 2번 명령은 operand를 로드하기 위한 명령에 불과하고 실제 연산은 3번 명령에서만 이루어진다는 것을 알 수 있다. 즉, 1, 2번 로드 명령의 오버헤드를 잘 처리하면 큰 성능 개선을 노려볼 수 있을 것이다.

 

3번 버전 커널에서는 2번 버전 커널과 마찬가지로 tiling을 통해서 성능을 개선할 것이다. 다만 2번 버전에서는 local memory에 대해서 타일링을 했다면, 여기에서는 register 레벨에서 타일링을 한다. 만약 한 개의 thread가 행렬 C의 연속적인 행들에 놓인 8개의 원소들에 대해서 계산을 수행한다면, 행렬 A에 대한 local memory 접근을 8배 줄일 수 있다. (이와 같이 각 쓰레드에서 계산하는 원소들의 개수를 WPT(Work-Per-Thread)라고 하자.) 아래의 그림은 WPT = 3인 경우를 시각화한 것이다. 단, 이해가 쉽도록 tiling은 적용하지 않은 상태이다.

 

 

다음은 register-level tiling을 적용한 새로운 버전의 커널이다. 주목할 부분은 다음과 같다.

 

  • 중간 누적값(intermediate_val)이 WPT개의 레지스터에 저장된다.
  • 각 thread는 행렬 A와 B의 WPT개의 원소를 local memory로 로드한다.
  • 각 thread는 타일마다 WPT*32번의 누적합을 수행한다. 이 때 WPT번 반복하는 가장 안 쪽의 loop에서 Asub의 각 인덱스 데이터에 접근 하지 않아서 local memory load로 인한 오버헤드를 크게 단축 가능하다.
  • 각 thread에서 계산한 WPT개의 값들은 최종적으로 행렬 C에 저장된다.

 

// 커널 버전 3

#define TS    32
#define WPT   8

__kernel void sgemm(__global float *A, __global float *B, __global float *C, int M, int N, int K) {
  // Thread identifiers
  const int row = get_local_id(0);   // local row ID (TS/WPT IDs) 
  const int col = get_local_id(1);   // local col ID (0~31)
  const int global_row = TS * get_group_id(0) + row;   // row ID of C (0~M)
  const int global_col = TS * get_group_id(1) + col;   // col ID of C (0~N)

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

  float intermediate_val[WPT];
  for (int w=0; w<WPT; w++) {
    intermediate_val[w] = 0.0f;
  }
  
  const int num_tiles = K / TS;
  for (int t = 0; t < num_tiles; t++) {
    for(int w = 0; w < WPT; w++) {
      const int t_row = TS * t + row;
      const int t_col = TS * t + col;
      Asub[row + w*RTS][col] = A[(global_row + w*RTS) * K + t_col];
      Bsub[row + w*RTS][col] = B[(t_row + w*RTS) * N + global_col];
    }

    barrier(CLK_LOCAL_MEM_FENCE);

    for (int k = 0; k < TS; k++) {
      for (int w = 0; w < WPT; w++) {
        intermediate_val[w] += Asub[row + w*RTS][k] * Bsub[k][col];
      }
    }

    barrier(CLK_LOCAL_MEM_FENCE);
  }

  for (int w=0; w<WPT; w++) {
    C[(global_row + w*RTS)*N + global_col] = intermediate_val[w];
  }
}

 

각 thread에 할당된 work-item의 수가 증가(kernal granularity 증가)했기 때문에 kernel launch 시에 지정하는 전체 쓰레드의 개수는 감소한다.

 

size_t gws[2] = {M/WPT, N};
size_t lws[2] = {TS/WPT, TS};

 

Register-level tiliing을 적용한 커널 버전 3의 PTX assembly code를 보면 다음과 같다. 앞에서 살펴본 것과 같이 load 명령 대비 FMA 명령의 비율이 크게 증가한 것을 알 수 있다.

 

ld.shared.f32   %f82, [%r101+4];
ld.shared.f32   %f83, [%r102];
fma.rn.f32  %f91, %f83, %f82, %f67;
ld.shared.f32   %f84, [%r101+516];
fma.rn.f32  %f92, %f83, %f84, %f69;
ld.shared.f32   %f85, [%r101+1028];
fma.rn.f32  %f93, %f83, %f85, %f71;
ld.shared.f32   %f86, [%r101+1540];
fma.rn.f32  %f94, %f83, %f86, %f73;
ld.shared.f32   %f87, [%r101+2052];
fma.rn.f32  %f95, %f83, %f87, %f75;
ld.shared.f32   %f88, [%r101+2564];
fma.rn.f32  %f96, %f83, %f88, %f77;
ld.shared.f32   %f89, [%r101+3076];
fma.rn.f32  %f97, %f83, %f89, %f79;
ld.shared.f32   %f90, [%r101+3588];
fma.rn.f32  %f98, %f83, %f90, %f81;

 

성능 (FLOPS)

이 글에서 성능 측정에 사용한 GPU는 NVIDIA GeForce GTX 1080이며, M = K = N = 4096 옵션에 대해서 성능을 측정했다. 

 

GLOPS: 544

 

 

4. Vector Type의 사용

앞에서는 각 thread에 할당되는 work-item의 granularity를 증대시키는 방식으로 성능을 개선할 수 있었다. Thread 마다 할당되는 work-item의 크기를 늘리는 것은 vector data type을 적용하는 것을 통해서도 가능하다. 여기에서는 vector type으로 single precision float이 4개가 들어가는 float4 타입의 vector를 사용할 것이다. (WIDTH = 4)

 

우선 kernel launch시에 index space는 WIDTH에 맞게 재설정 해준다.

 

size_t gws[2] = {M, N/WIDTH};
size_t lws[2] = {TS, TS/WIDTH};

 

다음은 vector type이 적용된 커널 버전 4의 코드이다. 데이터 타입이 float4 vector이기 때문에 각 인덱스는 4로 나누어진 값으로 사용되어야 한다.

 

__kernel void sgemm(__global float4 *A, __global float4 *B, __global float4 *C, int M, int N, int K) {
  // Thread identifiers
  const int row = get_local_id(0);   // local row ID (0~31) 
  const int col = get_local_id(1);   // local col ID (max: TS/WIDTH)
  const int global_row = TS * get_group_id(0) + row;   // row ID of C (0~M)
  const int global_col = (TS/WIDTH) * get_group_id(1) + col;   // col ID of C (0~N/WIDTH)

  __local float4 Asub[TS][TS/WIDTH];
  __local float4 Bsub[TS][TS/WIDTH];

  float4 intermediate_val = { 0.0f, 0.0f, 0.0f, 0.0f };
  
  const int num_tiles = K / TS;
  for (int t = 0; t < num_tiles; t++) {
    const int t_row = TS * t + row;
    const int t_col = (TS/WIDTH) * t + col;
    Asub[row][col] = A[global_row * (K/WIDTH) + t_col];
    Bsub[row][col] = B[t_row * (N/WIDTH) + global_col];

    barrier(CLK_LOCAL_MEM_FENCE);

    float4 vecA, vecB;
    float valA;
    for (int k = 0; k < TS/WIDTH; k++) {
      vecA = Asub[row][k];
      for (int w = 0; w < WIDTH; w++) {
        vecB = Bsub[WIDTH*k + w][col];

        switch(w) {
          case 0: valA = vecA.x; break;
          case 1: valA = vecA.y; break;
          case 2: valA = vecA.z; break;
          case 3: valA = vecA.w; break;
        }
        
        intermediate_val.x += vecB.x * valA;
        intermediate_val.y += vecB.y * valA;
        intermediate_val.z += vecB.z * valA;
        intermediate_val.w += vecB.w * valA;
      }
    }

    barrier(CLK_LOCAL_MEM_FENCE);
  }

  C[global_row*(N/WIDTH) + global_col] = intermediate_val;
}

 

성능 (FLOPS)

이 글에서 성능 측정에 사용한 GPU는 NVIDIA GeForce GTX 1080이며, M = K = N = 4096 옵션에 대해서 성능을 측정했다. 

 

GLOPS: 801

 

만약 vector 타입으로 float4 대신 float8을 사용한다면 work-item의 granularity를 더욱 증가시킬 수 있다. 다음은 위의 코드에서 float4를 float8로 변경한 코드이다.

 

#define NUM_WORK_ITEM   32
#define VECTOR_WIDTH    8 


__kernel void sgemm(__global float8 *A, __global float8 *B, __global float8 *C, int M, int N, int K) {
  // Thread identifiers
  const int row = get_local_id(0);   // row ID in a tile [0:32) 
  const int col = get_local_id(1);   // col ID in a tile [0:32/8)
  const int global_row = NUM_WORK_ITEM * get_group_id(0) + row;   // row ID of C [0:M)
  const int global_col = (NUM_WORK_ITEM/VECTOR_WIDTH) * get_group_id(1) + col;   // col ID of C [0:N/8)

  __local float8 tileA[NUM_WORK_ITEM][NUM_WORK_ITEM/VECTOR_WIDTH];
  __local float8 tileB[NUM_WORK_ITEM][NUM_WORK_ITEM/VECTOR_WIDTH];

  float8 intermediate_val = { 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f };
  
  const int num_tiles = K/NUM_WORK_ITEM;
  for (int t = 0; t < num_tiles; t++) {
    const int t_row = NUM_WORK_ITEM * t + row;
    const int t_col = (NUM_WORK_ITEM/VECTOR_WIDTH) * t + col;
    tileA[row][col] = A[global_row * (K/VECTOR_WIDTH) + t_col];
    tileB[row][col] = B[t_row * (N/VECTOR_WIDTH) + global_col];

    barrier(CLK_LOCAL_MEM_FENCE);

    float8 vecA, vecB;
    float valA;
    for (int k = 0; k < NUM_WORK_ITEM/VECTOR_WIDTH; k++) {
      vecA = tileA[row][k];
      for (int w = 0; w < VECTOR_WIDTH; w++) {
        vecB = tileB[VECTOR_WIDTH*k + w][col];

        switch(w) {
          case 0: valA = vecA.s0;  break;
          case 1: valA = vecA.s1;  break;
          case 2: valA = vecA.s2;  break;
          case 3: valA = vecA.s3;  break;
          case 4: valA = vecA.s4;  break;
          case 5: valA = vecA.s5;  break;
          case 6: valA = vecA.s6;  break;
          case 7: valA = vecA.s7;  break;
        }

        intermediate_val += vecB * valA;
      }
    }

    barrier(CLK_LOCAL_MEM_FENCE);
  }

  C[global_row*(N/VECTOR_WIDTH) + global_col] = intermediate_val;
}

 

성능 (FLOPS)

이 글에서 성능 측정에 사용한 GPU는 NVIDIA GeForce GTX 1080이며, M = K = N = 4096 옵션에 대해서 성능을 측정했다. 

 

GLOPS: 1067 GFLOPS

 

추가: Zero-padding

Work-group size인 32로 나누어 떨어지지 않는 경우에도 계산 결과의 correctness를 보장하기 위해서는 행렬의 alignment를 32의 배수로 맞춰줄 필요가 있다. 이를 위한 방법 중 하나는 32의 배수 크기의 행과 열을 갖도록 zero-padding을 넣어주는 것이다. 

 

 

Zero-padding을 넣어주는 커널을 따로 만들어도 되고, host source code에서 패딩 처리를 해주어도 상관 없다. Host에서 처리하는 경우엔 kernel 및 buffer 생성 이전, 그리고 host to device memory copy 이전에 패딩을 추가해주고, device to host memcpy 이후에 패딩을 제거하면 된다. 패딩을 추가할 땐 32의 배수로 alignment가 맞춰지도록 알맞은 크기를 넣어주면 된다.