Если вам нужно ускорить обработку NumPy или просто сократить использование памяти, попробуйте компилятор Numba just-in-time. С его помощью можно писать код на языке Python, который во время выполнения компилируется в машинный код. Это позволяет получить прирост скорости, сопоставимый с приростом, который можно получить на C, Fortran или Rust.

По крайней мере, так считается в теории. На практике же код на Numba может быть не быстрее, чем эквивалент NumPy.

Но если лучше понимать, как работают процессоры, можно добиться большего. Эти знания помогут более широко использовать любой компилируемый язык.

В этой статье мы:

• Рассмотрим простую задачу обработки изображений.

• Попытаемся (поначалу безуспешно) ускорить ее с помощью Numba.

• Рассмотрим, почему современные процессоры такие быстрые, и каковы возможности компиляторов.

• Опираясь на полученные знания, мы скорректируем наш код так, чтобы он выполнялся в 25 раз быстрее по сравнению с первоначальной версией.

Задача: удаление шума из изображения

Представьте, что вы получаете 16-разрядное изображение с цифрового микроскопа, и вас интересуют только яркие участки изображения. Темные области не несут полезной информации и содержат много шума. Чтобы очистить эти области, нам нужно установить все значения ниже определенного порога в полный черный цвет, т.е. в 0.

Вот как это можно сделать с помощью NumPy:

image[image < 1000] = 0

Чтобы посмотреть, насколько быстро это работает, я сгенерировал изображение:

import numpy as np
rng = np.random.default_rng(12345)
noise = rng.integers(0, high=1000, size=(4096, 4096), dtype=np.uint16)
signal = rng.integers(0, high=5000, size=(4096, 4096), dtype=np.uint16)
image = noise | signal

А затем замерил время с помощью %timeit, доступной в Jupyter и IPython, с отключенным turboboost на моем процессоре. Код изменяет изображение, поэтому мы работаем с копией, а копирование добавляет дополнительные расходы, поэтому мы также измеряем эти расходы:

>>> %timeit image2 = image.copy()
5.84 ms ± 189 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
>>> %timeit image2 = image.copy(); image2[image2 < 1000] = 0
54.2 ms ± 184 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Для выполнения этого простейшего алгоритма с помощью NumPy требуется около 48 мс (54 мс - 6 мс). Можем ли мы добиться большего?

Попытка № 0: исходный алгоритм, переведенный в Numba

Numba позволяет писать код на подмножестве языка Python и компилировать его в машинный код. Чтобы реализовать наш алгоритм, для начала сделаем следующее:

from numba import njit

@njit
def remove_noise_numba_0(arr, noise_level):
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            if arr[i, j] < noise_level:
                arr[i, j] = 0

remove_noise_numba_0(image.copy(), 1000)

Убедимся, что запустили эту функцию один раз, чтобы в дальнейшем при подсчете времени не учитывать время ее компиляции.

Измерение этого времени с помощью %timeit показывает, что то оно составляет 46 мс, что не намного лучше, чем в нашем коде на NumPy.

Попытка №1: Выбор хороших типов

Хотя Numba визуально выглядит как Python, лучше всего воспринимать его как C или другой статически-типизированный компилируемый язык. Он будет генерировать код специально для типов передаваемых данных.

Мы знаем тип массива: двумерный массив беззнаковых 16-битных целых чисел. А вот тип уровня noise_level Numba должен угадать, поскольку мы передали целое число Python, и он захочет преобразовать его в примитивный тип, эквивалентный dtype NumPy.

Однако мы знаем, что он также должен быть беззнаковым 16-битным целым числом (уровень шума не может быть отрицательным или больше, чем самое яркое пятно), поэтому можно подправить код:

@njit
def remove_noise_numba_1(arr, noise_level):
    noise_level = arr.dtype.type(noise_level)
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            if arr[i, j] < noise_level:
                arr[i, j] = 0

remove_noise_numba_1(image.copy(), 1000)

Хотя цель не заключалась в ускорении работы — целью была корректность — это изменение позволило увеличить скорость до 33 мс. Можем ли мы добиться большего?

Чтобы ответить на этот вопрос, давайте рассмотрим, как работают процессоры.

Лучшая «ментальная» модель для современных процессоров

Во-первых, все, о чем мы здесь говорим, относится к одному ядру процессора. Мы полностью игнорируем возможность распараллеливания на несколько ядер и просто предполагаем, что на одном ядре работает один поток.

Первая «ментальная» модель процессора состоит в том, что он считывает машинные инструкции и затем выполняет их по очереди.

