Fast Graph-based Higher-Order Clustering Statistics on the GPU

기술 요약: GPU 기반의 고차 클러스터링 통계량 계산 가속화

문제 정의
은하의 공간적 분포는 원시 밀도장과 중력 진화에 관한 결정적인 정보를 담고 있다. 2점 상관 함수(2pCF)는 이러한 분포를 분석하는 표준 도구로 오랫동안 사용되어 왔으나, 후기 비선형 중력 진화와 은하 편향(bias)에 의해 생성된 비가우스적 구조를 포착하기에는 불충분하다. 고차 통계량, 구체적으로 3점(3pCF) 및 4점(4pCF) 상관 함수는 편향 파라미터의 퇴화(degeneracy)를 깨뜨리고 원시 비가우스성(primordial non-Gaussianity) 및 패리티 위반(parity violation)을 조사하는 등 보완적인 정보를 제공한다. 그러나 대규모 우주론적 카탈로그에 대해 이러한 통계량을 계산하는 것은 계산적으로 매우 부담스럽다. 브루트 포스(brute-force) 접근 방식은 $O(N^k)$의 복잡도를 가지며, 트리 기반 방법은 이를 $O(N m^{k-1})$(여기서 $m$은 평균 이웃 수)로 줄여주지만, 현대적인 서베이 규모(예: DESI)에서 $k=3$ 및 $k=4$인 경우 비용이 여전히 높다. 또한, 기존 코드들은 패리티 분해를 포함한 연결된 4pCF의 직접적인 구성 공간(configuration-space) 구현이나, 장치 메모리를 초과하는 그래프를 처리할 수 있는 효율적인 GPU 가속 기능이 부족한 경우가 많다.

방법론
본 논문은 $k$-d 트리에 기반한 그래프 데이터베이스 접근 방식을 사용하는 알고리즘인 GRAMSCI(GRAph Made Statistics for Cosmological Information)의 버전 2를 제시한다. 핵심 방법론은 초기 좌표로부터 압축되고 정렬된 인접 구조(adjacency structure)를 구축한 다음, 원시 데이터를 다시 방문하지 않고도 $N$-튜플을 계산하도록 쿼리하는 것이다.

이번 업데이트에서는 다섯 가지 주요 방법론적 진보를 도입하였다:

1. 머지-워크(Merge-Walk) 알고리즘: 허브-스포크(hub-spoke) 쌍당 $O(m \log m)$의 비용을 발생시키던 기존의 이진 탐색 내부 루프를 머지-워크 알고리즘으로 대체하였다. 인접 리스트의 정렬된 순서를 활용함으로써, 알고리즘은 집합 교집합을 찾기 위해 선형 패스(linear pass)를 수행하여 복잡도를 $O(m)$으로 줄인다.
2. 패리티 분해 4pCF: 새로운 추정기는 4pCF 신호를 패리티-짝수($\zeta_+$) 및 패리티-홀수($\zeta_-$) 채널로 분리한다. 이는 은하 사면체의 부호 있는 부피(signed volume)를 계산함으로써 달성된다. 부동 소수점 노이즈가 패리티-홀수 채널에 침투하는 것을 방지하기 위해, 퇴화된 부피(사실상 0인 경우)를 가진 사면체는 결정론적으로 홀수 채널에서 0의 가중치를 할당한다.
3. 연결된 4pCF 추정: 코드는 이제 동일한 그래프와 빈닝(binning) 상에서 연결되지 않은 기여분(2pCF의 곱)을 내부적으로 계산한다. 이를 통해 전체 4pCF에서 연결되지 않은 항을 직접 차감하여 연결된(trispectrum) 신호를 분리할 수 있으며, 이는 원시 비가우스성을 연구하는 데 필수적이다.
4. 파이썬 인터페이스: 가벼운 래퍼(wrapper)를 통해 Fortran 엔진을 NumPy 배열으로부터 직접 호출할 수 있으며, 파일 I/O를 처리하고 구조화된 결과 객체를 반환한다.
5. Out-of-Core 지원을 갖춘 GPU 포팅: 전체 쿼리 엔진(3pCF, 4pCF 및 패리티 4pCF)을 OpenACC를 사용하여 NVIDIA GPU로 포팅하였다. 구현에는 압축 희소 행렬(CSR) 데이터 레이아웃이 사용되었으며, 하나의 스레드 블록이 하나의 허브를 처리하고 벡터 레인이 내부 이웃 루프를 처리하는 "갱-앤-벡터(gang-and-vector)" 매핑 전략을 채택하였다. 장치 메모리보다 큰 그래프(예: BAO 규모의 서베이)를 처리하기 위해, 에지(edge) 배열을 윈도우 단위로 나누어 정확한 계산을 유지하면서 그래프를 청크(chunk) 단위로 처리하는 out-of-core 타일링 기법을 사용한다.

