[인공지능 하드웨어] 6 - Deep Learning Optimization(Convolution Lowering, Systolic Array)

# 들어가며

Matrix Tiling등의 기법을 사용한 GPU Performance Optimization을 중점적으로 알아본다. 컴퓨터 구조 관련한 배경지식이 있으면 이해하기 편하므로 글을 읽다가 이해가 되지 않는다면 관련 블로그 포스트나 유튜브를 찾아보는 것이 좋다. 이전 글은 다음 링크의 블로그 포스트를 참고하면 된다.

[인공지능 하드웨어] 5 - GPU Performance Optimization(Matrix Tiling, Tensor Core)

---

# Convolution Lowering

GeMM(General Matrix to Matrix Multiplication)은 이미 다양한 라이브러리나 하드웨어 가속기에 기반해서 매우매우 다양한 방법으로 빠르게 수행할 수 있다. 수많은 Implementation 또한 존재한다.
비슷하게 Convolution 연산은 굉장히 시간이 많이 드는 작업이다. 그렇다면 Convolution을 행렬곱 즉, GeMM으로 바꾸면 GPU 아키텍쳐의 이점을 최대한을 살릴 수 있지 않을까? 이것이 Convolution Lowering의 아이디어이다.

 

다음은 Convolution 연산을 할 때 Convolution Lowering을 적용하기 위해서 Im2Col 방식을 사용한 모습이다. Im2Col은 말 그대로 Image to Column 즉, 이미지를 한 열으로 만들겠다는 뜻이다. 

 

Convolution 연산의 방식을 생각 해보자. O0를 만들기 위해서 필요한 연산은 [D0, D1, D3, D4] * [F0, F1, F2, F3] Element-wise 연산이다. 이걸 그냥 왼쪽은 행으로 만들고, 오른쪽은 열로 만든거다. 어라? 근데 인덱스가 다르다 왜지? 그건 그냥 옛날 컨볼루션은 순서를 바꿔서 진행했기 때문이다. 요즘은 픽셀에 딱딱 맞춰서 커널을 적용시키지만 예전에는 컨볼루션 연산 할 때 픽셀을 뒤집어서 적용했었다.

 

 

Im2Col 이미지

 

아래 그림은 도식화 해서 Im2Col 행렬을 그린 것이다. 한 면씩 해석 해보자.
Ho, Wo는 output 이미지의 사이즈이다. Ho와 Wo가 정해지는 방법은 커널이 input 이미지에 대해서 Stride나 Padding을 고려하여 가로와 세로로 각각 몇번씩 순회 하는지에 따라서 결정 될 것이다.

CI, Hk, Wk는 입력채널, 커널의 높이, 커널의 너비이다. 입력 채널은 이전 층에서 온 이미지의 채널 수이고, 커널의 높이와 너비는 커널의 크기만큼 인풋 이미지를 펴주기 때문이다.

Co는 커널의 갯수로 output의 채널이 된다.

두 행렬은 Ci, Hk, Wk로 행과 열이 맞추어져 있으며, 맞추어지는 과정은 위의 Im2Col 순서를 생각하면 지극히 자연스럽다.

 

Im2Col 행렬 도식화 이미지

 

조금 더 자세한 예시를 살펴보자. 아래 그림은 NVIDIA에서 발표한 CuDNN에 수록되어 있던 그림이다. Input 이미지의 Ci, H, W가 3, 3, 3이고 filter의 Co, Ci, H, W가 2, 3, 2, 2일때 Im2Col을 적용한 모습이다. output의 Co, H, W가 정상적으로 2, 2, 2로 나타나는 것을 볼 수 있다.

Im2Col Detailed Explanation

 

Convolution Lowering은 그 원리도 간단하고 성능도 강력하지만, 몇가지의 단점이 있다.

• 같은 데이터가 복제되어 여러번 저장된다. 위의 예시를 보면 알겠지만, 커널 크기만큼 영역을 행렬로 변환하는 과정에서 동일한 데이터가 여러번 중복 저장되는 경우가 있다.
• 메모리 사용량이 증가한다. 데이터가 중복저장되면 당연히 메모리 사용량이 증가한다.
• 계산 비효율성이 증가한다. 계산 과정에서 같은 데이터가 중복되어 계산되므로 당연한 것이다.

cuBLAS와 cuDNN은 똑같이 Convolution Lowering을 지원하지만, 그 방식이 약간은 다르다. cuBLAS는 애초에 행렬을 Im2Col로 변환해서 메모리에 저장한다. cuDNN은 메모리에 원본 행렬 형태로 저장했다가 내부적으로 Convolution Lowering을 돌려서 결과 데이터를 사용한다. 그 이후에 연산은 동일하다.

cuBLAS vs cuDNN

 

## Convolution Lowering in Hardware

하드웨어에서 Convolution Lowering이 어떻게 일어나는지 조금 자세히 알아보도록 하자.

 

아래 그림은 1D 즉 1차원 데이터상에서 일어나는 컨볼루션을 하드웨어 수준에서 도식화 한 것이다.

다음과 같은 3x1짜리 가중치 행렬이 있을 때, 컨볼루션의 결과 행렬인 y를 만드는 과정이다.

W1

W2

W3

1D Convolution in Hardware

 

아래 그림은 2D 즉 2차원 데이터상에 일어나는 컨볼루션을 하드웨어 관점에서 그린 그림이다. 커널이 3x3의 크기를 가질 때, 다음과 같이 별도의 버퍼나 메모리에 저장되면 3개씩 데이터를 읽어와서 연결한다. 연결된 벡터는 Systaolic Array같은 계산 유닛에 제공된다.

