[scheduling] Use ASAP/ALAP bounds in the SDC scheduler

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",

xls/scheduling/asap_scheduler.cc

@@ -143,7 +143,8 @@ absl::StatusOr<ScheduleCycleMap> ASAPScheduler::Schedule(
 
 absl::StatusOr<sched::ScheduleBounds> ASAPScheduler::ComputeBounds(
     std::optional<int64_t> pipeline_stages, int64_t clock_period_ps,
-    std::optional<int64_t> worst_case_throughput) {
+    std::optional<int64_t> worst_case_throughput, bool get_helpful_error,
+    int64_t max_upper_bound) {
   XLS_RET_CHECK(std::holds_alternative<FunctionBase*>(graph_.ir_scope()));
   auto* f = std::get<FunctionBase*>(graph_.ir_scope());
   // TODO(allight): This actually creates a copy of graph_ since it needs to
@@ -158,9 +159,9 @@ absl::StatusOr<sched::ScheduleBounds> 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<sched::ScheduleBounds> 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;
 }

xls/scheduling/asap_scheduler.h

@@ -51,6 +51,14 @@ class ASAPScheduler : public Scheduler {
       SchedulingFailureBehavior failure_behavior,
       std::optional<int64_t> 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<sched::ScheduleBounds> ComputeBounds(
+      std::optional<int64_t> pipeline_stages, int64_t clock_period_ps,
+      std::optional<int64_t> 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<const SchedulingConstraint> constraints() const {
@@ -68,12 +76,8 @@ class ASAPScheduler : public Scheduler {
       absl::Status&& orig_status, std::optional<int64_t> pipeline_stages,
       int64_t clock_period_ps, std::optional<int64_t> worst_case_throughput);
 
-  // Exposed to allow for Min-cut and random to be built on top of this.
-  absl::StatusOr<sched::ScheduleBounds> ComputeBounds(
-      std::optional<int64_t> pipeline_stages, int64_t clock_period_ps,
-      std::optional<int64_t> 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<int64_t> schedule_length);

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<double>(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<int64_t> 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<ScheduleCycleMap> 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<sched::ScheduleBounds> 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<sched::ScheduleBounds> 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<ScheduleCycleMap> 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<ScheduleCycleMap> 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_) {

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<Node*, int64_t>;
 
   static constexpr double kInfinity = std::numeric_limits<double>::infinity();
-  static constexpr double kMaxStages = (1 << 20);
 
  public:
+  static constexpr double kMaxStages = (1 << 20);
+
   SDCSchedulingModel(
       const ScheduleGraph& graph, const DelayMap& delay_map,
       std::optional<int64_t> 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<int64_t> pipeline_length);
   void MinimizePipelineLength();
 
@@ -319,13 +323,37 @@ 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<Node*, int64_t>& delay_map)
+        : DelayEstimator("precomputed_delay_estimator"),
+          delay_map_(delay_map) {}
+
+    absl::StatusOr<int64_t> 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<Node*, int64_t>& delay_map_;
+  };
+
   DelayMap delay_map_;
   ::operations_research::math_opt::SolverType solver_type_;
   ::operations_research::math_opt::SolveParameters solve_parameters_;
   SDCSchedulingModel model_;
   std::unique_ptr<operations_research::math_opt::IncrementalSolver> solver_;
   std::optional<double> 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