주요 결과
저자들은 새로운 코드를 기존의 CPU 구현 및 해석적 주입(analytic injections) 모델과 벤치마킹하였다:

• 알고리즘 가속: 머지-워크 커널은 64 스레드 CPU 노드 대비 3pCF에서 $1.5\times \sim 2.0\times$, 4pCF에서 $2.8\times \sim 3.8\times$의 속도 향상을 제공하며, 최대 스케일($r_{max}$)과 이웃 수가 증가함에 따라 이 이득은 더욱 커진다.
• GPU 성능: 단일 소비자용 GPU(NVIDIA RTX 3090 Ti)에서, 코드는 64 스레드 CPU 노드 대비 3pCF에서 $2.6\times$, 4pCF에서 $9\times$의 속도 향상을 달성하였다. GPU 성능은 부동 소수점 처리량보다는 메모리 대역폭에 따라 스케일링된다.
• 확장성: Out-of-core 타일링 기법을 통해 장치 메모리를 훨씬 초과하는 그래프를 처리할 수 있다. 특정 테스트에서 24 GB 그래픽 카드에서 $9 \times 10^9$개의 에지를 가진 3pCF를 처리했을 때, 가상적인 전량 상주 실행(all-resident run) 대비 약 $20\%$의 오버헤드만을 기록하였다.
• 정확성: GPU 커널은 빈닝된 카운트에서 CPU 참조 결과와 정확히 일치한다. 패리티 위반에 대한 해석적 신호 주입 테스트 및 패리티 대칭 시뮬레이션(Quijote)에 대한 영 검정(null test)을 통한 검증은 추정기가 편향되지 않았으며 퇴화된 기하학적 구조를 올바르게 처리함을 확인시켜 준다.
• 응용: 코드는 DESI DR1 LRG 샘플(300만 개의 은하)에 적용되었다. 측정된 3pCF와 연결된 4pCF는 EZmock 앙상블 예측치를 추종하며, 3pCF는 기대되는 BAO 특징을 보여주고 연결된 4pCF는 비가우스적 신호를 분리해낸다.

의의
본 논문은 이러한 업데이트가 우주론에서 고차 통계량의 실질적인 유용성을 크게 확장한다고 주장한다. 소비자용 하드웨어에서 4pCF의 계산 비용을 거의 10배 가까이 줄임으로써, 이 코드는 공분산 추정을 위한 수천 개의 모크(mock) 구현이 필요했던 과거의 제약을 넘어 DESI와 같은 대규모 카탈로그 전체의 BAO 범위에서 연결된 4pCF를 평가하는 것을 가능하게 한다. 패리티 분해 추정기의 포함은 구면 조화 분석(spherical-harmonic analysis)과는 독립적인, 패리티 위반 탐색을 위한 직접적인 구성 공간 교차 검증을 제공한다. 마지막으로, 외부 입력 없이 연결된 4pCF를 내부적으로 계산할 수 있는 능력은 빈닝 및 쌍(pair) 세트의 일관성을 보장하며, 이는 비가우스적 트리스펙트럼 신호를 분리하는 데 매우 중요하다. 저자들은 본 코드가 공개적으로 사용 가능하며 현재 및 미래의 대규모 구조 서베이에 즉시 적용될 준비가 되었음을 결론짓는다.