1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_COMMON_RUNTIME_PLACER_H_ 17 #define TENSORFLOW_COMMON_RUNTIME_PLACER_H_ 18 19 #include <string> 20 #include <unordered_map> 21 22 #include "tensorflow/core/common_runtime/device_set.h" 23 #include "tensorflow/core/graph/graph.h" 24 #include "tensorflow/core/lib/core/status.h" 25 #include "tensorflow/core/platform/macros.h" 26 #include "tensorflow/core/platform/types.h" 27 #include "tensorflow/core/public/session_options.h" 28 #include "tensorflow/core/util/device_name_utils.h" 29 30 namespace tensorflow { 31 32 // A placement algorithm that assigns the nodes of the given Graph to 33 // devices the given DeviceSet, respecting the following constraints: 34 // 35 // 1. Existing device assignments remain unchanged. 36 // 2. Requested (partial or complete) device specifications given by device name 37 // for each node are granted. 38 // 3. Nodes connected by edges of a reference type are colocated on 39 // the same device. 40 // 4. Given nodes "A" and "B", if node "B" has a colocation group 41 // "@loc:A", nodes "A" and "B" will be colocated on the same device. 42 // 43 // The implementation builds a constraint graph with the same set of 44 // nodes, and edges that represent colocation constraints between 45 // nodes. Each connected component in the resulting constraint graph 46 // is then assigned to a set of valid devices. 47 // 48 // Run() will finally assign the device to each node given the list of 49 // possible devices. 50 // 51 // TODO(mrry): "Soft" constraints, such as "place node 'x' as close as 52 // possible to node 'y' while respecting the other constraints"? 53 // TODO(mrry): Create a common interface for this and the other 54 // placement algorithms so that they may be injected into the graph 55 // builder. 56 class Placer { 57 public: 58 // A map from graph node names to numerical IDs (in a Graph object). 59 typedef std::unordered_map<string, int> NodeNameToIdMap; 60 61 // Creates an instance of the Placer algorithm for the given 62 // Graph "graph" (nodes in which may or may not be assigned) on the 63 // given DeviceSet "devices". 64 // 65 // The "graph", and "devices" pointer arguments 66 // are borrowed by this Placer, and must outlive it. 67 Placer(Graph* graph, const DeviceSet* devices, const SessionOptions* options); 68 69 Placer(Graph* graph, const DeviceSet* devices); 70 71 ~Placer(); 72 73 // Assigns each node in this Placer's graph to a device in its 74 // set of devices. 75 // 76 // This method is not thread-safe. 77 // Run() may be invoked at most once. 78 Status Run(); 79 80 private: 81 // Returns true if the device type of 'candidate_device_name' is 82 // found in 'devices'. 83 bool CanAssignToDevice(const string& candidate_device_name, 84 const std::vector<Device*>& devices) const; 85 86 // Assigns 'node's devices to 'assigned_device', and logs the 87 // placement if the SessionOptions entry in 'options_' requests it. 88 void AssignAndLog(int assigned_device, Node* node) const; 89 void LogDeviceAssignment(const Node* node) const; 90 91 Graph* const graph_; // Not owned. 92 const DeviceSet* const devices_; // Not owned. 93 const SessionOptions* options_; // Not owned. 94 const bool log_device_placement_; 95 96 TF_DISALLOW_COPY_AND_ASSIGN(Placer); 97 }; 98 99 } // namespace tensorflow 100 101 #endif // TENSORFLOW_COMMON_RUNTIME_PLACER_H_ 102
