Повторение массивов: почему np.repeat быстрее итераторов в Rust

Почему np.repeat в NumPy быстрее наивной реализации в Rust: роль memcpy и SIMD, преимущество непрерывных копий над индексацией при повторении массивов.

Реализуя пуассоновский бутстрэп в Rust, легко поддаться желанию сравнить ключевые строительные блоки с NumPy. Один из них — функция повторения: получив массив и вес для каждого элемента, нужно сформировать новый массив, где каждое значение воспроизводится столько раз, сколько указано его весом. Наивная реализация на Rust с итераторами может неожиданно работать медленнее, чем np.repeat из NumPy, и, что еще удивительнее, np.repeat нередко опережает, казалось бы, равносильные приемы NumPy вроде срезов или np.take.

Постановка задачи

Операция, к которой мы стремимся, выглядит так: значения [1, 2, 3] с весами [1, 2, 0] должны дать [1, 2, 2]. Идиоматичное решение на Rust обычно опирается на итераторы и repeat_n. Ниже — прямой перенос этой идеи.

pub fn inflate(values: &[f64], times: &[u64]) -> Vec<f64> {
    let out: Vec<f64> = values
        .iter()
        .zip(times.iter())
        .flat_map(|(&v, &k)| std::iter::repeat_n(v, k as usize))
        .collect();

    out
}

Со стороны NumPy быстрый бенчмарк показывает масштаб различий. Построить список индексов и вызвать np.take логически эквивалентно повторению, однако в этом сценарии np.repeat примерно в три раза быстрее срезов/np.take.

import timeit
import numpy as np

N = 100_000
src = np.random.rand(N)

wts = np.random.poisson(1, len(src))

# Precompute indices so Python looping doesn’t contaminate the timing
idx = []
for w in wts:
    for i in range(w):
        idx.append(i)

print(timeit.timeit(lambda: np.repeat(src, wts), number=1_000))
print(timeit.timeit(lambda: np.take(src, idx), number=1_000))

Откуда на самом деле разрыв

Короткий ответ: NumPy использует широкие загрузки и записи, а наивная версия на Rust — нет. Итераторный код на Rust сводится к скалярным перемещениям, то есть в «горячих» циклах обрабатывает по одному 64-битному значению за раз. Здесь не задействован memcpy и нет циклового использования SIMD-регистров. Уже этого достаточно, чтобы ограничить пропускную способность.

NumPy же, когда возможно, применяет memcpy во вложенных циклах. Для базовых непрерывных нативных массивов этот путь подходит. Реализации memcpy агрессивно оптимизированы и могут задействовать SIMD под капотом на тех платформах, где это окупается. Для достаточно длинных участков данных это означает куда более широкие передачи на инструкцию и заметно лучшую устойчивую пропускную способность, чем у скалярных копирований. На современных CPU с широкими векторными блоками такая разница легко объясняет ускорение в разы.

Почему же np.repeat настолько быстрее срезов или np.take? Индексация в этом контексте плохо сочетается с SIMD. В случае np.take реализации приходится читать и сами данные, и индексы; шаблон доступа перестает быть непрерывной копией. Хотя на некоторых x86-64 есть инструкции gather, на практике они уступают пакетным загрузкам/записям, особенно для 64-битных элементов. И даже если gather доступен, дополнительный трафик памяти для индексов не позволяет обогнать прямые непрерывные копии. Проще говоря, операции на индексах хуже дружат с SIMD, чем массовые копирования, и это отчетливо видно в замерах.

Что делать вместо этого

Если удается сохранить операцию в виде последовательности непрерывных копий, так и поступайте. Именно это делает np.repeat быстрым: когда можно, он вырождается в массовое копирование. В Rust стоит стремиться к подходу, который повторяет такое поведение, а не перебирать скаляры; чем ближе вы к передачам в духе memcpy по непрерывным участкам, тем больше выигрываете от широких загрузок/записей. И обязательно собирайте и замеряйте оптимизированные сборки: для честного сравнения производительности имеет значение режим release.

Почему это важно

В задачах обработки данных и численных вычислений часто сталкиваются два на вид равнозначных приема: индексация источника и копирование непрерывных фрагментов. Первый приводит к квазислучайным шаблонам доступа и дополнительным чтениям индексов — процессор не может полноценно загрузить векторные блоки. Второй позволяет задействовать хорошо настроенные реализации memcpy, которые эффективно используют SIMD и подсистему памяти. Понимание этой разницы помогает не гоняться за микрооптимизациями не там, где нужно, и выбирать формы алгоритмов, которые «нравятся» железу.

Выводы

Преимущество NumPy в данном случае дает использование memcpy, то есть широких векторизованных загрузок и записей. Наивный итераторный подход в Rust выполняет скалярные перемещения и расплачивается пропускной способностью. Стратегии на индексах — такие как срезы или np.take — медленнее, потому что не могут задействовать тот же путь непрерывного копирования и вынуждены обрабатывать индексы вместе с данными. Когда есть возможность, выбирайте решения, которые позволяют рантайму копировать непрерывную память крупными блоками, избегайте лишней индексации и всегда измеряйте в оптимизированных сборках.