Статья про хаскелл «Быстрее, чем C++; медленнее, чем PHP» подтолкнула к действию. Давно собирался написать статью о numba и о сравнении её быстродействия с си. В данной статье — чуть более подробный разбор этой ситуации (часть 2) и рекомендации по «приручению» numba (часть 1). В комментариях к этой статье упомянули о библиотеке numba и о том, что она магическим образом может приблизить скорость выполнения кода на питоне к скорости на си.
Разгонять python с переменным успехом стали чуть ли не с первых дней его существования: shedskin, nutika, pythran, parakeet, theano, cython, pypy, numba. Главным недостатком питона принято считать его скорость. Cython (не путать с cpython) — довольно сильно отличается семантически от обычного питона. На сегодняшний день самыми востребованными являются последние три. Что касается pypy (альтернативная реализация транслятора python с использованием jit-компиляции) и numba (библиотека для транскомпиляции кода в llvm) – они пошли разными путями. Фактически это отдельный язык — некий гибрид си и python. В numba же исходили из того, что чаще всего требует ускорения (cpu bound) — математические вычисления, соответственно, они выделили часть языка, связанную с вычислениями и начали разгонять её, постепенно увеличивая «охват» (например, до недавнего времени не было поддержки строк, сейчас она появилась). В pypy изначально была заявлена поддержка всех конструкций python. Numpy поддерживается (с незначительными ограничениями) и в pypy, и в numba. Соответственно, в numba разгоняется не вся программа, а отдельные функции, это позволяет совместить высокую скорость и обратную совместимость с библиотеками, которые numba (пока) не поддерживает.
Моё знакомство с Numba началось в 2015 году вот с этого вопроса на stackoverflow про скорость умножения матриц на питоне: Efficient outer product in python
╔═══════════╦═══════════╦═════════╗
║ method ║ time(ms)* ║ version ║
╠═══════════╬═══════════╬═════════╣
║ numba ║ 9.77 ║ 0.16.0 ║
║ np.outer ║ 9.79 ║ 1.9.1 ║
║ cython ║ 10.1 ║ 0.21.2 ║
║ parakeet ║ 11.6 ║ 0.23.2 ║
║ pypy ║ 16.36 ║ 2.4.0 ║
║ np.einsum ║ 16.6 ║ 1.9.1 ║
║ theano ║ 17.4 ║ 0.6.0 ║
╚═══════════╩═══════════╩═════════╝
* less time = faster
С тех пор произошло много событий в каждой из библиотек, но качественно картина в отношении numba/cython/pypy не изменилась: numba обгоняет cython за счёт использования нативных процессорных инструкций (cython не умеет jit), а pypy – за счёт более эффективного выполнения байткода llvm.
Использую numba в научных вычислениях и при обучении питону в НГУ.
как ускорять
Чтобы ускорить функцию, надо перед её определением вписать декоратор njit:
from numba import njit @njit
def f(n): s = 0 for i in range(n): s += i return s
Для экспериментов рекомендую взять функцию посложнее, поскольку в данном случае numba распознаёт сумму арифметической прогрессии(!) и вычисляет её за O(1). после этого f(n) начинает выполняться быстрее – бывает, на пару порядков.
Смысл в том, что в этом режиме можно использовать неподдерживаемые нумбой операции: нумба на большой скорости доходит до первой такой операции, затем замедляется и до конца функции исполнение продолжается с обычной питоновской скоростью, даже если больше в функции ничего «запретного» не встречается (т.н. Раньше был актуален режим просто @jit (а не @njit). Сейчас от @jit постепенно отказываются, рекомендуется всегда пользоватся @njit (или в полной форме @jit(nopython=True)): в этом режиме нумба ругается исключениями на такие места – всё равно лучше их переписать, чтобы не потерять в скорости. object mode), что, очевидно, нерационально.
что умеет разгонять
Все операторы, функции и классы делятся в отношении нумбы на две части: те, которые нумба «понимает» и те, которые она «не понимает». В разогнанных функциях можно использовать только часть функционала питона и нумбы.
В документации по numba есть два таких списка (с примерами):
- подмножество функционала питона, знакомое нумбе и
- подмножество функционала numpy, знакомое нумбе.
Из примечательного в этих списках:
- нумба «понимает» питоновские списки с быстрым (амортизированное O(1)) добавлением в конец, которые «не понимает» numpy (правда, только однородные – из элементов одного типа),
- numpy'евские массивы, которые отсутствуют в базовом питоне. Понимает также
- кортежи (tuples): они могут, как и в обычном питоне, содержать элементы разных типов.
- с недавних пор str и bytes, правда, только в качестве входных параметров, создавать их (пока?) нельзя.
Никакие другие библиотеки (в частности, scipy и pandas) она не понимает совсем.
Но даже того подмножества языка, которое она понимает, достаточно, чтобы разогнать большую часть кода для научных приложений, на которые numba в первую очередь и ориентирована.
важно!
Из разогнанных функций можно вызывать только разогнанные, не разогнанные нельзя.
(хотя разогнанные функции можно вызывать и из разогнанных и из не разогнанных).
globals
=> Не используйте глобальные переменные в разогнанных функциях (кроме констант). В разогнанных функциях глобальные переменные становятся константами: их значение фиксируется на момент компиляции функции (пример).
сигнатуры
сигнатуры. В нумбе каждой функции сопоставляется один или несколько типов входных и выходных аргументов, т.н. При запуске с другими типами аргументов будут создаваться новые сигнатуры и новые бинарники (старые при этом сохраняются). При первом вызове функции сигнатура формируется и автоматически компилируется соответствующий бинарный код функции. Так что надо либо
– «прогревать кэш», запуская с небольшими размерами входных массивов, либо
– указывать аргумент @jit(cache=True) для сохранения скомпилированного кода на диск с автоматической его загрузкой при последующих запусках программы (правда на практике на сегодняшний день этот первый запуск всё равно немного медленнее, чем последующие, но быстрее, чем без cache=True). Таким образом, «выход на режим» по скорости исполнения для каждой сигнатуры наступает начиная со второго запуска с этими типами аргументов.
Сигнатуры можно задавать вручную: Есть ещё третий способ.
from numba import int16, int32 @njit(int32(int16, int16))
def f(x, y): return x + y f.signatures
[(int16, int16)]
Сигнатур может быть несколько, порядок их следования имеет значение. При запуске функции с сигнатурой, указанной в декораторе, уже первый запуск будет быстрым: компиляция произойдёт в тот момент, когда питон увидит определение функции, а не при первом запуске.
Авторы numba предупреждают о том, что синтаксис указания типов может измениться в будущем, @jit/@njit без сигнатур – более безопасный в этом плане вариант. Предупреждение: этот последний способ не future-safe.
f.signatures начинают показывать сигнатуры только тогда, когда питон о них узнает, то есть после первого вызова функции, либо если они заданы вручную.
Кроме f.signatures сигнатуры можно посмотреть через f.inspect_types() – кроме типов входных параметров эта функция покажет типы выходных параметров, а также типы всех локальных переменных.
Кроме типов входных и выходных параметров, есть возможность вручную указать типы локальных переменных:
from numba import int16, int32 @njit(int32(int16, int16), locals=)
def f(x, y): z = y + 10 return x + z
int
Есть ещё типы int_ (а также float_), используя которые вы предоставляете нумбе возможность выбрать оптимальную (с её точки зрения) ширину поля. В нумбе у целых чисел нет длинной арифметики как в «просто» питоне, но есть стандартные типы различной ширины от int8 до int64.
классы
Поддержка классов (@jitclass) вообще есть, но пока она экспериментальная, так что лучше пока избегать их использования (на текущий момент, по моему опыту, с ними сильно медленнее, чем без них).
custom dtypes
Они работают с той же скоростью, что и обычные массивы numpy, их чуть удобнее индексировать (например, a['y2'] более читаемо, чем a[3]). В numba поддерживается некая альтернатива классам из numpy – структурные массивы (structured array), или, иначе говоря, пользовательские dtype'ы. Но в целом их поддержка в numba оставляет желать лучшего, и некоторые очевидные даже в numpy операции с ними в нумбе записываются достаточно нетривиально. Интересно, что в numba, в отличие от numpy, наряду с обычным синтаксисом a['y2'] допускается более лаконичный a.y2.
GPU
С этим пока разбирался мало. Умеет выполнять разогнанный код на GPU, причём в отличие от того же, например, pycuda или pytorch, не только на nvidia, но и на amd'шных карточках. Там получилась сопоставимая с С скорость. Вот статья на хабре 2016 года Сравнение производительности GPU-расчетов на Python и C.
дополнительные возможности разгона
-
ahead-of-time компиляция: В нумбе есть режим обычной (то есть не jit) компиляции (документация), но этот режим является не основным, я с ним не разбирался.
-
@jit(nopython=True, parallel=True)автоматически распараллеливает выполнение задачи на несколько ядер CPU (документация). параллельное выполнение. -
Для избежания race conditions необходимо использовать синхронизацию потоков. освобождение GIL: функции, декорированные
@jit(nogil=True)исполняются параллельно на нескольких ядрах.
установка
Еще пару лет назад были проблемы с установкой, сейчас всё разрешилось: одинаково хорошо устанавливается и через pip, и через conda; llvm подтягивается и устанавливается автоматически.
документация
Она есть, но в ней есть не всё. Нумбе до сих пор не хватает толковой документации.
оптимизация
Есть некоторая непредсказуемость при оптимизации кода вручную: unpythonic код зачастую работает быстрее, чем pythonic.
Оно правда длинновато и частично устарело (например, строки уже поддерживаются), но общее представление получить помогает. Заинтересовавшимся темой могу порекомендовать видео мастер-класса по numba с конференции scipy 2017 (есть исходники на гитхабе).
Во второй части рассмотрим применение numba на примере кода из упомянутой в начале статьи.
