[Optimizer] Use bitset for qubit masks (#220)

* [Optimizer] Use bitset for qubit masks

* code format

Author: xumingkuan

src/quartz/math/bitset.cpp

@@ -0,0 +1,57 @@
+#include "bitset.h"
+
+#include <cassert>
+
+namespace quartz {
+
+Bitset::Bitset(size_t n) { a.assign((n + BLOCK - 1) / BLOCK, 0); }
+
+Bitset::reference::reference(Bitset::value_t *a, int x)
+    : pos(a + (x >> LOGBLOCK)), dig(((value_t)1) << (x & BLOCK1)) {}
+Bitset::reference::operator bool() const { return *pos & dig; }
+bool Bitset::reference::operator~() const { return ~*pos & dig; }
+Bitset::reference &Bitset::reference::operator=(bool x) {
+  if (x)
+    *pos |= dig;
+  else
+    *pos &= MASK ^ dig;
+  return *this;
+}
+Bitset::reference &Bitset::reference::operator=(const Bitset::reference &x) {
+  if (*x.pos & x.dig)
+    *pos |= dig;
+  else
+    *pos &= MASK ^ dig;
+  return *this;
+}
+Bitset::reference &Bitset::reference::flip() {
+  *pos ^= dig;
+  return *this;
+}
+
+void Bitset::flip(int x) { a[x >> LOGBLOCK] ^= ((value_t)1) << (x & BLOCK1); }
+bool Bitset::operator[](int x) const {
+  return (a[x >> LOGBLOCK] >> (x & BLOCK1)) & 1;
+}
+Bitset::reference Bitset::operator[](int x) { return {a.data(), x}; }
+Bitset Bitset::operator^(const Bitset &b) const {
+  assert(a.size() == b.a.size());
+  Bitset result(a.size() * BLOCK);
+  for (int i = 0; i < (int)a.size(); i++) {
+    result.a[i] = a[i] ^ b.a[i];
+  }
+  return result;
+}
+bool Bitset::operator==(const Bitset &b) const {
+  return a == b.a;  // std::vector<>::operator== does the job!
+}
+
+std::size_t BitsetHash::operator()(const Bitset &x) const {
+  std::hash<size_t> hash_fn;
+  std::size_t result = x.a.size();
+  for (auto &val : x.a) {
+    result = result * 17 + hash_fn(val);
+  }
+  return result;
+}
+}  // namespace quartz

src/quartz/math/bitset.h

@@ -0,0 +1,44 @@
+#pragma once
+
+#include <vector>
+
+namespace quartz {
+
+class Bitset;
+struct BitsetHash {
+ public:
+  std::size_t operator()(const Bitset &x) const;
+};
+class Bitset {
+ public:
+  using value_t = unsigned long long;
+  Bitset() = default;
+  Bitset(size_t n);
+  class reference {
+   public:
+    value_t *pos;
+    value_t dig;
+    reference() = default;
+    reference(value_t *a, int x);
+    operator bool() const;
+    bool operator~() const;
+    reference &operator=(bool x);
+    reference &operator=(const reference &x);
+    reference &flip();
+  };
+  void flip(int x);
+  bool operator[](int x) const;
+  reference operator[](int x);
+  Bitset operator^(const Bitset &b) const;
+  bool operator==(const Bitset &b) const;
+  friend std::size_t BitsetHash::operator()(const Bitset &x) const;
+
+ private:
+  static const int LOGBLOCK = 6;
+  static const int BLOCK = (1 << LOGBLOCK);
+  static const int BLOCK1 = BLOCK - 1;
+  static const value_t MASK = ((value_t)-1);
+  std::vector<value_t> a;
+};
+
+}  // namespace quartz

src/quartz/tasograph/tasograph.cpp

@@ -1,5 +1,6 @@
 #include "tasograph.h"
 
+#include "quartz/math/bitset.h"
 #include "substitution.h"
 
 #include <cassert>
@@ -994,7 +995,7 @@ void Graph::rotation_merging(GateType target_rotation) {
     assert(false);
   }
   // Step 1: calculate the bitmask of each operator
-  std::unordered_map<Pos, uint64_t, PosHash> bitmasks;
+  std::unordered_map<Pos, QubitMaskType, PosHash> bitmasks;
   std::unordered_map<Pos, int, PosHash> pos_to_qubits;
   std::queue<Op> todos;
 
@@ -1003,7 +1004,10 @@ void Graph::rotation_merging(GateType target_rotation) {
     if (it.first.ptr->tp == GateType::input_qubit) {
       todos.push(it.first);
       int qubit_idx = input_qubit_op_2_qubit_idx[it.first];
-      bitmasks[Pos(it.first, 0)] = 1 << qubit_idx;
+      if (bitmasks.count(Pos(it.first, 0)) == 0) {
+        bitmasks[Pos(it.first, 0)] = Bitset(get_num_qubits());
+      }
+      bitmasks[Pos(it.first, 0)][qubit_idx] = true;
       pos_to_qubits[Pos(it.first, 0)] = qubit_idx;
     } else if (it.first.ptr->tp == GateType::input_param) {
       todos.push(it.first);
@@ -1106,7 +1110,8 @@ void Graph::rotation_merging(GateType target_rotation) {
 
     // Step 4: merge rotations with the same bitmasks on the same qubit
     std::unordered_map<
-        int, std::unordered_map<uint64_t, std::unordered_set<Pos, PosHash>>>
+        int, std::unordered_map<QubitMaskType, std::unordered_set<Pos, PosHash>,
+                                QubitMaskHash>>
         qubit_2_bm_2_pos;
     for (const auto &pos : covered) {
       if (pos.op.ptr->tp == GateType::cx) {

src/quartz/utils/utils.h

@@ -1,5 +1,7 @@
 #pragma once
 
+#include "quartz/math/bitset.h"
+
 #include <complex>
 #include <filesystem>
 
@@ -22,6 +24,8 @@ using CircuitSeqHashType = unsigned long long;
 using PhaseShiftIdType = int;
 using EquivalenceHashType = std::pair<unsigned long long, int>;
 using InputParamMaskType = unsigned long long;
+using QubitMaskType = quartz::Bitset;  // for rotation merging
+using QubitMaskHash = quartz::BitsetHash;
 
 using namespace std::complex_literals;  // so that we can write stuff like 1.0i