NumPy 2.5에서 np.searchsorted를 최대 25배 빠르게 만들기

AI 요약

np.searchsorted는 정렬된 배열에서 각 검색 키가 들어갈 위치를 찾는 함수입니다. 히스토그램 계산과 구간 조회 등에 쓰이며 SciPy와 scikit-learn을 비롯한 과학 계산 생태계에도 영향을 줍니다. 이 글은 여러 키를 한꺼번에 검색하는 상황에서 NumPy의 이진 탐색을 개선한 과정과 NumPy 2.5에 반영된 구현을 설명합니다.

여러 검색을 묶어 진행하기

기존 구현은 키마다 이진 탐색을 순차적으로 수행합니다. 각 단계에서 중간 위치의 값을 읽고 비교한 뒤 다음에 읽을 위치를 정하므로, 앞선 메모리 접근이 끝나야 다음 접근을 시작하는 의존 관계가 생깁니다. 배열이 CPU 캐시보다 커지면 단계마다 캐시 미스가 발생할 수 있고, 다음 접근을 미리 처리하기도 어렵습니다.

저자는 먼저 NumPy 배열 연산으로 여러 검색의 상태를 동시에 갱신하는 버전을 만들었습니다. 각 키에 lo와 hi를 두고 검색 구간을 좁히는 방식입니다. 하지만 키마다 수렴하는 시점이 달라 아직 끝나지 않은 검색을 active 마스크로 추적해야 합니다. 이 마스크를 없애기 위해 모든 검색을 같은 횟수만큼 진행하도록 알고리즘을 바꿉니다. 각 구간이 매 단계 절반가량으로 줄어들게 하면 배열 크기 n에 대해 ceil(log2(n))회 뒤 모든 검색이 끝납니다. 이 변경만으로도 최대 2배 빨라졌습니다.

검색 상태를 줄여 C++로 옮기기

벡터화한 Python 구현은 독립된 검색을 동시에 진행해 메모리 대기 시간을 숨깁니다. 반면 검색마다 lo와 hi를 저장하므로 쿼리 수에 비례하는 추가 메모리가 필요합니다. NumPy의 네이티브 구현에 그대로 옮기기에는 부담이 있습니다.

글은 구간을 왼쪽 경계 base와 공통 길이 length로 표현합니다. 모든 쿼리의 구간 길이가 같으므로 각 쿼리의 hi를 따로 저장하지 않아도 됩니다. 쿼리별 상태는 base 하나만 두고, length는 전체 검색에서 공유합니다. C++ 구현은 첫 반복에서 모든 키가 같은 중앙값을 비교한다는 점도 활용합니다. 중앙값을 한 번만 읽고 결과 배열에 초기 상태를 바로 기록해, 별도의 초기화 쓰기와 읽기를 줄입니다. 빈 배열은 결과를 모두 0으로 처리합니다.

벤치마크와 적용 범위

벤치마크는 Apple M1 Pro와 메모리 32GB를 갖춘 MacBook Pro에서 진행했습니다. 검색 키는 10,000개로 고정하고, 정렬 배열은 최대 2³⁰개 int32 값, 즉 4GiB까지 늘렸습니다. 각 실험은 50회 반복해 최솟값을 기록했습니다. 새 구현은 NumPy 2.4보다 최대 25배 빨랐고, 벡터화한 Python 구현보다 작은 배열에서 최대 2배 빠릅니다. 배열이 커질수록 두 구현의 차이는 줄어듭니다. C++ 구현은 쿼리별 검색 상태를 위한 추가 공간 없이 O(1) 메모리만 더 씁니다.

JAX, TensorFlow, PyTorch와 비교한 실험에서는 NumPy가 경쟁력 있는 성능을 보였습니다. NumPy와 JAX는 여러 검색을 묶어 메모리 지연을 감추고, TensorFlow와 PyTorch는 검색 작업을 CPU 스레드에 나눠 처리합니다. 실험에서는 병렬 처리를 8코어로 제한하고 키 수를 20,000개로 늘렸습니다. 배열이 CPU 캐시를 벗어나면 구현별 성능이 비슷해졌습니다. 글은 스레드마다 검색 묶음을 처리하는 방식을 시험해볼 만하다고 언급하지만, 메모리 대역폭이 포화되면 코어를 더 늘려도 처리량이 늘지 않을 수 있다고 설명합니다. Eytzinger 배열 배치도 향후 탐색 방향으로 제시합니다.

Lobsters 반응

