Update dominator pass to use bit vectors

Result is >100x speedup in some cases

Author: arshajii

codon/cir/analyze/dataflow/cfg.h

@@ -330,6 +330,9 @@ class CFGraph {
   /// Constructs a control-flow graph.
   explicit CFGraph(const BodiedFunc *f);
 
+  /// @return number of blocks in this CFG
+  auto size() const { return blocks.size(); }
+
   /// @return an iterator to the first block
   auto begin() { return util::raw_ptr_adaptor(blocks.begin()); }
   /// @return an iterator beyond the last block

codon/cir/analyze/dataflow/dominator.cpp

@@ -2,42 +2,81 @@
 
 #include "dominator.h"
 
+#include "codon/cir/llvm/llvm.h"
+
 namespace codon {
 namespace ir {
 namespace analyze {
 namespace dataflow {
 
 void DominatorInspector::analyze() {
-  auto changed = true;
+  const auto numBlocks = cfg->size();
+  std::unordered_map<id_t, size_t> mapping;        // id -> sequential id
+  std::vector<id_t> mappingInv(numBlocks);         // sequential id -> id
+  std::vector<llvm::BitVector> bitvecs(numBlocks); // sequential id -> bitvector
+
+  mapping.reserve(numBlocks);
+  size_t next = 0;
+  for (auto *blk : *cfg) {
+    auto id = blk->getId();
+    mapping[id] = next;
+    mappingInv[next] = id;
+    ++next;
+  }
+
+  // Initialize: all blocks dominate themselves; others start with universal set
+  for (auto *blk : *cfg) {
+    auto id = mapping[blk->getId()];
+    llvm::BitVector &bv = bitvecs[id];
+    bv.resize(numBlocks, true);
+    if (blk == cfg->getEntryBlock()) { // entry block only dominated by itself
+      bv.reset();
+      bv.set(id);
+    }
+  }
+
+  // Run simple domination algorithm
+  bool changed = true;
   while (changed) {
     changed = false;
     for (auto *blk : *cfg) {
-      auto init = false;
-      std::set<id_t> old = sets[blk->getId()];
-      std::set<id_t> working;
+      auto id = mapping[blk->getId()];
+      llvm::BitVector old = bitvecs[id];
+      llvm::BitVector newSet;
+
+      if (blk->predecessors_begin() == blk->predecessors_end())
+        newSet.resize(numBlocks);
 
+      bool first = true;
       for (auto it = blk->predecessors_begin(); it != blk->predecessors_end(); ++it) {
-        auto &predDoms = sets[(*it)->getId()];
-        if (!init) {
-          init = true;
-          working = std::set<id_t>(predDoms.begin(), predDoms.end());
+        auto predId = mapping[(*it)->getId()];
+        const auto &predSet = bitvecs[predId];
+        if (first) {
+          newSet = predSet;
+          first = false;
+        } else {
+          newSet &= predSet;
         }
-
-        std::set<id_t> newWorking;
-        std::set_intersection(working.begin(), working.end(), predDoms.begin(),
-                              predDoms.end(),
-                              std::inserter(newWorking, newWorking.begin()));
-        working = newWorking;
       }
 
-      working.insert(blk->getId());
+      newSet.set(id); // a block always dominates itself
 
-      if (working != old) {
+      if (newSet != old) {
+        bitvecs[id] = newSet;
         changed = true;
-        sets[blk->getId()] = working;
       }
     }
   }
+
+  // Map back to canonical id
+  sets.reserve(numBlocks);
+  for (unsigned id = 0; id < numBlocks; id++) {
+    auto &bv = bitvecs[id];
+    auto &set = sets[mappingInv[id]];
+    for (auto n = bv.find_first(); n != -1; n = bv.find_next(n)) {
+      set.insert(mappingInv[n]);
+    }
+  }
 }
 
 bool DominatorInspector::isDominated(const Value *v, const Value *dominator) {
@@ -52,7 +91,8 @@ bool DominatorInspector::isDominated(const Value *v, const Value *dominator) {
     return dDist <= vDist;
   }
 
-  return sets[vBlock->getId()].find(dBlock->getId()) != sets[vBlock->getId()].end();
+  auto &set = sets[vBlock->getId()];
+  return set.find(dBlock->getId()) != set.end();
 }
 
 const std::string DominatorAnalysis::KEY = "core-analyses-dominator";

codon/cir/analyze/dataflow/dominator.h

@@ -2,8 +2,8 @@
 
 #pragma once
 
-#include <set>
 #include <unordered_map>
+#include <unordered_set>
 #include <utility>
 
 #include "codon/cir/analyze/analysis.h"
@@ -17,7 +17,7 @@ namespace dataflow {
 /// Helper to query the dominators of a particular function.
 class DominatorInspector {
 private:
-  std::unordered_map<id_t, std::set<id_t>> sets;
+  std::unordered_map<id_t, std::unordered_set<id_t>> sets;
   CFGraph *cfg;
 
 public:

codon/cir/llvm/llvm.h

@@ -2,6 +2,7 @@
 
 #pragma once
 
+#include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"