1 GraphInfo g_info;
 2 if (!opts.control_flow_added) {
 3   // Add the "code" for distributed execution of control flow. Code is
 4   // added only for the frames that are placed on multiple devices. The
 5   // new graph is an equivalent transformation of the original graph and
 6   // has the property that it can be subsequently partitioned arbitrarily
 7   // (down to the level of individual device) for distributed execution.
 8   status = AddControlFlow(opts, g, &g_info);
 9   if (!status.ok()) return status;
10 }

1 // MemoryType is used to describe whether input or output Tensors of
2 // an OpKernel should reside in "Host memory" (e.g., CPU memory) or
3 // "Device" Memory (CPU memory for CPU devices, GPU memory for GPU
4 // devices).
5 enum MemoryType {
6   DEVICE_MEMORY = 0,
7   HOST_MEMORY = 1,
8 };

 1 #define REGISTER_GPU_KERNEL(type)                                2   REGISTER_KERNEL_BUILDER(Name("Reshape")                        3                               .Device(DEVICE_GPU)                4                               .HostMemory("shape")               5                               .TypeConstraint<type>("T")         6                               .TypeConstraint<int32>("Tshape"),  7                           ReshapeOp);                            8   REGISTER_KERNEL_BUILDER(Name("Reshape")                        9                               .Device(DEVICE_GPU)               10                               .HostMemory("shape")              11                               .TypeConstraint<type>("T")        12                               .TypeConstraint<int64>("Tshape"), 13                           ReshapeOp);

 1 bool DataTypeAlwaysOnHost(DataType dt) {
 2   // Includes DT_STRING and DT_RESOURCE.
 3   switch (dt) {
 4     case DT_STRING:
 5     case DT_STRING_REF:
 6     case DT_RESOURCE:
 7       return true;
 8     default:
 9       return false;
10   }
11 }

 1       // Check whether there is already a send/recv pair transferring
 2       // the same tensor/control from the src to dst partition.
 3       const bool on_host = IsDstInputOnHost(edge, g_info);
 4       DupRecvKey key{src->id(), edge->src_output(), dst_graph, on_host};
 5       auto iter = dup_recv.find(key);
 6       if (iter != dup_recv.end()) {
 7         // We found one. Reuse the data/control transferred already.
 8         const string& recv_node_name = iter->second.recv->name();
 9         if (edge->IsControlEdge()) {
10           AddInput(dst_def, recv_node_name, Graph::kControlSlot);
11         } else {
12           AddInput(dst_def, recv_node_name, 0);
13         }
14         ref_control_inputs.push_back(recv_node_name);
15 
16         // We want the start_time for the recv to be the smallest of the start
17         // times of it‘s consumers. So we update this whenever we use a recv,
18         // and write it out to the attribute at the end of the subroutine
19         if (iter->second.start_time > recv_start_time) {
20           iter->second.start_time = recv_start_time;
21         }
22         continue;
23       }

 1   const FunctionLibraryDefinition* flib_def = opts.flib_def;
 2   if (flib_def == nullptr) {
 3     flib_def = &g->flib_def();
 4   }
 5 
 6   // Set versions, function library and send/recv incarnation.
 7   for (auto& it : *partitions) {
 8     GraphDef* gdef = &it.second;
 9     *gdef->mutable_versions() = g->versions();
10     // Prune unreachable functions from `flib_def` before adding them to `gdef`.
11     *gdef->mutable_library() = flib_def->ReachableDefinitions(*gdef).ToProto();
12 
13     // Traverse the graph to fill every send/recv op‘s incarnation
14     // information.
15     SetIncarnation(opts, gdef);
16   }

 1       // Need to split edge by placing matching send/recv nodes on
 2       // the src/dst sides of the edge.
 3       NodeDef* send = AddSend(opts, g_info, src_graph, edge, send_from,
 4                               send_start_time, &status);
 5       if (!status.ok()) return status;
 6 
 7       NodeDef* real_recv = nullptr;
 8       NodeDef* recv =
 9           AddRecv(opts, g_info, dst_graph, edge, &real_recv, &status);
10       if (!status.ok()) return status;

// Fix up the control flow edge.
// NOTE(yuanbyu): ‘real_recv‘ must be the real recv node.
if (src_graph == dst_graph) {
  // For same device send/recv, add a control edge from send to recv.
  // This prevents the asynchronous recv kernel from being scheduled
  // before the data is available.
  AddInput(real_recv, send->name(), Graph::kControlSlot);
}