diff --git a/xls/scheduling/BUILD b/xls/scheduling/BUILD index 7597040c0e..71b6ae29fd 100644 --- a/xls/scheduling/BUILD +++ b/xls/scheduling/BUILD @@ -184,6 +184,7 @@ cc_library( srcs = ["sdc_scheduler.cc"], hdrs = ["sdc_scheduler.h"], deps = [ + ":asap_scheduler", ":schedule_bounds", ":schedule_graph", ":schedule_util", diff --git a/xls/scheduling/asap_scheduler.cc b/xls/scheduling/asap_scheduler.cc index 51fb332db1..cfde7f4caf 100644 --- a/xls/scheduling/asap_scheduler.cc +++ b/xls/scheduling/asap_scheduler.cc @@ -143,7 +143,8 @@ absl::StatusOr ASAPScheduler::Schedule( absl::StatusOr ASAPScheduler::ComputeBounds( std::optional pipeline_stages, int64_t clock_period_ps, - std::optional worst_case_throughput) { + std::optional worst_case_throughput, bool get_helpful_error, + int64_t max_upper_bound) { XLS_RET_CHECK(std::holds_alternative(graph_.ir_scope())); auto* f = std::get(graph_.ir_scope()); // TODO(allight): This actually creates a copy of graph_ since it needs to @@ -158,9 +159,9 @@ absl::StatusOr ASAPScheduler::ComputeBounds( VLOG(5) << " constraints: " << absl::StrJoin(constraints_, ", "); XLS_ASSIGN_OR_RETURN( - auto bounds, - sched::ScheduleBounds::Create(graph_, clock_period_ps, delay_estimator_, - worst_case_throughput, constraints_)); + auto bounds, sched::ScheduleBounds::Create( + graph_, clock_period_ps, delay_estimator_, + worst_case_throughput, constraints_, max_upper_bound)); // Add first and last stage constraints. using LastStageConstraint = sched::ScheduleBounds::NodeSchedulingConstraint::LastStageConstraint; @@ -173,9 +174,12 @@ absl::StatusOr ASAPScheduler::ComputeBounds( absl::Status tighten_bounds_status = TightenBounds(bounds, f, pipeline_stages); if (!tighten_bounds_status.ok()) { - return GenerateHelpfulError(std::move(tighten_bounds_status), - pipeline_stages, clock_period_ps, - worst_case_throughput); + if (get_helpful_error) { + return GenerateHelpfulError(std::move(tighten_bounds_status), + pipeline_stages, clock_period_ps, + worst_case_throughput); + } + return tighten_bounds_status; } return bounds; } diff --git a/xls/scheduling/asap_scheduler.h b/xls/scheduling/asap_scheduler.h index cf46b17304..ab1c3ec3cc 100644 --- a/xls/scheduling/asap_scheduler.h +++ b/xls/scheduling/asap_scheduler.h @@ -51,6 +51,14 @@ class ASAPScheduler : public Scheduler { SchedulingFailureBehavior failure_behavior, std::optional worst_case_throughput = std::nullopt) override; + // An alternate interface that returns the full ASAP/ALAP bounds; mostly used + // for other schedulers to build on top of this. + absl::StatusOr ComputeBounds( + std::optional pipeline_stages, int64_t clock_period_ps, + std::optional worst_case_throughput, + bool get_helpful_error = true, + int64_t max_upper_bound = sched::ScheduleBounds::kDefaultMaxUpperBound); + const ScheduleGraph& graph() const { return graph_; } DelayEstimator& delay_estimator() const { return delay_estimator_; } absl::Span constraints() const { @@ -68,12 +76,8 @@ class ASAPScheduler : public Scheduler { absl::Status&& orig_status, std::optional pipeline_stages, int64_t clock_period_ps, std::optional worst_case_throughput); - // Exposed to allow for Min-cut and random to be built on top of this. - absl::StatusOr ComputeBounds( - std::optional pipeline_stages, int64_t clock_period_ps, - std::optional worst_case_throughput); - // Helper to tighten bounds using the ASAP/ALAP bounds. + // Exposed as `protected` so the random scheduler can build on top of this. static absl::Status TightenBounds(sched::ScheduleBounds& bounds, FunctionBase* f, std::optional schedule_length); diff --git a/xls/scheduling/sdc_scheduler.cc b/xls/scheduling/sdc_scheduler.cc index a97c3c751e..18a4c46437 100644 --- a/xls/scheduling/sdc_scheduler.cc +++ b/xls/scheduling/sdc_scheduler.cc @@ -49,6 +49,7 @@ #include "xls/ir/op.h" #include "xls/ir/state_element.h" #include "xls/ir/type.h" +#include "xls/scheduling/asap_scheduler.h" #include "xls/scheduling/schedule_bounds.h" #include "xls/scheduling/schedule_graph.h" #include "xls/scheduling/schedule_util.h" @@ -985,6 +986,19 @@ absl::Status SDCSchedulingModel::SetWorstCaseThroughput( return absl::OkStatus(); } +void SDCSchedulingModel::SetNodeBounds(Node* node, int64_t lower_bound, + int64_t upper_bound) { + model_.set_lower_bound(cycle_var_.at(node), lower_bound); + model_.set_upper_bound( + cycle_var_.at(node), + std::min(kMaxStages, static_cast(upper_bound))); +} + +void SDCSchedulingModel::RemoveNodeBounds(Node* node) { + model_.set_lower_bound(cycle_var_.at(node), 0.0); + model_.set_upper_bound(cycle_var_.at(node), kMaxStages); +} + void SDCSchedulingModel::SetPipelineLength( std::optional pipeline_length) { if (pipeline_length.has_value()) { @@ -1022,12 +1036,20 @@ absl::Status SDCSchedulingModel::AddSlackVariables( "infeasible_per_state_backedge_slack_pool must be positive; was ", *infeasible_per_state_backedge_slack_pool)); } - // Add slack variables to all relevant constraints. + + // Remove any bounds on the cycle variables, relying entirely on the + // `last_stage_` bound that we'll add slack to below. + for (const auto& [node, var] : cycle_var_) { + model_.set_lower_bound(var, 0.0); + model_.set_upper_bound(var, kMaxStages); + } // Remove any pre-existing objective, and declare that we'll be minimizing // our new objective. model_.Minimize(0); + // Add slack variables to all relevant constraints. + // First, try to minimize the depth of the pipeline. We assume users are // most willing to relax this; i.e., they care about throughput more than // latency. @@ -1337,7 +1359,9 @@ SDCScheduler::SDCScheduler( solver_type_(solver_type), solve_parameters_(std::move(solve_parameters)), model_(graph, delay_map_, initiation_interval, sdc_solution_tolerance, - arc_worst_case_throughput, default_arc_worst_case_throughput) {} + arc_worst_case_throughput, default_arc_worst_case_throughput), + delay_map_as_estimator_(delay_map_), + asap_(model_.graph(), delay_map_as_estimator_) {} absl::Status SDCScheduler::Initialize() { XLS_ASSIGN_OR_RETURN(solver_, math_opt::NewIncrementalSolver( @@ -1351,6 +1375,7 @@ absl::Status SDCScheduler::AddConstraints( for (const SchedulingConstraint& constraint : constraints) { XLS_RETURN_IF_ERROR(model_.AddSchedulingConstraint(constraint)); } + XLS_RETURN_IF_ERROR(asap_.AddConstraints(constraints)); return absl::OkStatus(); } @@ -1401,6 +1426,60 @@ absl::StatusOr SDCScheduler::Schedule( VLOG(5) << " Configured with: " << (check_feasibility_ ? "check feasibility" : "minimize dynamic throughput"); + + // If possible, use the ASAP scheduler to compute ASAP/ALAP bounds, saving + // time by restricting the range of values the LP solver needs to consider. + // + // This also lets us use the ASAP scheduler to compute the tightest possible + // pipeline length we can target. + absl::StatusOr bounds = asap_.ComputeBounds( + pipeline_stages, clock_period_ps, worst_case_throughput, + /*get_helpful_error=*/false, + /*max_upper_bound=*/SDCSchedulingModel::kMaxStages); + + if (bounds.ok()) { + if (check_feasibility_) { + // We're just checking feasibility and the ASAP scheduler worked; we're + // already done! We can just return the ASAP schedule. + ScheduleCycleMap schedule; + schedule.reserve(model_.graph().nodes().size()); + for (const ScheduleNode& schedule_node : model_.graph().nodes()) { + schedule[schedule_node.node] = bounds->lb(schedule_node.node); + } + return schedule; + } + + // Apply the ASAP bounds to the cycle variables, saving the solver some work + for (const auto& [node, _] : model_.GetCycleVars()) { + model_.SetNodeBounds(node, bounds->lb(node), bounds->ub(node)); + } + + // If the user didn't specify a pipeline length, the ASAP scheduler will + // have produced the tightest possible bound - so we should use that. + // + // NOTE: Despite its name, the ASAP scheduler has not been verified to be + // able to produce a true ASAP schedule in the presence of + // sufficiently-complex constraints. Therefore, this is only safe if + // the ASAP scheduler gives the same pipeline length when only def-use + // constraints are in play; otherwise, we let ourselves fall through + // to a full model check. + if (!pipeline_stages.has_value()) { + ASAPScheduler asap_for_length(model_.graph(), asap_.delay_estimator()); + absl::StatusOr unrestricted_asap_bounds = + asap_for_length.ComputeBounds(std::nullopt, clock_period_ps, + worst_case_throughput, false); + if (unrestricted_asap_bounds.ok() && + unrestricted_asap_bounds->max_lower_bound() == + bounds->max_lower_bound()) { + pipeline_stages = bounds->max_lower_bound() + 1; + model_.SetPipelineLength(*pipeline_stages); + } + } + } else { + VLOG(1) << "Proceeding with default bounds; ASAP bound computation failed: " + << bounds.status(); + } + model_.SetClockPeriod(clock_period_ps); // TODO(allight): Having the II be sort of held in the model is a footgun // since it can be not clear what the II being targeted is. For now just force @@ -1412,7 +1491,8 @@ absl::StatusOr SDCScheduler::Schedule( model_.SetPipelineLength(pipeline_stages); if (!pipeline_stages.has_value() && !check_feasibility_) { - // Find the minimum feasible pipeline length. + // The ASAP scheduler must have failed to compute bounds. + // We still need to find the minimum feasible pipeline length. model_.MinimizePipelineLength(); XLS_ASSIGN_OR_RETURN( const math_opt::SolveResult result_with_minimized_pipeline_length, @@ -1427,6 +1507,15 @@ absl::StatusOr SDCScheduler::Schedule( model_.ExtractPipelineLength( result_with_minimized_pipeline_length.variable_values())); model_.SetPipelineLength(min_pipeline_length); + pipeline_stages = min_pipeline_length; + } + + if (!bounds.ok() && pipeline_stages.has_value()) { + // The ASAP scheduler failed to compute bounds, but we can at least bound + // each cycle variable to be < `pipeline_stages` to save some solver time. + for (const auto& [node, _] : model_.GetCycleVars()) { + model_.SetNodeBounds(node, 0, *pipeline_stages - 1); + } } if (check_feasibility_) { diff --git a/xls/scheduling/sdc_scheduler.h b/xls/scheduling/sdc_scheduler.h index 0563c96b2b..7748c2f2d8 100644 --- a/xls/scheduling/sdc_scheduler.h +++ b/xls/scheduling/sdc_scheduler.h @@ -29,9 +29,9 @@ #include "absl/status/statusor.h" #include "absl/types/span.h" #include "xls/estimators/delay_model/delay_estimator.h" -#include "xls/ir/function_base.h" #include "xls/ir/node.h" #include "xls/ir/nodes.h" +#include "xls/scheduling/asap_scheduler.h" #include "xls/scheduling/schedule_graph.h" #include "xls/scheduling/scheduler.h" #include "xls/scheduling/scheduling_options.h" @@ -47,9 +47,10 @@ class SDCSchedulingModel { using DelayMap = absl::flat_hash_map; static constexpr double kInfinity = std::numeric_limits::infinity(); - static constexpr double kMaxStages = (1 << 20); public: + static constexpr double kMaxStages = (1 << 20); + SDCSchedulingModel( const ScheduleGraph& graph, const DelayMap& delay_map, std::optional initiation_interval, double sdc_solution_tolerance, @@ -83,6 +84,9 @@ class SDCSchedulingModel { return initiation_interval_; } + void SetNodeBounds(Node* node, int64_t lower_bound, int64_t upper_bound); + void RemoveNodeBounds(Node* node); + void SetPipelineLength(std::optional pipeline_length); void MinimizePipelineLength(); @@ -319,6 +323,24 @@ class SDCScheduler final : public Scheduler { const operations_research::math_opt::SolveResult& result, SchedulingFailureBehavior failure_behavior); + class PrecomputedDelayEstimator : public DelayEstimator { + public: + explicit PrecomputedDelayEstimator( + const absl::flat_hash_map& delay_map) + : DelayEstimator("precomputed_delay_estimator"), + delay_map_(delay_map) {} + + absl::StatusOr GetOperationDelayInPs(Node* node) const override { + if (auto it = delay_map_.find(node); it != delay_map_.end()) { + return it->second; + } + return 0; + } + + private: + const absl::flat_hash_map& delay_map_; + }; + DelayMap delay_map_; ::operations_research::math_opt::SolverType solver_type_; ::operations_research::math_opt::SolveParameters solve_parameters_; @@ -326,6 +348,12 @@ class SDCScheduler final : public Scheduler { std::unique_ptr solver_; std::optional dynamic_throughput_objective_weight_; bool check_feasibility_ = false; + + // This holds the ASAP scheduler used to give the SDC solver bounds on cycle + // numbers for each node, along with a representation of our delay map as a + // DelayEstimator. + PrecomputedDelayEstimator delay_map_as_estimator_; + ASAPScheduler asap_; }; } // namespace xls