[Explicit State Access] Use state_element directly instead of state_read when determining mutual exclusion zones in pipeline_register_insertion.

xls/codegen_v_1_5/BUILD

@@ -295,8 +295,11 @@ cc_test(
         "//xls/codegen:codegen_options",
         "//xls/common:xls_gunit_main",
         "//xls/common/status:matchers",
+        "//xls/ir:bits",
+        "//xls/ir:channel",
         "//xls/ir:function_builder",
         "//xls/ir:ir_test_base",
+        "//xls/ir:value",
         "//xls/passes:optimization_pass",
         "//xls/passes:pass_base",
         "//xls/scheduling:scheduling_options",

xls/codegen_v_1_5/convert_to_block_test.cc

@@ -18,8 +18,11 @@
 #include "gtest/gtest.h"
 #include "xls/codegen/codegen_options.h"
 #include "xls/common/status/matchers.h"
+#include "xls/ir/bits.h"
+#include "xls/ir/channel.h"
 #include "xls/ir/function_builder.h"
 #include "xls/ir/ir_test_base.h"
+#include "xls/ir/value.h"
 #include "xls/passes/optimization_pass.h"
 #include "xls/passes/pass_base.h"
 #include "xls/scheduling/scheduling_options.h"
@@ -66,5 +69,33 @@ TEST_F(ConvertToBlockTest, SimpleFunction) {
   // TODO: https://github.com/google/xls/issues/3356 - assert stuff.
 }
 
+TEST_F(ConvertToBlockTest, ProcWithExplicitStateAccessNextValueStateElement) {
+  auto p = CreatePackage();
+  TokenlessProcBuilder pb(NewStyleProc(), TestName(), "tkn", p.get());
+  XLS_ASSERT_OK_AND_ASSIGN(SendChannelInterface * out_ch_interface,
+                           pb.AddOutputChannel("out_ch", p->GetBitsType(32)));
+
+  XLS_ASSERT_OK_AND_ASSIGN(StateElement * se,
+                           pb.UnreadStateElement("state", Value(UBits(0, 32))));
+  BValue state_read = pb.StateRead(se);
+  BValue current = pb.Identity(state_read);
+  BValue add_val = pb.Add(current, pb.Literal(UBits(1, 32)));
+
+  pb.Send(out_ch_interface, add_val);
+
+  pb.Next(se, add_val);
+
+  XLS_ASSERT_OK_AND_ASSIGN(Proc * top, pb.Build());
+  XLS_ASSERT_OK(p->SetTop(top));
+
+  OptimizationContext opt_context;
+  PassResults pass_results;
+  XLS_ASSERT_OK(ConvertToBlock(
+      p.get(),
+      codegen_options().clock_name("clk").reset("rst", false, false, false),
+      scheduling_options().opt_level(0), &delay_estimator_, &opt_context,
+      &pass_results));
+}
+
 }  // namespace
 }  // namespace xls::codegen

xls/codegen_v_1_5/pipeline_register_insertion_pass.cc

@@ -45,6 +45,7 @@
 #include "xls/ir/nodes.h"
 #include "xls/ir/op.h"
 #include "xls/ir/package.h"
+#include "xls/ir/proc.h"
 #include "xls/ir/register.h"
 #include "xls/ir/source_location.h"
 #include "xls/ir/state_element.h"
