[논문 리뷰] The CoRa Tensor Compiler: Compilation for Ragged Tensors with Minimal Padding

Pratik Fegade, Tianqi Chen외 2명|arXiv (Cornell University)|2021. 10. 19.

Parallel Computing and Optimization Techniques인용 수 8

한 줄 요약

CoRa는 최소한의 패딩을 사용하여 불규칙한 텐서 연산자를 컴iles하는 텐서 컴파일러로, 불규칙한 데이터 형상에 최적화된 코드를 생성함으로써 CPU 및 GPU에서 효율적인 실행을 가능하게 한다. CoRa는 NVIDIA V100 GPU에서 PyTorch 대비 1.6배의 기하평균 성능 향상을 달성했고, 64코어 ARM CPU에서 TensorFlow 대비 1.37배의 성능 향상을 보였다. 이는 트랜스포머의 멀티헤드 어텐션 연산에서의 성능 향상이다.

ABSTRACT

There is often variation in the shape and size of input data used for deep learning. In many cases, such data can be represented using tensors with non-uniform shapes, or ragged tensors. Due to limited and non-portable support for efficient execution on ragged tensors, current deep learning frameworks generally use techniques such as padding and masking to make the data shapes uniform and then offload the computations to optimized kernels for dense tensor algebra. Such techniques can, however, lead to a lot of wasted computation and therefore, a loss in performance. This paper presents CoRa, a tensor compiler that allows users to easily generate efficient code for ragged tensor operators targeting a wide range of CPUs and GPUs. Evaluating CoRa on a variety of operators on ragged tensors as well as on an encoder layer of the transformer model, we find that CoRa (i)performs competitively with hand-optimized implementations of the operators and the transformer encoder and (ii) achieves, over PyTorch, a 1.6X geomean speedup for the encoder on an Nvidia GPU and a 1.86X geomean speedup for the multi-head attention module used in transformers on an ARM CPU.

연구 동기 및 목표

• 변형된 형상의 데이터를 처리할 때 패딩과 마스킹으로 인한 성능 저하 문제를 해결하기 위해.
• CPU 및 GPU와 같은 다양한 하드웨어 기반에서 불규칙한 텐서 연산자에 대해 효율적이고 이식 가능한 컴파일을 가능하게 하기 위해.
• 불규칙한 텐서 연산을 위한 고성능 코드 생성을 자동화함으로써 수작업 최적화된 커널의 필요성을 줄이기 위해.
• 기존의 밀도 높은 텐서 및 희소 텐서 컴파일러가 불규칙한 텐서 연산을 효과적으로 표현하거나 최적화하지 못하는 한계를 극복하기 위해.
• CPU 및 GPU 백엔드를 모두 지원하는 통합적이고 확장 가능한 컴파일러 스택을 제공하여 형상의 비정규성으로 인한 성능 손실를 최소화하기 위해.

제안 방법

• 변형된 크기의 텐서 슬라이스와 비정규적인 루프 범위를 표현할 수 있는 도메인 특화 언어(DSL)를 설계하기 위해.
• 고수준의 연산자 기술서를 최적화된 LLVM 또는 CUDA 코드로 변환하는 코드 생성 파이프라인을 도입하여, 형상 인식 스케줄링을 통해 패딩을 피하기 위해.
• 타일링, 융합, 병렬화와 같은 루프 변환 기법을 활용하여 비정규적인 루프 범위를 처리하면서도 데이터 국소성과 명령어 수준의 병렬성을 유지하기 위해.
• 제약 조건 해결을 위해 Z3를 통합하고, 저수준 코드 생성을 위해 LLVM을 활용하여 메모리 접근 패턴과 커널 실행 구성 설정의 자동 최적화를 가능하게 하기 위해.
• CPU 및 GPU 백엔드를 모두 지원하기 위해 하드웨어 인식 최적화를 적용한 타겟별 코드 생성을 수행하며, GPU에서는 레지스터 블로킹 및 워프 수준의 프리미티브를 포함하기 위해.
• 모듈러 아키텍처를 사용하여 연산자 기술서와 코드 생성을 분리함으로써 새로운 연산자 및 하드웨어 플랫폼에 대한 쉽게 확장 가능하도록 하기 위해.

실험 결과

연구 질문

• RQ1 패딩이나 마스킹에 의존하지 않고도 텐서 컴파일러가 불규칙한 텐서 연산자에 대해 고성능 코드를 생성할 수 있는가?
• RQ2 CoRa에서 자동 생성된 코드의 성능은 다양한 하드웨어 플랫폼에서 수작업 최적화된 구현과 비교해 어떻게 되는가?
• RQ3 트랜스포머 모델에서 패딩 기반 접근 방식에 비해 CoRa는 계산 낭비를 어느 정도 줄일 수 있는가?
• RQ4 CoRa는 이식성과 낮은 구현 노력에도 불구하고 CPU 및 GPU 워크로드에서 경쟁 가능한 성능을 달성할 수 있는가?
• RQ5 CoRa는 멀티헤드 어텐션 및 트랜스포머 인코더 레이어와 같은 복잡한 실세계 연산자에서 얼마나 효과적으로 최적화할 수 있는가?

주요 결과

• CoRa는 트랜스포머 인코더 레이어에서 NVIDIA V100 GPU에서 PyTorch 대비 1.6배의 기하평균 성능 향상을 달성하여 수작업 최적화된 커널과 경쟁 가능한 성능을 보였다.
• 64코어 ARM Neoverse N1 CPU에서 CoRa는 멀티헤드 어텐션 모듈에서 TensorFlow 대비 1.37배의 기하평균 성능 향상을 달성했으며, 기준 대비 뚜렷한 성능 향상을 보였다.
• CoRa가 생성한 커널의 성능은 수작업 최적화된 구현과 경쟁 가능했으며, 이는 자동 코드 생성이 전문가 수준의 최적화를 따라할 수 있음을 시사한다.
• 패딩을 피함으로써 낭비되는 계산을 줄였으며, 특히 패딩 오버헤드가 상당해지는 대용량 배치 환경에서 FLOP 분석을 통해 그 효과가 뚜렷하게 드러났다.
• GPU와 같은 고도로 병렬화된 아키텍처에서도 성능을 저하시키지 않고 불규칙한 루프 구조를 효과적으로 처리함으로써 성능을 유지했다.
• 8개의 NLP 데이터셋에 대한 평가를 통해 CoRa의 성능가 다양한 시퀀스 길이와 데이터 분포에서 안정적이며, 일반화 능력이 검증되었다.