[ir] Extend controlled Stages to track two more signals

We missed some signals that are critical in actual implementation:

1. The external signal `outputs_ready`, which triggers the Stage's consuming I/O (`receive` & `state_read`), and
2. The internal signal `active_inputs_valid`, which (when combined with `inputs_valid`) triggers the Stage's committing I/O (`send` & `next_value`).

PiperOrigin-RevId: 836704396

Author: ericastor

xls/ir/block.cc

@@ -351,16 +351,14 @@ std::string Block::DumpIr() const {
 
       const Stage& stage = stages_[stage_idx];
       CHECK(stage.IsControlled());
-      absl::StrAppendFormat(
-          &res, "  controlled_stage(%s%s) {\n", stage.inputs_valid()->GetName(),
-          stage.contains(stage.outputs_valid())
-              ? ""
-              : absl::StrCat(", outputs_valid=",
-                             stage.outputs_valid()->GetName()));
+      absl::StrAppendFormat(&res, "  controlled_stage(%s, %s) {\n",
+                            stage.inputs_valid()->GetName(),
+                            stage.outputs_ready()->GetName());
       for (Node* node : staged_nodes[stage_idx]) {
-        absl::StrAppendFormat(&res, "    %s%s\n",
-                              node == stage.outputs_valid() ? "ret " : "",
-                              node->ToString());
+        absl::StrAppend(
+            &res, "    ", (node == stage.outputs_valid() ? "ret " : ""),
+            (node == stage.active_inputs_valid() ? "active_inputs_valid " : ""),
+            node->ToString(), "\n");
       }
       absl::StrAppend(&res, "  }\n");
     }
@@ -1134,6 +1132,12 @@ absl::StatusOr<bool> Block::RemoveNodeFromStage(Node* node) {
     return false;
   }
   int64_t stage_index = it->second;
+  Stage& stage = stages_[stage_index];
+  if (stage.active_inputs_valid() == node || stage.outputs_valid() == node) {
+    return absl::InvalidArgumentError(absl::StrFormat(
+        "Node %s has implicit uses in stage %d and cannot be removed.",
+        node->GetName(), stage_index));
+  }
   node_to_stage_.erase(it);
   stages_[stage_index].erase(node);
   return true;

xls/ir/block.h

