[opt] Migrate more passes to take advantage of PartialInformation

xls/ir/interval_ops.cc

@@ -141,33 +141,41 @@ IntervalSet FromTernary(TernarySpan tern, int64_t max_interval_bits) {
     // Need to extend the x-s to avoid creating too many intervals.
     lsb_xs = (x_locations.front() - tern.cbegin()) + 1;
     x_locations.pop_front();
+
+    // Make sure to include any contiguous X's in the trailing unknown region,
+    // maintaining that `lsb_xs` points to the first known bit (that we retain),
+    // and `x_locations` only includes the X's above that.
+    while (!x_locations.empty() &&
+           lsb_xs == (x_locations.front() - tern.cbegin())) {
+      ++lsb_xs;
+      x_locations.pop_front();
+    }
   }
 
-  IntervalSet is(tern.size());
-  if (x_locations.empty()) {
-    // All bits from 0 -> lsb_xs are unknown.
-    Bits high_bits = ternary_ops::ToKnownBitsValues(tern.subspan(lsb_xs));
-    is.AddInterval(Interval::Closed(
+  // Capture the input ternary above the last lsb_x.
+  TernarySpan prefix = tern.subspan(lsb_xs);
+
+  if (x_locations.empty() || lsb_xs == tern.size()) {
+    // All bits from 0 -> lsb_xs are unknown, and everything above it is known.
+    Bits high_bits = ternary_ops::ToKnownBitsValues(prefix);
+    return IntervalSet::Of({Interval::Closed(
         bits_ops::UMax(lb, bits_ops::Concat({high_bits, Bits(lsb_xs)})),
         bits_ops::UMin(ub,
-                       bits_ops::Concat({high_bits, Bits::AllOnes(lsb_xs)}))));
-    is.Normalize();
-    return is;
+                       bits_ops::Concat({high_bits, Bits::AllOnes(lsb_xs)})))});
   }
 
