Refactor NEON gate row helpers

Author: XIN XIE

lib/simulator_neon.h

@@ -317,98 +317,103 @@ class SimulatorNEON final : public SimulatorBase {
     return 0;
   }
 
-  static unsigned SIMDRegisterSize() { return 4; }
+  static constexpr unsigned SIMDRegisterSize() {
+    return sizeof(float32x4_t) / sizeof(float32_t);
+  }
 
  private:
+  struct Complex {
+    float32x4_t re;
+    float32x4_t im;
+  };
+
+  template <unsigned Size, typename GateCoeff>
+  static Complex ApplyGateRow(
+      const float32x4_t* state_re, const float32x4_t* state_im,
+      GateCoeff gate_coeff) {
+    auto gate = gate_coeff(0);
+    auto re = vmulq_f32(state_re[0], gate.re);
+    auto im = vmulq_f32(state_re[0], gate.im);
+    re = vfmsq_f32(re, state_im[0], gate.im);
+    im = vfmaq_f32(im, state_im[0], gate.re);
+
+    for (unsigned in = 1; in < Size; ++in) {
+      // Complex MAC: out += gate * state.
+      // re += state.re * gate.re - state.im * gate.im
+      // im += state.re * gate.im + state.im * gate.re
+      gate = gate_coeff(in);
+      re = vfmaq_f32(re, state_re[in], gate.re);
+      im = vfmaq_f32(im, state_re[in], gate.im);
+      re = vfmsq_f32(re, state_im[in], gate.im);
+      im = vfmaq_f32(im, state_im[in], gate.re);
+    }
+
+    return Complex{re, im};
+  }
+
+  static void StoreStateAmplitudeRow(
+      fp_type* state_block, const uint64_t* state_offsets,
+      unsigned output_basis, const Complex& output_amplitudes) {
+    const auto addr_re = state_block + state_offsets[output_basis];
+    const auto addr_im = addr_re + SIMDRegisterSize();
+    vst1q_f32(addr_re, output_amplitudes.re);
+    vst1q_f32(addr_im, output_amplitudes.im);
+  }
+
   template <unsigned H>
   void ApplyGateH(
       const std::vector<unsigned>& qs, const fp_type* matrix,
       State& state) const {
-    auto f = [](unsigned n, unsigned m, uint64_t i, const fp_type* v,
-                const uint64_t* ms, const uint64_t* xss, fp_type* rstate) {
+    auto f = [](unsigned n, unsigned m, uint64_t i,
+                const fp_type* gate_matrix, const uint64_t* masks,
+                const uint64_t* state_offsets, fp_type* state_data) {
       constexpr unsigned hsize = 1 << H;
 
-      float32x4_t rs[hsize];
-      float32x4_t is[hsize];
+      float32x4_t state_re[hsize];
+      float32x4_t state_im[hsize];
 
       i *= 4;
 
-      uint64_t ii = i & ms[0];
+      uint64_t ii = i & masks[0];
       for (unsigned j = 1; j <= H; ++j) {
         i *= 2;
-        ii |= i & ms[j];
+        ii |= i & masks[j];
       }
 
-      auto p0 = rstate + 2 * ii;
+      auto block = state_data + 2 * ii;
 
       for (unsigned k = 0; k < hsize; ++k) {
-        rs[k] = vld1q_f32(p0 + xss[k]);
-        is[k] = vld1q_f32(p0 + xss[k] + 4);
+        state_re[k] = vld1q_f32(block + state_offsets[k]);
+        state_im[k] = vld1q_f32(block + state_offsets[k] + 4);
       }
 
-      unsigned k = 0;
-
-      for (; k + 1 < hsize; k += 2) {
-        const fp_type* v0 = v + 2 * k * hsize;
-        const fp_type* v1 = v0 + 2 * hsize;
-
-        float32x4_t ru0 = vdupq_n_f32(v0[0]);
-        float32x4_t iu0 = vdupq_n_f32(v0[1]);
-        float32x4_t rn0 = vmulq_f32(rs[0], ru0);
-        float32x4_t in0 = vmulq_f32(rs[0], iu0);
-        rn0 = vfmsq_f32(rn0, is[0], iu0);
-        in0 = vfmaq_f32(in0, is[0], ru0);
-
-        float32x4_t ru1 = vdupq_n_f32(v1[0]);
-        float32x4_t iu1 = vdupq_n_f32(v1[1]);
-        float32x4_t rn1 = vmulq_f32(rs[0], ru1);
-        float32x4_t in1 = vmulq_f32(rs[0], iu1);
-        rn1 = vfmsq_f32(rn1, is[0], iu1);
-        in1 = vfmaq_f32(in1, is[0], ru1);
-
-        for (unsigned l = 1; l < hsize; ++l) {
-          ru0 = vdupq_n_f32(v0[2 * l]);
-          iu0 = vdupq_n_f32(v0[2 * l + 1]);
-          rn0 = vfmaq_f32(rn0, rs[l], ru0);
-          in0 = vfmaq_f32(in0, rs[l], iu0);
-          rn0 = vfmsq_f32(rn0, is[l], iu0);
-          in0 = vfmaq_f32(in0, is[l], ru0);
-
-          ru1 = vdupq_n_f32(v1[2 * l]);
-          iu1 = vdupq_n_f32(v1[2 * l + 1]);
-          rn1 = vfmaq_f32(rn1, rs[l], ru1);
-          in1 = vfmaq_f32(in1, rs[l], iu1);
-          rn1 = vfmsq_f32(rn1, is[l], iu1);
-          in1 = vfmaq_f32(in1, is[l], ru1);
-        }
-
-        vst1q_f32(p0 + xss[k], rn0);
-        vst1q_f32(p0 + xss[k] + 4, in0);
-        vst1q_f32(p0 + xss[k + 1], rn1);
-        vst1q_f32(p0 + xss[k + 1] + 4, in1);
+      auto load_gate_row = [](const fp_type* gate_row) {
+        return [gate_row](unsigned in) {
+          return Complex{
+              vdupq_n_f32(gate_row[2 * in]),
+              vdupq_n_f32(gate_row[2 * in + 1]),
+          };
+        };
+      };
+
+      unsigned out = 0;
+      for (; out + 1 < hsize; out += 2) {
+        const fp_type* gate_row0 = gate_matrix + 2 * out * hsize;
+        const fp_type* gate_row1 = gate_row0 + 2 * hsize;
+
+        auto out0 = ApplyGateRow<hsize>(state_re, state_im, load_gate_row(gate_row0));
+        auto out1 = ApplyGateRow<hsize>(state_re, state_im, load_gate_row(gate_row1));
+
+        StoreStateAmplitudeRow(block, state_offsets, out, out0);
+        StoreStateAmplitudeRow(block, state_offsets, out + 1, out1);
       }
 
-      for (; k < hsize; ++k) {
-        const fp_type* vk = v + 2 * k * hsize;
-
-        float32x4_t ru = vdupq_n_f32(vk[0]);
-        float32x4_t iu = vdupq_n_f32(vk[1]);
-        float32x4_t rn = vmulq_f32(rs[0], ru);
-        float32x4_t in = vmulq_f32(rs[0], iu);
-        rn = vfmsq_f32(rn, is[0], iu);
-        in = vfmaq_f32(in, is[0], ru);
+      for (; out < hsize; ++out) {
+        const fp_type* gate_row = gate_matrix + 2 * out * hsize;
 
-        for (unsigned l = 1; l < hsize; ++l) {
-          ru = vdupq_n_f32(vk[2 * l]);
-          iu = vdupq_n_f32(vk[2 * l + 1]);
-          rn = vfmaq_f32(rn, rs[l], ru);
-          in = vfmaq_f32(in, rs[l], iu);
-          rn = vfmsq_f32(rn, is[l], iu);
-          in = vfmaq_f32(in, is[l], ru);
-        }
+        auto out_row = ApplyGateRow<hsize>(state_re, state_im, load_gate_row(gate_row));
 
-        vst1q_f32(p0 + xss[k], rn);
-        vst1q_f32(p0 + xss[k] + 4, in);
+        StoreStateAmplitudeRow(block, state_offsets, out, out_row);
       }
     };
 
@@ -435,8 +440,8 @@ class SimulatorNEON final : public SimulatorBase {
       constexpr unsigned hsize = 1 << H;
       constexpr unsigned lsize = 1 << L;
 
-      float32x4_t rs[gsize];
-      float32x4_t is[gsize];
+      float32x4_t state_re[gsize];
+      float32x4_t state_im[gsize];
 
       i *= 4;
 
@@ -451,49 +456,41 @@ class SimulatorNEON final : public SimulatorBase {
       for (unsigned k = 0; k < hsize; ++k) {
         unsigned k2 = lsize * k;
 
-        rs[k2] = vld1q_f32(p0 + xss[k]);
-        is[k2] = vld1q_f32(p0 + xss[k] + 4);
+        state_re[k2] = vld1q_f32(p0 + xss[k]);
+        state_im[k2] = vld1q_f32(p0 + xss[k] + 4);
 
         if (L == 1) {
-          rs[k2 + 1] =
-              q0 == 0 ? vrev64q_f32(rs[k2]) : vextq_f32(rs[k2], rs[k2], 2);
-          is[k2 + 1] =
-              q0 == 0 ? vrev64q_f32(is[k2]) : vextq_f32(is[k2], is[k2], 2);
+          state_re[k2 + 1] =
+              q0 == 0 ? vrev64q_f32(state_re[k2])
+                      : vextq_f32(state_re[k2], state_re[k2], 2);
+          state_im[k2 + 1] =
+              q0 == 0 ? vrev64q_f32(state_im[k2])
+                      : vextq_f32(state_im[k2], state_im[k2], 2);
         } else if (L == 2) {
-          rs[k2 + 1] = vextq_f32(rs[k2], rs[k2], 1);
-          is[k2 + 1] = vextq_f32(is[k2], is[k2], 1);
-          rs[k2 + 2] = vextq_f32(rs[k2], rs[k2], 2);
-          is[k2 + 2] = vextq_f32(is[k2], is[k2], 2);
-          rs[k2 + 3] = vextq_f32(rs[k2], rs[k2], 3);
-          is[k2 + 3] = vextq_f32(is[k2], is[k2], 3);
+          state_re[k2 + 1] = vextq_f32(state_re[k2], state_re[k2], 1);
+          state_im[k2 + 1] = vextq_f32(state_im[k2], state_im[k2], 1);
+          state_re[k2 + 2] = vextq_f32(state_re[k2], state_re[k2], 2);
+          state_im[k2 + 2] = vextq_f32(state_im[k2], state_im[k2], 2);
+          state_re[k2 + 3] = vextq_f32(state_re[k2], state_re[k2], 3);
+          state_im[k2 + 3] = vextq_f32(state_im[k2], state_im[k2], 3);
         }
       }
 
-      uint64_t j = 0;
+      auto load_gate_row = [](const fp_type* gate_row) {
+        return [gate_row](unsigned in) {
+          return Complex{
+              vld1q_f32(gate_row + 8 * in),
+              vld1q_f32(gate_row + 8 * in + 4),
+          };
+        };
+      };
 
       for (unsigned k = 0; k < hsize; ++k) {
-        float32x4_t wre = vld1q_f32(w + 4 * j);
-        float32x4_t wim = vld1q_f32(w + 4 * (j + 1));
-        float32x4_t rn = vmulq_f32(rs[0], wre);
-        float32x4_t in = vmulq_f32(rs[0], wim);
-        rn = vfmsq_f32(rn, is[0], wim);
-        in = vfmaq_f32(in, is[0], wre);
+        const fp_type* gate_row = w + 8 * k * gsize;
+        auto out_row = ApplyGateRow<gsize>(
+            state_re, state_im, load_gate_row(gate_row));
 
-        j += 2;
-
-        for (unsigned l = 1; l < gsize; ++l) {
-          wre = vld1q_f32(w + 4 * j);
-          wim = vld1q_f32(w + 4 * (j + 1));
-          rn = vfmaq_f32(rn, rs[l], wre);
-          in = vfmaq_f32(in, rs[l], wim);
-          rn = vfmsq_f32(rn, is[l], wim);
-          in = vfmaq_f32(in, is[l], wre);
-
-          j += 2;
-        }
-
-        vst1q_f32(p0 + xss[k], rn);
-        vst1q_f32(p0 + xss[k] + 4, in);
+        StoreStateAmplitudeRow(p0, xss, k, out_row);
       }
     };