[논문 리뷰] A Learned Performance Model for Tensor Processing Units

Samuel J. Kaufman, Phitchaya Mangpo Phothilimthana외 5명|arXiv (Cornell University)|2020. 08. 03.

Parallel Computing and Optimization Techniques참고 문헌 26인용 수 33

한 줄 요약

본 논문은 그래프 신경망을 사용하여 커널 런타임을 예측하는 TPU용 학습된 성능 모델을 제시하고, 타일 크기 선택과 연산자 융합에서 분석 모델보다 우수하며 하드웨어 액세스가 제한될 때 자동 튜닝에 도움이 된다.

ABSTRACT

Accurate hardware performance models are critical to efficient code generation. They can be used by compilers to make heuristic decisions, by superoptimizers as a minimization objective, or by autotuners to find an optimal configuration for a specific program. However, they are difficult to develop because contemporary processors are complex, and the recent proliferation of deep learning accelerators has increased the development burden. We demonstrate a method of learning performance models from a corpus of tensor computation graph programs for Tensor Processing Unit (TPU) instances. We show that our learned model outperforms a heavily-optimized analytical performance model on two tasks -- tile-size selection and operator fusion -- and that it helps an autotuner discover faster programs in a setting where access to TPUs is limited or expensive.

연구 동기 및 목표

• 실제 하드웨어에 대한 의존도를 줄이기 위해 TPU에서 텐서 프로그램의 정확한 런타임 예측을 제공합니다.
• 과도한 특징 공학 없이 다양한 텐서 프로그램과 최적화 작업에 일반화합니다.
• 다른 컴파일러 최적화 작업에 대해 재타깃 가능한 성능 모델링을 가능하게 합니다.
• TPU 접근이 제한되거나 비용이 비쌀 때 자동 튜닝 효율의 향상을 입증합니다.

제안 방법

• 연산 노드와 텐서-플로우 간선을 가진 데이터 흐름 그래프로 텐서 프로그램을 표현합니다.
• 근접 이웃에서 노드 임베딩을 계산하기 위해 Graph Neural Network(GraphSAGE)를 사용합니다.
• 노드 임베딩 위에 시퀀스 모델(LSTM 또는 Transformer)을 선택적으로 결합하여 장기 의존성을 포착합니다.
• 연산 코드 임베딩과 노드별/커널 특징을 포함하여 모델 입력 벡터를 형성합니다.
• 타일 크기 선택을 위한 쌍 순위 손실과 융합을 위한 제곱 오차 손실(대상 로그 변환)을 사용하여 절대 런타임을 예측하도록 학습합니다.
• 수동으로 튜닝된 분석 모델과 비교하고 학습된 모델을 XLA 자동 튜너에 통합하여 구성 탐색을 안내합니다.

실험 결과

연구 질문

• RQ1 학습된 모델이 타일 크기 선택과 연산자 융합을 안내할 만큼 TPU 커널 런타임을 정확하게 예측할 수 있나요?
• RQ2 최소한의 특징 공학으로 그래프 기반 표현이 보지 않은 텐서 프로그램과 서로 다른 하드웨어 세대(TPU v2/v3)에서도 일반화되나요?
• RQ3 학습된 모델을 자동 튜너에 통합하면 빠른 구성을 찾는 데 필요한 실제 하드웨어 평가 수를 줄일 수 있나요?

주요 결과

• 학습된 모델은 타일 크기와 융합 작업에 대해 각각 참측값에 대해 96.3%와 95.5%의 정확도를 달성합니다.
• 타일 크기와 융합 작업에서 각각 최적의 수동 튜닝 분석 모델보다 2.4%와 26.6% 더 나은 정확도를 제공합니다.
• TPU v3에서 학습된 모델은 타일 크기 평가에 대해 평균 오차 3.8%를 달성하고 평균 상관계수는 약간 낮은 0.65이다.
• 발생 연구는 에지 방향, 정적 성능 특징, 커널 임베딩 선택이 융합 작업 정확도에 상당한 영향을 주고, 타일 크기 작업 정확도에는 덜하지만 영향을 준다.
• 학습된 모델을 XLA 자동 튜너에 통합하면 하드웨어 접근이 제한되거나 비용이 비쌀 때 효율적인 구성의 더 빠른 발견이 가능해진다.