Function InitMemory

oneflow/core/auto_parallel/auto_memory.cpp:393–511

Use two function

// Use two function
void InitMemory(const OpGraph& op_graph, SbpGraph* sbp_graph, bool nccl_use_compute_stream) {
  // Generate topological data structure for each sbp node
  HashMap<const OpNode*, TopoStruct> op_node2topo_struct;
  std::vector<TopoStruct*> topo_structs;
  std::vector<TopoStruct*> ordered_topo_structs;

  // Traverse all the nodes in the sbp graph
  for (const auto& sbp_node : sbp_graph->GetNodeList()) {
    auto* op_node = sbp_node->GetOperatorNode();
    CHECK(op_node != nullptr)
        << "No proxy node allow at this status. InitMemory() should be run before sbp collector!";
    op_node2topo_struct.insert({op_node, TopoStruct(sbp_node)});
    topo_structs.push_back(&op_node2topo_struct.at(op_node));
  }

  // Construct the map from a lbi to its id, consumers, blob size
  HashMap<LogicalBlobId, int32_t> lbi2id;
  std::vector<std::vector<TopoStruct*>> id2consumer_topo_structs;
  std::vector<int64_t> id2blob_size;

  StraightenOpNodes(op_node2topo_struct, &topo_structs, &lbi2id, &id2consumer_topo_structs,
                    &id2blob_size, &ordered_topo_structs);

  // Mark the memory support, which contains two part:
  // All the non-reusable memory and those blobs which is a part of the maximum reusable memory
  int64_t max_reusable_memory = 0;
  int64_t curr_reusable_memory = 0;
  std::vector<int32_t> id2count(id2blob_size.size(), -1);
  // Blobs born, increase count and memory
  auto GenerateBlobs = [&](TopoStruct* topo_struct) {
    const auto& curr_operator = topo_struct->op_node->op();
    if (topo_struct->is_reusable) {
      for (const auto& obn : curr_operator.output_bns()) {
        const LogicalBlobId& lbi = curr_operator.BnInOp2Lbi(obn);
        int32_t index = lbi2id.at(lbi);
        // Reusable blobs born
        curr_reusable_memory += id2blob_size[index];
        id2count[index] = id2consumer_topo_structs[index].size();
      }
    }
  };
  // Blobs die, decrease count and memory
  auto KillBlobs = [&](TopoStruct* topo_struct) {
    const auto& curr_operator = topo_struct->op_node->op();
    // Those reusable blobs who do not have a consumer would die immediately
    // For example:
    // register_num: 1, op_name:
    // "model.cls_head.loss_func.lm_loss-sparse_softmax_cross_entropy_ms-231-split_softmax_reduce_max_device_stage",
    // blob_name: "mask_0", shape { dim: 2048 dim: 21248 },
    // data_type: kBool, time_shape { dim: 1 dim: 1 }, enable_reuse_mem: true,
    // alloc_before_actor: 369, free_after_actor: 369
    if (topo_struct->is_reusable) {
      for (const auto& obn : curr_operator.output_bns()) {
        const LogicalBlobId& lbi = curr_operator.BnInOp2Lbi(obn);
        int32_t index = lbi2id.at(lbi);
        // Do not have consumer
        if (id2count[index] == 0) {
          // Reusable blobs die