[codegen] Support "ignorable" gates to signal load-enable gating

xls/codegen/module_builder.cc

@@ -1247,6 +1247,13 @@ absl::StatusOr<IndexableExpression*> ModuleBuilder::EmitGate(
   XLS_ASSIGN_OR_RETURN(LogicRef * ref,
                        DeclareVariable(gate->GetName(), gate->GetType()));
 
+  if (gate->gate_type() == GateType::kIgnorableGate) {
+    // Ignorable gates don't perform zero-clamping; they can be evaluated as
+    // identity operations.
+    XLS_RETURN_IF_ERROR(Assign(ref, data, gate->GetType()));
+    return ref;
+  }
+
   // Emit the gate as an AND of the (potentially replicated) condition and the
   // data. For example:
   //

xls/codegen/verilog_conversion.cc

@@ -652,6 +652,13 @@ class BlockGenerator {
       std::optional<OpOverride> op_override =
           options_.GetOpOverride(node->op());
 
+      // Ignorable-value gates are emitted as identity, so they don't need to
+      // support overrides.
+      if (node->op() == Op::kGate &&
+          node->As<Gate>()->gate_type() == GateType::kIgnorableGate) {
+        op_override = std::nullopt;
+      }
+
       if (op_override.has_value()) {
         std::vector<NodeRepresentation> inputs;
         for (const Node* operand : node->operands()) {

xls/codegen_v_1_5/BUILD

@@ -61,6 +61,7 @@ cc_library(
     hdrs = ["block_conversion_pass.h"],
     deps = [
         "//xls/codegen:codegen_options",
+        "//xls/estimators/delay_model:delay_estimator",
         "//xls/passes:optimization_pass",
         "//xls/passes:pass_base",
         "//xls/scheduling:pipeline_schedule_cc_proto",

xls/codegen_v_1_5/block_conversion_pass.h

@@ -16,6 +16,7 @@
 #define XLS_CODEGEN_V_1_5_BLOCK_CONVERSION_PASS_H_
 
 #include "xls/codegen/codegen_options.h"
+#include "xls/estimators/delay_model/delay_estimator.h"
 #include "xls/passes/optimization_pass.h"
 #include "xls/passes/pass_base.h"
 #include "xls/scheduling/pipeline_schedule.pb.h"
@@ -28,6 +29,9 @@ namespace xls::codegen {
 struct BlockConversionPassOptions : public PassOptionsBase {
   verilog::CodegenOptions codegen_options;
   PackageScheduleProto package_schedule;
+
+  // The delay estimator used for scheduling & schedule-checking purposes.
+  const DelayEstimator* delay_estimator = nullptr;
 };
 
 struct BlockConversionContext {

xls/codegen_v_1_5/convert_to_block.cc

@@ -58,6 +58,7 @@ absl::Status ConvertToBlock(
   BlockConversionPassOptions options{
       .codegen_options = std::move(codegen_options),
       .package_schedule = std::move(schedule),
+      .delay_estimator = delay_estimator,
   };
 
   XLS_ASSIGN_OR_RETURN(std::unique_ptr<BlockConversionPass> pipeline,

xls/codegen_v_1_5/pipeline_register_insertion_pass.cc

@@ -96,6 +96,23 @@ absl::StatusOr<RegisterRead*> CreatePipelineRegister(
                        block->AddRegister(name, node->GetType(), reset_value));
 
   Node* load_enable = stage_done;
+
+  if (node->Is<Gate>() &&
+      node->As<Gate>()->gate_type() == GateType::kIgnorableGate) {
+    // Ignorable gates are used to indicate that the node is conditionally
+    // visible. We can strengthen the logic for gating the pipeline register by
+    // ANDing the stage done signal with the gate's predicate, reducing the
+    // number of times the register needs to be updated.
+    Gate* gate = node->As<Gate>();
+    Node* predicate = gate->condition();
+    node = gate->data();
+    XLS_ASSIGN_OR_RETURN(
+        load_enable,
+        block->MakeNode<NaryOp>(node->loc(),
+                                absl::MakeConstSpan({stage_done, predicate}),
+                                Op::kAnd));
+  }
+
   if (block->GetResetPort().has_value() &&
       options.codegen_options.reset().has_value() &&
       !options.codegen_options.reset()->reset_data_path()) {
@@ -109,9 +126,9 @@ absl::StatusOr<RegisterRead*> CreatePipelineRegister(
     }
     XLS_ASSIGN_OR_RETURN(
         load_enable,
-        block->MakeNode<NaryOp>(node->loc(),
-                                absl::MakeConstSpan({stage_done, reset_active}),
-                                Op::kOr));
+        block->MakeNode<NaryOp>(
+            node->loc(), absl::MakeConstSpan({load_enable, reset_active}),
+            Op::kOr));
   }
   // NOTE: The RegisterWrite is added to the block, but not to the stage. Its
   //       `load_enable` depends on `outputs_ready`, which comes from outside
@@ -139,7 +156,13 @@ absl::StatusOr<RegisterRead*> CreatePipelineRegister(
 absl::StatusOr<Node*> AddPipelineRegisterFor(
     Node* node, int64_t stage_index, Node* stage_done, ScheduledBlock* block,
     const BlockConversionPassOptions& options) {
-  std::string base_name = PipelineSignalName(node->GetName(), stage_index);
+  Node* name_source = node;
+  if (node->Is<Gate>() &&
+      node->As<Gate>()->gate_type() == GateType::kIgnorableGate) {
+    name_source = node->As<Gate>()->data();
+  }
+  std::string base_name =
+      PipelineSignalName(name_source->GetName(), stage_index);
 
   // As a special case, check if the node is a tuple
   // containing types that are of zero-width.  If so, separate them out so

xls/ir/function.cc

@@ -135,17 +135,12 @@ std::string Function::DumpIr(const IrAnnotator& annotate) const {
    private:
     Node* return_value_;
   };
-  if (IsScheduled()) {
-    absl::StrAppend(&res,
-                    DumpFunctionBaseNodes(IrAnnotatorJoiner(
-                        SkipParamsAnnotator{}, AddRetAnnotator{return_value()},
-                        IrAnnotatorRef(annotate))));
-  } else {
-    absl::StrAppend(
-        &res,
-        DumpFunctionBaseNodes(IrAnnotatorJoiner(
-            SkipParamsAnnotator{}, AddRetAnnotator{return_value()},
-            IrAnnotatorRef(annotate), TopoSortAnnotator(!is_block_source_))));
+  absl::StrAppend(
+      &res,
+      DumpFunctionBaseNodes(IrAnnotatorJoiner(
+          SkipParamsAnnotator{}, AddRetAnnotator{return_value()},
+          IrAnnotatorRef(annotate), TopoSortAnnotator(!is_block_source_))));
+  if (!IsScheduled()) {
     // Need to add the 'ret' specially if the ret is also a parameter.
     if (return_value() != nullptr && return_value()->op() == Op::kParam) {
       absl::StrAppendFormat(&res, "  ret %s\n",

xls/ir/function_base.cc

@@ -64,10 +64,33 @@ namespace {
 // sections may be empty after it runs.
 class StageSectioner : public DfsVisitorWithDefault {
  public:
-  explicit StageSectioner(const FunctionBase* fb)
-      : fb_(fb), sections_(fb->stages().size() * 2 + 1) {}
+  struct Section {
+    std::optional<int64_t> stage_index = std::nullopt;
+    std::list<Node*> nodes;
 
-  absl::Status Run() {
+    bool IsStage() const { return stage_index.has_value(); }
+
+    bool empty() const { return nodes.empty(); }
+
+    auto begin() const { return nodes.begin(); }
+    auto end() const { return nodes.end(); }
+
+    auto cbegin() const { return nodes.cbegin(); }
+    auto cend() const { return nodes.cend(); }
+
+    void push_back(Node* node) { nodes.push_back(node); }
+  };
+
+  explicit StageSectioner(const FunctionBase* fb) : fb_(fb) {
+    sections_.reserve(fb->stages().size() * 2 + 1);
+    for (int i = 0; i < fb_->stages().size(); i++) {
+      sections_.push_back(Section{.stage_index = std::nullopt});
+      sections_.push_back(Section{.stage_index = i});
+    }
+    sections_.push_back(Section{.stage_index = std::nullopt});
+  }
+
+  absl::Status Run(std::optional<std::vector<Node*>> order = std::nullopt) {
     // Capture stages and their deps.
     for (int i = 0; i < fb_->stages().size(); i++) {
       const Stage& stage = fb_->stages()[i];
@@ -107,6 +130,26 @@ class StageSectioner : public DfsVisitorWithDefault {
         XLS_RETURN_IF_ERROR(node->Accept(this));
       }
     }
+
+    // If provided, we use the given order to sort within each section.
+    if (order.has_value()) {
+      std::optional<absl::flat_hash_map<Node*, int>> node_to_index;
+      node_to_index.emplace();
+      node_to_index->reserve(order->size());
+      for (int i = 0; i < order->size(); ++i) {
+        node_to_index->emplace(order->at(i), i);
+      }
+
+      for (Section& section : sections_) {
+        CHECK(absl::c_all_of(section.nodes, [&](Node* node) {
+          return node_to_index->contains(node);
+        })) << "Some nodes not found in order";
+        section.nodes.sort([&](Node* a, Node* b) {
+          return node_to_index->at(a) < node_to_index->at(b);
+        });
+      }
+    }
+
     return absl::OkStatus();
   }
 
@@ -126,11 +169,11 @@ class StageSectioner : public DfsVisitorWithDefault {
     return absl::OkStatus();
   }
 
-  const std::vector<std::list<Node*>>& sections() const { return sections_; }
+  const std::vector<Section>& sections() const { return sections_; }
 
  private:
   const FunctionBase* fb_;
-  std::vector<std::list<Node*>> sections_;
+  std::vector<Section> sections_;
   int64_t current_stage_;
   absl::flat_hash_set<Node*> handled_;
 };
@@ -266,15 +309,12 @@ std::string FunctionBase::DumpFunctionBaseNodes(
     const IrAnnotator& annotate) const {
   std::string res;
   if (IsScheduled()) {
-    CHECK(!annotate.NodeOrder(const_cast<FunctionBase*>(this)))
-        << "Cannot use custom node order for scheduled entities";
     StageSectioner sectioner(this);
-    CHECK(sectioner.Run().ok());
-    for (const std::list<Node*>& section : sectioner.sections()) {
-      if (section.empty()) {
-        continue;
-      }
-
+    // In a scheduled entity, we use the annotator's order to sort within each
+    // section.
+    CHECK(sectioner.Run(annotate.NodeOrder(const_cast<FunctionBase*>(this)))
+              .ok());
+    for (const StageSectioner::Section& section : sectioner.sections()) {
       struct ScheduledAnnotations : public IrAnnotator {
        public:
         ScheduledAnnotations(Node* outputs_valid, Node* active_inputs_valid)
@@ -294,9 +334,8 @@ std::string FunctionBase::DumpFunctionBaseNodes(
         Node* active_inputs_valid_;
       };
 
-      Node* first = *section.begin();
-      if (IsStaged(first)) {
-        const Stage& stage = stages()[*GetStageIndex(first)];
+      if (section.IsStage()) {
+        const Stage& stage = stages()[*section.stage_index];
         IrAnnotatorJoiner joiner(
             ScheduledAnnotations(stage.outputs_valid(),
                                  stage.active_inputs_valid()),

xls/ir/function_base.h

@@ -102,6 +102,10 @@ class Stage {
     active_outputs_.erase(node);
   }
 
+  bool empty() const {
+    return active_inputs_.empty() && logic_.empty() && active_outputs_.empty();
+  }
+
   inline auto begin() {
     return iter::chain(active_inputs_, logic_, active_outputs_).begin();
   }

xls/ir/function_builder.cc

@@ -1821,8 +1821,8 @@ BValue BuilderBase::Cover(BValue condition, std::string_view label,
                              /*original_label=*/std::nullopt, name);
 }
 
-BValue BuilderBase::Gate(BValue condition, BValue data, const SourceInfo& loc,
-                         std::string_view name) {
+BValue BuilderBase::Gate(BValue condition, BValue data, GateType gate_type,
+                         const SourceInfo& loc, std::string_view name) {
   if (ErrorPending()) {
     return BValue();
   }
@@ -1834,7 +1834,8 @@ BValue BuilderBase::Gate(BValue condition, BValue data, const SourceInfo& loc,
             condition.GetType()->ToString()),
         loc);
   }
-  return AddNode<xls::Gate>(loc, condition.node(), data.node(), name);
+  return AddNode<xls::Gate>(loc, condition.node(), data.node(), gate_type,
+                            name);
 }
 
 BValue TokenlessProcBuilder::MinDelay(int64_t delay, const SourceInfo& loc,

xls/ir/function_builder.h

@@ -653,6 +653,16 @@ class BuilderBase {
   // Adds a gate operation. The output of the operation is `data` if `cond` is
   // true and zero-valued otherwise. Gates are side-effecting.
   BValue Gate(BValue condition, BValue data,
+              const SourceInfo& loc = SourceInfo(),
+              std::string_view name = "") {
+    return Gate(condition, data, GateType::kZeroGate, loc, name);
+  }
+
+  // Adds a gate operation with a specified gate type. The output of the
+  // operation is `data` if `cond` is true and otherwise depends on the gate
+  // type. Gates are side-effecting.
+  // are side-effecting.
+  BValue Gate(BValue condition, BValue data, GateType gate_type,
               const SourceInfo& loc = SourceInfo(), std::string_view name = "");
 
   // Add a receive operation. The type of the data value received is

xls/ir/ir_parser.cc

@@ -1351,8 +1351,16 @@ absl::StatusOr<BValue> Parser::ParseNode(
       break;
     }
     case Op::kGate: {
+      std::optional<IdentifierString>* gate_type_str =
+          arg_parser.AddOptionalKeywordArg<IdentifierString>("gate_type");
       XLS_ASSIGN_OR_RETURN(operands, arg_parser.Run(/*arity=*/2));
-      bvalue = fb->Gate(operands[0], operands[1], *loc, node_name);
+      if (gate_type_str->has_value()) {
+        XLS_ASSIGN_OR_RETURN(GateType gate_type,
+                             ParseGateType(gate_type_str->value().value));
+        bvalue = fb->Gate(operands[0], operands[1], gate_type, *loc, node_name);
+      } else {
+        bvalue = fb->Gate(operands[0], operands[1], *loc, node_name);
+      }
       break;
     }
     case Op::kInstantiationInput: {

xls/ir/node.cc

@@ -639,6 +639,13 @@ std::string Node::ToStringInternal(bool include_operand_types) const {
       }
       break;
     }
+    case Op::kGate: {
+      const Gate* gate = As<Gate>();
+      if (gate->gate_type() == GateType::kIgnorableGate) {
+        args.push_back("gate_type=ignorable");
+      }
+      break;
+    }
     case Op::kNext: {
       const Next* next = As<Next>();
       std::string param_name = next->has_state_read()

xls/ir/nodes.cc

@@ -26,6 +26,7 @@
 #include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
 #include "absl/types/span.h"
 #include "xls/common/math_util.h"
 #include "xls/common/status/ret_check.h"
@@ -1319,9 +1320,20 @@ bool InstantiationInput::IsDefinitelyEqualTo(const Node* other) const {
          port_name_ == other->As<InstantiationInput>()->port_name_;
 }
 
+absl::StatusOr<GateType> ParseGateType(std::string_view gate_type) {
+  if (gate_type == "zero") {
+    return GateType::kZeroGate;
+  } else if (gate_type == "ignorable") {
+    return GateType::kIgnorableGate;
+  }
+  return absl::InvalidArgumentError(
+      absl::StrCat("Unknown gate type: ", gate_type));
+}
+
 Gate::Gate(const SourceInfo& loc, Node* condition, Node* data,
-           std::string_view name, FunctionBase* function)
-    : Node(Op::kGate, data->GetType(), loc, name, function) {
+           GateType gate_type, std::string_view name, FunctionBase* function)
+    : Node(Op::kGate, data->GetType(), loc, name, function),
+      gate_type_(gate_type) {
   CHECK(IsOpClass<Gate>(op_))
       << "Op `" << op_ << "` is not a valid op for Node class `Gate`.";
   AddOperand(condition);
@@ -1331,8 +1343,8 @@ Gate::Gate(const SourceInfo& loc, Node* condition, Node* data,
 absl::StatusOr<Node*> Gate::CloneInNewFunction(
     absl::Span<Node* const> new_operands, FunctionBase* new_function) const {
   XLS_RET_CHECK_EQ(operand_count(), new_operands.size());
-  return new_function->MakeNodeWithName<Gate>(loc(), new_operands[0],
-                                              new_operands[1], GetNameView());
+  return new_function->MakeNodeWithName<Gate>(
+      loc(), new_operands[0], new_operands[1], gate_type(), GetNameView());
 }
 
 SliceData Concat::GetOperandSliceData(int64_t operandno) const {