| #include <torch/csrc/python_headers.h> |
| |
| #include <torch/csrc/jit/argument_spec.h> |
| #include <torch/csrc/jit/export.h> |
| #include <torch/csrc/jit/ir.h> |
| #include <torch/csrc/jit/passes/python_print.h> |
| #include <torch/csrc/jit/passes/shape_analysis.h> |
| #include <torch/csrc/jit/pybind.h> |
| #include <torch/csrc/jit/python_tracer.h> |
| #include <torch/csrc/utils/auto_gil.h> |
| #include <torch/csrc/utils/pybind.h> |
| #include <torch/csrc/utils/python_strings.h> |
| |
| #include <iostream> |
| #include <sstream> |
| |
| namespace torch { |
| namespace jit { |
| |
| using c10::Type; |
| |
| std::string getPythonName(const PyObject* obj_) { |
| AutoGIL gil; |
| // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast) |
| PyObject* obj = const_cast<PyObject*>(obj_); |
| auto v = py::getattr(obj, "__name__", py::str("<python_value>")); |
| // if this was a autograd.Function recover the name of the class |
| return py::str(v); |
| } |
| |
| std::ostream& printPyObject(std::ostream& out, const THPObjectPtr& obj) { |
| AutoGIL gil; |
| // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast) |
| auto pyobj = py::handle(const_cast<PyObject*>(obj.get())); |
| if (py::isinstance<py::tuple>(pyobj)) { |
| // This special-case for printing tuples handles a problem where |
| // str((2L, 3L)) outputs "(2L, 3L)" in Python 2 but "(2, 3)" |
| // in Python 3. In order to suppress the L-suffix, we must |
| // manually print the string ourselves, calling str() on the |
| // sub-elements. |
| // |
| // This is a fairly fragile fix (What if you have nested tuples |
| // in tuples? What if you have dictionaries?) but it seems to hit |
| // the cases that are triggered in practice in onnx-pytorch. Revisit |
| // this code if this is not the case. |
| // |
| // By the way, one non-solution for this problem is to monkeypatch |
| // tuple.__str__; this doesn't work because Python doesn't allow |
| // monkeypatching methods of built-in types. |
| auto pytuple = pyobj.cast<py::tuple>(); |
| out << "("; |
| size_t i = 0; |
| for (const auto& o : pytuple) { |
| if (i > 0) { |
| out << ", "; |
| } |
| THPObjectPtr str(py::str(o).release().ptr()); |
| out << THPUtils_unpackString(str.get()); |
| i++; |
| } |
| if (i == 1) { |
| out << ","; |
| } |
| out << ")"; |
| return out; |
| } else { |
| return out << THPUtils_unpackString(py::str(pyobj).ptr()); |
| } |
| } |
| |
| // execute a Python function, used for Ops we can't optimize but that we want to |
| // optimize around |
| struct ConcretePythonOp : public PythonOp { |
| ConcretePythonOp(Graph* graph) : PythonOp(graph) {} |
| std::string name() const override { |
| AutoGIL gil; |
| if (auto autograd = autogradFunction()) { |
| return getPythonName(autograd->get()); |
| } else { |
| return getPythonName(pyobj.get()); |
| } |
| } |
| void cloneFrom(Node* other_) override { |
| Node::cloneFrom(other_); |
| auto other = other_->cast<PythonOp>(); |
| this->cconv = other->cconv; |
| Py_INCREF(other->pyobj.get()); |
| this->pyobj = THPObjectPtr(other->pyobj.get()); |
| for (auto& sa : other->scalar_args) { |
| Py_INCREF(sa.get()); |
| this->scalar_args.emplace_back(sa.get()); |
| } |
| } |
| Node* allocNewInstance(Graph* g) override { |
| return new ConcretePythonOp(g); |
| } |
| // recover the autograd.Function instance, if this PythonOp's function |
| // was originally SomeFunction.apply |
| // used in ONNX for discovering symbolics |
| c10::optional<THPObjectPtr> autogradFunction() const override { |
| AutoGIL gil; |
| // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast) |
| py::handle obj = const_cast<PyObject*>(pyobj.get()); |
| |
| auto r = py::getattr(obj, "__self__", py::none()); |
| if (r.is_none()) |
| return c10::nullopt; |
| |
| auto apply = py::getattr(r, "apply", py::none()); |
| if (apply.is_none()) |
| return c10::nullopt; |
| |
| auto c = PyObject_RichCompareBool(apply.ptr(), obj.ptr(), Py_NE); |
| if (PyErr_Occurred()) |
| throw py::error_already_set(); |
| if (c) |
| return c10::nullopt; |
| |
| return THPObjectPtr(r.release().ptr()); |
| } |
| |
| void writeScalars(std::ostream& out) const override { |
| out << "("; |
| int i = 0; |
| for (auto& scalar : scalar_args) { |
| if (i++ > 0) |
| out << ", "; |
| printPyObject(out, scalar); |
| } |
| out << ")"; |
| } |
| }; |
| |
| PythonOp* pythonAllocPythonOp(Graph* g) { |
| return new ConcretePythonOp(g); |
| } |
| |
| void initPythonIRBindings(PyObject* module_) { |
| setAllocPythonOp(pythonAllocPythonOp); |
| |
| auto m = py::handle(module_).cast<py::module>(); |
| #define GS(name) def(#name, &Graph ::name) |
| py::class_<Graph, std::shared_ptr<Graph>>(m, "Graph") |
| .def(py::init<>()) |
| .def( |
| "__repr__", |
| [](Graph& g) { |
| std::stringstream ss; |
| ss << g; |
| return ss.str(); |
| }) |
| .def( |
| "propagate_shapes", |
| [](std::shared_ptr<Graph> g, |
| std::vector<at::Tensor> inputs, |
| bool with_grad) { |
| setInputTypes( |
| *g, |
| ArgumentSpec(with_grad, fmap<IValue>(inputs), inputs.size())); |
| PropagateInputShapes(g); |
| }) |
| .def( |
| "_export_onnx", |
| [](const std::shared_ptr<Graph> g, |
| const std::vector<at::Tensor>& initializers, |
| int64_t onnx_opset_version, |
| bool defer_weight_export, |
| ::torch::onnx::OperatorExportTypes operator_export_type) { |
| std::string graph; |
| RawDataExportMap export_map; |
| std::tie(graph, export_map) = export_onnx( |
| g, |
| initializers, |
| onnx_opset_version, |
| defer_weight_export, |
| operator_export_type); |
| std::unordered_map<std::string, py::bytes> |
| python_serialized_export_map; |
| for (auto& kv : export_map) { |
| auto t = kv.second; |
| size_t copy_bytes = t.type().elementSizeInBytes() * t.numel(); |
| // TODO: this is an unecessary copy. In theory we can directly |
| // return the map from identifier to Tensor, but we need some API |
| // in Python to get raw `bytes` containing the raw tensor data. |
| python_serialized_export_map[kv.first] = |
| py::bytes(static_cast<const char*>(t.data_ptr()), copy_bytes); |
| } |
| return std::make_tuple( |
| py::bytes(graph), python_serialized_export_map); |
| }, |
| py::arg("initializers"), |
| py::arg("onnx_opset_version") = 0, |
| py::arg("defer_weight_export") = false, |
| py::arg("operator_export_type") = |
| ::torch::onnx::OperatorExportTypes::ONNX) |
| .def( |
| "_pretty_print_onnx", |
| [](const std::shared_ptr<Graph> g, |
| const std::vector<at::Tensor>& initializers, |
| int64_t onnx_opset_version, |
| bool defer_weight_export, |
| ::torch::onnx::OperatorExportTypes operator_export_type, |
| bool google_printer) { |
| return pretty_print_onnx( |
| g, |
| initializers, |
| onnx_opset_version, |
| defer_weight_export, |
| operator_export_type, |
| google_printer); |
| }, |
| py::arg("initializers"), |
| py::arg("onnx_opset_version") = 0, |
| py::arg("defer_weight_export") = false, |
| py::arg("operator_export_type") = |
| ::torch::onnx::OperatorExportTypes::ONNX, |
| py::arg("google_printer") = false) |
| .def( |
| "inputs", |
| [](Graph& g) { |
| return py::make_iterator(g.inputs().begin(), g.inputs().end()); |
| }) |
| .def( |
| "outputs", |
| [](Graph& g) { |
| return py::make_iterator(g.outputs().begin(), g.outputs().end()); |
| }) |
| // TODO: Iterator invalidation might make this hazardous |
| .def( |
| "nodes", |
| [](Graph& g) { |
| return py::make_iterator(g.nodes().begin(), g.nodes().end()); |
| }) |
| .def("addInput", [](Graph& g) { return g.addInput(); }) |
| .def("copy", [](Graph& g) { return g.copy(); }) |
| .GS(eraseInput) |
| .GS(registerOutput) |
| .def( |
| "create", |
| [](Graph& g, const char* str) { |
| return g.create(Symbol::fromQualString(str)); |
| }) |
| .def( |
| "create", |
| [](Graph& g, const char* str, size_t noutputs) { |
| return g.create(Symbol::fromQualString(str), noutputs); |
| }) |
| .def( |
| "create", |
| [](Graph& g, const char* str, const std::vector<Value*>& inputs) { |
| return g.create(Symbol::fromQualString(str), inputs); |
| }) |
| .def( |
| "create", |
| [](Graph& g, |
| const char* str, |
| const std::vector<Value*>& inputs, |
| size_t noutputs) { |
| return g.create(Symbol::fromQualString(str), inputs, noutputs); |
| }) |
| .def("param_node", [](Graph& g) { return g.block()->param_node(); }) |
| .def("return_node", [](Graph& g) { return g.block()->return_node(); }) |
| .def( |
| "pretty_print", |
| [](Graph& g) { |
| std::ostringstream oss; |
| g.prettyPrint(oss); |
| return oss.str(); |
| }) |
| .GS(createFusionGroup) |
| .def( |
| "createClone", |
| [](Graph& g, Node* n, py::object fn) { |
| return g.createClone( |
| n, [&](Value* e) { return fn(e).cast<Value*>(); }); |
| }) |
| .GS(appendNode) |
| .GS(prependNode) |
| .GS(lint) |
| .GS(insertNode); |
| #undef GS |
| |
| #define VS(name) def(#name, &Value ::name) |
| py::class_<Value, std::unique_ptr<Value, py::nodelete>>(m, "Value") |
| .def( |
| "__repr__", |
| [](Value& n) { |
| std::stringstream ss; |
| ss << n.uniqueName() << " defined in (" << *n.node() << ")"; |
| return ss.str(); |
| }) |
| .VS(type) |
| .VS(setType) |
| .VS(inferTypeFrom) |
| // skip owningGraph because it returns a raw pointer to a otherwise |
| // std::shared_ptr stored graph object, and would cause a double free |
| .VS(unique) |
| .VS(uniqueName) |
| .VS(setUniqueName) |
| .VS(offset) |
| .VS(uses) |
| .VS(replaceAllUsesWith) |
| .def("node", [](Value& v) { return v.node(); }) |
| .def( |
| "setTypeAs", |
| [](Value* node, Value* other) { |
| node->setType(other->type()); |
| return node; |
| }) |
| .VS(copyMetadata) |
| .VS(isTensor); |
| |
| #undef VS |
| |
| py::class_<Block, std::unique_ptr<Block, py::nodelete>>(m, "Block") |
| .def("nodes", [](Block& b) { |
| return py::make_iterator(b.nodes().begin(), b.nodes().end()); |
| }); |
| |
| #define NS(name) def(#name, &Node ::name) |
| py::class_<Node, std::unique_ptr<Node, py::nodelete>>(m, "Node") |
| .def( |
| "__repr__", |
| [](Node& n) { |
| std::stringstream ss; |
| ss << n; |
| return ss.str(); |
| }) |
| .def( |
| "getSourceLocation", |
| [](Node& n) -> py::object { |
| std::stringstream ss; |
| if (auto sl = n.getSourceLocation()) { |
| sl->highlight(ss); |
| return py::str(ss.str()); |
| } else { |
| return py::none(); |
| } |
| }) |
| .def("hasMultipleOutputs", [](Node& n) { return n.outputs().size() > 1; }) |
| .def("outputsSize", [](Node& n) { return n.outputs().size(); }) |
| .NS(kind) |
| .def( |
| "inputs", |
| [](Node& n) { |
| return py::make_iterator(n.inputs().begin(), n.inputs().end()); |
| }) |
| .def( |
| "outputs", |
| [](Node& n) { |
| return py::make_iterator(n.outputs().begin(), n.outputs().end()); |
| }) |
| .def("output", [](Node& n) { return n.output(); }) |
| .NS(addInput) |
| .NS(replaceInput) |
| .NS(replaceInputWith) |
| .NS(replaceAllUsesWith) |
| .NS(insertBefore) |
| .NS(insertAfter) |
| .NS(moveAfter) |
| .NS(moveBefore) |
| .NS(removeInput) |
| .NS(removeAllInputs) |
| .NS(destroy) |
| .NS(hasUses) |
| .NS(eraseOutput) |
| .NS(addOutput) |
| .NS(scopeName) |
| .NS(isNondeterministic) |
| .def( |
| "blocks", |
| [](Node& n) { |
| return py::make_iterator(n.blocks().begin(), n.blocks().end()); |
| }) |
| .NS(addBlock) |
| |
| #define AS(name) def(#name, &Attributes<Node>::name) |
| // methods from Attributes |
| .AS(copyAttributes) |
| .AS(hasAttributes) |
| #undef AS |
| #define AS(name) def(#name, &Attributes<Node>::name##S) |
| // The default method names take Symbol, but the string conversion for |
| // Symbol you to qualify with attr::. This is not very user friendly |
| // for attributes, so expose the string variants instead. |
| .AS(hasAttribute) |
| .AS(kindOf) |
| .AS(removeAttribute) |
| .AS(attributeNames) |
| #undef AS |
| #define CREATE_ACCESSOR(Kind, method) \ |
| def(#method "_", \ |
| [](Node& n, const char* name, Kind##Attr::ValueType v) { \ |
| return n.method##_(Symbol::attr(name), std::move(v)); \ |
| }) \ |
| .def(#method, [](Node& n, const char* name) { \ |
| return n.method(Symbol::attr(name)); \ |
| }) |
| .CREATE_ACCESSOR(Float, f) |
| .CREATE_ACCESSOR(Floats, fs) |
| .CREATE_ACCESSOR(String, s) |
| .CREATE_ACCESSOR(Strings, ss) |
| .CREATE_ACCESSOR(Int, i) |
| .CREATE_ACCESSOR(Ints, is) |
| .CREATE_ACCESSOR(Graph, g) |
| .CREATE_ACCESSOR(Graphs, gs) |
| #undef CREATE_ACCESSOR |
| // Tensor (t_) -- manually written to unwrap the variable into a tensor. |
| .def( |
| "t_", |
| [](Node& n, const char* name, torch::autograd::Variable v) { |
| return n.t_(Symbol::attr(name), v.data()); |
| }) |
| .def( |
| "t", |
| [](Node& n, const char* name) { |
| return torch::autograd::make_variable( |
| n.t(Symbol::attr(name)), /*requires_grad=*/false); |
| }) |
| // Tensors (ts_) -- manually written to unwrap variables into tensors. |
| .def( |
| "ts_", |
| [](Node& n, |
| const char* name, |
| std::vector<torch::autograd::Variable> vs) { |
| std::vector<at::Tensor> tensors; |
| tensors.reserve(vs.size()); |
| for (auto& variable : vs) { |
| tensors.push_back(variable.data()); |
| } |
| return n.ts_(Symbol::attr(name), std::move(tensors)); |
| }) |
| .def( |
| "ts", |
| [](Node& n, const char* name) { |
| auto tensors = n.ts(Symbol::attr(name)); |
| std::vector<torch::autograd::Variable> variables; |
| variables.reserve(tensors.size()); |
| for (auto& tensor : tensors) { |
| variables.push_back(torch::autograd::make_variable( |
| std::move(tensor), /*requires_grad=*/false)); |
| } |
| return variables; |
| }) |
| .def( |
| "z_", |
| [](Node& n, const char* name, at::Tensor v) { |
| return n.t_( |
| Symbol::attr(name), autograd::Variable(v.view({})).data()); |
| }) |
| .def( |
| "z", |
| [](Node& n, const char* name) { return n.t(Symbol::attr(name)); }) |
| .def( |
| "zs_", |
| [](Node& n, const char* name, TensorsAttr::ValueType v) { |
| for (auto& i : v) { |
| i = autograd::Variable(i.view({})).data(); |
| } |
| return n.ts_(Symbol::attr(name), std::move(v)); |
| }) |
| .def( |
| "zs", |
| [](Node& n, const char* name) { return n.ts(Symbol::attr(name)); }) |
| .def( |
| "pyobj", |
| [](Node& n) { |
| return py::handle(n.expect<PythonOp>()->pyobj.get()) |
| .cast<py::object>(); |
| }) |
| .def("cconv", [](Node& n) { return n.expect<PythonOp>()->cconv; }) |
| .def("pyname", [](Node& n) { return n.expect<PythonOp>()->name(); }) |
| .def("scalar_args", [](Node& n) { |
| auto op = n.expect<PythonOp>(); |
| auto scalars = py::list(); |
| auto append = scalars.attr("append"); |
| for (auto& arg : op->scalar_args) { |
| append(py::handle(arg.get())); |
| } |
| return scalars; |
| }); |
| |
| using ::c10::Type; |
| py::class_<Type, std::shared_ptr<Type>>(m, "Type") |
| .def("__repr__", [](Type& t) { return t.python_str(); }) |
| .def( |
| "str", |
| [](Type& t) { |
| std::ostringstream s; |
| s << t; |
| return s.str(); |
| }) |
| .def("kind", [](const Type& t) { return typeKindToString(t.kind()); }) |
| .def("dim", [](const Type& t) { return t.expect<TensorType>()->dim(); }) |
| .def( |
| "sizes", |
| [](Type& t) { return t.expect<CompleteTensorType>()->sizes(); }) |
| .def( |
| "strides", |
| [](Type& t) { return t.expect<CompleteTensorType>()->strides(); }) |
| .def( |
| "contiguous", |
| [](Type& t) { |
| return std::static_pointer_cast<Type>( |
| t.expect<CompleteTensorType>()->contiguous()); |
| }) |
| .def( |
| "scalarType", |
| [](Type& t) { |
| return toString(t.expect<TensorType>()->scalarType()); |
| }) |
| .def( |
| "__eq__", |
| [](std::shared_ptr<Type>& self, std::shared_ptr<Type>& other) { |
| return *self == *other; |
| }) |
| .def( |
| "isSubtypeOf", |
| [](std::shared_ptr<Type>& self, std::shared_ptr<Type> other) { |
| return self->isSubtypeOf(other); |
| }); |
| |
| py::class_<NumberType, Type, std::shared_ptr<NumberType>>(m, "NumberType") |
| .def_static("get", &NumberType::get); |
| py::class_<IntType, Type, std::shared_ptr<IntType>>(m, "IntType") |
| .def_static("get", &IntType::get); |
| py::class_<FloatType, Type, std::shared_ptr<FloatType>>(m, "FloatType") |
| .def_static("get", &FloatType::get); |
| py::class_<DynamicType, Type, std::shared_ptr<DynamicType>>(m, "DynamicType") |
| .def_static("get", &DynamicType::get); |
| py::class_<BoolType, Type, std::shared_ptr<BoolType>>(m, "BoolType") |
| .def_static("get", &BoolType::get); |
| |
| py::class_<TupleType, Type, std::shared_ptr<TupleType>>(m, "TupleType") |
| .def( |
| py::init([](std::vector<TypePtr> a) { return TupleType::create(a); })) |
| .def("elements", [](TupleType& self) { |
| std::vector<TypePtr> types; |
| for (const auto& type : self.elements()) { |
| types.push_back(type); |
| } |
| return types; |
| }); |
| py::class_<ListType, Type, std::shared_ptr<ListType>>(m, "ListType") |
| .def( |
| py::init([](TypePtr a) { return ListType::create(a); })) |
| .def_static("ofInts", &ListType::ofInts) |
| .def_static("ofTensors", &ListType::ofTensors) |
| .def("getElementType", &ListType::getElementType); |
| |
| py::class_<Use>(m, "Use") |
| .def_readonly("user", &Use::user) |
| .def_readonly("offset", &Use::offset); |
| } |
| } // namespace jit |
| } // namespace torch |