공대생의 공부흔적

[컴퓨터구조#17] Domain Specific Architecture - TPU 본문

Computer Architecture

[컴퓨터구조#17] Domain Specific Architecture - TPU

생대공 2024. 6. 10. 14:30

이번 글에서는 domain specific 아키텍처에 대해 알아보자. 

목차

  1. DSA 소개
  2. 예시 - TPU

1. DSA 소개

무어의 법칙(칩 내 트랜지스터 점점 더 많아짐)은 CPU, GPU에서 다음과 같은 일반적인 개선사항들을 가능하게 했다.

  • 깊은 메모리 위계, 넓은 SIMD unit, 깊은 파이프라인, 브랜치 예측, OoO 실행, speculative prefetching, Multithreading, Multiprocessing

이는 아키텍처를 인식하지 않는 소프트웨어에서 성능을 뽑아내는 데 목표가 있었다.

이와 달리 DSA는 domain에 specific하게 설계된 아키텍처이다. 최근에는 거의 머신러닝을 타겟으로 하지만, 주요 연산 커널을 정의할 수 있다면 어떤 응용이든 가능하다. DSA를 위한 몇 가지 가이드라인은 다음과 같다.

  • 데이터 이동을 최소화하기 위해 지정된 메모리 사용
  • 더 많은 계싼 unit이나 더 큰 메모리에 자원 투자
  • 각 domain에 맞는 가장 쉬운 병렬성 형태 사용
  • domain에 필요한 데이터 크기와 종류를 가장 간단한 형태로 줄임
  • domain-specific한 프로그래밍 언어 사용

 

2. 예시 - TPU

DSA의 한 예시로 DNN에 specific한 TPU를 살펴보자.

TPU는 Tensor Processing Unit의 약자로, 구글의 DNN ASIC이다. 256*256 8-bit MatMul unit(systolic array)와, 큰 SW-managed scratchpad 메모리(GPU의 shared memory와 비슷)가 존재한다. 또한 PCIe bus에 대한 coprocessor이다. 행렬곱 unit이 8비트인 이유는 inference에만 이를 사용했기 때문이다. 이후에 나온 아키텍처에서는 training으로도 확장하여 더 많은 비트를 사용했다.

TPU ISA

  • Read_Host_Memory
  • Read_Weights
  • MatrixMatrixMultiply/Convolve
  • Activate
  • Write_Host_Memory

on-chip scratchpad buffer는 activation에 사용되며 기존 칩의 29% 영역을 차지했으나 이후에는 크기를 줄이고 개수를 늘리는 방향으로 개선되었다. 행렬곱 unit은 칩 면적의 24%를 차지했다.

Systolic array는 이전 글의 Volta 아키텍처에 있었던 Tensor Core(행렬곱 가속기)와는 다른 개념이다. Systolic array 내 각 PE는 작고 가벼우며 weight와 partial sum을 저장하는 작은 register, 계산 unit ... 등으로 이루어져 있다.

위 가이드라인이 TPU에 어떻게 적용되는지 살펴보자.

  • 데이터 이동을 최소화하기 위해 지정된 메모리 사용 → 24MB 지정 버퍼, 4MB accumulator 버퍼
  • 더 많은 계싼 unit이나 더 큰 메모리에 자원 투자 → 메모리의 60%, 250개의 arithmetic unit
  • 각 domain에 맞는 가장 쉬운 병렬성 형태 사용 → 2D SIMD 병렬성 활용 (MatMul)
  • domain에 필요한 데이터 크기와 종류를 가장 간단한 형태로 줄임 → 8비트 정수
  • domain-specific한 프로그래밍 언어 사용 → TensorFlow