NumPy의 벡터화된 배열 연산이 왜 항상 충분하지 않은지, eager 평가와 중간 배열 할당의 비용, 그리고 NumExpr·JAX·Rust 같은 대안이 어떻게 성능을 바꾸는지 살펴본다.
2026년 2월 17일 게시
(후속 글 보기)
나는 몇 달 동안 McFACTS 시뮬레이션 프로젝트에 간헐적으로 기여해 왔고, 현재는 시뮬레이션의 핵심 ‘비즈니스 로직’에 대한 꽤 큰 규모의 최적화들을 테스트하고 병합하는 작업을 진행 중이다. 이 로직은 주로 긴 NumPy 배열 연산 시퀀스로 이루어져 있다.
지금 파이프라인에 올라와 있는 최적화들은 거의 전부 기본적으로 똑같아서, 아예 코드를 대신 써 주는 매크로를 만들까 생각했을 정도다. 지난 몇 달 동안 이런 최적화가 구체적으로 어떻게 작동하는지, 그리고 그 뒤에 있는 직관이 무엇인지 내가 배운 것을 압축해 정리해 볼 가치가 있겠다고 생각했다.
이 중 어느 것도 특별히 난해하거나 하드웨어 종속적인 내용은 아니고, 마지막에는 약간의 x86 어셈블리도 보겠지만, 그건 주로 설명을 위한 것이고 우리가 쌓아 올리는 직관을 검증하는 데 도움이 된다.
이 글은 특히 과학 계산용 Python 코드를 작성하는 사람들에게 유용할 수 있다. 학부 때 NumPy 배열은 아주 빠르다고 배웠는데도 프로그램이 느리게 돌아가서 고민하고 있다면 더욱 그렇다.
하지만 먼저, 고통스러울 정도로 억지인 비유부터 하자:
Alice, Bob, Carol은 British Columbia 숲에서 일하는 벌목 엔지니어다. 매일같이 그들은 TreeChopper 30000을 운용하는데, 이 산업의 기적은 원목을 그 자리에서 즉시 건축용 목재로 바꿔 준다. 하지만 단점이 하나 있다. TreeChopper는 크고 무겁고 옮기기가 몹시 어렵다.
Alice는 버튼을 누르고, TreeChopper가 5초 만에 나무 한 그루를 산산이 찢는 모습을 지켜본 뒤, 다음으로 베어진 나무 줄기로 기계를 옮기는 데 30분을 쓴다. 실질 작업 대 이동 비율: 1:360.
Bob은 같은 버튼을 누른 뒤 트럭에 올라타 다음 나무 줄기를 TreeChopper 쪽으로 끌고 오며, 이 과정에 5분만 든다. 비율: 1:60.
Carol은 병렬 트럭 운송 작업을 구축해 새로 베어진 줄기들이 TreeChopper로 계속 들어오도록 만들었다. 근처에는 작업 대기 줄기가 쌓여 있고, 줄기 사이의 비가동 시간은 겨우 10초다. 비율: 1:2.
Toronto의 한 실험실에서 Dave는 TreeChopper 30000.0.1의 첫 설계를 스케치하고 있다. 이 버전은 최대 8개의 나무 줄기를 동시에 처리할 수 있는 초광폭 투입 레인을 갖고 있다.
Alice, Bob, Carol은 모두 이 새 기계를 단일 투입 버전을 쓰던 것과 정확히 같은 방식으로 사용할 것이며, 작업 흐름은 전혀 바뀌지 않는다.
다음 코드는 무슨 일을 할까?
import numpy as np
a = np.array([...]) b = np.array([...]) c = np.array([...]) d = np.array([...])
e = a * b f = e / c g = f**d
방금 Data Science 101을 마쳤다면 ‘벡터화된 배열 연산!’이라고 말할 것이고, 그 말은 맞다. NumPy의 배열 산술은 리스트를 순회하며 같은 연산을 수행하는 방식보다 훨씬, 훨씬 더 뛰어난 성능을 제공하며, 더 짧고 읽기도 쉽다.
a = [...] b = [...] c = [...] d = [...]
e = [a*b for a, b, in zip(a, b)] f = [e/c for e, c, in zip(e, c)] g = [f**d for f, d, in zip(f, d)]
이제 이 코드를 보고 ‘왜 단계를 이렇게 쪼개나요? 필요 이상으로 길고, 아마 더 느리기까지 할 텐데요’라고 말할 수도 있다!
import numpy as np
a = np.array([...]) b = np.array([...]) c = np.array([...]) d = np.array([...])
g = ((a*b)/c)**d
이걸 한 줄로 옮기면 더 짧아지고, 아마 더 빨라지지 않을까?
첫 번째는 분명히 맞지만, 두 번째는 적어도 의미 있는 수준에서는 그렇지 않다는 사실이 의외일 수 있다. 내부적으로는 거의 똑같다.
이것은 NumPy의 eager 식 평가의 특징이다. NumPy는 이런 연산을 왼쪽에서 오른쪽으로 처리하며, a*b 같은 유효한 연산을 발견하는 순간 즉시 수행한다. 가상의 ‘lazy’ 평가는 줄 끝에 도달할 때까지 기다렸다가 식에 여러 연산이 포함되어 있음을 보고 그것들을 한꺼번에 수행하려고 시도하겠지만, 그건 훨씬 더 복잡한 일이고 NumPy의 방식은 아니다.
다음과 같은 한 줄을 넘기든
g = ((a*b)/c)**d
혹은
e = a * b f = e / c g = f**d
NumPy는 이것을 세 개의 벡터화된 배열 연산으로 분해한다. 한 줄 버전에 대한 디스어셈블 결과는 다음과 같다:
9 178 LOAD_FAST 0 (a)
180 LOAD_FAST 1 (b)
182 BINARY_OP 5 (*)
186 LOAD_FAST 2 (c)
188 BINARY_OP 11 (/)
192 LOAD_FAST 3 (d)
194 BINARY_OP 8 (**)
198 STORE_FAST 4 (g)
그리고 여러 줄 버전은 다음과 같다.
10 178 LOAD_FAST 0 (a)
180 LOAD_FAST 1 (b)
182 BINARY_OP 5 (*)
186 STORE_FAST 4 (e)
11 188 LOAD_FAST 4 (e)
190 LOAD_FAST 2 (c)
192 BINARY_OP 11 (/)
196 STORE_FAST 5 (f)
12 198 LOAD_FAST 5 (f)
200 LOAD_FAST 3 (d)
202 BINARY_OP 8 (**)
206 STORE_FAST 6 (g)
뒤쪽이 더 길어 보이지만, 그 차이는 전적으로 중간 변수들을 저장하고 다시 불러오기 위해 추가된 네 개의 STORE_FAST / LOAD_FAST 명령 때문이다. 이 작업도 시간이 들긴 하지만, 매번 몇 나노초 수준일 뿐이다.
반대로, 실제로 유용한 작업은 모두 BINARY_OP 명령에서 이루어진다. 그럼 이 명령은 실제로 무엇을 할까?
먼저 BINARY_OP는 왼쪽 피연산자의 타입을 검사한다. 우리 경우에는 그것이 항상 np.ndarray[np.float64]라고 가정하자. 그러면 PyNumberMethods에서 적절한 슬롯(*에 대해서는 nb_multiply, /에 대해서는 nb_true_divide, **에 대해서는 nb_power)을 찾아 그 슬롯을 호출한다. 예를 들면 np.ndarray.__mul__() 같은 것이다. 그 다음 왼쪽과 오른쪽 피연산자의 타입과 기타 필요한 단계들을 확인하는데, 예를 들어 차원이 broadcast 가능한지 검사하고, np.ufunc 루프를 선택하고, 출력 배열을 힙에 할당한 뒤, 실제로 원소별 계산을 수행한다.
스택 변수를 연달아 저장했다가 다시 불러오는 행위는 아마 왼쪽 피연산자의 타입을 검사하는 것과 비슷한 시간밖에 걸리지 않을 것이다. 배열이 아주 작더라도 전체 실행 시간에서 극히 미미한 부분이다.
반대로, 이와 같은 과정을 순수 리스트 구현으로 처리하면 모든 설정 작업을 원소마다 반복해야 하고, 각 원소마다 포인터 추적도 해야 한다. 반면 배열 구현은 이 설정을 이항 연산마다 한 번만 하면 되고, 원소들은 고르게 배치된 연속적인 선형 메모리에 놓여 있다.
이런 숨겨진 복잡성들이 g = ((a*b)/c)**d 같은 깔끔하고 단순한 문법을 뒷받침하며, 엄청난 양의 포인터 추적을 피하게 해 준다. 대부분의 목적에는 괜찮다. 하지만 최고 성능이 필요하다면, 이것만으로는 한참 부족하다.
몇 문단 전에 내가 광적으로 대문자를 썼던 걸 기억하는가? 모든 단 하나의 BINARY\_OP가 힙에 배열 하나를 할당한다. 연산자 하나가 감당하기에는 엄청난 책임이다.
분명 우리가 원하는 것은 이것을 한 번의 패스로 융합하는 것이다. 대략 이런 식으로 말이다.
a_arr = np.array([...]) b_arr = np.array([...]) c_arr = np.array([...]) d_arr = np.array([...])
out_arr = np.zeros(n)
for i, (a_scalar, b_scalar, c_scalar, d_scalar), in enumerate(zip(a_arr, b_arr, c_arr, d_arr)): out_arr[i] = ((a_scalar*b_scalar)/c_scalar)**d_scalar
배열 하나만 미리 할당하고, 이 모든 배열을 병렬로 단 한 번만 순회하면서, 중간 할당 없이 내부 계산을 수행할 수 있다. 많은 할당을 한 번의 할당으로 바꾸고, 모든 배열을 병렬로 한 번만 훑고, 내부 계산을 곧장 수행하는 것이다.
이건 잘 작동해야 하지 않을까?
이론적으로는 그렇다. 실제로는, 위와 같이 구현하면 아주 작은 배열을 제외하고는 벡터화된 연산보다 훨씬 훨씬 느리다.
여기서 내부적으로 무슨 일이 일어나고 있을까?
반복마다 우리는 여러 원소를 꺼내야 한다. 이 원소들은 연속된 배열에서 오기 때문에 캐시의 이점을 누리긴 하지만, 각 C float는 np.float64 객체로 감싸져야 하고, 다시 힙에 할당되어야 한다.
뭐라고, 이미 크기가 정해져 있는 숫자 타입이 스택에 저장될 거라고 생각했나? 어리석군! 전에 ~다른 블로그 글에서~ 지난 수업에서 말했듯, Python의 모든 것은 O.B.J.E.C.T.다. 이건 Reference-Counted Heap Allocated C****Struct의 약자다! 아니, 좋은 약자는 아니다. Guido가 무슨 생각이었는지는 나도 모르겠다.
그 다음에는 앞에서 본 것과 똑같은 이항 연산 준비 작업을 해야 하는데, 이번에는 배열당 한 번이 아니라 원소마다 해야 한다. Python은 이 값들이 원래 NumPy 배열에서 왔다는 사실을 모르고, 신경도 쓰지 않기 때문이다. 우리는 이제 완전히 순수 Python 루프 영역에 들어와 있으며, 그 말은 NumPy의 멋진 보장들은 무효가 되고 그 최적화도 접근할 수 없다는 뜻이다.
그 다음 실제로 계산을 수행해 출력을 만들고 나면, 그 결과 역시 배열에 삽입되기 전에 힙 할당 객체로 변환되어야 한다. 반면 벡터화된 연산은 이런 원소별 중간 할당 없이 처리할 수 있었다.
우리는 배열 단위로 이루어지는 소수의 선행 할당을, 원소 단위로 이루어지는 훨씬 더 많은 작은 할당으로 바꿔 버린 셈이다.
아, 그리고 벡터화된 연산에는 존재하지 않는 Python 루프 오버헤드도 있다. 벡터화된 연산은 SIMD 최적화도 활용할 수 있지만, 원소 하나씩 도는 Python 루프는 거기까지 가기엔 수십 겹의 층을 더 지나야 한다.
이건 내가 예전에 길게 써 온 주제와 맞닿아 있다. 단순하다는 것은 사실 굉장히 복잡한 일이며, 관용적인 Python의 단순함은 여러 겹의 간접 참조, 힙 할당, 런타임 타입 검사 위에 세워져 있다. 그래서 표준 라이브러리에 배열 구현이 없는 대신 패키지로 그걸 다시 사 와야 할 정도다.
그렇다면 우리는 실제로 무엇을 해서 이걸 더 빠르게 만들 수 있을까?
NumExpr Python 패키지는 NumPy의 벡터화된 배열 연산보다 훨씬 덜 알려진 대안을 제공한다. 전체 배열을 한 번에 처리하거나 값을 개별적으로 순회하는 대신, 배열을 청크 단위로 나눈다. 수학적 변환을 원시 코드로 표현하는 대신 문자열로 표현하고, 이를 다시 유효한 Python 식으로 컴파일한다(문서를 봐도 이 식이 NumPy처럼 eager 평가되는지, 아니면 다른 방식을 쓰는지는 분명하지 않다).
당신의 코드가 정말 ((a*b)/c)**d 같은 형태로 분해된다면, NumExpr는 아주 적절한 선택일 수 있다.
하지만 단점도 있다. 첫째, 이 모든 것이 별도의 가상 머신 안에서 이루어지기 때문에 프로그램에 추가적인 무게가 실리고, 과정도 훨씬 덜 관찰 가능해진다.
둘째, 중간에 다른 함수를 호출하는 함수를 가속하려고 하거나, 단순한 수학식으로 표현되지 않는 연산을 수행한다면 답이 없다. 그런 경우에는 더 일반적인 해법이 필요하다.
Numba는 Python용 JIT 컴파일러다. 기존 코드에 @njit만 붙이면 루프가 자동으로 융합되는 모습을 볼 수 있다!
안타깝게도 나는 Numba를 써 본 적이 없다. 마지막으로 이걸 쓰자고 제안했을 때 공동 작업자가 내가 연어로 그녀의 얼굴을 후려친 것처럼 쳐다봤기 때문이다.
다행히 JAX가 있고, 사실상 거의 같은 일을 하면서 모든 면에서 더 낫다.
심지어 GPU에서도 실행할 수 있다!
배열을 불변으로 생각해야 하긴 하지만, 나는 대체로 그 점을 장점으로 본다.
하지만 GPU 없이도 더더더 많은 성능이 필요하다면 어떨까?
Python을 아예 버리고 저수준 언어로 다시 작성하라.
C로 구현하면 대략 이런 모습이 될 수 있다(설명용일 뿐이며, 안전하지도 않고 컴파일되는지도 확인하지 않았다):
static PyObject *_foo(PyObject *self, PyObject *args) { npy_intp npts = 0; npy_intp i = 0;
// 입력과 출력을 위한 Py_objects PyObject *a_obj = NULL, *b_obj = NULL, *c_obj = NULL, *d_obj = NULL; PyObject *out_obj = NULL; // PyArray 객체들 PyArrayObject *a_arr = NULL, *b_arr = NULL, *c_arr = NULL, *d_arr = NULL; PyArrayObject *out_arr = NULL; // 데이터 포인터들 double *a_ptr, *b_ptr, *c_ptr, *d_ptr, *out_ptr;
// 네 개의 입력 인수 파싱 if (!PyArg_ParseTuple(args, "OOOO", &a_obj, &b_obj, &c_obj, &d_obj)) { PyErr_SetString(PyExc_TypeError, "입력 파싱 오류"); return NULL; }
// 배열 객체 생성 a_arr = (PyArrayObject *)PyArray_FROM_O(a_obj); b_arr = (PyArrayObject *)PyArray_FROM_O(b_obj); c_arr = (PyArrayObject *)PyArray_FROM_O(c_obj); d_arr = (PyArrayObject *)PyArray_FROM_O(d_obj);
// 차원이 맞는지 확인 npts = PyArray_DIM(a_arr, 0);
// 출력 배열 생성 out_obj = PyArray_ZEROS(1, &npts, NPY_DOUBLE, 0); out_arr = (PyArrayObject *)out_obj;
Py_BEGIN_ALLOW_THREADS for (i = 0; i < npts; i++) { a_ptr = PyArray_GETPTR1(a_arr, i); b_ptr = PyArray_GETPTR1(b_arr, i); c_ptr = PyArray_GETPTR1(c_arr, i); d_ptr = PyArray_GETPTR1(d_arr, i); out_ptr = PyArray_GETPTR1(out_arr, i); *out_ptr = pow(((*a_ptr * *b_ptr) / *c_ptr), *d_ptr); } Py_END_ALLOW_THREADS
// 입력 배열 해제 if (a_arr) { Py_DECREF(a_arr); } if (b_arr) { Py_DECREF(b_arr); } if (c_arr) { Py_DECREF(c_arr); } if (d_arr) { Py_DECREF(d_arr); }
return out_obj;
error: if (a_arr) { Py_DECREF(a_arr); } if (b_arr) { Py_DECREF(b_arr); } if (c_arr) { Py_DECREF(c_arr); } if (d_arr) { Py_DECREF(d_arr); } if (out_obj) { Py_DECREF(out_obj); } if (out_arr) { Py_DECREF(out_arr); } PyErr_SetString(PyExc_RuntimeError, "C 오류"); return NULL; }
C에 익숙하지 않다면 얼핏 봐서는 알기 어렵지만, 이것은 우리가 앞서 시도한 것과 같은 ‘멍청한 루프’ 접근이다. 다만 이번에는 제대로 작동한다!
수동 메모리 관리를 피하기 위해 같은 연산 사이사이에 수많은 할당과 검사를 끼워 넣는 Python과 달리, C는 모든 것을 수동으로 처리한다. 그 대가는 훨씬 더 많은 보일러플레이트와, 일이 잘못될 수 있는 더 많은 기회다. 뭔가 해제를 빼먹고 메모리 누수를 내기 정말 쉽다.
실제로라면 나는 그냥 Rust를 쓸 것이다. Rust는 약간의 unsafe가 필요하긴 하지만 훨씬 적은 보일러플레이트로 더 많은 안전성을 준다.
#[pyfunction] pub fn foo<'py>( py: Python<'py>, // numpy 바인딩을 통해 배열 객체 받기 a_arr: PyReadonlyArray1<f64>, b_arr: PyReadonlyArray1<f64>, c_arr: PyReadonlyArray1<f64>, d_arr: PyReadonlyArray1<f64>, ) -> Bound<'py, PyArray1<f64>> {
// 배열을 읽기 전용 슬라이스로 받기 let a_slice = a_arr.as_slice().unwrap(); let b_slice = b_arr.as_slice().unwrap(); let c_slice = c_arr.as_slice().unwrap(); let d_slice = d_arr.as_slice().unwrap();
// 올바른 길이의 새 배열을 만들고 가변 슬라이스로 받기 let out_arr = unsafe { PyArray1::new(py, a_slice.len(), false)}; let out_slice = unsafe {out_arr.as_slice_mut().unwrap()};
// 우리의 멍청한 루프는 원소별로 순회하면서 // 중간 할당 없이 수학 연산을 하나의 단계로 융합한다 for (i, (((a, b), c), d)) in a_slice.iter() .zip(b_slice) .zip(c_slice) .zip(d_slice) .enumerate() {
// 계산을 수행하고 결과를 out 슬라이스의 해당 인덱스에 직접 기록 out_slice[i] = ((a*b)/c).powf(*d); } // out 배열을 Python으로 반환 out_arr }
내 기준에서는 이것이 C 코드보다 훨씬 더 읽기 쉽고, 의도도 아주 잘 드러난다. 모든 배열의 읽기 전용 슬라이스를 받고, 길이가 같은 새 배열 하나(정확히 1개)를 Python 힙에 할당하고, 그 새 배열의 가변 슬라이스를 받은 다음, 돌처럼 단순한 무할당 루프를 돌린다. 실제 계산은 전부 스택에서 수행되고, 그 위에 SIMD도 조금 얹힌다!
하지만 내 말만 믿지는 말고, 어셈블리를 읽어 보자! 아래는 같은 로직을 순수 Rust로 단순화한 버전이다. Compiler Explorer가 불평하지 않도록 하기 위해서다.
pub fn foo( a_arr: &[f64], b_arr: &[f64], c_arr: &[f64], d_arr: &[f64], ) -> Vec<f64> {
// 여기서는 가변 슬라이스를 받을 필요가 없다는 점에 주목하자. // 이 배열들은 모두 Rust가 제어하는 메모리에 있기 때문이다. let mut out = vec![0.0; a_arr.len()];
for (i, (((a, b), c), d)) in a_arr.iter() .zip(b_arr) .zip(c_arr) .zip(d_arr) .enumerate() {
out[i] = ((a*b)/c).powf(*d); } out }
개발 빌드에서는 이 코드가 약 1700줄의 x86 어셈블리로 컴파일된다. opt-level-3를 주면 무려 150줄까지 내려간다! Rust 1.93.0 기준이다.
벡터화된 루프는 더 짧다. 전체가 이렇다:
.LBB0_9:
movupd xmm1, xmmword ptr [r13 + r15]
movupd xmm0, xmmword ptr [rbx + r15]
mulpd xmm0, xmm1
movupd xmm1, xmmword ptr [rcx + r15]
divpd xmm0, xmm1
movapd xmmword ptr [rsp + 80], xmm0
movups xmm1, xmmword ptr [rdx + r15]
movaps xmmword ptr [rsp + 96], xmm1
mov r12, rbx
mov rbx, rdx
call rbp
movapd xmmword ptr [rsp + 32], xmm0
movapd xmm0, xmmword ptr [rsp + 80]
unpckhpd xmm0, xmm0
movaps xmm1, xmmword ptr [rsp + 96]
movhlps xmm1, xmm1
call rbp
mov rdx, rbx
mov rbx, r12
mov rcx, qword ptr [rsp + 24]
mov rax, qword ptr [rsp]
movapd xmm1, xmmword ptr [rsp + 32]
unpcklpd xmm1, xmm0
movupd xmmword ptr [rax + r15], xmm1
add r15, 16
cmp r14, r15
jne .LBB0_9
mov rsi, qword ptr [rsp + 56]
mov rdi, qword ptr [rsp + 16]
이걸 뜯어보자.
여기서는 movupd (mov e u naligned p acked d oubles)를 사용해 배열 A와 B에서 128비트(64비트 float 2개)를 로드한다. xmmword는 128비트라는 점에 주목하자.
movupd xmm1, xmmword ptr [r13 + r15] ; A에서 double 2개 로드
movupd xmm0, xmmword ptr [rbx + r15] ; B에서 double 2개 로드
mulpd xmm0, xmm1 ; xmm0 = a * b (packed)
movupd xmm1, xmmword ptr [rcx + r15] ; C에서 double 2개 로드
divpd xmm0, xmm1 ; xmm0 /= c (packed)
그 다음 mulpd (mul tiply p acked d oubles)를 이용해 위아래 double을 동시에 곱해 결과를 xmm0에 저장하고, 배열 C에서 packed double 2개를 xmm1로 더 불러온 뒤, divpd (div ide p acked d oubles)로 같은 방식의 나눗셈을 수행한다.
우리는 (a*b)/c 단계를 단 다섯 개의 어셈블리 명령으로, 그것도 두 원소에 대해 처리하고 있다! 이건 Python에서 똑같은 멍청한 루프를 시도했던 것보다 당연히 훨씬 낫다. 매번 그 많은 암묵적 할당을 하지 않기 때문이다. 하지만 같은 하드웨어에서 NumPy의 벡터화된 배열 연산보다도 더 낫다. NumPy도 SIMD를 활용하고 있긴 하지만 eager 평가를 하기 때문에 temp /= c에 도달하기 전에 이미 힙 할당을 한 번 통째로 수행해 버린다.
temp ** d도 같은 단계에서 수행할 수는 있지만, 이 하드웨어에서는 벡터화되지 않는다.
훨씬 앞에서 우리는 Rust의 거듭제곱 함수를 rbp 레지스터에 저장해 두었다.
mov rbp, qword ptr [rip + pow@GOTPCREL]
그 다음 이전 연산의 결과를 함수 지역 스택에 저장하고(곧 이유가 나온다), temp**d를 준비하기 위해 D에서 packed 값 두 개를 같은 스택에 로드한다. 여기 일부는 doubles 대신 singles를 쓰고 있지만 별 영향은 없어 보인다. 여기서는 동등해야 한다.
movapd xmmword ptr [rsp + 80], xmm0 ; xmm0를 지역 스택에 저장 (16 bytes)
movups xmm1, xmmword ptr [rdx + r15] ; D에서 double 2개 로드
movaps xmmword ptr [rsp + 96], xmm1 ; ... 그리고 지역 스택에 저장, x86은 mem-to-mem 이동을 허용하지 않기 때문
mov r12, rbx ; callee-managed 레지스터에 보관하기 위한 약간의 정리
mov rbx, rdx
이제 자리를 비웠으니 실제로 거듭제곱 연산을 수행할 수 있다. 하지만 한 번에 스칼라 하나씩만 가능하다.
call rbp는 64비트 거듭제곱 연산을 수행한다. 하지만 128비트 xmm 레지스터에서 작동하므로, 각 레지스터의 첫 64비트를 넘어가는 부분은 무시한다. 작업이 끝나면 xmm0의 하위 비트에 저장된 결과를 꺼내 스택에 저장하고, 스택에 저장해 두었던 이전 xmm0 값을 다시 불러온 다음, _unpckhpd_를 사용해 상위 double을 하위 위치로 옮긴다. xmm1에 저장된 지수에도 똑같이 한 뒤 call rbp를 한 번 더 호출한다.
call rbp ; xmm0[0] ** xmm1[0]
movapd xmmword ptr [rsp + 32], xmm0 ; result[0]를 스택에 저장
movapd xmm0, xmmword ptr [rsp + 80] ; 이전 temp를 스택에서 다시 가져옴
unpckhpd xmm0, xmm0 ; 상위 비트를 하위 위치로 이동
movaps xmm1, xmmword ptr [rsp + 96] ; xmm0에 대해 한 것과 같은 작업
movhlps xmm1, xmm1
call rbp ; xmm0[1] ** xmm1[1]
여기서부터는 결과를 적절한 위치에 저장하고, 레지스터를 되돌리고, 루프를 계속하는 작업들이다.
그리고 다시 처음으로 돌아가, 원래의 numpy 벡터화 계산에서 생성된 바이트코드를 다시 보자.
def foo(): a = np.array([...]) b = np.array([...]) c = np.array([...]) d = np.array([...])
return ((a*b)/c)**d
12 178 LOAD_FAST 0 (a)
180 LOAD_FAST 1 (b)
182 BINARY_OP 5 (*)
186 LOAD_FAST 2 (c)
188 BINARY_OP 11 (/)
192 LOAD_FAST 3 (d)
194 BINARY_OP 8 (**)
198 RETURN_VALUE
불러오고, 불러오고, 연산하고, 불러오고, 연산하고, 불러오고, 연산하고, 반환한다.
개념적으로는 어셈블리와 그렇게까지 다르지 않다! 하지만 어셈블리는 하드웨어에 대한 명령을 반영하며, 각 명령은 나노초의 일부만에 실행되고 전부 스택 위에서 이루어진다. 반면 바이트코드는 하드웨어 위에 여러 층의 추상화가 쌓인 곳에 존재하며, 각 명령은 여러 단계의 포인터 간접 참조, 힙 할당, 타입 검사를 숨기고 있다.
NumPy를 쓰든 직접 어셈블리를 쓰든, 언어는 당신과 하드웨어 사이의 인터페이스이며, 당신이 ‘코드가 무엇을 하길 원하는지’를 더 정밀하게 명시할수록 더 빠르게 실행될 수 있다.
그렇다면 실제로는 얼마나 빠를까?
대충 한 빠른 벤치마크를 돌려 보면, 기본 NumPy 구현은 원소당 대략 ~7.3 나노초 정도가 걸린다.
리스트 컴프리헨션으로 같은 일을 하면 원소당 130-150 나노초 정도로, 약 20배 느리다.
Python의 ‘멍청한 루프’는 원소당 약 170ns 정도로, 23배 느리다. 별로다!
그리고 Rust 버전은… 원소당 ~7.7ns 정도로, 약 5% 느리다! 내가 예상했던 결과는 아니다!
왜 이런 걸까? 이유는 몇 가지 있다.
첫째, 다른 연산으로 실험해 보니, 거듭제곱 없이 (a*b)/c만 수행하면 상황이 달라진다.
NumPy: 0.847 ns/element
Rust: 0.575 / element
이제 Rust 구현이 47% 더 빠르다. 하지만 더 중요한 점은, 둘 다 거듭제곱 단계를 제거하면 약 8-15배 정도 빨라진다는 것이다. 알고 보니 pow가 함수 전체 실행 시간의 80-95+ %를 차지하고 있었다!
이건 두 번째 이유, 즉 실제 하드웨어로 이어진다.
앞서 한 어셈블리 분석은 Compiler Explorer가 제공하는 x86 하드웨어 기준이었지만, 나는 Apple Silicon에서 실행하고 있고, 여기서 NumPy는 많은 플랫폼별 SIMD 최적화를 활용할 수 있다. np.show_config()를 실행하면 다음이 나온다.
"SIMD Extensions": {
"baseline": [
"NEON",
"NEON_FP16",
"NEON_VFPV4",
"ASIMD"
],
"found": [
"ASIMDHP",
"ASIMDDP"
],
"not found": [
"ASIMDFHM"
]
}
ASIMD는 Apple의 Accelerate SIMD 프레임워크이며, 여러 기능 중 하나로 벡터화된 거듭제곱 함수를 제공한다.
하드웨어와의 이런 우수한 통합, 특히 거듭제곱 연산이 함수의 거의 모든 것을 지배할 때, 중간 메모리 할당이 있음에도 불구하고 단순한 NumPy 구현이 더 빠르게 된다.
(후속 글 보기)
교훈: 최적화에 대해 아무리 많이 생각해도, 서로 다른 해법을 실제로 벤치마크해야 한다. 이상적으로는 목표 하드웨어에서 해야 하지만, 많은 경우 그게 쉽지는 않을 수 있다.
이 사소한 예제에서 이 점을 확인했으니, 이제 McFACTS에서 실제로 최적화 대상이 된 함수를 한번 보자.
def analytical_kick_velocity( mass_1, mass_2, spin_1, spin_2, spin_angle_1, spin_angle_2): """ Compute the analytical gravitational wave recoil (kick) velocity for merging black hole binaries as in Akiba et al. 2024 (arXiv:2410.19881).
mass_1 : numpy.ndarray
Mass [M_sun] of object 1 with :obj:float type
mass_2 : numpy.ndarray
Mass [M_sun] of object 2 with :obj:float type
spin_1 : numpy.ndarray
Spin magnitude [unitless] of object 1 with :obj:float type
spin_2 : numpy.ndarray
Spin magnitude [unitless] of object 2 with :obj:float type
spin_angle_1 : numpy.ndarray
Spin angle [radian] of object 1 with :obj:float type
spin_angle_2 : numpy.ndarray
Spin angle [radian] of object 2 with :obj:float type
v_kick : np.ndarray
Kick velocity [km/s] of the remnant BH with :obj:float type
"""
# Akiba et al 2024 부록 A에 따라, mass_2는 쌍성계에서 더 무거운 BH여야 한다.
mask = mass_1 <= mass_2
m_1_new = np.where(mask, mass_1, mass_2) * u.solMass
m_2_new = np.where(mask, mass_2, mass_1)* u.solMass
spin_1_new = np.where(mask, spin_1, spin_2)
spin_2_new = np.where(mask, spin_2, spin_1)
spin_angle_1_new = np.where(mask, spin_angle_1, spin_angle_2)
spin_angle_2_new = np.where(mask, spin_angle_2, spin_angle_1)
# "perp"와 "par"는 각각 궤도 각운동량 축에 수직인 성분과 평행한 성분을 뜻한다.
# 쌍성계의 궤도 각운동량 축은 원반의 각운동량과 정렬되어 있다.
# 스핀의 perp 및 par 성분을 구한다:
spin_1_par = spin_1_new * np.cos(spin_angle_1_new)
spin_1_perp = spin_1_new * np.sin(spin_angle_1_new)
spin_2_par = spin_2_new * np.cos(spin_angle_2_new)
spin_2_perp = spin_2_new * np.sin(spin_angle_2_new)
# 질량비 q와 비대칭 질량비 eta를 구한다
# 정의는 Akiba et al. 2024 부록 A를 따른다:
q = m_1_new / m_2_new
eta = q / (1 + q)**2
# Akiba et al. 2024 식 A5 사용:
S = (2 * (spin_1_new + q**2 * spin_2_new)) / (1 + q)**2
# 정의는 Akiba et al. 2024 부록 A를 따른다:
xi = np.radians(145)
A = 1.2e4 * u.km / u.s
B = -0.93
H = 6.9e3 * u.km / u.s
V_11, V_A, V_B, V_C = 3678 * u.km / u.s, 2481 * u.km / u.s, 1793* u.km / u.s, 1507 * u.km / u.s
angle = rng.uniform(0.0, 2*np.pi, size=len(mass_1))
# Akiba et al. 2024 식 A2 사용:
v_m = A * eta**2 * np.sqrt(1 - 4 * eta) * (1 + B * eta)
# Akiba et al. 2024 식 A3 사용:
v_perp = (H * eta**2 / (1 + q)) * (spin_2_par - q * spin_1_par)
# Akiba et al. 2024 식 A4 사용:
v_par = ((16 * eta**2) / (1 + q)) * (V_11 + (V_A * S) + (V_B * S**2) + (V_C * S**3)) * \
np.abs(spin_2_perp - q * spin_1_perp) * np.cos(angle)
# Akiba et al. 2024 식 A1 사용:
v_kick = np.sqrt((v_m + v_perp * np.cos(xi))**2 +
(v_perp * np.sin(xi))**2 +
v_par**2)
v_kick = np.array(v_kick.value)
assert np.all(v_kick > 0), \
"v_kick에 <= 0 값이 있음"
assert np.isfinite(v_kick).all(), \
"유한성 검사 실패: v_kick"
return v_kick
이런 종류의 함수가 바로 우리가 앞서 이야기한 문제를 갖는다. 벡터화된 배열 연산이 잔뜩 있고, 모두 eager 평가되며 엄청난 수의 중간 배열을 할당한다.
우리는 rng 호출이 필요한 부분을 함수에서 떼어낼 수 있다.
def analytical_kick_velocity_optimized( ... ): angle = rng.uniform(0.0, 2*np.pi, size=len(mass_1))
v_kick = analytical_kick_velocity_helper(
mass_1,
mass_2,
spin_1,
spin_2,
spin_angle_1,
spin_angle_2,
angle
)
return v_kick
그리고 나머지는 Rust 쪽 헬퍼 함수에 넘긴다.
#[pyfunction] pub fn analytical_kick_velocity_helper<'py>( py: Python<'py>, mass1_arr: PyReadonlyArray1<f64>, mass2_arr: PyReadonlyArray1<f64>, spin1_arr: PyReadonlyArray1<f64>, spin2_arr: PyReadonlyArray1<f64>, spin_angle1_arr: PyReadonlyArray1<f64>, spin_angle2_arr: PyReadonlyArray1<f64>, angle_arr: PyReadonlyArray1<f64>, ) -> FloatArray1<'py> {
let m1_slice = mass1_arr.as_slice().unwrap(); let m2_slice = mass2_arr.as_slice().unwrap(); let s1_slice = spin1_arr.as_slice().unwrap(); let s2_slice = spin2_arr.as_slice().unwrap(); let sa1_slice = spin_angle1_arr.as_slice().unwrap(); let sa2_slice = spin_angle2_arr.as_slice().unwrap(); let angle_slice = angle_arr.as_slice().unwrap();
let out_arr = unsafe { PyArray1::new(py, m1_slice.len(), false)}; let out_slice = unsafe { out_arr.as_slice_mut().unwrap()};
let xi: f64 = 145.0f64.to_radians(); let const_a: f64 = 1.2e4; let const_b: f64 = -0.93; let const_h: f64 = 6.9e3; let v_11: f64 = 3678.0; let v_a: f64 = 2481.0; let v_b: f64 = 1793.0; let v_c: f64 = 1507.0;
for (i, ((((((m1, m2), s1), s2), sa1), sa2), angle)) in m1_slice.iter() .zip(m2_slice) .zip(s1_slice) .zip(s2_slice) .zip(sa1_slice) .zip(sa2_slice) .zip(angle_slice) .enumerate() {
// 교체 처리 (Akiba et al. 부록 A: mass_2가 더 무거워야 함) // 단순한 변수 섀도잉을 사용해 순수하게 스택 위에서 교체한다 let (loc_m1, loc_m2, loc_s1, loc_s2, loc_a1, loc_a2) = if m1 <= m2 { (m1, m2, s1, s2, sa1, sa2) } else { (m2, m1, s2, s1, sa2, sa1) };
// 스핀 성분 let (s1_sin, s1_cos) = loc_a1.sin_cos(); let (s2_sin, s2_cos) = loc_a2.sin_cos();
let s1_par = loc_s1 * s1_cos; let s1_perp = loc_s1 * s1_sin; let s2_par = loc_s2 * s2_cos; let s2_perp = loc_s2 * s2_sin;
// 질량비 let q = loc_m1 / loc_m2; let q_sq = q * q; let one_plus_q = 1.0 + q; let one_plus_q_sq = one_plus_q * one_plus_q;
let eta = q / one_plus_q_sq; let eta_sq = eta * eta;
// Akiba 식 A5 let s_big = (2.0 * (loc_s1 + q_sq * loc_s2)) / one_plus_q_sq; let s_big_sq = s_big * s_big; let s_big_cu = s_big_sq * s_big;
// Akiba 식 A2 (v_m) let term_sqrt = (1.0 - 4.0 * eta).sqrt(); let v_m = const_a * eta_sq * term_sqrt * (1.0 + const_b * eta);
// Akiba 식 A3 (v_perp) let v_perp_mag = (const_h * eta_sq / one_plus_q) * (s2_par - q * s1_par);
// Akiba 식 A4 (v_par) // 참고: Python 코드는 np.abs(spin_2_perp - q * spin_1_perp)를 사용했다 let term_v = v_11 + (v_a * s_big) + (v_b * s_big_sq) + (v_c * s_big_cu); let spin_diff_perp = (s2_perp - q * s1_perp).abs();
let v_par = ((16.0 * eta_sq) / one_plus_q) * term_v * spin_diff_perp * angle.cos();
// Akiba 식 A1 (총 킥) let (xi_sin, xi_cos) = xi.sin_cos();
let term_1 = v_m + v_perp_mag * xi_cos; let term_2 = v_perp_mag * xi_sin;
out_slice[i] = (term_1.powi(2) + term_2.powi(2) + v_par.powi(2)).sqrt(); }
out_arr }
실제 성능은 어떨까?
순수 Python 버전은 벤치마크에서 약 912 마이크로초가 나오고, Rust 버전은 4.3 마이크로초 정도다. 무려 210배 빠르다.
게다가 Rust로 가속된 함수의 절반 정도, 즉 2.1 마이크로초만이 실제 Rust 헬퍼 함수 안에서 쓰이고, 나머지는 난수 생성에 쓰인다. 같은 조건끼리 비교하면, NumPy 배열을 사용하는 것과 비교해 사실상 430배 정도 빨라진 셈이다.
일반적으로, eager 평가되는 단계들을 많이 융합해 중간 할당을 제거할수록 성능 격차는 더 커지고, 금세 NumPy의 우수한 SIMD 통합이 주는 이점을 압도하게 된다. 물론 원한다면 우리도 하드웨어별 SIMD 최적화를 직접 구현해서 더더욱 빠르게 만들 수 있겠지만, McFACTS에서는 현재 여기서 멈춘다. 912 마이크로초에서 4 마이크로초로 내려가면서, 이 함수는 전체 시뮬레이션 실행 시간의 약 3%를 차지하던 것에서 사실상 보이지 않는 수준이 되었다.
Powered by mataroa.blog.