Skip to content

Commit

Permalink
SIMD: Split math functions from SIMD_Common.hpp (kokkos#6487)
Browse files Browse the repository at this point in the history 
* SIMD: Split math functions from SIMD_Common.hpp

* Make Kokkos_SIMD_Common_Math.hpp self-contained

* Enforce header inclusion order

* Clean up comment in Kokkos_SIMD_Common_Math.hpp
  • Loading branch information
@masterleinad
masterleinad committed Oct 11, 2023
1 parent 8420c2f commit f511dca
Show file tree
Hide file tree
Showing 8 changed files with 283 additions and 225 deletions.
2 changes: 2 additions & 0 deletions simd/src/Kokkos_SIMD.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,8 @@
#pragma GCC diagnostic pop
#endif

#include <Kokkos_SIMD_Common_Math.hpp>

namespace Kokkos {
namespace Experimental {

Expand Down
5 changes: 5 additions & 0 deletions simd/src/Kokkos_SIMD_AVX2.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,11 @@

#include <immintrin.h>

#ifdef KOKKOS_SIMD_COMMON_MATH_HPP
#error \
"Kokkos_SIMD_AVX2.hpp must be included before Kokkos_SIMD_Common_Math.hpp!"
#endif

// FIXME_HIP ROCm 5.6 and 5.7 can't compile with the intrinsic used here.
#if defined(__HIPCC__) && (HIP_VERSION_MAJOR == 5) && \
((HIP_VERSION_MINOR == 6) || (HIP_VERSION_MINOR == 7))
Expand Down
5 changes: 5 additions & 0 deletions simd/src/Kokkos_SIMD_AVX512.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,11 @@

#include <immintrin.h>

#ifdef KOKKOS_SIMD_COMMON_MATH_HPP
#error \
"Kokkos_SIMD_AVX512.hpp must be included before Kokkos_SIMD_Common_Math.hpp!"
#endif

namespace Kokkos {
namespace Experimental {

Expand Down
224 changes: 0 additions & 224 deletions simd/src/Kokkos_SIMD_Common.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
#ifndef KOKKOS_SIMD_COMMON_HPP
#define KOKKOS_SIMD_COMMON_HPP

#include <cmath>
#include <cstring>

#include <Kokkos_Core.hpp>
Expand Down Expand Up @@ -329,230 +328,7 @@ template <class T, class Abi>
return a == simd_mask<T, Abi>(false);
}

template <typename T, typename Abi>
[[nodiscard]] KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION T
hmin(const_where_expression<simd_mask<T, Abi>, simd<T, Abi>> const& x) {
auto const& v = x.impl_get_value();
auto const& m = x.impl_get_mask();
auto result = Kokkos::reduction_identity<T>::min();
for (std::size_t i = 0; i < v.size(); ++i) {
if (m[i]) result = Kokkos::min(result, v[i]);
}
return result;
}

template <class T, class Abi>
[[nodiscard]] KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION T
hmax(const_where_expression<simd_mask<T, Abi>, simd<T, Abi>> const& x) {
auto const& v = x.impl_get_value();
auto const& m = x.impl_get_mask();
auto result = Kokkos::reduction_identity<T>::max();
for (std::size_t i = 0; i < v.size(); ++i) {
if (m[i]) result = Kokkos::max(result, v[i]);
}
return result;
}

template <class T, class Abi>
[[nodiscard]] KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION T
reduce(const_where_expression<simd_mask<T, Abi>, simd<T, Abi>> const& x, T,
std::plus<>) {
auto const& v = x.impl_get_value();
auto const& m = x.impl_get_mask();
auto result = Kokkos::reduction_identity<T>::sum();
for (std::size_t i = 0; i < v.size(); ++i) {
if (m[i]) result += v[i];
}
return result;
}

} // namespace Experimental

template <class T, class Abi>
[[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION Experimental::simd<T, Abi> min(
Experimental::simd<T, Abi> const& a, Experimental::simd<T, Abi> const& b) {
Experimental::simd<T, Abi> result;
for (std::size_t i = 0; i < Experimental::simd<T, Abi>::size(); ++i) {
result[i] = Kokkos::min(a[i], b[i]);
}
return result;
}

#ifdef KOKKOS_ENABLE_DEPRECATED_CODE_4
namespace Experimental {
template <class T, class Abi>
[[nodiscard]] KOKKOS_DEPRECATED KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION
Experimental::simd<T, Abi>
min(Experimental::simd<T, Abi> const& a,
Experimental::simd<T, Abi> const& b) {
return Kokkos::min(a, b);
}
} // namespace Experimental
#endif

template <class T, class Abi>
[[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION Experimental::simd<T, Abi> max(
Experimental::simd<T, Abi> const& a, Experimental::simd<T, Abi> const& b) {
Experimental::simd<T, Abi> result;
for (std::size_t i = 0; i < Experimental::simd<T, Abi>::size(); ++i) {
result[i] = Kokkos::max(a[i], b[i]);
}
return result;
}

#ifdef KOKKOS_ENABLE_DEPRECATED_CODE_4
namespace Experimental {
template <class T, class Abi>
[[nodiscard]] KOKKOS_DEPRECATED KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION
Experimental::simd<T, Abi>
max(Experimental::simd<T, Abi> const& a,
Experimental::simd<T, Abi> const& b) {
return Kokkos::max(a, b);
}
} // namespace Experimental
#endif

// fallback implementations of <cmath> functions.
// individual Abi types may provide overloads with more efficient
// implementations.
// These are not in the Experimental namespace because their double
// overloads are not either

#ifdef KOKKOS_ENABLE_DEPRECATED_CODE_4
#define KOKKOS_IMPL_SIMD_UNARY_FUNCTION(FUNC) \
template <class T, class Abi> \
[[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION Experimental::simd<T, Abi> FUNC( \
Experimental::simd<T, Abi> const& a) { \
Experimental::simd<T, Abi> result; \
for (std::size_t i = 0; i < Experimental::simd<T, Abi>::size(); ++i) { \
result[i] = Kokkos::FUNC(a[i]); \
} \
return result; \
} \
namespace Experimental { \
template <class T, class Abi> \
[[nodiscard]] KOKKOS_DEPRECATED KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION \
simd<T, Abi> \
FUNC(simd<T, Abi> const& a) { \
return Kokkos::FUNC(a); \
} \
}
#else
#define KOKKOS_IMPL_SIMD_UNARY_FUNCTION(FUNC) \
template <class T, class Abi> \
[[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION Experimental::simd<T, Abi> FUNC( \
Experimental::simd<T, Abi> const& a) { \
Experimental::simd<T, Abi> result; \
for (std::size_t i = 0; i < Experimental::simd<T, Abi>::size(); ++i) { \
result[i] = Kokkos::FUNC(a[i]); \
} \
return result; \
}
#endif

KOKKOS_IMPL_SIMD_UNARY_FUNCTION(abs)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(exp)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(exp2)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(log)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(log10)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(log2)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(sqrt)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(cbrt)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(sin)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(cos)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(tan)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(asin)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(acos)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(atan)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(sinh)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(cosh)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(tanh)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(asinh)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(acosh)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(atanh)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(erf)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(erfc)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(tgamma)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(lgamma)

#ifdef KOKKOS_ENABLE_DEPRECATED_CODE_4
#define KOKKOS_IMPL_SIMD_BINARY_FUNCTION(FUNC) \
template <class T, class Abi> \
[[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION Experimental::simd<T, Abi> FUNC( \
Experimental::simd<T, Abi> const& a, \
Experimental::simd<T, Abi> const& b) { \
Experimental::simd<T, Abi> result; \
for (std::size_t i = 0; i < Experimental::simd<T, Abi>::size(); ++i) { \
result[i] = Kokkos::FUNC(a[i], b[i]); \
} \
return result; \
} \
namespace Experimental { \
template <class T, class Abi> \
[[nodiscard]] KOKKOS_DEPRECATED KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION \
simd<T, Abi> \
FUNC(simd<T, Abi> const& a, simd<T, Abi> const& b) { \
Kokkos::FUNC(a, b); \
} \
}
#else
#define KOKKOS_IMPL_SIMD_BINARY_FUNCTION(FUNC) \
template <class T, class Abi> \
[[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION Experimental::simd<T, Abi> FUNC( \
Experimental::simd<T, Abi> const& a, \
Experimental::simd<T, Abi> const& b) { \
Experimental::simd<T, Abi> result; \
for (std::size_t i = 0; i < Experimental::simd<T, Abi>::size(); ++i) { \
result[i] = Kokkos::FUNC(a[i], b[i]); \
} \
return result; \
}
#endif

KOKKOS_IMPL_SIMD_BINARY_FUNCTION(pow)
KOKKOS_IMPL_SIMD_BINARY_FUNCTION(hypot)
KOKKOS_IMPL_SIMD_BINARY_FUNCTION(atan2)
KOKKOS_IMPL_SIMD_BINARY_FUNCTION(copysign)

#ifdef KOKKOS_ENABLE_DEPRECATED_CODE_4
#define KOKKOS_IMPL_SIMD_TERNARY_FUNCTION(FUNC) \
template <class T, class Abi> \
[[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION Experimental::simd<T, Abi> FUNC( \
Experimental::simd<T, Abi> const& a, \
Experimental::simd<T, Abi> const& b, \
Experimental::simd<T, Abi> const& c) { \
Experimental::simd<T, Abi> result; \
for (std::size_t i = 0; i < Experimental::simd<T, Abi>::size(); ++i) { \
result[i] = Kokkos::FUNC(a[i], b[i], c[i]); \
} \
return result; \
} \
namespace Experimental { \
template <class T, class Abi> \
[[nodiscard]] KOKKOS_DEPRECATED KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION \
simd<T, Abi> \
FUNC(simd<T, Abi> const& a, simd<T, Abi> const& b, \
simd<T, Abi> const& c) { \
return Kokkos::FUNC(a, b, c); \
} \
}
#else
#define KOKKOS_IMPL_SIMD_TERNARY_FUNCTION(FUNC) \
template <class T, class Abi> \
[[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION Experimental::simd<T, Abi> FUNC( \
Experimental::simd<T, Abi> const& a, \
Experimental::simd<T, Abi> const& b, \
Experimental::simd<T, Abi> const& c) { \
Experimental::simd<T, Abi> result; \
for (std::size_t i = 0; i < Experimental::simd<T, Abi>::size(); ++i) { \
result[i] = Kokkos::FUNC(a[i], b[i], c[i]); \
} \
return result; \
}
#endif

KOKKOS_IMPL_SIMD_TERNARY_FUNCTION(fma)
KOKKOS_IMPL_SIMD_TERNARY_FUNCTION(hypot)

} // namespace Kokkos

Expand Down

0 comments on commit f511dca

Please sign in to comment.