[ir] Optimize interval_ops::FromTernary

xls/data_structures/inline_bitmap.h

@@ -358,14 +358,14 @@ class InlineBitmap {
   void Overwrite(const InlineBitmap& other, int64_t cnt, int64_t w_offset = 0,
                  int64_t r_offset = 0);
 
+  static constexpr int64_t kWordBits = 64;
+  static constexpr int64_t kWordBytes = 8;
+
  private:
   FRIEND_TEST(InlineBitmapTest, MaskForWord);
   friend uint64_t GetWordBitsAtForTest(const InlineBitmap& ib,
                                        int64_t bit_offset);
 
-  static constexpr int64_t kWordBits = 64;
-  static constexpr int64_t kWordBytes = 8;
-
   // Gets the kWordBits bits following bit_offset with 'Get(bit_offset)' being
   // the LSB, Get(bit_offset + 1) being the next lsb etc.
   int64_t GetWordBitsAt(int64_t bit_offset) const;

xls/ir/bits_ops.cc

@@ -707,25 +707,26 @@ Bits LongestCommonPrefixMSB(absl::Span<const Bits> bits_span) {
     CHECK_EQ(bits.bit_count(), input_size);
   }
 
-  int64_t first_difference = -1;
-  std::vector<Bits::Iterator> iterators;
-  iterators.reserve(bits_span.size());
-  for (const Bits& bits : bits_span) {
-    iterators.push_back(bits.end() - 1);
-  }
-  for (; iterators[0] >= bits_span[0].begin(); --iterators[0]) {
-    for (auto& it : absl::MakeSpan(iterators).subspan(1)) {
-      if (*it-- != *iterators[0]) {
-        first_difference = std::distance(bits_span[0].begin(), iterators[0]);
-        break;
-      }
+  // Find the most-significant bit where any of the bitmaps differ, working one
+  // word at a time from MSB to LSB.
+  const int64_t word_count = bits_span[0].bitmap().word_count();
+  for (int64_t i = word_count - 1; i >= 0; --i) {
+    uint64_t diff = 0;
+    // Since the words are padded with zeros on the left, we can just XOR each
+    // with the first to find differing bits; the most significant 1 in the
+    // result is the first bit after the common prefix.
+    uint64_t base_word = bits_span[0].bitmap().GetWord(i);
+    for (auto it = bits_span.begin() + 1; it != bits_span.end(); ++it) {
+      diff |= base_word ^ it->bitmap().GetWord(i);
     }
-    if (first_difference != -1) {
-      break;
+    if (diff != 0) {
+      int64_t common_prefix_start =
+          1 + xls::FloorOfLog2(diff) + i * InlineBitmap::kWordBits;
+      return bits_span[0].Slice(common_prefix_start,
+                                input_size - common_prefix_start);
     }
   }
-  return bits_span[0].Slice(first_difference + 1,
-                            input_size - first_difference - 1);
+  return bits_span[0];
 }
 
 }  // namespace bits_ops

xls/ir/interval_ops.cc

@@ -143,31 +143,36 @@ IntervalSet FromTernary(TernarySpan tern, int64_t max_interval_bits) {
     x_locations.pop_front();
   }
 
-  IntervalSet is(tern.size());
-  if (x_locations.empty()) {
+  // If the leading component has any trailing unknown bits, we can just extend
+  // the unknown region to include them.
+  while (lsb_xs < tern.size() && ternary_ops::IsUnknown(tern[lsb_xs])) {
+    ++lsb_xs;
+  }
+
+  // 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.
-    Bits high_bits = ternary_ops::ToKnownBitsValues(tern.subspan(lsb_xs));
-    is.AddInterval(Interval::Closed(
+    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) {
@@ -1571,7 +1576,7 @@ IntervalSet Eq(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() && b.IsPrecise()) {

xls/ir/partial_information.cc

@@ -645,6 +645,20 @@ PartialInformation& PartialInformation::Shra(const PartialInformation& other) {
 
 PartialInformation& PartialInformation::JoinWith(
     const PartialInformation& other) {
+  CHECK_EQ(bit_count_, other.bit_count_);
+
+  if (IsImpossible() || other.IsUnconstrained()) {
+    return *this;
+  }
+  if (other.IsImpossible()) {
+    MarkImpossible();
+    return *this;
+  }
+  if (IsUnconstrained()) {
+    *this = other;
+    return *this;
+  }
+
   if (ternary_.has_value()) {
     if (other.ternary_.has_value() &&
         !ternary_ops::TryUpdateWithUnion(*ternary_, *other.ternary_)) {
@@ -676,6 +690,21 @@ PartialInformation& PartialInformation::JoinWith(
 
 PartialInformation& PartialInformation::MeetWith(
     const PartialInformation& other) {
+  CHECK_EQ(bit_count_, other.bit_count_);
+
+  if (IsUnconstrained() || other.IsImpossible()) {
+    return *this;
+  }
+  if (other.IsUnconstrained()) {
+    ternary_ = std::nullopt;
+    range_ = std::nullopt;
+    return *this;
+  }
+  if (IsImpossible()) {
+    *this = other;
+    return *this;
+  }
+
   if (ternary_.has_value() && other.ternary_.has_value()) {
     ternary_ops::UpdateWithIntersection(*ternary_, *other.ternary_);
   } else {
@@ -709,6 +738,69 @@ std::ostream& operator<<(std::ostream& os,
 }
 
 void PartialInformation::ReconcileInformation() {
+  // If the range is empty, just standardize on the "impossible" state.
+  if (range_.has_value() && range_->IsEmpty()) {
+    MarkImpossible();
+    return;
+  }
+
+  // Explicitly remove fully-unconstrained "information".
+  if (range_.has_value() && range_->IsMaximal()) {
+    range_ = std::nullopt;
+  }
+  if (ternary_.has_value() && ternary_ops::AllUnknown(*ternary_)) {
+    ternary_ = std::nullopt;
+  }
+
+  // If both are unset, we're done; this is the "unconstrained" state.
+  if (!ternary_.has_value() && !range_.has_value()) {
+    return;
+  }
+
+  // If only one of the two is set, then we can just set the other one from it
+  // directly, with no iteration needed.
+  if (!range_.has_value()) {
+    range_ = interval_ops::FromTernary(*ternary_);
+    if (range_.has_value() && range_->IsMaximal()) {
+      range_ = std::nullopt;
+    }
+    return;
+  }
+  if (!ternary_.has_value()) {
+    ternary_ = interval_ops::ExtractTernaryVector(*range_);
+    if (ternary_.has_value() && ternary_ops::AllUnknown(*ternary_)) {
+      ternary_ = std::nullopt;
+    }
+    return;
+  }
+
+  // From here on out, we know that both `range_` and `ternary_` are set, and
+  // both will only be made more specific as we go.
+
+  // If our information is already contradictory, this is just impossible.
+  if (!interval_ops::CoversTernary(*range_, *ternary_)) {
+    MarkImpossible();
+    return;
+  }
+  // Otherwise, we should never reach an impossible state, as we're only
+  // transferring information from one representation to the other.
+
+  // If either side has a precise value, use that to set the other side; we know
+  // they're compatible, so we don't need to check again.
+  if (std::optional<Bits> precise_value = range_->GetPreciseValue();
+      precise_value.has_value()) {
+    ternary_ = ternary_ops::BitsToTernary(*precise_value);
+    return;
+  }
+  if (ternary_ops::IsFullyKnown(*ternary_)) {
+    range_ = IntervalSet::Precise(ternary_ops::ToKnownBitsValues(*ternary_));
+    return;
+  }
+
+  // Repeatedly transfer information from one representation to the other until
+  // no more information can be transferred. Since they're compatible and
+  // neither is unconstrained, this will always converge to a non-empty but not
+  // unconstrained state.
   bool changed = true;
   int64_t iterations = 0;
   while (changed) {
@@ -730,52 +822,17 @@ void PartialInformation::ReconcileInformation() {
       return;
     }
     changed = false;
-    if (range_.has_value() && range_->IsMaximal()) {
-      range_ = std::nullopt;
-    }
-    if (ternary_.has_value() && ternary_ops::AllUnknown(*ternary_)) {
-      ternary_ = std::nullopt;
-    }
 
-    if (range_.has_value() && range_->IsEmpty()) {
-      // Standardize the "impossible" state.
-      MarkImpossible();
-      return;
-    }
+    // Transfer as much information as we can into the ternary.
+    TernaryVector range_ternary = interval_ops::ExtractTernaryVector(*range_);
+    CHECK(ternary_ops::TryUpdateWithUnion(*ternary_, range_ternary, &changed));
 
-    // Check whether `ternary_` and `range_` are in conflict before we go any
-    // further; if they are, mark the combination impossible.
-    if (ternary_.has_value() && range_.has_value() &&
-        !interval_ops::CoversTernary(*range_, *ternary_)) {
-      MarkImpossible();
-      return;
-    }
-
-    if (range_.has_value()) {
-      // Transfer as much information as we can into the ternary.
-      TernaryVector range_ternary = interval_ops::ExtractTernaryVector(*range_);
-      if (ternary_.has_value()) {
-        CHECK(ternary_ops::TryUpdateWithUnion(*ternary_, range_ternary,
-                                              &changed));
-      } else if (!ternary_ops::AllUnknown(range_ternary)) {
-        ternary_ = std::move(range_ternary);
-        changed = true;
-      }
-    }
-
-    if (ternary_.has_value()) {
-      // Transfer as much information as we can into the range.
-      IntervalSet ternary_range = interval_ops::FromTernary(*ternary_);
-      if (range_.has_value()) {
-        IntervalSet new_range = IntervalSet::Intersect(ternary_range, *range_);
-        CHECK(!new_range.IsEmpty());
-        changed = changed || (new_range != *range_);
-        range_ = std::move(new_range);
-      } else if (!ternary_range.IsMaximal()) {
-        range_ = std::move(ternary_range);
-        changed = true;
-      }
-    }
+    // Transfer as much information as we can into the range.
+    IntervalSet ternary_range = interval_ops::FromTernary(*ternary_);
+    IntervalSet new_range = IntervalSet::Intersect(ternary_range, *range_);
+    CHECK(!new_range.IsEmpty());
+    changed = changed || (new_range != *range_);
+    range_ = std::move(new_range);
   }
 }