1/**
2 * Copyright (c) Glow Contributors. See CONTRIBUTORS file.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16#ifndef GLOW_IR_GRAPH_SCHEDULER_H
17#define GLOW_IR_GRAPH_SCHEDULER_H
18
19#include "glow/IR/IR.h"
20
21#include <unordered_map>
22
23namespace glow {
24
25/// Specifies the kind of graph scheduling to perform.
26enum class SchedulerKind {
27 /// This is a heuristics that tries to minimize memory usage.
28 ChildMemSizeBased,
29 /// Performs a standard topological search
30 TopologicalSortBased,
31};
32
33class Scheduler {
34protected:
35 /// Graph being processed.
36 Function &G_;
37 /// Scheduled nodes.
38 NodesPtrList &scheduled_;
39
40public:
41 Scheduler(Function &G, NodesPtrList &scheduled)
42 : G_(G), scheduled_(scheduled) {}
43
44 virtual ~Scheduler() = default;
45
46 // Create a linear execution schedule for a graph.
47 virtual void schedule() = 0;
48
49 NodesPtrList &getSchedule() { return scheduled_; }
50};
51
52/// This is a scheduler based on the generalized the paper "Generalizations of
53/// the Sethi-Ullman algorithm for register allocation" by Andrew W. Appel and
54/// Kenneth J. Supowit.
55/// http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.54.319&rep=rep1&type=pdf
56///
57/// The idea is to give more priority and schedule earlier those child nodes
58/// that free more memory after their computation.
59class ChildMemSizeBasedScheduler : public Scheduler {
60 /// Required number of bytes to hold the results of a given node.
61 std::unordered_map<const Node *, int64_t> resultMemSize_;
62 /// Max number of bytes required during the computation of a given node.
63 std::unordered_map<const Node *, int64_t> maxMemSize_;
64
65 /// \returns true if a node \p N is scheduled already.
66 bool isScheduled(const Node *N) const;
67
68 /// Computes the amount of memory required to keep the result
69 /// of each node.
70 void computeNodeResultsMemorySize();
71
72 /// Computes the max amount of memory required during the computation
73 /// of children for each node.
74 void computeNodeComputationMaxMemorySize();
75
76 /// Order children by (maxSize - resultSize). It gives more
77 /// priority to the nodes that free more memory after
78 /// their computation.
79 void orderChildNodesAndSchedule(Node *N);
80
81 void scheduleNodes();
82
83public:
84 ChildMemSizeBasedScheduler(Function &G, NodesPtrList &Schedule)
85 : Scheduler(G, Schedule) {}
86
87 ~ChildMemSizeBasedScheduler() override = default;
88
89 void schedule() override;
90};
91
92/// This is a simple scheduler based on topological sort based
93/// on post order traversal.
94
95class TopologicalSortBasedScheduler : public Scheduler {
96 std::unordered_map<const Node *, int64_t> indegree_;
97
98public:
99 TopologicalSortBasedScheduler(Function &G, NodesPtrList &Schedule)
100 : Scheduler(G, Schedule) {}
101
102 ~TopologicalSortBasedScheduler() override = default;
103
104 void schedule() override;
105};
106
107Scheduler *createScheduler(SchedulerKind schedulerKind, Function &G,
108 NodesPtrList &scheduled);
109
110} // namespace glow
111
112#endif // GLOW_IR_GRAPH_SCHEDULER_H
113