Эта модель проста для понимания и достаточна для многих целей, но если вы хотите писать более быстрый код, вам нужна более точная модель. Модель «одна инструкция за раз» предполагает, что выполнение большего количества инструкций замедлит код в зависимости от количества добавленных инструкций. Это не всегда так. Напротив, различные инструкции могут оказывать совершенно разное влияние на производительность.

Современные процессоры могут выполнять инструкции параллельно

На самом деле современные процессоры, которые можно встретить в настольных компьютерах, ноутбуках и серверах, могут выполнять несколько инструкций одновременно, если эти инструкции не мешают друг другу. Рассмотрим следующий код — представьте, что он написан на языке C или другом компилируемом языке, так чтобы выражения однозначно отображались на инструкции процессора:

a = a + 1;
b = b + 1;

Эти две строки кода не связаны друг с другом. Если их перевести в машинный код, то процессор вполне может выполнить обе строки более или менее одновременно. Опять же, это не имеет никакого отношения к нескольким ядрам процессора или нескольким потокам: это все один поток на одном ядре.

Теперь рассмотрим, что происходит в следующем коде:

a = a + 1;
if (a > 100) {
    b = b + 1;
}

В данном примере значение b зависит от значения a. Это означает, что третья строка кода не может быть выполнена до тех пор, пока не будет выполнена первая строка, а затем будет выполнено сравнение. Больше никакого параллелизма, и наш код внезапно стал намного медленнее!

Современные процессоры спекулятивно выполняют и предсказывают ветвления

Для того чтобы вернуть часть недостающего параллелизма, современные процессоры выполняют более поздний код и предсказывают, какая ветвь будет выбрана при том или ином условии — исходя из допущения, что большинство кода на самом деле достаточно предсказуемо. Рассмотрим:

int a = 0;
int b = 0;
while (true) {
    a = a + 1
    if (a > 100) {
        b = b + 1;
    }
}

В течение первых ста итераций b никогда не будет увеличиваться. После этого b будет увеличиваться всегда. За исключением переходного периода, процессор, скорее всего, может предсказать, что произойдет дальше, и наугад выбрать, какой код выполнить.

Что произойдет, если процессор угадает неправильно? Это отменит работу, что может быть затратно. А это значит, что неудачные предсказания ветвлений могут существенно снизить производительность.

Современные процессоры имеют специальные инструкции для выполнения параллельной работы: SIMD

В дополнение к прозрачно распараллеливаемым инструкциям современные процессоры имеют инструкции SIMD: Single Instruction Multiple Data. Идея заключается в том, чтобы выполнять одну операцию на множестве элементов массива с помощью одной инструкции.

Например, здесь мы вручную в цикле добавляем 100 к четырем записям массива:

int array[4] = {1, 2, 3, 4};
for (int i=0; i < 4; i++) {
    array[i] += 100;
}

Вместо этого с помощью SIMD мы можем использовать специальную инструкцию процессора для добавления 100 ко всем четырем элементам массива одновременно. Даже если эта инструкция работает немного медленнее, поскольку мы меняем четыре инструкции (и, возможно, часть цикла) на одну инструкцию, в целом можно ожидать, что код будет выполняться намного быстрее.

Поддержка этих специальных инструкций возможна потому, что мы выполняем одну и ту же операцию над всеми четырьмя значениями, в данном случае добавляем константу. Добавление двух массивов друг к другу — это опять же одна и та же операция, выполняемая со всеми элементами. Ключ — в названии: одна инструкция, много данных.

Можно использовать SIMD-инструкции явно, но Numba этого не поддерживает, и, кроме того, это обычно делает код специфичным для конкретного процессора. Процессоры ARM, используемые в новых компьютерах Mac, имеют другие SIMD-инструкции, в отличие от процессоров x86-64, и даже процессоры x86-64 имеют разные SIMD-инструкции в зависимости от модели.

В качестве альтернативы можно надеяться, что ваш компилятор достаточно умен, чтобы автоматически использовать SIMD в соответствующих местах. В случае Numba, поскольку он компилируется «на лету» для вашего компьютера, он может использовать любые SIMD-инструкции, которые поддерживает ваш процессор.

Помощь компилятору: линейный код — быстрый код

Независимо от того, идет ли речь об автоматическом параллелизме инструкций процессора или о параллелизме на основе SIMD, операторы if и другие условные конструкции представляют собой проблему.

• При параллелизме инструкций условный код означает, что процессору придется угадывать, что произойдет дальше, и он может ошибиться.

