Получите скорость C ++ / Fortran для ваших функций с Numba

Это третий пост из серии, которую я пишу. Все посты здесь:

  1. Ускорьте свои алгоритмы. Часть 1 - PyTorch
  2. Ускорьте свои алгоритмы, часть 2 - Numba
  3. Ускорьте свои алгоритмы. Часть 3 - Распараллеливание
  4. Ускорьте свои алгоритмы. Часть 4 - Dask

И это относится к блокнотам Jupyter, доступным здесь:

[Github-SpeedUpYourAlgorithms] и [Kaggle]

Показатель

  1. Вступление
  2. Почему Нумба?
  3. Как работает Numba?
  4. Использование основных функций numba (просто @jit it!)
  5. Обертка @vectorize
  6. Запуск ваших функций на GPU
  7. Дальнейшее чтение
  8. использованная литература
NOTE:
This post goes with Jupyter Notebook available in my Repo on Github:[SpeedUpYourAlgorithms-Numba]

1. Введение

Numba - это компилятор точно в срок для python, т.е. всякий раз, когда вы вызываете функцию python, весь или часть вашего кода преобразуется в машинный код «точно в срок »выполнения, и затем он будет работать на скорости вашего машинного кода! Он спонсируется Anaconda Inc и поддерживается многими другими организациями.

С помощью Numba вы можете ускорить выполнение всех ваших вычислительных и сложных функций Python (например, циклов). Также есть поддержка библиотеки numpy! Таким образом, вы также можете использовать numpy в своих вычислениях и ускорить общие вычисления, поскольку циклы в python очень медленные. Вы также можете использовать многие функции математической библиотеки стандартной библиотеки Python, такие как sqrt и т. Д. Полный список всех совместимых функций смотрите здесь.

2. Почему Нумба?

["Источник"]

Итак, почему онумба? Когда есть много других компиляторов вроде cython или любых других подобных компиляторов или чего-то вроде pypy.

По той простой причине, что здесь вам не нужно выходить из зоны комфорта написания кода на Python. Да, вы прочитали правильно, вам не нужно вообще изменять свой код для базового ускорения, которое сравнимо с ускорением, которое вы получаете от аналогичного кода Cython с определениями типов. Разве это не здорово?

Вам просто нужно добавить знакомую функциональность Python, декоратор (оболочку) вокруг ваших функций. Также в разработке оболочка для класса.

Итак, вам просто нужно добавить декоратор, и все готово. например:

from numba import jit
@jit
def function(x):
    # your loop or numerically intensive computations
    return x

Он по-прежнему выглядит как чистый код Python, не так ли?

3. Как работает нумба?

Numba генерирует оптимизированный машинный код из чистого кода Python с использованием инфраструктуры компилятора LLVM. Скорость выполнения кода с использованием numba сопоставима со скоростью выполнения аналогичного кода на C, C ++ или Fortran.

Вот как компилируется код:

["Источник"]

Сначала берется функция Python, оптимизируется и преобразуется в промежуточное представление Numba, а затем после вывода типа, который похож на вывод типа Numpy (так что python float - это float64), она преобразуется в интерпретируемый код LLVM. Затем этот код передается в оперативный компилятор LLVM для выдачи машинного кода.

Вы можете сгенерировать код во время выполнения или импортировать время на CPU (по умолчанию) или GPU, как вам удобнее.

4. Использование основных функций numba (просто @jit it!)

Кусок торта!

Для лучшей производительности numba рекомендует использовать аргумент nopython = True с вашей jit-оболочкой, при этом он вообще не будет использовать интерпретатор Python. Или вы также можете использовать @njit. Если ваша оболочка с nopython = True выйдет из строя с ошибкой, вы можете использовать простую @jit оболочку, которая скомпилирует часть вашего кода, выполнит цикл, который он может компилировать, и превратит их в функции, чтобы скомпилировать в машинный код и передать остальное интерпретатору python. < br /> Итак, вам просто нужно сделать:

from numba import njit, jit
@njit      # or @jit(nopython=True)
def function(a, b):
    # your loop or numerically intensive computations
    return result

При использовании @jit убедитесь, что в вашем коде есть что-то, что numba может компилировать, например, цикл с интенсивными вычислениями, возможно, с библиотеками (numpy) и функциями, которые он поддерживает. В противном случае он ничего не сможет скомпилировать.

Чтобы добавить вишенку, numba также кэширует функции после первого использования в качестве машинного кода. Таким образом, после первого раза это будет еще быстрее, потому что не нужно снова компилировать этот код, учитывая, что вы используете те же типы аргументов, которые использовали раньше.

И если ваш код распараллеливаемый, вы также можете передать parallel = True в качестве аргумента, но он должен использоваться вместе с nopython = True. На данный момент он работает только на ЦП.

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

from numba import jit, int32
@jit(int32(int32, int32))
def function(a, b):
    # your loop or numerically intensive computations
    return result
# or if you haven't imported type names
# you can pass them as string
@jit('int32(int32, int32)')
def function(a, b):
    # your loop or numerically intensive computations
    return result

Теперь ваша функция будет принимать только два int32 и возвращать int32. Благодаря этому у вас будет больше контроля над своими функциями. Вы даже можете передать несколько функциональных подписей, если хотите.

Вы также можете использовать другие обертки, предоставляемые numba:

  1. @Vectorize: позволяет использовать скалярные аргументы как numpy ufuncs,
  2. @Guvectorize: производит обобщенные ufuncs NumPy,
  3. @Stencil: объявить функцию как ядро ​​для операции, подобной трафарету,
  4. @Jitclass: для классов, поддерживающих jit,
  5. @Cfunc: объявить функцию для использования в качестве собственного обратного вызова (для вызова из C / C ++ и т. Д.),
  6. @Overload: зарегистрируйте собственную реализацию функции для использования в режиме nopython, например @overload(scipy.special.j0).

