From 4afaa9c45756fe46bfbcd016fd0ed7928afefafd Mon Sep 17 00:00:00 2001
From: Eric Astor <epastor@google.com>
Date: Tue, 4 Aug 2026 07:14:43 -0700
Subject: [PATCH] [scheduling] Optimize AddCycleConstraints when re-applying
 scheduling

When trying to preserve an existing schedule, we previously did a quadratic-complexity scan for existing or contradictory constraints; we now just populate a hash map for linear-time lookup.

PiperOrigin-RevId: 959004301
---
 xls/scheduling/pipeline_schedule.h         |  2 +
 xls/scheduling/pipeline_scheduling_pass.cc | 56 +++++++++++++++-------
 2 files changed, 41 insertions(+), 17 deletions(-)

diff --git a/xls/scheduling/pipeline_schedule.h b/xls/scheduling/pipeline_schedule.h
index 1d5e4e7c88..e7d5babd98 100644
--- a/xls/scheduling/pipeline_schedule.h
+++ b/xls/scheduling/pipeline_schedule.h
@@ -84,6 +84,8 @@ class PipelineSchedule {
   // been placed in this schedule.
   int64_t cycle(const Node* node) const { return cycle_map_.at(node); }
 
+  bool empty() const { return cycle_map_.empty(); }
+
   // Returns the nodes scheduled in the given cycle. The node order is
   // guaranteed to be topological.
   absl::Span<Node* const> nodes_in_cycle(int64_t cycle) const;
diff --git a/xls/scheduling/pipeline_scheduling_pass.cc b/xls/scheduling/pipeline_scheduling_pass.cc
index 245cabb6cc..9ca9927ebe 100644
--- a/xls/scheduling/pipeline_scheduling_pass.cc
+++ b/xls/scheduling/pipeline_scheduling_pass.cc
@@ -43,28 +43,50 @@ namespace {
 // Adds cycle constraints from a PipelineSchedule into SchedulingOptions.
 absl::Status AddCycleConstraints(const PipelineSchedule& schedule,
                                  SchedulingOptions& scheduling_options) {
+  if (schedule.empty()) {
+    return absl::OkStatus();
+  }
+
+  absl::flat_hash_map<Node*, int64_t> existing_constraints;
+  for (const SchedulingConstraint& constraint :
+       scheduling_options.constraints()) {
+    if (std::holds_alternative<NodeInCycleConstraint>(constraint)) {
+      const NodeInCycleConstraint& nic =
+          std::get<NodeInCycleConstraint>(constraint);
+      auto [it, inserted] =
+          existing_constraints.try_emplace(nic.GetNode(), nic.GetCycle());
+      if (!inserted && it->second != nic.GetCycle()) {
+        return absl::InvalidArgumentError(absl::StrFormat(
+            "Node %s has conflicting cycle constraints: constrained to both %d "
+            "and %d",
+            nic.GetNode()->GetName(), it->second, nic.GetCycle()));
+      }
+    }
+  }
+
   for (int64_t c = 0; c < schedule.length(); ++c) {
     for (Node* node : schedule.nodes_in_cycle(c)) {
-      bool already_constrained = false;
-      for (const SchedulingConstraint& constraint :
-           scheduling_options.constraints()) {
-        if (std::holds_alternative<NodeInCycleConstraint>(constraint)) {
-          const NodeInCycleConstraint& nic =
-              std::get<NodeInCycleConstraint>(constraint);
-          if (nic.GetNode() == node) {
-            if (nic.GetCycle() != c) {
-              return absl::InvalidArgumentError(absl::StrFormat(
-                  "Schedule contradicts cycle constraints for node %s: "
-                  "scheduled for cycle %d, but constrained to cycle %d",
-                  node->GetName(), c, nic.GetCycle()));
-            }
-            already_constrained = true;
-          }
-        }
+      if (IsUntimed(node)) {
+        continue;
       }
-      if (!already_constrained && !IsUntimed(node)) {
+
+      auto it = existing_constraints.find(node);
+      if (it == existing_constraints.end()) {
+        // No existing constraint; add this one.
         scheduling_options.add_constraint(NodeInCycleConstraint(node, c));
+        continue;
+      }
+
+      // Otherwise, there's an existing constraint; we either match it, or fail
+      // with a contradiction.
+      if (it->second == c) {
+        // Already constrained to this cycle.
+        continue;
       }
+      return absl::InvalidArgumentError(
+          absl::StrFormat("Schedule contradicts cycle constraints for node %s: "
+                          "scheduled for cycle %d, but constrained to cycle %d",
+                          node->GetName(), c, it->second));
     }
   }
   return absl::OkStatus();