Numba njit делает вызов функции, включающей длинные выражения, очень медленно

Я пишу программу конечного объема для решения невязких, сжимаемых уравнений Эйлера. В рамках этого я выполняю так называемый процесс Коши-Ковалевской. Фрагмент кода приведен ниже

def cauchyKovalevskaya(drho_x: npt.NDArray[np.float64], dq_x: npt.NDArray[np.float64], dE_x: npt.NDArray[np.float64], k: int, gamma: float) -> npt.NDArray:
    if k == 0:
        return np.array([drho_x[0], dq_x[0], dE_x[0]], dtype=np.float64)
    
    elif k == 1:
        rho_t = -dq_x[1]

        q_t = -dq_x[0]**2*drho_x[1]*gamma/(2*drho_x[0]**2) + 3*dq_x[0]**2*drho_x[1]/(2*drho_x[0]**2) + dq_x[0]*dq_x[1]*gamma/drho_x[0] - 3*dq_x[0]*dq_x[1]/drho_x[0] - dE_x[1]*gamma + dE_x[1]

        E_t = -dq_x[0]**3*drho_x[1]*gamma/drho_x[0]**3 + dq_x[0]**3*drho_x[1]/drho_x[0]**3 + 3*dq_x[0]**2*dq_x[1]*gamma/(2*drho_x[0]**2) - 3*dq_x[0]**2*dq_x[1]/(2*drho_x[0]**2) - dq_x[0]*dE_x[1]*gamma/drho_x[0] + dq_x[0]*drho_x[1]*dE_x[0]*gamma/drho_x[0]**2 - dq_x[1]*dE_x[0]*gamma/drho_x[0]

        return np.array([rho_t, q_t, E_t], dtype=np.float64)
    
    elif k == 2:
        

        rho_tt = dq_x[0]**2*drho_x[2]*gamma/(2*drho_x[0]**2) - 3*dq_x[0]**2*drho_x[2]/(2*drho_x[0]**2) - dq_x[0]**2*drho_x[1]**2*gamma/drho_x[0]**3 + 3*dq_x[0]**2*drho_x[1]**2/drho_x[0]**3 + 2*dq_x[0]*dq_x[1]*drho_x[1]*gamma/drho_x[0]**2 - 6*dq_x[0]*dq_x[1]*drho_x[1]/drho_x[0]**2 - dq_x[0]*dq_x[2]*gamma/drho_x[0] + 3*dq_x[0]*dq_x[2]/drho_x[0] - dq_x[1]**2*gamma/drho_x[0] + 3*dq_x[1]**2/drho_x[0] + dE_x[2]*gamma - dE_x[2]

        q_tt = dq_x[0]**3*drho_x[2]*gamma**2/(2*drho_x[0]**3) + dq_x[0]**3*drho_x[2]*gamma/drho_x[0]**3 - 7*dq_x[0]**3*drho_x[2]/(2*drho_x[0]**3) - 3*dq_x[0]**3*drho_x[1]**2*gamma**2/(2*drho_x[0]**4) - 3*dq_x[0]**3*drho_x[1]**2*gamma/drho_x[0]**4 + 21*dq_x[0]**3*drho_x[1]**2/(2*drho_x[0]**4) + 5*dq_x[0]**2*dq_x[1]*drho_x[1]*gamma**2/(2*drho_x[0]**3) + 8*dq_x[0]**2*dq_x[1]*drho_x[1]*gamma/drho_x[0]**3 - 45*dq_x[0]**2*dq_x[1]*drho_x[1]/(2*drho_x[0]**3) - dq_x[0]**2*dq_x[2]*gamma**2/(2*drho_x[0]**2) - 5*dq_x[0]**2*dq_x[2]*gamma/(2*drho_x[0]**2) + 6*dq_x[0]**2*dq_x[2]/drho_x[0]**2 - dq_x[0]*dq_x[1]**2*gamma**2/drho_x[0]**2 - 5*dq_x[0]*dq_x[1]**2*gamma/drho_x[0]**2 + 12*dq_x[0]*dq_x[1]**2/drho_x[0]**2 + 3*dq_x[0]*dE_x[2]*gamma/drho_x[0] - 3*dq_x[0]*dE_x[2]/drho_x[0] - dq_x[0]*drho_x[1]*dE_x[1]*gamma**2/drho_x[0]**2 - 2*dq_x[0]*drho_x[1]*dE_x[1]*gamma/drho_x[0]**2 + 3*dq_x[0]*drho_x[1]*dE_x[1]/drho_x[0]**2 - dq_x[0]*drho_x[2]*dE_x[0]*gamma**2/drho_x[0]**2 + dq_x[0]*drho_x[2]*dE_x[0]*gamma/drho_x[0]**2 + 2*dq_x[0]*drho_x[1]**2*dE_x[0]*gamma**2/drho_x[0]**3 - 2*dq_x[0]*drho_x[1]**2*dE_x[0]*gamma/drho_x[0]**3 + dq_x[1]*dE_x[1]*gamma**2/drho_x[0] + 2*dq_x[1]*dE_x[1]*gamma/drho_x[0] - 3*dq_x[1]*dE_x[1]/drho_x[0] - 2*dq_x[1]*drho_x[1]*dE_x[0]*gamma**2/drho_x[0]**2 + 2*dq_x[1]*drho_x[1]*dE_x[0]*gamma/drho_x[0]**2 + dq_x[2]*dE_x[0]*gamma**2/drho_x[0] - dq_x[2]*dE_x[0]*gamma/drho_x[0]

        E_tt = dq_x[0]**4*drho_x[2]*gamma**2/(4*drho_x[0]**4) + 2*dq_x[0]**4*drho_x[2]*gamma/drho_x[0]**4 - 9*dq_x[0]**4*drho_x[2]/(4*drho_x[0]**4) - dq_x[0]**4*drho_x[1]**2*gamma**2/drho_x[0]**5 - 8*dq_x[0]**4*drho_x[1]**2*gamma/drho_x[0]**5 + 9*dq_x[0]**4*drho_x[1]**2/drho_x[0]**5 + dq_x[0]**3*dq_x[1]*drho_x[1]*gamma**2/drho_x[0]**4 + 37*dq_x[0]**3*dq_x[1]*drho_x[1]*gamma/(2*drho_x[0]**4) - 39*dq_x[0]**3*dq_x[1]*drho_x[1]/(2*drho_x[0]**4) - 7*dq_x[0]**3*dq_x[2]*gamma/(2*drho_x[0]**3) + 7*dq_x[0]**3*dq_x[2]/(2*drho_x[0]**3) - 21*dq_x[0]**2*dq_x[1]**2*gamma/(2*drho_x[0]**3) + 21*dq_x[0]**2*dq_x[1]**2/(2*drho_x[0]**3) - dq_x[0]**2*dE_x[2]*gamma**2/(2*drho_x[0]**2) + 3*dq_x[0]**2*dE_x[2]*gamma/drho_x[0]**2 - 3*dq_x[0]**2*dE_x[2]/(2*drho_x[0]**2) + dq_x[0]**2*drho_x[1]*dE_x[1]*gamma**2/(2*drho_x[0]**3) - 15*dq_x[0]**2*drho_x[1]*dE_x[1]*gamma/(2*drho_x[0]**3) + 3*dq_x[0]**2*drho_x[1]*dE_x[1]/drho_x[0]**3 - dq_x[0]**2*drho_x[2]*dE_x[0]*gamma**2/(2*drho_x[0]**3) - 3*dq_x[0]**2*drho_x[2]*dE_x[0]*gamma/(2*drho_x[0]**3) + 3*dq_x[0]**2*drho_x[1]**2*dE_x[0]*gamma**2/(2*drho_x[0]**4) + 9*dq_x[0]**2*drho_x[1]**2*dE_x[0]*gamma/(2*drho_x[0]**4) - dq_x[0]*dq_x[1]*dE_x[1]*gamma**2/drho_x[0]**2 + 8*dq_x[0]*dq_x[1]*dE_x[1]*gamma/drho_x[0]**2 - 3*dq_x[0]*dq_x[1]*dE_x[1]/drho_x[0]**2 - dq_x[0]*dq_x[1]*drho_x[1]*dE_x[0]*gamma**2/drho_x[0]**3 - 7*dq_x[0]*dq_x[1]*drho_x[1]*dE_x[0]*gamma/drho_x[0]**3 + 2*dq_x[0]*dq_x[2]*dE_x[0]*gamma/drho_x[0]**2 + 2*dq_x[1]**2*dE_x[0]*gamma/drho_x[0]**2 + dE_x[0]*dE_x[2]*gamma**2/drho_x[0] - dE_x[0]*dE_x[2]*gamma/drho_x[0] + dE_x[1]**2*gamma**2/drho_x[0] - dE_x[1]**2*gamma/drho_x[0] - drho_x[1]*dE_x[0]*dE_x[1]*gamma**2/drho_x[0]**2 + drho_x[1]*dE_x[0]*dE_x[1]*gamma/drho_x[0]**2


        return np.array([rho_tt, q_tt, E_tt], dtype=np.float64)


    else:
        raise Exception("Invalid order")

