[Explicit State Access] Guard state read references from next in ir_minimizer_main

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

@@ -761,14 +761,20 @@ absl::Status CloneNodesIntoBlockHandler::HandleNextValue(Node* node,
   Proc* proc = function_base_->AsProcOrDie();
   Next* next = node->As<Next>();
   StateElement* state_element =
-      next->state_read()->As<StateRead>()->state_element();
+      next->has_state_read()
+          ? next->state_read()->As<StateRead>()->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"
@@ -259,6 +260,64 @@ 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);
+  XLS_ASSERT_OK_AND_ASSIGN(ReceiveChannelInterface * in,
+                           pb.AddInputChannel("top_in", u32));
+  XLS_ASSERT_OK_AND_ASSIGN(SendChannelInterface * out,
+                           pb.AddOutputChannel("top_out", u32));
+
+  XLS_ASSERT_OK_AND_ASSIGN(StateElement * 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/contrib/mlir/tools/xls_translate/xls_translate_to_mlir.cc

@@ -1600,10 +1600,12 @@ absl::StatusOr<Operation*> translateProc(::xls::Proc& xls_proc,
 
     if (n->Is<::xls::Next>()) {
       auto xls_next = n->As<::xls::Next>();
-      auto state_elem_name = xls_next->state_read()
-                                 ->As<::xls::StateRead>()
-                                 ->state_element()
-                                 ->name();
+      auto state_elem_name = xls_next->has_state_read()
+                                 ? xls_next->state_read()
+                                       ->As<::xls::StateRead>()
+                                       ->state_element()
+                                       ->name()
+                                 : xls_next->state_element()->name();
       next_value_ops[state_elem_name].push_back(xls_next);
       continue;
     }

xls/dev_tools/dev_passes/decompose_dataflow_pass.cc

@@ -417,7 +417,13 @@ class DecomposeDataflowVisitor final : public DataflowVisitor<Node*> {
     std::vector<Node*> new_args;
     new_args.reserve(node->operands().size());
     for (Node* arg : node->operands()) {
-      if (node->Is<Next>() && arg == node->As<Next>()->state_read()) {
+      if (node->Is<Next>() &&
+          arg == (node->As<Next>()->has_state_read()
+                      ? node->As<Next>()->state_read()
+                      : node->function_base()
+                            ->AsProcOrDie()
+                            ->GetStateReadByStateElement(
+                                node->As<Next>()->state_element()))) {
         // Don't decompose state reads. Leave pass-throughs alone and never
         // touch the original read.
         new_args.push_back(arg);

xls/dev_tools/dev_passes/decompose_dataflow_pass_test.cc

@@ -118,6 +118,21 @@ TEST_F(DecomposeDataflowPassTest, ProcTest) {
                         m::ArrayIndex(m::Literal(), {m::Literal(3)})));
 }
 
+TEST_F(DecomposeDataflowPassTest, DecoupledProcPassThroughTest) {
+  auto p = CreatePackage();
+  TokenlessProcBuilder pb(NewStyleProc(), "p", "tkn", p.get());
+  XLS_ASSERT_OK_AND_ASSIGN(StateElement * state_element,
+                           pb.UnreadStateElement("st", Value(UBits(42, 32)),
+                                                 /*non_synthesizable=*/false));
+  BValue st = pb.StateRead(state_element);
+  pb.Next(state_element, st);
+
+  XLS_ASSERT_OK_AND_ASSIGN(Proc * proc, pb.Build());
+  ASSERT_TRUE(proc->uses_decoupled_next());
+
+  EXPECT_THAT(Run(p.get()), IsOkAndHolds(false));
+}
+
 TEST_F(DecomposeDataflowPassTest, OneHotSelect) {
   auto p = CreatePackage();
   FunctionBuilder fb("f", p.get());

xls/dev_tools/extract_state_element.cc

@@ -17,6 +17,7 @@
 #include <memory>
 #include <optional>
 #include <string>
+#include <variant>
 #include <vector>
 
 #include "absl/algorithm/container.h"
@@ -123,13 +124,16 @@ absl::Status ExtractSegmentInto(ProcBuilder& pb, Proc* original,
         XLS_ASSIGN_OR_RETURN(chan, new_pkg->GetChannel(r->channel_name()));
       } else {
         XLS_ASSIGN_OR_RETURN(auto orig_chan_ref, r->GetChannelRef());
-        Channel* orig_chan = std::get<Channel*>(orig_chan_ref);
-        XLS_ASSIGN_OR_RETURN(
-            Type * map_ty, new_pkg->MapTypeFromOtherPackage(orig_chan->type()));
+        XLS_ASSIGN_OR_RETURN(Type * map_ty, new_pkg->MapTypeFromOtherPackage(
+                                                ChannelRefType(orig_chan_ref)));
+        absl::Span<const Value> initial_values = {};
+        if (std::holds_alternative<Channel*>(orig_chan_ref)) {
+          initial_values = std::get<Channel*>(orig_chan_ref)->initial_values();
+        }
         XLS_ASSIGN_OR_RETURN(
-            chan, new_pkg->CreateStreamingChannel(
-                      r->channel_name(), ChannelOps::kReceiveOnly, map_ty,
-                      orig_chan->initial_values()));
+            chan, new_pkg->CreateStreamingChannel(r->channel_name(),
+                                                  ChannelOps::kReceiveOnly,
+                                                  map_ty, initial_values));
         if (r->is_blocking()) {
           if (r->predicate()) {
             old_to_new[n] =
@@ -188,22 +192,27 @@ absl::Status ExtractSegmentInto(ProcBuilder& pb, Proc* original,
       }
     } else if (n->Is<Next>()) {
       Next* nxt = n->As<Next>();
-      if (absl::c_contains(
-              state_elements,
-              nxt->state_read()->As<StateRead>()->state_element())) {
-        if (nxt->predicate()) {
+      if (absl::c_contains(state_elements, nxt->state_element())) {
+        std::optional<BValue> pred =
+            nxt->predicate().has_value()
+                ? std::make_optional(
+                      BValue(old_to_new.at(*nxt->predicate()), &pb))
+                : std::nullopt;
+        std::string name = nxt->HasAssignedName() ? nxt->GetName() : "";
+        if (nxt->has_state_read()) {
+          old_to_new[n] = pb.Next(BValue(old_to_new.at(nxt->state_read()), &pb),
+                                  BValue(old_to_new.at(nxt->value()), &pb),
+                                  pred, nxt->label(), nxt->loc(), name)
+                              .node();
+        } else {
+          StateRead* state_read =
+              original->GetStateReadByStateElement(nxt->state_element());
+          StateElement* new_se =
+              old_to_new.at(state_read)->As<StateRead>()->state_element();
           old_to_new[n] =
-              pb.Next(BValue(old_to_new[nxt->state_read()], &pb),
-                      BValue(old_to_new[nxt->value()], &pb),
-                      BValue(old_to_new[*nxt->predicate()], &pb), nxt->label(),
-                      nxt->loc(), nxt->HasAssignedName() ? nxt->GetName() : "")
+              pb.Next(new_se, BValue(old_to_new.at(nxt->value()), &pb), pred,
+                      nxt->label(), nxt->loc(), name)
                   .node();
-        } else {
-          old_to_new[n] = pb.Next(BValue(old_to_new[nxt->state_read()], &pb),
-                                  BValue(old_to_new[nxt->value()], &pb),
-                                  std::nullopt, nxt->label(), nxt->loc(),
-                                  nxt->HasAssignedName() ? nxt->GetName() : "")
-                              .node();
         }
       }
       // Non-extracted nexts can be dropped.

xls/dev_tools/extract_state_element_test.cc

@@ -96,5 +96,40 @@ TEST_F(ExtractStateElementTest, SendState) {
   ExpectEqualToGoldenFile(TestFilePath(TestName()), new_pkg->DumpIr());
 }
 
+TEST_F(ExtractStateElementTest, SendStateDecoupledNext) {
+  auto p = CreatePackage();
+  TokenlessProcBuilder pb(NewStyleProc{}, TestName(), "tkn", p.get());
+  XLS_ASSERT_OK_AND_ASSIGN(ReceiveChannelInterface * chan,
+                           pb.AddInputChannel("inp_chan", p->GetBitsType(32)));
+  XLS_ASSERT_OK_AND_ASSIGN(StateElement * state_element_a,
+                           pb.UnreadStateElement("a", Value(UBits(1, 32)),
+                                                 /*non_synthesizable=*/false));
+  XLS_ASSERT_OK_AND_ASSIGN(StateElement * state_element_b,
+                           pb.UnreadStateElement("b", Value(UBits(1, 32)),
+                                                 /*non_synthesizable=*/false));
+  XLS_ASSERT_OK_AND_ASSIGN(StateElement * state_element_c,
+                           pb.UnreadStateElement("c", Value(UBits(1, 32)),
+                                                 /*non_synthesizable=*/false));
+  XLS_ASSERT_OK_AND_ASSIGN(StateElement * state_element_d,
+                           pb.UnreadStateElement("d", Value(UBits(1, 32)),
+                                                 /*non_synthesizable=*/false));
+  BValue a = pb.StateRead(state_element_a);
+  BValue b = pb.StateRead(state_element_b);
+  BValue c = pb.StateRead(state_element_c);
+  BValue d = pb.StateRead(state_element_d);
+  pb.Next(state_element_a, pb.Add(a, b));
+  pb.Next(state_element_b, pb.Add(c, a));
+  pb.Next(state_element_c, pb.Add(a, d));
+  pb.Next(state_element_d,
+          pb.Add(pb.TupleIndex(pb.Receive(chan, pb.Literal(Value::Token())), 1),
+                 c));
+  XLS_ASSERT_OK_AND_ASSIGN(Proc * proc, pb.Build());
+  XLS_ASSERT_OK_AND_ASSIGN(
+      auto new_pkg,
+      ExtractStateElementsInNewPackage(proc, {state_element_a, state_element_b},
+                                       /*send_state_values=*/true));
+  RecordProperty("ir", new_pkg->DumpIr());
+  ExpectEqualToGoldenFile(TestFilePath(TestName()), new_pkg->DumpIr());
+}
 }  // namespace
 }  // namespace xls

xls/dev_tools/ir_minimizer_main.cc

@@ -134,6 +134,9 @@ ABSL_FLAG(bool, can_remove_params, false,
           "Whether parameters can be removed during the minimization process. "
           "If the test executable interprets the IR using a fixed set of "
           "arguments, parameters should not be removed.");
+ABSL_FLAG(
+    bool, can_remove_state_elements, true,
+    "Whether state elements can be removed during the minimization process.");
 ABSL_FLAG(int64_t, failed_attempt_limit, 256,
           "Failed simplification attempts (in a row) before we conclude we're "
           "done reducing.");
@@ -444,10 +447,15 @@ absl::StatusOr<bool> RemoveDeadParameters(FunctionBase* f) {
     absl::flat_hash_set<StateElement*> invariant_state_elements(
         p->StateElements().begin(), p->StateElements().end());
     for (Next* next : p->next_values()) {
-      if (next->value() != next->state_read()) {
+      if (next->value() !=
+          (next->has_state_read()
+               ? next->state_read()
+               : p->GetStateReadByStateElement(next->state_element()))) {
         // This state param is not actually invariant.
         invariant_state_elements.erase(
-            next->state_read()->As<StateRead>()->state_element());
+            next->has_state_read()
+                ? next->state_read()->As<StateRead>()->state_element()
+                : next->state_element());
       }
     }
     for (StateElement* invariant : invariant_state_elements) {
@@ -763,7 +771,16 @@ absl::StatusOr<SimplificationResult> SimplifyReturnValue(
 template <typename NodeT, typename... Args>
 absl::StatusOr<Node*> SafeReplaceUsesWithNew(Node* target, Args... v) {
   auto is_state_read_of_next = [&](Node* n) {
-    return n->Is<Next>() && n->As<Next>()->state_read() == target;
+    if (!n->Is<Next>()) {
+      return false;
+    }
+    Next* next = n->As<Next>();
+    if (next->has_state_read()) {
+      return next->state_read() == target;
+    } else {
+      return n->function_base()->AsProcOrDie()->GetStateReadByStateElement(
+                 next->state_element()) == target;
+    }
   };
   if (!target->Is<StateRead>() ||
       !absl::c_any_of(target->users(), is_state_read_of_next)) {
@@ -802,9 +819,11 @@ absl::StatusOr<SimplificationResult> RunRandomPass(
     // Only can inline from the top.
     passes.push_back(std::make_unique<InliningPass>());
   }
-  passes.push_back(std::make_unique<ProcStateArrayFlatteningPass>());
-  passes.push_back(std::make_unique<ProcStateTupleFlatteningPass>());
-  passes.push_back(std::make_unique<ProcStateOptimizationPass>());
+  if (absl::GetFlag(FLAGS_can_remove_state_elements)) {
+    passes.push_back(std::make_unique<ProcStateArrayFlatteningPass>());
+    passes.push_back(std::make_unique<ProcStateTupleFlatteningPass>());
+    passes.push_back(std::make_unique<ProcStateOptimizationPass>());
+  }
 
   int64_t pass_no = absl::Uniform<int64_t>(rng, 0, passes.size());
   PassResults results;