[논문 리뷰] Minimum Cost Loop Nests for Contraction of a Sparse Tensor with a Tensor Network

Raghavendra Kanakagiri, Edgar Solomonik|arXiv (Cornell University)|2023. 07. 11.

Tensor decomposition and applicationsMathematics인용 수 3

한 줄 요약

이 논문은 스parse 텐서 컨트랙션과 텐서 네트워크(SpTTN)를 위한 가장 효율적인 루프 네스트를 자동으로 식별하고 실행하는 프레임워크인 SpTTN-Cyclops를 제시한다. 동적 프로그래밍을 활용하여 버퍼 크기 및 캐시 미스와 같은 비용 지표를 최소화한다. 이는 일반적인 라이브러리보다 뛰어나고, 데이터에 독립적인 희소성 패턴을 활용해 텐서 분해 및 완성 워크로드에서 특화된 코드와 유사한 성능을 달성한다.

ABSTRACT

Sparse tensor decomposition and completion are common in numerous applications, ranging from machine learning to computational quantum chemistry. Typically, the main bottleneck in optimization of these models are contractions of a single large sparse tensor with a network of several dense matrices or tensors (SpTTN). Prior works on high-performance tensor decomposition and completion have focused on performance and scalability optimizations for specific SpTTN kernels. We present algorithms and a runtime system for identifying and executing the most efficient loop nest for any SpTTN kernel. We consider both enumeration of such loop nests for autotuning and efficient algorithms for finding the lowest cost loop-nest for simpler metrics, such as buffer size or cache miss models. Our runtime system identifies the best choice of loop nest without user guidance, and also provides a distributed-memory parallelization of SpTTN kernels. We evaluate our framework using both real-world and synthetic tensors. Our results demonstrate that our approach outperforms available generalized state-of-the-art libraries and matches the performance of specialized codes.

연구 동기 및 목표

• 텐서 분해 및 완성에서 스parse 텐서 컨트랙션과 희소하지 않은 텐서 네트워크(SpTTN)가 주요 계산 커널이 되는 성능 저하 문제를 해결하기 위해.
• 사용자 지시 없이도 가능한 루프 순서의 조합 폭발 문제를 해결하면서도 가장 효율적인 루프 네스트를 자동으로 선택하기 위해.
• 기존 고성능 텐서 라이브러리와 통합된 일반화되고 자동 튜닝된 프레임워크를 통해 고성능, 확장성 있고 이식 가능한 SpTTN 실행을 가능하게 하기 위해.
• 데이터에 독립적인 희소성 패턴을 활용하고 BLAS 커널의 효율적 사용을 가능하게 하여 수작업 최적화된 특화된 코드의 성능을 따라하거나 초월하기 위해.

제안 방법

• 각 노드가 루프이고 자식 노드가 내포된 루프인 트리로 루프 네스트를 표현함으로써, 체계적인 비용 모델링을 가능하게 한다.
• 지수적 공간을 효율적으로 탐색하기 위해 동적 프로그래밍을 적용함으로써, 가능한 루프 순서의 수를 O((m!)^N)에서 O(N^3 2^m m)로 감소시킨다.
• 루프 네스트 트리 구조에 기반한 주요 성능 저하 요인인 버퍼 크기와 캐시 미스에 대한 비용 모델을 수립한다.
• 분산 메모리 병렬 처리를 가능하게 하기 위해 CTF 런타임 시스템과 통합한다.
• 고성능을 위해 BLAS-1 및 BLAS-2 커널을 고수준으로 활용할 수 있도록 하면서도 비용 함수를 최소화하는 루프 순서를 자동으로 선택한다.
• 단일 스parse 텐서를 동시에 여러 개의 밀도 높은 텐서와 컨트랙션할 수 있도록 사이클롭스 텐서 프레임워크를 확장한다.

실험 결과

연구 질문

• RQ1 메모리 프로파일과 캐시 미스를 최소화하는 데 최적의 루프 네스트 구조는 무엇인가요?
• RQ2 완전한 열거 없이도 SpTTN 커널의 광범위한 루프 순서 순열 공간을 효율적으로 탐색하기 위해 동적 프로그래밍을 사용할 수 있을까요?
• RQ3 자동 생성된 루프 네스트의 성능은 텐서 분해 및 완성 워크로드에서 수작업 최적화된 특화된 구현과 비교해 어떻게 되나요?
• RQ4 중간 텐서의 차원 크기 제한이 루프 네스트의 효율성과 BLAS 커널 활용도에 어떤 영향을 미치나요?
• RQ5 일반화된 프레임워크가 다양한 SpTTN 커널에서 이식 가능성을 유지하면서도 특화된 코드와 유사한 성능을 달성할 수 있을까요?

주요 결과

• 작고 매우 희소한 텐서(N=40, 희소성 0.1%)에서 TACO 대비 534배의 성능 향상을 기록하여, 저차원 영역에서 뚜렷한 성능 향상을 입증했다.
• 실제 워크로드와 합성 SpTTN 커널 전반에 걸쳐 TACO 및 SparseLNR와 같은 일반 라이브러리보다 뛰어난 성능을 보였다.
• 스칼라 중간 텐서보다 크기가 T 및 S×T인 중간 텐서를 사용하는 루프 네스트가 더 높은 메모리 사용에도 불구하고 BLAS-1 및 BLAS-2 커널의 더 나은 활용도 덕분에 더 뛰어난 성능을 보였다.
• SpTTN-Cyclops가 선택한 루프 순서는 무작위로 샘플링한 25%의 루프 순서 중에서 가장 높은 성능을 기록하여 효과적인 자동 튜닝을 입증했다.
• MTTKRP, TTMc, TTTP 커널에서 특화된 코드의 성능을 그대로 따라하거나 초월하여, 이 프레임워크의 일반성과 효율성을 검증했다.
• 동적 프로그래밍을 사용함으로써 탐색 공간의 복잡도를 계승적으로 감소시켜, SpTTN 루프 네스트의 실용적인 자동 튜닝을 가능하게 했다.