Fixed missing producer consumer system reset

include/aidge/backend/OperatorImpl.hpp

@@ -63,6 +63,12 @@ public:
      */
     virtual void updateConsummerProducer();
 
+    /**
+     * @brief Reset the Consummer Producer system.
+     *
+     */
+    virtual void resetConsummerProducer();
+
     virtual ~OperatorImpl() = default;
 
 protected:

include/aidge/operator/Operator.hpp

@@ -153,6 +153,8 @@ public:
 
     virtual void updateConsummerProducer();
 
+    virtual void resetConsummerProducer();
+
     virtual void forward();
 
     virtual void backward();

include/aidge/scheduler/Scheduler.hpp

@@ -58,11 +58,7 @@ public:
     ~SequentialScheduler() = default;
 
     void generateScheduling(bool verbose = false);
-    inline void resetScheduling() {
-        mScheduling.clear();
-        mStaticSchedule.clear();
-        mStaticScheduleStep = 0;
-    }
+    void resetScheduling();
 
     /**
      * Generate the memory layout for the current static scheduling.

python_binding/backend/pybind_OperatorImpl.cpp

@@ -102,6 +102,15 @@ public:
 
         );
     }
+    void resetConsummerProducer() override {
+        PYBIND11_OVERRIDE_NAME(
+            void,
+            OperatorImpl,
+            "reset_consummer_producer",
+            resetConsummerProducer,
+
+        );
+    }
 };
 
 void init_OperatorImpl(py::module& m){
@@ -116,6 +125,7 @@ void init_OperatorImpl(py::module& m){
     .def("get_nb_consumed_data", &OperatorImpl::getNbConsumedData)
     .def("get_nb_produced_data", &OperatorImpl::getNbProducedData)
     .def("update_consummer_producer", &OperatorImpl::updateConsummerProducer)
+    .def("reset_consummer_producer", &OperatorImpl::resetConsummerProducer)
     ;
 }
 }

src/backend/OperatorImpl.cpp

@@ -68,6 +68,11 @@ void Aidge::OperatorImpl::updateConsummerProducer(){
     }
 }
 
+void Aidge::OperatorImpl::resetConsummerProducer(){
+    std::fill(mNbConsumedData.begin(), mNbConsumedData.end(), 0);
+    std::fill(mNbProducedData.begin(), mNbProducedData.end(), 0);
+}
+
 void Aidge::OperatorImpl::forward() {
     AIDGE_THROW_OR_ABORT(std::runtime_error, "forward() not implemented");
 }

src/operator/Operator.cpp

@@ -59,6 +59,10 @@ void Aidge::Operator::updateConsummerProducer(){
     AIDGE_ASSERT(mImpl != nullptr, "updateConsummerProducer(): an implementation is required for {}!", type());
     mImpl->updateConsummerProducer();
 }
+void Aidge::Operator::resetConsummerProducer(){
+    AIDGE_ASSERT(mImpl != nullptr, "resetConsummerProducer(): an implementation is required for {}!", type());
+    mImpl->resetConsummerProducer();
+}
 
 void Aidge::Operator::runHooks() const {
     for (auto& hook : mHooks) {

src/scheduler/Scheduler.cpp

@@ -301,6 +301,16 @@ void Aidge::SequentialScheduler::generateScheduling(bool verbose) {
     }
 }
 
+void Aidge::SequentialScheduler::resetScheduling() {
+    for (auto node : mGraphView->getNodes()) {
+        node->getOperator()->resetConsummerProducer();
+    }
+
+    mStaticSchedule.clear();
+    mStaticScheduleStep = 0;
+    mScheduling.clear();
+}
+
 /**
  * This version is a simplified version without special handling of concatenation.
 */
@@ -423,7 +433,6 @@ void Aidge::SequentialScheduler::connectInputs(std::vector<std::shared_ptr<Aidge
 }
 
 
-// TODO: handle multiple inputs/outputs
 void Aidge::SequentialScheduler::forward(bool forwardDims, bool verbose, std::vector<std::shared_ptr<Aidge::Tensor>> data) {
     
     // Collect all data input of the graph (that are producers)
@@ -461,6 +470,9 @@ void Aidge::SequentialScheduler::forward(bool forwardDims, bool verbose, std::ve
     fmt::print("\n");
 
     ++mStaticScheduleStep;
+    if (mStaticScheduleStep == mStaticSchedule.size()) {
+        mStaticScheduleStep = 0;
+    }
 }
 
 void Aidge::SequentialScheduler::saveSchedulingDiagram(const std::string& fileName) const {