[논문 리뷰] Ansor: Generating High-Performance Tensor Programs for Deep Learning

Lianmin Zheng, Chengfan Jia외 10명|arXiv (Cornell University)|2020. 06. 11.

Parallel Computing and Optimization Techniques참고 문헌 53인용 수 82

한 줄 요약

Ansor는 텐서 프로그램 생성 프레임워크로, 자동으로 큰 위계적 탐색 공간을 구축하고, 전체 프로그램을 샘플링하며, 학습된 비용 모델로 진화적 탐색을 사용하여 CPU와 GPU에서 DNN 서브그래프를 최적화하고, 최첨단 방법에 비해 상당한 속도 향상을 달성합니다.

ABSTRACT

High-performance tensor programs are crucial to guarantee efficient execution of deep neural networks. However, obtaining performant tensor programs for different operators on various hardware platforms is notoriously challenging. Currently, deep learning systems rely on vendor-provided kernel libraries or various search strategies to get performant tensor programs. These approaches either require significant engineering effort to develop platform-specific optimization code or fall short of finding high-performance programs due to restricted search space and ineffective exploration strategy. We present Ansor, a tensor program generation framework for deep learning applications. Compared with existing search strategies, Ansor explores many more optimization combinations by sampling programs from a hierarchical representation of the search space. Ansor then fine-tunes the sampled programs with evolutionary search and a learned cost model to identify the best programs. Ansor can find high-performance programs that are outside the search space of existing state-of-the-art approaches. In addition, Ansor utilizes a task scheduler to simultaneously optimize multiple subgraphs in deep neural networks. We show that Ansor improves the execution performance of deep neural networks relative to the state-of-the-art on the Intel CPU, ARM CPU, and NVIDIA GPU by up to $3.8 imes$, $2.6 imes$, and $1.7 imes$, respectively.

연구 동기 및 목표

• 다양한 하드웨어 플랫폼에서 DNN을 위한 고성능 텐서 프로그램의 필요성을 강조한다.
• 텐서 프로그램의 방대한 탐색 공간을 자동으로 생성하고 탐색함으로써 수작업 템플릿 엔지니어링을 제거한다.
• 스케줄링 및 자원 인식 선택을 통해 DNN 서브그래프의 엔드-투-엔드 최적화를 가능하게 한다.
• Ansor가 CPU 및 GPU 대상을 대상으로 최첨단 방법과 비교하여 우수한 성능을 달성함을 입증한다.

제안 방법

• 고수준 구조(스케치)와 저수준 세부사항(주석)을 분리하는 계층적 탐색 공간을 도입한다.
• 계산 그래프를 나타내는 DAG에 대한 유도 기반 열거를 통해 자동으로 스케치를 생성한다.
• 스케치를 무작위로 주석 달아 평가 및 튜닝을 위한 완전한 프로그램을 얻는다.
• 학습된 비용 모델을 사용한 진화적 탐색을 적용하여 샘플링된 프로그램을 미세 조정한다.
• 엔드-투-엔드 DNN 성능을 가장 크게 향상시키는 서브그래프를 우선순위화하기 위해 기울기 하강 기반 작업 스케줄러를 사용한다.
• 그래디언트 부스팅 트리로 학습된 비용 모델을 활용하여 프로그램 처리량을 예측하고 탐색을 안내한다.

실험 결과

연구 질문

• RQ1 Ansor가 수작업 템플릿 없이도 텐서 프로그램의 크고 포괄적인 탐색 공간을 자동으로 생성할 수 있는가?
• RQ2 학습된 비용 모델과 함께 진화적 탐색이 CPU 및 GPU 백엔드 전반에서 고성능 프로그램을 효과적으로 식별하는가?
• RQ3 스케줄러를 통한 서브그래프의 우선순위 지정이 연산자별 최적화와 비교하여 엔드-투-엔드 DNN 성능을 향상시키는가?
• RQ4 여러 하드웨어 플랫폼에서 Ansor의 성능이 최첨단 탐색 기반 프레임워크와 어떻게 비교되는가?

주요 결과

• Ansor는 Intel CPU, ARM CPU 및 NVIDIA GPU에서 각각 최대 3.8×, 2.6×, 1.7×의 DNN 실행 성능 향상을 달성한다.
• Ansor가 찾은 가장 성능이 좋은 프로그램 대부분은 기존 접근법의 탐색 공간 밖에 있다.
• Ansor는 더 큰 탐색 공간에도 불구하고 더 효율적으로 탐색하고 짧은 시간에 더 높은 성능의 프로그램을 찾는다.
• 탐색 시간을 한 차원 더 적게 들여도 최첨단 성능에 일치할 수 있다.
• Ansor는 수작업 템플릿이 아닌 수학적 정의만으로도 새로운 연산자에 자동 확장을 가능하게 한다.