Function InitInOutTopoStructs

oneflow/core/auto_parallel/auto_memory.cpp:264–295

}

void InitInOutTopoStructs(std::vector<TopoStruct*>* topo_structs) {
  // Generate the map from operator names to topological structure
  HashMap<std::string, TopoStruct*> op_name2topo_structs;
  for (auto& topo_struct : *topo_structs) {
    op_name2topo_structs[topo_struct->op_node->op().op_name()] = topo_struct;
  }

  // Traverse the topological structures
  for (auto& this_topo_struct : *topo_structs) {
    auto& node = this_topo_struct->op_node;
    // Initialize input nodes for edges with data
    node->ForEachNodeOnInEdge([&](OpNode* in) {
      // Since we might be looking at a sub-graph of the operator graph.
      // We need to check if the op_node exists in the sub-graph.
      auto it = op_name2topo_structs.find(in->op().op_name());
      if (it != op_name2topo_structs.end()) {
        this_topo_struct->in_topo_structs.insert(it->second);
        it->second->out_topo_structs.insert(this_topo_struct);
      }
    });
    // Initialize input nodes for control edges
    for (const auto& ctrl_in_op_name : node->op().op_conf().ctrl_in_op_name()) {
      auto it = op_name2topo_structs.find(ctrl_in_op_name);
      if (it != op_name2topo_structs.end()) {
        auto& ctrl_in_topo_struct = it->second;
        this_topo_struct->in_topo_structs.insert(ctrl_in_topo_struct);
        // Initialize output nodes for this control edge simultaneously
        ctrl_in_topo_struct->out_topo_structs.insert(this_topo_struct);
      }
    }
  }
}

void ComputeLayer(std::vector<TopoStruct*>* topo_structs) {
  int32_t max_min_layer = -1;