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
|
#include <torch/csrc/jit/runtime/interpreter.h>
#include <torch/csrc/python_headers.h>
#include <torch/csrc/autograd/edge.h>
#include <torch/csrc/autograd/function.h>
#include <torch/csrc/autograd/profiler.h>
#include <torch/csrc/autograd/variable.h>
#include <torch/csrc/jit/ir/ir.h>
#include <torch/csrc/jit/python/pybind_utils.h>
#include <torch/csrc/jit/python/python_ir.h>
#include <torch/csrc/jit/runtime/custom_operator.h>
#include <torch/csrc/jit/runtime/graph_executor.h>
#include <torch/csrc/jit/runtime/operator.h>
#include <typeinfo>
#include <pybind11/pybind11.h>
#include <torch/csrc/Exceptions.h>
#include <torch/csrc/autograd/python_engine.h>
#include <torch/csrc/autograd/python_variable.h>
#include <torch/csrc/jit/python/pybind.h>
#include <torch/csrc/utils/pybind.h>
namespace py = pybind11;
namespace torch {
namespace jit {
namespace {
// Note: const_cast is used twice below to acquire a handle to a pyobject.
Operation createPythonOperation(const Node* op_) {
pybind11::gil_scoped_acquire gil;
const ConcretePythonOp* op = static_cast<const ConcretePythonOp*>(op_);
const py::function func = py::reinterpret_borrow<const py::function>(
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
py::handle(const_cast<ConcretePythonOp*>(op)->pyobj.get()));
size_t num_inputs = 0;
for (auto arg_type : op->cconv) {
if (arg_type == 'd')
num_inputs++;
}
AT_ASSERT(op->outputs().size() == 1);
return [=](Stack& stack) {
pybind11::gil_scoped_acquire gil;
py::tuple py_inputs(op->cconv.size());
size_t i = 0;
size_t next_scalar = 0;
size_t next_tensor = 0;
for (auto arg_type : op->cconv) {
if (arg_type == 'c') {
py_inputs[i] = py::reinterpret_borrow<const py::object>(
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
const_cast<ConcretePythonOp*>(op)
->scalar_args[next_scalar++]
.get());
} else if (arg_type == 'd') {
py_inputs[i] =
toPyObject(std::move(peek(stack, next_tensor, num_inputs)));
next_tensor++;
}
i++;
}
drop(stack, num_inputs);
try {
py::object py_output(func(*py_inputs));
stack.push_back(returnToIValue(op->output()->type(), py_output));
} catch (py::error_already_set& e) {
throw std::runtime_error(e.what());
}
};
}
c10::AliasAnalysisKind aliasAnalysisIsSpecialCase() {
return AliasAnalysisKind::INTERNAL_SPECIAL_CASE;
}
RegisterOperators reg({Operator(
prim::PythonOp,
createPythonOperation,
aliasAnalysisIsSpecialCase())});
} // namespace
} // namespace jit
} // namespace torch
|
