[논문 리뷰] Compiler-Level Matrix Multiplication Optimization for Deep Learning

Huaqing Zhang, Xiaolin Cheng외 2명|arXiv (Cornell University)|2019. 09. 23.

Parallel Computing and Optimization Techniques참고 문헌 34인용 수 6

한 줄 요약

이 논문은 딥러닝 워크로드에서 행렬 곱셈(GEMM)을 위한 컴파일러 수준 최적화 기법 두 가지—탐욕적 최적 우선 탐색(G-BFS)과 이웃성 액터-보상 비평가(N-A2C)—을 제안한다. 마르코프 결정 과정(MDP)에서 이웃 구성 관계를 활용함으로써, 두 방법 모두 XGBoost 및 RNN 기반 튜너보다 각각 24%와 40% 낮은 GEMM 계산 시간을 달성하였으며, 1024×1024 행렬에서 구성 공간의 0.1%만 탐색하였다.

ABSTRACT

An important linear algebra routine, GEneral Matrix Multiplication (GEMM), is a fundamental operator in deep learning. Compilers need to translate these routines into low-level code optimized for specific hardware. Compiler-level optimization of GEMM has significant performance impact on training and executing deep learning models. However, most deep learning frameworks rely on hardware-specific operator libraries in which GEMM optimization has been mostly achieved by manual tuning, which restricts the performance on different target hardware. In this paper, we propose two novel algorithms for GEMM optimization based on the TVM framework, a lightweight Greedy Best First Search (G-BFS) method based on heuristic search, and a Neighborhood Actor Advantage Critic (N-A2C) method based on reinforcement learning. Experimental results show significant performance improvement of the proposed methods, in both the optimality of the solution and the cost of search in terms of time and fraction of the search space explored. Specifically, the proposed methods achieve 24% and 40% savings in GEMM computation time over state-of-the-art XGBoost and RNN methods, respectively, while exploring only 0.1% of the search space. The proposed approaches have potential to be applied to other operator-level optimizations.

연구 동기 및 목표

• 다양한 하드웨어에서 하위 최적의 GEMM 커널 생성으로 인해 발생하는 딥러닝 추론 및 훈련의 성능 저하 문제를 해결하기 위해.
• 딥러닝 프레임워크에서 GEMM 연산자의 수동적이고 하드웨어 특화된 튜닝에 대한 의존도를 줄이기 위해.
• 다양한 하드웨어 플랫폼 간에 고성능 GEMM 구성에 대한 효율적이고 자동화된 탐색을 가능하게 하기 위해.
• 대규모이고 복잡한 구성 공간에서 텐서 연산자 최적화를 위한 탐색 효율성과 해의 품질을 향상시키기 위해.

제안 방법

• G-BFS 방법은 히우리스틱 기반의 탐욕적 탐색 전략을 사용하여 GEMM 구성 공간 내 이웃 구성에 대해 탐색하며, 성능 향상이 기대되는 구성에 우선순위를 부여한다.
• N-A2C 방법은 이웃 기반 액션 공간을 사용하는 액터-크리틱 프레임워크를 활용해 강화학습을 적용하여 구성 선택을 반복적으로 개선한다.
• 두 방법 모두 GEMM 최적화를 마르코프 결정 과정(MDP)으로 모델링하며, 타일링 및 레이아웃 변환 기반 상태 전이를 정의한다.
• 구성 공간은 다중 수준의 타일링 파라미터로 정의되며, 상태는 행렬 A, B, C에 대한 타일 크기의 튜플로 표현된다.
• N-A2C 방법은 정책 네트워크 학습을 위해 측정된 커널 실행 시간 기반 보상 신호를 사용하는 반면, G-BFS는 다음 단계 우선순위를 정하기 위해 히우리스틱 점수를 사용한다.
• 두 접근 방식 모두 다양한 크기의 GEMM 워크로드에 대해 NVIDIA Titan XP GPU에서 TVM 컴파일러 프레임워크 내에서 평가되었다.

실험 결과

연구 질문

• RQ1 G-BFS와 같은 히우리스틱 기반 탐색 방법이 최소한의 탐색 비용으로 고성능 커널을 찾는 데 있어 GEMM 구성 공간을 효율적으로 탐색할 수 있는가?
• RQ2 N-A2C와 같은 강화학습 접근 방식이 기존 최첨단 튜닝 방법보다 해의 품질과 탐색 효율성 측면에서 뛰어나게 성능을 발휘할 수 있는가?
• RQ3 G-BFS와 N-A2C의 성능는 XGBoost 및 RNN 기반 튜너와 비교해 계산 시간 감소율과 구성 공간 탐색 비율 측면에서 어떻게 다른가?
• RQ4 이러한 방법들이 2048×2048과 같은 더 큰 행렬 크기로 확장될 경우, 탐색 효율성과 해의 품질 측면에서 얼마나 잘 작동하는가?

주요 결과

• N-A2C 방법은 1024×1024 행렬에서 구성 공간의 0.1%만 탐색함에도 불구하고 RNN 기반 튜너보다 40% 낮은 GEMM 계산 시간을 달성하였다.
• G-BFS 방법은 동일한 0.1% 구성 공간 탐색 제약 조건 하에서 XGBoost 튜너보다 계산 시간을 24% 감소시켰다.
• 두 방법 모두 XGBoost 및 RNN 방법보다 최적의 구성 설정을 더 빨리 발견하였으며, N-A2C는 다단계 계획 기능 덕분에 더 큰 행렬에서 뛰어난 성능를 보였다.
• 10회의 시험에서 두 방법 모두 성능 변동성이 낮아 안정적이고 신뢰할 수 있는 튜닝 행동을 보였다.
• N-A2C 방법은 더 큰 행렬 크기(예: 2048×2048)에서 G-BFS를 능가하여 복잡한 탐색 공간에서 순차적 결정의 이점을 입증하였다.
• 결과적으로 G-BFS와 N-A2C가 대규모 구성 공간을 가진 텐서 연산자 컴파일러 수준 최적화에 효과적이고 일반화 가능한 방법임을 확인하였다.