From 07d99bae97f3542d458c2c7d7898a157155e9ac9 Mon Sep 17 00:00:00 2001
From: Sukrit Kalra <sukritkalra@gmail.com>
Date: Mon, 20 May 2024 12:36:43 -0700
Subject: [PATCH] Clean up code and fix TetriSched scheduler.

---
 main.py                                       | 22 ++++-----
 .../include/tetrisched/OptimizationPasses.hpp | 13 ++++--
 .../include/tetrisched/Scheduler.hpp          |  6 +--
 schedulers/tetrisched/python/TetrischedPy.cpp | 14 +++---
 .../tetrisched/src/CapacityConstraint.cpp     | 41 ++++++++---------
 .../tetrisched/src/OptimizationPasses.cpp     | 34 +++-----------
 schedulers/tetrisched/src/Scheduler.cpp       | 33 ++++++-------
 schedulers/tetrisched_scheduler.py            | 46 ++++++++++---------
 8 files changed, 92 insertions(+), 117 deletions(-)

diff --git a/main.py b/main.py
index 39a73db6..77e21767 100644
--- a/main.py
+++ b/main.py
@@ -64,12 +64,6 @@
     "Sets the mode in which the log file is opened. If 'append', the log file is "
     "opened in append mode, and if 'write', the log file is opened in write mode. ",
 )
-flags.DEFINE_multi_enum(
-    "opt_passes",
-    [], # default
-    ["CRITICAL_PATH_PASS", "DYNAMIC_DISCRETIZATION_PASS", "CAPACITY_CONSTRAINT_PURGE_PASS"], # choices
-    help="Specify the optimization passes that needs to be enabled once the STRL is generated, default: ['CRITICAL_PATH_PASS', 'CAPACITY_CONSTRAINT_PURGE_PASS']",
-)
 flags.DEFINE_string(
     "csv_file_name",
     None,
@@ -477,12 +471,6 @@
     "If `True`, the scheduler is allowed to batch tasks "
     "that share a WorkProfile together.",
 )
-# flags.DEFINE_bool(
-#     "scheduler_enable_optimization_pass",
-#     False,
-#     "If `True`, the scheduler runs pre/post-translation optimization passes"
-#     "when registering STRL expression.",
-# )
 flags.DEFINE_bool(
     "scheduler_selective_rescheduling",
     False,
@@ -501,6 +489,16 @@
     "The percentage of critical path duration until which the scheduler will try "
     "placing the TaskGraph, and drop the TaskGraph if it cannot be placed after.",
 )
