Fix Python GIL lock handling (Fixes #6524, Fixes #5631) (#6525)

* Fix Python GIL lock handling (Fixes #6524, Fixes #5631)

As disscussed in #6523 (comment)
and later in #6524,
pybind11 v2.8.1 added some defensive checks that fail for halide,
namely in `python_tutorial_lesson_04_debugging_2`
and `python_tutorial_lesson_05_scheduling_1`.

#6524 (comment) notes:
> * Python calls a Halide-JIT-generated function , which runs with the GIL held.
> * Halide runtime spawns worker threads.
> * The worker threads try to call pybind11's py::print function to emit traces.
> * Pybind11 complains, correctly, that the worker thread doesn't hold the GIL.
>
> Trying to acquire the GIL hangs, because the main thread is still holding it. I tried teaching the main thread to release the GIL (as suggested in #5631), but I still saw hangs when I tried this.

I have tried, and just dropping the lock before calling into halide,
or just acquiring it in `halide_python_print` doesn't work,
we need to do both.

I have verified that the two tests fail without this fix,
and pass with it.

python_bindings/src/PyError.cpp

@@ -10,6 +10,7 @@ void halide_python_error(JITUserContext *, const char *msg) {
 }
 
 void halide_python_print(JITUserContext *, const char *msg) {
+ py::gil_scoped_acquire acquire;
  py::print(msg, py::arg("end") = "");
 }
 

python_bindings/src/PyFunc.cpp

@@ -117,6 +117,7 @@ void define_func(py::module &m) {
  .def(
  "realize",
  [](Func &f, Buffer<> buffer, const Target &target) -> void {
+ py::gil_scoped_release release;
  f.realize(buffer, target);
  },
  py::arg("dst"), py::arg("target") = Target())
@@ -125,14 +126,20 @@ void define_func(py::module &m) {
  .def(
  "realize",
  [](Func &f, std::vector<Buffer<>> buffers, const Target &t) -> void {
+ py::gil_scoped_release release;
  f.realize(Realization(buffers), t);
  },
  py::arg("dst"), py::arg("target") = Target())
 
  .def(
  "realize",
  [](Func &f, const std::vector<int32_t> &sizes, const Target &target) -> py::object {
- return realization_to_object(f.realize(sizes, target));
+ std::optional<Realization> r;
+ {
+ py::gil_scoped_release release;
+ r = std::move(f.realize(sizes, target));
+ }
+ return realization_to_object(*r);
  },
  py::arg("sizes") = std::vector<int32_t>{}, py::arg("target") = Target())
 

python_bindings/src/PyPipeline.cpp

@@ -96,20 +96,27 @@ void define_pipeline(py::module &m) {
 
  .def(
  "realize", [](Pipeline &p, Buffer<> buffer, const Target &target) -> void {
+ py::gil_scoped_release release;
  p.realize(Realization(buffer), target);
  },
  py::arg("dst"), py::arg("target") = Target())
 
  // This will actually allow a list-of-buffers as well as a tuple-of-buffers, but that's OK.
  .def(
  "realize", [](Pipeline &p, std::vector<Buffer<>> buffers, const Target &t) -> void {
+ py::gil_scoped_release release;
  p.realize(Realization(buffers), t);
  },
  py::arg("dst"), py::arg("target") = Target())
 
  .def(
  "realize", [](Pipeline &p, std::vector<int32_t> sizes, const Target &target) -> py::object {
- return realization_to_object(p.realize(std::move(sizes), target));
+ std::optional<Realization> r;
+ {
+ py::gil_scoped_release release;
+ r = std::move(p.realize(std::move(sizes), target));
+ }
+ return realization_to_object(*r);
  },
  py::arg("sizes") = std::vector<int32_t>{}, py::arg("target") = Target())
 

src/PythonExtensionGen.cpp

@@ -347,7 +347,9 @@ void PythonExtensionGen::compile(const LoweredFunc &f) {
  // Python already converted this.
  }
  }
- dest << " int result = " << f.name << "(";
+ dest << " int result;\n";
+ dest << " Py_BEGIN_ALLOW_THREADS\n";
+ dest << " result = " << f.name << "(";
  for (size_t i = 0; i < args.size(); i++) {
  if (i > 0) {
  dest << ", ";
@@ -359,6 +361,7 @@ void PythonExtensionGen::compile(const LoweredFunc &f) {
  }
  }
  dest << ");\n";
+ dest << " Py_END_ALLOW_THREADS\n";
  release_buffers();
  dest << R"INLINE_CODE(
  if (result != 0) {