• В случае SIMD эти инструкции, как правило, должны выполнять одни и те же действия со всеми элементами небольшого массива. Условный код значительно усложняет выполнение одних и тех же действий.

Чтобы ускорить код, мы должны постараться сделать его как можно более линейным, без операторов if или других условий ветвления. Это одновременно и повысит вероятность того, что компилятор сможет определить возможности использования SIMD, и облегчит процессору параллельное выполнение результирующих инструкций.

Проверка кода с помощью perf

Мы узнали, что условный код затрудняет компилятору написание быстрого кода, а процессору — быстрое выполнение полученного кода. Вот наш текущий код Numba:

@njit
def remove_noise_numba_1(arr, noise_level):
    noise_level = arr.dtype.type(noise_level)
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            if arr[i, j] < noise_level:
                arr[i, j] = 0

Там определенно есть условие! Код иногда записывает в массив, а иногда нет, в зависимости от сравнения.

Мы можем измерить часть этого влияния с помощью инструмента perf в Linux, используя следующий декоратор для подключения к текущему процессу и получения данных о работе процессора в течение интервала, к которому он подключен:

from subprocess import Popen
from contextlib import contextmanager
from os import getpid
from time import sleep
from signal import SIGINT

@contextmanager
def perf_stat():
    p = Popen(["perf", "stat", "-p", str(getpid())])
    sleep(0.5)
    yield
    p.send_signal(SIGINT)

Можнозапустить это в нашем коде:

image2 = image.copy()
with perf_stat():
    remove_noise_numba_1(image2, 1000)

В результате получается много информации, но для наших целей нам интересны два числа: процент ненужного спекулятивного выполнения и процент плохих предсказаний.

...
52.8% Bad Speculation
52.8% Branch Mispredict
...

Мы видим высокий процент ошибок в прогнозировании — и это неудивительно, если вспомнить, что делает наш алгоритм: сравнивает кучу зашумленных чисел со значением и смотрит, они больше или меньше порогового значения. Это очень трудно предсказать для изображений шума.

Попытка №2: Избавление от условных элементов

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

После нескольких минут отладки у меня получился следующий код. Он дает те же результаты, но без ветвления. К сожалению, понять его гораздо сложнее: он опирается на кучу трюков, связанных с тем, как целые числа представляются в памяти процессора.

@njit
def remove_noise_numba_2(arr, noise_level):
    noise_level = arr.dtype.type(noise_level)
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            res = arr[i, j] - noise_level
            mask_to_zero_if_wrapped = -(res <= arr[i, j])
            res = (res & mask_to_zero_if_wrapped) + (
                   noise_level & mask_to_zero_if_wrapped)
            arr[i, j] = res

remove_noise_numba_2(image.copy(), 1000)

Если бы процессоры действительно выполняли по одной инструкции за раз, то эта версия, скорее всего, была бы медленнее, чем наша предыдущая попытка: например, мы выполняем запись для всех элементов массива. Но в настоящее время процессоры работают не так, и поэтому эта версия намного быстрее.

Фактически, время работы составляет всего 4 мс.

И если мы измерим это с помощью perf, то увидим, что мы больше не страдаем от ошибок предсказания ветвей:

...
2.4% Bad Speculation
2.4% Branch Mispredict
...

Попытка №3: Помочь компилятору помочь нам

Избавившись от условных обозначений, мы значительно ускорили наш код — ценой непонятного, уродливого и трудноразрешимого кода. Можно ли сделать лучше?

С математической точки зрения две наши версии одинаковы. Можно предположить, что компиляторы будут стараться избегать условных обозначений, если очевидно, что это замедлит работу процессора. И компиляторы, конечно, очень хорошо умеют превращать неэффективные математические выражения в эффективные.

Почему же компилятор не смог этого сделать в данном случае? Давайте еще раз рассмотрим неэффективную версию:

@njit
def remove_noise_numba_1(arr, noise_level):
    noise_level = arr.dtype.type(noise_level)
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            if arr[i, j] < noise_level:
                arr[i, j] = 0

Проблема в том, что мы делаем только условную запись: для чисел, превышающих уровень шума, мы ничего не делаем. Моя “ментальная” модель авторов компиляторов такова, что, хотя они с удовольствием удаляют ненужные считывания памяти, они откажутся добавлять дополнительные записи в память, о которых вы не просили.

Так что если мы будем писать в память в обоих случаях, а в остальном просто сохраним условие?

@njit
def remove_noise_numba_3(arr, noise_level):
    noise_level = arr.dtype.type(noise_level)
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            arr[i, j] = (
                arr[i, j] if arr[i, j] >= noise_level
                else 0
            )