Существуют также утверждения для k = 3, 4 и 5, которые я для краткости исключил. Достаточно сказать, что эти строки становятся все длиннее и длиннее (более 80 000 символов в строке) с множеством подобных обращений к памяти, степеней и умножений, как указано выше.

Когда я использую njit с этим кодом, а) он работает медленнее, чем без njit (для k = 5 разница составляет от 500 до 600 раз), и б) он становится все медленнее и медленнее, чем выше k, для которого я включаю операторы elif, независимо от для которого k я оцениваю функцию. Насколько я могу судить, в функции не должно быть ничего, что ухудшало бы производительность при использовании Numba.

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

Nin17 26.06.2024 13:20

@Nin17 Nin17 этот вызов функции является небольшой частью гораздо большего кода. Я не фиксирую время самого вызова функции, но если бы я мог сказать вам, сколько времени потребуется для выполнения всего кода с njit и без него, для случая, когда k = 5, без njit — около минуты, а с njit — 600 минут.

Siddharth Yajaman 26.06.2024 14:22

Можно ли упростить уравнения, чтобы уменьшить количество операций? Можете ли вы определить часто используемые термины для повторного использования (например, drho_x_inv = 1/drho_x[0], чтобы не вычислять одно и то же значение снова и снова)? Тот факт, что он скомпилирован, не обязательно означает, что вам сойдет с рук написание «наивного» кода (т. е. вам, возможно, придется провести некоторые собственные оптимизации).

