[논문 리뷰] Sparse Tensor Algebra Optimizations with Workspaces

Fredrik Kjølstad, Willow Ahrens외 2명|arXiv (Cornell University)|2018. 02. 28.

Tensor decomposition and applications참고 문헌 24인용 수 3

한 줄 요약

이 논문은 희소 텐서 대수를 위한 새로운 중간 언어인 구체적 인덱스 표기법과, 밀집 임시 텐서(워크스페이스)를 사용해 희소 계산을 가속화하는 워크스페이스 최적화를 소개한다. 루프 불변 연산을 이전에 수행하고, 비용이 많이 드는 희소 삽입 및 병합 루프를 제거함으로써, MKL 및 SPLATT와 같은 수작업 최적화 라이브러리와 비슷한 성능을 달성하며, MTTKRP와 같은 커널에서 이전에 지원되지 않았던 희소 결과를 가능하게 한다.

ABSTRACT

This paper shows how to optimize sparse tensor algebraic expressions by introducing temporary tensors, called workspaces, into the resulting loop nests. We develop a new intermediate language for tensor operations called concrete index notation that extends tensor index notation. Concrete index notation expresses when and where sub-computations occur and what tensor they are stored into. We then describe the workspace optimization in this language, and how to compile it to sparse code by building on prior work in the literature. We demonstrate the importance of the optimization on several important sparse tensor kernels, including sparse matrix-matrix multiplication (SpMM), sparse tensor addition (SpAdd), and the matricized tensor times Khatri-Rao product (MTTKRP) used to factorize tensors. Our results show improvements over prior work on tensor algebra compilation and brings the performance of these kernels on par with state-of-the-art hand-optimized implementations. For example, SpMM was not supported by prior tensor algebra compilers, the performance of MTTKRP on the nell-2 data set improves by 35%, and MTTKRP can for the first time have sparse results.

연구 동기 및 목표

• 희소 텐서 대수 컴파일러의 성능 격차를 해소하기 위해, 일시적 밀집 텐서(워크스페이스)를 사용하는 최적화를 가능하게 하기 위해.
• 이전에 비가능하거나 느렸던 커널들인 MTTKRP 및 SpMM와 같이 희소 결과를 내는 커널을 지원하기 위해, 희소 텐서 코드 생성을 일반화하기 위해.
• 계산 중에 조건부 병합 루프와 고비용의 희소 삽입을 제거하여 성능을 향상시키기 위해.
• 강력한 최적화가 가능한 형식적이고 조합 가능한 중간 표현을 제공하여, 희소 코드 생성 이전에 집요한 최적화를 가능하게 하기 위해.

제안 방법

• 루프 순서, 부분 계산 배치, 워크스페이스 사용을 명시하는 텐서 인덱스 표기법의 확장인 구체적 인덱스 표기법을 도입한다.
• 표현식을 재작성하여 부분 표현을 밀집 임시 텐서(워크스페이스)에서 사전에 계산함으로써 런타임 비용을 줄이는 워크스페이스 최적화를 적용한다.
• 스케줄링 구조를 사용해 워크스페이스 삽입을 요청함으로써, 루프 불변 코드 이동을 가능하게 하고 반복적인 희소 삽입을 피한다.
• 이전 연구(Kjolstad 등, 2017)에서 제안한 반복 그래프로 구체적 인덱스 표기법을 내림차순 변환하여 효율적인 희소 코드 생성을 가능하게 한다.
• SpMM, SpAdd, MTTKRP와 같은 핵심 커널에 최적화를 적용하여 일반성과 성능 향상을 입증한다.
• 워크스페이스를 사용해 다중 방향 병합 루프를 랜덤 액세스 패턴으로 대체함으로써 조건부 분기 수를 줄이고 데이터 국소성을 향상시킨다.

실험 결과

연구 질문

• RQ1 밀집 일시 텐서(워크스페이스)를 사용해 희소 텐서 대수를 어떻게 최적화할 수 있을까?
• RQ2 워크스페이스 최적화는 고수준 중간 언어에서 표현되고 컴파일될 수 있는가? 이는 부분 계산이 언제이고 어디서 발생하는지를 정확히 제어할 수 있는가?
• RQ3 이 최적화는 MTTKRP 및 SpMM와 같은 커널에서 이전에 지원되지 않았던 희소 결과를 가능하게 할 수 있는가?
• RQ4 최적화된 코드의 성능은 MKL 및 SPLATT와 같은 수작업 최적화된 최신 라이브러리와 비교해 어떻게 되는가?
• RQ5 이 최적화는 연구된 사례를 넘어서 광범위한 희소 텐서 커널 클래스에 일반화될 수 있는가?

주요 결과

• 워크스페이스 최적화는 이전 텐서 대수 컴파일 기술 대비 nell-2 데이터셋에서 MTTKRP 성능을 35% 향상시켰다.
• 이전 텐서 대수 컴파일러에서 지원되지 않았던 SpMM가 이제 효율적으로 컴파일되며, 수작업 최적화 구현과 경쟁 가능한 성능을 보였다.
• 이 방법은 MTTKRP에 대해 처음으로 희소 결과를 가능하게 하여, 지원되는 텐서 연산의 범위를 넓혔다.
• 최적화로 인해 고비용의 희소 삽입과 조건부 병합 루프가 제거되었으며, 이는 효율적인 밀집 워크스페이스 랜덤 액세스로 대체되었다.
• 최종 코드는 MKL, Eigen, SPLATT와 같은 최신 라이브러리 수준의 성능을 달성하여, 이 방법의 효과성을 입증했다.