On destroy - acquire gil

Author: tvorogme

tvm-python/PyGlobal.cpp

@@ -2,69 +2,78 @@
 
 #include <memory>
 #include <mutex>
+#include "third-party/pybind11/include/pybind11/pybind11.h"
 
 #include "PyGlobal.h"
 
 std::mutex scheduler_init_mutex;
 bool scheduler_running = false;
 static std::unique_ptr<td::actor::Scheduler> thread_local_scheduler;
 static std::unique_ptr<td::thread> scheduler_thread;
+namespace py = pybind11;
 
 namespace pyglobal {
-    class Runner : public td::actor::Actor {
-    public:
-        explicit Runner(std::function<void()> f) : f_(std::move(f)) {
-        }
+class Runner : public td::actor::Actor {
+public:
+  explicit Runner(std::function<void()> f)
+    : f_(std::move(f)) {
+  }
 
-        void start_up() override {
-          f_();
-          stop();
-        }
+  ~Runner() override {
+    py::gil_scoped_acquire gil;
+    f_ = {};
+  }
 
-    private:
-        std::function<void()> f_;
-    };
 
+  void start_up() override {
+    f_();
+    stop();
+  }
 
-    void init_thread_scheduler(int thread_count) {
-      std::lock_guard<std::mutex> lock(scheduler_init_mutex);
-      if (!thread_local_scheduler) {
-        thread_local_scheduler = std::unique_ptr<td::actor::Scheduler>(new td::actor::Scheduler({thread_count}));
-        scheduler_running = true;
+private:
+  std::function<void()> f_;
+};
 
-        scheduler_thread = std::make_unique<td::thread>([&] {
-            thread_local_scheduler->run();
-        });
-      }
-    }
 
-    td::actor::Scheduler *get_thread_scheduler() {
-      if (!scheduler_running) {
-        init_thread_scheduler(6);
-      }
-      return thread_local_scheduler.get();
-    }
+void init_thread_scheduler(int thread_count) {
+  std::lock_guard<std::mutex> lock(scheduler_init_mutex);
+  if (!thread_local_scheduler) {
+    thread_local_scheduler = std::unique_ptr<td::actor::Scheduler>(new td::actor::Scheduler({thread_count}));
+    scheduler_running = true;
 
-    void stop_scheduler_thread() {
-      if (scheduler_running) {
-        std::lock_guard<std::mutex> lock(scheduler_init_mutex);
-        if (thread_local_scheduler) {
-          thread_local_scheduler->run_in_context_external([] {
-              td::actor::SchedulerContext::get()->stop();
-          });
-        }
-        if (scheduler_thread) {
-          scheduler_thread->join();
-          scheduler_thread.reset();
-        }
-        thread_local_scheduler.reset();
-        scheduler_running = false;
-      }
-    }
+    scheduler_thread = std::make_unique<td::thread>([&] {
+      thread_local_scheduler->run();
+    });
+  }
+}
 
-    void execute_async(std::function<void()> f) {
-      get_thread_scheduler()->run_in_context_external([&] {
-          td::actor::create_actor<pyglobal::Runner>("executeasync", std::move(f)).release();
+td::actor::Scheduler* get_thread_scheduler() {
+  if (!scheduler_running) {
+    init_thread_scheduler(6);
+  }
+  return thread_local_scheduler.get();
+}
+
+void stop_scheduler_thread() {
+  if (scheduler_running) {
+    std::lock_guard<std::mutex> lock(scheduler_init_mutex);
+    if (thread_local_scheduler) {
+      thread_local_scheduler->run_in_context_external([] {
+        td::actor::SchedulerContext::get()->stop();
       });
     }
+    if (scheduler_thread) {
+      scheduler_thread->join();
+      scheduler_thread.reset();
+    }
+    thread_local_scheduler.reset();
+    scheduler_running = false;
+  }
+}
+
+void execute_async(std::function<void()> f) {
+  get_thread_scheduler()->run_in_context_external([&] {
+    td::actor::create_actor<pyglobal::Runner>("executeasync", std::move(f)).release();
+  });
 }
+}

tvm-python/PyTVM.cpp

@@ -35,8 +35,10 @@ void PyTVM::set_c7(PyStackEntry x) {
 
 void PyTVM::log(const std::string &log_string, int level) {
   if (log_level >= level && level == LOG_INFO) {
+    py::gil_scoped_acquire gil;
     py::print("INFO: " + log_string);
   } else if (log_level >= level && level == LOG_DEBUG) {
+    py::gil_scoped_acquire gil;
     py::print("DEBUG: " + log_string);
   }
 }
@@ -129,6 +131,7 @@ class PythonLogger : public td::LogInterface {
       }
 
       if (!muted) {
+        py::gil_scoped_acquire gil;
         py::print(slice.str());
       }
     }