[delay_info_main] Add state critical path for procs

Since designing a proc can require understanding the delay cost of computing the state before the next activation can proceed, we add the critical path for each state element to the output of `delay_info_main`.

Example of an added section:
```
Critical path for state element __state:
      4ps (+  0ps): __state_next: () = next_value(param=__state, value=send.24, id=25)
      4ps (+  1ps): send.24: token = send(tok__3: token, bit_slice.16: bits[32], channel=serialized_decomposer__result4, id=24)
      3ps (+  1ps): tok__3: token = send(tok__2: token, bit_slice.14: bits[32], channel=serialized_decomposer__result3, id=23)
      2ps (+  1ps): tok__2: token = send(tok__1: token, bit_slice.12: bits[32], channel=serialized_decomposer__result2, id=22)
      1ps (+  1ps): tok__1: token = send(tok: token, bit_slice.10: bits[32], channel=serialized_decomposer__result1, id=21)
      0ps (+  0ps): tok: token = after_all(__state: token, input_tok: token, id=9)
      0ps (+  0ps): __state: token = state_read(state_element=__state, id=2)
```

Fixes #1770

PiperOrigin-RevId: 705533866

xls/estimators/delay_model/BUILD

@@ -137,7 +137,9 @@ cc_library(
         "//xls/common/status:status_macros",
         "//xls/ir",
         "//xls/ir:op",
+        "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",

xls/estimators/delay_model/analyze_critical_path.cc

@@ -21,7 +21,9 @@
 #include <utility>
 #include <vector>
 
+#include "absl/algorithm/container.h"
 #include "absl/container/flat_hash_map.h"
+#include "absl/functional/any_invocable.h"
 #include "absl/log/check.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
@@ -40,7 +42,9 @@ namespace xls {
 
 absl::StatusOr<std::vector<CriticalPathEntry>> AnalyzeCriticalPath(
     FunctionBase* f, std::optional<int64_t> clock_period_ps,
-    const DelayEstimator& delay_estimator) {
+    const DelayEstimator& delay_estimator,
+    absl::AnyInvocable<bool(Node*)> source_filter,
+    absl::AnyInvocable<bool(Node*)> sink_filter) {
   struct NodeEntry {
     Node* node;
 
@@ -65,15 +69,27 @@ absl::StatusOr<std::vector<CriticalPathEntry>> AnalyzeCriticalPath(
   std::optional<NodeEntry> latest_entry;
 
   for (Node* node : TopoSort(f)) {
+    if (!source_filter(node) &&
+        !absl::c_any_of(node->operands(), [&](Node* operand) {
+          return node_entries.contains(operand);
+        })) {
+      // This node is neither a source nor on a path from a source.
+      continue;
+    }
     NodeEntry& entry = node_entries[node];
     entry.node = node;
 
     // The maximum delay from any path up to but not including `node`.
     int64_t max_path_delay = 0;
     for (Node* operand : node->operands()) {
-      int64_t operand_path_delay = node_entries.at(operand).critical_path_delay;
+      auto it = node_entries.find(operand);
+      if (it == node_entries.end()) {
+        // This operand is neither a source nor on a path from a source.
+        continue;
+      }
+      int64_t operand_path_delay = it->second.critical_path_delay;
       if (operand_path_delay >= max_path_delay) {
-        max_path_delay = node_entries.at(operand).critical_path_delay;
+        max_path_delay = operand_path_delay;
         entry.critical_path_predecessor = operand;
       }
     }
@@ -94,13 +110,17 @@ absl::StatusOr<std::vector<CriticalPathEntry>> AnalyzeCriticalPath(
     }
     entry.critical_path_delay = max_path_delay + entry.node_delay;
 
+    if (!sink_filter(node)) {
+      continue;
+    }
     if (!latest_entry.has_value() ||
         latest_entry->critical_path_delay <= entry.critical_path_delay) {
       latest_entry = entry;
     }
   }
 
-  // `latest_entry` has no value for empty FunctionBases.
+  // `latest_entry` has no value for empty FunctionBases or if the source & sink
+  // filters removed all nodes.
   if (!latest_entry.has_value()) {
     return std::vector<CriticalPathEntry>();
   }

xls/estimators/delay_model/analyze_critical_path.h

@@ -21,6 +21,7 @@
 #include <string>
 #include <vector>
 
+#include "absl/functional/any_invocable.h"
 #include "absl/status/statusor.h"
 #include "absl/types/span.h"
 #include "xls/estimators/delay_model/delay_estimator.h"
@@ -60,7 +61,9 @@ struct CriticalPathEntry {
 // the front of the returned vector.
 absl::StatusOr<std::vector<CriticalPathEntry>> AnalyzeCriticalPath(
     FunctionBase* f, std::optional<int64_t> clock_period_ps,
-    const DelayEstimator& delay_estimator);
+    const DelayEstimator& delay_estimator,
+    absl::AnyInvocable<bool(Node*)> source_filter = [](Node*) { return true; },
+    absl::AnyInvocable<bool(Node*)> sink_filter = [](Node*) { return true; });
 
 // Returns a string representation of the critical-path. Includes delay
 // information for each node as well as cumulative delay.

xls/fdo/synthesized_delay_diff_utils.h

@@ -37,6 +37,7 @@ namespace synthesis {
 // synthesis tool reports.
 struct SynthesizedDelayDiff {
   std::vector<CriticalPathEntry> critical_path;
+  std::vector<CriticalPathEntry> state_critical_path;
   int64_t xls_delay_ps = 0;
   int64_t synthesized_delay_ps = 0;
 };

xls/tools/BUILD

@@ -926,6 +926,7 @@ cc_binary(
         "//xls/ir",
         "//xls/ir:ir_parser",
         "//xls/ir:op",
+        "//xls/ir:state_element",
         "//xls/scheduling:pipeline_schedule",
         "//xls/scheduling:pipeline_schedule_cc_proto",
         "@com_google_absl//absl/flags:flag",

xls/tools/delay_info_main.cc

@@ -43,8 +43,10 @@
 #include "xls/fdo/synthesizer.h"
 #include "xls/ir/ir_parser.h"
 #include "xls/ir/node.h"
+#include "xls/ir/nodes.h"
 #include "xls/ir/op.h"
 #include "xls/ir/package.h"
+#include "xls/ir/state_element.h"
 #include "xls/ir/topo_sort.h"
 #include "xls/scheduling/pipeline_schedule.h"
 #include "xls/scheduling/pipeline_schedule.pb.h"
@@ -202,6 +204,37 @@ absl::Status RealMain(std::string_view input_path) {
       }
       std::cout << "\n";
     }
+    if (top->IsProc()) {
+      Proc* proc = top->AsProcOrDie();
+      for (StateElement* state_element : proc->StateElements()) {
+        std::cout << absl::StrFormat("# Critical path for state element %s:\n",
+                                     state_element->name());
+        XLS_ASSIGN_OR_RETURN(
+            std::vector<CriticalPathEntry> state_critical_path,
+            AnalyzeCriticalPath(
+                top, /*clock_period_ps=*/std::nullopt, *delay_estimator,
+                [&](Node* node) {
+                  return node->Is<StateRead>() &&
+                         node->As<StateRead>()->state_element() ==
+                             state_element;
+                },
+                /*sink_filter=*/
+                [&](Node* node) {
+                  if (node->Is<StateRead>() &&
+                      node->As<StateRead>()->state_element() == state_element) {
+                    return true;
+                  }
+                  if (node->Is<Next>()) {
+                    return node->As<Next>()->state_read() ==
+                           proc->GetStateRead(state_element);
+                  }
+                  return node ==
+                         proc->GetNextStateElement(
+                             *proc->GetStateElementIndex(state_element));
+                }));
+        std::cout << CriticalPathToString(state_critical_path) << "\n";
+      }
+    }
   }
 
   std::cout << "# Delay of all nodes:\n";