1fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar/* Copyright 2016 The TensorFlow Authors. All Rights Reserved. 2fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 3fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh SivakumarLicensed under the Apache License, Version 2.0 (the "License"); 4fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumaryou may not use this file except in compliance with the License. 5fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh SivakumarYou may obtain a copy of the License at 6fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 7fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar http://www.apache.org/licenses/LICENSE-2.0 8fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 9fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh SivakumarUnless required by applicable law or agreed to in writing, software 10fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumardistributed under the License is distributed on an "AS IS" BASIS, 11fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh SivakumarWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh SivakumarSee the License for the specific language governing permissions and 13fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumarlimitations under the License. 14fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar==============================================================================*/ 15fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 16fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar#ifndef TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_SCHEDULER_H_ 17fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar#define TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_SCHEDULER_H_ 18fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 19fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar#include <deque> 207149a2e2e2f549035f23e21224ee41afe8df3876A. Unique TensorFlower#include <functional> 21fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar#include <map> 22fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar#include <unordered_map> 23fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar#include <vector> 24fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 25fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar#include "tensorflow/core/common_runtime/device.h" 26fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar#include "tensorflow/core/common_runtime/device_set.h" 277149a2e2e2f549035f23e21224ee41afe8df3876A. Unique TensorFlower#include "tensorflow/core/graph/costmodel.h" 28fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 29fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumarnamespace tensorflow { 30fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 31fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumarclass SlackAnalysis { 32fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar public: 33fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar SlackAnalysis(const Graph* g, const CostModel* cost_model); 34fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 35fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar ~SlackAnalysis() {} 36fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 37fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar // Compute the earliest possible start time for each node, based on 38fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar // a given cost model. 'asap_time' is indexed by node id. 39fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar Microseconds ComputeAsap(std::vector<Microseconds>* asap_times); 40fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 41fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar // Compute the latest possible start time for each node, based on 42fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar // a given cost model. 'alap_time' is indexed by node id. 43fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar Microseconds ComputeAlap(std::vector<Microseconds>* alap_times); 44fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 45fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar // Compute the "slack" of each node. 'slacks' is indexed by node id. 46fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar void ComputeSlack(std::vector<int64>* slacks); 47fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 48fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar private: 49fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar const Graph* graph_; 50fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar const CostModel* cost_model_; 51fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 52fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar TF_DISALLOW_COPY_AND_ASSIGN(SlackAnalysis); 53fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar}; 54fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 55fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumarclass GreedyScheduler { 56fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar public: 57fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar struct Sim { 58fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar int degree_parallelism; 59fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar int num_running; 60fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar std::vector<Node*> ready_nodes; 61fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar }; 62fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 63fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar struct Event { 64fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar Node* node; 65fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar Microseconds time; 66fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar bool is_completion; 67fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 68fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar bool operator<(const Event& other) const { return time < other.time; } 69fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar }; 70fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 71fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar GreedyScheduler(const DeviceSet* devices, const CostModel* cost_model, 72fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar const Graph* g, std::vector<int64>* priority); 73fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 74fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar ~GreedyScheduler(); 75fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 76fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar // Computes the start time of each node given the priorities of 77fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar // the nodes. 78fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar Microseconds ComputeSchedule(std::vector<Microseconds>* start_times); 79fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 80fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar private: 81fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar // Returns the ready node with the highest priority for a sim. 82fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar Node* GetNodeWithHighestPriority(const std::vector<Node*>& nodes); 83fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 84fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar const DeviceSet* devices_; 85fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar const CostModel* cost_model_; 86fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar const Graph* graph_; 87fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar std::vector<int64>* priority_; 88fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar std::unordered_map<string, Sim*> device_states_; 89fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 90fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar TF_DISALLOW_COPY_AND_ASSIGN(GreedyScheduler); 91fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar}; 92fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 93fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumarclass PriorityScheduler { 94fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar public: 95fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar PriorityScheduler(const DeviceSet* devices, const CostModel* cost_model, 96fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar const Graph* g); 97fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 98fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar ~PriorityScheduler() {} 99fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 100fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar // Computes a schedule of the ideal start time for each node. 101fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar // Returns the makespan (the total running time). 102fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar Microseconds ComputeSchedule(std::vector<Microseconds>* start_times); 103fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 104fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar // Computes a schedule and assigns priorities to the nodes based on 105fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar // the schedule. Returns the makespan. 106fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar Microseconds AssignPriorities(std::vector<int64>* priorities); 107fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 108fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar private: 109fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar const DeviceSet* devices_; 110fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar const CostModel* cost_model_; 111fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar const Graph* graph_; 112fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 113fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar TF_DISALLOW_COPY_AND_ASSIGN(PriorityScheduler); 114fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar}; 115fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 116fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar} // namespace tensorflow 117fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar 118fafcc249f07342479674c1c24acfb1bf473ffdaeSuharsh Sivakumar#endif // TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_SCHEDULER_H_ 119