[논문 리뷰] A High-Performance Sparse Tensor Algebra Compiler in Multi-Level IR

Ruiqin Tian, Luanzheng Guo외 3명|arXiv (Cornell University)|2021. 02. 09.

Parallel Computing and Optimization Techniques참고 문헌 65인용 수 10

한 줄 요약

COMET는 MLIR 기반의 고성능 희소 텐서 대칭 컴파일러로, 아인슈타인 표기법 유사 추상화를 사용하는 도메인 특화 언어(DSL)를 통해 혼합 희소-밀도 텐서 연산을 위한 최적화된 커널을 자동으로 생성한다. 네 개의 차원 속성을 통해 다양한 저장 형식을 지원하며, 데이터 재정렬을 통해 국소성을 향상시켜 SpMV에서 TACO 대비 최대 20.92배, SpMM에서 6.39배, TTM에서 13.9배의 성능 향상을 달성한다.

ABSTRACT

Tensor algebra is widely used in many applications, such as scientific computing, machine learning, and data analytics. The tensors represented real-world data are usually large and sparse. There are tens of storage formats designed for sparse matrices and/or tensors and the performance of sparse tensor operations depends on a particular architecture and/or selected sparse format, which makes it challenging to implement and optimize every tensor operation of interest and transfer the code from one architecture to another. We propose a tensor algebra domain-specific language (DSL) and compiler infrastructure to automatically generate kernels for mixed sparse-dense tensor algebra operations, named COMET. The proposed DSL provides high-level programming abstractions that resemble the familiar Einstein notation to represent tensor algebra operations. The compiler performs code optimizations and transformations for efficient code generation while covering a wide range of tensor storage formats. COMET compiler also leverages data reordering to improve spatial or temporal locality for better performance. Our results show that the performance of automatically generated kernels outperforms the state-of-the-art sparse tensor algebra compiler, with up to 20.92x, 6.39x, and 13.9x performance improvement, for parallel SpMV, SpMM, and TTM over TACO, respectively.

연구 동기 및 목표

• 다양한 저장 형식과 아키텍처에서 희소 텐서 커널을 수동으로 최적화하는 과제를 해결한다.
• 아인슈타인 표기법을 닮은 DSL을 통해 텐서 대칭 연산을 위한 고성능 프로그래밍을 가능하게 한다.
• MLIR 프레임워크 통합을 통해 성능 이식성과 확장성을 제공한다.
• 다양한 저장 형식에서 혼합 희소-밀도 텐서 연산을 위한 효율적인 코드 생성을 지원한다.
• 컴pile 파이프라인 내에서 자동화된 데이터 재정렬 기법을 통해 국소성과 성능을 향상시킨다.

제안 방법

• 아인슈타인 표기법 유사 구문을 사용하는 고수준 텐서 대칭 표현식을 지원하는 도메인 특화 언어(DSL)를 설계한다.
• 도메인 특화 최적화와 아키텍처 특화 최적화를 분리하기 위해 MLIR의 다층 중간 표현(IR)을 기반으로 컴파일러를 구축한다.
• 텐서 저장 형식을 네 가지 차원 속성으로 표현한다: 밀도, 압축 유일, 압축 비유일, 싱글턴.
• 차원 속성을 분석하여 비영 요소 반복을 위한 효율적이고 형식 인식 커널을 생성하는 코드 생성 알고리즘을 구현한다.
• 커널 생성 과정에서 공간적 및 시간적 국소성을 향상시키기 위해 데이터 재정렬 알고리즘을 통합한다.
• MLIR의 계층 내에서 고수준 IR에서 아키텍처에 종속되지 않는 최적화와 저수준에서 레지스터/메모리 최적화를 적용한다.

실험 결과

연구 질문

• RQ1 희소 텐서 저장 형식의 광범위한 범위를 지원하면서도 고성능을 유지할 수 있는 텐서 대칭 컴파일러는 어떻게 설계할 수 있는가?
• RQ2 아인슈타인 표기법 기반의 DSL은 희소 텐서 계산에서 프로그래머의 생산성을 얼마나 향상시킬 수 있는가?
• RQ3 MLIR 기반 컴파일링은 다양한 아키텍처와 형식 간의 성능 이식성과 확장성을 어떻게 보장할 수 있는가?
• RQ4 자동화된 데이터 재정렬 기법은 희소 텐서 연산의 국소성과 커널 성능 향상에 얼마나 효과적인가?
• RQ5 SpMV, SpMM, TTM와 같은 핵심 희소 텐서 커널에서 기존 컴파일러인 TACO에 비해 어떤 성능 향상을 달성할 수 있는가?

주요 결과

• COMET는 CPU에서 희소 행렬-벡터 곱셈(SpMV)에서 TACO 대비 최대 20.92배의 성능 향상을 달성한다.
• COMET는 CPU에서 희소 텐서-행렬 곱셈(SpMM)에서 TACO 대비 6.39배의 성능 향상을 보였다.
• COMET는 CPU에서 텐서 변환(TTM) 연산에서 TACO 대비 13.9배 빠른 성능을 발휘했다.
• 데이터 재정렬의 사용은 국소성을 크게 향상시켜 관측된 성능 향상에 기여했다.
• 네 가지 차원 속성을 통한 형식에 종속되지 않는 설계 덕분에 COMET는 COO, CSR, DCSR, ELLPACK, CSF 및 모드-유사 형식을 하드코딩하지 않고도 일반적인 형식을 지원할 수 있었다.
• MLIR 기반 아키텍처는 모듈성, 확장성, 이식성을 보장하여 향후 GPU와 같은 이종 아키텍처 지원을 용이하게 한다.