From 5b593e73593c327f0c569b5f09f51885317f9661 Mon Sep 17 00:00:00 2001 From: Eric Astor Date: Sun, 28 Jun 2026 05:31:11 -0700 Subject: [PATCH] [scheduling] Use ASAP/ALAP bounds in the SDC scheduler By providing tighter initial ranges for the cycle variables in the LP model, the solver saves (apparently significant) time. This also lets us use the ASAP scheduler to compute the minimum possible pipeline length, if we need it. If the ASAP scheduler fails to converge (which is possible for complex constraints, but rare), we fall back to using the SDC scheduler directly. PiperOrigin-RevId: 939372611 --- xls/scheduling/BUILD | 1 + xls/scheduling/asap_scheduler.cc | 18 +++--- xls/scheduling/asap_scheduler.h | 14 +++-- xls/scheduling/sdc_scheduler.cc | 95 +++++++++++++++++++++++++++++++- xls/scheduling/sdc_scheduler.h | 32 ++++++++++- 5 files changed, 143 insertions(+), 17 deletions(-) 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