remove_noise_numba_3(image.copy(), 1000)

Да, условие осталось, но это условие внутри довольно простого вычисления, которое не влияет на запись в память. Мы по-прежнему пишем по одному и тому же адресу памяти по обе стороны условия. Умный компилятор заметит, что существует эквивалентное математическое выражение, не требующее условия, и переключится на его использование, если посчитает, что это ускорит работу кода.

При выполнении этой версии время выполнения составляет 2 мс — это самая быстрая версия! Какое бы выражение ни нашел компилятор, оно даже более эффективно, чем грязное битовое жонглирование, что мы делали вручную.

Что мы узнали на данный момент

Вот производительность всех рассмотренных нами версий:

Нам удалось получить код, который работает в 25 раз быстрее, но при этом остается вполне читаемым! Мы добились этого за счет:

1. Перешли на Numba, чтобы иметь больше контроля над выполнением.

2. Убедились, что все типы явно указаны и одинаковы (в данном случае np.uint16).

3. Удалили условные записи; это позволило компилятору перейти к линейной реализации без ветвлений.

Никакой многопоточности или многопроцессорности не было: это все однопоточный код на одном ядре процессора.

Результаты могут различаться в зависимости от версии компилятора и процессора

С одной стороны, все компиляторы стараются написать код, который будет эффективен на целевом процессоре — поэтому аналогичные подходы, скорее всего, будут работать и в других языках, не только в Numba. На самом деле эта статья была частично вдохновлена статьей Алекса "matklad" Кладова, в которой используется язык Rust.

При этом разные версии компиляторов и их реализации ведут себя по-разному, меняя способы оптимизации. Это может иметь значение. Так, код, который автоматически оптимизируется в одной версии, может регрессировать в другой. Если скорость критична, важно иметь контрольные показатели, чтобы отлавливать регрессии до их развертывания.

При предварительной компиляции, используемой в C, Rust, Cython и т.д., компилятор обычно ориентируется на подмножество доступных функций, чтобы код был переносим между моделями процессоров. Чтобы максимизировать производительность при работе только на конкретном оборудовании, можно выбрать целевой процессор или определенный набор функций. Numba компилируется непосредственно во время работы компьютера, поэтому она ориентирована на конкретные возможности текущего процессора компьютера. Это означает, что разные компьютеры могут давать разные результаты.

Вот пример: читатель Надав Хореш (Nadav Horesh) протестировал код из этой статьи на процессоре Xeon, который поддерживает SIMD-инструкции AVX-512, в отличие от моего i7-12700K. На его конкретной установке Numba получила значительное ускорение на начальной версии, remove_noise_numba_0(), по сравнению с NumPy, причем никаких дополнительных изменений не потребовалось. По всей видимости, это ускорение было вызвано тем, что Numba использовала SIMD-инструкции AVX-512, а не машинный код с ветвлениями, сгенерированный на моем компьютере.

Какие из вариантов, основанных на Numba, в итоге использовали SIMD?

Для SIMD с плавающей запятой процессор Intel моего компьютера может сообщить, сколько инструкций было выполнено (доступно, например, через perf stat), но в данном примере мы используем целочисленный SIMD, и эта статистика недоступна. Однако можно получить диагностику использования SIMD для компиляции Numba just-in-time, предварительно выполнив эту процедуру:

import llvmlite.binding as llvm
llvm.set_option('', '--debug-only=loop-vectorize')

Насколько я могу судить, все попытки на базе Numba использовали тот или иной уровень SIMD. Но SIMD — это только один из способов ускорить код. Есть и другие узкие места, а в первой версии (по крайней мере, на моем компьютере) использовалось много ветвлений.

Подробнее

• Книга Computer Systems: A Programmer's Perspective (3-е издание) — дает отличное введение в принципы работы аппаратного обеспечения, ориентированное на потребности людей, пишущих программы.

• В репозитории примеров Numba есть SIMD notebook, в котором рассматриваются детали повышения вероятности того, что SIMD будет работать и выполняться быстро.

---

Материал подготовлен в преддверии старта курса Python Developer.

Недавно в рамках этого курса прошел открытый урок «Работа с пакетами в Python с помощью pip и poetry». На этом занятии участники разобрали оба пакетных менеджера, поговорили про основные сценарии использования, обсудили различия, а также узнали в каких случаях можно обойтись решением попроще, а в каких случаях требуется более продвинутый подход. Если тема для вас актуально, делимся записью этого урока.