Create BddQueryEngine per-function in benchmark schedule printing.

The BddQueryEngine is associated with a single function. When printing schedule information for a PackageSchedule containing multiple functions, a new BddQueryEngine must be created and populated for each function in the package.

Fixes: #3637
PiperOrigin-RevId: 886998620

Author: allight

xls/dev_tools/BUILD

@@ -70,6 +70,14 @@ cc_binary(
     ],
 )
 
+sh_test(
+    name = "benchmark_main_test",
+    srcs = ["benchmark_main_test.sh"],
+    data = [
+        "//xls/examples:proc_network_opt_ir_benchmark",
+    ],
+)
+
 pytype_strict_contrib_test(
     name = "tool_timeout_test",
     srcs = ["tool_timeout_test.py"],

xls/dev_tools/benchmark_main.cc

@@ -416,9 +416,7 @@ absl::Status PrintScheduleInfo(FunctionBase* f,
   return absl::OkStatus();
 }
 
-absl::Status PrintScheduleInfo(FunctionBase* f,
-                               const PackageSchedule& package_schedule,
-                               const BddQueryEngine& bdd_query_engine,
+absl::Status PrintScheduleInfo(const PackageSchedule& package_schedule,
                                const DelayEstimator& delay_estimator,
                                std::optional<int64_t> clock_period_ps) {
   if (package_schedule.GetSchedules().size() == 1) {
@@ -436,9 +434,12 @@ absl::Status PrintScheduleInfo(FunctionBase* f,
   }
   for (auto& [function_base, schedule] : package_schedule.GetSchedules()) {
     std::cout << "\n\nFunction: " << function_base->name() << "\n";
-    XLS_RETURN_IF_ERROR(PrintScheduleInfo(function_base, schedule,
-                                          bdd_query_engine, delay_estimator,
-                                          clock_period_ps));
+    // NB query engines are associated with a specific function, so we need to
+    // create a new one for each function.
+    BddQueryEngine query_engine(BddQueryEngine::kDefaultPathLimit);
+    XLS_RETURN_IF_ERROR(query_engine.Populate(function_base).status());
+    XLS_RETURN_IF_ERROR(PrintScheduleInfo(function_base, schedule, query_engine,
+                                          delay_estimator, clock_period_ps));
   }
   return absl::OkStatus();
 }
@@ -858,7 +859,7 @@ absl::Status RealMain(std::string_view path) {
           synthesizer.get()));
 
       XLS_RETURN_IF_ERROR(PrintScheduleInfo(
-          f, package_schedule, query_engine, defaulted_delay_estimator,
+          package_schedule, defaulted_delay_estimator,
           scheduling_options_flags_proto.has_clock_period_ps()
               ? std::make_optional(
                     scheduling_options_flags_proto.clock_period_ps())

xls/dev_tools/benchmark_main_test.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+MY_BENCHMARK_FILE="./xls/examples/proc_network_opt_ir_benchmark.sh"
+
+die () {
+  echo "ERROR: $1"
+  exit 1
+}
+
+# Run the benchmark
+${MY_BENCHMARK_FILE} || die "Failed in ${MY_BENCHMARK_FILE}"
+
+echo "PASS"

xls/examples/BUILD

@@ -1382,11 +1382,32 @@ xls_dslx_test(
 xls_dslx_ir(
     name = "proc_network_ir",
     dslx_top = "Initiator",
-    ir_conv_args = {"lower_to_proc_scoped_channels": "true"},
+    ir_conv_args = {
+        "lower_to_proc_scoped_channels": "true",
+        "default_fifo_config": "depth: 2, bypass: true, " +
+                               "register_push_outputs: false, register_pop_outputs: false",
+    },
     ir_file = "proc_network.ir",
     library = ":proc_network_dslx",
 )
 
+xls_ir_opt_ir(
+    name = "proc_network_opt_ir",
+    src = "proc_network.ir",
+)
+
+xls_benchmark_ir(
+    name = "proc_network_opt_ir_benchmark",
+    src = ":proc_network_opt_ir",
+    codegen_args = {
+        "module_name": "ProcName",
+        "delay_model": "unit",
+        "pipeline_stages": "2",
+        "reset": "rst",
+    },
+    synthesize = True,
+)
+
 xls_dslx_ir(
     name = "proc_network_proc_scoped_ir",
     dslx_top = "Initiator",