cProfile 在调用 numba jit 函数时会增加大量开销

cProfile adds significant overhead when calling numba jit functions

将纯 Python 空操作函数与装饰有 @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 函数的开销很大。 这扩展到 "real" 代码:我有一个简单地调用我的昂贵计算的函数(计算是 numba-JIT 编译的),并且 cProfile 报告包装函数大约占总数的三分之一时间。

我不介意 cProfile 增加一些开销,但如果它在增加开销的位置上存在大量不一致,那将不是很有帮助。有谁知道为什么会发生这种情况,是否有什么办法可以解决,and/or 是否有其他分析工具不会与 numba 交互不良?

当您创建一个 numba 函数时,您实际上创建了一个 numba Dispatcher 对象。这个对象 "re-directs" "call" 到 boring_numba 到正确的(就类型而言)内部 "jitted" 函数。因此,即使您创建了一个名为 boring_numba 的函数 - 该函数未被调用,被调用的是基于您的函数的编译函数

这样你就可以看到函数 boring_numba 在分析过程中被调用了(即使它没有,调用的是 CPUDispatcher.__call__Dispatcher 对象需要挂钩进入当前线程状态并检查是否有 profiler/tracer 运行ning 如果 "yes" 它看起来像 boring_numba 是 called.This 最后一步是导致开销,因为它必须为 boring_numba.

伪造 "Python stack frame"

有点技术性:

当您调用 numba 函数 boring_numba 时,它实际上调用 Dispatcher_Call which is a wrapper around call_cfunc,这是主要区别:当您有分析器时 运行ning 处理分析器的代码构成大多数函数调用(如果没有 profiler/tracer,只需比较 if (tstate->use_tracing && tstate->c_profilefunc) 分支和 运行ning 的 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);
}

我假设这个额外的代码(如果分析器是 运行ning)会在你进行 cProfile-ing 时减慢函数速度。

有点不幸的是,当您 运行 分析器时,numba 函数会增加如此多的开销,但如果您在 numba 函数中做任何实质性的事情,那么减速实际上几乎可以忽略不计。 如果您还要在 numba 函数中移动 for 循环,那么更是如此。

如果您注意到 numba 函数(有或没有分析器 运行ning)花费太多时间,那么您可能调用它的频率太高了。然后你应该检查你是否真的可以在 numba 函数中移动循环或者将包含循环的代码包装在另一个 numba 函数中。

注意:所有这些都是(有点)推测,我实际上并没有用调试符号构建 numba 并分析 C 代码以防分析器 运行ning。然而,如果有一个分析器运行ning,那么操作的数量使这看起来非常合理。所有这些都假设 numba 0.39,不确定这是否也适用于过去的版本。