Ускорьте время выполнения с помощью молниеносного кода Mojo
Делиться
Для специалистов по данным, инженеров машинного обучения и разработчиков программного обеспечения оптимизация производительности каждого бита кода может быть критически важным фактором. Если вы используете Python, вы наверняка знаете о некоторых его недостатках в этом отношении. Python считается медленным языком, и вы, вероятно, слышали, что во многом это связано с механизмом глобальной блокировки интерпретатора (GIL).
Что есть, то есть, но что с этим делать? Есть несколько способов решить эту проблему при программировании на Python, особенно если вы используете достаточно актуальную версию Python.
- В самых последних версиях Python реализована возможность запуска кода без использования GIL.
- Для выполнения обработки чисел мы можем использовать высокопроизводительные сторонние библиотеки, такие как NumPy.
- В настоящее время в язык встроено множество методов параллельной и одновременной обработки.
Ещё один метод, который мы можем использовать, — это вызов других высокопроизводительных языков из Python для критичных по времени фрагментов кода. Именно это мы и рассмотрим в этой статье, где я покажу, как вызывать код Mojo из Python.
Вы когда-нибудь слышали о Mojo? Если нет, вот небольшой экскурс в историю.
Mojo — относительно новый язык системного уровня, разработанный компанией Modular Inc. (компания, занимающаяся инфраструктурой искусственного интеллекта, основанная в 2022 году легендой разработки компиляторов Крисом Латтнером, известным создателем LLVM и Swift, и бывшим руководителем Google TPU Тимом Дэвисом ) и впервые представленный публике в мае 2023 года .
Он родился из простой проблемы: низкой производительности Python, о которой мы говорили ранее. Mojo решает эту проблему напрямую, внедряя надмножество синтаксиса Python в конвейер компилятора на основе LLVM/MLIR, который обеспечивает абстракции с нулевой стоимостью, статическую типизацию, управление памятью на основе владения, автоматическую векторизацию и бесшовную генерацию кода для ускорителей ЦП и ГП.
Ранние тесты, продемонстрированные при запуске, выполняли рабочие нагрузки с высокой плотностью ядра в 35 000 раз быстрее, чем обычный Python, доказывая, что Mojo может сравниться — или превзойти — производительность C/CUDA, позволяя разработчикам оставаться на знакомой «питоновской» территории.
Однако всегда есть камень преткновения — инертность людей, нежелание полностью переходить на новый язык. Я тоже один из таких людей, поэтому был рад, когда узнал, что несколько недель назад появилась возможность вызывать код Mojo напрямую из Python.
Означает ли это, что мы получаем лучшее из обоих миров: простоту Python и производительность Mojo?
Для проверки этих утверждений мы напишем код на чистом Python. Затем для каждого из них мы также напишем версию на NumPy и, наконец, версию на Python, которая перенесёт часть вычислений в модуль Mojo. В заключение мы сравним время выполнения различных функций.
Увидим ли мы значительный прирост производительности? Читайте дальше, чтобы узнать.
Настройка среды разработки
Для разработки я буду использовать WSL2 Ubuntu для Windows. Рекомендуется настраивать новую среду разработки для каждого проекта. Обычно я использую Conda, но, поскольку все, включая их бабушек, переходят на новый менеджер пакетов UV , я попробую его. Установить UV можно несколькими способами.
$ curl -LsSf https://astral.sh/uv/install.sh | sh или… $ pip install uv
Далее инициализируйте проект.
$ uv init mojo-test $ cd mojo-test $ uv venv $ source .venv/bin/activate Инициализированный проект `mojo-test` в `/home/tom/projects/mojo-test` (mojo-test) $ cd mojo-test (mojo-test) $ ls -al итого 28 drwxr-xr-x 3 tom tom 4096 27 июня 09:20 . drwxr-xr-x 15 том том 4096 27 июня 09:20 .. drwxr-xr-x 7 том том 4096 27 июня 09:20 .git -rw-r—r— 1 том том 109 27 июня 09:20 .gitignore -rw-r—r— 1 том том 5 27 июня 09:20 .python-version -rw-r—r— 1 том том 0 27 июня 09:20 README.md -rw-r—r— 1 том том 87 27 июня 09:20 main.py -rw-r—r— 1 том том 155 27 июня 09:20 pyproject.toml
Теперь добавьте все необходимые нам внешние библиотеки.
(mojo-test) $ uv pip install modular numpy matplotlib
Как работает вызов Mojo из Python?
Предположим, у нас есть следующая простая функция Mojo, которая принимает переменную Python в качестве аргумента и прибавляет к её значению два. Например,
# mojo_func.mojo # fn add_two(py_obj: PythonObject) вызывает -> Python var n = Int(py_obj) возвращает n + 2
Когда Python пытается загрузить add_two, он ищет функцию PyInit_add_two() . В PyInit_add_two() нам нужно объявить все функции и типы Mojo, которые можно вызвать из Python с помощью библиотеки PythonModuleBuilder . Таким образом, наш код Mojo в окончательном виде будет выглядеть примерно так.
из python import PythonObject из python.bindings import PythonModuleBuilder из os import abort @export fn PyInit_mojo_module() -> PythonObject: try: var m = PythonModuleBuilder(«mojo_func») m.def_function[add_two](«add_two», docstring=»Добавить 2 к n») return m.finalize() except e: return abort[PythonObject](String(«Ошибка при создании модуля Python Mojo:», e)) fn add_two(py_obj: PythonObject) raises -> PythonObject: var n = Int(py_obj) n + 2
Для корректной работы кода Python требуется дополнительный шаблонный код, как показано здесь.
import max.mojo.importer import sys sys.path.insert(0, «») import mojo_func print(mojo_func.add_two(5)) # Должен вывести 7
Примеры кода
Для каждого из моих примеров я покажу три разные версии кода. Одна будет написана на чистом Python, вторая будет использовать NumPy для ускорения, а третья будет заменять вызовы Mojo там, где это необходимо.
Обратите внимание, что вызов кода Mojo из Python находится на ранней стадии разработки. Ожидаются значительные изменения в API и эргономике.
Пример 1 — Вычисление множества Мандельброта
В нашем первом примере мы вычислим и выведем на экран множество Мандельброта. Это довольно затратный с точки зрения вычислений процесс, и, как мы увидим, реализация на чистом Python занимает значительное время.
Для этого примера нам понадобится всего четыре файла.
1/ mandelbrot_pure_py.py
# mandelbrot_pure_py.py def compute(width, height, max_iters): «»»Генерирует изображение множества Мандельброта с помощью чистого Python.»»» image = [[0] * width for _ in range(height)] for row in range(height): for col in range(width): c = complex(-2.0 + 3.0 * col / width, -1.5 + 3.0 * row / height) z = 0 n = 0 while abs(z) <= 2 and n < max_iters: z = z*z + cn += 1 image[row][col] = n return image
2/ mandelbrot_numpy.py
# mandelbrot_numpy.py import numpy as np def compute(width, height, max_iters): «»»Генерирует множество Мандельброта с помощью NumPy для векторизованных вычислений.»»» x = np.linspace(-2.0, 1.0, width) y = np.linspace(-1.5, 1.5, height) c = x[:, np.newaxis] + 1j * y[np.newaxis, :] z = np.zeros_like(c, dtype=np.complex128) image = np.zeros(c.shape, dtype=int) for n in range(max_iters): not_diverged = np.abs(z) <= 2 image[not_diverged] = nz[not_diverged] = z[not_diverged]**2 + c[not_diverged] image[np.abs(z) <= 2] = max_iters return image.T
3/ mandelbrot_mojo.mojo
# mandelbrot_mojo.mojo from python import PythonObject, Python from python.bindings import PythonModuleBuilder from os import abort from complex import ComplexFloat64 # Это основная логика, которая будет быстро работать в Mojo fn compute_mandel_pixel(c: ComplexFloat64, max_iters: Int) -> Int: var z = ComplexFloat64(0, 0) var n: Int = 0 while n < max_iters: # abs(z) > 2 то же самое, что и z.norm() > 4, что быстрее, если z.norm() > 4.0: break z = z * z + cn += 1 return n # Это функция, которую вызовет Python fn mandelbrot_mojo_compute(width_obj: PythonObject, height_obj: PythonObject, max_iters_obj: PythonObject) raises -> PythonObject: var width = Int(width_obj) var height = Int(height_obj) var max_iters = Int(max_iters_obj) # Мы построим список Python в Mojo для возврата результатов var image_list = Python.list() for row in range(height): # Мы создаем вложенный список для представления 2D-изображения var row_list = Python.list() for col in range(width): var c = ComplexFloat64( -2.0 + 3.0 * col / width, -1.5 + 3.0 * row / height ) var n = compute_mandel_pixel(c, max_iters) row_list.append(n) image_list.append(row_list) return image_list # Это специальная функция, которая «экспортирует» нашу функцию Mojo в Python @export fn PyInit_mandelbrot_mojo() -> PythonObject: try: var m = PythonModuleBuilder(«mandelbrot_mojo») m.def_function[mandelbrot_mojo_compute](«compute», «Генерирует множество Мандельброта.») return m.finalize() except e: return abort[PythonObject](«Ошибка создания модуля mandelbrot_mojo»)
4/ main.py
Это вызовет остальные три программы, а также позволит нам построить график Мандельброта в Jupyter Notebook. Я покажу график только один раз. Поверьте мне на слово, он был построен правильно во всех трёх запусках кода.
# main.py (финальная версия с визуализацией) import time import numpy as np import sys import matplotlib.pyplot as plt # Теперь импортируем pyplot # — Настройка Mojo — try: import max.mojo.importer except ImportError: print(«Импортер Mojo не найден. Убедитесь, что MODULAR_HOME и PATH заданы правильно.») sys.exit(1) sys.path.insert(0, «») # — Импортируем наши модули — import mandelbrot_pure_py import mandelbrot_numpy import mandelbrot_mojo # — Функция визуализации — def visualize_mandelbrot(image_data, title=»Набор Мандельброта»): «»»Отображает данные набора Мандельброта в виде изображения с помощью Matplotlib.»»» print(f»Отображение изображения для: {title}») plt.figure(figsize=(10, 8)) # 'hot', 'inferno' и 'plasma' — отличные цветовые карты для этого plt.imshow(image_data, cmap='hot', interpolation='bicubic') plt.colorbar(label=»Iterations») plt.title(title) plt.xlabel(«Width») plt.ylabel(«Height») plt.show() # — Test Runner — def run_test(name, compute_func, *args): «»»Вспомогательная функция для запуска и измерения времени теста.»»» print(f»Running {name} version…») start_time = time.time() # Вычислительная функция возвращает данные изображения result_data = compute_func(*args) duration = time.time() — start_time print(f»-> {name} version took: {duration:.4f} seconds») # Возвращает данные для их визуализации return result_data if __name__ == «__main__»: WIDTH, HEIGHT, MAX_ITERS = 800, 600, 5000 print(«Запуск сравнения производительности Mandelbrot…») print(«-» * 40) # Запуск теста Pure Python py_image = run_test(«Pure Python», mandelbrot_pure_py.compute, WIDTH, HEIGHT, MAX_ITERS) visualize_mandelbrot(py_image, «Pure Python Mandelbrot») print(«-» * 40) # Запуск теста NumPy np_image = run_test(«NumPy», mandelbrot_numpy.compute, WIDTH, HEIGHT, MAX_ITERS) # Раскомментируйте строку ниже, если хотите увидеть график #visualize_mandelbrot(np_image, «NumPy Mandelbrot») print(«-» * 40) # Запуск теста Mojo mojo_list_of_lists = run_test(«Mojo», mandelbrot_mojo.compute, WIDTH, HEIGHT, MAX_ITERS) # Преобразуем список списков Mojo в массив NumPy для визуализации mojo_image = np.array(mojo_list_of_lists) # раскомментируйте строку ниже, если хотите увидеть график #visualize_mandelbrot(mojo_image, «Mojo Mandelbrot») print(«-» * 40) print(«Сравнение завершено.»)
Наконец, вот что получилось.
Итак, это впечатляющее начало для Mojo. Он оказался почти в 20 раз быстрее реализации на чистом Python и в 5 раз быстрее кода на NumPy.
Пример 2 — Численное интегрирование
В этом примере мы выполним численное интегрирование, используя правило Симпсона, чтобы определить значение sin(X) в интервале от 0 до π. Напомним, что правило Симпсона — это метод вычисления приближённого значения интеграла и определяется как:
∫ ≈ (h/3) * [f(x₀) + 4f(x₁) + 2f(x₂) + 4f(x₃) + … + 2f(xₙ-₂) + 4f(xₙ-₁) + f(xₙ)]
Где:
- h — ширина каждой ступеньки.
- Веса: 1, 4, 2, 4, 2, …, 4, 1.
- Первая и последняя точки имеют вес 1 .
- Точки с нечетными индексами имеют вес 4 .
- Точки с четными индексами имеют вес 2 .
Истинное значение интеграла, который мы пытаемся вычислить, равно двум . Давайте проверим, насколько точны (и быстры) наши методы.
Нам снова понадобятся четыре файла.
1/integration_pure_py.py
# integration_pure_py.py import math def compute(start, end, n): «»»Вычисляет определенный интеграл sin(x) с использованием правила Симпсона.»»» if n % 2 != 0: n += 1 # Правило Симпсона требует четного количества интервалов h = (end — start) / n integral = math.sin(start) + math.sin(end) for i in range(1, n, 2): integral += 4 * math.sin(start + i * h) for i in range(2, n, 2): integral += 2 * math.sin(start + i * h) integral *= h / 3 return integral
2/integration_numpy
# integration_numpy.py import numpy as np def compute(start, end, n): «»»Вычисляет определенный интеграл sin(x) с помощью NumPy.»»» if n % 2 != 0: n += 1 x = np.linspace(start, end, n + 1) y = np.sin(x) # Применяем веса правила Симпсона: 1, 4, 2, 4, …, 2, 4, 1 integral = (y[0] + y[-1] + 4 * np.sum(y[1:-1:2]) + 2 * np.sum(y[2:-1:2])) h = (end — start) / n
3/integration_mojo.mojo
# integration_mojo.mojo from python import PythonObject, Python from python.bindings import PythonModuleBuilder from os import abort from math import sin # Примечание: здесь используется ключевое слово 'fn', так как оно совместимо со всеми версиями. fn compute_integral_mojo(start_obj: PythonObject, end_obj: PythonObject, n_obj: PythonObject) raises -> PythonObject: # Пересечение моста происходит ОДИН РАЗ в начале. var start = Float64(start_obj) var end = Float64(end_obj) var n = Int(n_obj) if n % 2 != 0: n += 1 var h = (end — start) / n # Все вычисления ниже выполняются над НАТИВНЫМИ типами Mojo. Взаимодействие с Python отсутствует. var integral = sin(start) + sin(end) # Первый цикл для членов '4 * f(x)' var i_1: Int = 1 while i_1 < n: integral += 4 * sin(start + i_1 * h) i_1 += 2 # Второй цикл для членов '2 * f(x)' var i_2: Int = 2 while i_2 < n: integral += 2 * sin(start + i_2 * h) i_2 += 2 integral *= h / 3 # Пересечение моста происходит ОДИН РАЗ в конце. return Python.float(integral) @export fn PyInit_integration_mojo() -> PythonObject: try: var m = PythonModuleBuilder(«integration_mojo») m.def_function[compute_integral_mojo](«compute», «Вычисляет определенный интеграл в Mojo.») return m.finalize() except e: return abort[PythonObject](«ошибка создания модуля integration_mojo»)
4/ main.py
import time import sys import numpy as np # — Настройка Mojo — try: import max.mojo.importer except ImportError: print(«Импортер Mojo не найден. Убедитесь, что ваша среда настроена правильно.») sys.exit(1) sys.path.insert(0, «») # — Импорт наших модулей — import integration_pure_py import integration_numpy import integration_mojo # — Исполнитель тестов — def run_test(name, compute_func, *args): print(f»Запуск версии {name}…») start_time = time.time() result = compute_func(*args) duration = time.time() — start_time print(f»-> {name} version took: {duration:.4f} seconds») print(f» Result: {result}») # — Выполнение основного теста — if __name__ == «__main__»: # Используйте очень большое количество шагов для выделения производительность цикла START = 0.0 END = np.pi NUM_STEPS = 100_000_000 # 100 миллионов шагов print(f»Вычисление интеграла sin(x) от {START} до {END:.2f} с {NUM_STEPS:,} шагов…») print(«-» * 50) run_test(«Pure Python», integration_pure_py.compute, START, END, NUM_STEPS) print(«-» * 50) run_test(«NumPy», integration_numpy.compute, START, END, NUM_STEPS) print(«-» * 50) run_test(«Mojo», integration_mojo.compute, START, END, NUM_STEPS) print(«-» * 50) print(«Сравнение завершено.»)
И каковы результаты?
Вычисление интеграла sin(x) от 0,0 до 3,14 за 100 000 000 шагов… ————————————————— Запуск версии Pure Python… -> Версия Pure Python заняла: 4,9484 секунды Результат: 2,00000000000000346 ————————————————— Запуск версии NumPy… -> Версия NumPy заняла: 0,7425 секунды Результат: 1,9999999999999998 ————————————————— Запуск версии Mojo… -> Версия Mojo заняла: 0,8902 секунды Результат: 2,0000000000000346 ————————————————— Сравнение завершено.
Интересно, что на этот раз код NumPy оказался немного быстрее кода Mojo, а его итоговое значение оказалось точнее. Это подчёркивает ключевую концепцию высокопроизводительных вычислений: компромисс между векторизацией и JIT-компилированными циклами .
Преимущество NumPy заключается в его способности векторизовать операции. Он выделяет большой блок памяти, а затем вызывает высокооптимизированный, предварительно скомпилированный код на C, использующий современные возможности процессоров, такие как SIMD, для одновременного выполнения функции sin() над миллионами значений. Такая «пакетная обработка» невероятно эффективна.
Mojo, с другой стороны, берёт наш простой цикл while и JIT-компилирует его в высокоэффективный машинный код. Хотя это позволяет избежать большого начального выделения памяти, характерного для NumPy, в данном конкретном случае мощь векторизации NumPy даёт ему небольшое преимущество.
Пример 3 — Сигмовидная функция
Сигмоидальная функция — важная концепция в ИИ, поскольку она является краеугольным камнем двоичной классификации.
Также известная как логистическая функция, она определяется следующим образом.
Сигмоидальная функция принимает любой действительный входной сигнал x и плавно «сжимает» его в открытый интервал (0, 1). Проще говоря, независимо от того, что передано сигмоидальной функции, она всегда вернёт значение в диапазоне от 0 до 1.
Так, например,
S(-197865) = 0 S(-2) = 0,0009111 S(3) = 0,9525741 S(10776,87) = 1
Это делает его идеальным для представления определенных вещей, таких как вероятности.
Поскольку код Python проще, мы можем включить его в скрипт сравнительного анализа, поэтому на этот раз у нас будет только два файла.
sigmoid_mojo.mojo
from python import Python, PythonObject from python.bindings import PythonModuleBuilder from os import abort from math import exp from time import perf_counter # ———————————————————————- # Быстрая процедура Mojo (без вызовов Python внутри) # ———————————————————————- fn sigmoid_sum(n: Int) -> (Float64, Float64): # детерминированное заполнение, размер вычисляется один раз var data = List[Float64](length = n, fill = 0.0) for i in range(n): data[i] = (Float64(i) / Float64(n)) * 10.0 — 5.0 # [-5, +5] var t0: Float64 = perf_counter() var total: Float64 = 0.0 for x in data: # одинарный плотный цикл total += 1.0 / (1.0 + exp(-x)) var elapsed: Float64 = perf_counter() — t0 return (total, elapsed) # ———————————————————————- # Видимая Python обертка # ———————————————————————- fn py_sigmoid_sum(n_obj: PythonObject) raises -> PythonObject: var n: Int = Int(n_obj) # проверяет аргумент var (tot, secs) = sigmoid_sum(n) # Самый безопасный контейнер: создает список Python (автоматически упаковывает скаляры) var out = Python.list() out.append(tot) out.append(secs) return out # -> PythonObject # ———————————————————————- # Инициализатор модуля (имя должно совпадать: PyInit_sigmoid_mojo) # ———————————————————————- @export fn PyInit_sigmoid_mojo() -> PythonObject: try: var m = PythonModuleBuilder(«sigmoid_mojo») m.def_function[py_sigmoid_sum]( «sigmoid_sum», «Return [total_sigmoid, elapsed_seconds]» ) return m.finalize() except e: # если что-то возникает, выдаем Python настоящую ImportError return abort[PythonObject](«error Creating Sigmoid_mojo module»)
main.py
# bench_sigmoid.py import time, math, numpy as np N = 50_000_000 # ————————— pure-Python ———————————- py_data = [(i / N) * 10.0 — 5.0 for i in range(N)] t0 = time.perf_counter() py_total = sum(1 / (1 + math.exp(-x)) for x in py_data) print(f»Pure-Python : {time.perf_counter()-t0:6.3f} s — Σσ={py_total:,.1f}») # ————————— NumPy —————————————— np_data = np.linspace(-5.0, 5.0, N, dtype=np.float64) t0 = time.perf_counter() np_total = float(np.sum(1 / (1 + np.exp(-np_data)))) print(f»NumPy : {time.perf_counter()-t0:6.3f} s — Σσ={np_total:,.1f}») # ————————— Mojo —————————————— import max.mojo.importer # устанавливает .mojo import hook import sigmoid_mojo # компилирует и загружает общий объект mj_total, mj_secs = sigmoid_mojo.sigmoid_sum(N) print(f»Mojo : {mj_secs:6.3f} s — Σσ={mj_total:,.1f}»)
Вот что получилось.
$ python sigmoid_bench.py Pure-Python: 1,847 с — Σσ=24 999 999,5 NumPy: 0,323 с — Σσ=25 000 000,0 Mojo: 0,150 с — Σσ=24 999 999,5
Выходные данные Σσ=… показывают сумму всех вычисленных значений сигмоиды. Теоретически она должна быть точно равна входному значению N, делённому на 2, поскольку N стремится к бесконечности.
Но, как мы видим, реализация mojo обеспечивает прирост производительности более чем в 2 раза по сравнению с и без того быстрым кодом NumPy и более чем в 12 раз быстрее базовой реализации Python.
Неплохо.
Краткое содержание
В этой статье рассматривается новая захватывающая возможность вызова высокопроизводительного кода Mojo непосредственно из Python для ускорения ресурсоёмких вычислительных задач. Mojo, относительно новый язык системного программирования от Modular, обещает производительность уровня C с использованием привычного синтаксиса Python, стремясь преодолеть исторические ограничения скорости Python.
Чтобы проверить это обещание, мы провели сравнительный анализ трех сценариев, требующих больших вычислительных затрат: генерация множества Мандельброта, численное интегрирование и функция вычисления сигмоиды, реализовав каждый из них на чистом Python, оптимизированном NumPy и гибридном подходе Python-Mojo.
Результаты демонстрируют нюансы производительности для алгоритмов с большим количеством циклов, где данные могут обрабатываться исключительно с использованием нативных типов Mojo. Mojo может значительно превосходить как чистый Python, так и даже высокооптимизированный код NumPy. Однако мы также увидели, что в задачах, идеально соответствующих векторизованным, предварительно скомпилированным функциям C NumPy, NumPy может сохранять небольшое преимущество перед Mojo.
Это исследование показывает, что хотя Mojo является мощным новым инструментом для ускорения Python, достижение максимальной производительности требует продуманного подхода к минимизации накладных расходов на «переход по мосту» между двумя языковыми средами выполнения.
Как всегда, при рассмотрении возможностей повышения производительности кода, тестируйте, тестируйте и ещё раз тестируйте. Это станет окончательным критерием того, стоит ли это делать или нет.
Источник: towardsdatascience.com