• @itamarst — 글에서 말하는 ‘벡터화’는 Python 생태계에서 쓰는 뜻입니다. 컴파일 언어로 작성한 확장을 이용해 묶음 연산을 수행한다는 의미입니다.
  • @addison — 알려주셔서 감사합니다. 몰랐으면 이 뜻을 알아채지 못했을 것 같습니다. Python 생태계가 왜 이 정의를 중심으로 발전했나요?
  • @JulianWgs — Python 생태계 밖에서 컴퓨팅 분야의 ‘벡터화’를 어떤 뜻으로 쓰시나요? Python에서 쓰는 뜻도 생태계 밖에서 사용됩니다.
  • @tomsmeding — 성능 중심 컴퓨팅에서 벡터화는 한 CPU 코어 안에서도 병렬성을 활용해 속도를 높이려고 알고리즘을 SIMD, 즉 벡터 CPU 명령어로 표현하는 뜻입니다. 제어 흐름이 거의 없고 데이터 묶음을 다루는 알고리즘이라면 숫자 너비, 하드웨어, 메모리 병목 여부에 따라 2~8배 이상 빨라질 수 있습니다. Python 생태계의 관습은 잘 모르지만, C 컴파일러는 자동 벡터화를 하고 Python 인터프리터는 하지 않으므로 C 확장으로 옮겨야 한다는 점이 연결고리일 것 같습니다. 벡터화가 없어도 컴파일된 C는 Python보다 빠르지만 PyPy 같은 JIT를 쓰면 차이가 작아지거나 사라질 수도 있습니다.
  • @JulianWgs — 맞는 말씀입니다. 벡터화는 수학에서 벡터를 다루듯 목록이나 배열 단위로 연산한다는 뜻으로도 쓰입니다. 스칼라 함수 하나를 n차원 NumPy 배열에 적용하는 np.vectorize가 예입니다. 여기서는 성능보다 표현을 간결하게 하는 게 목적입니다.
  • @tomsmeding — 맞습니다. 또 다른 흔한 용례네요. np.vectorize가 추가 처리를 하지 않고 반복문을 실행하는 방식이라면 HPC 관점에서 제가 말하는 벡터화는 아닙니다. 제가 벡터화라고 부르는 것은 jax.vmap입니다. 함수를 추적해 AST를 얻은 뒤 한 차원 높은 배열을 직접 처리하도록 바꿉니다. 보통 성능이 좋아지지만 항상 그런 것은 아닙니다.
  • @athas — 이 용어는 실제 벡터 프로세서보다 훨씬 앞선 APL의 벡터화 개념에서 온 것으로 봅니다. 제가 알기로 NumPy는 APL에서 영감을 받았습니다.
  • @tomsmeding — 그렇군요. 박사 학위 심사 전날 밤에 좋은 이야기를 들었습니다.
• @tomsmeding — 이런 글에서 문제 정의를 바꾸면서 그 사실을 말하지 않는 점이 마음에 들지 않습니다. 이진 탐색을 25배 빠르게 만든 게 아니라 여러 키를 묶은 이진 탐색을 25배 빠르게 만든 것입니다. 유용한 개선이지만 같은 문제는 아닙니다.
  • @itamarst — 제가 알기로 searchsorted()는 검색할 키 배열을 받습니다.
  • @tomsmeding — 아, 제가 잘못 읽었네요. 이제 이해했습니다. 감사합니다.
  • @amcandio — 맞습니다. 간결한 제목과 정확한 제목 사이의 균형도 있습니다. ‘여러 키를 조회할 때 np.searchsorted를 최대 25배 빠르게 만들기’라고 쓸 수도 있습니다.
• @amcandio — 공유해주셔서 감사합니다. 궁금한 점이 있으면 저자에게 질문해주세요.
  • @metahost — 어떻게 이 문제와 해결책을 발견했는지 궁금합니다.
  • @amcandio — 지난해 HPC를 더 배우고 싶어서 여러 블로그와 책을 읽었습니다. Algorithmica를 특히 추천합니다. 배운 내용이 있으면 적용할 곳을 찾습니다. NumPy에서 간단히 실험하다가 최적화 기회를 발견했습니다.
• @jamescampbell — 이런 글을 읽을 때마다 왜 이 작업에 NumPy를 쓰는지 의문입니다. 이제는 Dask 같은 걸 쓰면 되지 않나요?
  • @amcandio — @itamarst가 말했듯 Dask는 NumPy 위에서 실행되므로 생태계 비교에 넣지 않았습니다. 멀티스레딩으로 비슷한 성능을 내는 라이브러리를 포함했습니다. 알고리즘의 주된 병목은 메모리 접근입니다. 단일 스레드의 묶음 검색만으로도 메모리 수준 병렬성을 충분히 만들 수 있어 같은 머신에서 스레드를 더해도 처리량이 크게 늘지 않을 수 있습니다. 다만 직접 측정해볼 만합니다. 여러 머신을 쓰는 경우는 다른 이야기입니다.
  • @jamescampbell — 이해했습니다.
  • @itamarst — Dask의 숫자 배열 라이브러리는 NumPy 위에서 작동합니다. 그래서 NumPy를 단일 스레드에서도 더 빠르게 만드는 일이 중요합니다.

원문: Scientific Python / 번역·요약: Trawling