add optimization of sampling with awareness of memory allocation

Qiskit · Jul 14, 2020 · 341cd8f · 341cd8f
1 parent 9b5c146
commit 341cd8f
Showing 1 changed file with 43 additions and 16 deletions.
diff --git a/src/simulators/statevector/qubitvector.hpp b/src/simulators/statevector/qubitvector.hpp
@@ -1980,30 +1980,57 @@ reg_t QubitVector<data_t>::sample_measure(const std::vector<double> &rnds) const
       }
     } // end omp parallel
 
-    #pragma omp parallel if (num_qubits_ > omp_threshold_ && omp_threads_ > 1) num_threads(omp_threads_)
-    {
-      #pragma omp for
-      for (int_t i = 0; i < SHOTS; ++i) {
-        double rnd = rnds[i];
-        double p = .0;
-        int_t sample = 0;
-        for (uint_t j = 0; j < idxs.size(); ++j) {
-          if (rnd < (p + idxs[j])) {
-            break;
-          }
-          p += idxs[j];
-          sample += loop;
+    // accumulate indices
+    for (uint_t i = 1; i < INDEX_END; ++i)
+      idxs[i] += idxs[i - 1];
+
+    // reduce rounding error
+    double correction = 1.0 / idxs[INDEX_END - 1];
+    for (int_t i = 1; i < INDEX_END - 1; ++i)
+      idxs[i] *= correction;
+
+    idxs[INDEX_END - 1] = 1.0;
+
+    auto rnds_ = rnds;
+
+    std::sort(rnds_.begin(), rnds_.end());
+
+    #pragma omp parallel for if (num_qubits_ > omp_threshold_ && omp_threads_ > 1) num_threads(omp_threads_)
+    for (int_t i = 0; i < INDEX_END; ++i) {
+      uint_t start_sample_idx = 0;
+      if (i != 0) {
+        for (uint_t sample_idx = 0; sample_idx < SHOTS; ++sample_idx) {
+          if (rnds_[sample_idx] < idxs[i - 1])
+            continue;
+          start_sample_idx = sample_idx;
+          break;
         }
+      }
 
-        for (; sample < END - 1; ++sample) {
+      uint_t end_sample_idx = SHOTS;
+      if (i != (INDEX_END - 1)) {
+        for (uint_t sample_idx = start_sample_idx; sample_idx < SHOTS; ++sample_idx) {
+          if (rnds_[sample_idx] < idxs[i])
+            continue;
+          end_sample_idx = sample_idx;
+          break;
+        }
+      }
+
+      auto sample = loop * i;
+      double p = (i == 0) ? 0.0 : idxs[i - 1];
+
+      for (uint_t sample_idx = start_sample_idx; sample_idx < end_sample_idx; ++sample_idx) {
+        auto rnd = rnds_[sample_idx];
+        for (; sample < loop * (i + 1) - 1; ++sample) {
           p += probability(sample);
           if (rnd < p){
             break;
           }
         }
-        samples[i] = sample;
+        samples[sample_idx] = sample;
       }
-    } // end omp parallel
+    }
   }
   return samples;
 }