-  TernaryVector vec(tern.size() - lsb_xs, TernaryValue::kKnownZero);
-  // Copy input ternary from after the last lsb_x.
-  std::copy(tern.cbegin() + lsb_xs, tern.cend(), vec.begin());
-
-  Bits high_lsb = Bits::AllOnes(lsb_xs);
-  Bits low_lsb(lsb_xs);
-  for (const Bits& v : ternary_ops::AllBitsValues(vec)) {
-    is.AddInterval(
-        Interval::Closed(bits_ops::UMax(lb, bits_ops::Concat({v, low_lsb})),
-                         bits_ops::UMin(ub, bits_ops::Concat({v, high_lsb}))));
+  std::vector<Interval> intervals;
+  intervals.reserve(uint64_t{1} << x_locations.size());
+  Bits lsbs_low(lsb_xs);
+  Bits lsbs_high = Bits::AllOnes(lsb_xs);
+  for (const Bits& v : ternary_ops::AllBitsValues(prefix)) {
+    // Since prefix's LSB is known (see above), the intervals we create here
+    // will never abut.
+    intervals.push_back(
+        Interval::Closed(bits_ops::UMax(lb, bits_ops::Concat({v, lsbs_low})),
+                         bits_ops::UMin(ub, bits_ops::Concat({v, lsbs_high}))));
   }
-  is.Normalize();
-  return is;
+  return IntervalSet::UnsafeFromNormalized(tern.size(), std::move(intervals));
 }
 
 bool CoversTernary(const Interval& interval, TernarySpan ternary) {
@@ -1509,7 +1517,7 @@ IntervalSet Xor(const IntervalSet& a, const IntervalSet& b) {
 IntervalSet AndReduce(const IntervalSet& a) {
   if (a.IsEmpty()) {
     // If the input is empty, so is the output.
-    return IntervalSet(a.BitCount());
+    return IntervalSet(/*bit_count=*/1);
   }
   // Unless the intervals cover max, the and_reduce of the input must be 0.
   if (!a.CoversMax()) {
@@ -1526,7 +1534,7 @@ IntervalSet AndReduce(const IntervalSet& a) {
 IntervalSet OrReduce(const IntervalSet& a) {
   if (a.IsEmpty()) {
     // If the input is empty, so is the output.
-    return IntervalSet(a.BitCount());
+    return IntervalSet(/*bit_count=*/1);
   }
   // Unless the intervals cover 0, the or_reduce of the input must be 1.
   if (!a.CoversZero()) {
@@ -1542,7 +1550,7 @@ IntervalSet OrReduce(const IntervalSet& a) {
 IntervalSet XorReduce(const IntervalSet& a) {
   if (a.IsEmpty()) {
     // If the input is empty, so is the output.
-    return IntervalSet(a.BitCount());
+    return IntervalSet(/*bit_count=*/1);
   }
   // XorReduce determines the parity of the number of 1s in a bitstring.
   // Incrementing a bitstring always outputs in a bitstring with a different
@@ -1594,7 +1602,7 @@ IntervalSet ULt(const IntervalSet& a, const IntervalSet& b) {
   CHECK_EQ(a.BitCount(), b.BitCount());
   if (a.IsEmpty() || b.IsEmpty()) {
     // If the input is empty, so is the output.
-    return IntervalSet(a.BitCount());
+    return IntervalSet(/*bit_count=*/1);
   }
   if (a.IsPrecise() && a.GetPreciseValue() == Bits::AllOnes(a.BitCount())) {
     // If a is all ones, then it is not less than any value.
@@ -1623,7 +1631,7 @@ IntervalSet SLt(const IntervalSet& a, const IntervalSet& b) {
   CHECK_EQ(a.BitCount(), b.BitCount());
   if (a.IsEmpty() || b.IsEmpty()) {
     // If the input is empty, so is the output.
-    return IntervalSet(a.BitCount());
+    return IntervalSet(/*bit_count=*/1);
   }
   CHECK(a.IsNormalized());
   CHECK(b.IsNormalized());

xls/ir/interval_ops_test.cc

@@ -216,6 +216,12 @@ TEST(IntervalOpsTest, FromTernarySegmentsExtended) {
             FromRanges({{0b10101000, 0b11111101}}, 8));
 }
 
+TEST(IntervalOpsTest, FromTernaryAdjacentUnknownBits) {
+  // Test case where advancing lsb_xs encounters adjacent unknown bits.
+  EXPECT_EQ(FromTernaryString("0b0XX0XX", /*max_unknown_bits=*/1),
+            FromRanges({{0b000000, 0b011011}}, 6));
+}
+
 TEST(IntervalOpsTest, ExactResultsForSmallRanges) {
   // Only 8 possible multiplies so try them all.
   IntervalSet lhs = FromRanges({{1234, 1235}}, 64);
@@ -1507,6 +1513,45 @@ FUZZ_TEST(IntervalOpsTest, OneHotZ3Fuzz)
     .WithDomains(IntervalDomain(8),
                  fuzztest::ElementOf({LsbOrMsb::kLsb, LsbOrMsb::kMsb}));
 
+TEST(IntervalOpsTest, EmptyAndReduce) {
+  EXPECT_EQ(AndReduce(IntervalSet(/*bit_count=*/5)), FromRanges({}, 1));
+}
+
+void AndReduceZ3Fuzz(absl::Span<std::pair<int64_t, int64_t> const> lhs) {
+  UnaryOpFuzz(
+      "and_reduce",
+      [&](FunctionBuilder& fb, BValue l) { return fb.AndReduce(l); },
+      [&](const auto& l) { return AndReduce(l); }, lhs,
+      /*bits=*/8);
+}
+FUZZ_TEST(IntervalOpsTest, AndReduceZ3Fuzz).WithDomains(IntervalDomain(8));
+
+TEST(IntervalOpsTest, EmptyOrReduce) {
+  EXPECT_EQ(OrReduce(IntervalSet(/*bit_count=*/5)), FromRanges({}, 1));
+}
+
+void OrReduceZ3Fuzz(absl::Span<std::pair<int64_t, int64_t> const> lhs) {
+  UnaryOpFuzz(
+      "or_reduce",
+      [&](FunctionBuilder& fb, BValue l) { return fb.OrReduce(l); },
+      [&](const auto& l) { return OrReduce(l); }, lhs,
+      /*bits=*/8);
+}
+FUZZ_TEST(IntervalOpsTest, OrReduceZ3Fuzz).WithDomains(IntervalDomain(8));
+
+TEST(IntervalOpsTest, EmptyXorReduce) {
+  EXPECT_EQ(XorReduce(IntervalSet(/*bit_count=*/5)), FromRanges({}, 1));
+}
+
+void XorReduceZ3Fuzz(absl::Span<std::pair<int64_t, int64_t> const> lhs) {
+  UnaryOpFuzz(
+      "xor_reduce",
+      [&](FunctionBuilder& fb, BValue l) { return fb.XorReduce(l); },
+      [&](const auto& l) { return XorReduce(l); }, lhs,
+      /*bits=*/8);
+}
+FUZZ_TEST(IntervalOpsTest, XorReduceZ3Fuzz).WithDomains(IntervalDomain(8));
+
 TEST(IntervalOpsTest, ReduceIntervalFragmentation) {
   // Create a set with 10 separate intervals.
   IntervalSet lhs = FromValues({0, 2, 4, 6, 8, 10, 12, 14, 16, 18}, 8);

xls/ir/partial_information.cc

@@ -204,6 +204,22 @@ int64_t PartialInformation::KnownLeadingSignBits() const {
   return 1 + bit_count_ - interval_ops::MinimumSignedBitCount(*range_);
 }
 
+int64_t PartialInformation::KnownLeadingBits() const {
+  if (IsImpossible()) {
+    return 0;
+  }
+  if (IsUnconstrained()) {
+    return 0;
+  }
+  if (!ternary_) {
+    return 0;
+  }
+  return std::find_if(
+             ternary_->rbegin(), ternary_->rend(),
+             [](TernaryValue v) { return v == TernaryValue::kUnknown; }) -
+         ternary_->rbegin();
+}
+
 int64_t PartialInformation::MaxPopCount() const {
   if (IsImpossible()) {
     return 0;

xls/ir/partial_information.h

@@ -140,6 +140,9 @@ class PartialInformation {
   // Gets the number of leading sign bits known.
   int64_t KnownLeadingSignBits() const;
 
+  // Gets the number of leading (high bit) bits known.
+  int64_t KnownLeadingBits() const;
+
   // Returns an upper bound on the popcount of any value that can satisfy this
   // PartialInformation.
   int64_t MaxPopCount() const;

xls/passes/BUILD

@@ -946,6 +946,7 @@ cc_library(
     srcs = ["back_propagate_range_analysis.cc"],
     hdrs = ["back_propagate_range_analysis.h"],
     deps = [
+        ":query_engine",
         ":range_query_engine",
         "//xls/common/status:ret_check",
         "//xls/common/status:status_macros",
@@ -1816,6 +1817,7 @@ cc_library(
         "//xls/data_structures:leaf_type_tree",
         "//xls/ir",
         "//xls/ir:bits",
+        "//xls/ir:bits_ops",
         "//xls/ir:interval",
         "//xls/ir:interval_ops",
         "//xls/ir:interval_set",
@@ -2470,8 +2472,10 @@ xls_pass(
         "//xls/ir",
         "//xls/ir:bits",
         "//xls/ir:interval_ops",
+        "//xls/ir:interval_set",
         "//xls/ir:node_util",
         "//xls/ir:op",
+        "//xls/ir:partial_info",
         "//xls/ir:state_element",
         "//xls/ir:ternary",
         "//xls/ir:type",
@@ -4446,9 +4450,6 @@ cc_library(
         ":partial_info_query_engine",
         ":predicate_state",
         ":query_engine",
-        ":range_query_engine",
-        ":ternary_query_engine",
-        ":union_query_engine",
         "//xls/common/status:ret_check",
         "//xls/common/status:status_macros",
         "//xls/data_structures:leaf_type_tree",
@@ -4457,14 +4458,12 @@ cc_library(
         "//xls/ir:bits",
         "//xls/ir:bits_ops",
         "//xls/ir:interval",
-        "//xls/ir:interval_ops",
         "//xls/ir:interval_set",
         "//xls/ir:op",
         "//xls/ir:partial_info",
         "//xls/ir:state_element",
         "//xls/ir:ternary",
         "//xls/ir:type",
-        "//xls/ir:value",
         "//xls/ir:value_utils",
         "@abseil-cpp//absl/algorithm:container",
         "@abseil-cpp//absl/container:btree",
@@ -4476,6 +4475,7 @@ cc_library(
         "@abseil-cpp//absl/status:statusor",
         "@abseil-cpp//absl/strings",
         "@abseil-cpp//absl/strings:str_format",
+        "@abseil-cpp//absl/types:optional_ref",
         "@abseil-cpp//absl/types:span",
     ],
 )
@@ -4488,12 +4488,14 @@ cc_test(
         ":range_query_engine",
         "//xls/common:xls_gunit_main",
         "//xls/common/status:matchers",
+        "//xls/data_structures:leaf_type_tree",
         "//xls/ir",
         "//xls/ir:bits",
         "//xls/ir:channel",
         "//xls/ir:channel_ops",
         "//xls/ir:function_builder",
         "//xls/ir:ir_test_base",
+        "//xls/ir:ternary",
         "//xls/ir:value",
         "@abseil-cpp//absl/log:check",
         "@abseil-cpp//absl/strings:str_format",

xls/passes/back_propagate_range_analysis.cc

@@ -41,6 +41,7 @@
 #include "xls/ir/op.h"
 #include "xls/ir/topo_sort.h"
 #include "xls/ir/type.h"
+#include "xls/passes/query_engine.h"
 #include "xls/passes/range_query_engine.h"
 
 namespace xls {
@@ -54,7 +55,7 @@ namespace {
 // based on that node and (2) the inputs to that node have updated information.
 class BackPropagate : public DfsVisitorWithDefault {
  public:
-  explicit BackPropagate(const RangeQueryEngine& query_engine,
+  explicit BackPropagate(const QueryEngine& query_engine,
                          absl::flat_hash_map<Node*, IntervalSet> givens)
       : query_engine_(query_engine), result_(std::move(givens)) {
     for (const auto& [node, _] : result_) {
@@ -240,11 +241,7 @@ class BackPropagate : public DfsVisitorWithDefault {
     if (result_.contains(node)) {
       return result_[node];
     }
-    if (query_engine_.HasExplicitIntervals(node)) {
-      // Try to avoid allocating LTTs needlessly.
-      return query_engine_.GetIntervalSetTreeView(node)->Get({});
-    }
-    return query_engine_.GetIntervalSetTree(node).Get({});
+    return query_engine_.GetIntervals(node).Get({});
   }
 
   // Merge the given 'new_data' with the already known facts about the given
@@ -254,7 +251,7 @@ class BackPropagate : public DfsVisitorWithDefault {
     XLS_RET_CHECK(node->GetType()->IsBits());
     XLS_RET_CHECK(new_data.IsNormalized());
     if (!result_.contains(node)) {
-      result_[node] = query_engine_.GetIntervalSetTree(node).Get({});
+      result_[node] = GetIntervals(node);
     }
     IntervalSet old_data = std::move(result_[node]);
     result_[node] = IntervalSet::Intersect(old_data, new_data);
@@ -590,7 +587,7 @@ class BackPropagate : public DfsVisitorWithDefault {
   }
 
   // Underlying query-engine providing base ranges.
-  const RangeQueryEngine& query_engine_;
+  const QueryEngine& query_engine_;
   // Set of all givens and any calculated refined ranges.
   absl::flat_hash_map<Node*, IntervalSet> result_;
   // Set of nodes which we have updated data for which might be possible to
@@ -602,7 +599,7 @@ class BackPropagate : public DfsVisitorWithDefault {
 
 absl::StatusOr<absl::flat_hash_map<Node*, IntervalSet>>
 PropagateGivensBackwards(
-    const RangeQueryEngine& engine, FunctionBase* function,
+    const QueryEngine& engine, FunctionBase* function,
     absl::flat_hash_map<Node*, IntervalSet> givens,
     std::optional<absl::Span<Node* const>> reverse_topo_sort) {
   XLS_RET_CHECK(!givens.empty());
@@ -629,7 +626,7 @@ PropagateGivensBackwards(
 }
 
 absl::StatusOr<absl::flat_hash_map<Node*, IntervalSet>>
-PropagateOneGivenBackwards(const RangeQueryEngine& engine, Node* node,
+PropagateOneGivenBackwards(const QueryEngine& engine, Node* node,
                            const Bits& given) {
   return PropagateOneGivenBackwards(engine, node, IntervalSet::Precise(given));
 }

xls/passes/back_propagate_range_analysis.h

@@ -25,7 +25,7 @@
 #include "xls/ir/function_base.h"
 #include "xls/ir/interval_set.h"
 #include "xls/ir/node.h"
-#include "xls/passes/range_query_engine.h"
+#include "xls/passes/query_engine.h"
 
 namespace xls {
 
@@ -40,22 +40,22 @@ namespace xls {
 // another).
 absl::StatusOr<absl::flat_hash_map<Node*, IntervalSet>>
 PropagateGivensBackwards(
-    const RangeQueryEngine& engine, FunctionBase* function,
+    const QueryEngine& engine, FunctionBase* function,
     absl::flat_hash_map<Node*, IntervalSet> given,
     std::optional<absl::Span<Node* const>> reverse_topo_sort = std::nullopt);
 
 // Helper to analyze a function by back-propagating range information.
 //
 // Returns the data extracted from analyzing the given computation.
 inline absl::StatusOr<absl::flat_hash_map<Node*, IntervalSet>>
-PropagateOneGivenBackwards(const RangeQueryEngine& engine, Node* node,
+PropagateOneGivenBackwards(const QueryEngine& engine, Node* node,
                            const IntervalSet& given) {
   return PropagateGivensBackwards(engine, node->function_base(),
                                   {{node, given}});
 }
 
 absl::StatusOr<absl::flat_hash_map<Node*, IntervalSet>>
-PropagateOneGivenBackwards(const RangeQueryEngine& engine, Node* node,
+PropagateOneGivenBackwards(const QueryEngine& engine, Node* node,
                            const Bits& given);
 
 }  // namespace xls

xls/passes/bdd_query_engine.cc

@@ -670,16 +670,12 @@ std::unique_ptr<QueryEngine> BddQueryEngine::SpecializeGiven(
           value_knowledge.intervals->AsView(),
           [&](Type*, const IntervalSet& intervals,
               absl::Span<const int64_t> tree_index) -> absl::Status {
-            std::vector<SaturatingBddNodeIndex> bits;
-            bits.reserve(intervals.BitCount());
-            for (int64_t i = 0; i < intervals.BitCount(); ++i) {
-              std::optional<BddNodeIndex> bit =
-                  GetBddNode(TreeBitLocation(node, i, tree_index));
-              if (!bit.has_value()) {
-                return absl::OkStatus();
-              }
-              bits.push_back(*bit);
+            std::optional<SharedBddTree> info = GetInfo(node);
+            if (!info.has_value()) {
+              return absl::OkStatus();
             }
+            absl::Span<const SaturatingBddNodeIndex> bits =
+                info->Get(tree_index);
 
             SaturatingBddNodeVector in_interval_checks;
             for (const Interval& interval : intervals.Intervals()) {