261 lines
12 KiB
C++
261 lines
12 KiB
C++
//@HEADER
|
|
// ************************************************************************
|
|
//
|
|
// Kokkos v. 4.0
|
|
// Copyright (2022) National Technology & Engineering
|
|
// Solutions of Sandia, LLC (NTESS).
|
|
//
|
|
// Under the terms of Contract DE-NA0003525 with NTESS,
|
|
// the U.S. Government retains certain rights in this software.
|
|
//
|
|
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://kokkos.org/LICENSE for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//@HEADER
|
|
|
|
#ifndef KOKKOS_SIMD_COMMON_MATH_HPP
|
|
#define KOKKOS_SIMD_COMMON_MATH_HPP
|
|
|
|
#include <Kokkos_Core.hpp> // Kokkos::min, etc.
|
|
|
|
namespace Kokkos {
|
|
|
|
namespace Experimental {
|
|
|
|
template <class T, class Abi>
|
|
class simd;
|
|
|
|
template <class T, class Abi>
|
|
class simd_mask;
|
|
|
|
template <class M, class T>
|
|
class const_where_expression;
|
|
|
|
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.
|
|
|
|
#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
|
|
|
|
#endif
|