[논문 리뷰] Polyhedral Specification and Code Generation of Sparse Tensor Contraction with Co-Iteration

Tuowen Zhao, Tobi Popoola외 3명|arXiv (Cornell University)|2022. 08. 25.

Parallel Computing and Optimization Techniques인용 수 6

한 줄 요약

이 논문은 희소 텐서 곱셈을 위한 다각형 프레임워크 확장을 제안하며, 레이아웃 공간과 계산 공간을 분리하여 유연한 레이아웃 명세와 SMT로 합성된 find 알고리즘을 통한 효율적 공반복(iteration)을 가능하게 한다. 이는 TACO 대비 평균 1.63배의 성능 향상을 달성하며, 최적화된 find 알고리즘을 사용할 경우 최대 2.72배까지 향상된다. 이로써 기존에 구현이 어려웠던 레이아웃, 예를 들어 BCSR나 대각선 계산도 지원할 수 있다.

ABSTRACT

This paper presents a code generator for sparse tensor contraction computations. It leverages a mathematical representation of loop nest computations in the sparse polyhedral framework (SPF), which extends the polyhedral model to support non-affine computations, such as arise in sparse tensors. SPF is extended to perform layout specification, optimization, and code generation of sparse tensor code: 1) we develop a polyhedral layout specification that decouples iteration spaces for layout and computation; and, 2) we develop efficient co-iteration of sparse tensors by combining polyhedra scanning over the layout of one sparse tensor with the synthesis of code to find corresponding elements in other tensors through an SMT solver. We compare the generated code with that produced by a state-of-the-art tensor compiler, TACO. We achieve on average 1.63$ imes$ faster parallel performance than TACO on sparse-sparse co-iteration and describe how to improve that to 2.72$ imes$ average speedup by switching the find algorithms. We also demonstrate that decoupling iteration spaces of layout and computation enables additional layout and computation combinations to be supported.

연구 동기 및 목표

• 기존의 희소 텐서 컴파일러들이 계산 공간과 레이아웃 공간이 밀접하게 결합되어 있어 BCSR나 대각선 계산과 같은 복잡한 레이아웃을 지원하지 못하는 한계를 해결하기 위해.
• 논리적 좌표 공간과 물리적 레이아웃 반복 공간을 분리함으로써 다수의 희소 텐서에 대한 일반적인 공반복을 가능하게 하기 위해.
• 비형식적 함수와 SMT 솔버를 사용하여 레이아웃 명세 및 공반복 지원 기능을 확장한 희소 다각형 프레임워크(SPF)를 확장하기 위해.
• 레이아웃 변환과 공반복 패턴을 모두 지원하는 최적화된 아키텍처 인식 코드를 생성하기 위해, 특히 향후 작업으로 예정된 비결합 병합(merge) 패턴을 포함하여.
• 다양한 하드웨어, 특히 SIMD와 텐서 코어를 포함한 하드웨어에 대해 데이터 레이아웃과 계산을 공최적화함으로써 성능 이식성(성능 포트폴리오)을 향상시키기 위해.

제안 방법

• 레이아웃을 물리적 위치와 논리적 좌표 간의 수학적 관계로 모델링하여 반복 공간을 분리함으로써, 레이아웃과 계산을 별도로 최적화할 수 있도록 한다.
• 희소 계산에서 비선형 접근을 지원하기 위해 다각형 모델 내에서 런타임 의존성에 기반한 인덱스 배열을 비형식적 함수로 표현한다.
• SMT 솔버를 사용하여 한 텐서의 레이아웃 인덱스를 다른 텐서의 비제로 요소에 매핑하는 find 알고리즘을 합성한다. 이를 통해 공반복이 가능해진다.
• 한 텐서의 레이아웃에 대해 다각형 스캐닝을 수행하고, 합성된 코드를 통해 다른 텐서의 대응하는 비제로 요소를 찾는 방식으로, 명시적 병합 격자(merge lattice)가 필요 없도록 한다.
• 레이아웃 수준와 계산 인덱스 간 일对일 매핑이 필요 없도록 하여 블록 형식(예: BCSR)과 대각선 계산과 같은 복잡한 레이아웃을 지원한다.
• 레이아웃 매핑과 계산 변환을 조합하여 최적화된 코드를 생성함으로써, 차단(blocking), 재정렬(reordering), 벡터화(vectorization)와 같은 최적화와의 수직적 조합(orthogonal composition)을 가능하게 한다.

실험 결과

연구 질문

• RQ1 다각형 프레임워크는 계산 반복 공간과 분리된 상태에서 희소 텐서 레이아웃의 공식적 명세를 지원할 수 있는가?
• RQ2 공식적이고 수학적으로 탄탄한 프레임워크를 사용하여 다수의 희소 텐서에 대한 공반복을 효율적으로 표현하고 컴파일할 수 있는가?
• RQ3 SMT로 합성된 find 알고리즘이 공반복 성능에서 수동으로 작성하거나 템플릿 기반 접근 방식보다 뛰어나게 성능을 낼 수 있는가?
• RQ4 이전에 지원되지 않았던 형식, 예를 들어 BCSR나 대각선 연산과 같은 형식을 지원하기 위해, 레이아웃과 계산을 얼마나 잘 공최적화할 수 있는가?
• RQ5 다양한 희소 텐서 워크로드에서 최신 기술의 텐서 컴파일러인 TACO와 비교해 생성된 코드의 성능은 어떠한가?

주요 결과

• 제안된 프레임워크는 벤치마크 세트에서 병렬 희소-희소 공반복에서 TACO 대비 평균 1.63배의 성능 향상을 달성한다.
• 최적화된 find 알고리즘으로 전환함으로써, 프레임워크는 TACO 대비 최대 2.72배의 평균 성능 향상을 달성하여 알고리즘 합성의 성능 영향을 입증한다.
• 레이아웃과 계산 공간의 분리 덕분에 이전에 지원되지 않았던 레이아웃, 예를 들어 블록 압축 희소 행렬(BCSR)을 지원할 수 있게 되었으며, TACO는 이를 표현할 수 없다.
• 프레임워크는 TACO의 수준 기반 형식 모델에서 구현이 불가능한, 다수의 레벨에서 공통된 루프 인덱스가 필요한 대각선 계산과 같은 패턴도 지원한다.
• SMT로 합성된 find 알고리즘의 사용 덕분에 병합 격자나 사전 계산된 인덱스 매핑이 필요 없이 효율적인 공반복이 가능해졌다.
• 이 접근 방식은 블로킹, 벡터화와 같은 최적화와의 수직적 조합이 가능하여 현대 아키텍처에서 성능 향상을 이끌 수 있다.