[Explicit State Access] Support multiple State Reads nodes per State Element in proc_state_legalization_pass

Only apply state_read predicate legalization (incorporating next_value predicates) for state elements with exactly one state_read.

Prevent nullptr from attempting to fetch State Reads from Next Nodes constructed from a State Element.

PiperOrigin-RevId: 910694399

Author: NL02

xls/scheduling/BUILD

@@ -459,6 +459,7 @@ cc_library(
         "//xls/ir:value",
         "//xls/passes:pass_base",
         "//xls/solvers:z3_ir_translator",
+        "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:btree",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",

xls/scheduling/proc_state_legalization_pass.cc

@@ -20,6 +20,7 @@
 #include <variant>
 #include <vector>
 
+#include "absl/algorithm/container.h"
 #include "absl/container/btree_set.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/log/check.h"
@@ -53,9 +54,17 @@ namespace {
 absl::StatusOr<bool> LegalizeStateReadPredicate(
     Proc* proc, StateElement* state_element,
     const SchedulingPassOptions& options) {
-  StateRead* state_read = proc->GetStateReadByStateElement(state_element);
+  absl::Span<StateRead* const> reads =
+      proc->GetStateReadsByStateElement(state_element);
   const absl::btree_set<Next*, Node::NodeIdLessThan>& next_values =
       proc->next_values(state_element);
+
+  // Only legalize predicate if there is exactly one state_read.
+  if (reads.size() != 1) {
+    return false;
+  }
+
+  StateRead* state_read = reads.front();
   if (!state_read->predicate().has_value() || next_values.empty()) {
     // Already unconditional, or no explicit `next_value`s; nothing to do.
     return false;
@@ -66,7 +75,7 @@ absl::StatusOr<bool> LegalizeStateReadPredicate(
   predicates.reserve(1 + next_values.size());
   predicates_set.reserve(next_values.size());
   for (Next* next : next_values) {
-    if (next->state_read() == next->value()) {
+    if (state_read == next->value()) {
       // This is a no-op next_value; we will narrow it to the case where the
       // state read is active instead.
       continue;
@@ -215,8 +224,12 @@ absl::StatusOr<bool> AddMutualExclusionAsserts(
 absl::StatusOr<bool> AddWriteWithoutReadAsserts(
     Proc* proc, StateElement* state_element,
     const SchedulingPassOptions& options) {
-  StateRead* state_read = proc->GetStateReadByStateElement(state_element);
-  if (!state_read->predicate().has_value()) {
+  absl::Span<StateRead* const> reads =
+      proc->GetStateReadsByStateElement(state_element);
+
+  if (absl::c_any_of(reads, [](StateRead* state_read) {
+        return !state_read->predicate().has_value();
+      })) {
     return false;
   }
 
@@ -226,7 +239,23 @@ absl::StatusOr<bool> AddWriteWithoutReadAsserts(
     return false;
   }
 
-  std::vector<Node*> predicate_list;
+  std::vector<Node*> read_predicates;
+  read_predicates.reserve(reads.size());
+  for (StateRead* read : reads) {
+    read_predicates.push_back(*read->predicate());
+  }
+
+  Node* any_read_active;
+  if (reads.size() == 1) {
+    any_read_active = *reads[0]->predicate();
+  } else {
+    XLS_ASSIGN_OR_RETURN(
+        any_read_active,
+        proc->MakeNodeWithName<NaryOp>(
+            SourceInfo(), read_predicates, Op::kOr,
+            absl::StrCat("__", state_element->name(), "__any_read_active")));
+  }
+
   for (Next* next : next_values) {
     XLS_RET_CHECK(next->predicate().has_value());
     XLS_ASSIGN_OR_RETURN(
@@ -239,8 +268,7 @@ absl::StatusOr<bool> AddWriteWithoutReadAsserts(
         Node * no_write_without_read,
         proc->MakeNodeWithName<NaryOp>(
             SourceInfo(),
-            absl::MakeConstSpan({*state_read->predicate(), next_not_triggered}),
-            Op::kOr,
+            absl::MakeConstSpan({any_read_active, next_not_triggered}), Op::kOr,
             absl::StrCat("__", state_element->name(), "__no_next_", next->id(),
                          "_without_read")));
     std::string label = absl::StrCat("__", state_element->name(), "__next_",
@@ -291,7 +319,18 @@ absl::StatusOr<bool> AddDefaultNextValue(Proc* proc,
                                          StateElement* state_element,
                                          const SchedulingPassOptions& options) {
   absl::btree_set<Node*, Node::NodeIdLessThan> predicates;
-  StateRead* state_read = proc->GetStateReadByStateElement(state_element);
+  absl::Span<StateRead* const> reads =
+      proc->GetStateReadsByStateElement(state_element);
+  StateRead* state_read = nullptr;
+  for (StateRead* read : reads) {
+    if (!read->predicate().has_value()) {
+      state_read = read;
+      break;
+    }
+  }
+  if (state_read == nullptr) {
+    state_read = reads.front();
+  }
   for (Next* next : proc->next_values(state_element)) {
     if (next->predicate().has_value()) {
       predicates.insert(*next->predicate());

xls/scheduling/proc_state_legalization_pass_test.cc

@@ -549,6 +549,67 @@ TEST_P(ProcStateLegalizationPassTest,
                                  m::Literal(0))))))));
 }
 
+TEST_P(ProcStateLegalizationPassTest,
+       ProcWithMultipleStateReadsAndDefaultNextValue) {
+  auto p = CreatePackage();
+  ProcBuilder pb("p", p.get());
+  XLS_ASSERT_OK_AND_ASSIGN(
+      StateElement * x_se,
+      pb.proc()->AppendUnreadStateElement("x", Value(UBits(0, 32))));
+  BValue x = pb.StateRead(x_se);
+
+  XLS_ASSERT_OK_AND_ASSIGN(
+      StateElement * y_se,
+      pb.proc()->AppendUnreadStateElement("y", Value(UBits(0, 32))));
+
+  BValue x_even =
+      pb.Eq(pb.UMod(x, pb.Literal(UBits(2, 32))), pb.Literal(UBits(0, 32)));
+  BValue x_multiple_of_3 =
+      pb.Eq(pb.UMod(x, pb.Literal(UBits(3, 32))), pb.Literal(UBits(0, 32)));
+
+  BValue y_read1 = pb.StateRead(y_se, x_even);
+  BValue y_read2 = pb.StateRead(y_se, x_multiple_of_3);
+
+  BValue y_val1 = pb.Add(y_read1, pb.Literal(UBits(3, 32)));
+  BValue y_val2 = pb.Add(y_read2, pb.Literal(UBits(2, 32)));
+
+  BValue y_cond = pb.Select(x_even, {y_val2, y_val1});
+
+  BValue cond_write =
+      pb.Eq(pb.UMod(x, pb.Literal(UBits(4, 32))), pb.Literal(UBits(0, 32)));
+  pb.Next(y_se, pb.Add(y_cond, pb.Literal(UBits(1, 32))), cond_write);
+
+  pb.Next(x_se, pb.Add(x, pb.Literal(UBits(1, 32))));
+
+  XLS_ASSERT_OK_AND_ASSIGN(Proc * proc, pb.Build());
+  XLS_ASSERT_OK(p->SetTop(proc));
+
+  ASSERT_THAT(Run(proc), IsOkAndHolds(true));
+
+  // 1. Verify the generated safety assertions
+  std::vector<Node*> asserts;
+  absl::c_copy_if(proc->nodes(), std::back_inserter(asserts),
+                  [](Node* node) { return node->Is<Assert>(); });
+
+  auto x_even_matcher =
+      m::Eq(m::UMod(m::StateRead("x"), m::Literal(2)), m::Literal(0));
+  auto x_multiple_of_3_matcher =
+      m::Eq(m::UMod(m::StateRead("x"), m::Literal(3)), m::Literal(0));
+  auto cond_write_matcher =
+      m::Eq(m::UMod(m::StateRead("x"), m::Literal(4)), m::Literal(0));
+
+  EXPECT_THAT(asserts,
+              Contains(m::Assert(
+                  _, m::Or(m::Or(x_even_matcher, x_multiple_of_3_matcher),
+                           m::Not(cond_write_matcher)))));
+
+  // 2. Verify the generated default next-state feedback loop
+  EXPECT_THAT(
+      proc->next_values(y_se),
+      Contains(m::Next(m::StateRead("y"), m::StateRead("y"),
+                       m::And(x_even_matcher, m::Not(cond_write_matcher)))));
+}
+
 INSTANTIATE_TEST_SUITE_P(ProcStateLegalizationPassTestSuite,
                          ProcStateLegalizationPassTest,
                          testing::Values(false, true));