Change how new FSM generates next state values.

- Reduce build times, fixing many timeouts
- Reduce the size of generated IR by orders of magnitude when there are many states, as with deeply nested loops
- Improve QoR

1. Caching of phi conditions in GeneratePhiCondition() reduces the size of the IR considerably.

2. The caching from #1 helps to enable the consolidation of entries added to extra_next_state_values, which ultimately feed into a priority select for the next value of the state element. The XLS back-end does much better with this reduced form, with fewer bits in the priority select, fed by ORs over multiple conditions.

PiperOrigin-RevId: 863297433

Author: Sean Purser-Haskell

xls/contrib/xlscc/BUILD

@@ -202,6 +202,7 @@ cc_library(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/types:span",
         "@llvm-project//clang:ast",
     ],
 )

xls/contrib/xlscc/generate_fsm.cc

@@ -19,6 +19,7 @@
 #include <limits>
 #include <optional>
 #include <string>
+#include <tuple>
 #include <utility>
 #include <vector>
 
@@ -32,6 +33,7 @@
 #include "absl/status/statusor.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/str_join.h"
+#include "absl/types/span.h"
 #include "clang/include/clang/AST/Decl.h"
 #include "xls/common/math_util.h"
 #include "xls/common/status/status_macros.h"
@@ -635,6 +637,8 @@ NewFSMGenerator::GenerateNewFSMInvocation(
   XLS_ASSIGN_OR_RETURN(layout,
                        LayoutNewFSM(func, state_element_for_static, body_loc));
 
+  absl::flat_hash_map<PhiConditionCacheKey, TrackedBValue> generated_conditions;
+
   const int64_t num_slice_index_bits =
       xls::CeilOfLog2(1 + xls_func->slices.size());
 
@@ -707,7 +711,7 @@ NewFSMGenerator::GenerateNewFSMInvocation(
   XLS_ASSIGN_OR_RETURN(
       phi_elements_by_param_node_id,
       GeneratePhiConditions(layout, state_element_by_jump_slice_index, pb,
-                            body_loc));
+                            body_loc, generated_conditions));
 
   // The value from the current activation's perspective,
   // either outputted from invoke or state element.
@@ -763,6 +767,30 @@ NewFSMGenerator::GenerateNewFSMInvocation(
   TrackedBValue after_activation_transition =
       pb.Literal(xls::UBits(0, 1), body_loc);
 
+  // Sort by Node ID and StateElement name for determinism.
+  struct StateElementAndNodeLessThan {
+    bool operator()(const std::tuple<xls::StateElement*, xls::Node*>& a,
+                    const std::tuple<xls::StateElement*, xls::Node*>& b) const {
+      const auto& [a_elem, a_node] = a;
+      const auto& [b_elem, b_node] = b;
+      if (a_elem->name() != b_elem->name()) {
+        return a_elem->name() < b_elem->name();
+      }
+      return a_node->id() < b_node->id();
+    }
+  };
+
+  struct NodeIdLessThan {
+    bool operator()(const xls::Node* a, const xls::Node* b) const {
+      return a->id() < b->id();
+    }
+  };
+
+  absl::btree_map<std::tuple<xls::StateElement*, xls::Node*>,
+                  absl::btree_set<xls::Node*, NodeIdLessThan>,
+                  StateElementAndNodeLessThan>
+      next_value_conditions_by_state_element_and_value;
+
   for (int64_t slice_index = 0; slice_index < func.slices.size();
        ++slice_index) {
     const bool is_last_slice = (slice_index == func.slices.size() - 1);
@@ -1026,44 +1054,65 @@ NewFSMGenerator::GenerateNewFSMInvocation(
         if (state.slice_index != slice_index) {
           continue;
         }
-
         absl::btree_set<int64_t> jumped_from_slice_indices_this_state;
         for (const JumpInfo& jump_info : state.jumped_from_slice_indices) {
           jumped_from_slice_indices_this_state.insert(jump_info.from_slice);
         }
 
         XLS_ASSIGN_OR_RETURN(
             TrackedBValue state_active_condition,
-            GeneratePhiCondition(from_jump_slice_indices,
-                                 jumped_from_slice_indices_this_state,
-                                 state_element_by_jump_slice_index, pb,
-                                 state.slice_index, body_loc));
+            GeneratePhiCondition(
+                from_jump_slice_indices, jumped_from_slice_indices_this_state,
+                state_element_by_jump_slice_index, pb, state.slice_index,
+                body_loc, generated_conditions));
 
         TrackedBValue next_value_condition =
             pb.And(state_active_condition, jump_condition, body_loc,
                    /*name=*/GetIRStateName(state));
 
         for (const ContinuationValue* continuation_out : state.values_to_save) {
           // Generate next values for state elements
-          NextStateValue next_value = {
-              .priority = 0,
-              .value = value_by_continuation_value.at(continuation_out),
-              .condition = next_value_condition,
-          };
-
           xls::StateElement* state_elem =
               state_element_by_continuation_value.at(continuation_out)
                   .node()
                   ->As<xls::StateRead>()
                   ->state_element();
 
+          std::tuple<xls::StateElement*, xls::Node*> key = {
+              state_elem,
+              value_by_continuation_value.at(continuation_out).node()};
+
           // Generate next values
-          extra_next_state_values.insert({state_elem, next_value});
+          next_value_conditions_by_state_element_and_value[key].insert(
+              next_value_condition.node());
         }
       }
     }
   }
 
