[Rust] iterator 성능 벤치

iterator는 Rust가 저렴한 추상화라고 자랑하는 주요 슬로건 중 하나다. zero cost abstractions라고 부른다.
근데 저비용으로 최적화가 된다고는 주장하는데, 그 방식과 원리에 대해서는 딱히 뚜렷한 설명을 해주지는 않더라.

그러면 진짜로 iterator가 for 루프와 대등한 속도를 내는걸까? 궁금해서 실험부터 해봤다.

pub struct Timer {
    start: std::time::Instant,
}

impl Timer {
    pub fn new() -> Timer {
        Timer {
            start: std::time::Instant::now(),
        }
    }

    pub fn elapsed(&self) -> std::time::Duration {
        self.start.elapsed()
    }

    pub fn elapsed_as_millis(&self) -> u128 {
        self.elapsed().as_millis()
    }

    pub fn elapsed_as_secs(&self) -> u64 {
        self.elapsed().as_secs()
    }
}

fn do_iterator_test(input: Vec<i32>) -> Vec<f64> {
    input.iter().map(|&x| x as f64).collect()
}

fn do_loop_test(input: Vec<i32>) -> Vec<f64> {
    let mut result = Vec::new();
    result.reserve(input.len());
    for i in 0..input.len() {
        result.push(input[i] as f64);
    }
    result
}

fn main() {
    const N: usize = 1_000_000_000;

    let mut loop_test = vec![1; N];
    let mut iter_test = vec![1; N];
    for i in 0..N {
        loop_test[i] = i as i32;
        iter_test[i] = i as i32;
    }

    let timer = Timer::new();
    let sum_1 = do_loop_test(loop_test);
    println!("loop: {}ms", timer.elapsed_as_millis());

    let timer = Timer::new();
    let sum_2 = do_iterator_test(iter_test);
    println!("iter: {}ms", timer.elapsed_as_millis());

    assert_eq!(sum_1, sum_2);
}

그냥 i32 벡터를 f64 벡터로 치환하는 간단한 코드를 loop 버전과 iterator 버전으로 각각 작성해봤다.
나는 성능이 비슷하게 나오거나 일반 루프가 근소하게 빠른 결과를 예상했다.

첫번째 테스트환경은 Ubuntu 22, x86_64, RAM 32GB였다.

하지만 의외로 iterator의 처리속도가 훨씬 빨랐다.

release로 빌드해도 그렇고, 실행 순서를 반대로 바꿔도 다름이 없었다. 캐시의 영향도 아니라는 것이다.

fn do_loop_test(input: Vec<i32>) -> Vec<f64> {
    let mut result = vec![0_f64; input.len()];
    let mut i = 0;
    while i < input.len() {
        result[i] = input[i] as f64;
        i += 1;
    }
    result
}

vector push를 제거하고 미리 할당하는 방식으로 변경하니 좀 대등해졌으나, 그래도 항상 iterator가 미세하게 빨랐다.

그럼 다른 플랫폼에서는 어떨까?
Mac 환경에서도 테스트해봤다.
스펙은 M2 Pro Ventura 13, RAM 16GB였다.

Mac에서의 벤치마크 결과는 좀 더 혼란스러웠다.
iter 버전을 먼저 실행하면 iter가 더 빠르고, loop 버전을 먼저 실행하면 loop가 더 빨랐다. 근데 아무튼 엄청 큰 차이까지 벌어지지는 않았다.

LLVM-IR 까보기

어셈블리 수준까지 가기에는 너무 고통스럽고, llvm이 자기들 마음대로 최적화를 가하는 것도 있기 때문에 Rust가 최적화를 하는 부분을 보고 싶다면 llvm-ir을 까보는 것이 가장 명확한 방법이다.

우선 llvm-ir을 뽑아봤다.

그러면 이렇게 주석이 친절하게 달려있기 때문에 보는 것이 그렇게 어렵지만은 않다.

루프 버전 함수의 경우에는 다음과 같이 고전적인 형태로 루프 구문을 구성한다.

i 변수를 0으로 깔고

반복을 위한 length 값을 할당하고

뺑뺑돌면서 작을 때만 bb5 레이블로 이동한다.

좀 장황하긴 한데, result[i] = input[i] as f64가 변환된 결과라고 보면 된다.
오버플로 체크도 하고, 바운드 체크해서 패닉 핸들링을 하는 것도 있고, 뭐가 인스트럭션이 많은게 그냥 봐도 느껴질 것이다.

그에 반해 iterator로 생성된 코드는

비교적 간단하다.
바운드 체크도 없고?, iteration core 함수를 2개 호출하는 코드로 생성이 된다.
iterator 코드의 경우에는 반복의 형태가 정형화된 채로 고정되어있다보니 좀더 간결하고 과격한 형태로 최적화가 가능한 것 같다.

물론 인스트럭션이 많더라도 llvm이 알아서 최적화를 해줄 수도 있지만, 이 정도로 차이가 난다면 결과에도 유의미한 차이가 발생할 수 있을 것이다. 작은 변수 하나만 있어도 무언가 나쁜 영향을 줄 수 있다면 최적화 수단을 하나 포기할 수 있는 것이기 때문이다.

filter추가해보기

위에서 테스트했던 코드는 타입변환 정도만 하는 아주 간단한 코드였다. 그러면 체이닝을 좀더 더해보면 어떨까?

map 후에 filter만 한번 더해보는 형태로 코드를 변경해봤다.

이 경우에는 iteration이 약간 느린 정도로 최적화가 되었다.
테스트 환경은 Mac M2다.

하지만 이 버전의 iterator 코드에는 사실 조금 비효율적인 부분이 있다.
filter를 먼저 한 뒤에 map을 해도 되는데 그 반대로 하고 있기 때문이다. 그러면 filter로 걸러질 대상에 대해서도 불필요한 map 반복을 돌 것이다.

이걸 다시 순서만 조정해줘도

조금 더 동등한 수준까지 최적화가 되는 것 같다. 물론 왔다갔다 하긴 하는데, iterator가 빠르게 나올 때도 있다.

아무튼 미묘하긴 하지만 Rust는 map과 filter로 구성된 복합 iteration의 경우에도 어느정도 for 루프와 동일한 형태로 최적화를 해준다. 그래서 for 버전의 코드와 완벽하게 동일한 바이너리를 방출할 때도 많다.
이런걸 Haskell에서는 fusion이라고 부르더라.

참조
https://doc.rust-lang.org/book/ch13-04-performance.html
https://ntietz.com/blog/rusts-iterators-optimize-footgun/