[Explicit State Access] Guard direct references to state read and support decoupled next in decompose_dataflow_pass

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/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/testdata/extract_state_element_test_SendStateDecoupledNext.ir

@@ -0,0 +1,19 @@
+package SendStateDecoupledNext
+
+chan a_value_chan(bits[32], id=1, kind=streaming, ops=send_only, flow_control=ready_valid, strictness=proven_mutually_exclusive)
+chan b_value_chan(bits[32], id=2, kind=streaming, ops=send_only, flow_control=ready_valid, strictness=proven_mutually_exclusive)
+chan c_chan(bits[32], id=3, kind=streaming, ops=receive_only, flow_control=ready_valid, strictness=proven_mutually_exclusive)
+
+top proc SendStateDecoupledNext(a: bits[32], b: bits[32], init={1, 1}) {
+  literal.1: token = literal(value=token, id=1)
+  receive.10: (token, bits[32]) = receive(literal.1, channel=c_chan, id=10)
+  a: bits[32] = state_read(state_element=a, id=3)
+  b: bits[32] = state_read(state_element=b, id=6)
+  c: bits[32] = tuple_index(receive.10, index=1, id=11)
+  add.16: bits[32] = add(a, b, id=16)
+  add.17: bits[32] = add(c, a, id=17)
+  a__1: token = send(literal.1, a, channel=a_value_chan, id=5)
+  b__1: token = send(literal.1, b, channel=b_value_chan, id=8)
+  next_value.21: () = next_value(state_element=a, value=add.16, id=21)
+  next_value.22: () = next_value(state_element=b, value=add.17, id=22)
+}

xls/dslx/ir_convert/testdata/ir_converter_test_ExplicitStateAccessArray.ir

@@ -28,7 +28,7 @@ top proc __test_module__Array_0_next<>(__state: bits[32][4], init={[0, 1, 2, 3]}
   tuple.3: () = tuple(id=3)
   tuple_index.20: token = tuple_index(counted_for.19, index=0, id=20)
   assert.24: token = assert(__token, or.23, message="State element written before read in same activation.", id=24)
-  next_value.25: () = next_value(param=__state, value=new_array, predicate=literal.2, id=25)
+  next_value.25: () = next_value(state_element=__state, value=new_array, predicate=literal.2, id=25)
   tuple.26: () = tuple(id=26)
   tuple.27: () = tuple(id=27)
 }

xls/dslx/ir_convert/testdata/ir_converter_test_ExplicitStateAccessBool.ir

@@ -12,7 +12,7 @@ top proc __test_module__Bool_0_next<>(__state: bits[1], init={0}) {
   tuple.3: () = tuple(id=3)
   false_val: bits[1] = identity(__state, id=5)
   assert.9: token = assert(__token, or.8, message="State element written before read in same activation.", id=9)
-  next_value.10: () = next_value(param=__state, value=literal.6, predicate=literal.2, id=10)
+  next_value.10: () = next_value(state_element=__state, value=literal.6, predicate=literal.2, id=10)
   tuple.11: () = tuple(id=11)
   tuple.12: () = tuple(id=12)
 }

xls/dslx/ir_convert/testdata/ir_converter_test_ExplicitStateAccessConditionalWrite.ir

