auto_gil.h

#pragma once

// RAII structs to acquire and release Python's global interpreter lock (GIL)

#include <torch/csrc/python_headers.h>

// Acquires the GIL on construction
struct AutoGIL {
  AutoGIL() : gstate(PyGILState_Ensure()) {
  }
  ~AutoGIL() {
    PyGILState_Release(gstate);
  }

  PyGILState_STATE gstate;
};

// Releases the GIL on construction
struct AutoNoGIL {
  AutoNoGIL() : save(PyEval_SaveThread()) {
  }
  ~AutoNoGIL() {
    PyEval_RestoreThread(save);
  }

  PyThreadState* save;
};

// Runs the function without the GIL
template<typename F>
inline void with_no_gil(F f) {
  AutoNoGIL no_gil;
  f();
}