Method AddToCostGraphDef

tensorflow/core/graph/costmodel.cc:478–543

}

void CostModel::AddToCostGraphDef(const Graph* graph,
                                  CostGraphDef* cost_graph) const {
  std::vector<const Edge*> inputs;
  std::vector<const Edge*> control_inputs;
  for (const Node* n : graph->nodes()) {
    CostGraphDef::Node* cnode = cost_graph->add_node();
    cnode->set_name(n->name());
    cnode->set_device(n->assigned_device_name());
    cnode->set_id(Id(n));

    inputs.clear();
    inputs.resize(n->num_inputs(), nullptr);
    control_inputs.clear();
    for (const Edge* e : n->in_edges()) {
      if (e->IsControlEdge()) {
        control_inputs.push_back(e);
      } else {
        inputs[e->dst_input()] = e;
      }
    }
    std::sort(control_inputs.begin(), control_inputs.end(),
              [this](Edge const* a, Edge const* b) {
                return Id(a->src()) < Id(b->src());
              });

    for (const Edge* e : inputs) {
      CostGraphDef::Node::InputInfo* input_info = cnode->add_input_info();
      input_info->set_preceding_node(Id(e->src()));
      input_info->set_preceding_port(e->src_output());
    }

    for (int i = 0; i < n->num_outputs(); i++) {
      CostGraphDef::Node::OutputInfo* output_info = cnode->add_output_info();
      int64 alloc_id = AllocationId(n, i);
      int64 alias_to_input = -1;
      for (const Edge* e : inputs) {
        int64 input_alloc_id = AllocationId(e->src(), e->src_output());
        if (input_alloc_id == alloc_id) {
          alias_to_input = e->dst_input();
          break;
        }
      }
      output_info->set_alias_input_port(alias_to_input);
      output_info->set_dtype(MaxMemoryType(n, i));
      *output_info->mutable_shape() = MaxMemoryShape(n, i);
      if (alias_to_input < 0 && IsPersistentTensor(n, alloc_id)) {
        output_info->set_size(0);
      } else {
        output_info->set_size(MaxMemorySize(n, i).value());
      }
    }

    for (const Edge* e : control_inputs) {
      cnode->add_control_input(Id(e->src()));
    }

    cnode->set_temporary_memory_size(TempMemorySize(n).value());
    cnode->set_persistent_memory_size(PersistentMemorySize(n).value());