Optimize Decompose()/Dispose()

unitaryfund · Nov 2, 2024 · 7ea9e0e · 7ea9e0e
1 parent 641f639
commit 7ea9e0e
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 13 deletions.
diff --git a/src/common/qengine.cl b/src/common/qengine.cl
@@ -610,6 +610,10 @@ void kernel decomposeprob(global cmplx* stateVec, constant bitCapIntOcl4* bitCap
         }
 
         partStateProb[lcv] = partProb;
+        const real1 prob = partStateProb[lcv];
+        if (prob > REAL1_EPSILON) {
+            partStateAngle[lcv] /= prob;
+        }
     }
 
     for (bitCapIntOcl lcv = ID; lcv < remainderPower; lcv += Nthreads) {
@@ -618,13 +622,6 @@ void kernel decomposeprob(global cmplx* stateVec, constant bitCapIntOcl4* bitCap
             remainderStateAngle[lcv] /= prob;
         }
     }
-
-    for (bitCapIntOcl lcv = ID; lcv < partPower; lcv += Nthreads) {
-        const real1 prob = partStateProb[lcv];
-        if (prob > REAL1_EPSILON) {
-            partStateAngle[lcv] /= prob;
-        }
-    }
 }
 
 void kernel decomposeamp(

diff --git a/src/qengine/state.cpp b/src/qengine/state.cpp
@@ -1200,19 +1200,18 @@ void QEngineCPU::DecomposeDispose(bitLenInt start, bitLenInt length, QEngineCPUP
                 partStateProb[lcv] += nrm;
                 remainderStateAngle[k] += arg(amp) * nrm;
             }
+
+            const real1 prob = partStateProb[lcv];
+            if (prob > amplitudeFloor) {
+                partStateAngle[lcv] /= prob;
+            }
         });
         par_for(0U, remainderPower, [&](const bitCapIntOcl& lcv, const unsigned& cpu) {
             const real1 prob = remainderStateProb[lcv];
             if (prob > amplitudeFloor) {
                 remainderStateAngle[lcv] /= prob;
             }
         });
-        par_for(0U, partPower, [&](const bitCapIntOcl& lcv, const unsigned& cpu) {
-            const real1 prob = partStateProb[lcv];
-            if (prob > amplitudeFloor) {
-                partStateAngle[lcv] /= prob;
-            }
-        });
     } else {
         par_for(0U, remainderPower, [&](const bitCapIntOcl& lcv, const unsigned& cpu) {
             bitCapIntOcl j = lcv & pow2MaskOcl(start);