google TPU 논문 리뷰: In-Datacenter Performance Analysis of a Tensor Processing Unit (1) - TPU 구조 및 명령어

구글이 2017년에 발표한 TPU 논문 리뷰를 해보려고 한다. 이 논문은 첫번째 버전의 TPU에 대한 설명을 다루고 있고, 현재에는 TPU는 4번째 버전까지 발표되었다. 해당 논문과 구글 블로그의 글을 참고하여 논문을 정리하였다.

다음과 같은 구성으로 글을 써보려고 한다.
1. TPU origin : TPU 왜 만들게 되었는지
2. TPU architecture : TPU의 구성 및 동작 원리
3. CPU, GPU, TPU 연산량, 구조 비교
4. performance : CPU, GPU와 비교한 TPU의 성능

 

TPU Origin

논문에서 최근(2013년 기준)에 DNN 모델들이 요구하는 datacenter의 연산량이 매우 증가했다고 한다. 예를 들어 하루에 3분씩 사람들이 음성 검색을 한다면 (구글의) datacenter의 계산량이 두배가 증가한다고 한다. datacenter는 많은 서버가 모여있는 시설이라고 생각하면 될 것 같다.
이로 인해 구글은 빠른 inference(추론)를 위한 custom ASIC이 필요했고 GPU보다 cost-performance(가격 대비 성능)가 10배 정도 좋은 성능을 가지도록 만들어내는 것을 목표로 15개월 만에 TPU를 설계했다. ASIC은 범용 용도가 아닌 특정 용도에 맞게 맞춤 제작된 집적 회로를 말한다. CPU 같은 칩은 게임을 돌릴 수도 있고, 계산을 할 수도 있고, 동영상도 재생할 수 있는 등 매우 다양한 일을 할 수 있는 범용 용도로 만들어진 칩인 반면에, TPU는 DNN의 계산만을 위한 칩으로, 동영상을 재생한다거나 다른 프로그램을 실행할 수 없는 특정 용도의 칩, ASIC이다.

 

 

TPU Architecture

TPU의 구조

위의 그림은 논문에 실린 TPU의 구조 그림이다. 저 각각의 네모들을 unit이라고 부를 수 있는데, 노란색 unit들은 연산을 수행하고, 파란색 unit은 data buffer로, 연산에 필요한 값들을 미리 가져와서 저장해놓고 연산할때 제공해주거나 혹은 연산의 결과값을 저장해준다. buffer는 어떤 데이터를 잠깐 저장해주는 메모리라고 보면 된다. control unit들은 각 unit들이 제대로 수행되도록 도와주는 unit이라고 할 수 있고, 연두색은 off-chip에 있는 메모리들이라고 보면 된다. off-chip이니 접근할때 꽤 많은 시간과 많은 에너지가 소모될 것이다. 각 unit들에 대해 살펴보자.

 

 

Unit 소개

• Matrix Multiply Unit(MXU)
  • 주요 연산이 일어나는 곳이다. 주요 연산이라 하면 DNN의 연산식에서 i

    

    nput과 weight값을 곱하고 그 값들의 합을 구하는 연산을 말한다. (오른쪽 사진의 연산식 참고)
  • 256x256 개의 MACs으로 구성된다, 8-bit 부호 있거나 없는 정수의 연산을 한다. weight값과 input값을 곱하니까 곱의 결과값은 16-bit가 된다.
  • 1 clock cycle만에 256-elements의 한번 부분합을 계산한다.
  • Systolic Array로 구성되어 있다. Systolic Array에 대해서는 여기서 더 알아볼 수 있다.
  • TPU의 경우 Systolic Array 중에서 weight stationary로 구현되었다.
• Unified Buffer(UB)
  • matrix multiply unit의 연산에 필요한 input값을 CPU memory에서 읽어와서 가지고 있었다가 필요할때 matrix multiply unit에 넘겨준다. 또한 matrix multiply unit의 연산값에 accumulator와 activation을 거친 결과 값을 저장한다.
• Systolic Data Setup
  • Unified Buffer로부터 입력값 또는 중간 결과값을 불러와 해당 순서로 입력 데이터를 전달하는 timing을 조절하는 역할을 하며 Weight FIFO또한 비슷한 역할을 수행한다. 
• Weight FIFO
  • weight값들은 Weight Memory라고 불리는 off-chip 8GiB DRAM에서 on-chip Weight FIFO로 읽혀진 후에 Multiply Unit에 사용된다. 
  • Unified Buffer로 weight 값들을 보낸 후에 matrix multiply unit으로 보내져 연산에 사용된다. 
  • CPU host memory 혹은 unified buffer에서 data를 가져오는 것은 programmable DMA controller가 한다.
• Accumulators
  • 16-bit의 곱 결과들을 모아준다. 
  • 32-bits, 256-elements, 4096 (4MiB)
  • 4096은 peak performance(최고 성능)에 도달하기 위해 실험을 통해 설정한 값이다.
• Activation Unit(AU)
  • DNN의 모델을 보면 각 layer마다 값에 비선형성을 부여하기 위해서 activation function을 적용하는데, 그와 같은 역할을 해준다고 할 수 있다. accumulator의 결과값에 activation function을 적용해주어 비선형성을 부여해준다. 

 

TPU Instruction

MXU, UB, AU를 사용하여 연산을 수행하기 위해서 정의된 TPU 명령어는 총 12개로 구성되어있는데, 그 중 중요한 5개를 보자면

• Read_Host_Memory: data를 CPU host memory로부터 Unified Buffer로 읽어온다.
• Read_Weights: weight를 Weight Memory에서 Weight FIFO(matrix multiply unit의 input)로 읽어온다.
• MatrixMultiply/Convolve: Unifed Buffer로부터 accumulator까지 matrix multiply unit이 matrix multiply 혹은 convolution을 수행하게끔 한다. matrix operation은 x*256개의 input을 matrix unit에 넣어 x pipeline동안 256x256의 constant weight input값들과 곱해져서 연산된다.(x값은 최대 256까지 가능)
• Activate: accumulator의 연산 결과에 Sigmoid나 ReLU와 같은 인공신경망의 비선형함수를 적용한 값을 unified buffer에 저장하고 convolution의 pooling 연산을 수행하기도 한다.
• Write_Host_Memory: unified buffer의 data를 CPU host memory로 쓴다.

이 외에도 alternate host memory read/write, set configuration, two versions of synchronization, interrupt host, debug-tag, nop, halt의 명령어가 존재한다. 그리고 matrix unit의 utilization을 높이기 위해 (자주 busy 상태로 만들기 위해) 4-stages pipeline을 사용하여 명령어를 수행한다. CISC로 명령어를 구현했음에도 불구하고 하는 연산이 매우 단순하여 MatrixMultiply의 명령어를 실행하는데 오랜 시간이 걸리지 않는다. 덕분에 각 명령어들의 실행시간(CPI, Cycles Per Instruction)이 차이가 많이 나지 않고 pipeline을 하는데 큰 문제가 되지 않는다. 반면에, CPU의 경우에는 워낙 다양한 연산을 하다보니까 각 명령어의 CPI의 차이가 많이 나는 경우가 있어서 CISC로 구현하였을 경우 파이프라인을 실행하는데 어려움이 있을 수 있다.

 

※ 공부하고 있는 단계라 내용에 부족한 부분이 많습니다. 조언이나 지적은 감사히 받아들이고 수정하겠습니다.

 

 

 

[구글 TPU 논문 리뷰]

(2) CPU, GPU, TPU 연산량, 구조 비교
(3) performance : CPU, GPU와 비교한 TPU의 성능