// returns a vector `ret` such that transposing by `ret` is equivalent // to transposing by `t1` and then by `t2` std::vector<int64_t> composeTransposes(const std::vector<int64_t> & t1, const std::vector<int64_t> & t2) { JIT_ASSERT(t1.size() == t2.size()); std::vector<int64_t> ret; for (size_t i = 0; i < t1.size(); i++) { JIT_ASSERT( t1[i] < int64_t(t2.size())); JIT_ASSERT(t2[t1[i]] < int64_t(t2.size())); ret.push_back(t2[t1[i]]); } return ret; }
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_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))); }
static std::array<int64_t, 2> as_array(at::IntList sizes) { JIT_ASSERT(sizes.size() == 2); std::array<int64_t, 2> arr; arr[0] = sizes[0]; arr[1] = sizes[1]; return arr; }
void check_block(const Block *b) { for (auto input : b->inputs()) { check_value(input); JIT_ASSERT(input->node()->kind_ == kParam); } for (auto n : b->nodes()) { JIT_ASSERT(n->kind_ != kParam); JIT_ASSERT(n->kind_ != kReturn); check_node(n); } JIT_ASSERT(b->output_->kind() == kReturn); 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)); }
// TODO: I'm not entirely sure why this can't be in the header... bool micropb_callback_string_from_tensor(pb_ostream_t *stream, const pb_field_t *field, void * const *arg) { at::Tensor* t = static_cast<at::Tensor*>(*arg); JIT_ASSERT(t->is_contiguous()); // Packed array format! pb_encode_tag_for_field(stream, field); pb_encode_string(stream, (pb_byte_t*)(t->data_ptr()), t->type().elementSizeInBytes()*t->numel()); return true; }
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(); }
std::vector<Value*> Method::emit_call_to(SourceRange loc, Method & callee, ArrayRef<Value*> inputs) { ensureTensors(loc, inputs); JIT_ASSERT(!executor); try { callee.ensure_defined(); } catch (RecursiveMethodCallError&) { throw ErrorReport(loc) << " method '" << callee.name() << "' is called recursively involving this call site. Recursive calls are not supported"; } auto fn = callee.graph(); ensureSizeMatches(loc, callee.num_inputs(), inputs.size(), "inputs"); std::vector<Value*> all_inputs = inputs; // parameters to callee method (which become parameters to _this_ method // if they were not already) for(at::Tensor* member : callee.member_inputs) { all_inputs.push_back(get_or_add_parameter(member)); } return inlineCallTo(*graph(), *callee.graph(), all_inputs); }
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); } }
void fuseBroadcast(std::shared_ptr<Graph>& graph) { for(auto n : graph->nodes()) { // Can't fuse into nodes that don't support broadcasting if (!isBroadcasting(n)) continue; // If the node already broadcasts, can't "rebroadcast" // TODO: Actually, maybe you can, if there is a broadcast for some // dims, and then another broadcast for the rest. But this will // never happen in practice so I didn't implement it. if (n->hasAttribute(kbroadcast) && n->i(kbroadcast)) continue; JIT_ASSERT(!n->hasAttribute(kaxis)); auto input_index = n->inputs().size() - 1; auto* expanded_rhs = n->input(input_index)->node(); // The expanded_rhs input isn't actually an expand, so no fusion available if (expanded_rhs->kind() != kExpand) continue; auto* unexpanded_rhs = expanded_rhs->input(); // We need to know what the type pre-expand is. We should basically // always have this information (because expands are only ever traced, // not generated from symbolic), but if for some reason we don't // have it, we need to skip. if (!unexpanded_rhs->isTensor()) continue; // Not all broadcasts are supported by ONNX broadcast. if (!fusibleExpandTo(unexpanded_rhs->type()->expect<TensorType>()->sizes(), // from expanded_rhs->output()->type()->expect<TensorType>()->sizes()) // to ) continue; n->replaceInput(input_index, unexpanded_rhs); n->i_(kbroadcast, 1); if (!expanded_rhs->hasUses()) { expanded_rhs->destroy(); } } }
void insert(const Node * n) { JIT_ASSERT(!contains(n)); nodes.insert(n); }
void insert(const Value * v) { JIT_ASSERT(!contains(v)); values.insert(v); }
// 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(kSubgraph)->lint(); IR_END() }