2D Convolution in Hardware

2D Convolution in Hardware

# Systolic Array

## TPU Architecture

아래 보이는 구조는 Google에서 2017년에 발표한 TPU의 구조이다. 노란색으로 표시된 부분은 실제로 계산이 일어나는 부분이고, 하늘색으로 표시된 부분은 계산에 필요한 입출력값과 Weight들을 DDR3 메모리와 TPU사이에서 옮기거나 중간 결과를 임시로 저장하는 부분이다. 빨간색은 명령어 수행을 위한 제어 로직이고, 초록색은 TPU 외부와 연결되는 DDR3 메모리와 PCIe 인터페이스이다.

2017, Google, In-Datacenter Performance Analysis of a Tensor Processing Unit

 

이 TPU에서 가장 중요한 부분은 오른쪽의 Matrix Multiply Unit(MMU) 부분인데, 이 부분에서 행렬곱이 Systolic Array로 구성되어 동작하기 때문이다. 이 Systolic Array를 사용하는 TPU의 장점 및 특징은 다음과 같다.

• 하나의 데이터를 여러번 재사용해야 하는 행렬곱에 매우 뛰어난 성능을 보인다.
• Training은 지원되지 않으며 Inferencing에 특화되어 있다.(TPU-v2부터는 지원된다고 한다.)
• TPU에는 256x256 크기의 Systolic Array를 사용한다. 그 중에서도 Weight Stationary로 구성된다.
•  각 셀은 PE(Processing Element)라고 하며, MACs(Multiply/Accumulate Units)라고도 한다. 각 셀은 FPA연산을 지원하여 곱셈과 누적합을 동시 수행한다.

Matrix Multiply 주변에 있는 유닛들의 역할에 대해서 간략하게 살펴보도록 하자.

• Unified Buffer는 MMU 연산에 필요한 input값을 CPU Memory에서 읽어서 갖고 있다가 MMU에 넘겨준다. 또한, MMU에서 Activation과 Normalize 등을 거친 결과값을 저장하기도 한다.
• Systolic Data Setup은 MMU로 연산에 필요한 값을 전달하는 타이밍을 조절한다. Stationary 방식에 따라서 다르게 전달할 수 있다. 또한 Im2Col 방식으로 포매팅 하기도 한다.
• Weight FIFO는 Systolic Data Setup과 비슷한 역할을 하지만 Input Data가 아닌 Weight 값들을 가져온다.

## Output Stationary

Systolic Array에는 몇가지 방식이 있다. 그 중 첫번째로 Output Stationary이다. Output 즉, 행렬 연산의 결과값이 각 PE에 남는다(Stationary)고 해서 Output Stationary이다. 아래 그림을 보면 단번에 이해할 수 있다. 

Output Stationary 출처: make a gif

 

아래 그림을 보면 조금 더 자세히 이해 할 수 있다. 먼저, 왼쪽과 위쪽의 유닛에 의해서 input data와 weight data가 입맛에 맞게 정렬된다.

 

2 사이클이 지난 후의 모습이다. input data와 weight data는 각각 오른쪽, 아래로 흐르면서 곱셈을 하고 그 결과를 각 PE에 누적합하는 것까지 1 사이클에 수행한다.

 

최종적으로 연산이 완료되면 다음과 같이 Output들이 MMU에 남게된다.

## Weight Stationary

두번째로 Weight Stationary이다. Google TPU에는 Systolic Array 연산 중에서 이 방식이 들어가게 되며, 아래 그림과 같은 절차를 따르게 된다. 이때, 중요한 것은 Data는 오른쪽으로 전파되며, Partial Sum즉, 누적합은 아래로 전파되어 최종적으로 Output이 차례대로 아래쪽에 출력된다.

 

아래 GIF를 보면 단번에 이해할 수 있다. 이때 중요한 것은 Weight를 Pre-load 즉 각 PE에 맞게 저장해 놓은 다음에 연산을 시작한다는 것이다. 따라서 Weight를 MMU로 Load하는 시간이 추가적으로 발생하지만 8사이클만에 3x3 연산이 완료된다.

Weight Stationary

 

아래 그림을 보면 더욱 자세하게 이해할 수 있다. 먼저 MMU에 Weight를 먼저 Load해 놓는다. 그 후 Systolic Data Setup 유닛에 의해서 값들이 정렬되어 대기한다.

 

2 사이클이 지난 시점이다. Data들은 오른쪽으로, 누적합들이 아래로 전파되고 있는 것을 볼 수 있다.

 

4 사이클이 지난 시점이다. Output이 모두 완료되어 아래의 Accumulator 유닛을 통과하면 Activation, Pooling 등의 유닛을 거쳐서 최종적으로 신경망의 1개 layer의 task가 완료된 결과를 얻을 수 있을 것이다. 결과는 다시 Unified Buffer로 이동하여 다음 연산에 쓰인다. 이것이 전체적인 TPU의 구조이다.

---

 

[인공지능 하드웨어] 7 - Winograd Convolution

국민대학교 권은지 교수님의 인공지능 하드웨어 강의를 수강하며 정리한 내용입니다.

Slides Provided

임베디드 시스템 프로그래밍 - Eunhyeok Park

딥러닝 최적화 - Eunhyeok Park

References

https://www.donghyun53.net/%EA%B5%AC%EA%B8%80-tpu-%EB%93%A4%EC%97%AC%EB%B3%B4%EA%B8%B0/

https://justmajor.tistory.com/7