Add a ReserveNodes method to DfsVisitor.

xls/data_structures/binary_decision_diagram.cc

@@ -344,6 +344,14 @@ std::string BinaryDecisionDiagram::ToStringDnf(BddNodeIndex expr,
   return result;
 }
 
+namespace {
+#ifdef NDEBUG
+constexpr bool kDebug = false;
+#else
+constexpr bool kDebug = true;
+#endif
+}  // namespace
+
 absl::StatusOr<std::vector<BddNodeIndex>> BinaryDecisionDiagram::GarbageCollect(
     absl::Span<BddNodeIndex const> roots, double gc_threshold) {
   if (nodes_size_ * gc_threshold < prev_nodes_size_) {
@@ -400,12 +408,12 @@ absl::StatusOr<std::vector<BddNodeIndex>> BinaryDecisionDiagram::GarbageCollect(
     }
   }
   // build a new free-list prefix.
-  for (auto vs : iter::sliding_window(
-           iter::chain(dead, std::vector<BddNodeIndex>{BddNodeIndex(
-                                 free_node_head_.value())}),
-           2)) {
-    BddNodeIndex first = vs[0];
-    BddNodeIndex second = vs[1];
+  // Don't use iter::sliding_window because this is faster.
+  for (int i = 0; i < dead.size(); ++i) {
+    BddNodeIndex first = dead[i];
+    BddNodeIndex second = i + 1 < dead.size()
+                              ? dead[i + 1]
+                              : BddNodeIndex(free_node_head_.value());
     nodes_[first.value()] = FreeListNode(FreeListNode::Index(second.value()));
   }
   free_node_head_ = FreeListNode::Index(dead.front().value());
@@ -439,24 +447,26 @@ absl::StatusOr<std::vector<BddNodeIndex>> BinaryDecisionDiagram::GarbageCollect(
   }
   nodes_size_ = cnt_live;
   prev_nodes_size_ = nodes_size_;
-  auto head = free_node_head_;
-  auto it = dead.begin();
-  InlineBitmap debug_seen(nodes_.size());
-  while (head.value() < nodes_.size() &&
-         head != FreeListNode::kNextIsConsecutive) {
-    CHECK(!debug_seen.Get(head.value()))
-        << " loop in free list at " << head.value();
-    debug_seen.Set(head.value());
-    if (it != dead.end()) {
-      CHECK_EQ(head, FreeListNode::Index(it->value()));
-      ++it;
-    }
-    CHECK(std::holds_alternative<FreeListNode>(nodes_[head.value()]));
-    auto cur = std::get<FreeListNode>(nodes_[head.value()]);
-    if (cur.raw_next() == FreeListNode::kNextIsConsecutive) {
-      break;
+  if constexpr (kDebug) {
+    auto head = free_node_head_;
+    auto it = dead.begin();
+    InlineBitmap debug_seen(nodes_.size());
+    while (head.value() < nodes_.size() &&
+           head != FreeListNode::kNextIsConsecutive) {
+      CHECK(!debug_seen.Get(head.value()))
+          << " loop in free list at " << head.value();
+      debug_seen.Set(head.value());
+      if (it != dead.end()) {
+        CHECK_EQ(head, FreeListNode::Index(it->value()));
+        ++it;
+      }
+      CHECK(std::holds_alternative<FreeListNode>(nodes_[head.value()]));
+      auto cur = std::get<FreeListNode>(nodes_[head.value()]);
+      if (cur.raw_next() == FreeListNode::kNextIsConsecutive) {
+        break;
+      }
+      head = cur.next_free(head);
     }
-    head = cur.next_free(head);
   }
   return std::move(dead);
 }

xls/ir/dfs_visitor.h

@@ -30,6 +30,12 @@ class DfsVisitor {
  public:
   virtual ~DfsVisitor() = default;
 
+  // Let the visitor know that 'count' nodes are likely to be visited in the
+  // near future. This is purely advisory to allow the visitor to reserve space
+  // to avoid reallocations. The visitor is free to ignore this hint. The user
+  // is free to call this method or not.
+  virtual void ReserveNodes(int64_t count) { visited_.reserve(count); }
+
   virtual absl::Status HandleAdd(BinOp* add) = 0;
   virtual absl::Status HandleAfterAll(AfterAll* after_all) = 0;
   virtual absl::Status HandleMinDelay(MinDelay* min_delay) = 0;

xls/passes/BUILD

@@ -1337,6 +1337,7 @@ cc_library(
     srcs = ["predicate_dominator_analysis.cc"],
     hdrs = ["predicate_dominator_analysis.h"],
     deps = [
+        ":optimization_pass",
         ":predicate_state",
         "//xls/common:strong_int",
         "//xls/ir",
@@ -3193,6 +3194,7 @@ cc_test(
     name = "predicate_dominator_analysis_test",
     srcs = ["predicate_dominator_analysis_test.cc"],
     deps = [
+        ":optimization_pass",
         ":predicate_dominator_analysis",
         ":predicate_state",
         "//xls/common:xls_gunit_main",
@@ -4080,6 +4082,7 @@ cc_library(
     hdrs = ["bit_provenance_analysis.h"],
     deps = [
         ":dataflow_visitor",
+        ":optimization_pass",
         ":query_engine",
         "//xls/common/status:ret_check",
         "//xls/common/status:status_macros",

xls/passes/bit_provenance_analysis.cc

@@ -34,6 +34,7 @@
 #include "xls/ir/nodes.h"
 #include "xls/ir/type.h"
 #include "xls/passes/dataflow_visitor.h"
+#include "xls/passes/optimization_pass.h"
 #include "xls/passes/query_engine.h"
 
 namespace xls {
@@ -301,6 +302,13 @@ BitProvenanceAnalysis::CreatePrepopulated(FunctionBase* func) {
   XLS_RETURN_IF_ERROR(result.Populate(func));
   return result;
 }
+/* static */ absl::StatusOr<BitProvenanceAnalysis>
+BitProvenanceAnalysis::CreatePrepopulated(FunctionBase* func,
+                                          OptimizationContext& context) {
+  BitProvenanceAnalysis result;
+  XLS_RETURN_IF_ERROR(result.Populate(func, context));
+  return result;
+}
 
 BitProvenanceAnalysis::BitProvenanceAnalysis()
     : visitor_{std::make_unique<internal::BitProvenanceVisitor>()} {}
@@ -320,6 +328,14 @@ absl::Status BitProvenanceAnalysis::Populate(FunctionBase* func) {
   XLS_RETURN_IF_ERROR(func->Accept(visitor_.get()));
   return absl::OkStatus();
 }
+absl::Status BitProvenanceAnalysis::Populate(FunctionBase* func,
+                                             OptimizationContext& context) {
+  for (Node* node : context.TopoSort(func)) {
+    XLS_RETURN_IF_ERROR(node->VisitSingleNode(visitor_.get()));
+    visitor_->MarkVisited(node);
+  }
+  return absl::OkStatus();
+}
 
 absl::StatusOr<TreeBitLocation> BitProvenanceAnalysis::GetSource(
     const TreeBitLocation& bit) const {

xls/passes/bit_provenance_analysis.h

@@ -28,6 +28,7 @@
 #include "absl/types/span.h"
 #include "xls/data_structures/leaf_type_tree.h"
 #include "xls/ir/node.h"
+#include "xls/passes/optimization_pass.h"
 #include "xls/passes/query_engine.h"
 
 namespace xls {
@@ -157,6 +158,8 @@ class BitProvenanceAnalysis {
   // invalid if the function is modified.
   static absl::StatusOr<BitProvenanceAnalysis> CreatePrepopulated(
       FunctionBase* func);
+  static absl::StatusOr<BitProvenanceAnalysis> CreatePrepopulated(
+      FunctionBase* func, OptimizationContext& context);
 
   // constructors and destructors need to be declared here and implemented in
   // the .cc file to avoid the compiler inserting constructors and destructors
@@ -170,6 +173,7 @@ class BitProvenanceAnalysis {
   BitProvenanceAnalysis& operator=(BitProvenanceAnalysis&& other);
 
   absl::Status Populate(FunctionBase* func);
+  absl::Status Populate(FunctionBase* func, OptimizationContext& context);
 
   // Get the tree-bit-location which provides the original source of the given
   // bit.

xls/passes/dataflow_simplification_pass.cc

@@ -86,11 +86,15 @@ absl::StatusOr<bool> DataflowSimplificationPass::RunOnFunctionBaseInternal(
     FunctionBase* func, const OptimizationPassOptions& options,
     PassResults* results, OptimizationContext& context) const {
   NodeSourceDataflowVisitor visitor;
-  XLS_RETURN_IF_ERROR(func->Accept(&visitor));
+  visitor.ReserveNodes(func->node_count());
+  for (Node* node : context.TopoSort(func)) {
+    XLS_RETURN_IF_ERROR(node->VisitSingleNode(&visitor));
+  }
   bool changed = false;
   // Hashmap from the LTT<NodeSource> of a node to the Node*. If two nodes have
   // the same LTT<NodeSource> they are necessarily equivalent.
   absl::flat_hash_map<LeafTypeTreeView<NodeSource>, Node*> source_map;
+  source_map.reserve(func->node_count());
   for (Node* node : context.TopoSort(func)) {
     LeafTypeTreeView<NodeSource> source = visitor.GetValue(node);
     VLOG(3) << absl::StrFormat("Considering `%s`: %s", node->GetName(),

xls/passes/dataflow_visitor.h

@@ -84,6 +84,10 @@ namespace xls {
 template <typename LeafT>
 class DataflowVisitor : public DfsVisitorWithDefault {
  public:
+  void ReserveNodes(int64_t count) override {
+    DfsVisitorWithDefault::ReserveNodes(count);
+    map_.reserve(count);
+  }
   absl::Status HandleArray(Array* array) override {
     std::vector<LeafTypeTreeView<LeafT>> elements;
     for (Node* operand : array->operands()) {

xls/passes/narrowing_pass.cc

@@ -2258,7 +2258,7 @@ absl::StatusOr<bool> NarrowingPass::RunOnFunctionBaseInternal(
   XLS_ASSIGN_OR_RETURN(AliasingQueryEngine query_engine,
                        GetQueryEngine(f, RealAnalysis(options), context));
 
-  PredicateDominatorAnalysis pda = PredicateDominatorAnalysis::Run(f);
+  PredicateDominatorAnalysis pda = PredicateDominatorAnalysis::Run(f, context);
   SpecializedQueryEngines sqe(RealAnalysis(options), pda, query_engine);
 
   NarrowVisitor narrower(sqe, RealAnalysis(options), options,

xls/passes/predicate_dominator_analysis.cc

@@ -16,8 +16,10 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <functional>
 #include <optional>
 #include <ostream>
+#include <utility>
 #include <vector>
 
 #include "absl/container/flat_hash_map.h"
@@ -27,6 +29,7 @@
 #include "xls/ir/node.h"
 #include "xls/ir/nodes.h"
 #include "xls/ir/topo_sort.h"
+#include "xls/passes/optimization_pass.h"
 #include "xls/passes/predicate_state.h"
 
 namespace xls {
@@ -97,15 +100,16 @@ class AnalysisHelper {
       .previous = kRootPredicateId,
       .distance_to_root = 0};
 
-  explicit AnalysisHelper(FunctionBase* func) : function_(func) {}
+  AnalysisHelper(FunctionBase* func, OptimizationContext& context)
+      : function_(func), context_(context) {}
 
   absl::flat_hash_map<Node*, PredicateState> Analyze() {
     CHECK(node_states_.empty());
     CHECK(predicate_stacks_.empty());
     node_states_.reserve(function_->node_count());
     predicate_stacks_.push_back(kRootPredicateStackNode);
     // Run in reverse topo sort order. Handle users before the values they use.
-    for (Node* node : ReverseTopoSort(function_)) {
+    for (Node* node : context_.ReverseTopoSort(function_)) {
       HandleNode(node);
     }
 
@@ -259,15 +263,17 @@ class AnalysisHelper {
 
  private:
   FunctionBase* function_;
+  OptimizationContext& context_;
   // Map of node to the predicate list head they are guarded by.
   absl::flat_hash_map<Node*, PredicateStackId> node_states_;
   // Map from 'PredicateStackId' to the predicate node.
   std::vector<PredicateStackNode> predicate_stacks_;
 };
 }  // namespace
 
-PredicateDominatorAnalysis PredicateDominatorAnalysis::Run(FunctionBase* f) {
-  AnalysisHelper helper(f);
+PredicateDominatorAnalysis PredicateDominatorAnalysis::Run(
+    FunctionBase* f, OptimizationContext& context) {
+  AnalysisHelper helper(f, context);
   return PredicateDominatorAnalysis(helper.Analyze());
 }
 

xls/passes/predicate_dominator_analysis.h

@@ -20,6 +20,7 @@
 #include "absl/container/flat_hash_map.h"
 #include "xls/ir/function_base.h"
 #include "xls/ir/node.h"
+#include "xls/passes/optimization_pass.h"
 #include "xls/passes/predicate_state.h"
 
 namespace xls {
@@ -38,7 +39,8 @@ class PredicateDominatorAnalysis {
   PredicateDominatorAnalysis& operator=(PredicateDominatorAnalysis&&) = default;
 
   // Execute this analysis and return results.
-  static PredicateDominatorAnalysis Run(FunctionBase* f);
+  static PredicateDominatorAnalysis Run(FunctionBase* f,
+                                        OptimizationContext& context);
 
   // Returns a single element of the common predicate dominators which is
   // closest to the node (ie the last predicate which gates the use of this