Skip to content

Commit

Permalink
Clean up code and fix TetriSched scheduler.
Browse files Browse the repository at this point in the history
  • Loading branch information
@sukritkalra
sukritkalra committed May 20, 2024
1 parent ca9d241 commit 07d99ba
Show file tree
Hide file tree
Showing 8 changed files with 92 additions and 117 deletions.
22 changes: 10 additions & 12 deletions main.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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(
Expand Down
13 changes: 9 additions & 4 deletions schedulers/tetrisched/include/tetrisched/OptimizationPasses.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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();
}
};
Expand Down Expand Up @@ -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,
Expand All @@ -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_
6 changes: 1 addition & 5 deletions schedulers/tetrisched/include/tetrisched/Scheduler.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
14 changes: 6 additions & 8 deletions schedulers/tetrisched/python/TetrischedPy.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down Expand Up @@ -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",
Expand Down
41 changes: 20 additions & 21 deletions schedulers/tetrisched/src/CapacityConstraint.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down
34 changes: 7 additions & 27 deletions schedulers/tetrisched/src/OptimizationPasses.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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) {
Expand Down Expand Up @@ -1129,47 +1129,27 @@ 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));
break;
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));
}
}

Expand Down
33 changes: 15 additions & 18 deletions schedulers/tetrisched/src/Scheduler.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>();
Expand All @@ -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);
Expand Down Expand Up @@ -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));
Expand All @@ -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() << ".")
Expand All @@ -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));
Expand Down Expand Up @@ -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);
}
Expand Down
Loading

0 comments on commit 07d99ba

Please sign in to comment.