Add support for dense adjacency-matrix transitive closure calculation.

This can be faster for some transitive closure calculations due to not requiring reallocation.

PiperOrigin-RevId: 718528395

xls/data_structures/BUILD

@@ -97,22 +97,28 @@ cc_library(
     name = "transitive_closure",
     hdrs = ["transitive_closure.h"],
     deps = [
+        ":inline_bitmap",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/types:span",
     ],
 )
 
 cc_test(
     name = "transitive_closure_test",
     srcs = ["transitive_closure_test.cc"],
     deps = [
+        ":inline_bitmap",
         ":transitive_closure",
+        "//xls/common:math_util",
         "//xls/common:xls_gunit_main",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/random",
         "@com_google_absl//absl/random:bit_gen_ref",
         "@com_google_absl//absl/random:distributions",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
         "@com_google_benchmark//:benchmark",
         "@com_google_googletest//:gtest",
     ],

xls/data_structures/transitive_closure.h

@@ -15,37 +15,117 @@
 #ifndef XLS_DATA_STRUCTURES_TRANSITIVE_CLOSURE_H_
 #define XLS_DATA_STRUCTURES_TRANSITIVE_CLOSURE_H_
 
+#include <cstdint>
+#include <vector>
+
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
+#include "absl/types/span.h"
+#include "xls/data_structures/inline_bitmap.h"
 
 namespace xls {
 
-template <typename V>
-using HashRelation = absl::flat_hash_map<V, absl::flat_hash_set<V>>;
-
+namespace internal {
 // Compute the transitive closure of a relation.
-template <typename V>
-HashRelation<V> TransitiveClosure(HashRelation<V> relation) {
+template <typename Relation>
+void TransitiveClosure(Relation relation) {
   // Warshall's algorithm; https://cs.winona.edu/lin/cs440/ch08-2.pdf
   // modified in the typical way to avoid unnecessary copies of the expanded
   // relation. It's safe to update the relation as we go, since at each stage k,
   // i relates to k via nodes < k iff i relates to k via nodes <= k, and
   // similarly for k relating to j.
-  for (const auto& [k, from_k] : relation) {
-    for (auto& [i, from_i] : relation) {
+  //
+  // Callbacks are used to avoid having to deal with the complication of
+  // enumerating the values in a consistent way with both the array and map
+  // formulations.
+  relation.ForEachKeyValue([&](const auto& k, const auto& from_k) {
+    relation.ForEachKeyValue([&](const auto& i, auto& from_i) {
       if (i == k) {
         // Updating would be a no-op, so skip it.
-        continue;
+        return;
       }
-      if (from_i.contains(k)) {
+      if (relation.Contains(from_i, k)) {
         // i relates to k (via nodes < k), so:
         //   for any j where k relates to j (via nodes < k),
         //     i relates to j (via nodes <= k).
-        from_i.insert(from_k.begin(), from_k.end());
+        relation.UnionInPlace(from_i, from_k);
       }
+    });
+  });
+}
+
+template <typename V>
+class HashRelation {
+ public:
+  explicit HashRelation(
+      absl::flat_hash_map<V, absl::flat_hash_set<V>>& relation)
+      : relation_(relation) {}
+  template <typename F>
+  void ForEachKeyValue(F f) const {
+    for (auto& [j, from_j] : relation_) {
+      f(j, from_j);
+    }
+  }
+  bool Contains(const absl::flat_hash_set<V>& vs, const V& v) const {
+    return vs.contains(v);
+  }
+  void UnionInPlace(absl::flat_hash_set<V>& i,
+                    const absl::flat_hash_set<V>& k) const {
+    i.insert(k.begin(), k.end());
+  }
+
+ private:
+  absl::flat_hash_map<V, absl::flat_hash_set<V>>& relation_;
+};
+
+class DenseIdRelation {
+ public:
+  explicit DenseIdRelation(absl::Span<InlineBitmap> relation)
+      : relation_(relation) {}
+  template <typename F>
+  void ForEachKeyValue(F f) const {
+    for (int64_t i = 0; i < relation_.size(); ++i) {
+      f(i, relation_[i]);
     }
   }
-  return relation;
+  bool Contains(const InlineBitmap& vs, int64_t v) const { return vs.Get(v); }
+  void UnionInPlace(InlineBitmap& i, const InlineBitmap& k) const {
+    i.Union(k);
+  }
+
+ private:
+  absl::Span<InlineBitmap> relation_;
+};
+
+}  // namespace internal
+
+template <typename V>
+using HashRelation = absl::flat_hash_map<V, absl::flat_hash_set<V>>;
+
+// Compute the transitive closure of a relation represented as an explicit
+// adjacency list.
+template <typename V>
+HashRelation<V> TransitiveClosure(HashRelation<V> v) {
+  internal::TransitiveClosure(internal::HashRelation<V>(v));
+  return v;
+}
+
+// TODO(allight): Using a more efficient bitmap format like croaring might give
+// a speedup here.
+using DenseIdRelation = absl::Span<InlineBitmap>;
+// Compute the transitive closure of a relation represented as a boolean
+// adjacency matrix.
+inline DenseIdRelation TransitiveClosure(DenseIdRelation v) {
+  internal::TransitiveClosure(internal::DenseIdRelation(v));
+  return v;
+}
+
+// Compute the transitive closure of a relation represented as a boolean
+// adjacency matrix.
+inline std::vector<InlineBitmap> TransitiveClosure(
+    std::vector<InlineBitmap> v) {
+  TransitiveClosure(absl::MakeSpan(v));
+  return v;
 }
 
 }  // namespace xls

xls/data_structures/transitive_closure_test.cc

@@ -14,6 +14,7 @@
 
 #include "xls/data_structures/transitive_closure.h"
 
+#include <cstdint>
 #include <random>
 #include <string>
 #include <vector>
@@ -25,11 +26,16 @@
 #include "absl/container/flat_hash_set.h"
 #include "absl/random/bit_gen_ref.h"
 #include "absl/random/distributions.h"
+#include "absl/random/random.h"
 #include "absl/strings/str_cat.h"
+#include "absl/types/span.h"
+#include "xls/common/math_util.h"
+#include "xls/data_structures/inline_bitmap.h"
 
 namespace xls {
 namespace {
 
+using ::testing::ElementsAre;
 using ::testing::UnorderedElementsAre;
 
 using V = std::string;
@@ -49,6 +55,24 @@ TEST(TransitiveClosureTest, Simple) {
   EXPECT_FALSE(tc.contains("qux"));
 }
 
+TEST(TransitiveClosureTest, SimpleDense) {
+  std::vector<InlineBitmap> rel{
+      InlineBitmap::FromBitsLsbIs0({false, false, true, false, false}),
+      InlineBitmap::FromBitsLsbIs0({false, false, false, false, false}),
+      InlineBitmap::FromBitsLsbIs0({false, true, false, false, false}),
+      InlineBitmap::FromBitsLsbIs0({true, false, false, false, false}),
+      InlineBitmap::FromBitsLsbIs0({false, true, false, false, true}),
+  };
+  std::vector<InlineBitmap> tc = TransitiveClosure(rel);
+  EXPECT_THAT(
+      tc, ElementsAre(
+              InlineBitmap::FromBitsLsbIs0({false, true, true, false, false}),
+              InlineBitmap::FromBitsLsbIs0({false, false, false, false, false}),
+              InlineBitmap::FromBitsLsbIs0({false, true, false, false, false}),
+              InlineBitmap::FromBitsLsbIs0({true, true, true, false, false}),
+              InlineBitmap::FromBitsLsbIs0({false, true, false, false, true})));
+}
+
 HashRelation<V> RandomRelation(std::vector<V> nodes, double p,
                                absl::BitGenRef rng) {
   HashRelation<V> rel;
@@ -79,5 +103,36 @@ void BM_RandomRelation(benchmark::State& state) {
 }
 BENCHMARK(BM_RandomRelation)->Range(5, 500);
 
+std::vector<InlineBitmap> RandomDenseRelation(int64_t node_cnt,
+                                              absl::BitGenRef rng) {
+  int64_t byte_cnt = CeilOfRatio(node_cnt, int64_t{8});
+  std::vector<InlineBitmap> res;
+  res.reserve(node_cnt);
+  for (int64_t i = 0; i < node_cnt; ++i) {
+    std::vector<uint8_t> bytes;
+    bytes.reserve(byte_cnt);
+    for (int64_t j = 0; j < byte_cnt; ++j) {
+      bytes.push_back(absl::Uniform<uint8_t>(rng));
+    }
+    res.push_back(InlineBitmap::FromBytes(node_cnt, bytes));
+  }
+  return res;
+}
+
+void BM_RandomDenseRelation(benchmark::State& state) {
+  absl::BitGen rng;
+  for (auto _ : state) {
+    state.PauseTiming();
+    // Avoid the copy being counted, since it is a non-trivial part of the
+    // timing.
+    std::vector<InlineBitmap> rel = RandomDenseRelation(state.range(0), rng);
+    absl::Span<InlineBitmap> span = absl::MakeSpan(rel);
+    state.ResumeTiming();
+    span = TransitiveClosure(span);
+    benchmark::DoNotOptimize(span);
+    benchmark::DoNotOptimize(rel);
+  }
+}
+BENCHMARK(BM_RandomDenseRelation)->Range(500, 10000);
 }  // namespace
 }  // namespace xls