Method Run

tensorflow/compiler/jit/async_io_conversion_pass.cc:373–430

}

Status AsyncIoConversionPassImpl::Run() {
  if (VLOG_IS_ON(1)) {
    DumpGraphToFile("before_async_io_conversion_pass", *graph_, flib_def_);
  }

  // Collect output edges into `cluster_out_edges_`.
  CollectOutputEdgesForClusters();

  // For each cluster, insert AsyncOut operations.
  for (auto& p : cluster_out_edges_) {
    const string& cluster_name = p.first;
    const std::vector<const Edge*>& cluster_outputs = p.second;

    absl::flat_hash_set<OutputTensor, OutputTensor::Hash>
        cluster_outputs_to_convert;
    for (const Edge* e : cluster_outputs) {
      if (e->IsControlEdge()) {
        // Do not support control edges now.
        continue;
      }

      if (IsConvertibleAndProfitable(*e, enable_full_async_io_)) {
        cluster_outputs_to_convert.insert({e->src(), e->src_output()});
      }
    }

    if (cluster_outputs_to_convert.empty()) {
      // No output edges to be converted for the current cluster.
      continue;
    }

    std::vector<string> async_out_tensor_names;
    std::vector<Node*> async_out_sends;
    std::vector<Node*> async_out_recvs;
    for (const OutputTensor& o : cluster_outputs_to_convert) {
      // Remove unnecessary input control edges to users of `n:oidx`.
      //
      // TODO: generalize it by replacing the control edge also with
      // AsyncOutSend and AsyncOutRecv.
      RemoveUnnecessaryControlEdgesToUsersOf(o);

      // n:oidx -> AsyncOutSend
      //           AsyncOutRecv -> uses of `n:oidx`
      InsertAsyncOutSendAndRecv(o, &async_out_tensor_names, &async_out_sends,
                                &async_out_recvs);
    }

    // Insert AsyncOutInit and AsyncOutDone.
    InsertAsyncOutInitAndDone(cluster_name, async_out_tensor_names,
                              async_out_sends, async_out_recvs);
  }

  if (VLOG_IS_ON(1)) {
    DumpGraphToFile("async_io_conversion_pass", *graph_, flib_def_);
  }

  return Status::OK();
}