[논문 리뷰] SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning

Zihao Ye, Ruihang Lai외 3명|arXiv (Cornell University)|2022. 07. 11.

Advanced Neural Network Applications인용 수 5

한 줄 요약

SparseTIR는 스퍼스 딥 러닝 컴파일을 위한 조합 가능한 중간 표현을 도입하여, 조합 가능한 형식과 변환을 통해 효율적인 최적화를 가능하게 한다. 이는 스퍼스 연산자에서 메이커 라이브러리 대비 1.05–7.45×의 성능 향상을 달성하고, 종합적인 GNN 추론에서 최대 40.18×의 성능 향상을 이룬다. 이는 동적으로 최적의 스퍼스 형식과 하드웨어에 맞는 최적화를 조합함으로써 달성된다.

ABSTRACT

Sparse tensors are rapidly becoming critical components of modern deep learning workloads. However, developing high-performance sparse operators can be difficult and tedious, and existing vendor libraries cannot satisfy the escalating demands from new operators. Sparse tensor compilers simplify the development of operators, but efficient sparse compilation for deep learning remains challenging because a single sparse format cannot maximize hardware efficiency, and single-shot compilers cannot keep up with latest hardware and system advances. In this paper, we observe that the key to addressing both these challenges is to leverage composable formats and composable transformations. We propose SparseTIR, a sparse tensor compilation abstraction that offers composable formats and composable transformations for deep learning workloads. SparseTIR constructs a search space over these composable components for performance tuning. With these improvements, SparseTIR obtains consistent performance speedups vs vendor libraries on GPUs for single operators: 1.20-2.34x for GNN operators, 1.05-2.98x for sparse attention operators, and 0.56-7.45x for sparse convolution operators. SparseTIR also accelerates end-to-end GNNs by 1.08-1.52x for GraphSAGE training, and 4.20-40.18x for RGCN inference.

연구 동기 및 목표

• 새로운 다양한 워크로드를 지원하지 못하는 기존 메이커 라이브러리로 인해 스퍼스 연산자 개발의 고성능 달성을 어렵게 하는 문제를 해결한다.
• 다양한 하드웨어나 변화하는 딥 러닝 워크로드에 대응할 수 없는 단일 형식 스퍼스 컴파일러의 한계를 극복한다.
• 형식과 스케줄링 결정을 분리함으로써 스퍼스 연산자에 대한 효율적이고 자동화된 성능 튜닝을 가능하게 한다.
• 조합 가능한 컴파일 추상화를 통해 GNN과 스퍼스 트랜스포머와 같은 복잡한 딥 러닝 워크로드의 종단 간 가속을 지원한다.
• 기존의 밀도 텐서 컴파일러 최적화를 재사용하면서도 새로운 스퍼스 전용 변환을 가능하게 하는 확장성 있고 스케일 수 있는 프레임워크를 제공한다.

제안 방법

• 스퍼스 형식 기술과 계산 기술을 분리하는 조합 가능한 중간 표현(IR)인 SparseTIR을 도입한다.
• 스퍼스 텐서의 각 부분이 국소적 스퍼스 패턴에 따라 다른 최적의 형식에 저장될 수 있도록 형식의 조합 가능성을 보장한다.
• 컴파일을 재사용 가능하고 조합 가능한 프로그램 변환의 시퀀스로 분해함으로써 변환의 조합 가능성을 실현한다.
• 기존의 밀도 텐서 컴파일러에서 유래한 루프 수준 추상화(예: 텐서화, 타일링)를 활용하고, 이를 스퍼스 데이터 처리에 맞게 확장한다.
• 사용자 정의 가능한 검색 시스템을 통해 조합 가능한 형식과 변환에 대한 검색 공간을 구성하여 자동 성능 튜닝을 이끈다.
• 기존 컴파일러에서 유래한 하드웨어 특화 최적화(예: 텐서 코어, GPU 매핑)를 통합하여 현대 가속기에서 스퍼스 커널의 성능을 가속화한다.

실험 결과

연구 질문

• RQ1 조합 가능한 형식은 다양한 스퍼스 딥 러닝 워크로드에서 성능 이식성과 하드웨어 활용도를 향상시킬 수 있는가?
• RQ2 조합 가능한 변환은 밀도 텐서 컴파일러 최적화의 효율적 재사용을 스퍼스 컴파일링에 가능하게 할 수 있는가?
• RQ3 조합 가능한 구성 요소에 기반한 검색 기반 튜닝 시스템이 새로운 스퍼스 연산자에서 메이커 최적화 라이브러리보다 뛰어난 성능을 낼 수 있는가?
• RQ4 SparseTIR는 기존 라이브러리 대비 종단 간 GNN 학습 및 추론을 얼마나 빠르게 가속화할 수 있는가?
• RQ5 형식의 조합 가능성이 스퍼스 커널에서 텐서 코어와 메모리 계층과 같은 하드웨어 기능을 더 효과적으로 활용하는 데 어떻게 기여하는가?

주요 결과

• GNN SpMM 연산자에서 SparseTIR는 메이커 라이브러리 대비 1.20–2.34×의 성능 향상을 달성했으며, 특히 불규칙한 스퍼스 패턴에서 가장 높은 성능 향상을 보였다.
• 스퍼스 어텐션 연산자에서는 1.05–2.98×의 성능 향상을 기록하여 중간 정도의 스퍼스성과 구조적 스퍼스성에서도 효과적임을 입증했다.
• 스퍼스 컨볼루션 연산자에서는 최대 7.45×의 성능 향상을 달성했으며, 포인트 클라우드 기반 워크로드에서 메이커 라이브러리보다 뚜렷한 성능 우월성을 보였다.
• 종합적인 GraphSAGE 학습은 1.08–1.52× 빠르게 되었으며, 전체 모델 학습 파이프라인에서 일관된 성능 향상을 보였다.
• RGCN 추론에서는 4.20×에서 40.18×까지의 놀라운 성능 향상을 기록했으며, 이는 복잡하고 이질적인 스퍼스성에서 조합 최적화의 이점을 잘 보여준다.
• 산출물은 다양한 GPU(예: RTX 3070, V100)에서 재현 가능하며, L2 캐시 플래시 기능을 활성화한 경우 논문에서 보고한 수치와 매우 유사한 결과를 보였다.