@@ -39,10 +39,10 @@ top proc __test_module__main_0_next<>(__state: (bits[32], bits[1]), init={(0, 0)
   tuple.3: () = tuple(id=3)
   literal.6: bits[32] = literal(value=0, id=6)
   literal.12: bits[32] = literal(value=0, id=12)
-  next_value.22: () = next_value(param=__state, value=tuple.17, predicate=and.18, id=22)
+  next_value.22: () = next_value(state_element=__state, value=tuple.17, predicate=and.18, id=22)
   literal.24: bits[32] = literal(value=0, id=24)
   assert.38: token = assert(assert.34, or.37, message="State element written after write in same activation.", id=38)
   or.39: bits[1] = or(and.18, and.31, id=39)
-  next_value.40: () = next_value(param=__state, value=tuple.29, predicate=and.31, id=40)
+  next_value.40: () = next_value(state_element=__state, value=tuple.29, predicate=and.31, id=40)
   sel.42: () = sel(eq.11, cases=[tuple.41, tuple.23], id=42)
 }

xls/dslx/ir_convert/testdata/ir_converter_test_ExplicitStateAccessLabeledReadAndWrite.ir

@@ -13,7 +13,7 @@ top proc __test_module__main_0_next<>(__state: bits[32], init={0}) {
   y: bits[32] = add(x, literal.6, id=7)
   tuple.3: () = tuple(id=3)
   assert.10: token = assert(__token, or.9, message="State element written before read in same activation.", id=10)
-  next_value.11: () = next_value(param=__state, value=y, predicate=literal.2, label="main_write", id=11)
+  next_value.11: () = next_value(state_element=__state, value=y, predicate=literal.2, label="main_write", id=11)
   tuple.12: () = tuple(id=12)
   tuple.13: () = tuple(id=13)
 }

xls/dslx/ir_convert/testdata/ir_converter_test_ExplicitStateAccessMatch.ir

@@ -45,14 +45,14 @@ top proc __test_module__main_0_next<>(__state: bits[32], init={0}) {
   tuple.36: () = tuple(id=36)
   tuple.50: () = tuple(id=50)
   tuple.3: () = tuple(id=3)
-  next_value.17: () = next_value(param=__state, value=add.11, predicate=and.13, id=17)
+  next_value.17: () = next_value(state_element=__state, value=add.11, predicate=and.13, id=17)
   tuple.18: () = tuple(id=18)
-  next_value.34: () = next_value(param=__state, value=umul.23, predicate=and.25, id=34)
+  next_value.34: () = next_value(state_element=__state, value=umul.23, predicate=and.25, id=34)
   tuple.35: () = tuple(id=35)
   literal.37: bits[1] = literal(value=1, id=37)
   assert.46: token = assert(assert.42, or.45, message="State element written after write in same activation.", id=46)
   or.47: bits[1] = or(or.33, and.39, id=47)
-  next_value.48: () = next_value(param=__state, value=current, predicate=and.39, id=48)
+  next_value.48: () = next_value(state_element=__state, value=current, predicate=and.39, id=48)
   tuple.49: () = tuple(id=49)
   priority_sel.52: () = priority_sel(concat.51, cases=[tuple.19, tuple.36], default=tuple.50, id=52)
 }

xls/dslx/ir_convert/testdata/ir_converter_test_ExplicitStateAccessMatchMultipleWrites.ir

@@ -51,16 +51,16 @@ top proc __test_module__main_0_next<>(__state: bits[1], init={1}) {
   tuple.29: () = tuple(id=29)
   tuple.58: () = tuple(id=58)
   tuple.3: () = tuple(id=3)
-  next_value.14: () = next_value(param=__state, value=literal.8, predicate=and.10, id=14)
+  next_value.14: () = next_value(state_element=__state, value=literal.8, predicate=and.10, id=14)
   tuple.15: () = tuple(id=15)
-  next_value.27: () = next_value(param=__state, value=literal.16, predicate=and.18, id=27)
+  next_value.27: () = next_value(state_element=__state, value=literal.16, predicate=and.18, id=27)
   tuple.28: () = tuple(id=28)
   eq.31: bits[1] = eq(literal.30, val, id=31)
-  next_value.43: () = next_value(param=__state, value=literal.32, predicate=and.34, id=43)
+  next_value.43: () = next_value(state_element=__state, value=literal.32, predicate=and.34, id=43)
   tuple.44: () = tuple(id=44)
   assert.54: token = assert(assert.50, or.53, message="State element written after write in same activation.", id=54)
   or.55: bits[1] = or(or.42, and.47, id=55)
-  next_value.56: () = next_value(param=__state, value=literal.45, predicate=and.47, id=56)
+  next_value.56: () = next_value(state_element=__state, value=literal.45, predicate=and.47, id=56)
   tuple.57: () = tuple(id=57)
   priority_sel.60: () = priority_sel(concat.59, cases=[tuple.29], default=tuple.58, id=60)
 }

xls/dslx/ir_convert/testdata/ir_converter_test_ExplicitStateAccessMultipleBranchingLabeledReads.ir

@@ -42,15 +42,15 @@ top proc __test_module__main_0_next<>(__val: bits[32], __switch: bits[1], init={
   or.46: bits[1] = or(not.45, literal.2, id=46)
   not.44: bits[1] = not(switch_val, id=44)
   tuple.3: () = tuple(id=3)
-  next_value.15: () = next_value(param=__val, value=add.10, predicate=and.11, label="EvenWrite", id=15)
+  next_value.15: () = next_value(state_element=__val, value=add.10, predicate=and.11, label="EvenWrite", id=15)
   tuple.16: () = tuple(id=16)
   or.24: bits[1] = or(and.6, and.19, id=24)
   or.39: bits[1] = or(and.11, and.31, id=39)
-  next_value.40: () = next_value(param=__val, value=add.29, predicate=and.31, label="OddWrite", id=40)
+  next_value.40: () = next_value(state_element=__val, value=add.29, predicate=and.31, label="OddWrite", id=40)
   tuple.41: () = tuple(id=41)
   sel.43: () = sel(switch_val, cases=[tuple.42, tuple.17], id=43)
   assert.47: token = assert(assert.38, or.46, message="State element written before read in same activation.", id=47)
-  next_value.48: () = next_value(param=__switch, value=not.44, predicate=literal.2, id=48)
+  next_value.48: () = next_value(state_element=__switch, value=not.44, predicate=literal.2, id=48)
   tuple.49: () = tuple(id=49)
   tuple.50: () = tuple(id=50)
 }

xls/dslx/ir_convert/testdata/ir_converter_test_ExplicitStateAccessMultipleStates.ir

@@ -36,12 +36,12 @@ top proc __test_module__main_0_next<>(__state_0: (bits[32], bits[32]), __state_1
   tuple.3: () = tuple(id=3)
   literal.13: bits[32] = literal(value=1, id=13)
   literal.17: bits[32] = literal(value=0, id=17)
-  next_value.27: () = next_value(param=__state_0, value=new_a, predicate=literal.2, id=27)
+  next_value.27: () = next_value(state_element=__state_0, value=new_a, predicate=literal.2, id=27)
   tuple.28: () = tuple(id=28)
-  next_value.32: () = next_value(param=__state_1, value=new_b, predicate=literal.2, id=32)
+  next_value.32: () = next_value(state_element=__state_1, value=new_b, predicate=literal.2, id=32)
   tuple.33: () = tuple(id=33)
   assert.38: token = assert(assert.31, or.37, message="State element written before read in same activation.", id=38)
-  next_value.39: () = next_value(param=__state_2, value=add.35, predicate=literal.2, id=39)
+  next_value.39: () = next_value(state_element=__state_2, value=add.35, predicate=literal.2, id=39)
   tuple.40: () = tuple(id=40)
   tuple.41: () = tuple(id=41)
 }