Update overloads

PennyLaneAI · Mar 27, 2024 · 68eb3cf · 68eb3cf
1 parent f87bdc1
commit 68eb3cf
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 6 deletions.
diff --git a/...ightning/core/src/simulators/lightning_qubit/gates/cpu_kernels/avx_common/AVX2Concept.hpp b/...ightning/core/src/simulators/lightning_qubit/gates/cpu_kernels/avx_common/AVX2Concept.hpp
@@ -103,13 +103,29 @@ template <typename T> struct AVX2Concept {
         }
     }
 
+    PL_FORCE_INLINE
+    static void stream_(PrecisionT *p, IntrinsicType value) {
+        if constexpr (std::is_same_v<PrecisionT, float>) {
+            _mm256_stream_ps(p, value);
+        } else if (std::is_same_v<PrecisionT, double>) {
+            _mm256_stream_pd(p, value);
+        } else {
+            static_assert(std::is_same_v<PrecisionT, float> ||
+                          std::is_same_v<PrecisionT, double>);
+        }
+    }
+
     PL_FORCE_INLINE
     static void store(std::complex<PrecisionT> *p, IntrinsicType value) {
-        #ifdef PL_LQ_KERNEL_AVX_STREAMING
+        store(reinterpret_cast<PrecisionT *>(p), value);
+    }
+    PL_FORCE_INLINE
+    static void store(PrecisionT *p, IntrinsicType value) {
+#ifdef PL_LQ_KERNEL_AVX_STREAMING
         store_(p, value);
-        #else
+#else
         stream_(p, value);
-        #endif
+#endif
     }
 
     PL_FORCE_INLINE

diff --git a/...htning/core/src/simulators/lightning_qubit/gates/cpu_kernels/avx_common/AVX512Concept.hpp b/...htning/core/src/simulators/lightning_qubit/gates/cpu_kernels/avx_common/AVX512Concept.hpp
@@ -104,13 +104,30 @@ template <typename T> struct AVX512Concept {
         }
     }
 
+    PL_FORCE_INLINE
+    static void stream_(PrecisionT *p, IntrinsicType value) {
+        if constexpr (std::is_same_v<PrecisionT, float>) {
+            _mm512_stream_ps(p, value);
+        } else if (std::is_same_v<PrecisionT, double>) {
+            _mm512_stream_pd(p, value);
+        } else {
+            static_assert(std::is_same_v<PrecisionT, float> ||
+                          std::is_same_v<PrecisionT, double>);
+        }
+    }
+
     PL_FORCE_INLINE
     static void store(std::complex<PrecisionT> *p, IntrinsicType value) {
-        #ifdef PL_LQ_KERNEL_AVX_STREAMING
+        store(reinterpret_cast<PrecisionT *>(p), value);
+    }
+
+    PL_FORCE_INLINE
+    static void store(PrecisionT *p, IntrinsicType value) {
+#ifdef PL_LQ_KERNEL_AVX_STREAMING
         store_(p, value);
-        #else
+#else
         stream_(p, value);
-        #endif
+#endif
     }
 
     PL_FORCE_INLINE