[opt] [xlscc] Refactor DataFlowLazyNodeInfo::MergeInfos to reduce copies

The `MergeInfos` method in DataFlowLazyNodeInfo now takes a Span of Spans, instead of a flat `absl::Span<const Info>`. This avoids creating large temporary vectors by allowing merging across multiple input spans directly.

PiperOrigin-RevId: 887004771

Author: ericastor

xls/contrib/xlscc/continuations.cc

@@ -197,10 +197,12 @@ xls::LeafTypeTree<ParamSet> SourcesSetNodeInfo::ComputeInfoTreeForNode(
 }
 
 ParamSet SourcesSetNodeInfo::MergeInfos(
-    const absl::Span<const ParamSet>& infos) const {
+    absl::Span<const absl::Span<const ParamSet>> spans) const {
   ParamSet ret;
-  for (const ParamSet& info : infos) {
-    ret.insert(info.begin(), info.end());
+  for (const auto& span : spans) {
+    for (const ParamSet& info : span) {
+      ret.insert(info.begin(), info.end());
+    }
   }
   return ret;
 }
@@ -239,10 +241,12 @@ xls::LeafTypeTree<NodeSourceSet> SourcesSetTreeNodeInfo::ComputeInfoTreeForNode(
 }
 
 NodeSourceSet SourcesSetTreeNodeInfo::MergeInfos(
-    const absl::Span<const NodeSourceSet>& infos) const {
+    absl::Span<const absl::Span<const NodeSourceSet>> spans) const {
   NodeSourceSet ret;
-  for (const NodeSourceSet& info : infos) {
-    ret.insert(info.begin(), info.end());
+  for (const auto& span : spans) {
+    for (const NodeSourceSet& info : span) {
+      ret.insert(info.begin(), info.end());
+    }
   }
   return ret;
 }

xls/contrib/xlscc/translator_types.h

@@ -1266,7 +1266,7 @@ class SourcesSetNodeInfo
       xls::Node* node) const override final;
 
   ParamSet MergeInfos(
-      const absl::Span<const ParamSet>& infos) const override final;
+      absl::Span<const absl::Span<const ParamSet>> spans) const override final;
 };
 
 typedef absl::flat_hash_set<xls::NodeSource> NodeSourceSet;
@@ -1291,8 +1291,8 @@ class SourcesSetTreeNodeInfo
   xls::LeafTypeTree<NodeSourceSet> ComputeInfoTreeForNode(
       xls::Node* node) const override final;
 
-  NodeSourceSet MergeInfos(
-      const absl::Span<const NodeSourceSet>& infos) const override final;
+  NodeSourceSet MergeInfos(absl::Span<const absl::Span<const NodeSourceSet>>
+                               spans) const override final;
 };
 
 class OptimizationContext {

xls/passes/data_flow_node_info.h

@@ -81,7 +81,8 @@ class DataFlowLazyNodeInfo : public LazyNodeInfo<Info> {
   // Only used if compute_tree_for_source = true.
   virtual xls::LeafTypeTree<Info> ComputeInfoTreeForNode(Node* node) const = 0;
 
-  virtual Info MergeInfos(const absl::Span<const Info>& infos) const = 0;
+  virtual Info MergeInfos(
+      absl::Span<const absl::Span<const Info>> spans) const = 0;
 
   explicit DataFlowLazyNodeInfo(bool compute_tree_for_source,
                                 bool default_info_source,
@@ -115,7 +116,7 @@ class DataFlowLazyNodeInfo : public LazyNodeInfo<Info> {
 
   Info GetSingleInfoForNode(Node* node) {
     SharedLeafTypeTree<Info> info = LazyNodeInfo<Info>::GetInfo(node);
-    return MergeInfos(info.elements());
+    return MergeInfos({info.elements()});
   }
 
  private:
@@ -409,16 +410,15 @@ class DataFlowLazyNodeInfo : public LazyNodeInfo<Info> {
       ret = ComputeInfoForNode(node);
     } else {
       // Merge all operand infos
-      absl::InlinedVector<Info, 1> infos_in;
+      absl::InlinedVector<absl::Span<const Info>, 1> spans;
+      spans.reserve(operand_infos_out.size());
       for (int64_t op = 0; op < operand_infos_out.size(); ++op) {
-        for (int64_t i = 0; i < operand_infos_out[op]->elements().size(); ++i) {
-          infos_in.push_back(operand_infos_out[op]->elements()[i]);
-        }
+        spans.push_back(operand_infos_out[op]->elements());
       }
 
-      CHECK(!infos_in.empty());
+      CHECK(!spans.empty());
 
-      ret = MergeInfos(infos_in);
+      ret = MergeInfos(spans);
     }
 
     absl::InlinedVector<Info, 1> infos;
@@ -528,12 +528,15 @@ class DataFlowLazyNodeInfo : public LazyNodeInfo<Info> {
     }
 
     absl::InlinedVector<Info, 1> infos_out;
+    absl::InlinedVector<absl::Span<const Info>, 1> spans_this_elem;
+    spans_this_elem.reserve(operand_infos.size());
     for (int64_t e = 0; e < first_info->elements().size(); ++e) {
-      absl::InlinedVector<Info, 1> infos_this_elem;
+      spans_this_elem.clear();
       for (int64_t op = 0; op < operand_infos.size(); ++op) {
-        infos_this_elem.push_back(operand_infos.at(op)->elements().at(e));
+        spans_this_elem.push_back(
+            operand_infos.at(op)->elements().subspan(e, 1));
       }
-      infos_out.push_back(MergeInfos(infos_this_elem));
+      infos_out.push_back(MergeInfos(spans_this_elem));
     }
     return LeafTypeTree<Info>::CreateFromVector(first_info->type(),
                                                 std::move(infos_out));

xls/passes/data_flow_node_info_test.cc

@@ -69,10 +69,12 @@ class TestParamCountInfo
     return xls::LeafTypeTree<int64_t>();
   }
   int64_t MergeInfos(
-      const absl::Span<const int64_t>& infos) const override final {
+      absl::Span<const absl::Span<const int64_t>> spans) const override final {
     int64_t result = 0;
-    for (int64_t info : infos) {
-      result += info;
+    for (const auto& span : spans) {
+      for (int64_t info : span) {
+        result += info;
+      }
     }
     return result;
   }
@@ -116,11 +118,13 @@ class TestNodeSourceInfo
     return result;
   }
 
-  NodeSourceSet MergeInfos(
-      const absl::Span<const NodeSourceSet>& infos) const override final {
+  NodeSourceSet MergeInfos(absl::Span<const absl::Span<const NodeSourceSet>>
+                               spans) const override final {
     NodeSourceSet ret;
-    for (const NodeSourceSet& info : infos) {
-      ret.insert(info.begin(), info.end());
+    for (const auto& span : spans) {
+      for (const NodeSourceSet& info : span) {
+        ret.insert(info.begin(), info.end());
+      }
     }
     return ret;
   }