[논문 리뷰] GE-SpMM: General-purpose Sparse Matrix-Matrix Multiplication on GPUs for Graph Neural Networks

Guyue Huang, Guohao Dai외 2명|arXiv (Cornell University)|2020. 07. 07.

Advanced Graph Neural Networks참고 문헌 27인용 수 6

한 줄 요약

GE-SpMM는 그래프 신경망(GNNs)에 최적화된 일반 목적의 희소 행렬-행렬 곱셈(SpMM) 커널을 GPU에서 제공하며, 압축 희소 행렬 형식(CSR)을 사용하여 사전 처리 오버헤드를 제거한다. Coalesced Row Caching와 Coarse-grained Warp Merging를 도입하여 메모리 액세스를 최적화하고 중복 로딩을 줄여, cuSPARSE 대비 최대 1.41배의 성능 향상을 달성하였으며, DGL 및 PyG와 같은 GNN 프레임워크에서 최대 3.67배의 종단 간 성능 향상을 이룬다.

ABSTRACT

Graph Neural Networks (GNNs) have achieved significant improvements in various domains. Sparse Matrix-Matrix multiplication (SpMM) is a fundamental operator in GNNs, which performs a multiplication between a sparse matrix and a dense matrix. Accelerating SpMM on parallel hardware like GPUs can face the following challenges: From the GNN application perspective, the compatibility needs to be considered. General GNN algorithms require SpMM-like operations (e.g., pooling) between matrices, which are not supported in current high-performance GPU libraries (e.g., Nvidia cuSPARSE). Moreover, the sophisticated preprocessing in previous implementations will lead to heavy data format conversion overheads in GNN frameworks. From the GPU hardware perspective, optimizations in SpMV (Sparse Matrix-Vector) designs on GPUs do not apply well to SpMM. SpMM exposes the column-wise parallelism in the dense output matrix, but straightforward generalization from SpMV leads to inefficient, uncoalesced access to sparse matrix in global memory. The sparse row data can be reused among GPU threads, which is neither possible in SpMM designs inherited from SpMV. To tackle these challenges, we propose GE-SpMM. GE-SpMM performs SpMM-like operation on sparse matrices represented in the most common Compressed Sparse Row (CSR) format, so it can be embedded in GNN frameworks with no preprocessing overheads and support general GNN algorithms. We introduce the Coalesced Row Caching method to process columns in parallel and ensure coalesced access to sparse matrix data. We also present the Coarse-grained Warp Merging to reduce redundant data loading among GPU warps. Experiments on a real-world graph dataset show that GE-SpMM achieves up to 1.41X speedup over Nvidia cuSPARSE and up to 1.81X over GraphBLAST. We also embed GE-SpMM in GNN frameworks and get up to 3.67X speedup over popular GNN models like GCN and GraphSAGE.

연구 동기 및 목표

• 기존 GPU 라이브러리인 cuSPARSE와 같이 일반 SpMM 유사 연산(예: 풀링)을 지원하지 않아 GNN 프레임워크 성능에 제약이 되는 문제를 해결한다.
• GNN 프레임워크에서 비용이 많이 드는 데이터 형식 변환 오버헤드를 제거하기 위해 원천적으로 CSR 형식을 지원함으로써 사전 처리 없이 원활한 통합을 가능하게 한다.
• GPU SpMM에서 비효율적인 메모리 액세스 패턴을 해결하기 위해 열 기반 병렬성과 워프 간 중복 데이터 로딩을 줄이기 위해 기술을 적용한다.
• 표준이 아닌 감소 연산(예: max 풀링)을 사용하는 다양한 GNN 모델에 대한 고성능 가속을 가능하게 한다. 이는 메이저 라이브러리에서 원천적으로 지원되지 않는 기능이다.

제안 방법

• 사전 처리나 형식 변환 없이 CSR 형식의 희소 행렬을 직접 소비할 수 있는 CSR 네이티브 SpMM 커널을 설계하여 GNN 파ip라인의 오버헤드를 제거한다.
• 행렬 열을 병렬로 처리하고 공유된 메모리 액세스 패턴을 보장함으로써 글로벌 메모리 대역폭 활용도를 향상시키기 위해 Coalesced Row Caching를 도입한다.
• GPU 워프 간 희소 행 데이터를 재사용하여 중복된 메모리 로딩을 줄임으로써 중복 로딩을 최소화하기 위해 Coarse-grained Warp Merging를 구현한다.
• GPU 아키텍처의 특성을 고려하여 행 기반 재사용과 열 기반 병렬성을 동시에 활용하도록 스레드 및 메모리 액세스 패턴을 최적화한다.
• 감소 연산을 추상화함으로써 일반 SpMM 유사 연산(예: max 풀링)을 지원함으로써 다양한 GNN 모델에의 통합 유연성을 확보한다.
• DGL 및 PyG와 같은 주요 GNN 프레임워크에 커널을 통합하여 종단 간 학습 워크로드에서 실제 성능 향상을 평가한다.

실험 결과

연구 질문

• RQ1 사전 처리나 형식 변환 없이도 표준 SpMM와 SpMM 유사 연산(예: max 풀링)을 모두 지원할 수 있는 일반 목적의 SpMM 커널을 설계할 수 있는가?
• RQ2 SpMM에서 GPU 메모리 액세스 패턴을 어떻게 최적화할 수 있는가? 특히 워프 간 중복 데이터 로딩을 줄이고 공유된 액세스 패턴을 확보할 수 있는가?
• RQ3 CSR 네이티브 SpMM 커널이 GNN 워크로드에서 기존 고성능 라이브러리인 cuSPARSE 및 GraphBLAST에 비해 얼마나 뛰어난 성능을 보일 수 있는가?
• RQ4 DGL 및 PyG와 같은 실제 GNN 프레임워크에 SpMM 커널을 통합했을 때, 다양한 모델 아키텍처와 데이터셋에서 어떤 성능 향상을 기대할 수 있는가?

주요 결과

• GE-SpMM는 실제 그래프 데이터셋에서 Nvidia cuSPARSE 대비 최대 1.41배, GraphBLAST 대비 최대 1.81배의 성능 향상을 달성한다.
• Coalesced Row Caching 기술만으로도 열 기반 처리 중 효율적인 공유된 메모리 액세스를 가능하게 하여 최대 1.25배의 성능 향상을 기록한다.
• Coarse-grained Warp Merging 기술은 GPU 워프 간 중복 데이터 로딩을 줄여 최대 1.51배의 성능 향상을 기여한다.
• 종단 간 GNN 학습에서 DGL에 통합된 GE-SpMM는 CUDA 실행 시간을 최대 3.67배 감소시키며, PyG에선 최대 2.10배 감소시킨다. 이는 모델 구성에 따라 달라진다.
• GraphSAGE-pool에서 max 풀링과 같은 SpMM 유사 연산의 경우, Pubmed 그래프에서 DGL의 네이티브 구현 대비 GE-SpMM가 최대 6.15배의 성능 향상을 기록한다.
• GTX 1080Ti 및 RTX 2080과 같은 다양한 GPU에서 동일한 성능 향상이 유지되며, 다양한 특징 벡터 길이와 모델 깊이에서도 일관된 성능 향상을 보인다.