@@ -340,7 +340,9 @@ class Block : public FunctionBase {
 
   bool HasImplicitUse(Node* node) const override {
     return absl::c_any_of(stages_, [node](const Stage& stage) {
-      return stage.inputs_valid() == node || stage.outputs_valid() == node;
+      return stage.inputs_valid() == node || stage.outputs_valid() == node ||
+             stage.active_inputs_valid() == node ||
+             stage.outputs_ready() == node;
     });
   }
 

xls/ir/block_test.cc

@@ -18,6 +18,7 @@
 #include <optional>
 #include <string>
 #include <string_view>
+#include <utility>
 #include <vector>
 
 #include "gmock/gmock.h"
@@ -1006,18 +1007,22 @@ class ScheduledBlockTest : public IrTestBase {
   struct TestBlock {
     ScheduledBlock* block;
     Node* iv0;
+    Node* or0;
+    Node* aiv0;
     Node* ov0;
   };
   absl::StatusOr<TestBlock> CreateScheduledBlock(Package* p) {
     ScheduledBlockBuilder bb("b", p);
     BValue iv0 = bb.Literal(UBits(1, 1), SourceInfo(), "iv0");
-    bb.StartStage(iv0);
+    BValue or0 = bb.Literal(UBits(1, 1), SourceInfo(), "or0");
+    bb.StartStage(iv0, or0);
     BValue x = bb.InputPort("x", p->GetBitsType(32));
     bb.OutputPort("out", x);
+    BValue aiv0 = bb.Literal(UBits(1, 1), SourceInfo(), "aiv0");
     BValue ov0 = bb.Literal(UBits(1, 1), SourceInfo(), "ov0");
-    bb.EndStage(ov0);
+    bb.EndStage(aiv0, ov0);
     XLS_ASSIGN_OR_RETURN(ScheduledBlock * block, bb.Build());
-    return TestBlock{block, iv0.node(), ov0.node()};
+    return TestBlock{block, iv0.node(), or0.node(), aiv0.node(), ov0.node()};
   }
 
   void ExpectIr(std::string_view got, std::string_view test_name) {
@@ -1044,23 +1049,38 @@ TEST_F(ScheduledBlockTest, StageAddAndClear) {
 
   XLS_ASSERT_OK_AND_ASSIGN(
       Node * iv1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Node * or1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Node * aiv1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
   XLS_ASSERT_OK_AND_ASSIGN(
       Node * ov1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
-  block->AddStage(Stage(iv1, ov1));
+  Stage stage1(iv1, or1, aiv1, ov1);
+  stage1.AddNode(aiv1);
+  stage1.AddNode(ov1);
+  block->AddStage(std::move(stage1));
+
   EXPECT_EQ(block->stages().size(), 2);
 
   XLS_ASSERT_OK_AND_ASSIGN(
       Node * iv2, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Node * or2, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Node * aiv2, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
   XLS_ASSERT_OK_AND_ASSIGN(
       Node * ov2, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
-  block->AddStage(Stage(iv2, ov2));
+  Stage stage2(iv2, or2, aiv2, ov2);
+  stage2.AddNode(aiv2);
+  stage2.AddNode(ov2);
+  block->AddStage(std::move(stage2));
   EXPECT_EQ(block->stages().size(), 3);
 
   block->ClearStages();
   EXPECT_TRUE(block->stages().empty());
 
   // Re-stage nodes to satisfy the verifier on destruction.
-  block->AddStage(Stage(tb.iv0, tb.ov0));
+  block->AddStage(Stage(tb.iv0, tb.or0, tb.aiv0, tb.ov0));
   for (Node* node : block->nodes()) {
     if (!block->IsStaged(node)) {
       XLS_ASSERT_OK(block->AddNodeToStage(0, node).status());
@@ -1075,9 +1095,16 @@ TEST_F(ScheduledBlockTest, GetStageIndex) {
 
   XLS_ASSERT_OK_AND_ASSIGN(
       Node * iv1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Node * or1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Node * aiv1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
   XLS_ASSERT_OK_AND_ASSIGN(
       Node * ov1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
-  block->AddStage(Stage(iv1, ov1));  // Stage 1
+  Stage stage1(iv1, or1, aiv1, ov1);
+  stage1.AddNode(aiv1);
+  stage1.AddNode(ov1);
+  block->AddStage(std::move(stage1));
 
   XLS_ASSERT_OK_AND_ASSIGN(
       Node * lit1,
@@ -1094,9 +1121,16 @@ TEST_F(ScheduledBlockTest, AddNodeToStage) {
   ScheduledBlock* block = tb.block;
   XLS_ASSERT_OK_AND_ASSIGN(
       Node * iv1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Node * or1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Node * aiv1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
   XLS_ASSERT_OK_AND_ASSIGN(
       Node * ov1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
-  block->AddStage(Stage(iv1, ov1));  // Stage 1
+  Stage stage1(iv1, or1, aiv1, ov1);
+  stage1.AddNode(aiv1);
+  stage1.AddNode(ov1);
+  block->AddStage(std::move(stage1));
 
   XLS_ASSERT_OK_AND_ASSIGN(
       Node * lit, block->MakeNode<Literal>(SourceInfo(), Value(UBits(10, 32))));
@@ -1114,9 +1148,16 @@ TEST_F(ScheduledBlockTest, MakeNodeInStage) {
   ScheduledBlock* block = tb.block;
   XLS_ASSERT_OK_AND_ASSIGN(
       Node * iv1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Node * or1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Node * aiv1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
   XLS_ASSERT_OK_AND_ASSIGN(
       Node * ov1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
-  block->AddStage(Stage(iv1, ov1));  // Stage 1
+  Stage stage1(iv1, or1, aiv1, ov1);
+  stage1.AddNode(aiv1);
+  stage1.AddNode(ov1);
+  block->AddStage(std::move(stage1));
 
   XLS_ASSERT_OK_AND_ASSIGN(
       Literal * literal,
@@ -1130,9 +1171,16 @@ TEST_F(ScheduledBlockTest, MakeNodeWithNameInStage) {
   ScheduledBlock* block = tb.block;
   XLS_ASSERT_OK_AND_ASSIGN(
       Node * iv1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Node * or1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Node * aiv1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
   XLS_ASSERT_OK_AND_ASSIGN(
       Node * ov1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
-  block->AddStage(Stage(iv1, ov1));  // Stage 1
+  Stage stage1(iv1, or1, aiv1, ov1);
+  stage1.AddNode(aiv1);
+  stage1.AddNode(ov1);
+  block->AddStage(std::move(stage1));
 
   XLS_ASSERT_OK_AND_ASSIGN(
       Literal * literal, block->MakeNodeWithNameInStage<Literal>(
@@ -1154,14 +1202,28 @@ TEST_F(ScheduledBlockTest, CloneScheduledBlock) {
   ScheduledBlock* block = tb.block;
   XLS_ASSERT_OK_AND_ASSIGN(
       Node * iv1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Node * or1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Node * aiv1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
   XLS_ASSERT_OK_AND_ASSIGN(
       Node * ov1, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
-  block->AddStage(Stage(iv1, ov1));
+  Stage stage1(iv1, or1, aiv1, ov1);
+  stage1.AddNode(aiv1);
+  stage1.AddNode(ov1);
+  block->AddStage(std::move(stage1));
   XLS_ASSERT_OK_AND_ASSIGN(
       Node * iv2, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Node * or2, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Node * aiv2, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
   XLS_ASSERT_OK_AND_ASSIGN(
       Node * ov2, block->MakeNode<Literal>(SourceInfo(), Value(UBits(1, 1))));
-  block->AddStage(Stage(iv2, ov2));
+  Stage stage2(iv2, or2, aiv2, ov2);
+  stage1.AddNode(aiv2);
+  stage1.AddNode(ov2);
+  block->AddStage(std::move(stage2));
 
   XLS_ASSERT_OK(block->MakeNodeWithNameInStage<Literal>(
       0, SourceInfo(), Value(UBits(1, 32)), "my_x"));

xls/ir/function_base.cc

@@ -94,6 +94,13 @@ absl::StatusOr<Stage> Stage::Clone(
   if (inputs_valid_ != nullptr) {
     XLS_ASSIGN_OR_RETURN(cloned_stage.inputs_valid_, map_node(inputs_valid_));
   }
+  if (outputs_ready_ != nullptr) {
+    XLS_ASSIGN_OR_RETURN(cloned_stage.outputs_ready_, map_node(outputs_ready_));
+  }
+  if (active_inputs_valid_ != nullptr) {
+    XLS_ASSIGN_OR_RETURN(cloned_stage.active_inputs_valid_,
+                         map_node(active_inputs_valid_));
+  }
   if (outputs_valid_ != nullptr) {
     XLS_ASSIGN_OR_RETURN(cloned_stage.outputs_valid_, map_node(outputs_valid_));
   }

xls/ir/function_base.h

@@ -52,11 +52,16 @@ class Proc;
 // Represents a pipeline stage after scheduling.
 class Stage {
  public:
-  Stage(Node* inputs_valid, Node* outputs_valid)
-      : inputs_valid_(inputs_valid), outputs_valid_(outputs_valid) {
+  Stage(Node* inputs_valid, Node* outputs_ready, Node* active_inputs_valid,
+        Node* outputs_valid)
+      : inputs_valid_(inputs_valid),
+        outputs_ready_(outputs_ready),
+        active_inputs_valid_(active_inputs_valid),
+        outputs_valid_(outputs_valid) {
+    CHECK_EQ(inputs_valid_ == nullptr, outputs_ready_ == nullptr);
     CHECK_EQ(inputs_valid_ == nullptr, outputs_valid_ == nullptr);
   }
-  Stage() : Stage(nullptr, nullptr) {}
+  Stage() : Stage(nullptr, nullptr, nullptr, nullptr) {}
 
   Stage(const Stage& other) = default;
   Stage& operator=(const Stage& other) = default;
@@ -94,19 +99,37 @@ class Stage {
       const absl::flat_hash_map<Node*, Node*>& node_mapping) const;
 
   bool IsControlled() const {
-    return inputs_valid_ != nullptr && outputs_valid_ != nullptr;
+    return inputs_valid_ != nullptr && outputs_ready_ != nullptr &&
+           active_inputs_valid_ != nullptr && outputs_valid_ != nullptr;
   }
 
   // Returns the node that signals whether it would be valid for this stage to
-  // execute; i.e., that all passive inputs are updated for the next activation,
-  // and all active inputs are valid.
+  // execute; i.e., that all passive inputs are updated for the next activation.
   Node* inputs_valid() const { return inputs_valid_; }
 
   void set_inputs_valid(Node* inputs_valid) { inputs_valid_ = inputs_valid; }
 
-  // Returns the node that signals whether the stage's outputs are valid to be
-  // read; i.e., that all logic nodes are updated by the current activation, and
-  // all active outputs have completed execution.
+  // Returns the node that signals whether it is safe for this stage to execute;
+  // i.e., that the receiver for all passive outputs will have space to store
+  // the data.
+  Node* outputs_ready() const { return outputs_ready_; }
+
+  void set_outputs_ready(Node* outputs_ready) {
+    outputs_ready_ = outputs_ready;
+  }
+
+  // Returns the node that signals whether all active inputs to this stage are
+  // valid; i.e., that all active receives & all actively-read state values have
+  // the correct values for this activation.
+  Node* active_inputs_valid() const { return active_inputs_valid_; }
+
+  void set_active_inputs_valid(Node* active_inputs_valid) {
+    active_inputs_valid_ = active_inputs_valid;
+  }
+
+  // Returns the node that signals whether it would be safe for this stage to be
+  // done executing; i.e., that all logic nodes are updated by the current
+  // activation, and all active outputs have completed execution.
   Node* outputs_valid() const { return outputs_valid_; }
 
   void set_outputs_valid(Node* outputs_valid) {
@@ -118,6 +141,8 @@ class Stage {
   absl::btree_set<Node*, Node::NodeIdLessThan> logic_;
   absl::btree_set<Node*, Node::NodeIdLessThan> active_outputs_;
   Node* inputs_valid_ = nullptr;
+  Node* outputs_ready_ = nullptr;
+  Node* active_inputs_valid_ = nullptr;
   Node* outputs_valid_ = nullptr;
 };
 

xls/ir/ir_parser.cc

@@ -3397,33 +3397,57 @@ absl::Status Parser::ParseControlledStage(
     absl::flat_hash_map<std::string, BValue>* name_to_value, Package* package) {
   XLS_RETURN_IF_ERROR(scanner_.DropKeywordOrError("controlled_stage"));
   XLS_RETURN_IF_ERROR(scanner_.DropTokenOrError(LexicalTokenType::kParenOpen));
-  XLS_ASSIGN_OR_RETURN(BValue iv, ParseAndResolveIdentifier(*name_to_value));
+  XLS_ASSIGN_OR_RETURN(BValue inputs_valid,
+                       ParseAndResolveIdentifier(*name_to_value));
+  XLS_RETURN_IF_ERROR(scanner_.DropTokenOrError(LexicalTokenType::kComma));
+  XLS_ASSIGN_OR_RETURN(BValue outputs_ready,
+                       ParseAndResolveIdentifier(*name_to_value));
   XLS_RETURN_IF_ERROR(scanner_.DropTokenOrError(LexicalTokenType::kParenClose));
   XLS_RETURN_IF_ERROR(scanner_.DropTokenOrError(LexicalTokenType::kCurlOpen));
 
-  builder->StartStage(iv);
+  builder->StartStage(inputs_valid, outputs_ready);
 
-  std::optional<BValue> ov;
+  std::optional<BValue> active_inputs_valid;
+  std::optional<BValue> outputs_valid;
   while (!scanner_.PeekTokenIs(LexicalTokenType::kCurlClose)) {
+    bool is_active_inputs_valid =
+        scanner_.TryDropKeyword("active_inputs_valid");
     bool is_ret = scanner_.TryDropKeyword("ret");
+    if (is_ret && !is_active_inputs_valid) {
+      is_active_inputs_valid = scanner_.TryDropKeyword("active_inputs_valid");
+    }
     XLS_ASSIGN_OR_RETURN(BValue result, ParseNode(builder, name_to_value));
     (*name_to_value)[result.node()->GetName()] = result;
+    if (is_active_inputs_valid) {
+      if (active_inputs_valid.has_value()) {
+        return absl::InvalidArgumentError(absl::StrFormat(
+            "More than one active_inputs_valid found in controlled_stage @ %s",
+            result.node()->loc().ToString()));
+      }
+      active_inputs_valid = result;
+    }
     if (is_ret) {
-      if (ov.has_value()) {
+      if (outputs_valid.has_value()) {
         return absl::InvalidArgumentError(
             absl::StrFormat("More than one ret found in controlled_stage @ %s",
                             result.node()->loc().ToString()));
       }
-      ov = result;
+      outputs_valid = result;
     }
   }
 
-  if (!ov.has_value()) {
+  if (!active_inputs_valid.has_value()) {
+    return absl::InvalidArgumentError(absl::StrFormat(
+        "No 'active_inputs_valid' node found in controlled_stage @ %s",
+        scanner_.PeekTokenOrDie().pos().ToHumanString()));
+  }
+  if (!outputs_valid.has_value()) {
     return absl::InvalidArgumentError(
-        "No 'ret' node found in controlled_stage");
+        absl::StrFormat("No 'ret' node found in controlled_stage @ %s",
+                        scanner_.PeekTokenOrDie().pos().ToHumanString()));
   }
 
-  builder->EndStage(*ov);
+  builder->EndStage(*active_inputs_valid, *outputs_valid);
 
   XLS_RETURN_IF_ERROR(scanner_.DropTokenOrError(LexicalTokenType::kCurlClose));
   return absl::OkStatus();