cProfile добавляет значительные накладные расходы при вызове функций numba jit

Сравните чистую функцию no-op Python с функцией no-op, украшенной @numba.jit, то есть:

import numba

@numba.njit
def boring_numba():
    pass

def call_numba(x):
    for t in range(x):
        boring_numba()

def boring_normal():
    pass

def call_normal(x):
    for t in range(x):
        boring_normal()

Если это время с %timeit, мы получаем следующее:

%timeit call_numba(int(1e7))
792 ms ± 5.51 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit call_normal(int(1e7))
737 ms ± 2.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Все совершенно разумно; есть небольшая накладная плата для функции numba, но не так много.

Если, однако, мы используем cProfile для профайла этого кода, мы получаем следующее:

cProfile.run('call_numba(int(1e7)); call_normal(int(1e7))', sort='cumulative')

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
     76/1    0.003    0.000    8.670    8.670 {built-in method builtins.exec}
        1    6.613    6.613    7.127    7.127 experiments.py:10(call_numba)
        1    1.111    1.111    1.543    1.543 experiments.py:17(call_normal)
 10000000    0.432    0.000    0.432    0.000 experiments.py:14(boring_normal)
 10000000    0.428    0.000    0.428    0.000 experiments.py:6(boring_numba)
        1    0.000    0.000    0.086    0.086 dispatcher.py:72(compile)

cProfile считает, что при вызове функции numba существует огромная накладная стоимость. Это распространяется на "настоящий" код: у меня была функция, которая просто называлась моим дорогостоящим вычислением (вычисление было numba-JIT-скомпилировано), а cProfile сообщил, что функция-обертка занимает около трети от общего времени.

Я не возражаю против того, что cProfile добавляет немного накладных расходов, но если он массово несовместим с тем, где он добавляет, что накладные расходы это не очень полезно. Кто-нибудь знает, почему это происходит, есть ли что-нибудь, что можно сделать по этому поводу, и/или если есть какие-либо альтернативные инструменты профилирования, которые плохо взаимодействуют с numba?

Ответ 1

При создании функции Numba вы на самом деле создать Numba Dispatcher объект. Этот объект "перенаправляет" "вызов" на boring_numba на правильную (насколько это касается типов) внутреннюю функцию "jitted". Поэтому, даже если вы создали функцию boring_numba - эта функция не вызывается, то, что называется, является скомпилированной функцией, основанной на вашей функции.

Просто вы видите, что boring_numba функция boring_numba (даже если это не так, что называется CPUDispatcher.__call__) во время профилирования объекта Dispatcher необходимо подключиться к текущему состоянию потока и проверить, работает ли профайлер/трассировщик и если "да", это заставляет его выглядеть как boring_numba Этот последний шаг - это то, что навлекает накладные расходы, потому что он должен подделать "стек стека Python" для boring_numba.

Немного более технический:

Когда вы вызываете функцию boring_numba она фактически вызывает Dispatcher_Call которая является оберткой вокруг call_cfunc и вот основное различие: когда у вас есть профайлер, выполняющий код, связанный с профилировщиком, составляет большую часть вызова функции (просто сравните if (tstate->use_tracing && tstate->c_profilefunc) веткой else которая запущена, если нет профилировщика/трассировщика):

static PyObject *
call_cfunc(DispatcherObject *self, PyObject *cfunc, PyObject *args, PyObject *kws, PyObject *locals)
{
    PyCFunctionWithKeywords fn;
    PyThreadState *tstate;
    assert(PyCFunction_Check(cfunc));
    assert(PyCFunction_GET_FLAGS(cfunc) == METH_VARARGS | METH_KEYWORDS);
    fn = (PyCFunctionWithKeywords) PyCFunction_GET_FUNCTION(cfunc);
    tstate = PyThreadState_GET();
    if (tstate->use_tracing && tstate->c_profilefunc)
    {
        /*
         * The following code requires some explaining:
         *
         * We want the jit-compiled function to be visible to the profiler, so we
         * need to synthesize a frame for it.
         * The PyFrame_New() constructor doesn't do anything with the 'locals' value if the 'code's
         * 'CO_NEWLOCALS' flag is set (which is always the case nowadays).
         * So, to get local variables into the frame, we have to manually set the 'f_locals'
         * member, then call 'PyFrame_LocalsToFast', where a subsequent call to the 'frame.f_locals'
         * property (by virtue of the 'frame_getlocals' function in frameobject.c) will find them.
         */
        PyCodeObject *code = (PyCodeObject*)PyObject_GetAttrString((PyObject*)self, "__code__");
        PyObject *globals = PyDict_New();
        PyObject *builtins = PyEval_GetBuiltins();
        PyFrameObject *frame = NULL;
        PyObject *result = NULL;

        if (!code) {
            PyErr_Format(PyExc_RuntimeError, "No __code__ attribute found.");
            goto error;
        }
        /* Populate builtins, which is required by some JITted functions */
        if (PyDict_SetItemString(globals, "__builtins__", builtins)) {
            goto error;
        }
        frame = PyFrame_New(tstate, code, globals, NULL);
        if (frame == NULL) {
            goto error;
        }
        /* Populate the 'fast locals' in 'frame' */
        Py_XDECREF(frame->f_locals);
        frame->f_locals = locals;
        Py_XINCREF(frame->f_locals);
        PyFrame_LocalsToFast(frame, 0);
        tstate->frame = frame;
        C_TRACE(result, fn(PyCFunction_GET_SELF(cfunc), args, kws));
        tstate->frame = frame->f_back;

    error:
        Py_XDECREF(frame);
        Py_XDECREF(globals);
        Py_XDECREF(code);
        return result;
    }
    else
        return fn(PyCFunction_GET_SELF(cfunc), args, kws);
}

Я предполагаю, что этот дополнительный код (в случае профайлера) замедляет функцию, когда вы выполняете cProfile-ing.

Немного неудачно, что функция numba добавляет столько накладных расходов при запуске профилировщика, но замедление фактически будет практически незначительным, если вы сделаете что-то существенное в функции numba. Если бы вы также переместили цикл for в функцию numba, тем более.

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

Примечание. Все это (небольшая) спекуляция, я на самом деле не строю numba с символами отладки и не профилировал C-Code в случае, если работает профайлер. Однако количество операций в случае, если работает профайлер, делает это очень правдоподобным. И все это предполагает numba 0.39, не уверен, что это относится и к прошлым версиям.