jared 26.06.2024 15:49

Компиляторы @jared могут сделать это сами, хотя это помогает Numba генерировать меньше кода для LLVM (который заполняется количеством сгенерированных инструкций) и делать это быстрее (потому что Numba написана на Python, что медленно), только ускоряя сборку в конец, а не время выполнения (особенно с fastmath=True).

Jérôme Richard 26.06.2024 21:03

При всем уважении, (цит.) «(...) строки становятся все длиннее и длиннее (80000+ символов в строке)» на самом деле является самоочевидным антипаттерном неправильного подхода к коду для производительности. Пересмотренный код должен избегать тысяч в основном бесполезных операций, чтобы алгебра выполняла свою работу перед кодированием повторяющихся (дорогих) операций, которые будут происходить много раз, вместо правильно сокращенных формул с умными заранее вычисленными константами. Если бы целью была производительность, код никогда не должен был бы выглядеть так. После более чем 20 лет разработки проектов со сверхнизкой задержкой и высокой пропускной способностью я осмелюсь заявить, что сокращенный код (на FPGA) побеждает.

user3666197 26.06.2024 22:14

Как генерируется этот код? Это результат какой-то генерации кода, может быть, SymPy? В большинстве генераторов кода, таких как Sympy, есть общее исключение подвыражений, позволяющее автоматически избегать таких огромных кодов, если есть общие подвыражения. docs.sympy.org/latest/modules/simplify/…

max9111 27.06.2024 10:39

@jared, да, я уверен, что мог бы сделать какие-то упрощения, и переделаю свой код, чтобы он учитывал эти вещи

Siddharth Yajaman 27.06.2024 12:27

@user3666197 user3666197, в идеале я бы выполнил некоторые предварительные вычисления, как вы предложили, но, например, как только я достигну чего-то вроде k = 5, мне придется предварительно вычислить значительно больше значений. Тем не менее, я обязательно поищу способы сделать код более эффективным.

Siddharth Yajaman 27.06.2024 12:31

@ max9111 max9111, да, это создано с помощью SymPy! Я постараюсь применить исключение общих подвыражений, как вы предлагаете, чтобы сделать код более эффективным.

Siddharth Yajaman 27.06.2024 12:33
Структурированный массив Numpy
Структурированный массив Numpy
Однако в реальных проектах я чаще всего имею дело со списками, состоящими из нескольких типов данных. Как мы можем использовать массивы numpy, чтобы...
1
9
87
2
Перейти к ответу Данный вопрос помечен как решенный

Ответы 2

Ответ принят как подходящий

TL;DR: вам следует избегать предоставления такого большого кода непосредственно в Numba и компилировать вызывающую функцию с помощью Numba. Если вы не можете, единственная возможная оптимизация — это предварительно выделить выходные данные (или фактически не использовать расширение Numba, а C напрямую, без проверки входных данных — что небезопасно).


Оптимизация времени компиляции

Во-первых, скомпилированный код довольно большой, особенно если применить ту же логику «для k = 3, 4 и 5». Это проблема по двум причинам. Во-первых, чем больше компилируемый код, тем медленнее его сборка. Во-вторых, большие скомпилированные коды по-прежнему не оптимальны, поскольку собственный скомпилированный код необходимо извлечь из кэша/ОЗУ, декодировать и затем выполнить. Первые два шага являются дорогостоящими, особенно если код большой. Таким образом, циклы иногда могут работать быстрее из-за двух первых накладных расходов.

Чтобы быть уверенным, что время компиляции не является проблемой в середине вашего теста/приложения (из-за ленивой компиляции), вы можете быстро скомпилировать функцию, предоставив целевую сигнатуру функции. Обратите внимание, что Numba в настоящее время не заботится о типах аннотаций, таких как npt.NDArray[np.float64]. Вот пример подписи:

@nb.njit('(float64[::1], float64[::1], float64[::1], int64, float64)')

На моей машине сборка функций занимает 3,74 секунды. Это очень много для такой простой вычислительной функции, и на разбор и семантический анализ кода Numba Python придется потратить значительное время. Хотя это время оплачивается только один раз, когда функция определена благодаря энергичной компиляции, столь сильно замедлять инициализацию вашей программы для меня нецелесообразно. Главное – избегать повторного выражения. Обратите внимание, что это не сильно влияет на время выполнения скомпилированной функции, поскольку компиляторы (например, Numba) могут анализировать повторяющиеся термины и избегать повторных вычислений. Также обратите внимание, что это делает код более читабельным (избегая некрасивых строк длиной > 1700 символов). Вот немного лучшая версия:

