[Explicit State Access] Update xls/passes code to use GetStateReads and GetStateReadsByStateElement

This feature isn't usable as it will fail on passes with more than one state read. This CL is focused more on the refactor.

PiperOrigin-RevId: 895997589

Author: NL02

xls/ir/ir_matcher.cc

@@ -628,6 +628,11 @@ bool StateReadMatcher::MatchAndExplain(
     *listener << " has incorrect label";
     return false;
   }
+  if (predicate_.has_value() &&
+      !predicate_->MatchAndExplain(node->As<xls::StateRead>()->predicate(),
+                                   listener)) {
+    return false;
+  }
   return true;
 }
 
@@ -646,6 +651,12 @@ void StateReadMatcher::DescribeTo(::std::ostream* os) const {
     label_->DescribeTo(&ss);
     additional_fields.push_back(ss.str());
   }
+  if (predicate_.has_value()) {
+    std::stringstream ss;
+    ss << "predicate=";
+    predicate_->DescribeTo(&ss);
+    additional_fields.push_back(ss.str());
+  }
   DescribeToHelper(os, additional_fields);
 }
 

xls/ir/ir_matcher.h

@@ -1269,16 +1269,20 @@ inline ::testing::Matcher<const ::xls::Node*> OutputPort(
 //   EXPECT_THAT(x, m::StateRead());
 //   EXPECT_THAT(x, m::StateRead("x"));
 //   EXPECT_THAT(x, m::StateRead(HasSubstr("substr")));
+//   EXPECT_THAT(x, m::StateRead("x", /*predicate=*/m::Param("pred")));
 //
 class StateReadMatcher : public NodeMatcher {
  public:
   explicit StateReadMatcher(
       std::optional<::testing::Matcher<const std::string>> state_element_name,
       std::optional<::testing::Matcher<const std::optional<std::string>&>>
-          label = std::nullopt)
+          label = std::nullopt,
+      std::optional<::testing::Matcher<std::optional<Node*>>> predicate =
+          std::nullopt)
       : NodeMatcher(Op::kStateRead, /*operands=*/{}),
         state_element_name_(std::move(state_element_name)),
-        label_(std::move(label)) {}
+        label_(std::move(label)),
+        predicate_(std::move(predicate)) {}
 
   bool MatchAndExplain(const Node* node,
                        ::testing::MatchResultListener* listener) const override;
@@ -1287,6 +1291,7 @@ class StateReadMatcher : public NodeMatcher {
  private:
   std::optional<::testing::Matcher<const std::string>> state_element_name_;
   std::optional<::testing::Matcher<const std::optional<std::string>&>> label_;
+  std::optional<::testing::Matcher<std::optional<Node*>>> predicate_;
 };
 
 template <typename T>
@@ -1324,6 +1329,14 @@ inline ::testing::Matcher<const ::xls::Node*> StateRead() {
   return ::xls::op_matchers::StateReadMatcher(std::nullopt, std::nullopt);
 }
 
+inline ::testing::Matcher<const ::xls::Node*> StateRead(
+    ::testing::Matcher<const std::string> name,
+    ::testing::Matcher<const std::optional<std::string>&> label,
+    ::testing::Matcher<std::optional<Node*>> predicate) {
+  return ::xls::op_matchers::StateReadMatcher(std::move(name), std::move(label),
+                                              std::move(predicate));
+}
+
 // Next matcher. Supported forms:
 //
 //   EXPECT_THAT(x, m::Next());

xls/passes/BUILD

@@ -2507,6 +2507,7 @@ xls_pass(
         "//xls/ir:type",
         "//xls/ir:value",
         "//xls/ir:value_utils",
+        "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:btree",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
@@ -2565,6 +2566,7 @@ xls_pass(
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
     ],
 )
 
@@ -2576,6 +2578,7 @@ xls_pass(
     deps = [
         ":optimization_pass",
         ":pass_base",
+        "//xls/common/status:ret_check",
         "//xls/common/status:status_macros",
         "//xls/ir",
         "//xls/ir:bits",

xls/passes/array_untuple_pass.cc

@@ -119,17 +119,32 @@ absl::StatusOr<absl::flat_hash_set<Node*>> FindExternalGroups(
     // Don't mess with params that are only used in identity updates. Would
     // infinite loop otherwise since we don't remove these very often.
     for (StateElement* state_element : f->AsProcOrDie()->StateElements()) {
-      StateRead* state_read =
-          f->AsProcOrDie()->GetStateReadByStateElement(state_element);
-      if (absl::c_all_of(state_read->users(), [&](Node* n) -> bool {
-            if (n->Is<Next>()) {
+      absl::Span<StateRead* const> state_reads =
+          f->AsProcOrDie()->GetStateReadsByStateElement(state_element);
+      bool all_reads_identity = true;
+      for (StateRead* state_read : state_reads) {
+        if (!absl::c_all_of(state_read->users(), [&](Node* n) -> bool {
+              if (!n->Is<Next>()) {
+                return false;
+              }
               Next* nxt = n->As<Next>();
-              return nxt->state_read() == nxt->value() &&
-                     nxt->state_read() == state_read;
-            }
-            return false;
-          })) {
-        excluded.insert(groups.Find(state_read));
+              if (nxt->state_element() != state_element) {
+                return false;
+              }
+              if (!nxt->value()->Is<StateRead>()) {
+                return false;
+              }
+              return nxt->value()->As<StateRead>()->state_element() ==
+                     state_element;
+            })) {
+          all_reads_identity = false;
+          break;
+        }
+      }
+      if (all_reads_identity) {
+        for (StateRead* state_read : state_reads) {
+          excluded.insert(groups.Find(state_read));
+        }
       }
     }
   }

xls/passes/array_untuple_pass_test.cc

@@ -458,6 +458,59 @@ TEST_F(ArrayUntuplePassTest, ProcStateArrayImplicitNext) {
                             m::StateElement(_, m::Type("bits[3][4]"))}));
 }
 
+TEST_F(ArrayUntuplePassTest, ProcStateArrayActiveMultiReadOptimization) {
+  auto p = CreatePackage();
+  XLS_ASSERT_OK_AND_ASSIGN(
+      auto chan_resp, p->CreateStreamingChannel("resp", ChannelOps::kSendOnly,
+                                                p->GetBitsType(3)));
+  ProcBuilder pb(TestName(), p.get());
+  BValue state = pb.StateElement(
+      "foo", ValueBuilder::ArrayB({
+                 ValueBuilder::Tuple({ValueBuilder::Bits(UBits(0, 1)),
+                                      ValueBuilder::Bits(UBits(1, 3))}),
+             }));
+
+  // Read #1: Extract the inner tuple value out of the array
+  BValue extracted_tuple = pb.ArrayIndex(state, {pb.Literal(UBits(0, 1))});
+  BValue extracted_data = pb.TupleIndex(extracted_tuple, 1);
+
+  // Send required by equivalence checker to observe proc mutations
+  pb.Send(chan_resp, pb.Literal(Value::Token()), extracted_data);
+
+  // Read #2: Perform an independent update on the state
+  BValue updated_tuple = pb.Tuple({pb.Literal(UBits(1, 1)), extracted_data});
+  BValue next_state =
+      pb.ArrayUpdate(state, updated_tuple, {pb.Literal(UBits(0, 1))});
+
+  XLS_ASSERT_OK_AND_ASSIGN(Proc * pr, pb.Build({next_state}));
+  solvers::z3::ScopedVerifyProcEquivalence svpe(pr, /*activation_count=*/2,
+                                                /*include_state=*/false);
+  ScopedRecordIr sri(p.get());
+  ASSERT_THAT(RunPass(p.get()), IsOkAndHolds(true));
+  EXPECT_THAT(pr->StateElements(),
+              IsSupersetOf({m::StateElement(_, m::Type("bits[1][1]")),
+                            m::StateElement(_, m::Type("bits[3][1]"))}));
+}
+
+TEST_F(ArrayUntuplePassTest, ProcStateArrayIdentityUpdateOnly) {
+  auto p = CreatePackage();
+  ProcBuilder pb(TestName(), p.get());
+  BValue state = pb.StateElement(
+      "foo", ValueBuilder::ArrayB({
+                 ValueBuilder::Tuple({ValueBuilder::Bits(UBits(0, 1)),
+                                      ValueBuilder::Bits((UBits(1, 3)))}),
+             }));
+  pb.Next(state, state);
+
+  XLS_ASSERT_OK(pb.Build().status());
+  ScopedRecordIr sri(p.get());
+
+  ArrayUntuplePass pass;
+  PassResults res;
+  OptimizationContext ctx;
+  ASSERT_THAT(pass.Run(p.get(), {}, &res, ctx), IsOkAndHolds(false));
+}
+
 void IrFuzzArrayUntuple(FuzzPackageWithArgs fuzz_package_with_args) {
   ArrayUntuplePass pass;
   OptimizationPassChangesOutputs(std::move(fuzz_package_with_args), pass);

xls/passes/canonicalization_pass_test.cc

@@ -49,9 +49,12 @@ namespace xls {
 namespace {
 
 using ::absl_testing::IsOkAndHolds;
+using ::testing::_;
 using ::testing::ElementsAre;
+using ::testing::Eq;
 using ::testing::IsEmpty;
 using ::testing::Optional;
+using ::testing::UnorderedElementsAre;
 
 class CanonicalizePassTest : public IrTestBase {
  protected:
@@ -348,10 +351,13 @@ TEST_F(CanonicalizePassTest, StateReadWithAlwaysTruePredicate) {
                              /*read_predicate=*/pb.Literal(UBits(1, 1)));
   pb.Next(x, pb.Literal(UBits(1, 32)));
   XLS_ASSERT_OK_AND_ASSIGN(Proc * proc, pb.Build());
-  EXPECT_THAT(proc->GetStateRead(0)->predicate(), Optional(m::Literal(1)));
+  EXPECT_THAT(
+      proc->GetStateReads(0).front(),
+      UnorderedElementsAre(m::StateRead(_, _, Optional(m::Literal(1)))));
 
   EXPECT_THAT(Run(p.get()), IsOkAndHolds(true));
-  EXPECT_EQ(proc->GetStateRead(0)->predicate(), std::nullopt);
+  EXPECT_THAT(proc->GetStateReads(0).front(),
+              UnorderedElementsAre(m::StateRead(_, _, Eq(std::nullopt))));
 }
 
 void IrFuzzCanonicalization(FuzzPackageWithArgs fuzz_package_with_args) {

xls/passes/conditional_specialization_pass_test.cc

@@ -1266,14 +1266,16 @@ TEST_F(ConditionalSpecializationPassTest, StateReadSpecialization) {
       Run(proc, /*use_bdd=*/true, /*optimize_for_best_case_throughput=*/true),
       IsOkAndHolds(true));
 
-  EXPECT_THAT(
-      proc->GetStateReadByStateElement(*proc->GetStateElementByName("counter0"))
-          ->predicate(),
-      Optional(m::Not(m::StateRead("index"))));
-  EXPECT_THAT(
-      proc->GetStateReadByStateElement(*proc->GetStateElementByName("counter1"))
-          ->predicate(),
-      Optional(m::StateRead("index")));
+  EXPECT_THAT(proc->GetStateReadsByStateElement(
+                      *proc->GetStateElementByName("counter0"))
+                  .front()
+                  ->predicate(),
+              Optional(m::Not(m::StateRead("index"))));
+  EXPECT_THAT(proc->GetStateReadsByStateElement(
+                      *proc->GetStateElementByName("counter1"))
+                  .front()
+                  ->predicate(),
+              Optional(m::StateRead("index")));
 }
 
 TEST_F(ConditionalSpecializationPassTest, HarderStateReadSpecialization) {
@@ -1303,16 +1305,20 @@ TEST_F(ConditionalSpecializationPassTest, HarderStateReadSpecialization) {
       Run(proc, /*use_bdd=*/true, /*optimize_for_best_case_throughput=*/true),
       IsOkAndHolds(true));
 
+  EXPECT_THAT(proc->GetStateReadsByStateElement(
+                      *proc->GetStateElementByName("counter0"))
+                  .front()
+                  ->predicate(),
+              Optional(m::Eq(m::StateRead("index"), m::Literal(0))));
+  EXPECT_THAT(proc->GetStateReadsByStateElement(
+                      *proc->GetStateElementByName("counter1"))
+                  .front()
+                  ->predicate(),
+              Optional(m::Eq(m::StateRead("index"), m::Literal(1))));
   EXPECT_THAT(
-      proc->GetStateReadByStateElement(*proc->GetStateElementByName("counter0"))
-          ->predicate(),
-      Optional(m::Eq(m::StateRead("index"), m::Literal(0))));
-  EXPECT_THAT(
-      proc->GetStateReadByStateElement(*proc->GetStateElementByName("counter1"))
-          ->predicate(),
-      Optional(m::Eq(m::StateRead("index"), m::Literal(1))));
-  EXPECT_THAT(
-      proc->GetStateReadByStateElement(*proc->GetStateElementByName("counter2"))
+      proc->GetStateReadsByStateElement(
+              *proc->GetStateElementByName("counter2"))
+          .front()
           ->predicate(),
       Optional(m::And(
           // High bit of index is set

xls/passes/proc_state_array_flattening_pass.cc

@@ -20,6 +20,7 @@
 #include "absl/log/log.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
+#include "absl/types/span.h"
 #include "xls/common/math_util.h"
 #include "xls/common/status/ret_check.h"
 #include "xls/common/status/status_macros.h"
@@ -131,10 +132,18 @@ absl::StatusOr<bool> SimplifyProcState(Proc* proc,
   Value new_init_value = Value::Tuple(old_init_value.elements());
 
   ArrayToTupleStateTransformer transformer;
-  XLS_RETURN_IF_ERROR(proc->TransformStateElement(
-                              proc->GetStateReadByStateElement(state_element),
-                              new_init_value, transformer)
-                          .status());
+  absl::Span<StateRead* const> state_reads =
+      proc->GetStateReadsByStateElement(state_element);
+
+  XLS_ASSIGN_OR_RETURN(StateRead * new_state_read,
+                       proc->TransformStateElement(
+                           state_reads.front(), new_init_value, transformer));
+  for (int64_t i = 0; i < state_reads.size(); ++i) {
+    XLS_ASSIGN_OR_RETURN(
+        Node * replacement,
+        transformer.TransformStateRead(proc, new_state_read, state_reads[i]));
+    XLS_RETURN_IF_ERROR(state_reads[i]->ReplaceUsesWith(replacement));
+  }
 
   std::vector<Next*> old_next_values(proc->next_values(state_element).begin(),
                                      proc->next_values(state_element).end());

xls/passes/proc_state_bits_shattering_pass.cc

@@ -27,6 +27,7 @@
 #include "absl/strings/str_format.h"
 #include "absl/types/span.h"
 #include "cppitertools/reversed.hpp"
+#include "xls/common/status/ret_check.h"
 #include "xls/common/status/status_macros.h"
 #include "xls/ir/bits.h"
 #include "xls/ir/node.h"
@@ -117,7 +118,12 @@ absl::StatusOr<bool> MaybeSplitStateElements(
     // to use STL set intersection algorithms.
     std::vector<int64_t> split_ends;
     bool could_benefit_from_splitting = false;
-    StateRead* state_read = proc->GetStateReadByStateElement(state_element);
+    absl::Span<StateRead* const> state_reads =
+        proc->GetStateReadsByStateElement(state_element);
+    XLS_RET_CHECK_EQ(state_reads.size(), 1)
+        << "ProcStateBitsShatteringPass only supports one StateRead per "
+           "StateElement for now.";
+    StateRead* state_read = state_reads.front();
     for (Next* next : proc->next_values(state_element)) {
       if (next->value() == state_read) {
         // This is a no-op next-value; it doesn't affect whether or not it's

xls/passes/proc_state_narrowing_pass.cc

@@ -161,7 +161,12 @@ absl::StatusOr<bool> ProcStateNarrowingPass::RunOnProcInternal(
               << state_element->type()->ToString();
       continue;
     }
-    StateRead* state_read = proc->GetStateReadByStateElement(state_element);
+    absl::Span<StateRead* const> state_reads =
+        proc->GetStateReadsByStateElement(state_element);
+    XLS_RET_CHECK_EQ(state_reads.size(), 1)
+        << "ProcStateNarrowingPass only supports one StateRead per "
+           "StateElement for now.";
+    StateRead* state_read = state_reads.front();
     std::optional<SharedLeafTypeTree<TernaryVector>> ternary =
         qe.GetTernary(state_read);
     if (!ternary) {