Olivier BICHLER · a138240f · 84a56eb6 · 894b770e · 69351afb · 83a05b9a
--- a/include/aidge/scheduler/Scheduler.hpp

+ 66

− 9
+++ b/include/aidge/scheduler/Scheduler.hpp

+ 66

− 9
 @@ -29,7 +29,21 @@ class Node;
 class GraphView;

 class SequentialScheduler {
-private:
+protected:
+    struct StaticSchedulingElement {
+        StaticSchedulingElement(
+            std::shared_ptr<Node> node_,
+            size_t early_ = static_cast<size_t>(-1),
+            size_t late_ = static_cast<size_t>(-1))
+            : node(node_), early(early_), late(late_) {}
+
+        std::shared_ptr<Node> node;
+        size_t early;
+        size_t late;
+        std::vector<std::shared_ptr<StaticSchedulingElement>> earlierThan;
+        std::vector<std::shared_ptr<StaticSchedulingElement>> laterThan;
+    };
+
    struct SchedulingElement {
        SchedulingElement(
            std::shared_ptr<Node> node_,
 @@ -55,9 +69,19 @@ public:
    {
        // ctor
    };
-    ~SequentialScheduler() = default;
+    virtual ~SequentialScheduler() = default;
+
+    /**
+     * Generate full static scheduling of the GraphView.
+     * For each node, an earliest and latest possible execution logical step
+     * is specified. Nodes that may be scheduled at the same logical step have
+     * no data dependency and can be run in parallel.
+    */
+    void generateScheduling();

-    void generateScheduling(bool verbose = false);
+    /**
+     * Reset all scheduling and associated nodes producer consumer.
+    */
    void resetScheduling();

    /**
 @@ -77,7 +101,13 @@ public:
    /**
     * @brief Run the provided Computational Graph with a batch of data
     */
-    void forward(bool forwardDims = true, bool verbose = false, std::vector<std::shared_ptr<Aidge::Tensor>> data = {});
+    virtual void forward(bool forwardDims = true, bool verbose = false, std::vector<std::shared_ptr<Aidge::Tensor>> data = {});
+
+    /**
+     * @brief Save in a Markdown file the static scheduling with early and late relative order for the nodes.
+     * @param fileName Name of the generated file.
+     */
+    void saveStaticSchedulingDiagram(const std::string& fileName) const;

    /**
     * @brief Save in a Markdown file the order of layers execution.
 @@ -89,14 +119,26 @@ public:
     * @brief Return a vector of Node ordered by the order they are called by the scheduler
     * @return std::vector<std::shared_ptr<Node>>
     */
-    inline std::vector<std::shared_ptr<Node>> getStaticScheduling(size_t step = 0) const noexcept {
-        return mStaticSchedule.at(step);
-    }
+    std::vector<std::shared_ptr<Node>> getStaticScheduling(size_t step = 0) const;
    inline std::shared_ptr<GraphView> getGraphView() const noexcept {
        return mGraphView;
    }

-private:
+protected:
+    /**
+     * Generate an initial base scheduling for the GraphView.
+     * The scheduling is entirely sequential and garanteed to be valid w.r.t.
+     * each node producer-consumer model.
+    */
+    std::vector<std::shared_ptr<StaticSchedulingElement>> generateBaseScheduling() const;
+
+    /**
+     * Fill-in early and late scheduling step from initial base scheduling.
+     * For each node, specifies the earliest and latest possible execution 
+     * logical step.
+    */
+    void generateEarlyLateScheduling(std::vector<std::shared_ptr<StaticSchedulingElement>>& schedule) const;
+
    /**
     * @brief Set of layers receiving an input from currently processing layers
     *
 @@ -114,9 +156,24 @@ private:
    /** @brief List of SchedulingElement (i.e: Nodes with their computation time) */
    std::vector<SchedulingElement> mScheduling;
    /** @brief List of nodes ordered by their */
-    std::vector<std::vector<std::shared_ptr<Node>>> mStaticSchedule;
+    std::vector<std::vector<std::shared_ptr<StaticSchedulingElement>>> mStaticSchedule;
    size_t mStaticScheduleStep = 0;
 };
+
+/**
+ * Multi-threaded parallel scheduler with dynamic scheduling.
+*/
+class ParallelScheduler : public SequentialScheduler {
+public:
+    ParallelScheduler(std::shared_ptr<GraphView> graphView, std::shared_ptr<Node> upperNode = nullptr)
+        : SequentialScheduler(graphView, upperNode)
+    {
+        // ctor
+    };
+    ~ParallelScheduler() = default;
+
+    virtual void forward(bool forwardDims = true, bool verbose = false, std::vector<std::shared_ptr<Aidge::Tensor>> data = {});
+};
 } // namespace Aidge

 #endif /* AIDGE_SCHEDULER_H_ */