def cauchyKovalevskaya(drho_x: npt.NDArray[np.float64], dq_x: npt.NDArray[np.float64], dE_x: npt.NDArray[np.float64], k: int, gamma: float) -> npt.NDArray:
    g = gamma
    r0, r1, r2 = drho_x
    q0, q1, q2 = dq_x
    e0, e1, e2 = dE_x

    r0p2, r1p2 = r0**2, r1**2
    q0p2, q1p2 = q0**2, q1**2
    e0p2, e1p2 = e0**2, e1**2
    gp2 = g**2
    r0p3, r0p4 = r0**3, r0**4
    q0p3, q0p4 = q0**3, q0**4

    if k == 0:
        return np.array([r0, q0, e0], dtype=np.float64)
    
    elif k == 1:
        rho_t = -q1
        q_t = -q0p2*r1*g/(2*r0p2) + 3*q0p2*r1/(2*r0p2) + q0*q1*g/r0 - 3*q0*q1/r0 - e1*g + e1
        E_t = -q0p3*r1*g/r0p3 + q0p3*r1/r0p3 + 3*q0p2*q1*g/(2*r0p2) - 3*q0p2*q1/(2*r0p2) - q0*e1*g/r0 + q0*r1*e0*g/r0p2 - q1*e0*g/r0
        return np.array([rho_t, q_t, E_t], dtype=np.float64)
    
    elif k == 2:
        rho_tt = q0p2*r2*g/(2*r0p2) - 3*q0p2*r2/(2*r0p2) - q0p2*r1p2*g/r0p3 + 3*q0p2*r1p2/r0p3 + 2*q0*q1*r1*g/r0p2 - 6*q0*q1*r1/r0p2 - q0*q2*g/r0 + 3*q0*q2/r0 - q1p2*g/r0 + 3*q1p2/r0 + e2*g - e2
        q_tt = q0p3*r2*gp2/(2*r0p3) + q0p3*r2*g/r0p3 - 7*q0p3*r2/(2*r0p3) - 3*q0p3*r1p2*gp2/(2*r0p4) - 3*q0p3*r1p2*g/r0p4 + 21*q0p3*r1p2/(2*r0p4) + 5*q0p2*q1*r1*gp2/(2*r0p3) + 8*q0p2*q1*r1*g/r0p3 - 45*q0p2*q1*r1/(2*r0p3) - q0p2*q2*gp2/(2*r0p2) - 5*q0p2*q2*g/(2*r0p2) + 6*q0p2*q2/r0p2 - q0*q1p2*gp2/r0p2 - 5*q0*q1p2*g/r0p2 + 12*q0*q1p2/r0p2 + 3*q0*e2*g/r0 - 3*q0*e2/r0 - q0*r1*e1*gp2/r0p2 - 2*q0*r1*e1*g/r0p2 + 3*q0*r1*e1/r0p2 - q0*r2*e0*gp2/r0p2 + q0*r2*e0*g/r0p2 + 2*q0*r1p2*e0*gp2/r0p3 - 2*q0*r1p2*e0*g/r0p3 + q1*e1*gp2/r0 + 2*q1*e1*g/r0 - 3*q1*e1/r0 - 2*q1*r1*e0*gp2/r0p2 + 2*q1*r1*e0*g/r0p2 + q2*e0*gp2/r0 - q2*e0*g/r0
        E_tt = q0p4*r2*gp2/(4*r0p4) + 2*q0p4*r2*g/r0p4 - 9*q0p4*r2/(4*r0p4) - q0p4*r1p2*gp2/r0**5 - 8*q0p4*r1p2*g/r0**5 + 9*q0p4*r1p2/r0**5 + q0p3*q1*r1*gp2/r0p4 + 37*q0p3*q1*r1*g/(2*r0p4) - 39*q0p3*q1*r1/(2*r0p4) - 7*q0p3*q2*g/(2*r0p3) + 7*q0p3*q2/(2*r0p3) - 21*q0p2*q1p2*g/(2*r0p3) + 21*q0p2*q1p2/(2*r0p3) - q0p2*e2*gp2/(2*r0p2) + 3*q0p2*e2*g/r0p2 - 3*q0p2*e2/(2*r0p2) + q0p2*r1*e1*gp2/(2*r0p3) - 15*q0p2*r1*e1*g/(2*r0p3) + 3*q0p2*r1*e1/r0p3 - q0p2*r2*e0*gp2/(2*r0p3) - 3*q0p2*r2*e0*g/(2*r0p3) + 3*q0p2*r1p2*e0*gp2/(2*r0p4) + 9*q0p2*r1p2*e0*g/(2*r0p4) - q0*q1*e1*gp2/r0p2 + 8*q0*q1*e1*g/r0p2 - 3*q0*q1*e1/r0p2 - q0*q1*r1*e0*gp2/r0p3 - 7*q0*q1*r1*e0*g/r0p3 + 2*q0*q2*e0*g/r0p2 + 2*q1p2*e0*g/r0p2 + e0*e2*gp2/r0 - e0*e2*g/r0 + e1p2*gp2/r0 - e1p2*g/r0 - r1*e0*e1*gp2/r0p2 + r1*e0*e1*g/r0p2
        return np.array([rho_tt, q_tt, E_tt], dtype=np.float64)

    else:
        raise Exception("Invalid order")

Он еще далек от совершенства, но я не могу найти общий шаблон для всех подвыражений (проще разработать выражение, чем факторизовать его).

