Pythran: как заставить работать код Python со скоростью С++
Хотите писать программы на Python, работающие со скоростью кода, написанного на С++? Достаточно добавить аннотацию Pythran!
Инструменты Python многогранны, и с их помощью "змеиный язык" можно легко разогнать до скорости С++. Как? Рассказываем.
Python – высокоуровневый универсальный язык, который почти так же легко читать и писать, как псевдокод. Но его главная проблема – низкая производительность. Это становится особенно проблематичным при работе с большими многомерными массивами. Решением стала библиотека NumPy, которая вместо стандартных объектов Python использует оптимизированные алгоритмические решения.
Pythran преобразует функции Python в нативный код. Библиотека берёт Python-модуль, аннотированный небольшим интерфейсным описанием, и превращает его в нативный модуль с тем же интерфейсом, но более быстрым. Pythran предназначен для эффективной компиляции программ с использованием нескольких ядер и SIMD-инструкций. Пакет поддерживает как вторую, так и третью версии Python, работает на Windows, Linux и macOS.
В качестве примера рассмотрим алгоритм оптимизации проблемы коммивояжера методом имитации отжига. Соответствующий код приведен на GitLab и в Jupyter-блокноте.
Задача коммивояжёра
Рассмотрим пример хорошо известной проблемы коммивояжёра. Известен список из N городов и расстояния между каждой парой городов. Нужно определить кратчайший маршрут посещения каждого города с возвращением в исходный пункт. Оптимизационная постановка задачи относится к классу NP-трудных задач и требует значительных вычислительных мощностей.
Для конкретики взят набор из 100 пунктов на пространстве [0, 1]x[0, 1].
Алгоритм имитации отжига
Алгоритм имитации отжига – это эвристический алгоритм для аппроксимации глобального экстремума функции. Он часто используется, когда пространство поиска дискретно. Алгоритм удобен тем, что не налагает никаких предварительных условий на оптимизируемую функцию.
Идея основана на имитации физического процесса, происходящего при кристаллизации вещества, в том числе при отжиге металлов. Предполагается, что атомы уже выстроились в кристаллическую решётку, но ещё допустимы переходы отдельных атомов из одной ячейки в другую. Процесс протекает при постепенно понижающейся температуре. Переход атома из одной ячейки в другую происходит с некоторой вероятностью, причём вероятность уменьшается с понижением температуры. Устойчивая кристаллическая решётка соответствует минимуму энергии атомов, поэтому атом либо переходит в состояние с меньшим уровнем энергии, либо остаётся на месте. В результате система стремится к общему глобальному минимуму.
Алгоритм чрезвычайно прост, гибок и применим к широкому кругу реальных проблем. В алгоритме можно адаптировать:
- Функцию энергии (состояний).
- Функцию изменения.
- Профиль снижения температуры в процессе отжига.
Наиболее важной составляющей является функция изменения, которая определяет процесс решения задачи. Можно рассматривать её как вид генетической мутации, которая либо сохраняется, либо отбрасывается. В случае задачи коммивояжера изменения на каждом шаге – это пути между случайными точками i и j.
В случае задачи коммивояжера уникальной сигнатурой состояния является кортеж городов, сопоставленных целым числам от 0 до N.
Аннотирование Pythran
Единственное дополнение, которое необходимо сделать в исходном коде, это аннотация. Аннотация – это строковый комментарий, начинающийся с # pythran. Pythran учитывает эту аннотацию для компиляции функции в нативный модуль. Пример аннотации:
# pythran export search_for_best(int, float list list, int, float, int, float, float)
def search_for_best(seed,
cities,
nb_step,
beta_mult=1.005,
accept_nb_step=100,
p1=0.2,
p2=0.6):
"""exported function"""
# etc
Аннотации Pythran можно ставить в любом месте кода. Предпочтительнее в начале файла или сразу над компилируемыми функциями. Такие аннотации имеют следующий синтаксис:
# pythran export function_name(argument_type*)
где function_name – это имя функции, определенной в модуле, а argument_type* – разделенные запятыми типы аргументов. Можно использовать составные аргументы в виде кортежей, например, (int,(float, str)), или списков. Вот полная грамматика:
argument_type = basic_type | (argument_type+) # this is a tuple | argument_type list # this is a list | argument_type set # this is a set | argument_type []+ # this is a ndarray, C-style | argument_type [::]+ # this is a strided ndarray | argument_type [:,...,:]+ # this is a ndarray, Cython style | argument_type [:,...,3]+ # this is a ndarray, some dimension fixed | argument_type:argument_type dict # this is a dictionary basic_type = bool | int | float | str | None | uint8 | uint16 | uint32 | uint64 | uintp | int8 | int16 | int32 | int64 | intp | float32 | float64 | float128 | complex64 | complex128 | complex256
Ключевым моментом является возможность передачи многомерных массивов, что позволяет обычным пользователям NumPy использовать свой код «как есть». Им лишь необходимо изолировать вычислительно интенсивные операции в функции с некоторованным фиксированным типом.
Установка Pythran
Pythran можно установить на трех основных операционных системах: Linux, macOS и Windows. Ниже предполагается, что в качестве среды установлена miniconda3.
Linux
Представлены шаги, которые работали на облачной виртуальной машине. Для Linux вы можете посмотреть соответствующий скрипт script_setup_linux.sh.
Устанавливаем компилятор:
# ubuntu $ sudo apt-get update $ sudo apt-get -y install g++
Устанавливаем Pythran в виртуальном окружении conda env:
$ conda create -n pythran python=3 $ source activate pythran $ pip install numpy $ pip install pythran
macOS
Инструкции для macOS проверены на HighSierra и Mojave. Соответствующий скрипт: script_setup_macos.sh.
Устанавливаем компилятор с homebrew:
$ brew install llvm
Устанавливаем Pythran в виртуальном окружении conda env:
$ conda create -n pythran python=3 -y $ source activate pythran $ conda install -c conda-forge pythran -y
Альтернативный способ установки pip install pythran, к сожалению, приводит к потере некоторых зависимостей, например, пакета blas. Поэтому для простоты рекомендуется использовать conda install.
На macOS Mojave может понадобиться установить SDK Headers for macOS 10.14. Для этого под sudo необходимо сделать
$ xcode-select --install
Далее пройдите к /Library/Developer/CommandLineTools/Packages/. Там вы увидите пакет macOS_SDK_headers_for_macOS_10.14.pkg. Установите его. Создайте файл ~/.pythranrc:
[compiler] cflags=-std=c++11 -fno-math-errno -w ldflags=-L/usr/local/opt/llvm/lib # from `brew info llvm` CC=/usr/local/opt/llvm/bin/clang # brew installed clang path CXX=/usr/local/opt/llvm/bin/clang++ # brew installed clang++ path
Windows 7
Установите компилятор. Для этого в Visual Studio Community 2017 нужно установить Desktop Development with C++.
Установка Pythran в виртуальном окружении conda env:
$ conda create -n pythran python=3 $ source activate pythran $ pip install numpy $ pip install pythran
Убедитесь, что файл ~/.pythranrc не существует, пуст или содержит следующее:
[compiler] defines= undefs= include_dirs= libs= library_dirs= cflags=/std:c++14 ldflags= blas=blas CC= CXX= ignoreflags=
Компиляция и запуск с Pythran
Компиляция
В качестве примера для Pythran были аннотированы две функции в представленных ниже двух py-файлах:
# regular compilation - under pythran env # In my case it's safe to use -Ofast (affects precision) $ pythran -Ofast -march=native tsp_compute_single_threaded.py
# compilation with omp - under pythran env # the compilation flags activate OMP and vectorization using # https://github.com/QuantStack/xsimd $ pythran -DUSE_XSIMD -fopenmp -march=native tsp_compute_multi_threaded_omp.py
В обоих случаях, если Pythran правильно установлен, при запуске кода создается модуль с названием вида:
- macOS: tsp_compute_[xxx].cpython-37m-darwin.so
- Linux: tsp_compute_[xxx].cpython-37m-x86_64-linux-gnu.so
- Windows: tsp_compute_[xxx].cp37-win_amd64.pyd
Запуск
Создаваемые модули ведут себя аналогично стандартным модулям Python, за исключением того, что работают быстрее. Если в каталоге присутствуют обычный модуль и модуль Pythran, то импортируется последний.
Параллелизм
Поиск возможностей ускорения кода не ограничивается вычислительными операциями. В нашем примере задача отжига является вероятностным поиском и может быть распараллелена. Здесь мы протестируем два метода:
- concurrent.futures на уровне Python
- OMP (Open Multi Processing) на уровне Pythran
OMP
OMP (Open Multi Processing) позволяет проводить мультипроцессинг через относительно простой API. В коде C++ параллелизм описывается с помощью OMP директив, вводимых комментариями, начинающимися с # pragma omp. Аналогично Pythran позволяет разработчикам на Python писать такие директивы в коде Python комментариями, начинающимися с # omp.
В качестве примера работы с OMP посмотрите tsp_compute_multi_threaded_omp.py.
concurrent.futures
Эффективной альтернативой параллелизма между CPU на уровне Python является использование параллельных фьючерсов. Посмотрите функцию search_concurrent в файле tsp_wrapper.py. Она использует ProcessPoolExecutor для создания подпроцессов и сбора результатов.
Демо-пакет и блокнот
Данный блокнот содержит пример пользовательского интерфейса к демо-пакету, написанному с использованием Pythran. Он позволяет:
- Компилировать функции tsp_compute_(single|multi)_threaded.py в pythran-модули. Важно: перезагрузите кернел после компиляции, чтобы загрузить обновленный модуль.
- Откатить модуль обратно к Python-модулю.
- Сгенерировать случайный набор городов и настроить параметры.
- Запустить оба типа представления: concurrent.futures и OMP; с проверкой сигнатуры и без нее.
- Визуализировать и сохранить результаты.
Ниже представлен пример трех лучших результатов расчета. Результаты найдены для одного набора случайных точек при параллельном запуске 32 процессов.
Производительность
Для оценки производительности поиск оптимального пути был запущен:
- в двух версиях: с concurrent.futures и OMP, без проверки подписи;
- на трех типах машин: ноутбук (macOS), стационарный компьютер (iMac), виртуальная машина (Linux).
laptop: macOS MacBook Pro 2017 2.3 GHz Intel Core i5 desktop: macOS iMac 2014 4.0 GHz Intel Core i7 VM GCP: Linux n1-highcpu-64 — Ubuntu:18.04
В тесте внимание сосредоточено на версии без проверки подписи. Таблица иллюстрирует результаты решения задачи с выполнением миллиона шагов алгоритма.
Таблица была создана с помощью запуска Jupyter-блокнота на трех машинах и аггрегирования результатов в общем блокноте.
Результаты дают примерно 16-кратный прирост на macOS и ~32-кратный прирост на виртуальной машине Linux. Добавление параллелизации в Pythran дает ~32-кратный прирост на ноутбуке, ~70-кратный на десктопе и ~800-кратный на Linux (подробнее читайте в оригинале публикации).
Упаковка / Распространение
Когда дело доходит до упаковки / распространения программного кода, интеллектуальный дизайн Pythran оказывается особенно удобным. В зависимости от ваших целей, программы можно распространять через:
- Исходный код Python (.py)
- Сгенерированный код C++ (.cpp)
- Родной тип модуля (.so/pyd)
Pythran предоставляет расширение distutils, которое позволяет стандартно использовать файл setup.py. Но если вы собираетесь распространять исходный Python код, необходимо, чтобы у пользователей был и Pythran, и компилятор.