[Explicit State Access] Guard state_read references in schedule_graph and clone_nodes_into_block

PiperOrigin-RevId: 951672414

Author: NL02

xls/codegen/BUILD

@@ -2066,6 +2066,7 @@ cc_test(
         "//xls/ir:function_builder",
         "//xls/ir:op",
         "//xls/ir:proc_elaboration",
+        "//xls/ir:value",
         "//xls/scheduling:pipeline_schedule",
         "//xls/scheduling:run_pipeline_schedule",
         "//xls/scheduling:scheduling_options",

xls/codegen/clone_nodes_into_block_handler.cc

@@ -760,15 +760,18 @@ absl::Status CloneNodesIntoBlockHandler::HandleNextValue(Node* node,
                                                          Stage stage) {
   Proc* proc = function_base_->AsProcOrDie();
   Next* next = node->As<Next>();
-  StateElement* state_element =
-      next->state_read()->As<StateRead>()->state_element();
+  StateElement* state_element = next->state_element();
   XLS_ASSIGN_OR_RETURN(int64_t index,
                        proc->GetStateElementIndex(state_element));
 
   StateRegister& state_register = *result_.state_registers.at(index);
+  bool equivalent_value =
+      next->has_state_read()
+          ? next->value() == next->state_read()
+          : next->value() == proc->GetStateReadByStateElement(state_element);
   state_register.next_values.push_back(
       {.stage = stage,
-       .value = next->value() == next->state_read()
+       .value = equivalent_value
                     ? std::nullopt
                     : std::make_optional(node_map_.at(next->value())),
        .predicate =

xls/codegen/synchronous_procs_test.cc

@@ -38,6 +38,7 @@
 #include "xls/ir/op.h"
 #include "xls/ir/package.h"
 #include "xls/ir/proc_elaboration.h"
+#include "xls/ir/value.h"
 #include "xls/scheduling/pipeline_schedule.h"
 #include "xls/scheduling/run_pipeline_schedule.h"
 #include "xls/scheduling/scheduling_options.h"
@@ -244,6 +245,60 @@ TEST_P(SynchronousProcsTest, ChainedProc) {
               absl_testing::IsOkAndHolds(outputs));
 }
 
+TEST_P(SynchronousProcsTest, DecoupledNextProc) {
+  Package package(TestBaseName());
+
+  TokenlessProcBuilder pb(NewStyleProc(), "my_proc", "tkn", &package);
+  Type* u32 = package.GetBitsType(32);
+  BReceiveChannel in = pb.AddInputChannel("top_in", u32);
+  BSendChannel out = pb.AddOutputChannel("top_out", u32);
+  BStateElement st_elem = pb.UnreadStateElement("accum", Value(UBits(0, 32)),
+                                                /*non_synthesizable=*/false);
+  BValue st = pb.StateRead(st_elem);
+  BValue received = pb.Receive(in);
+  BValue next_st = pb.Add(st, received);
+
+  pb.Next(st_elem, next_st);
+  pb.Send(out, next_st);
+
+  XLS_ASSERT_OK_AND_ASSIGN(Proc * top, pb.Build({}));
+  XLS_ASSERT_OK(package.SetTop(top));
+
+  ASSERT_TRUE(top->uses_decoupled_next());
+
+  XLS_ASSERT_OK_AND_ASSIGN(ProcElaboration elab,
+                           ProcElaboration::Elaborate(top));
+
+  XLS_ASSERT_OK_AND_ASSIGN(PackageSchedule package_schedule,
+                           RunSynchronousPipelineSchedule(
+                               &package, TestDelayEstimator(),
+                               SchedulingOptions().clock_period_ps(1), elab));
+
+  ResetProto reset_proto;
+  reset_proto.set_name("rst");
+  reset_proto.set_asynchronous(false);
+  reset_proto.set_active_low(false);
+
+  XLS_ASSERT_OK_AND_ASSIGN(
+      CodegenResult result,
+      ToPipelineModuleText(
+          package_schedule, &package,
+          BuildPipelineOptions()
+              .reset(reset_proto.name(), reset_proto.asynchronous(),
+                     reset_proto.active_low(), reset_proto.reset_data_path())
+              .use_system_verilog(UseSystemVerilog())
+              .emit_as_pipeline(false)));
+
+  ModuleSimulator simulator =
+      NewModuleSimulator(result.verilog_text, result.signature, {});
+  absl::flat_hash_map<std::string, std::vector<Bits>> inputs = {
+      {"top_in", {UBits(3, 32), UBits(10, 32), UBits(42, 32)}}};
+  absl::flat_hash_map<std::string, std::vector<Bits>> outputs = {
+      {"top_out", {UBits(3, 32), UBits(13, 32), UBits(55, 32)}}};
+  EXPECT_THAT(simulator.RunInputSeriesProc(inputs, {{"top_out", 3}}),
+              absl_testing::IsOkAndHolds(outputs));
+}
+
 INSTANTIATE_TEST_SUITE_P(PipelineGeneratorTestInstantiation,
                          SynchronousProcsTest,
                          testing::ValuesIn(kDefaultSimulationTargets),

xls/scheduling/schedule_graph.cc

@@ -308,10 +308,13 @@ absl::StatusOr<ScheduleGraph> ScheduleGraph::CreateSynchronousGraph(
     }
 
     for (Next* next : proc->next_values()) {
-      backedges.push_back(
-          ScheduleBackedge{.source = next,
-                           .destination = next->state_read(),
-                           .distance = LessThanInitiationInterval()});
+      backedges.push_back(ScheduleBackedge{
+          .source = next,
+          .destination =
+              next->has_state_read()
+                  ? next->state_read()
+                  : proc->GetStateReadByStateElement(next->state_element()),
+          .distance = LessThanInitiationInterval()});
     }
   }