[논문 리뷰] Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions

Nicolas Vasilache, Oleksandr Zinenko외 7명|arXiv (Cornell University)|2018. 02. 13.

Computational Physics and Python Applications참고 문헌 61인용 수 252

한 줄 요약

Tensor Comprehensions(TC)는 ML 워크로드를 위한 고성능 CUDA 커널을 생성하는 도메인 특화 언어이자 폴리헤드 JIT 컴파일러로, 프레임워크 간 통합과 연산자 융합 및 자동 튜닝을 가능하게 한다.

ABSTRACT

Deep learning models with convolutional and recurrent networks are now ubiquitous and analyze massive amounts of audio, image, video, text and graph data, with applications in automatic translation, speech-to-text, scene understanding, ranking user preferences, ad placement, etc. Competing frameworks for building these networks such as TensorFlow, Chainer, CNTK, Torch/PyTorch, Caffe1/2, MXNet and Theano, explore different tradeoffs between usability and expressiveness, research or production orientation and supported hardware. They operate on a DAG of computational operators, wrapping high-performance libraries such as CUDNN (for NVIDIA GPUs) or NNPACK (for various CPUs), and automate memory allocation, synchronization, distribution. Custom operators are needed where the computation does not fit existing high-performance library calls, usually at a high engineering cost. This is frequently required when new operators are invented by researchers: such operators suffer a severe performance penalty, which limits the pace of innovation. Furthermore, even if there is an existing runtime call these frameworks can use, it often doesn't offer optimal performance for a user's particular network architecture and dataset, missing optimizations between operators as well as optimizations that can be done knowing the size and shape of data. Our contributions include (1) a language close to the mathematics of deep learning called Tensor Comprehensions, (2) a polyhedral Just-In-Time compiler to convert a mathematical description of a deep learning DAG into a CUDA kernel with delegated memory management and synchronization, also providing optimizations such as operator fusion and specialization for specific sizes, (3) a compilation cache populated by an autotuner. [Abstract cutoff]

연구 동기 및 목표

• ML용 텐서 연산을 표현하기 위한 간결하고 수학 친화적인 언어를 제공한다.
• 메모리 관리 및 스케줄링을 갖춘 최적화된 GPU 코드로 텐서 컴프리헨션을 변환한다.
• TC를 PyTorch와 Caffe2에 연결하여 프레임워크 간 통합을 가능하게 한다.
• 자동 튜너를 통한 자동 튜닝으로 최적화 기회를 탐색한다.
• 기존 ML 프레임워크와의 실용적 통합 및 경쟁력 있는 성능을 보여준다.

제안 방법

• 하이 레벨 연산에 가까운 아이젠슈타인 합(sum) 표기법에 근접한 텐서 컴프리헨션(TC)을 정의한다.
• 메모리 관리 및 동기화를 포함해 TC를 CUDA 커널로 하향 변환하는 폴리헤드(polyhedral) JIT 컴파일러를 개발한다.
• 커널 융합 및 크기 특화에 대한 맞춤형 폴리헤드 최적화 흐름을 만든다.
• JIT 컴파일링과 컴파일 캐시를 활용하는 자동 튜닝 프레임워크를 구현한다.
• in-process 인터페이스와 ATen 비동기 텐서 라이브러리를 통해 TC를 PyTorch와 Caffe2에 통합한다.

실험 결과

연구 질문

• RQ1 TC가 일반적인 ML 연산자와 사용자가 정의한 연산자를 간결하고 안전하게 표현할 수 있는가?
• RQ2 표준 ML 커널에 대해 TC 기반 흐름이 GPU에서 벤더 라이브러리와 경쟁할 만한 성능을 달성하는가?
• RQ3 폴리헤드 기반 옵티마이저가 딥러닝 모델의 데이터 크기/모양에 대해 연산자 융합 및 최적화를 효과적으로 수행하는가?
• RQ4 TC가 PyTorch와 Caffe2와 같은 주류 프레임워크와 얼마나 잘 통합되어 생산 지향적 경로를 제공하는가?
• RQ5 자동 튜닝과 JIT 컴파일의 영향은 하드웨어별 성능 향상을 달성하는 데 어떤가?

주요 결과

• TC 흐름은 ML 워크로드에 관련된 커널에서 NVIDIA 라이브러리보다 최대 4배의 속도 향상을 제공합니다.
• Caffe2 및 PyTorch와의 엔드투엔드 TC 통합은 실용적이고 생산 지향적 적용 가능성을 보여줍니다.
• 특화된 폴리헤드 컴파일러가 긴 의존성을 갖는 딥러닝 커널에 대해 효과적인 최적화를 제공합니다.
• 자동 튜닝 및 코드 캐시로 비표준 크기와 배열에 대한 도메인 특화 최적화를 가능하게 한다.
• 프레임워크는 안전한 인플레이스 업데이트와 일반적인 계층(예: SGEMM, conv2d, maxpool)에 대한 간단하고 선언적인 정의를 지원합니다.
• 시스템은 핵심적으로 프레임워크에 구애받지 않으면서도 생산 환경과의 긴밀한 통합을 제공합니다.