Векторизация

Фотография «Вода и горы» — 600 на 305 пикселей, 183 тысячи точек. Перевести её в серый — по нескольку умножений и сложений на пиксель. Цикл Python делает это за 92 миллисекунды, numpy — меньше чем за миллисекунду, хотя арифметики столько же. Ниже четыре операции над этим снимком, и для каждой — гонка способов на настоящих замерах: цикл Python, numpy, цикл Go и ваш браузер. Выберите операцию и размер: одна фотография, снимок телефона или альбом из ста снимков.

Серый, гистограмма и размытие в numpy догоняют цикл Go, а дизеринг — нет: для него в numpy нет готовой операции, и цикл Python по массиву numpy оказывается втрое медленнее цикла по обычному списку. Почему так, разберём по порядку.

Идея

В цикле Python каждое число — отдельный объект где-то в памяти, и каждая итерация — десятки шагов интерпретатора: прочитать байт-код, узнать тип, найти сложение для этого типа, создать объект-результат. Сама арифметика — малая доля этой работы. Векторизация переносит цикл внутрь скомпилированной библиотеки: программа говорит «сложи эти два массива», и цикл по миллиону элементов крутится в машинном коде над плотным массивом чисел одного типа. Там процессор может применить и векторные команды, которые обрабатывают по 16–64 байта за раз.

Слово «векторизация» значит две разные вещи, и их полезно различать. Для программиста на Python, R или MATLAB это способ записи: целыми массивами вместо циклов. Для компилятора и процессора это способ исполнения: одна команда над несколькими числами. Первое почти всегда ведёт ко второму, но главный выигрыш даёт уже первое — уход от интерпретатора.

Куда уходит время цикла

Самая простая задача — посчитать пробелы в «Войне и мире». Это пять миллионов байт, и на каждый — одно сравнение.

Цикл Python тратит на байт около 40 наносекунд — на машине замера это больше сотни тактов на одно сравнение. Вызов bytes.count — тот же цикл, но на C, и он в двадцать раз быстрее. numpy сравнивает весь массив с пробелом векторными командами и считает единицы — ещё в семь раз быстрее: 0,3 наносекунды на байт, в 130 раз быстрее цикла Python. В Go обычный цикл — уже машинный код, 0,9 наносекунды на байт, а bytes.Count написан на ассемблере и сравнивает по 16 байт одной командой на arm64 и по 32 на amd64: ещё в 15 раз быстрее.

Все прямые на графике параллельны: время любого способа растёт как n. Векторизация не меняет асимптотику — она меняет постоянный множитель, иногда в сотни раз.

Код: что векторизуется, а что нет

Четыре операции из сцены — четыре разных случая. Листинг делает каждую дважды, циклом Python и целыми массивами numpy, и проверяет, что ответы совпадают до пикселя.

photo 600 x 305 = 183000 pixels, 549000 bytes of colour
brightness: mean 122.78, from 0 to 228
histogram: the most common brightness 211, 2384 pixels
3 x 3 blur of the inner pixels: sum 22272633
Floyd-Steinberg: 96337 black and 86663 white pixels

• Поэлементная операция — серый. Каждый пиксель независим: яркость = (299r + 587g + 114b + 500) div 1000. В numpy это та же формула, только над тремя столбцами массива сразу. Быстрее цикла в 56 раз.
• Подсчёт — гистограмма. Цикл «counts[v] += 1» пишет в ячейки вразброс, и записать его формулой над массивом нельзя. Но для этой операции в numpy есть готовая функция np.bincount: цикл тот же, только на C. Быстрее в 9 раз.
• Соседи — размытие 3 × 3. Новый пиксель — среднее девяти соседей. Если сдвинуть всю картинку на пиксель влево, вправо, вверх и вниз — девять срезов одного массива без копирования, — то сумма девяти срезов и есть сумма соседей для всех пикселей сразу. Быстрее цикла в 147 раз, и это лучший случай: работы в цикле много, а в numpy она уходит в девять сложений целых массивов.
• Зависимость — дизеринг Флойда — Стейнберга. Пиксель округляется до чёрного или белого, а ошибка округления уходит соседям справа и снизу. Значит, следующий пиксель зависит от того, как округлили предыдущий, и все пиксели сразу не посчитать. Операции «дизеринг» в numpy нет, и ничего не выигрывается.

Последний случай поучительнее всех. Если взять массив numpy и пройти его тем же циклом Python, станет не быстрее, а медленнее — в два с половиной раза и у дизеринга, и у гистограммы. Каждое обращение v[i] к массиву numpy создаёт новый объект-число numpy, а список Python хранит готовые объекты, и для чисел от −5 до 256 это вообще одни и те же заранее созданные объекты. numpy быстр только тогда, когда цикл внутри него, а не снаружи. Для таких операций нужен компилируемый язык — Go проходит дизеринг в 40 раз быстрее цикла Python — или компилятор для Python вроде Numba и Cython.

Векторные команды и сумма дробей

Сложить миллион дробей векторными командами — значит разложить их по полосам: восемь сумматоров, первый складывает числа 0, 8, 16, …, второй — 1, 9, 17, …, а в конце восемь частичных сумм складываются. Казалось бы, компилятор C или Go мог бы сделать так сам. Он этого не делает, и вот почему.

exact:            499774542.9491545
1 accumulator:    499774542.94915307  off by -1.43e-06
4 accumulators:   499774542.949157  off by +2.5e-06
8 accumulators:   499774542.94915396  off by -5.36e-07

