Add cache-friendly qubit index remapping

apps/qsim_base.cc

@@ -24,13 +24,14 @@
 #include "../lib/fuser_mqubit.h"
 #include "../lib/io_file.h"
 #include "../lib/operation.h"
+#include "../lib/qubit_remap.h"
 #include "../lib/run_qsim.h"
 #include "../lib/simmux.h"
 #include "../lib/util_cpu.h"
 
 constexpr char usage[] = "usage:\n  ./qsim_base -c circuit -d maxtime "
                          "-s seed -t threads -f max_fused_size "
-                         "-v verbosity -z\n";
+                         "-v verbosity -r -z\n";
 
 struct Options {
   std::string circuit_file;
@@ -39,6 +40,7 @@ struct Options {
   unsigned num_threads = 1;
   unsigned max_fused_size = 2;
   unsigned verbosity = 0;
+  bool cache_local_remap = false;
   bool denormals_are_zeros = false;
 };
 
@@ -47,7 +49,7 @@ Options GetOptions(int argc, char* argv[]) {
 
   int k;
 
-  while ((k = getopt(argc, argv, "c:d:s:t:f:v:z")) != -1) {
+  while ((k = getopt(argc, argv, "c:d:s:t:f:v:rz")) != -1) {
     switch (k) {
       case 'c':
         opt.circuit_file = optarg;
@@ -67,6 +69,9 @@ Options GetOptions(int argc, char* argv[]) {
       case 'v':
         opt.verbosity = std::atoi(optarg);
         break;
+      case 'r':
+        opt.cache_local_remap = true;
+        break;
       case 'z':
         opt.denormals_are_zeros = true;
         break;
@@ -91,7 +96,8 @@ bool ValidateOptions(const Options& opt) {
 
 template <typename StateSpace, typename State>
 void PrintAmplitudes(
-    unsigned num_qubits, const StateSpace& state_space, const State& state) {
+    unsigned num_qubits, const StateSpace& state_space, const State& state,
+    const qsim::qubit_remap::QubitMap& logical_to_physical = {}) {
   static constexpr char const* bits[8] = {
     "000", "001", "010", "011", "100", "101", "110", "111",
   };
@@ -100,7 +106,9 @@ void PrintAmplitudes(
   unsigned s = 3 - std::min(unsigned{3}, num_qubits);
 
   for (uint64_t i = 0; i < size; ++i) {
-    auto a = state_space.GetAmpl(state, i);
+    uint64_t physical_i =
+        qsim::qubit_remap::LogicalToPhysicalIndex(i, logical_to_physical);
+    auto a = state_space.GetAmpl(state, physical_i);
     qsim::IO::messagef("%s:%16.8g%16.8g%16.8g\n",
                        bits[i] + s, std::real(a), std::imag(a), std::norm(a));
   }
@@ -141,29 +149,37 @@ int main(int argc, char* argv[]) {
     unsigned num_threads;
   };
 
-  using Simulator = Factory::Simulator;
-  using StateSpace = Simulator::StateSpace;
+  using StateSpace = Factory::StateSpace;
   using State = StateSpace::State;
   using Fuser = MultiQubitGateFuser<IO>;
   using Runner = QSimRunner<IO, Fuser, Factory>;
 
-  StateSpace state_space = Factory(opt.num_threads).CreateStateSpace();
+  Factory factory(opt.num_threads);
+  StateSpace state_space = factory.CreateStateSpace();
   State state = state_space.Create(circuit.num_qubits);
 
   if (state_space.IsNull(state)) {
     IO::errorf("not enough memory: is the number of qubits too large?\n");
     return 1;
   }
 
+  qubit_remap::QubitMap logical_to_physical;
+
   state_space.SetStateZero(state);
 
   Runner::Parameter param;
   param.max_fused_size = opt.max_fused_size;
   param.seed = opt.seed;
   param.verbosity = opt.verbosity;
+  param.cache_local_remap = opt.cache_local_remap;
+
+  auto simulator = factory.CreateSimulator();
+  bool ok = Runner::Run(param, circuit, state_space, simulator, state,
+                        logical_to_physical);
 
-  if (Runner::Run(param, Factory(opt.num_threads), circuit, state)) {
-    PrintAmplitudes(circuit.num_qubits, state_space, state);
+  if (ok) {
+    PrintAmplitudes(circuit.num_qubits, state_space, state,
+                    logical_to_physical);
   }
 
   return 0;

lib/BUILD

@@ -50,6 +50,7 @@ cc_library(
         "operation_base.h",
         "parfor.h",
         "qtrajectory.h",
+        "qubit_remap.h",
         "run_qsim.h",
         "run_qsimh.h",
         "seqfor.h",
@@ -127,6 +128,7 @@ cuda_library(
         "operation_base.h",
         "parfor.h",
         "qtrajectory.h",
+        "qubit_remap.h",
         "run_qsim.h",
         "run_qsimh.h",
         "seqfor.h",
@@ -274,6 +276,7 @@ cc_library(
         "operation.h",
         "operation_base.h",
         "parfor.h",
+        "qubit_remap.h",
         "run_qsim.h",
         "seqfor.h",
         "simmux.h",
@@ -515,6 +518,18 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "qubit_remap",
+    hdrs = ["qubit_remap.h"],
+    deps = [
+        ":circuit",
+        ":fuser",
+        ":gate",
+        ":matrix",
+        ":operation",
+    ],
+)
+
 cc_library(
     name = "fuser_basic",
     hdrs = ["fuser_basic.h"],
@@ -560,6 +575,7 @@ cc_library(
         ":gate",
         ":gate_appl",
         ":operation_base",
+        ":qubit_remap",
         ":util",
     ],
 )

lib/fuser.h

@@ -15,6 +15,7 @@
 #ifndef FUSER_H_
 #define FUSER_H_
 
+#include <algorithm>
 #include <cstdint>
 #include <vector>
 
@@ -178,6 +179,58 @@ inline void CalculateFusedMatrix(FusedGate<FP>& gate) {
   }
 }
 
+/**
+ * Multiplies component gate matrices for a range of fused gates.
+ * @param gbeg, gend The iterator range [gbeg, gend) of fused gates.
+ */
+template <typename Iterator>
+inline void CalculateFusedMatrices(Iterator gbeg, Iterator gend) {
+  for (auto g = gbeg; g != gend; ++g) {
+    if (!g->ParentIsDecomposed()) {
+      CalculateFusedMatrix(*g);
+    }
+  }
+}
+
+/**
+ * Multiplies component gate matrices for a vector of fused gates.
+ * @param gates The vector of fused gates.
+ */
+template <typename FusedGate>
+inline void CalculateFusedMatrices(std::vector<FusedGate>& gates) {
+  CalculateFusedMatrices(gates.begin(), gates.end());
+}
+
+/**
+ * Rebuilds fused-gate qubit lists and matrices from component gates.
+ * @param fused_gates The vector of fused gates to rebuild.
+ */
+template <typename FusedGate>
+inline void RebuildFusedGates(std::vector<FusedGate>& fused_gates) {
+  for (auto& op : fused_gates) {
+    auto* fused_gate = OpGetAlternative<std::variant_alternative_t<0, FusedGate>>(
+        op);
+    if (fused_gate == nullptr) {
+      continue;
+    }
+
+    auto& qubits = fused_gate->qubits;
+    qubits.clear();
+
+    for (const auto& gate : fused_gate->gates) {
+      const auto& gate_qubits = OpQubits(gate);
+      qubits.insert(qubits.end(), gate_qubits.begin(), gate_qubits.end());
+    }
+
+    std::sort(qubits.begin(), qubits.end());
+    qubits.erase(std::unique(qubits.begin(), qubits.end()), qubits.end());
+
+    if (!fused_gate->ParentIsDecomposed()) {
+      CalculateFusedMatrix(*fused_gate);
+    }
+  }
+}
+
 }  // namespace qsim
 
 #endif  // FUSER_H_