From 8400cbf9bf17c73556b898a629788baf8675fb38 Mon Sep 17 00:00:00 2001
From: Dong Hun Lee <59181952+ldh4@users.noreply.github.com>
Date: Thu, 23 Mar 2023 16:32:08 -0600
Subject: [PATCH] simd: Fixed an incorrectly returning size for uint64_t in
 avx2 (#6004)

* Fixed a size error in simd uint64_t avx2

* Converted unit test to compile time checks

* Removed an unused variable
---
 simd/src/Kokkos_SIMD_AVX2.hpp |  2 +-
 simd/unit_tests/TestSIMD.cpp  | 29 +++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/simd/src/Kokkos_SIMD_AVX2.hpp b/simd/src/Kokkos_SIMD_AVX2.hpp
index 1732c33ca5..86b944efa5 100644
--- a/simd/src/Kokkos_SIMD_AVX2.hpp
+++ b/simd/src/Kokkos_SIMD_AVX2.hpp
@@ -804,7 +804,7 @@ class simd<std::uint64_t, simd_abi::avx2_fixed_size<4>> {
   KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION simd& operator=(simd const&) = default;
   KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION simd& operator=(simd&&) = default;
   KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION static constexpr std::size_t size() {
-    return 8;
+    return 4;
   }
   template <class U, std::enable_if_t<std::is_convertible_v<U, value_type>,
                                       bool> = false>
diff --git a/simd/unit_tests/TestSIMD.cpp b/simd/unit_tests/TestSIMD.cpp
index 7a4ecf19ed..92c77033b9 100644
--- a/simd/unit_tests/TestSIMD.cpp
+++ b/simd/unit_tests/TestSIMD.cpp
@@ -486,3 +486,32 @@ TEST(simd, device) {
   Kokkos::parallel_for(Kokkos::RangePolicy<Kokkos::IndexType<int>>(0, 1),
                        simd_device_functor());
 }
+
+TEST(simd, test_size) {
+#if defined(KOKKOS_ARCH_AVX512XEON)
+  constexpr auto width = 8;
+  using Abi = Kokkos::Experimental::simd_abi::avx512_fixed_size<width>;
+  static_assert(width ==
+                Kokkos::Experimental::simd<std::uint32_t, Abi>::size());
+
+#elif defined(KOKKOS_ARCH_AVX2)
+  constexpr auto width = 4;
+  using Abi            = Kokkos::Experimental::simd_abi::avx2_fixed_size<width>;
+
+#elif defined(__ARM_NEON)
+  constexpr auto width = 2;
+  using Abi            = Kokkos::Experimental::simd_abi::neon_fixed_size<width>;
+
+#else
+  constexpr auto width = 1;
+  using Abi            = Kokkos::Experimental::simd_abi::scalar;
+  static_assert(width ==
+                Kokkos::Experimental::simd<std::uint32_t, Abi>::size());
+#endif
+
+  static_assert(width == Kokkos::Experimental::simd<double, Abi>::size());
+  static_assert(width == Kokkos::Experimental::simd<std::int64_t, Abi>::size());
+  static_assert(width ==
+                Kokkos::Experimental::simd<std::uint64_t, Abi>::size());
+  static_assert(width == Kokkos::Experimental::simd<std::int32_t, Abi>::size());
+}