Ответы разные. Сложение дробей не ассоциативно: (10¹⁶ + 1) − 10¹⁶ в double даёт 0, потому что единица теряется при округлении суммы, а (10¹⁶ − 10¹⁶) + 1 даёт 1. Порядок сложений меняет результат, и компилятор, который обязан выполнять программу как написано, переставлять их не вправе. Компиляторы C делают это только с разрешения — с флагами вроде -ffast-math, — а целые числа складывают по полосам свободно: там порядок ничего не меняет. numpy, наоборот, складывает дроби по-своему намеренно — попарно, кусками, — и это не только быстрее, но обычно и точнее простого цикла.

Компилятор Go векторных команд для таких циклов не порождает, но полосы помогают и без них.

С одним сумматором Go тратит около наносекунды на число: каждое сложение ждёт результата предыдущего, а сложение дробей в процессоре длится несколько тактов. С четырьмя сумматорами — 0,24, с восемью — около 0,2: четыре или восемь независимых цепочек процессор ведёт одновременно. Это не векторные команды, а параллелизм внутри одного ядра, но причина выигрыша та же — независимость. И та же ловушка: восемь сумматоров, сложенные в массив [8]float64, дают 0,53 — вдвое хуже восьми отдельных переменных. Компилятор Go держит такой массив в памяти, а не в регистрах, и каждое сложение превращается в чтение и запись. Такое видно только по замеру.

Python в этом сравнении проигрывает всем: цикл — 65 наносекунд на число, встроенная sum — 11, numpy — 0,5.

Сколько это стоит

Асимптотика не меняется: векторная операция над n числами — те же O(n) действий, что у цикла. Меняется цена одного действия, и у неё есть свои подводные камни.

• Вызов стоит. Каждый вызов numpy — разбор аргументов, проверка типов, выделение памяти под результат — несколько микросекунд. На графике сумм это видно слева: на тысяче чисел numpy тратит 6,5 наносекунды на число, всего вдвое меньше sum, а на миллионе — 0,5. Векторизовать сто операций над массивами из десяти чисел бессмысленно; объединить их в одну операцию над массивом 100 × 10 — полезно.
• Временные массивы. Выражение (a * b + c) // d над массивами numpy создаёт два промежуточных массива размером с исходный и ещё один под результат. На миллионе чисел это незаметно, а на сотнях миллионов память кончается раньше, чем время. Во второй задаче статьи очевидная векторная запись просит почти гигабайт.
• Память, а не арифметика. Когда операция простая — сложить два массива, — numpy упирается не в процессор, а в то, как быстро данные приходят из памяти. Тогда два прохода по массиву вместо одного стоят вдвое дороже, и выгоднее одна операция, которая делает всё за проход. Об этом — статьи о кэше и о задержках.

Когда не поможет

• Каждый шаг зависит от предыдущего. Дизеринг, поколения клеточного автомата, шаги симуляции во времени, динамика, где следующая ячейка ждёт предыдущую. Векторизовать можно только вдоль того направления, где зависимости нет: в «Жизни» из первой задачи поколения идут строго по очереди, но внутри поколения все клетки считаются сразу.
• Ветвления по данным. «Если элемент больше порога — одно, иначе другое» превращается в маску: посчитать обе ветки для всех элементов и выбрать по маске. Если ветки дорогие, а одна из них почти никогда не нужна, это лишняя работа.
• Код и так скомпилирован. В Go, C, Rust или Java цикл — уже машинный код. Там «запись массивами» ради скорости не нужна, а векторные команды компилятор вставляет сам там, где может, — для целых чисел, но не для суммы дробей, — или их берут из библиотек, написанных на ассемблере.

Где работает

• Научные вычисления на Python. numpy, pandas, SciPy и PyTorch — это операции над целыми массивами, написанные на C, C++ и Fortran. Хороший код на них почти не содержит циклов по элементам.
• Изображения и звук. Фильтры — суммы сдвинутых копий сигнала, как размытие выше; яркость, контраст, наложение — поэлементные операции.
• Стандартные библиотеки. Поиск байта, сравнение и копирование памяти, подсчёт — memchr и memcpy в C, bytes.Count и bytes.IndexByte в Go — написаны векторными командами для каждого процессора отдельно.
• Машинное обучение. Умножение матриц — основа нейронных сетей — выполняется векторными командами процессора, а чаще видеокартой, где тысячи простых ядер работают по тому же принципу: одна команда — много данных.

История

Мысль писать вычисления над целыми массивами старше векторных процессоров. В 1962 году Кеннет Айверсон описал в книге «A Programming Language» язык APL, где сложение двух массивов, сумма по оси и выборка по маске — обычные операции языка, без циклов. От APL пошли языки J, K и q, а MATLAB, R и numpy сделали такую запись привычной для учёных и аналитиков.

Векторные команды появились в суперкомпьютерах. Cray-1 Сеймура Крея (1976) держал в векторных регистрах по 64 числа и обрабатывал их конвейером, по одному результату за такт. В массовые процессоры они пришли ради видео и звука: MMX в 1997 году, SSE в 1999-м, NEON на ARM в середине 2000-х, AVX в 2011-м.

numpy вырос из двух библиотек. Numeric написал Джим Хугунин в 1995 году, позже появилась Numarray, а в 2005–2006 годах Трэвис Олифант объединил их в numpy. В 2020 году авторы описали в журнале Nature, как библиотека стала основой научного Python — от снимка тени чёрной дыры до поиска гравитационных волн.

Задачи

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