Add OptimizationContext to analysis passes and visitors.

This change modifies several analysis and pass implementations to accept and utilize an OptimizationContext. This allows them to leverage pre-computed topological sorts from the context, avoiding redundant computations.

Specifically:
-   DataflowVisitor and BitProvenanceAnalysis are updated to use `OptimizationContext` for topological sorts.
-   PredicateDominatorAnalysis is updated to use `OptimizationContext` for reverse topological sorts.
-   DataflowSimplificationPass and SelectSimplificationPass now pass the `OptimizationContext` to the analyses they use.
-   Benchmarks for PredicateDominatorAnalysis are updated to test with and without sharing the `OptimizationContext`.
PiperOrigin-RevId: 861883547

Author: allight

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,14 @@ absl::StatusOr<bool> DataflowSimplificationPass::RunOnFunctionBaseInternal(
     FunctionBase* func, const OptimizationPassOptions& options,
     PassResults* results, OptimizationContext& context) const {
   NodeSourceDataflowVisitor visitor;
-  XLS_RETURN_IF_ERROR(func->Accept(&visitor));
+  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/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

@@ -18,6 +18,7 @@
 #include <cstdint>
 #include <optional>
 #include <ostream>
+#include <utility>
 #include <vector>
 
 #include "absl/container/flat_hash_map.h"
@@ -26,7 +27,7 @@
 #include "xls/ir/function_base.h"
 #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 +98,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 +261,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