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.

(cherry picked from commit b8eb22d)

python_bindings/src/PyError.cpp

@@ -10,6 +10,7 @@ void halide_python_error(void *, const char *msg) {
 }
 
 void halide_python_print(void *, 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) {