Replaces NcclReduce node with _NcclReduceRecv reusing one input of same device, adds one _NcclReduceSend for each other input.
| 29 | // Replaces NcclReduce node with _NcclReduceRecv reusing one input of same |
| 30 | // device, adds one _NcclReduceSend for each other input. |
| 31 | Status ReplaceReduce(Graph* graph, Node* node) { |
| 32 | string reduction; |
| 33 | TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), "reduction", &reduction)); |
| 34 | DataType dtype; |
| 35 | TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), "T", &dtype)); |
| 36 | int num_devices = node->num_inputs(); |
| 37 | string shared_name = node->name(); |
| 38 | auto make_builder = [&](StringPiece op_name, StringPiece suffix) { |
| 39 | return NodeBuilder(strings::StrCat(shared_name, suffix), op_name) |
| 40 | .Attr("reduction", reduction) |
| 41 | .Attr("num_devices", num_devices) |
| 42 | .Attr("shared_name", shared_name) |
| 43 | .Attr("T", dtype); |
| 44 | }; |
| 45 | std::vector<Node*> control_inputs; |
| 46 | for (const auto& edge : node->in_edges()) { |
| 47 | if (edge->IsControlEdge()) { |
| 48 | control_inputs.push_back(edge->src()); |
| 49 | } |
| 50 | } |
| 51 | std::vector<NodeBuilder::NodeOut> out_nodes; |
| 52 | for (const auto& edge : node->out_edges()) { |
| 53 | out_nodes.emplace_back(edge->dst(), edge->dst_input()); |
| 54 | } |
| 55 | int recv_dev = node->assigned_device_name_index(); |
| 56 | NodeBuilder recv_builder = |
| 57 | make_builder("_NcclReduceRecv", "Recv").ControlInputs(control_inputs); |
| 58 | bool recv_input_set = false; |
| 59 | int send_counter = 0; |
| 60 | for (const auto& edge : node->in_edges()) { |
| 61 | Node* src_node = edge->src(); |
| 62 | if (edge->IsControlEdge()) { |
| 63 | continue; |
| 64 | } |
| 65 | int send_dev = src_node->assigned_device_name_index(); |
| 66 | if (!recv_input_set && send_dev == recv_dev) { |
| 67 | recv_builder.Input(src_node); |
| 68 | recv_input_set = true; |
| 69 | continue; |
| 70 | } |
| 71 | auto send_builder = make_builder("_NcclReduceSend", |
| 72 | strings::StrCat("Send_", ++send_counter)) |
| 73 | .Input(src_node) |
| 74 | .ControlInputs(control_inputs); |
| 75 | Node* send_node = nullptr; |
| 76 | TF_RETURN_IF_ERROR(send_builder.Finalize(graph, &send_node)); |
| 77 | send_node->set_assigned_device_name_index(send_dev); |
| 78 | // Send nodes don't have any outputs and therefore have no data dependencies |
| 79 | // to the outputs of the graph. We add a control dependency to the receive |
| 80 | // node so that those 'dangling' nodes are run. |
| 81 | // TODO(b/67027412): Avoid these cross-device control edges. |
| 82 | for (const auto& out_node : out_nodes) { |
| 83 | graph->AddControlEdge(send_node, out_node.node); |
| 84 | } |
| 85 | } |
| 86 | if (!recv_input_set) { |
| 87 | return errors::InvalidArgument( |
| 88 | "No input tensor uses the same device as the NcclReduce op"); |
no test coverage detected