На моей машине сборка этой функции занимает 1,55 секунды (с тем же временем выполнения). Это сборка на 2,4 быстрее.

Вы можете еще больше уменьшить эти накладные расходы, используя флаг компиляции cache=True.


Настройка и тестирование

Вот настройка, используемая для измерения производительности вашего кода:

import numpy as np
import numpy.typing as npt
import numba as nb

%%time
@nb.njit('(float64[::1], float64[::1], float64[::1], int64, float64)')
def cauchyKovalevskaya(drho_x: npt.NDArray[np.float64], dq_x: npt.NDArray[np.float64], dE_x: npt.NDArray[np.float64], k: int, gamma: float) -> npt.NDArray:
    # [...]

drho_x = np.random.rand(3)
dq_x = np.random.rand(3)
dE_x = np.random.rand(3)
gamma = 42.0
%timeit -n 20_000 cauchyKovalevskaya(drho_x, dq_x, dE_x, 0, gamma)
%timeit -n 20_000 cauchyKovalevskaya(drho_x, dq_x, dE_x, 1, gamma)
%timeit -n 20_000 cauchyKovalevskaya(drho_x, dq_x, dE_x, 2, gamma)

Что касается времени выполнения, вот тайминги на моей машине (с тщательной компиляцией):

1.27 µs ± 8.18 ns per loop (mean ± std. dev. of 7 runs, 20000 loops each)
1.28 µs ± 16 ns per loop (mean ± std. dev. of 7 runs, 20000 loops each)
1.37 µs ± 24.4 ns per loop (mean ± std. dev. of 7 runs, 20000 loops each)

Анализ производительности и решения

>90% времени приходится на накладные расходы (например, создание выходного массива, проверку входных типов и преобразование их в собственные типы, вызов собственной функции из CPython), а не на вычисления. Предварительное выделение выходных данных помогает немного сократить время выполнения. При этом невозможно значительно сократить время выполнения, пока эта функция вызывается из CPython. Это основное узкое место. Вы можете скомпилировать вызывающую функцию с помощью Numba, чтобы значительно сократить эти накладные расходы.

На моей машине похоже, что вычисление занимает всего 0,1 мкс.

Обратите внимание, что после того, как вызывающая функция скомпилирована с помощью Numba, вы можете добавить флаг компиляции fastmath=True, чтобы ускорить математические вычисления. При этом этот параметр вреден, если значения могут быть NaN, -inf, +inf, ниже нормы или даже если ассоциативность операции имеет значение. Подробнее об этом читайте в этом посте. Я не рекомендую это в данном случае, поскольку не ожидаю огромного ускорения модифицированного кода.

