blob: 598d3fff2264d21c154bf6dba0f5970619e516cb [file] [log] [blame]
#ifndef NO_PYTHON
#include "torch/csrc/python_headers.h"
#endif
#include "ir.h"
#include "torch/csrc/autograd/function.h"
#include <iostream>
#include <unordered_map>
#include <unordered_set>
#include <set>
#include <stack>
#include <sstream>
#include <algorithm>
#include <string>
#ifndef NO_PYTHON
#include "torch/csrc/utils/auto_gil.h"
#include "torch/csrc/utils/python_strings.h"
#include "pybind11/pybind11.h"
namespace py = pybind11;
namespace torch { namespace jit {
std::string getPythonName(const PyObject* obj, bool is_legacy) {
AutoGIL gil;
if (is_legacy) {
return std::string(obj->ob_type->tp_name);
} else {
auto v = py::getattr(const_cast<PyObject*>(obj), "__name__", py::str("<python_value>"));
return py::str(v);
}
}
std::ostream& printPyObject(std::ostream & out, const THPObjectPtr& obj) {
AutoGIL gil;
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 (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());
}
}
std::string PythonOp::name() const {
return getPythonName(pyobj.get(),is_legacy);
}
void PythonOp::cloneFrom(Node * other_) {
Node::cloneFrom(other_);
auto other = other_->cast<PythonOp>();
this->cconv = other->cconv;
this->is_legacy = other->is_legacy;
Py_INCREF(other->pyobj.get());
this->pyobj = THPObjectPtr(other->pyobj.get());
this->var_flags = other->var_flags;
for(auto & sa : other->scalar_args) {
Py_INCREF(sa.get());
this->scalar_args.emplace_back(sa.get());
}
this->tracing_autograd_python_function =
other->tracing_autograd_python_function;
}
}} // namespace torch::jit
#else
namespace torch { namespace jit {
std::ostream& printPyObject(std::ostream & out, const THPObjectPtr& obj) {
throw std::runtime_error("Trying to print Python object from a C++ build");
}
std::string PythonOp::name() const {
throw std::runtime_error("Trying to call PythonOp::name from a C++ build");
return std::string();
}
}} // namespace torch::jit
#endif
namespace torch { namespace jit {
// Sigh, see https://stackoverflow.com/questions/8016780/undefined-reference-to-static-constexpr-char
constexpr Symbol CppOp::Kind;
constexpr Symbol PythonOp::Kind;
constexpr int max_tensor_display_size = 10;
void printValueRef(std::ostream & out, const Value * n) {
out << "%" << n->uniqueName();
}
template <typename T>
std::ostream& operator<<(std::ostream & out, const std::vector<T> & nodes) {
out << at::ArrayRef<T>{nodes};
return out;
}
template <typename T>
std::ostream& operator<<(std::ostream & out, const at::ArrayRef<T> & nodes) {
size_t i = 0;
for(auto n : nodes) {
if(i++ > 0)
out << ", ";
printValueRef(out, n);
}
return out;
}
std::string CppOp::name() const {
return fn->name();
}
struct const_value_list_with_types {
const std::vector<const Value*>& values;
bool use_newlines;
const_value_list_with_types(const std::vector<const Value*>& values, bool use_newlines = false)
: values(values), use_newlines(use_newlines) {}
};
std::ostream& operator<<(std::ostream & out, const_value_list_with_types l) {
size_t i = 0;
size_t prev_stage = 0;
for(auto n : l.values) {
if(i++ > 0) {
if (l.use_newlines) {
// TODO: Indent here is hard-coded for "graph(": un-hard-code it
out << "\n ";
if (n->stage() != prev_stage) {
out << "-------- stage " << n->stage() << " --------\n ";
prev_stage = n->stage();
}
} else {
out << ", ";
}
}
printValueRef(out, n);
out << " : ";
out << *n->type();
}
return out;
}
template<typename T>
void printPrimList(std::ostream & out, const std::vector<T> & items) {
out << "[";
int i = 0;
for(auto & item : items) {
if(i++ > 0)
out << ", ";
out << item;
}
out << "]";
}
void printAttributes(std::ostream & out, const Node * n) {
out << "[";
auto names = n->attributeNames();
int i = 0;
for(auto name : names) {
if(i++ > 0)
out << ", ";
// TODO: debugging mode to see the qualifier. We definitely
// don't want to print the qualifier since it should always
// be attribute, but you might be able to track down a weird
// bug by printing it out.
out << name.toUnqualString() <<"=";
switch(n->kindOf(name)) {
case AttributeKind::f:
out << n->f(name);
break;
case AttributeKind::fs:
printPrimList(out,n->fs(name));
break;
case AttributeKind::i:
out << n->i(name);
break;
case AttributeKind::is:
printPrimList(out,n->is(name));
break;
case AttributeKind::s:
out << n->s(name);
break;
case AttributeKind::ss:
printPrimList(out,n->ss(name));
break;
case AttributeKind::t:
{
at::Tensor t = n->t(name);
// 1-elem tensors are usually boxed scalars, so print them like it
if (t.numel() == 1) {
auto scalar = at::Scalar(t.view({})).local();
out << "{";
if (scalar.isFloatingPoint()) {
out << scalar.toDouble();
} else {
out << scalar.toLong();
}
out << "}";
} else if (t.numel() <= max_tensor_display_size) {
// TODO: This is awful code. Also it doesn't work on Windows.
std::ostringstream tensor_ss;
tensor_ss << t;
std::string tensor_s{tensor_ss.str()};
// Remove newlines
std::replace(tensor_s.begin(), tensor_s.end(), '\n', ' ');
out << tensor_s;
} else {
out << "<Tensor>";
}
break;
}
case AttributeKind::ts:
out << "[<Tensors>]";
break;
case AttributeKind::g:
out << "<Graph>";
break;
case AttributeKind::gs:
out << "[<Graphs>]";
break;
}
}
out << "]";
}
static std::ostream & indent(std::ostream & out, size_t level) {
for(size_t i = 0; i < level; ++i)
out << " ";
return out;
}
std::ostream& printNode(std::ostream & out, size_t level, const Node * n, std::vector<const Node*> * groups) {
auto outputs = n->outputs();
indent(out, level) << const_value_list_with_types(outputs);
out << " = ";
IR_IFM_CONST(n,PythonOp)
out << "^" << value->name();
out << "(";
int i = 0;
for (auto& scalar : value->scalar_args) {
if (i++ > 0)
out << ", ";
printPyObject(out, scalar);
}
out << ")";
IR_ELSEIFM_CONST(CppOp)
out << "CppOp[" << value->name() << "]";
IR_ELSE()
if(n->hasAttribute(attr::Subgraph)) {
if(groups) {
out << n->kind().toQualString() << "_" << groups->size();
groups->push_back(n);
} else {
out << n->kind().toQualString() << "[" << *n->g(attr::Subgraph) << "]";
}
} else {
out << n->kind().toQualString();
if(n->hasAttributes()) {
printAttributes(out,n);
}
}
IR_END()
out << "(" << n->inputs() << ")";
std::string scopeName = n->scopeName();
if (scopeName.empty()) {
out << "\n";
}
else {
out << ", ";
out << "scope: " << scopeName << "\n";
}
for(size_t i = 0; i < n->blocks().size(); ++i) {
auto b = n->blocks()[i];
indent(out, level + 1) << "block" << i << "(" << const_value_list_with_types(b->inputs(), false) << ") {\n";
for(auto n : b->nodes()) {
printNode(out, level + 2, n, groups);
}
indent(out, level + 2) << "-> (" << b->outputs() << ")\n";
indent(out, level + 1) << "}\n";
}
return out;
}
std::ostream& operator<<(std::ostream & out, const Node & n) {
return printNode(out, 0, &n, nullptr);
}
std::ostream& operator<<(std::ostream & out, const Graph & g) {
out << "graph(" << const_value_list_with_types(g.inputs(), true) << ") {\n";
std::vector<const Node*> groups;
size_t prev_stage = 0;
for(auto n : g.nodes()) {
if (n->stage() != prev_stage) {
out << " ---------------- stage " << n->stage() << " ----------------\n";
prev_stage = n->stage();
}
printNode(out, 1, n, &groups);
}
out << " return (" << g.outputs() << ");\n}\n";
size_t i = 0;
for(auto fg : groups) {
out << "with " << fg->kind().toQualString() << "_" <<i++ << " = " << *fg->g(attr::Subgraph);
}
/*
// Uncomment this to debug all_nodes issues
{
out << "\n";
out << "all_nodes:\n";
for (auto& n : g.all_nodes) {
printNode(out, const_cast<Node*>(n), nullptr);
}
}
*/
return out;
}
static void checkSameDevice(const Node* node) {
bool has_device = false;
int device;
auto checkValue = [&](const Value* v) {
if(TensorType* type = v->type()->cast<TensorType>()) {
if(!has_device) {
has_device = true;
device = type->device();
} else {
JIT_ASSERT(device == type->device());
}
}
};
for(auto input : node->inputs()) {
checkValue(input);
}
for(auto output : node->outputs()) {
checkValue(output);
}
}
using node_set = std::set<const Node*>;
#define ALL_OF(container) container.begin(), container.end()
// These functions purposely operate on the internal members directly, to force
// you to think about how the invariants change if you change the data
// representation (even if the external API does not change.)
// NB: This assert is written to assume you don't have any unattached
// nodes. Unattached nodes can occur while manipulations to the
// graph are occurring.
void Node::lint() const {
// Node invariants
// - if node should live in list, nodes_iter is consistent
// - Inputs are all marked as a use by the nodes they refer to
// - Stage is consistent (stage is >= all input stages)
// - Owning graph is non-null and consistent
// - The "Select" invariant, when the node is MultiReturn
//
// The handle invariant:
// If a node takes a handle as an input, it is always the
// LAST input of the node. There is at most one handle input.
{
size_t i = 0;
for (auto input : inputs_) {
// WARNING: O(n^2)
JIT_ASSERT(std::find(ALL_OF(input->uses_), Use(const_cast<Node*>(this), i)) != input->uses_.end());
JIT_ASSERT(stage_ >= input->stage_);
JIT_ASSERT(graph_->all_nodes.count(this) == 1);
// Handle invariant
if (i != inputs_.size() - 1) {
JIT_ASSERT(input->type()->kind() != TypeKind::HandleType);
}
i++;
}
}
for(auto o : outputs()) {
size_t i = 0;
for (auto use : o->uses()) {
// Use invariants
// - Use is consistent with inputs
// - Every user node is live (checked in Graph)
JIT_ASSERT(use.user->inputs_[use.offset] == o);
i++;
}
}
// Node subclass invariants
// - Return uses is zero
// - Param inputs is zero
// - Select inputs is one
// - Python operator cconv is correct
IR_IF(this,Constant)
JIT_ASSERT(inputs_.size() == 0);
IR_ELSEIF(Return)
JIT_ASSERT(outputs().size() == 0);
IR_ELSEIF(Param)
JIT_ASSERT(inputs_.size() == 0);
IR_ELSEIFM_CONST(PythonOp)
std::size_t n_scalars = 0, n_tensors = 0;
for (auto c : value->cconv) {
if (c == 's') {
n_scalars++;
} else if (c == 't') {
n_tensors++;
} else {
JIT_ASSERT(0);
}
JIT_ASSERT(static_cast<bool>(value->pyobj));
}
JIT_ASSERT(n_scalars == value->scalar_args.size());
JIT_ASSERT(n_tensors == inputs_.size());
IR_ELSEIFM_CONST(CppOp)
// TODO: add invariants
IR_ELSEIF(Eval)
// TODO: add invariants
// TODO: It's not good for these ops to be top-level, it makes cases longer.
IR_ELSEIF(FusionGroup)
checkSameDevice(value);
// TODO: Typecheck the parameters
value->g(attr::Subgraph)->lint();
IR_END()
}
// TODO: When lint fails, give better indication about which
// instruction triggered the failure.
void Graph::lint() const {
// Graph invariants
// Uncomment the following to see the graph
// std::cout << *const_cast<Graph*>(this);
// nodes
// - nodes_ is a valid topological ordering for inputs
// - No repeated nodes
// - Params and return do NOT occur in nodes
// - next_unique_ is greater than all uniques in graph
// - uniques in all_nodes are unique
// - every use will occur later in the topsort
struct LintScope {
LintScope() {}
LintScope(std::unique_ptr<LintScope> parent)
: parent(std::move(parent)) {}
bool contains(const Value * v) {
return values.count(v) > 0 || (parent && parent->contains(v));
}
bool contains(const Node * n) {
return nodes.count(n) > 0 || (parent && parent->contains(n));
}
void insert(const Value * v) {
JIT_ASSERT(!contains(v));
values.insert(v);
}
void insert(const Node * n) {
JIT_ASSERT(!contains(n));
nodes.insert(n);
}
std::unique_ptr<LintScope> parent;
private:
std::unordered_set<const Value*> values;
std::unordered_set<const Node*> nodes;
};
// Struct enables mutual recursion in linting methods.
// Putting it inside Graph::lint enables access to private Graph members
struct LintImpl {
LintImpl(const Graph & g)
: g(g)
, scope(new LintScope())
, all_nodes_set(ALL_OF(g.all_nodes)) {} // NB: all_nodes is *unordered*
const Graph & g;
std::unique_ptr<LintScope> scope;
std::unordered_set<size_t> seen_uniques;
std::unordered_map<const Node*, int64_t> anticipated_uses;
node_set all_nodes_set;
node_set sum_set;
void check_value(const Value* v) {
scope->insert(v);
auto b2 = seen_uniques.insert(v->unique());
JIT_ASSERT(b2.second); // insertion took place
JIT_ASSERT(v->unique() < g.next_unique_);
for (auto use : v->uses()) {
JIT_ASSERT(!scope->contains(use.user));
JIT_ASSERT(g.all_nodes.count(use.user) == 1);
anticipated_uses[use.user]++; // int default constructs to 0
}
}
void check_node(const Node* n) {
for (auto input : n->inputs_) {
if (!scope->contains(input)) {
JIT_ASSERTM(0, "%%%d not in scope", input->unique());
}
}
JIT_ASSERT(anticipated_uses[n] == static_cast<int64_t>(n->inputs_.size()));
anticipated_uses[n] = -1; // we saw the anticipated user!
scope->insert(n);
for(auto block : n->blocks()) {
std::unique_ptr<LintScope> new_scope(new LintScope(std::move(scope)));
scope = std::move(new_scope);
check_block(block);
scope = std::move(scope->parent);
}
size_t i = 0;
for(auto o : n->outputs()) {
JIT_ASSERT(o->node() == n);
JIT_ASSERT(i++ == o->offset_);
check_value(o);
}
n->lint();
}
void check_block(const Block *b) {
for (auto input : b->inputs()) {
check_value(input);
JIT_ASSERT(input->node()->kind_ == prim::Param);
}
for (auto n : b->nodes()) {
JIT_ASSERT(n->kind_ != prim::Param);
JIT_ASSERT(n->kind_ != prim::Return);
check_node(n);
}
JIT_ASSERT(b->output_->kind() == prim::Return);
check_node(b->output_);
// all_nodes
// - inputs_, output_ and nodes_ are all included in all_nodes
// - all_nodes does not contain dead nodes??? (likely to be temporarily
// suspended). Weaker: all_nodes contains all inputs and returns
// - only one return node???
node_set nodes_set(ALL_OF(b->nodes()));
node_set inputs_set {b->input_};
node_set output_set {b->output_};
// TODO: Make a more type safe std::includes wrapper which disallows use on
// non-ordered containers
JIT_ASSERT(std::includes(ALL_OF(all_nodes_set), ALL_OF(nodes_set)));
JIT_ASSERT(std::includes(ALL_OF(all_nodes_set), ALL_OF(inputs_set)));
JIT_ASSERT(std::includes(ALL_OF(all_nodes_set), ALL_OF(output_set)));
sum_set.insert(ALL_OF(nodes_set));
sum_set.insert(ALL_OF(inputs_set));
sum_set.insert(ALL_OF(output_set));
}
void check_graph() {
node_set all_nodes_set(ALL_OF(g.all_nodes)); // NB: all_nodes is *unordered*
check_block(g.block_);
for (auto kv : anticipated_uses) {
JIT_ASSERT(kv.second == -1);
}
// graph->stage() should be equal to max(node.stage for node in graph->nodes())
if (g.nodes().begin() == g.nodes().end()) {
JIT_ASSERT(g.stage() == 0);
} else {
JIT_ASSERT(g.stage() == g.nodes().rbegin()->stage());
}
JIT_ASSERT(std::includes(ALL_OF(sum_set), ALL_OF(all_nodes_set)));
}
};
LintImpl(*this).check_graph();
}
void Graph::dump() const {
std::cout << *this << "\n";
}
void LintGraph(std::shared_ptr<Graph>& graph) {
graph->lint();
}
void Block::cloneFrom(Block * src, std::function<Value*(Value*)> outer_map) {
std::unordered_map<Value*, Value*> local_map;
auto env = [&](Value * v) {
auto it = local_map.find(v);
if(it != local_map.end())
return it->second;
return outer_map(v);
};
auto graph = owningGraph();
for(auto input : src->inputs()) {
local_map[input] = this->addInput()->copyMetadata(input)->setStage(input->stage());
graph->setStage(std::max(graph->stage(), input->stage()));
}
for(auto node : src->nodes()) {
auto new_node = this->appendNode(graph->createClone(node, env));
new_node->setStage(node->stage());
graph->setStage(std::max(graph->stage(), node->stage()));
for(size_t i = 0; i < node->outputs().size(); ++i) {
auto oo = node->outputs()[i];
auto no = new_node->outputs()[i];
local_map[oo] = no;
no->copyMetadata(oo);
no->setStage(oo->stage());
}
}
for(auto output : src->outputs()) {
this->registerOutput(env(output));
}
}
std::shared_ptr<Graph> Graph::copy() {
auto new_g = std::make_shared<Graph>();
auto env = [](Value *) -> Value* {
barf("Graph::copy() encountered a use of a value not in scope. Run lint!");
};
new_g->block()->cloneFrom(this->block(), env);
return new_g;
}
}} // namespace torch::jit