+  for (auto& [key, or_nodes] :
+       next_value_conditions_by_state_element_and_value) {
+    xls::StateElement* state_elem = std::get<0>(key);
+    xls::Node* next_value_node = std::get<1>(key);
+    std::vector<NATIVE_BVAL> or_bvals;
+    for (xls::Node* or_node : or_nodes) {
+      or_bvals.push_back(NATIVE_BVAL(or_node, &pb));
+    }
+
+    TrackedBValue or_bval =
+        pb.Or(absl::MakeSpan(or_bvals), body_loc,
+              /*name=*/
+              absl::StrFormat("%s_v_%s_or_bval", state_elem->name(),
+                              next_value_node->GetName()));
+
+    NextStateValue next_value = {
+        .priority = 0,
+        .value = TrackedBValue(next_value_node, &pb),
+        .condition = or_bval,
+    };
+    extra_next_state_values.insert({state_elem, next_value});
+  }
+
   // Set next slice index
   const TrackedBValue finished_iteration =
       pb.Not(after_activation_transition, body_loc,
@@ -1098,8 +1147,16 @@ absl::StatusOr<TrackedBValue> NewFSMGenerator::GeneratePhiCondition(
     const absl::btree_set<int64_t>& jumped_from_slice_indices_this_state,
     const absl::flat_hash_map<int64_t, TrackedBValue>&
         state_element_by_jump_slice_index,
-    xls::ProcBuilder& pb, int64_t slice_index,
-    const xls::SourceInfo& body_loc) {
+    xls::ProcBuilder& pb, int64_t slice_index, const xls::SourceInfo& body_loc,
+    absl::flat_hash_map<PhiConditionCacheKey, TrackedBValue>&
+        phi_condition_cache) {
+  PhiConditionCacheKey key = {from_jump_slice_indices,
+                              jumped_from_slice_indices_this_state};
+
+  if (phi_condition_cache.contains(key)) {
+    return phi_condition_cache.at(key);
+  }
+
   TrackedBValue condition = pb.Literal(xls::UBits(1, 1), body_loc);
 
   // Include all jump slices in each condition
@@ -1110,6 +1167,7 @@ absl::StatusOr<TrackedBValue> NewFSMGenerator::GeneratePhiCondition(
         jumped_from_slice_indices_this_state.contains(from_jump_slice_index)
             ? 1
             : 0;
+
     TrackedBValue condition_part =
         pb.Eq(jump_state_element,
               pb.Literal(xls::UBits(active_value, 1), body_loc,
@@ -1125,6 +1183,7 @@ absl::StatusOr<TrackedBValue> NewFSMGenerator::GeneratePhiCondition(
                    absl::StrJoin(jumped_from_slice_indices_this_state, "_")));
   }
 
+  phi_condition_cache[key] = condition;
   return condition;
 }
 
@@ -1134,7 +1193,9 @@ NewFSMGenerator::GeneratePhiConditions(
     const NewFSMLayout& layout,
     const absl::flat_hash_map<int64_t, TrackedBValue>&
         state_element_by_jump_slice_index,
-    xls::ProcBuilder& pb, const xls::SourceInfo& body_loc) {
+    xls::ProcBuilder& pb, const xls::SourceInfo& body_loc,
+    absl::flat_hash_map<PhiConditionCacheKey, TrackedBValue>&
+        phi_condition_cache) {
   absl::flat_hash_map<int64_t, std::vector<PhiElement>>
       phi_elements_by_param_node_id;
 
@@ -1171,10 +1232,10 @@ NewFSMGenerator::GeneratePhiConditions(
 
       XLS_ASSIGN_OR_RETURN(
           TrackedBValue condition,
-          GeneratePhiCondition(from_jump_slice_indices,
-                               jumped_from_slice_indices_this_state,
-                               state_element_by_jump_slice_index, pb,
-                               state->slice_index, body_loc));
+          GeneratePhiCondition(
+              from_jump_slice_indices, jumped_from_slice_indices_this_state,
+              state_element_by_jump_slice_index, pb, state->slice_index,
+              body_loc, phi_condition_cache));
 
       PhiElement& phi_element = phi_elements.emplace_back();
       phi_element.value = state->current_inputs_by_input_param.at(param);
@@ -1204,13 +1265,47 @@ NewFSMGenerator::GenerateInputValueInContext(
   std::vector<TrackedBValue> phi_conditions;
   std::vector<TrackedBValue> phi_values;
 
+  // Sort by Node ID for determinism.
+  struct NodeIdLessThan {
+    bool operator()(const xls::Node* a, const xls::Node* b) const {
+      return a->id() < b->id();
+    }
+  };
+  struct BValueIdLessThan {
+    bool operator()(const TrackedBValue& a, const TrackedBValue& b) const {
+      return a.node()->id() < b.node()->id();
+    }
+  };
+  absl::btree_map<xls::Node*, absl::btree_set<TrackedBValue, BValueIdLessThan>,
+                  NodeIdLessThan>
+      conditions_by_value_node;
+
   phi_conditions.reserve(phi_elements.size());
   phi_values.reserve(phi_elements.size());
+
   for (const PhiElement& phi_element : phi_elements) {
-    phi_conditions.push_back(phi_element.condition);
     XLSCC_CHECK(value_by_continuation_value.contains(phi_element.value),
                 phi_element.value->output_node->loc());
-    phi_values.push_back(value_by_continuation_value.at(phi_element.value));
+
+    xls::Node* value_node =
+        value_by_continuation_value.at(phi_element.value).node();
+    conditions_by_value_node[value_node].insert(phi_element.condition);
+  }
+
+  for (auto& [value_node, or_nodes] : conditions_by_value_node) {
+    std::vector<NATIVE_BVAL> or_bvals;
+    or_bvals.reserve(or_nodes.size());
+    for (const TrackedBValue& or_node : or_nodes) {
+      or_bvals.push_back(or_node);
+    }
+
+    TrackedBValue or_bval =
+        pb.Or(absl::MakeSpan(or_bvals), body_loc,
+              /*name=*/
+              absl::StrFormat("%s_v_%s_or_bval", param->name(),
+                              value_node->GetName()));
+    phi_conditions.push_back(or_bval);
+    phi_values.push_back(TrackedBValue(value_node, &pb));
   }
 
   std::reverse(phi_conditions.begin(), phi_conditions.end());

xls/contrib/xlscc/generate_fsm.h

@@ -145,19 +145,27 @@ class NewFSMGenerator : public GeneratorBase {
     const ContinuationValue* value;
   };
 
+  typedef std::tuple<absl::btree_set<int64_t>, absl::btree_set<int64_t>>
+      PhiConditionCacheKey;
+
   absl::StatusOr<absl::flat_hash_map<int64_t, std::vector<PhiElement>>>
-  GeneratePhiConditions(const NewFSMLayout& layout,
-                        const absl::flat_hash_map<int64_t, TrackedBValue>&
-                            state_element_by_jump_slice_index,
-                        xls::ProcBuilder& pb, const xls::SourceInfo& body_loc);
+  GeneratePhiConditions(
+      const NewFSMLayout& layout,
+      const absl::flat_hash_map<int64_t, TrackedBValue>&
+          state_element_by_jump_slice_index,
+      xls::ProcBuilder& pb, const xls::SourceInfo& body_loc,
+      absl::flat_hash_map<PhiConditionCacheKey, TrackedBValue>&
+          phi_condition_cache);
 
   absl::StatusOr<TrackedBValue> GeneratePhiCondition(
       const absl::btree_set<int64_t>& from_jump_slice_indices,
       const absl::btree_set<int64_t>& jumped_from_slice_indices_this_state,
       const absl::flat_hash_map<int64_t, TrackedBValue>&
           state_element_by_jump_slice_index,
       xls::ProcBuilder& pb, int64_t slice_index,
-      const xls::SourceInfo& body_loc);
+      const xls::SourceInfo& body_loc,
+      absl::flat_hash_map<PhiConditionCacheKey, TrackedBValue>&
+          phi_condition_cache);
 
   absl::StatusOr<std::optional<TrackedBValue>> GenerateInputValueInContext(
       const xls::Param* param,