Function SwappingPass

tensorflow/core/grappler/optimizers/memory_optimizer.cc:1140–1261

}

bool SwappingPass(RewriterConfig::MemOptType optimization_level,
                  Cluster* cluster, GrapplerItem* item,
                  std::unordered_set<string>* skip_list) {
  std::unordered_map<NodeDef*, SwapInfo> nodes_to_swap;
  if (optimization_level == RewriterConfig::DEFAULT_MEM_OPT ||
      optimization_level == RewriterConfig::SWAPPING_HEURISTICS ||
      optimization_level == RewriterConfig::HEURISTICS) {
    // Use heuristics to figure out what needs to be swapped;
    IdentifySwappingCandidates(cluster, item, skip_list, &nodes_to_swap);
  }
  // Look for manual annotatations in the graph.
  for (auto& node : *item->graph.mutable_node()) {
    if (node.attr().count("_swap_to_host") != 0) {
      SwapInfo& swap_info = nodes_to_swap[&node];
      const AttrValue& val = node.attr().at("_swap_to_host");
      if (val.has_list()) {
        for (int64 input_id : val.list().i()) {
          swap_info.inputs_to_swap.push_back(input_id);
        }
      } else {
        int64 input_id = val.i();
        swap_info.inputs_to_swap.push_back(input_id);
      }
    }
  }
  if (nodes_to_swap.empty()) {
    // Nothing to do.
    return false;
  }

  // Estimate the size of the data to swap for each node.
  GraphProperties properties(*item);
  if (!properties
           .InferStatically(/*assume_valid_feeds=*/true,
                            /*aggressive_shape_inference=*/false,
                            /*include_tensor_values=*/false)
           .ok()) {
    return false;
  }
  for (auto& swap : nodes_to_swap) {
    const NodeDef* node = swap.first;
    const std::vector<OpInfo::TensorProperties>& props =
        properties.GetInputProperties(node->name());
    SwapInfo& swap_info = swap.second;
    int64 bytes_to_swap = 0;
    for (int64 input_id : swap_info.inputs_to_swap) {
      const OpInfo::TensorProperties& t = props[input_id];
      bytes_to_swap += CalculateTensorSize(t);
    }
    // Let's assume we're going to swap over PCIe running at 16 GBps.
    swap_info.time_to_swap = bytes_to_swap / 16;
  }

  std::unordered_map<const NodeDef*, Costs::NanoSeconds> execution_times;
  if (!EstimateEarliestExecutionTimes(*item, cluster, &execution_times).ok()) {
    return false;
  }