Numba также имеет компиляцию Ahead of time (AOT), которая создает скомпилированный модуль расширения, не зависящий от Numba. Но:

  1. Он позволяет использовать только обычные функции (не ufuncs),
  2. Вы должны указать подпись функции. Вы можете указать только один, многие указываются под разными именами.

Он также создает общий код для архитектурного семейства вашего процессора.

5. Обертка @vectorize

Используя оболочку @vectorize, вы можете преобразовать свои функции, которые работают только со скалярами, например, если вы используете библиотеку python math, которая работает только со скалярами, для работы с массивами. Это дает скорость, аналогичную скорости операций с множеством массивов (ufuncs). Например:

@vectorize
def func(a, b):
    # Some operation on scalars
    return result

Вы также можете передать этой оболочке аргумент target, который может иметь значение, равное parallel для распараллеливания кода, cuda для запуска кода на cuda / GPU.

@vectorize(target="parallel")
def func(a, b):
    # Some operation on scalars
    return result

Векторизация с помощью target = “parallel” или “cuda” обычно выполняется быстрее, чем реализация numpy, если ваш код достаточно ресурсоемкий или массив достаточно велик. В противном случае это приводит к накладным расходам времени на создание потоков и разделение элементов для разных потоков, что может быть больше, чем фактическое время вычислений для всего процесса. Итак, работа должна быть достаточно тяжелой, чтобы получить ускорение.

В этом отличном видео есть пример ускорения уравнения Навье-Стокса для вычислительной гидродинамики с помощью Numba:

6. Запуск ваших функций на GPU

Вы также можете передавать @jit как обертки для запуска функций на cuda / GPU. Для этого вам нужно будет импортировать cuda из библиотеки numba. Но запустить ваш код на GPU будет не так просто, как раньше. В нем есть некоторые начальные вычисления, которые необходимо выполнить для запуска функции в сотнях или даже тысячах потоков на GPU. Вы должны объявить и управлять иерархией сеток, блоков и потоков. И это не так уж и сложно.

Чтобы выполнить функцию на графическом процессоре, вы должны определить что-то под названием kernel function или device function. Во-первых, давайте посмотрим на kernel function.

Некоторые моменты, которые следует помнить о функциях ядра:

а) ядра явно объявляют свою иерархию потоков при вызове, то есть количество блоков и количество потоков на блок. Вы можете скомпилировать ядро ​​один раз и вызывать его несколько раз с разными размерами блоков и сеток.

б) ядра не могут возвращать значение. Итак, вам придется либо внести изменения в исходный массив, либо передать другой массив для хранения результата. Для вычисления скаляра вам нужно будет передать массив из 1 элемента.

# Defining a kernel function
from numba import cuda
@cuda.jit
def func(a, result):
    # Some cuda related computation, then
    # your computationally intensive code.
    # (Your answer is stored in 'result')

Итак, для запуска ядра вам нужно будет выполнить две вещи:

  1. Количество потоков на блок,
  2. Количество блоков.

Например:

threadsperblock = 32
blockspergrid = (array.size + (threadsperblock - 1)) // threadsperblock
func[blockspergrid, threadsperblock](array)

Функция ядра в каждом потоке должна знать, в каком потоке он находится, чтобы знать, за какие элементы массива она отвечает. Numba позволяет легко получить эти позиции элементов одним вызовом.

@cuda.jit
def func(a, result):
    pos = cuda.grid(1)  # For 1D array
    # x, y = cuda.grid(2) # For 2D array
    if pos < a.shape[0]:
        result[pos] = a[pos] * (some computation)

Чтобы сэкономить время, которое будет потрачено на копирование массива numpy на определенное устройство с последующим сохранением результата в массиве numpy, Numba предоставляет некоторые функции для объявления и отправки массивов на конкретное устройство, например: numba.cuda.device_array, numba.cuda.device_array_like, numba.cuda.to_device и т. Д. чтобы сэкономить время ненужных копий в процессор (если это не нужно).

С другой стороны, device function может быть вызван только изнутри устройства (ядром или другой функцией устройства). Плюс в том, что вы можете вернуть значение из device function. Итак, вы можете использовать это возвращаемое значение функции для вычисления чего-либо внутри kernel function или device function.

from numba import cuda
@cuda.jit(device=True)
def device_function(a, b):
    return a + b

Вам также следует изучить поддерживаемые функции библиотеки Numba cuda, здесь.

Numba также имеет реализации атомарных операций, генераторов случайных чисел, реализации разделяемой памяти (для ускорения доступа к данным) и т. Д. В своей библиотеке cuda.

Совместимость ctypes / cffi / cython:

  • cffi - Вызов функций CFFI поддерживается в режиме nopython.
  • ctypes - Вызов функций оболочки ctypes поддерживается в режиме nopython…
  • Экспортированные функции Cython вызываемы.

7. Дополнительная литература

  1. Https://nbviewer.jupyter.org/github/ContinuumIO/gtc2017-numba/tree/master/
  2. Https://devblogs.nvidia.com/seven-things-numba/
  3. Https://devblogs.nvidia.com/numba-python-cuda-acceleration/
  4. Https://jakevdp.github.io/blog/2015/02/24/optimizing-python-with-numpy-and-numba/
  5. Https://www.youtube.com/watch?v=1AwG0T4gaO0

8. Ссылки

  1. Http://numba.pydata.org/numba-doc/latest/user/index.html
  2. Https://github.com/ContinuumIO/gtc2018-numba
  3. Http://stephanhoyer.com/2015/04/09/numba-vs-cython-how-to-choose/
Suggestions and reviews are welcome.
Thank you for reading!

Подпись: