[Explicit State Access] Support decoupled nodes in proc_state_legalization.

xls/scheduling/BUILD

@@ -231,6 +231,8 @@ cc_test(
         "//xls/ir:function_builder",
         "//xls/ir:ir_test_base",
         "//xls/ir:value",
+        "@abseil-cpp//absl/status",
+        "@abseil-cpp//absl/status:status_matchers",
         "@abseil-cpp//absl/status:statusor",
         "@googletest//:gtest",
     ],
@@ -536,6 +538,7 @@ cc_test(
         "//xls/common/status:status_macros",
         "//xls/ir",
         "//xls/ir:bits",
+        "//xls/ir:channel",
         "//xls/ir:function_builder",
         "//xls/ir:ir_matcher",
         "//xls/ir:ir_test_base",

xls/scheduling/proc_state_legalization_pass.cc

@@ -61,6 +61,9 @@ absl::StatusOr<bool> LegalizeStateReadPredicate(
     // Already unconditional, or no explicit `next_value`s; nothing to do.
     return false;
   }
+  if (proc->uses_decoupled_next()) {
+    return false;
+  }
 
   std::vector<Node*> predicates;
   absl::flat_hash_set<Node*> predicates_set;
@@ -221,6 +224,10 @@ absl::StatusOr<bool> AddWriteWithoutReadAsserts(
     return false;
   }
 
+  if (proc->uses_decoupled_next()) {
+    return false;
+  }
+
   const absl::btree_set<Next*, Node::NodeIdLessThan>& next_values =
       proc->next_values(state_element);
   if (next_values.empty()) {
@@ -305,8 +312,11 @@ absl::StatusOr<bool> AddDefaultNextValue(Proc* proc,
   if (predicates.empty()) {
     // No explicit `next_value` node; leave the state element unchanged by
     // default.
+    Next::StateIdentifier state_identifier =
+        proc->uses_decoupled_next() ? Next::StateIdentifier(state_element)
+                                    : Next::StateIdentifier(state_read);
     XLS_RETURN_IF_ERROR(proc->MakeNodeWithName<Next>(
-                                state_read->loc(), /*state_read=*/state_read,
+                                state_read->loc(), state_identifier,
                                 /*value=*/state_read,
                                 /*predicate=*/state_read->predicate(),
                                 /*label=*/std::nullopt,
@@ -395,6 +405,9 @@ absl::StatusOr<bool> AddDefaultNextValue(Proc* proc,
       Node * default_predicate,
       NaryNorIfNeeded(proc, std::vector(predicates.begin(), predicates.end()),
                       /*name=*/"", state_read->loc()));
+  Next::StateIdentifier state_identifier =
+      proc->uses_decoupled_next() ? Next::StateIdentifier(state_element)
+                                  : Next::StateIdentifier(state_read);
   if (state_read->predicate().has_value()) {
     XLS_ASSIGN_OR_RETURN(
         default_predicate,
@@ -404,7 +417,7 @@ absl::StatusOr<bool> AddDefaultNextValue(Proc* proc,
             Op::kAnd));
   }
   XLS_RETURN_IF_ERROR(proc->MakeNodeWithName<Next>(
-                              state_read->loc(), /*state_read=*/state_read,
+                              state_read->loc(), state_identifier,
                               /*value=*/state_read,
                               /*predicate=*/default_predicate,
                               /*label=*/std::nullopt,

xls/scheduling/proc_state_legalization_pass_test.cc

@@ -27,6 +27,7 @@
 #include "xls/common/status/ret_check.h"
 #include "xls/common/status/status_macros.h"
 #include "xls/ir/bits.h"
+#include "xls/ir/channel.h"
 #include "xls/ir/function_builder.h"
 #include "xls/ir/ir_matcher.h"
 #include "xls/ir/ir_test_base.h"
@@ -552,6 +553,112 @@ TEST_P(ProcStateLegalizationPassTest,
                                  m::Literal(0))))))));
 }
 
+TEST_P(ProcStateLegalizationPassTest,
+       DecoupledUnconditionalReadAndNextNoDefaultNextValue) {
+  auto p = CreatePackage();
+  ProcBuilder pb("p", p.get());
+  XLS_ASSERT_OK_AND_ASSIGN(StateElement * x_element,
+                           pb.UnreadStateElement("x", Value(UBits(0, 32)),
+                                                 /*non_synthesizable=*/false));
+  BValue x_read = pb.StateRead(x_element);
+  BValue x_add = pb.Add(x_read, pb.Literal(UBits(1, 32)));
+  pb.Next(x_element, x_add);
+  XLS_ASSERT_OK_AND_ASSIGN(Proc * proc, pb.Build());
+  XLS_ASSERT_OK(p->SetTop(proc));
+  ASSERT_THAT(Run(proc), IsOkAndHolds(false));
+  EXPECT_EQ(proc->next_values().size(), 1);
+  EXPECT_THAT(
+      proc->next_values(x_element),
+      UnorderedElementsAre(m::NextWithStateElement(x_element, x_add.node())));
+}
+
+TEST_P(ProcStateLegalizationPassTest, DecoupledNoExplicitNextValueDefault) {
+  auto p = CreatePackage();
+  ProcBuilder pb("p", p.get());
+  XLS_ASSERT_OK_AND_ASSIGN(StateElement * x_element,
+                           pb.UnreadStateElement("x", Value(UBits(0, 32)),
+                                                 /*non_synthesizable=*/false));
+  BValue x_read = pb.StateRead(x_element);
+  BValue x_add = pb.Add(x_read, pb.Literal(UBits(1, 32)));
+  pb.Next(x_element, x_add);
+  XLS_ASSERT_OK_AND_ASSIGN(StateElement * y_element,
+                           pb.UnreadStateElement("y", Value(UBits(10, 32)),
+                                                 /*non_synthesizable=*/false));
+  BValue y_read = pb.StateRead(y_element);
+  XLS_ASSERT_OK_AND_ASSIGN(Proc * proc, pb.Build());
+  XLS_ASSERT_OK(p->SetTop(proc));
+  ASSERT_THAT(Run(proc), IsOkAndHolds(true));
+  EXPECT_EQ(proc->next_values().size(), 2);
+  EXPECT_THAT(
+      proc->next_values(x_element),
+      UnorderedElementsAre(m::NextWithStateElement(x_element, x_add.node())));
+  EXPECT_THAT(
+      proc->next_values(y_element),
+      UnorderedElementsAre(m::NextWithStateElement(y_element, y_read.node())));
+}
+
+TEST_P(ProcStateLegalizationPassTest, DecoupledPredicatedNextValue) {
+  auto p = CreatePackage();
+  ProcBuilder pb("p", p.get());
+  BValue cond = pb.StateElement("cond", Value(UBits(1, 1)));
+  XLS_ASSERT_OK_AND_ASSIGN(StateElement * x_element,
+                           pb.UnreadStateElement("x", Value(UBits(0, 32)),
+                                                 /*non_synthesizable=*/false));
+  BValue x_read = pb.StateRead(x_element, cond);
+  BValue incremented = pb.Add(x_read, pb.Literal(UBits(1, 32)));
+  BValue write_pred = pb.Eq(x_read, pb.Literal(UBits(0, 32)));
+  pb.Next(x_element, incremented, write_pred);
+  XLS_ASSERT_OK_AND_ASSIGN(Proc * proc, pb.Build());
+  XLS_ASSERT_OK(p->SetTop(proc));
+  ASSERT_THAT(Run(proc), IsOkAndHolds(true));
+  StateRead* read_node = proc->GetStateReadByStateElement(x_element);
+  EXPECT_EQ(*read_node->predicate(), cond.node());
+  EXPECT_THAT(proc->next_values(x_element),
+              UnorderedElementsAre(
+                  m::NextWithStateElement(x_element, incremented.node(),
+                                          write_pred.node()),
+                  m::NextWithStateElement(
+                      x_element, read_node,
+                      m::And(cond.node(), m::Not(write_pred.node())))));
+}
+
+TEST_P(ProcStateLegalizationPassTest, DecoupledMultiplePredicatedNextValues) {
+  auto p = CreatePackage();
+  ProcBuilder pb("p", p.get());
+  BValue cond = pb.StateElement("cond", Value(UBits(1, 1)));
+  XLS_ASSERT_OK_AND_ASSIGN(StateElement * x_element,
+                           pb.UnreadStateElement("x", Value(UBits(0, 32)),
+                                                 /*non_synthesizable=*/false));
+  BValue x_read = pb.StateRead(x_element, cond);
+  BValue incremented_1 = pb.Add(x_read, pb.Literal(UBits(1, 32)));
+  BValue write_pred_1 = pb.Eq(x_read, pb.Literal(UBits(0, 32)));
+  pb.Next(x_element, incremented_1, write_pred_1);
+
+  BValue incremented_2 = pb.Add(x_read, pb.Literal(UBits(2, 32)));
+  BValue write_pred_2 = pb.Eq(x_read, pb.Literal(UBits(1, 32)));
+  pb.Next(x_element, incremented_2, write_pred_2);
+
+  XLS_ASSERT_OK_AND_ASSIGN(Proc * proc, pb.Build());
+  XLS_ASSERT_OK(p->SetTop(proc));
+  ASSERT_THAT(Run(proc), IsOkAndHolds(true));
+  StateRead* read_node = proc->GetStateReadByStateElement(x_element);
+  EXPECT_EQ(*read_node->predicate(), cond.node());
+  EXPECT_THAT(proc->next_values(x_element),
+              UnorderedElementsAre(
+                  m::NextWithStateElement(x_element, incremented_1.node(),
+                                          write_pred_1.node()),
+                  m::NextWithStateElement(x_element, incremented_2.node(),
+                                          write_pred_2.node()),
+                  m::NextWithStateElement(
+                      x_element, read_node,
+                      m::And(cond.node(), m::Nor(write_pred_1.node(),
+                                                 write_pred_2.node())))));
+  std::vector<Node*> asserts;
+  absl::c_copy_if(proc->nodes(), std::back_inserter(asserts),
+                  [](Node* node) { return node->Is<Assert>(); });
+  EXPECT_EQ(asserts.size(), 1);
+}
+
 INSTANTIATE_TEST_SUITE_P(ProcStateLegalizationPassTestSuite,
                          ProcStateLegalizationPassTest,
                          testing::Values(false, true));

xls/scheduling/sdc_scheduler.cc

@@ -327,6 +327,10 @@ absl::Status SDCSchedulingModel::AddAllDefUseConstraints() {
             absl::StrFormat("state_read_before_write_%s_%s", read->GetName(),
                             next->GetName()));
 
+        if (!next->has_state_read()) {
+          XLS_RETURN_IF_ERROR(AddThroughputConstraint(read, next));
+        }
+
         VLOG(2) << "Setting state read-before-write constraint: "
                 << absl::StrFormat("cycle[%s] - cycle[%s] >= 0",
                                    next->GetName(), read->GetName());
@@ -347,10 +351,11 @@ absl::Status SDCSchedulingModel::AddDefUseConstraints(
 
   if (node->Is<StateRead>() && user.has_value() && user.value()->Is<Next>()) {
     Next* next = user.value()->As<Next>();
-    if (next->state_read() == node) {
+    if (next->has_state_read() && next->state_read() == node) {
       XLS_RETURN_IF_ERROR(AddThroughputConstraint(node->As<StateRead>(), next));
     }
-    if (next->value() != node && next->predicate() != node) {
+    if (next->has_state_read() && next->value() != node &&
+        next->predicate() != node) {
       XLS_RET_CHECK_EQ(next->state_read(), node);
       // We don't need to keep the param's value alive to this user, so no need
       // for a lifetime constraint.

xls/scheduling/sdc_scheduler_test.cc

@@ -17,7 +17,10 @@
 #include <cstdint>
 #include <optional>
 
+#include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
 #include "absl/status/statusor.h"
 #include "xls/common/status/matchers.h"
 #include "xls/estimators/delay_model/delay_estimator.h"
@@ -147,5 +150,57 @@ TEST_F(SDCSchedulerTest, WithIOConstraint) {
   EXPECT_EQ(cycle_map.at(send.node()), 1);
 }
 
+TEST_F(SDCSchedulerTest, DecoupledFeedbackLoopInfeasible) {
+  auto p = CreatePackage();
+  ProcBuilder pb(TestName(), p.get());
+  XLS_ASSERT_OK_AND_ASSIGN(StateElement * x_element,
+                           pb.UnreadStateElement("x", Value(UBits(0, 32)),
+                                                 /*non_synthesizable=*/false));
+  BValue x_read = pb.StateRead(x_element);
+  // Create a long chain of adds to increase the critical path length
+  // to 5. This means that for worst_case_throughput = 1, the constraint will be
+  // violated This test will fail for worst_case_throughput = 1, but pass for
+  // worst_case_throughput = 4
+
+  BValue add1 = pb.Add(x_read, pb.Literal(UBits(1, 32)));
+  BValue add2 = pb.Add(add1, pb.Literal(UBits(1, 32)));
+  BValue add3 = pb.Add(add2, pb.Literal(UBits(1, 32)));
+  BValue add4 = pb.Add(add3, pb.Literal(UBits(1, 32)));
+  BValue add5 = pb.Add(add4, pb.Literal(UBits(1, 32)));
+  pb.Next(x_element, add5);
+  XLS_ASSERT_OK_AND_ASSIGN(Proc * proc, pb.Build());
+  XLS_ASSERT_OK_AND_ASSIGN(ScheduleGraph graph,
+                           ScheduleGraph::Create(proc, {}));
+  TestDelayEstimator delay_estimator;
+  SchedulingOptions options;
+
+  // 1. Verify scheduling fails under tight worst_case_throughput = 1 (requires
+  // cycle difference <= 0, but min is 3)
+  XLS_ASSERT_OK_AND_ASSIGN(
+      auto infeasible_scheduler,
+      SDCScheduler::Create(graph, delay_estimator, options));
+  auto infeasible_result = infeasible_scheduler->Schedule(
+      std::nullopt, 2,
+      SchedulingFailureBehavior{.explain_infeasibility = false}, 1);
+  EXPECT_THAT(infeasible_result,
+              absl_testing::StatusIs(
+                  absl::StatusCode::kInternal,
+                  testing::HasSubstr("does not have an optimal solution")));
+
+  // 2. Verify scheduling succeeds under relaxed worst_case_throughput = 4
+  // (allows cycle difference <= 3, which matches min of 3)
+  XLS_ASSERT_OK_AND_ASSIGN(
+      auto feasible_scheduler,
+      SDCScheduler::Create(graph, delay_estimator, options));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      ScheduleCycleMap feasible_cycle_map,
+      feasible_scheduler->Schedule(
+          std::nullopt, 2,
+          SchedulingFailureBehavior{.explain_infeasibility = false}, 4));
+  EXPECT_EQ(feasible_cycle_map.at(*proc->next_values().begin()) -
+                feasible_cycle_map.at(x_read.node()),
+            3);
+}
+
 }  // namespace
 }  // namespace xls