| 736 | } |
| 737 | |
| 738 | Status BuildSwapPair(NodeDef* node, int input_to_swap, |
| 739 | const std::unordered_map<string, const NodeDef*>& name_map, |
| 740 | GraphDef* graph, |
| 741 | std::pair<NodeDef*, NodeDef*>* swap_pair) { |
| 742 | string task, device; |
| 743 | if (!DeviceNameUtils::SplitDeviceName(node->device(), &task, &device) || |
| 744 | !absl::StrContains(device, DEVICE_GPU)) { |
| 745 | return errors::InvalidArgument("Can't swap input ", input_to_swap, |
| 746 | " of node ", node->name(), |
| 747 | " since it is not on GPU"); |
| 748 | } |
| 749 | const OpDef* op_def; |
| 750 | TF_RETURN_IF_ERROR(OpRegistry::Global()->LookUpOpDef(node->op(), &op_def)); |
| 751 | DataType input_type; |
| 752 | TF_RETURN_IF_ERROR( |
| 753 | InputTypeForNode(*node, *op_def, input_to_swap, &input_type)); |
| 754 | if (IsRefType(input_type)) { |
| 755 | return errors::InvalidArgument("Can't swap input ", input_to_swap, |
| 756 | " of node ", node->name(), |
| 757 | " since it expects a reference"); |
| 758 | } |
| 759 | |
| 760 | string tensor_to_swap = strings::StrCat(node->name(), "_", input_to_swap); |
| 761 | string swap_out_name = strings::StrCat("swap_out_", tensor_to_swap); |
| 762 | string swap_in_name = strings::StrCat("swap_in_", tensor_to_swap); |
| 763 | if (name_map.find(swap_out_name) != name_map.end() || |
| 764 | name_map.find(swap_in_name) != name_map.end()) { |
| 765 | return errors::InvalidArgument("Input ", input_to_swap, " of node ", |
| 766 | node->name(), " is already swapped"); |
| 767 | } |
| 768 | |
| 769 | // Force the tensor to be copied to cpu. |
| 770 | NodeDef* swap_out_node = graph->add_node(); |
| 771 | swap_out_node->set_name(swap_out_name); |
| 772 | swap_out_node->set_op("_CopyFromGpuToHost"); |
| 773 | |
| 774 | // Force the tensor to be restored to the device. |
| 775 | NodeDef* swap_in_node = graph->add_node(); |
| 776 | swap_in_node->set_name(swap_in_name); |
| 777 | swap_in_node->set_op("_CopyFromHostToGpu"); |
| 778 | *swap_in_node->add_input() = swap_out_node->name(); |
| 779 | |
| 780 | // Colocate the swap_out_ and swap_in_ nodes with the node itself. |
| 781 | swap_out_node->set_device(node->device()); |
| 782 | swap_in_node->set_device(node->device()); |
| 783 | string coloc_group = strings::StrCat("loc@", tensor_to_swap); |
| 784 | (*swap_out_node->mutable_attr())["_class"].mutable_list()->add_s(coloc_group); |
| 785 | (*swap_in_node->mutable_attr())["_class"].mutable_list()->add_s(coloc_group); |
| 786 | (*node->mutable_attr())["_class"].mutable_list()->add_s(coloc_group); |
| 787 | |
| 788 | (*swap_in_node->mutable_attr())["T"].set_type(input_type); |
| 789 | (*swap_out_node->mutable_attr())["T"].set_type(input_type); |
| 790 | *swap_pair = std::make_pair(swap_out_node, swap_in_node); |
| 791 | |
| 792 | return Status::OK(); |
| 793 | } |
| 794 | |
| 795 | struct SwapInfo { |
no test coverage detected