From b736a518595798657dc119107cea72c1c4790d2a Mon Sep 17 00:00:00 2001 From: Alex Light Date: Wed, 28 Jan 2026 13:16:57 -0800 Subject: [PATCH] Add a ReserveNodes method to DfsVisitor. This allows visitors to pre-allocate space for visited nodes, potentially reducing reallocations. The DataflowSimplificationPass now uses this to reserve space based on the function's node count. PiperOrigin-RevId: 862366169 --- xls/ir/dfs_visitor.h | 6 ++++++ xls/passes/dataflow_simplification_pass.cc | 1 + xls/passes/dataflow_visitor.h | 4 ++++ 3 files changed, 11 insertions(+) diff --git a/xls/ir/dfs_visitor.h b/xls/ir/dfs_visitor.h index 1b70ca0696..9012dd26b6 100644 --- a/xls/ir/dfs_visitor.h +++ b/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; diff --git a/xls/passes/dataflow_simplification_pass.cc b/xls/passes/dataflow_simplification_pass.cc index 48e5114ba9..fa0376d285 100644 --- a/xls/passes/dataflow_simplification_pass.cc +++ b/xls/passes/dataflow_simplification_pass.cc @@ -86,6 +86,7 @@ absl::StatusOr DataflowSimplificationPass::RunOnFunctionBaseInternal( FunctionBase* func, const OptimizationPassOptions& options, PassResults* results, OptimizationContext& context) const { NodeSourceDataflowVisitor visitor; + visitor.ReserveNodes(func->node_count()); for (Node* node : context.TopoSort(func)) { XLS_RETURN_IF_ERROR(node->VisitSingleNode(&visitor)); } diff --git a/xls/passes/dataflow_visitor.h b/xls/passes/dataflow_visitor.h index 6efe5eb2d4..396961d0ea 100644 --- a/xls/passes/dataflow_visitor.h +++ b/xls/passes/dataflow_visitor.h @@ -84,6 +84,10 @@ namespace xls { template 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> elements; for (Node* operand : array->operands()) {