Merge pull request #6011 from ldh4/release-candidate-4.0.01

[4.0.01] Fix simd incorrect size return in avx2
sandialabs · Mar 24, 2023 · dc876ea · dc876ea
2 parents aa1f48f + 8400cbf
commit dc876ea
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 1 deletion.
diff --git 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
@@ -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());
+}