Python 与 TBB 结合时的死锁

Question

Here 我想提供一个完整的测试用例，显示一个简单的 TBB parallel_for 构造导致 Python 应用程序死锁。 Python 前端使用 pybind11 与 TBB 后端相结合：

void backend_tbb(vector<int>& result, std::function<int (int)>& callback)
{
    int nthreads = tbb::task_scheduler_init::default_num_threads();
    const char* cnthreads = getenv("TBB_NUM_THREADS");
    if (cnthreads) nthreads = std::max(1, atoi(cnthreads));

    tbb::task_group group;
    tbb::task_arena arena(nthreads, 1);
    tbb::task_scheduler_init init(nthreads);

    group.run( [&] {
        tbb::parallel_for(tbb::blocked_range<int>(0, result.size()),
            [&](const tbb::blocked_range<int>& range)
        {
            for (int i = range.begin(); i != range.end(); i++)
                result[i] = callback(i);
        });
    });

    arena.execute( [&] { group.wait(); });      
}

void backend_serial(vector<int>& result, std::function<int (int)>& callback)
{
    for (int i = 0; i < result.size(); i++)
        result[i] = callback(i);
}

PYBIND11_MODULE(python_tbb, m)
{
    pybind11::bind_vector<std::vector<int> >(m, "stdvectorint");

    m.def("backend_tbb", &backend_tbb, "TBB backend");
    m.def("backend_serial", &backend_serial, "Serial backend");
}

backend_tbb 未注释，应用程序无限死锁：

from python_tbb import *
import numpy as np

def callback(a) :
    return int(a) * 10

def main() :
    length = 10
    result1 = stdvectorint(np.zeros(length, np.int32))
    result2 = stdvectorint(np.zeros(length, np.int32))

    backend_serial(result1, callback)
        # XXX Uncomment this to get the program hang
    #backend_tbb(result2, callback)

    for i in range(length) :
        print("%d vs %d" % (result1[i], result2[i]))

if __name__ == "__main__" :
    main()

我试过gil_scoped_acquire/gil_scoped_release，但没有任何变化。据报道Similar solution 适用于 OpenMP 循环 - 但当我尝试对 TBB 做同样的事情时，再次没有运气。请对此案提出建议，谢谢！

Answer 1

问题是 TBB 任务在与 task_group 关联的 task_arena 实例中生成，但等待是在另一个名为 arena 的 task_arena 实例中完成的。这可能导致死锁。要解决此问题，请尝试将对 group.run() 的调用封装到 task_arena.execute() 中，就像对 group.wait() 所做的那样。

但是，在这种情况下，后一种包装似乎是多余的。因此，您可能希望将两个包装合二为一

arena.execute() {
   group.run( /* ... */ );
   group.wait();
}

在这个特定的例子中，这使得task_group 的使用变得不必要，因为主线程产生任务并立即加入以参与它们的执行，就像在tbb::parallel_for 中所做的那样。因此，task_group 只能被删除。