#python
Продолжаем подготовку к
Python-собесам!В прошлом посте мы перенесли тяжёлую функцию из
Python в C, используя механизм FFI. Сегодня мы научимся делать это более элегантно 💅 посредством Python-компиляторов, таких как Numba, Pythran или JAX. Они компилируют подмножество Python в машинный код для CPU или GPU, тем самым повышая производительность и оставаясь в рамках Python.У каждого из приведённых инструментов своё собственное поддерживаемое подмножество
Python, но базово можно считать, что все они способны компилировать функции, принимающие на вход, порождающие внутри себя и возвращающие:-
numpy-структуры данных, исключая object;- Python-примитивы (
int, float, complex, bool, str, ...);- реже - гомогенные числовые списки (
List[float], List[int], ...).Python-компиляторы работают следующим образом:- Вы пишете
Python-функцию и помечаете её компилируемой с помощью декоратора (Numba, JAX) или комментария (Pythran).- Далее, если компилятор работает как JIT (Just-in-Time), то он... ничего не делает 😄 (
Numba, JAX), а вот AOT-компилятор сразу приступает к «конвертации» функции в машинный код (Pythran). Недостатком второго подхода является то, что он требует задавать типы входных данных на этапе написания кода. Кстати, в Numba так тоже можно сделать.- При исполнении
Python-кода в случае AOT'а сразу запускается «бинарная» версия функции, а в случае JIT'а - функция компилируется при первом вызове, а при последующих уже используется её производительная версия. В связи с этим JIT'ованные функции важно «прогревать» - вызывать их с примером входных параметров при инициализации приложения, дабы исключить задержки компиляции при исполнении. Отмечу, что Numba и Pythran налагают ограничения только на типы входных данных, а JAX - ещё и на их размер, то есть для Numba и Pythran вызовы func([1, 2, 3]) и func([4, 5]) идентичны, а для JAX они, грубо говоря, приводят к двум разным функциям.Пара важных нюансов:
- Переходы между
Python и скомпилированным кодом дороги. Чтобы их минимизировать, следует делать компилируемые функции крупнее. Например, если в коде есть обработка данных в цикле, то лучше компилировать цикл целиком.- Чтобы не бить проект на отдельные компилируемые функции, можно использовать режим
jit в Numba: он скомпилирует лишь то, что сможет. К сожалению, его производительность непредсказуема.К примерам! Напишем функцию
fmul из прошлого поста на Numba и Pythran. А на JAX можете написать сами 😄.Numba:
Поставим
Numba:pip install numba
Затем создадим файл
numba_demo.py:import numpy as np
from numba import njit
@njit
def fmul_numba(a, b, out, n):
for i in range(n):
out[i] = a[i] * b[i]
out = np.empty(4, dtype=np.float64)
fmul_numba(np.array([0.0, 0.5, 1.0, 1.5]), np.array([1.0, 2.0, 3.0, 4.0]), out, 4)
print(out)
Запустим его:
python3 numba_demo.py
[0. 1. 3. 6.]
Pythran:
pip install pythran
fmul_pythran.py:# pythran export fmul_pythran(float[], float[], float[], int)
def fmul_pythran(a, b, out, n):
for i in range(n):
out[i] = a[i] * b[i]
pythran_demo.py:import numpy as np
from fmul_pythran import fmul_pythran
import sys
print(sys.modules["fmul_pythran"])
out = np.empty(4, dtype=np.float64)
fmul_pythran(np.array([0.0, 0.5, 1.0, 1.5]), np.array([1.0, 2.0, 3.0, 4.0]), out, 4)
print(out)
pythran fmul_pythran.py
ls
fmul_pythran.cpython-312-x86_64-linux-gnu.so fmul_pythran.py pythran_demo.py
python3 pythran_demo.py
<module 'fmul_pythran' from './fmul_pythran...so'>
[0. 1. 3. 6.]
Кстати, пару лет назад на одном из собесов мне пришлось писать
Numba-код, так что это даже более актуально, чем FFI 🤓.Сегодня мы познакомились с
Python-компиляторами, в следующем посте поговорим о том, как понять, какие именно функции нам надо компилировать! ⚙️🚀С уважением,
Михаил Масягин