+1 Правильно — накладные расходы имеют значение (код, созданный с помощью njit(), платит больше неснижаемых накладных расходов, чем технический JIT-скомпилированный код может сэкономить взамен). Поскольку динамические системы являются «последовательными» во времени (последовательность шагов во времени взаимно (линейно) зависима, вероятность векторизации или какого-либо другого повышения производительности равна нулю. И последнее, но не менее важное: старые системы управления. / HFT-трюки могут помочь: предварительно вычислить все константы + избегать DIV-ов со стоимостью 0(1) (постарайтесь найти любые внутренние зависимости для повторного использования частичных результатов на лету, чтобы не платить одни и те же затраты более одного раза). Никакой другой магии. или доступные трюки

user3666197 26.06.2024 21:20

Вопрос :
"Numba njit делает вызов функции, включающей длинные выражения, очень медленно"

При всем уважении, нет, давайте не будем винить стрелу, которую мы использовали, в том, что она не попала в рыбу.

Методы конечного объема обычно выполняют итерацию во времени в некоторых двумерных или трехмерных областях, поэтому опубликованный код будет вызываться огромное количество раз. Каким бы неэффективным ни был изначально написанный код, в принципе ничего, кроме k-переключенного return()-однострочника (хотя называть SLOC размером около 250 [кБ] однострочным не так уж и часто, но код «как есть» был таким разработанный), факт входа в очень много раз вызываемую часть вычислений конечного объема доставляет больше проблем снаружи, а затем был обнаружен внутри метода def-ed.

Накладные расходы имеют значение. Для многократно вызываемого кода ВСЕ накладные расходы имеют значение.

Хитрость, скомпилированная с помощью Numba JIT, потребует некоторых дополнительных затрат (первоначальная стоимость компиляции + дополнительные накладные расходы на каждый вызов), поэтому, если вы не спроектируете вызывающий код лучше, вы вскоре поймете, что продолжаете платить гораздо больше за дополнительные накладные расходы, а затем вы когда-либо получите любое (только потенциальное) ускорение.

Почему?

Помимо неснижаемых накладных расходов, вы также платите дополнительное время за то, что можно было бы сделать более разумно. Дополнительные накладные расходы, оплачиваемые за каждый return() для создания экземпляра ALAP нового нового объекта Numpy, его инициализации, выделения новой памяти для перемещения всего 3 скаляров, технически плохи. И последнее, но не менее важное: ваш расчетный код на самом деле повторяет столько временных фрагментов вычислений, что вообще не нужно было делать это таким образом и столько раз снова.

Как лучше?

Мы знаем, что можем выполнить в принципе около 30x ADD-, SUB-упс или около 10x MUL-уопов за время, когда вычисляется один единственный DIV-уоп, поэтому DIV-s являются наиболее дорогостоящими элементарными операциями (чем больше sqrt() или **0.5 и др.), поэтому, говоря о производительности, о ВЫСОКОЙ ПРОИЗВОДИТЕЛЬНОСТИ, мы можем избегать DIV-ов с помощью более умных MUL-ов.

Далее следует искусство алгебры (возможности см. в коде макета ниже),
не только для сокращения присваиваний длиной 80 000 байт, но и для более разумной настройки вычислений дерева зависимостей, чтобы избежать повторения и повторения вычислений значения, которые мы уже вычислили (не ждите волшебства от оптимизации компилятора, это умно, но для других целей, кроме преобразования 80 000 байт опубликованных уравнений, полученных по времени, в алгебраически эквивалентные, более эффективные вычисления).

Вопрос :
(цит.) "(...) строки становятся всё длиннее и длиннее (80000+ символов в строке)"

Длина линии не является единственной основной причиной проблемы, а их алгебраическая неэффективность.

Прилагается редакция кода с комментариями, которая демонстрирует цель рефакторинга - найти самоповторяющиеся фрагменты, перепроектировать использование переформулировки, позволяющей избежать DIV, и эффективно использовать MUL, чтобы вычислить то же самое, многократно повторно используя те значения, которые встречаются много раз внутри уравнений:

def cauchyKovalevskaya( drho_x: npt.NDArray[np.float64],
                          dq_x: npt.NDArray[np.float64],
                          dE_x: npt.NDArray[np.float64],
                             k: int,
                        #gamma: float
                             g: float
                              ) -> npt.NDArray:
    
    if k == 0:                                                          ## JIT/RET --^         non-MISRA-C
        return np.array([drho_x[0],
                           dq_x[0],
                           dE_x[0]
                           ],
                        dtype=np.float64)                               ## rather expensive to instantiate a fresh new instance, initialise and allocate memory ALAP so as to return just 3 scalars
    
    if k == 1:
        #g         = gamma
        gp2        = g**2
        
        r0, r1, r2 = drho_x
        q0, q1, q2 = dq_x
        e0, e1, e2 = dE_x
        
        r0p2, r1p2 = r0**2,   r1**2
        
        #0p3, r0p4 = r0**3,   r0**4
        r0p3, r0p4 = r0*r0p2, r0p2*r0p2
        
        q0p2, q1p2 = q0**2, q1**2
        e0p2, e1p2 = e0**2, e1**2
        
        
        q0p3, q0p4 = q0**3, q0**4

        ##--------------------------------------------------------------##---------------------------------------##-------------------------------------------------------------------------------( per-call consts )
        ##                                                              ##                                       ## var_q0___DIV_r0       1x DIV   VALUE REUSED: 8x
        ##                                                              ##                                       ## var_q0p2_DIV_r0p2     1x MUL                 3x
        ##                                                              ##                                       ## var_q0p3_DIV_r0p3     1x MUL                 ?  ( MUL var_q0p2_DIV_r0p2, var_q0___DIV_r0 )
        ##                                                              ##                                       ## 
        ##                                                              ##                                       ##                           var_q0p2_DIV_r0p2
        rho_t =             -q1                                         ##                                       ##                                    /
        q_t   =  (    -q0p2        *r1                *g / (2*r0p2)     ## ~~ constA * r1 * q0p2 / r0p2          ## constA ~ -g/2   \_____ ( r1 * q0p2/r0p2 )*((3-g)/2) ... constAB (( 3 - g )*0.5 )
                 + 3 * q0p2        *r1                   / (2*r0p2)     ## ~~ constB * r1 * q0p2 / r0p2          ## constB ~ +3/2   /
                 +     q0   *q1                       *g /    r0        ## ~~ constC *      q0*q1/ r0            ## constC ~ +g     \_____ ( q1 * q0  /r0   )*( g - 3 ) ... constCD
                 - 3 * q0   *q1                          /    r0        ## ~~ constD *      q0*q1/ r0            ## constD ~ -3     /
                 -                               e1   *g                ## ~~ constE *                   e1      ## constE ~ -g     \_____ e1 *               ( 1 - g ) ... constEF
                 +                               e1                     ## ~~ constF *                   e1      ## constF ~ +1     /
                   )                                                    ##---------------------------------------##-------------------------------------------------------------------------------
        ##                                                              ##                                       ## 
        ##                                                              ##                                       ##                      (  var_q0p2_DIV_r0p2 
        ##                                                              ##                                       ##                      *  var_q0___DIV_r0   )
        ##                                                              ##                                       ##                                  /
        E_t   =  (    -q0p3        *r1                *g /    r0p3      ## ~~ constG * r1 * q0p3 / r0p3          ## constG ~ -g     \_____ r1 * q0p3/r0p3 * ( 1 - g ) ... constEF
                 +     q0p3        *r1                   /    r0p3      ## ~~ constH * r1 * q0p3 / r0p3          ## constH ~ +1     /
                 + 3 * q0p2 *q1                       *g / (2*r0p2)     ## ~~ constI * q1 * q0p2 / r0p2          ## constI ~ +3g/2                    \____ q1 * q0p2 / r0p2 * ( 3/2 * ( g - 1 ) ) ... constIJ
                 - 3 * q0p2 *q1                          / (2*r0p2)     ## ~~ constJ * q1 * q0p2 / r0p2          ## constJ ~ -3/2 == -constB          /
                 -     q0                       *e1   *g /    r0        ## ~~ constK *      q0   / r0   *e1      ## constK ~ -g   == -constC == constE         g *               e1 * q0 / r0
                 +     q0          *r1       *e0      *g /    r0p2      ## ~~ constL * r1 * q0   / r0p2 *e0      ## constL ~ +g   ==  constC          \____    g * e0 / r0   * ( r1 * q0 / r0 - q1 )
                 -           q1              *e0      *g /    r0        ## ~~ constM * q1        / r0   *e0      ## constM ~ -g   == -constC == constE/
                   )
        return np.array([rho_t, q_t, E_t], dtype=np.float64)            ## rather expensive to instantiate a fresh new instance, initialise and allocate memory ALAP so as to return just 3 scalars
    
    elif k == 2:
        #g         = gamma
        gp2        = g**2
        
        r0, r1, r2 = drho_x
        q0, q1, q2 = dq_x
        e0, e1, e2 = dE_x
        
        r0p2, r1p2 = r0**2,   r1**2
        
        #0p3, r0p4 = r0**3,   r0**4
        #0p3, r0p4 = r0*r0p2, r0p2*r0p2
        
        q0p2, q1p2 = q0**2, q1**2
        e0p2, e1p2 = e0**2, e1**2
        
        q0p3, q0p4 = q0**3, q0**4

        ##--------------------------------------------------------------##---------------------------------------##-------------------------------------------------------------------------------( per-call consts )
        ##                                                              ##                                       ## var______DIV_r0       1x DIV   VALUE REUSED: 2x
        ##                                                              ##                                       ## var_q0___DIV_r0       1x MUL                 2x (+MUL var_q0___DIV_r0,   var______DIV_r0 ~ q0/r0p2 )
        ##                                                              ##                                       ## var_q0p2_DIV_r0p2     1x MUL                 2x
        ##                                                              ##                                       ## var_q0p3_DIV_r0p3     1x MUL                 ?  ( MUL var_q0p2_DIV_r0p2, var_q0___DIV_r0 )
        ##                                                              ##                                       ##                             
        ##                                                              ##                                       ##                          var_q0p2_DIV_r0p2
        ##                                                              ##                                       ##                                   /
        rho_tt = (     q0p2              *r2            *g / (2*r0p2)   ## ~~ const * r2 * q0p2 / r0p2           ## const   ~ +g/2 \_____ ( r2 * q0p2/r0p2 ) *((g-3)*0.5) ... constAB2
                 - 3 * q0p2              *r2               / (2*r0p2)   ## ~~ const * r2 * q0p2 / r0p2           ##         ~ -3/2 / 
                 -     q0p2        *r1p2                *g /    r0p3    ##                                                         \_____ ( r1p2*q0p2/r0p2/r0)*(  3 - g ) 
                 + 3 * q0p2        *r1p2                   /    r0p3    ##                                                         /                                     \
                 + 2 * q0   *q1    *r1                  *g /    r0p2    ##                                                         \_____ ( r1 * q0  /r0p2*q1)*(2*g - 6 ) \
                 - 6 * q0   *q1    *r1                     /    r0p2    ##                                                         /                                       \
                 -     q0       *q2                     *g /    r0      ##                                                         \_____ (      q0  /r0  *q2)*(  3 - g ) --+---- ( 3 - g ) * ( (...) + (...) + (...) )
                 + 3 * q0       *q2                        /    r0      ##                                                         /                                       /
                 -           q1p2                       *g /    r0      ##                                                         \_____ (      q1p2/r0     )*(  3 - g ) /
                 + 3 *       q1p2                          /    r0      ##                                                         /
                 +                                  e2  *g              ##                                                         \_____ (                e2)*(  g - 1 ) 
                 -                                  e2                  ##                                                         /
                   )                            
                                                
        q_tt   = (     q0p3              *r2           *gp2/ (2*r0p3)
                 +     q0p3              *r2           *g  /    r0p3
                 - 7 * q0p3              *r2               / (2*r0p3)
                 - 3 * q0p3        *r1p2               *gp2/ (2*r0p4)
                 - 3 * q0p3        *r1p2               *g  /    r0p4
                 +21 * q0p3        *r1p2                   / (2*r0p4)
                 + 5 * q0p2 *q1    *r1                 *gp2/ (2*r0p3)
                 + 8 * q0p2 *q1    *r1                 *g  /    r0p3
                 -45 * q0p2 *q1    *r1                     / (2*r0p3)
                 -     q0p2     *q2                    *gp2/ (2*r0p2)
                 - 5 * q0p2     *q2                    *g  / (2*r0p2)
                 + 6 * q0p2     *q2                        /    r0p2
                 -     q0   *q1p2                      *gp2/    r0p2
                 - 5 * q0   *q1p2                      *g  /    r0p2
                 +12 * q0   *q1p2                          /    r0p2
                 + 3 * q0                          *e2 *g  /    r0
                 - 3 * q0                          *e2     /    r0
                 -     q0           *r1         *e1    *gp2/    r0p2
                 - 2 * q0           *r1         *e1    *g  /    r0p2
                 + 3 * q0           *r1         *e1        /    r0p2
                 -     q0                *r2 *e0       *gp2/    r0p2
                 +     q0                *r2 *e0       *g  /    r0p2
                 + 2 * q0           *r1p2    *e0       *gp2/    r0p3
                 - 2 * q0           *r1p2    *e0       *g  /    r0p3
                 +          q1                  *e1    *gp2/    r0
                 + 2       *q1                  *e1    *g  /    r0
                 - 3       *q1                  *e1        /    r0
                 - 2       *q1      *r1      *e0       *gp2/    r0p2
                 + 2       *q1      *r1      *e0       *g  /    r0p2
                 +               q2          *e0       *gp2/    r0
                 -               q2          *e0       *g  /    r0
                   )
        E_tt   = (   q0p4*r2*gp2/(4*r0p4)
                 + 2*q0p4*r2*g/r0p4
                 - 9*q0p4*r2/(4*r0p4)
                 -   q0p4*r1p2*gp2/r0**5
                 - 8*q0p4*r1p2*g/r0**5
                 + 9*q0p4*r1p2/r0**5
                 +   q0p3*q1*r1*gp2/r0p4
                 +37*q0p3*q1*r1*g/(2*r0p4)
                 -39*q0p3*q1*r1/(2*r0p4)
                 - 7*q0p3*q2*g/(2*r0p3)
                 + 7*q0p3*q2/(2*r0p3)
                 -21*q0p2*q1p2*g/(2*r0p3)
                 +21*q0p2*q1p2/(2*r0p3)
                 -   q0p2*e2*gp2/(2*r0p2)
                 + 3*q0p2*e2*g/r0p2
                 - 3*q0p2*e2/(2*r0p2)
                 +   q0p2*r1*e1*gp2/(2*r0p3)
                 -15*q0p2*r1*e1*g/(2*r0p3)
                 + 3*q0p2*r1*e1/r0p3
                 -   q0p2*r2*e0*gp2/(2*r0p3)
                 - 3*q0p2*r2*e0*g/(2*r0p3)
                 + 3*q0p2*r1p2*e0*gp2/(2*r0p4)
                 + 9*q0p2*r1p2*e0*g/(2*r0p4)
                 -   q0*q1*e1*gp2/r0p2
                 + 8*q0*q1*e1*g/r0p2
                 - 3*q0*q1*e1/r0p2
                 -   q0*q1*r1*e0*gp2/r0p3
                 - 7*q0*q1*r1*e0*g/r0p3
                 + 2*q0*q2*e0*g/r0p2
                 + 2*q1p2*e0*g/r0p2
                 +     e0*e2*gp2/r0
                 -     e0*e2*g/r0
                 +     e1p2*gp2/r0
                 -     e1p2*g/r0
                 - r1*e0*e1*gp2/r0p2
                 + r1*e0*e1*g/r0p2
                   )
        
        return np.array([rho_tt, q_tt, E_tt], dtype=np.float64)         ## rather expensive to instantiate a fresh new instance, initialise and allocate memory ALAP so as to return just 3 scalars

    else:
        raise Exception("Invalid order")

Моей целью никогда не было обвинить Нумбу! Я хотел знать, что не так с моим кодом, из-за которого он так плохо оптимизируется. Но ваш ответ проливает свет на некоторые вещи, которые было бы весьма полезно изменить в коде.

Siddharth Yajaman 27.06.2024 12:36

Рад, что вы нашли это вдохновляющим. Красота высочайшего качества достигается за счет уважения к техническим деталям. Удачи в ваших исследованиях и следите за обновлениями, наша интересная работа всегда преподносит сюрпризы, @SiddharthYajaman

user3666197 27.06.2024 20:55

Поздравляем, что у вас хватило смелости факторизовать выражение (я даже не пытался)! При этом «в принципе мы можем выполнить около 30x ADD-, SUB-uops или около 10x MUL-uops за время, пока вычисляется один DIV-uop», это уже не соответствует действительности на основных процессорах (начиная с> 10 лет) на скалярные числа FP64. Например, что касается пропускной способности инструкций (что здесь наиболее важно), процессоры Intel могут выполнять 2 addsd/subsd/mulsd за цикл, а divsd занимает 3 цикла/элемент. Таким образом, умножение FP64 происходит так же быстро, как сложение/вычитание FP64, а деление всего в 6 раз медленнее (а не в 30).

Jérôme Richard 28.06.2024 02:12

Однако это утверждение приблизительно справедливо для целых чисел (не используется в коде OP): на основных процессорах x86 mul обычно в 4 раза медленнее, чем add/sub, а div > 10x (в зависимости от количества бит целого числа). -- часто медленнее для 64-битных).

Jérôme Richard 28.06.2024 02:29

Другие вопросы по теме