OpenMP loops hang with Python virtual function overrides

[jmikeowen]

I'm having a problem with pybind11 wrapped code using virtual function overrides in OpenMP sections of C++ code. An OpenMP loop will hang when run with more than one thread if there are calls to Python overrides in the body of the code. Since the macro PYBIND11_OVERRIDE acquires the GIL I thought these sorts of operations should be safe, but clearly I'm missing something.

Consider the following C++ code wrapped in pybind11:

#include "pybind11/pybind11.h"
#include "pybind11/functional.h"

namespace py = pybind11;
using namespace pybind11::literals;

#include <cstdio>

class A {
public:
  A()                               { printf("A::A()\n"); }    
  virtual ~A()                      { printf("A::~A()\n"); }   
  virtual void void_func() const    { printf("A::void_func()\n"); }
  virtual int int_func(int x) const { printf("A::int_func(%d)\n", x); return x + 1; }
};

void do_threaded_stuff(const A& a) {
  int sum = 0;
#pragma omp parallel for
  for (auto i = 0u; i < 10u; ++i) {
    a.void_func();
#pragma omp critical
    {
      sum = a.int_func(sum);
    }
  }
  printf("Final sum: %d\n", sum);
}

//------------------------------------------------------------------------------
// Trampoline class for A
//------------------------------------------------------------------------------
class PYB11TrampolineA: public A {
public:
  using A::A;
  virtual void void_func() const override {
    PYBIND11_OVERRIDE(void, A, void_func);
  }
  virtual int int_func(int x) const override {
    PYBIND11_OVERRIDE(int, A, int_func, x);
  }
};

//------------------------------------------------------------------------------
// Make the module
//------------------------------------------------------------------------------
PYBIND11_MODULE(virtual_override_thread, m) {
  py::class_<A, PYB11TrampolineA> obj(m, "A");
  obj.def(py::init<>());
  obj.def("void_func", (void (A::*)() const) &A::void_func);
  obj.def("int_func", (int (A::*)(int) const) &A::int_func);

  m.def("do_threaded_stuff", (void (*)(const A&)) &do_threaded_stuff, "a"_a);
}

If we exercise this code from Python (with more than one OpenMP thread) I see calls to do_threaded_stuff work fine with A objects (no Python overrides), but if I create a Python class inheriting from A and override these methods the loop in do_threaded_stuff will hang. Here is an example Python block that will reproduce the problem with the above Python module:

from virtual_override_thread import *

class B(A):
    def __init__(self):
        A.__init__(self)

    def void_func(self):
        print("B::void_func")

    def int_func(self, x):
        print("B::int_func({})".format(x))
        return x + 10

a = A()
do_threaded_stuff(a)    # OK

b = B()
do_threaded_stuff(b)    # Hang with OMP_NUM_THREADS > 1

How can I make this sort of pattern thread safe for OpenMP threaded code?