| /** |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| |
| #ifndef SINGA_CORE_SCHEDULER_H_ |
| #define SINGA_CORE_SCHEDULER_H_ |
| |
| #include <condition_variable> |
| #include <functional> |
| #include <mutex> |
| #include <string> |
| #include <thread> |
| #include <unordered_map> |
| #include <unordered_set> |
| #include <vector> |
| |
| #include "singa/core/common.h" |
| #include "singa/utils/safe_queue.h" |
| |
| using std::function; |
| using std::string; |
| using std::unordered_map; |
| using std::vector; |
| |
| namespace singa { |
| |
| class Node; |
| class Edge; |
| class Graph; |
| class Device; |
| class BlkInfo; |
| |
| typedef std::vector<Node *> NodeVec; |
| typedef std::vector<Edge *> EdgeVec; |
| typedef std::vector<Block *> BlockVec; |
| typedef std::unordered_set<Block *> BlockSet; |
| typedef std::function<void(Context *)> OpFunc; |
| typedef std::unordered_map<Block *, BlkInfo *> Blk2InfoMap; |
| typedef std::chrono::high_resolution_clock::time_point TimePoint; |
| |
| enum BlockType { kUnknow, kInput, kParam, kInter, kEnd }; |
| |
| class Node { |
| public: |
| Node(int id, OpFunc &&op, string op_name) |
| : id_(id), op_(std::move(op)), op_name_(op_name) {} |
| |
| void AddInEdge(Edge *in_edge); |
| void AddOutEdge(Edge *out_edge); |
| |
| // getters of Node |
| int id() const { return id_; } |
| string op_name() const { return op_name_; } |
| const EdgeVec &in_edges() const { return in_edges_; } |
| const EdgeVec &out_edges() const { return out_edges_; } |
| float time_elapsed() const { return time_elapsed_; } |
| |
| // time profiling |
| void time_elapsed_inc(float time) { time_elapsed_ += time; } |
| |
| private: |
| friend Graph; |
| |
| int id_; |
| OpFunc op_; |
| EdgeVec in_edges_; |
| EdgeVec out_edges_; |
| |
| string op_name_; |
| float time_elapsed_ = 0; |
| |
| #ifdef USE_CUDA |
| cudaEvent_t start_; |
| cudaEvent_t end_; |
| friend class CudaGPU; |
| #endif // USE_CUDA |
| }; |
| |
| class Edge { |
| public: |
| Edge(int id, Block *blk, Node *src_node, Node *dst_node) |
| : id_(id), blk_(blk), src_node_(src_node), dst_node_(dst_node) {} |
| |
| void SetBlock(Block *blk); |
| void SetSrcNode(Node *src_node); |
| void SetDstNode(Node *dst_node); |
| |
| // getters of Edge |
| int id() const { return id_; } |
| Block *block() const { return blk_; } |
| Node *src_node() const { return src_node_; } |
| Node *dst_node() const { return dst_node_; } |
| |
| private: |
| friend Graph; |
| |
| int id_; |
| Block *blk_; |
| Node *src_node_; |
| Node *dst_node_; |
| }; |
| |
| class BlkInfo { |
| public: |
| BlkInfo(int id, Block *blk, BlockType type = BlockType::kUnknow) |
| : id_(id), blk_(blk), type_(type), graph_ref_(0), write_edge_(nullptr) {} |
| |
| // getters of BlkInfo |
| int id() const { return id_; } |
| Block *block() const { return blk_; } |
| BlockType type() const { return type_; } |
| int graph_ref() const { return graph_ref_; } |
| Edge *write_edge() const { return write_edge_; } |
| const NodeVec &used_nodes() const { return used_nodes_; } |
| Node *used_node(const size_t idx) const; |
| |
| private: |
| friend Graph; |
| |
| int id_; |
| Block *blk_; |
| BlockType type_; |
| int graph_ref_; |
| Edge *write_edge_; // the edge of last node that writes data into blk |
| NodeVec used_nodes_; // the nodes that use this block(in order of execution) |
| }; |
| |
| class Graph { |
| public: |
| struct CBData { |
| Graph *graph_; |
| Node *node_; |
| |
| CBData(Graph *graph, Node *node) : graph_(graph), node_(node) {} |
| }; |
| |
| ~Graph(); |
| Graph(Device *device); |
| |
| void Reset(); |
| void Debug(); |
| void RunGraph(); |
| void RunInSerial(); |
| void PrintTimeProfiling(); |
| void AddOperation(OpFunc &&op, const BlockVec &read_blocks, |
| const BlockVec &write_blocks, string op_name = "no_name"); |
| |
| // getters of Graph |
| const NodeVec &nodes() const { return nodes_; } |
| const EdgeVec &edges() const { return edges_; } |
| const Blk2InfoMap &blocks() const { return blocks_; } |
| |
| const BlockSet &leaf_blocks() const { return leaf_blocks_; } |
| |
| bool dirty() const { return dirty_; } |
| const NodeVec &begin_nodes() const { return begin_nodes_; } |
| const std::vector<NodeVec> &next_nodes() const { return next_nodes_; } |
| const std::vector<BlockVec> &free_blocks() const { return free_blocks_; } |
| int iteration() const { return iteration_; } |
| |
| Node *node(const size_t idx) const; |
| Edge *edge(const size_t idx) const; |
| BlkInfo *block(Block *blk) const; |
| |
| Node *begin_node(const size_t idx) const; |
| const NodeVec &next_nodes(const size_t idx) const; |
| const BlockVec &free_blocks(const size_t idx) const; |
| |
| private: |
| void Analyze(); |
| void FreeLoop(); |
| void AnalyzeNodes(); |
| void AnalyzeEdges(); |
| void TimeProfilingDoExec(Node *curNode); |
| void AddSyncOp(function<void(Context *)> &&op, string op_name = "no_name"); |
| |
| void step() { iteration_++; } |
| void time_elapsed_inc(float time) { time_elapsed_ += time; } |
| void TakeStartTime(TimePoint &start); |
| void EvaluateTimeElapsed(const TimePoint &start); |
| |
| // static void CUDART_CB Callback(cudaStream_t stream, cudaError_t status, |
| // void *data); |
| |
| private: |
| Device *device_; |
| |
| // nodes, edges and blocks included in the calculation graph |
| NodeVec nodes_; |
| EdgeVec edges_; |
| Blk2InfoMap blocks_; |
| |
| // Leaf blocks written by the previous operations, used for sync op |
| BlockSet leaf_blocks_; |
| |
| // Computational graph analysis |
| bool dirty_ = false; |
| bool in_serial_ = false; |
| NodeVec begin_nodes_; |
| std::vector<NodeVec> next_nodes_; |
| std::vector<BlockVec> free_blocks_; |
| |
| // Time Profiling |
| int iteration_ = 0; |
| float time_elapsed_ = 0; |
| |
| SafeQueue<int> free_queue_; |
| }; |
| |
| /// Scheduling Tensor operations with dependency detection. |
| class Scheduler {}; |
| |
| } // namespace singa |
| #endif // SINGA_CORE_SCHEDULER_H_ |