@@ -192,32 +193,49 @@ absl::StatusOr<ConcurrentStageGroups> CalculateConcurrentStages(
   if (block->source() == nullptr || !block->source()->IsProc()) {
     return result;
   }
+  Proc* proc = block->source()->AsProcOrDie();
 
   // Find all the mutex regions created by unconditional state feedback.
   absl::flat_hash_map<StateElement*, int64_t> read_by_stage;
   absl::flat_hash_map<StateElement*, int64_t> first_write_stage;
-  for (Node* node : block->nodes()) {
-    if (!node->Is<Next>()) {
-      continue;
-    }
-    Next* next = node->As<Next>();
-    StateRead* state_read = node->As<Next>()->state_read()->As<StateRead>();
-    StateElement* state_element = state_read->state_element();
-    if (state_read->predicate().has_value()) {
+
+  // 1. Find the earliest unconditional read stage for each state element.
+  // During block lowering, StateRead nodes remain owned by the source Proc.
+  for (Node* node : proc->nodes()) {
+    if (node->Is<StateRead>()) {
+      StateRead* state_read = node->As<StateRead>();
       // If the state read is predicated, then it doesn't start a mutual
       // exclusion zone.
-      continue;
+      if (state_read->predicate().has_value()) {
+        continue;
+      }
+      StateElement* state_element = state_read->state_element();
+      XLS_ASSIGN_OR_RETURN(int64_t read_stage,
+                           block->GetStageIndex(state_read));
+      auto [it, inserted] =
+          read_by_stage.try_emplace(state_element, read_stage);
+      if (!inserted) {
+        it->second = std::min(it->second, read_stage);
+      }
     }
-    XLS_ASSIGN_OR_RETURN(read_by_stage[state_element],
-                         block->GetStageIndex(state_read));
-
-    XLS_ASSIGN_OR_RETURN(int64_t write_stage, block->GetStageIndex(next));
-    auto [it, inserted] =
-        first_write_stage.try_emplace(state_element, write_stage);
-    if (!inserted) {
-      it->second = std::min(it->second, write_stage);
+  }
+
+  // 2. Find the earliest write stage for each state element.
+  // Next nodes are added directly to the Block and are owned by the Block.
+  for (Node* node : block->nodes()) {
+    if (node->Is<Next>()) {
+      Next* next = node->As<Next>();
+      StateElement* state_element = next->state_element();
+      XLS_ASSIGN_OR_RETURN(int64_t write_stage, block->GetStageIndex(next));
+      auto [it, inserted] =
+          first_write_stage.try_emplace(state_element, write_stage);
+      if (!inserted) {
+        it->second = std::min(it->second, write_stage);
+      }
     }
   }
+
+  // 3. Mark mutual exclusion zones.
   for (const auto& [state_element, read_stage] : read_by_stage) {
     auto write_stage_it = first_write_stage.find(state_element);
     if (write_stage_it == first_write_stage.end()) {
@@ -253,8 +271,7 @@ absl::StatusOr<bool> PipelineRegisterInsertionPass::InsertPipelineRegisters(
   for (Node* node : block->nodes()) {
     if (node->Is<Next>()) {
       Next* next = node->As<Next>();
-      StateElement* state_element =
-          next->state_read()->As<StateRead>()->state_element();
+      StateElement* state_element = next->state_element();
       int64_t stage = *block->GetStageIndex(next);
       auto [it, inserted] =
           earliest_next_stage.try_emplace(state_element, stage);
@@ -351,8 +368,9 @@ absl::StatusOr<bool> PipelineRegisterInsertionPass::InsertPipelineRegisters(
             // corresponding StateRead, we need to avoid updating that specific
             // operand, since no data is being passed. For simplicity, we put it
             // back afterwards rather than actually avoiding the update.
-            bool restore_state_read =
-                user->Is<Next>() && user->As<Next>()->state_read() == node;
+            bool restore_state_read = user->Is<Next>() &&
+                                      user->As<Next>()->has_state_read() &&
+                                      user->As<Next>()->state_read() == node;
             user->ReplaceOperand(node, live_node);
             if (restore_state_read) {
               XLS_RETURN_IF_ERROR(user->As<Next>()->ReplaceOperandNumber(

xls/codegen_v_1_5/pipeline_register_insertion_pass_test.cc

@@ -450,6 +450,51 @@ TEST_F(PipelineRegisterInsertionPassTest, TestCombinedRegisters) {
               ElementsAre(m2::Register("p0_v0", m::Type("bits[32]"))));
 }
 
+TEST_F(PipelineRegisterInsertionPassTest,
+       TestCombinedRegistersWithStateElementNext) {
+  auto p = CreatePackage();
+  ScheduledBlockBuilder sbb(TestName(), p.get());
+  Proc* source;
+  {
+    std::unique_ptr<Proc> owned_source = std::make_unique<Proc>(
+        absl::StrCat("__", TestName(), "_source"), p.get());
+    source = owned_source.get();
+    sbb.SetSource(std::move(owned_source));
+  }
+  XLS_ASSERT_OK(sbb.block()->AddClockPort("clk"));
+  BValue x = sbb.InputPort("x", p->GetBitsType(32));
+
+  sbb.StartStage(sbb.Literal(UBits(1, 1)), sbb.Literal(UBits(1, 1)));
+  XLS_ASSERT_OK_AND_ASSIGN(Node * source_acc, source->AppendStateElement(
+                                                  "acc", Value(UBits(0, 32))));
+  BValue acc = sbb.SourceNode(source_acc);
+  sbb.AddStateReadToCurrentStage(acc);
+  BValue v0 = sbb.Add(x, acc, SourceInfo(), "v0");
+  sbb.EndStage(sbb.Literal(UBits(1, 1)), sbb.Literal(UBits(1, 1)));
+
+  sbb.StartStage(sbb.Literal(UBits(1, 1)), sbb.Literal(UBits(1, 1)));
+  sbb.EndStage(sbb.Literal(UBits(1, 1)), sbb.Literal(UBits(1, 1)));
+
+  sbb.StartStage(sbb.Literal(UBits(1, 1)), sbb.Literal(UBits(1, 1)));
+  BValue neg_v0 = sbb.Negate(v0);
+
+  // Associate Next directly with the StateElement.
+  StateElement* se = source_acc->As<StateRead>()->state_element();
+  sbb.Next(se, neg_v0);
+
+  sbb.OutputPort("out", neg_v0);
+  sbb.EndStage(sbb.Literal(UBits(1, 1)), sbb.Literal(UBits(1, 1)));
+
+  XLS_ASSERT_OK_AND_ASSIGN(ScheduledBlock * sb, sbb.Build());
+
+  BlockConversionPassOptions options;
+  options.codegen_options.register_merge_strategy(
+      verilog::CodegenOptions::RegisterMergeStrategy::kIdentityOnly);
+  EXPECT_THAT(Run(p.get(), options), IsOkAndHolds(true));
+  EXPECT_THAT(sb->GetRegisters(),
+              ElementsAre(m2::Register("p0_v0", m::Type("bits[32]"))));
+}
+
 TEST_F(PipelineRegisterInsertionPassTest, TestRegistersDontCombine) {
   auto p = CreatePackage();
   ScheduledBlockBuilder sbb(TestName(), p.get());

xls/scheduling/BUILD

@@ -460,6 +460,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",
@@ -488,7 +489,9 @@ cc_test(
         "//xls/ir:function_builder",
         "//xls/ir:ir_matcher",
         "//xls/ir:ir_test_base",
+        "//xls/ir:op",
         "//xls/ir:proc_conversion",
+        "//xls/ir:source_location",
         "//xls/ir:value",
         "//xls/passes:cse_pass",
         "//xls/passes:dce_pass",
@@ -497,6 +500,7 @@ cc_test(
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/types:span",
         "@googletest//:gtest",
     ],
 )