1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
|
#pragma once
#include <ATen/core/ivalue.h>
#include <pybind11/pybind11.h>
#include <torch/csrc/jit/python/pybind_utils.h>
#include <torch/csrc/python_headers.h>
#include <torch/csrc/utils/pybind.h>
namespace py = pybind11;
namespace c10::ivalue {
// concrete ivalue Holder that hold a py::object
struct C10_EXPORT ConcretePyObjectHolder final : PyObjectHolder {
public:
static c10::intrusive_ptr<PyObjectHolder> create(py::object py_obj) {
return c10::make_intrusive<ConcretePyObjectHolder>(std::move(py_obj));
}
static c10::intrusive_ptr<PyObjectHolder> create(const py::handle& handle) {
py::gil_scoped_acquire ag;
return c10::make_intrusive<ConcretePyObjectHolder>(
handle.cast<py::object>());
}
PyObject* getPyObject() override {
return py_obj_.ptr();
}
InferredType tryToInferType() override {
pybind11::gil_scoped_acquire ag;
return torch::jit::tryToInferType(py_obj_);
}
IValue toIValue(const TypePtr& type, std::optional<int32_t> N = std::nullopt)
override {
pybind11::gil_scoped_acquire ag;
return torch::jit::toIValue(py_obj_, type, N);
}
std::string toStr() override {
pybind11::gil_scoped_acquire ag;
return py::str(py_obj_);
}
std::vector<at::Tensor> extractTensors() override {
// We could implement this entirely in C++ via pybind11 but it turns out to
// be substantially slower. Namely, the total time taken by markCompleted on
// a CUDAFuture is 21.5us with this implementation, but goes up to 58.7us
// when using C++. The reason is unclear.
try {
pybind11::gil_scoped_acquire ag;
#if IS_PYBIND_2_13_PLUS
PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store<py::object>
storage;
auto& extractorFn =
storage
.call_once_and_store_result([]() -> py::object {
return py::module_::import("torch._jit_internal")
.attr("_extract_tensors");
})
.get_stored();
#else
static py::object& extractorFn = *new py::object(
py::module::import("torch._jit_internal").attr("_extract_tensors"));
#endif
return extractorFn(py_obj_).cast<std::vector<at::Tensor>>();
} catch (py::error_already_set& e) {
auto err = std::runtime_error(
c10::str("Cannot extract tensors from value: ", e.what()));
{
pybind11::gil_scoped_acquire ag;
e.restore();
PyErr_Clear();
}
throw std::runtime_error(err);
}
}
// Note [Destructing py::object]
// ~~~~~~~~~~~~~~~~~~~~~~~~~~
//
// (1) Why py_obj_ = py::none(); does not work. Because we also need to
// acquire GIL when destructing py::object of None that de-references None.
// https://docs.python.org/3/c-api/none.html#c.Py_RETURN_NONE
//
// https://stackoverflow.com/questions/15287590/why-should-py-increfpy-none-be-required-before-returning-py-none-in-c
//
// (2) Why we need to call dec_ref() explicitly. Because py::object of
// nullptr, on destruction, effectively does nothing because of it calls
// Py_XDECREF(NULL) underlying.
// https://docs.python.org/3/c-api/refcounting.html#c.Py_XDECREF
~ConcretePyObjectHolder() override {
pybind11::gil_scoped_acquire ag;
py_obj_.dec_ref();
// explicitly setting PyObject* to nullptr to prevent py::object's dtor to
// decref on the PyObject again.
py_obj_.ptr() = nullptr;
}
// explicit construction to avoid errornous implicit conversion and
// copy-initialization
explicit ConcretePyObjectHolder(py::object py_obj)
: py_obj_(std::move(py_obj)) {}
private:
py::object py_obj_;
};
} // namespace c10::ivalue
|
