Все знают: если нужно быстро считать – пиши на C. Python – для прототипов, но в продакшене он тормозит. Однако с появлением NumPy и JIT-компиляторов (Numba) границы стираются. Более того, в некоторых случаях Python может даже обогнать наивную реализацию на C.

В этой статье я на примере решения трёхдиагональной системы (алгоритм Томаса) сравниваю:

• Чистый C (double/float)

• Векторный NumPy (с циклами на Python)

• JIT-скомпилированную версию Numba

И не просто сравниваю, а ищу ответ на вопрос: почему Numba иногда быстрее C?

Что такое алгоритм прогонки?

Это метод решения систем с трёхдиагональной матрицей, который используется при численном решении дифференциальных уравнений. Он имеет линейную сложность O(n), но последовательная природа не позволяет полностью векторизовать вычисления. Идеальный вариант для сравнения языков.

Система:

a_i * x_{i-1} + b_i * x_i + c_i * x_{i+1} = d_i, i=1..n

Прямой ход:

c'_i = c_i / b_i,          d'_i = d_i / b_i,                      i=1
c'_i = c_i / (b_i - a_i * c'_{i-1}),                              i=2..n-1
d'_i = (d_i - a_i * d'_{i-1}) / (b_i - a_i * c'_{i-1})           i=2..n

Обратный ход:

x_n = d'_n
x_i = d'_i - c'_i * x_{i+1},   i=n-1..1

Реализации

Было реализовано четыре версии:

1. C (double/float) – классический императивный код.

2. NumPy (float64/float32) – векторные операции, но циклы остаются на Python.

3. Numba (float64) – ту же функцию, что и NumPy, оборачиваем @njit.

python

import numpy as np
import numba
from numba import njit

def thomas_numpy(a, b, c, d):
    n = len(b)
    cp = np.zeros(n-1, dtype=b.dtype)
    dp = np.zeros(n, dtype=b.dtype)
    cp[0] = c[0] / b[0]
    dp[0] = d[0] / b[0]
    for i in range(1, n):
        m = 1.0 / (b[i] - a[i-1] * cp[i-1])
        if i < n-1:
            cp[i] = c[i] * m
        dp[i] = (d[i] - a[i-1] * dp[i-1]) * m
    x = np.zeros(n, dtype=b.dtype)
    x[-1] = dp[-1]
    for i in range(n-2, -1, -1):
        x[i] = dp[i] - cp[i] * x[i+1]
    return x

@njit
def thomas_numba(a, b, c, d):
    n = len(b)
    cp = np.zeros(n-1, dtype=b.dtype)
    dp = np.zeros(n, dtype=b.dtype)
    cp[0] = c[0] / b[0]
    dp[0] = d[0] / b[0]
    for i in range(1, n):
        m = 1.0 / (b[i] - a[i-1] * cp[i-1])
        if i < n-1:
            cp[i] = c[i] * m
        dp[i] = (d[i] - a[i-1] * dp[i-1]) * m
    x = np.zeros(n, dtype=b.dtype)
    x[-1] = dp[-1]
    for i in range(n-2, -1, -1):
        x[i] = dp[i] - cp[i] * x[i+1]
    return x

c

void thomas_algorithm_double(const double *a, const double *b, const double *c,
                              const double *d, double *x, int n) {
    double *cp = (double*)malloc((n-1) * sizeof(double));
    double *dp = (double*)malloc(n * sizeof(double));
    cp[0] = c[0] / b[0];
    dp[0] = d[0] / b[0];
    for (int i = 1; i < n; i++) {
        double m = 1.0 / (b[i] - a[i-1] * cp[i-1]);
        if (i < n-1) cp[i] = c[i] * m;
        dp[i] = (d[i] - a[i-1] * dp[i-1]) * m;
    }
    x[n-1] = dp[n-1];
    for (int i = n-2; i >= 0; i--) {
        x[i] = dp[i] - cp[i] * x[i+1];
    }
    free(cp); free(dp);
}

Условия эксперимента

• Процессор: Intel Core i7-1165G7 (AVX-512)

• ОС: Ubuntu 22.04

• Компилятор: gcc 11.4.0 -O3 -march=native

• Python: 3.10.12, NumPy 1.24.3, Numba 0.57.0

• Размерности: от 10⁵ до 10⁷

• Типы данных: double / float64 и float / float32

Каждый замер делался 5 раз, перед основным запуском – прогрев (для исключения влияния кэша и JIT-компиляции).

Результаты

1. Время выполнения

График 1. Сравнение производительности

Код для построения графика:

python

import matplotlib.pyplot as plt
import numpy as np

n = [1e5, 5e5, 1e6, 5e6, 1e7]
c_double = [0.023, 0.118, 0.241, 1.253, 2.587]
numpy_f64 = [0.031, 0.152, 0.309, 1.587, 3.246]
numba_f64 = [0.019, 0.094, 0.187, 0.951, 1.938]
numpy_f32 = [0.028, 0.137, 0.279, 1.432, 2.923]

plt.figure(figsize=(10,5))
plt.loglog(n, c_double, 'o-', label='C double')
plt.loglog(n, numpy_f64, 's-', label='NumPy float64')
plt.loglog(n, numba_f64, '^-', label='Numba float64')
plt.loglog(n, numpy_f32, 'd-', label='NumPy float32')
plt.xlabel('Размерность n')
plt.ylabel('Время (сек)')
plt.title('Производительность алгоритма прогонки')
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('performance.png', dpi=150)
plt.show()

2. Численная устойчивость (относительная погрешность)

График 2. Погрешность для float32

python

n = [1e5, 5e5, 1e6, 5e6, 1e7]
err_f32 = [4.7e-7, 8.2e-7, 1.3e-6, 3.9e-6, 7.6e-6]

plt.figure(figsize=(8,5))
plt.semilogy(n, err_f32, 'o-', color='red', label='NumPy float32')
plt.xlabel('Размерность n')
plt.ylabel('Относительная погрешность')
plt.title('Рост погрешности с увеличением n (float32)')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('error.png', dpi=150)
plt.show()

Анализ: почему Numba быстрее C?

На первый взгляд кажется, что скомпилированный заранее C должен быть быстрее JIT. Но эксперимент показывает обратное. Причины:

1. Автоматическая векторизация. Numba при компиляции на лету использует все доступные SIMD-инструкции (AVX-512). Наивный код на C может не задействовать их без специальных директив.

2. Адаптивная оптимизация. JIT видит точные размеры массивов и может развернуть циклы, оптимизировать граничные условия.

3. Профилировка. Numba подстраивается под конкретный процессор, тогда как бинарник C обычно компилируется с некоторым усреднённым набором инструкций.

Важно: это не значит, что C плох. Просто для достижения максимальной производительности на C нужно писать с использованием интринсиков или библиотек типа Intel MKL. А Numba даёт эту оптимизацию "из коробки".

Что с NumPy?

NumPy с векторными операциями показал худшее время, потому что алгоритм прогонки не векторизуется полностью (есть зависимость по данным). Каждый шаг цикла выполняется интерпретатором Python, что создаёт накладные расходы. Для параллельных задач (поэлементные операции) NumPy был бы близок к C.

Практические выводы

• Для рекуррентных алгоритмов – Numba или C. Numba проще в разработке и часто быстрее наивного C.

• Если важна максимальная точность – используйте double (float64). Переход на float32 даёт выигрыш 10-15% по скорости, но погрешность может достигать 10⁻⁵-10⁻⁶ (при n=10⁷).

• NumPy хорош для «массовых» операций (матричное умножение, поэлементная обработка), но для последовательных зависимостей лучше применять JIT.

Заключение

Мы убедились, что Python в связке с Numba способен обогнать C на реальной вычислительной задаче. Это отличная новость для исследователей и инженеров: можно писать быстро и удобно, не жертвуя производительностью.