+flags.DEFINE_multi_enum(
+    "optimization_passes",
+    [],
+    [
+        "CRITICAL_PATH_PASS",
+        "CAPACITY_CONSTRAINT_PURGE_PASS",
+        "DYNAMIC_DISCRETIZATION_PASS",
+    ],
+    "Specify the optimizations that needs to be enabled once the STRL is generated.",
+)
 
 # Workload definition related flags.
 flags.DEFINE_integer(
diff --git a/schedulers/tetrisched/include/tetrisched/OptimizationPasses.hpp b/schedulers/tetrisched/include/tetrisched/OptimizationPasses.hpp
index 409d4101..49aec636 100644
--- a/schedulers/tetrisched/include/tetrisched/OptimizationPasses.hpp
+++ b/schedulers/tetrisched/include/tetrisched/OptimizationPasses.hpp
@@ -17,6 +17,7 @@ enum OptimizationPassType {
   POST_TRANSLATION_PASS = 1,
 };
 
+/// The types of OptimizationPasses implemented in the library.
 enum OptimizationPassCategory {
     CRITICAL_PATH_PASS = 0,
     DYNAMIC_DISCRETIZATION_PASS = 1,
@@ -35,7 +36,11 @@ struct OptimizationPassConfig {
     Time finerDiscretizationWindow = 5;
     std::string toString() const {
         std::stringstream ss;
-        ss << "{minDisc: " << minDiscretization << ", maxDisc: " << maxDiscretization << ", maxOccupancyThreshold: " << maxOccupancyThreshold << ", finerDiscretizationAtPrevSolution: " << finerDiscretizationAtPrevSolution << ", finerDiscretizationWindow: " << finerDiscretizationWindow << "\n";
+        ss << "{ minDisc: " << minDiscretization
+           << ", maxDisc: " << maxDiscretization
+           << ", maxOccupancyThreshold: " << maxOccupancyThreshold
+           << ", finerDiscretizationAtPrevSolution: " << finerDiscretizationAtPrevSolution
+           << ", finerDiscretizationWindow: " << finerDiscretizationWindow << std::endl;
         return ss.str();
     }
 };
@@ -184,12 +189,12 @@ class OptimizationPassRunner {
   bool debug;
   /// A list of optimization passes to run.
   std::vector<OptimizationPassPtr> registeredPasses;
-  
+
   OptimizationPassConfigPtr optConfig;
 
   public :
   /// Initialize the OptimizationPassRunner.
-  OptimizationPassRunner(bool debug = false, OptimizationPassConfigPtr optConfig = nullptr);
+  OptimizationPassRunner(OptimizationPassConfigPtr optConfig, bool debug = false);
 
   /// Run the pre-translation optimization passes on the given STRL expression.
   void runPreTranslationPasses(Time currentTime, ExpressionPtr strlExpression,
@@ -200,7 +205,7 @@ class OptimizationPassRunner {
                                 CapacityConstraintMapPtr capacityConstraints);
 
   /// Add the optimization passes
-  void addOptimizationPass(OptimizationPassCategory optPass);  
+  void addOptimizationPass(OptimizationPassCategory optPass);
 };
 }  // namespace tetrisched
 #endif  // _TETRISCHED_OPTIMIZATION_PASSES_HPP_
diff --git a/schedulers/tetrisched/include/tetrisched/Scheduler.hpp b/schedulers/tetrisched/include/tetrisched/Scheduler.hpp
index 015829a1..1eb08a1e 100644
--- a/schedulers/tetrisched/include/tetrisched/Scheduler.hpp
+++ b/schedulers/tetrisched/include/tetrisched/Scheduler.hpp
@@ -9,12 +9,8 @@ namespace tetrisched {
 
 /// The `SchedulerConfig` structure represents the configuration of the
 /// scheduler. This config is used to inform the choice of how the scheduler
-/// should schedule the STRL expression.
+/// should run the solver.
 struct SchedulerConfig {
-  /// If True, the scheduler will optimize the STRL expression.
-  bool optimize;
-
-  /// The configuration for the solver backend.
   std::optional<uint64_t> numThreads;
   std::optional<Time> totalSolverTimeMs;
   std::optional<Time> newSolutionTimeMs;
diff --git a/schedulers/tetrisched/python/TetrischedPy.cpp b/schedulers/tetrisched/python/TetrischedPy.cpp
index bea54bdd..395deb09 100644
--- a/schedulers/tetrisched/python/TetrischedPy.cpp
+++ b/schedulers/tetrisched/python/TetrischedPy.cpp
@@ -12,8 +12,8 @@ void defineBasicTypes(py::module_& tetrisched_m) {
   py::enum_<tetrisched::OptimizationPassCategory>(tetrisched_m, "OptimizationPassCategory")
       .value("CRITICAL_PATH_PASS", tetrisched::OptimizationPassCategory::CRITICAL_PATH_PASS)
       .value("DYNAMIC_DISCRETIZATION_PASS", tetrisched::OptimizationPassCategory::DYNAMIC_DISCRETIZATION_PASS)
-      .value("CAPACITY_CONSTRAINT_PURGE_PASS", tetrisched::OptimizationPassCategory::CAPACITY_CONSTRAINT_PURGE_PASS)
-      .export_values();
+      .value("CAPACITY_CONSTRAINT_PURGE_PASS", tetrisched::OptimizationPassCategory::CAPACITY_CONSTRAINT_PURGE_PASS);
+
   // Define the Partition type.
   py::class_<tetrisched::Partition, tetrisched::PartitionPtr>(
       tetrisched_m, "Partition", py::dynamic_attr())
@@ -64,23 +64,21 @@ void defineScheduler(py::module_& tetrisched_m) {
   py::class_<tetrisched::OptimizationPassConfig, tetrisched::OptimizationPassConfigPtr>(tetrisched_m, "OptimizationPassConfig")
       .def(py::init<>(), "Initializes an empty OptimizationPassConfig.")
       .def_readwrite("minDiscretization", &tetrisched::OptimizationPassConfig::minDiscretization,
-                     "the min discretization for Dynamic Discretiazation OPT pass.")
+                     "The minimum discretization for Dynamic Discretiazation OPT pass.")
       .def_readwrite("maxDiscretization", &tetrisched::OptimizationPassConfig::maxDiscretization,
-                     "the max discretization for Dynamic Discretiazation OPT pass.")
+                     "The maximum discretization for Dynamic Discretiazation OPT pass.")
       .def_readwrite("maxOccupancyThreshold", &tetrisched::OptimizationPassConfig::maxOccupancyThreshold,
-                     "the max occupancy threshold beyond which dynamic discretization is always min.")
+                     "The max occupancy threshold beyond which dynamic discretization is always min.")
       .def_readwrite("finerDiscretizationAtPrevSolution", &tetrisched::OptimizationPassConfig::finerDiscretizationAtPrevSolution,
                      "Whether to enabled finer discretization at solved solutions.")
       .def_readwrite("finerDiscretizationWindow", &tetrisched::OptimizationPassConfig::finerDiscretizationWindow,
                      "The window upto which finer discretization should be enabled around previously solved solutions.")
       .def("toString", &tetrisched::OptimizationPassConfig::toString, "Print String Representation");
+
   // Define the Config for the Scheduler.
   py::class_<tetrisched::SchedulerConfig, tetrisched::SchedulerConfigPtr>(
       tetrisched_m, "SchedulerConfig")
       .def(py::init<>(), "Initializes an empty SchedulerConfig.")
-      .def_readwrite("optimize", &tetrisched::SchedulerConfig::optimize,
-                     "If True, the scheduler will optimize the STRL "
-                     "expression.")
       .def_readwrite("numThreads", &tetrisched::SchedulerConfig::numThreads,
                      "The number of threads to use for the solver.")
       .def_readwrite("totalSolverTimeMs",
diff --git a/schedulers/tetrisched/src/CapacityConstraint.cpp b/schedulers/tetrisched/src/CapacityConstraint.cpp
index d0b1f704..57732ae4 100644
--- a/schedulers/tetrisched/src/CapacityConstraint.cpp
+++ b/schedulers/tetrisched/src/CapacityConstraint.cpp
@@ -253,15 +253,13 @@ void CapacityConstraintMap::registerUsageForDuration(
     // Time currentTime = startTime;
     Time currentTime = timeRangeToGranularities[granularityIndex].first.first;
     Time remainderTime = duration + (startTime - currentTime);
-    // if (print)
-    //   {
-    //     std::cout << "[CapConstraint] Registering usage for expression: " <<
-    //     exprName << " starting at "
-    //                      << startTime << " with duration " << duration
-    //                      << " at time: " << currentTime
-    //                      << " and remainder time: " << remainderTime << "."
-    //                      << std::endl;
-    //   }
+    if (print) {
+      std::cout <<
+        "[CapConstraint] Registering usage for expression: " <<
+        exprName << " starting at " << startTime << " with duration " <<
+        duration << " at time: " << currentTime << " and remainder time: " <<
+        remainderTime << "." << std::endl;
+    }
     while (remainderTime > 0) {
       auto& timeRange = timeRangeToGranularities[granularityIndex].first;
       auto& granularity = timeRangeToGranularities[granularityIndex].second;
@@ -289,18 +287,19 @@ void CapacityConstraintMap::registerUsageForDuration(
           currentTime += granularity;
           remainderTime = 0;
         }
-        // if(print){
-        //   std::cout << "\t[CapConstraintInside] Registering usage for
-        //   expression starting at "
-        //                    << startTime << " with duration " << duration
-        //                    << " at time: " << currentTime
-        //                    << " within the time range [" << timeRange.first
-        //                    << ", "
-        //                    << timeRange.second
-        //                    << "] with granularity: " << granularity
-        //                    << " and remainder time: " << remainderTime << "."
-        //                    << std::endl;
-        // }
+        if (print){
+          std::cout <<
+            "\t[CapConstraintInside] Registering usage for expression starting at "
+            << startTime << " with duration " << duration
+            << " at time: " << currentTime
+            << " within the time range [" << timeRange.first
+            << ", "
+            << timeRange.second
+            << "] with granularity: " << granularity
+            << " and remainder time: " << remainderTime << "."
+            << std::endl;
+        }
+
         if (remainderTime != 0 &&
             currentTime >= std::min(startTime + duration, timeRange.second) &&
             granularityIndex == timeRangeToGranularities.size() - 1) {
diff --git a/schedulers/tetrisched/src/OptimizationPasses.cpp b/schedulers/tetrisched/src/OptimizationPasses.cpp
index df22c278..39d97dcd 100644
--- a/schedulers/tetrisched/src/OptimizationPasses.cpp
+++ b/schedulers/tetrisched/src/OptimizationPasses.cpp
@@ -699,7 +699,7 @@ void DiscretizationSelectorOptimizationPass::runPass(
   //             - minOccupancyTime] << std::endl;
   // }
 
-  // changing the STRL expressions 
+  // changing the STRL expressions
   // for (auto &[discreteTimeRange, discreteGranularity] : timeRangeToGranularities)
   // Find Max expressions over NCK and remove NCK expressions with redundant
   // start times
@@ -716,7 +716,7 @@ void DiscretizationSelectorOptimizationPass::runPass(
         auto endTime = discreteTimeRange.second;
         // find ncks within startTime and endTime for this Max expr
         // std::vector<ExpressionPtr> ncksWithinTimeRange;
-        
+
         ExpressionPtr minStartTimeNckExpr = nullptr;
         ExpressionPtr prevSolutionNckExpr = nullptr;
         for (; child != expressionChildren.end(); ++child) {
@@ -1129,37 +1129,17 @@ void CapacityConstraintMapPurgingOptimizationPass::clean() { cliques.clear(); }
 
 /* Methods for OptimizationPassRunner */
 OptimizationPassRunner::OptimizationPassRunner(
-    bool debug, OptimizationPassConfigPtr optConfig)
-    : debug(debug), optConfig(optConfig)
-{
-  // Register the Critical Path optimization pass.
-  // registeredPasses.push_back(std::make_shared<CriticalPathOptimizationPass>());
-
-  // if (enableDynamicDiscretization) {
-  //   // Register the DiscretizationGenerator pass.
-  //   registeredPasses.push_back(
-  //       std::make_shared<DiscretizationSelectorOptimizationPass>(
-  //           minDiscretization, maxDiscretization, maxOccupancyThreshold,
-  //           finerDiscretizationAtPrevSolution, finerDiscretizationWindow));
-  // }
-
-  // // Register the CapacityConstraintMapPurging optimization pass.
-  // registeredPasses.push_back(
-  //     std::make_shared<CapacityConstraintMapPurgingOptimizationPass>());
-}
+    OptimizationPassConfigPtr optConfig, bool debug)
+    : debug(debug), optConfig(optConfig) { }
 
 void OptimizationPassRunner::addOptimizationPass(OptimizationPassCategory optPass) {
-  // Register the Critical Path optimization pass.
-
+  // Register the requested optimization pass.
   switch (optPass)
   {
     case OptimizationPassCategory::CRITICAL_PATH_PASS:
       registeredPasses.push_back(std::make_shared<CriticalPathOptimizationPass>());
       break;
     case OptimizationPassCategory::DYNAMIC_DISCRETIZATION_PASS:
-      if (optConfig == nullptr) {
-        optConfig = std::make_shared<OptimizationPassConfig>();
-      }
       registeredPasses.push_back(std::make_shared<DiscretizationSelectorOptimizationPass>(
           optConfig->minDiscretization, optConfig->maxDiscretization, optConfig->maxOccupancyThreshold,
           optConfig->finerDiscretizationAtPrevSolution, optConfig->finerDiscretizationWindow));
@@ -1167,9 +1147,9 @@ void OptimizationPassRunner::addOptimizationPass(OptimizationPassCategory optPas
     case OptimizationPassCategory::CAPACITY_CONSTRAINT_PURGE_PASS:
       registeredPasses.push_back(std::make_shared<CapacityConstraintMapPurgingOptimizationPass>());
       break;
-
     default:
-      break;
+      throw tetrisched::exceptions::RuntimeException(
+        "Unknown Optimization Pass Category: " + std::to_string((int) optPass));
   }
 }
 
diff --git a/schedulers/tetrisched/src/Scheduler.cpp b/schedulers/tetrisched/src/Scheduler.cpp
index 8aaa6019..babc7483 100644
--- a/schedulers/tetrisched/src/Scheduler.cpp
+++ b/schedulers/tetrisched/src/Scheduler.cpp
@@ -11,15 +11,14 @@
 #endif
 
 namespace tetrisched {
-  Scheduler::Scheduler(Time discretization, SolverBackendType solverBackend,
-                       std::string logDir, OptimizationPassConfigPtr optConfig)
-      : solverBackend(solverBackend),
-        discretization(discretization),
-        logDir(logDir)
-  {
-    // Initialize the solver backend.
-    switch (solverBackend)
-    {
+Scheduler::Scheduler(Time discretization, SolverBackendType solverBackend,
+                     std::string logDir, OptimizationPassConfigPtr optConfig)
+    : solverBackend(solverBackend),
+      discretization(discretization),
+      optimizationPasses(optConfig, false),
+      logDir(logDir) {
+  // Initialize the solver backend.
+  switch (solverBackend) {
 #ifdef _TETRISCHED_WITH_CPLEX_
     case SolverBackendType::CPLEX:
       solver = std::make_shared<CPLEXSolver>();
@@ -35,10 +34,8 @@ namespace tetrisched {
           "The solver backend type is not supported.");
   }
   solverModel = solver->getModel();
-  optimizationPasses = OptimizationPassRunner(
-      false, optConfig);
   solverConfig = std::make_shared<SolverConfig>();
-  }
+}
 
 void Scheduler::addOptimizationPass(OptimizationPassCategory optPass) {
   optimizationPasses.addOptimizationPass(optPass);
@@ -87,7 +84,7 @@ void Scheduler::registerSTRL(
   }
 
   // Run the Pre-Translation OptimizationPasses on this expression.
-  if (schedulerConfig->optimize) {
+  {
     TETRISCHED_SCOPE_TIMER(
         "Scheduler::registerSTRL::preTranslationOptimizationPasses," +
         std::to_string(currentTime));
@@ -97,7 +94,7 @@ void Scheduler::registerSTRL(
 
   {
     TETRISCHED_SCOPE_NECESSARY_TIMER("Scheduler::registerSTRL::parse," +
-                           std::to_string(currentTime));
+                                     std::to_string(currentTime));
     // Parse the ExpressionTree to populate the solver model.
     TETRISCHED_DEBUG("Beginning the parsing of the ExpressionTree rooted at "
                      << expression->getName() << ".")
@@ -108,7 +105,7 @@ void Scheduler::registerSTRL(
   }
 
   // Run the Post-Translation OptimizationPasses on this expression.
-  if (schedulerConfig->optimize) {
+  {
     TETRISCHED_SCOPE_TIMER(
         "Scheduler::registerSTRL::postTranslationOptimizationPasses," +
         std::to_string(currentTime));
@@ -146,21 +143,21 @@ void Scheduler::schedule(Time currentTime) {
   // Translate the model to the solver backend.
   {
     TETRISCHED_SCOPE_NECESSARY_TIMER("Scheduler::schedule::translateModel," +
-                           std::to_string(currentTime));
+                                     std::to_string(currentTime));
     this->solver->translateModel(solverConfig);
   }
 
   // Solve the model.
   {
     TETRISCHED_SCOPE_NECESSARY_TIMER("Scheduler::schedule::solveModel," +
-                           std::to_string(currentTime));
+                                     std::to_string(currentTime));
     solverSolution = this->solver->solveModel();
   }
 
   // Populate the results from the solver into the expression tree.
   {
     TETRISCHED_SCOPE_NECESSARY_TIMER("Scheduler::schedule::populateResults," +
-                           std::to_string(currentTime));
+                                     std::to_string(currentTime));
     if (solverSolution.has_value() && solverSolution.value()->isValid()) {
       this->expression.value()->populateResults(solverModel);
     }
diff --git a/schedulers/tetrisched_scheduler.py b/schedulers/tetrisched_scheduler.py
index a5170351..6cbeb425 100644
--- a/schedulers/tetrisched_scheduler.py
+++ b/schedulers/tetrisched_scheduler.py
@@ -181,7 +181,6 @@ def __init__(
         _flags: Optional["absl.flags"] = None,
         adaptive_discretization: bool = False,
         max_time_discretization: EventTime = EventTime(5, EventTime.Unit.US),
-        # dynamic_discretization: bool = False,
         max_occupancy_threshold: float = 0.8,
         finer_discretization_at_prev_solution: bool = False,
         finer_discretization_window: EventTime = EventTime(5, EventTime.Unit.US),
@@ -209,43 +208,47 @@ def __init__(
 
         # Values for STRL generation.
         self._use_windowed_choose = False
-        self._dynamic_discretization = True if "DYNAMIC_DISCRETIZATION_PASS" in self._flags.opt_passes else False
+        self._dynamic_discretization = (
+            "DYNAMIC_DISCRETIZATION_PASS" in self._flags.opt_passes
+        )
         self._adaptive_discretization = adaptive_discretization
         if self._dynamic_discretization or self._adaptive_discretization:
+            # Force windowed choose to be false since these passes don't work with
+            # windowed choose.
             self._use_windowed_choose = False
         self._max_discretization = max_time_discretization.to(EventTime.Unit.US)
-        
-        # optimization passes config 
+
+        # optimization passes config
         self._opt_configuration = tetrisched.OptimizationPassConfig()
         self._opt_configuration.minDiscretization = self._time_discretization.time
         self._opt_configuration.maxDiscretization = self._max_discretization.time
         self._opt_configuration.maxOccupancyThreshold = max_occupancy_threshold
-        self._opt_configuration.finerDiscretizationAtPrevSolution = finer_discretization_at_prev_solution
-        self._opt_configuration.finerDiscretizationWindow = finer_discretization_window.to(EventTime.Unit.US).time
-        
-        # print(f"----- [PYTHON] {self._opt_configuration.toString()}")
-        
+        self._opt_configuration.finerDiscretizationAtPrevSolution = (
+            finer_discretization_at_prev_solution
+        )
+        self._opt_configuration.finerDiscretizationWindow = (
+            finer_discretization_window.to(EventTime.Unit.US).time
+        )
+
         self._scheduler = tetrisched.Scheduler(
             self._time_discretization.time,
             tetrisched.backends.SolverBackendType.GUROBI,
-            self._log_dir, self._opt_configuration
+            self._log_dir,
+            self._opt_configuration,
         )
-        
-            
+
         self._use_task_graph_indicator_utility = self._goal == "max_goodput"
         self._previously_placed_reward_scale_factor = 1.0
-        self._enable_optimization_passes = (
-            len(self._flags.opt_passes) != 0
-        )
-        
+
         # add optimization passes
         if len(self._flags.opt_passes) == 0:
-            self._logger.info(f"Running With no Optimization Pass")
+            self._logger.info("Running with no Optimization Passes.")
         for opt_pass in self._flags.opt_passes:
-            self._logger.info(f"Enable Optimization Passes: {opt_pass}")
-            self._scheduler.addOptimizationPass(tetrisched.OptimizationPassCategory.__members__[opt_pass])
-            # self._scheduler.addOptimizationPass(tetrisched.OptimizationPassCategoryopt_pass)
-        
+            self._logger.info(f"Enabling the Optimization Pass: {opt_pass}")
+            self._scheduler.addOptimizationPass(
+                tetrisched.OptimizationPassCategory.__members__[opt_pass]
+            )
+
         self._selectively_choose_task_graphs_for_rescheduling = (
             _flags.scheduler_selective_rescheduling if _flags else False
         )
@@ -257,7 +260,6 @@ def __init__(
 
         # Scheduler configuration.
         self._scheduler_configuration = tetrisched.SchedulerConfig()
-        self._scheduler_configuration.optimize = self._enable_optimization_passes
         self._scheduler_configuration.totalSolverTimeMs = (
             # 1 minute interrupt by default.
             1 * 60 * 1000