[논문 리뷰] Accelerating SpMM Kernel with Cache-First Edge Sampling for Graph Neural Networks

Chien‐Yu Lin, Liang Luo외 1명|arXiv (Cornell University)|2021. 04. 21.

Advanced Graph Neural Networks참고 문헌 30인용 수 5

한 줄 요약

이 논문은 그래프를 GPU 공유 메모리에 맞게 엣지 샘플링을 통해 GNN 추론을 가속화하는 캐시 최우선 엣지 샘플링 메커니즘과 코드생산된 SpMM 커널인 ES-SpMM을 제안한다. 이는 계산량을 줄이고 캐시 국지성을 향상시킨다. 표준 GNN 벤치마크에서 정확도 손실 없이 cuSPARSE 대비 최대 4.35배 빠르게 작동하며, 정확도 손실이 1% 미만일 경우 45.3배 빠르게 작동한다.

ABSTRACT

Graph neural networks (GNNs), an emerging deep learning model class, can extract meaningful representations from highly expressive graph-structured data and are therefore gaining popularity for wider ranges of applications. However, current GNNs suffer from the poor performance of their sparse-dense matrix multiplication (SpMM) operator, even when using powerful GPUs. Our analysis shows that 95% of the inference time could be spent on SpMM when running popular GNN models on NVIDIA's advanced V100 GPU. Such SpMM performance bottleneck hinders GNNs' applicability to large-scale problems or the development of more sophisticated GNN models. To address this inference time bottleneck, we introduce ES-SpMM, a cache-first edge sampling mechanism and codesigned SpMM kernel. ES-SpMM uses edge sampling to downsize the graph to fit into GPU's shared memory. It thus reduces the computation cost and improves SpMM's cache locality. To evaluate ES-SpMM's performance, we integrated it with a popular GNN framework, DGL, and tested it using representative GNN models and datasets. Our results show that ES-SpMM outperforms the highly optimized cuSPARSE SpMM kernel by up to 4.35x with no accuracy loss and by 45.3x with less than a 1% accuracy loss.

연구 동기 및 목표

• GPU에서 느린 희소-밀집 행렬 곱셈(SpMM)으로 인한 GNN 추론의 성능 저하 문제를 해결하기 위해.
• 모델 정확도를 훼손하지 않으면서 SpMM 계산 시간을 단축하고 캐시 국지성을 향상시키기 위해.
• 엔드포인트 전처리 오버헤드를 제거하기 위해 엣지 샘플링을 SpMM 커널 내부에 직접 통합하기 위해.
• 사용자 코드를 수정하지 않고도 기존 GNN 프레임워크(DGL 등)와 원활하게 통합하기 위해.
• 다양한 GNN 워크로드에 적용 가능한 일반적인 고성능 SpMM 솔루션을 제공하기 위해.

제안 방법

• CPU 전처리를 제거하고 GPU에서 샘플링을 수행하는 커널 내부 엣지 샘플링 메커니즘을 도입하여 GPU 병렬성을 활용한다.
• 공유 메모리 크기를 샘플링 한계로 삼는 캐시 최우선 설계를 통해 샘플된 그래프가 완전히 공유 메모리에 들어오도록 보장한다.
• 공유 메모리에 샘플된 데이터를 효율적으로 로드하기 위해 연속적인 글로벌 메모리 액세스 패턴을 사용한다.
• 데이터셋 특성에 따라 속도와 정확도를 균형 잡는 두 가지 경량 샘플링 전략인 Bucket 및 FastRand를 활용한다.
• 형식 변환 오버헤드를 방지하기 위해 표준 CSR 형식을 지원하며, DGL에 대한 드롭인 대체가 가능하고 사용자 코드 변경 없이 구현된다.
• 샘플링 논리를 SpMM 커널과 코드생산하여 메모리 액세스와 계산 감소 측면에서 일관된 최적화를 이룬다.

실험 결과

연구 질문

• RQ1 SpMM 커널 내부에 엣지 샘플링을 통합함으로써 전처리 오버헤드를 제거하고 추론 속도를 향상시킬 수 있는가?
• RQ2 공유 메모리에 샘플된 그래프가 완전히 들어오도록 하는 캐시 최우선 샘플링 전략이 SpMM 성능을 크게 향상시킬 수 있는가?
• RQ3 커널 내부 엣지 샘플링을 통해 정확도 손실 최소화로 SpMM 성능을 얼마나 빠르게 향상시킬 수 있는가?
• RQ4 cuSPARSE와 같이 최적화된 SpMM 커널과 비교해 ES-SpMM의 속도와 정확도는 어떻게 되는가?
• RQ5 ES-SpMM은 사용자 수준 코드를 수정하지 않고도 기존 GNN 프레임워크에 원활하게 통합될 수 있는가?

주요 결과

• 표준 GNN 벤치마크에서 ES-SpMM은 정확도 손실 없이 매우 최적화된 cuSPARSE SpMM 커널 대비 최대 4.35배 빠른 성능을 달성한다.
• 정확도 손실이 1% 미만일 경우, ES-SpMM은 cuSPARSE 대비 45.3배 빠른 성능을 보이며 놀라운 성능-정확도 트레이드오프를 입증한다.
• 표준 GCN 모델에서 추론 시간의 대부분(최대 95%)이 SpMM에 소요되며, 이는 SpMM이 주요 성능 저하 요인임을 확인한다.
• CPU 기반 전처리 대비 커널 내부 샘플링 메커니즘이 오버헤드를 줄여 더 빠르고 효율적인 샘플링을 가능하게 한다.
• ES-SpMM은 GCN 및 GraphSAGE와 같은 다양한 GNN 모델과 Reddit, Cora, PubMed, Ppi 등의 데이터셋에서 높은 성능을 유지한다.
• DGL와의 통합은 투명하며 사용자 수준 코드 변경이 필요하지 않으며, 형식 변환 없이 표준 CSR 형식을 지원한다.