Added unit tests for reduction ops and few intel svml intrinsics

masterleinad · Oct 24, 2023 · ae0bd54 · ae0bd54
1 parent cf5a859
commit ae0bd54
Show file tree

Hide file tree

Showing 5 changed files with 360 additions and 9 deletions.
diff --git a/simd/src/Kokkos_SIMD_Scalar.hpp b/simd/src/Kokkos_SIMD_Scalar.hpp
@@ -315,13 +315,13 @@ class const_where_expression<simd_mask<T, simd_abi::scalar>,
       mem[static_cast<Integral>(index)] = static_cast<T>(m_value);
   }
 
-  [[nodiscard]] KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION value_type const&
-  impl_get_value() const {
+  [[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION value_type const& impl_get_value()
+      const {
     return m_value;
   }
 
-  [[nodiscard]] KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION mask_type const&
-  impl_get_mask() const {
+  [[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION mask_type const& impl_get_mask()
+      const {
     return m_mask;
   }
 };

diff --git a/simd/unit_tests/TestSIMD.cpp b/simd/unit_tests/TestSIMD.cpp
@@ -21,3 +21,4 @@
 #include <TestSIMD_Condition.hpp>
 #include <TestSIMD_GeneratorCtors.hpp>
 #include <TestSIMD_WhereExpressions.hpp>
+#include <TestSIMD_Reductions.hpp>
diff --git a/simd/unit_tests/include/SIMDTesting_Ops.hpp b/simd/unit_tests/include/SIMDTesting_Ops.hpp
@@ -209,4 +209,165 @@ class shift_left {
   }
 };
 
+class cbrt_op {
+ public:
+  template <typename T>
+  auto on_host(T const& a) const {
+#if defined(KOKKOS_ENABLE_DEPRECATED_CODE_4)
+    return Kokkos::Experimental::cbrt(a);
+#else
+    return Kokkos::cbrt(a);
+#endif
+  }
+  template <typename T>
+  auto on_host_serial(T const& a) const {
+    return Kokkos::cbrt(a);
+  }
+};
+
+class exp_op {
+ public:
+  template <typename T>
+  auto on_host(T const& a) const {
+#if defined(KOKKOS_ENABLE_DEPRECATED_CODE_4)
+    return Kokkos::Experimental::exp(a);
+#else
+    return Kokkos::exp(a);
+#endif
+  }
+  template <typename T>
+  auto on_host_serial(T const& a) const {
+    return Kokkos::exp(a);
+  }
+};
+
+class log_op {
+ public:
+  template <typename T>
+  auto on_host(T const& a) const {
+#if defined(KOKKOS_ENABLE_DEPRECATED_CODE_4)
+    return Kokkos::Experimental::log(a);
+#else
+    return Kokkos::log(a);
+#endif
+  }
+  template <typename T>
+  auto on_host_serial(T const& a) const {
+    return Kokkos::log(a);
+  }
+};
+
+class hmin {
+ public:
+  template <typename T>
+  auto on_host(T const& a) const {
+    return Kokkos::Experimental::hmin(a);
+  }
+  template <typename T>
+  auto on_host_serial(T const& a) const {
+    using DataType = typename T::value_type::value_type;
+
+    auto const& v = a.impl_get_value();
+    auto const& m = a.impl_get_mask();
+    auto result   = Kokkos::reduction_identity<DataType>::min();
+    for (std::size_t i = 0; i < v.size(); ++i) {
+      if (m[i]) result = Kokkos::min(result, v[i]);
+    }
+    return result;
+  }
+
+  template <typename T>
+  KOKKOS_INLINE_FUNCTION auto on_device(T const& a) const {
+    return Kokkos::Experimental::hmin(a);
+  }
+  template <typename T>
+  KOKKOS_INLINE_FUNCTION auto on_device_serial(T const& a) const {
+    using DataType = typename T::value_type::value_type;
+
+    auto const& v = a.impl_get_value();
+    auto const& m = a.impl_get_mask();
+    auto result   = Kokkos::reduction_identity<DataType>::min();
+    for (std::size_t i = 0; i < v.size(); ++i) {
+      if (m[i]) result = Kokkos::min(result, v[i]);
+    }
+    return result;
+  }
+};
+
+class hmax {
+ public:
+  template <typename T>
+  auto on_host(T const& a) const {
+    return Kokkos::Experimental::hmax(a);
+  }
+  template <typename T>
+  auto on_host_serial(T const& a) const {
+    using DataType = typename T::value_type::value_type;
+
+    auto const& v = a.impl_get_value();
+    auto const& m = a.impl_get_mask();
+    auto result   = Kokkos::reduction_identity<DataType>::max();
+    for (std::size_t i = 0; i < v.size(); ++i) {
+      if (m[i]) result = Kokkos::max(result, v[i]);
+    }
+    return result;
+  }
+
+  template <typename T>
+  KOKKOS_INLINE_FUNCTION auto on_device(T const& a) const {
+    return Kokkos::Experimental::hmax(a);
+  }
+  template <typename T>
+  KOKKOS_INLINE_FUNCTION auto on_device_serial(T const& a) const {
+    using DataType = typename T::value_type::value_type;
+
+    auto const& v = a.impl_get_value();
+    auto const& m = a.impl_get_mask();
+    auto result   = Kokkos::reduction_identity<DataType>::max();
+    for (std::size_t i = 0; i < v.size(); ++i) {
+      if (m[i]) result = Kokkos::max(result, v[i]);
+    }
+    return result;
+  }
+};
+
+class reduce {
+ public:
+  template <typename T>
+  auto on_host(T const& a) const {
+    using DataType = typename T::value_type::value_type;
+    return Kokkos::Experimental::reduce(a, DataType(0), std::plus<>());
+  }
+  template <typename T>
+  auto on_host_serial(T const& a) const {
+    using DataType = typename T::value_type::value_type;
+
+    auto const& v = a.impl_get_value();
+    auto const& m = a.impl_get_mask();
+    auto result   = Kokkos::reduction_identity<DataType>::sum();
+    for (std::size_t i = 0; i < v.size(); ++i) {
+      if (m[i]) result += v[i];
+    }
+    return result;
+  }
+
+  template <typename T>
+  KOKKOS_INLINE_FUNCTION auto on_device(T const& a) const {
+    using DataType = typename T::value_type::value_type;
+    return Kokkos::Experimental::reduce(a, DataType(0), std::plus<>());
+  }
+  template <typename T>
+  KOKKOS_INLINE_FUNCTION auto on_device_serial(T const& a) const {
+    using DataType = typename T::value_type::value_type;
+
+    auto const& v = a.impl_get_value();
+    auto const& m = a.impl_get_mask();
+    auto result   = Kokkos::reduction_identity<DataType>::sum();
+    for (std::size_t i = 0; i < v.size(); ++i) {
+      if (m[i]) result += v[i];
+    }
+    return result;
+  }
+};
+
 #endif
diff --git a/simd/unit_tests/include/TestSIMD_MathOps.hpp b/simd/unit_tests/include/TestSIMD_MathOps.hpp
@@ -61,13 +61,18 @@ void host_check_math_op_one_loader(UnaryOp unary_op, std::size_t n,
     simd_type arg;
     bool const loaded_arg = loader.host_load(args + i, nlanes, arg);
     if (!loaded_arg) continue;
-    auto computed_result = unary_op.on_host(arg);
 
-    decltype(computed_result) expected_result;
+    decltype(unary_op.on_host(arg)) expected_result;
     for (std::size_t lane = 0; lane < simd_type::size(); ++lane) {
-      if (lane < nlanes)
+      if (lane < nlanes) {
+        if constexpr (std::is_same_v<UnaryOp, cbrt_op> ||
+                      std::is_same_v<UnaryOp, exp_op> ||
+                      std::is_same_v<UnaryOp, log_op>)
+          arg[lane] = Kokkos::abs(arg[lane]);
         expected_result[lane] = unary_op.on_host_serial(T(arg[lane]));
+      }
     }
+    auto computed_result = unary_op.on_host(arg);
     host_check_equality(expected_result, computed_result, nlanes);
   }
 }
@@ -96,6 +101,13 @@ inline void host_check_all_math_ops(const DataType (&first_args)[n],
   // TODO: Place fallback implementations for all simd integer types
   if constexpr (std::is_floating_point_v<DataType>) {
     host_check_math_op_all_loaders<Abi>(divides(), n, first_args, second_args);
+
+#if defined(__INTEL_COMPILER) && \
+    (defined(KOKKOS_ARCH_AVX2) || defined(KOKKOS_ARCH_AVX512XEON))
+    host_check_math_op_all_loaders<Abi>(cbrt_op(), n, first_args);
+    host_check_math_op_all_loaders<Abi>(exp_op(), n, first_args);
+    host_check_math_op_all_loaders<Abi>(log_op(), n, first_args);
+#endif
   }
 }
 
@@ -282,8 +294,7 @@ TEST(simd, host_math_ops) {
 }
 
 TEST(simd, device_math_ops) {
-  Kokkos::parallel_for(Kokkos::RangePolicy<Kokkos::IndexType<int>>(0, 1),
-                       simd_device_math_ops_functor());
+  Kokkos::parallel_for(1, simd_device_math_ops_functor());
 }
 
 #endif