Skip to content

Commit

Permalink
stokhos: mask logic signature improvement (#12325)
Browse files Browse the repository at this point in the history
  • Loading branch information
uliegecsm committed Oct 3, 2023
1 parent e135b58 commit d8a6bd8
Show file tree
Hide file tree
Showing 2 changed files with 67 additions and 45 deletions.
2 changes: 1 addition & 1 deletion packages/stokhos/doc/Doxyfile
View file
Original file line number Diff line number Diff line change
Expand Up @@ -121,7 +121,7 @@ CLASS_DIAGRAMS = YES
# If the SOURCE_BROWSER tag is set to YES then a list of source files will
# be generated. Documented entities will be cross-referenced with these sources.

SOURCE_BROWSER = NO
SOURCE_BROWSER = YES

# Setting the INLINE_SOURCES tag to YES will include the body
# of functions and classes directly in the documentation.
Expand Down
110 changes: 66 additions & 44 deletions packages/stokhos/src/sacado/kokkos/vector/Stokhos_MP_Vector_MaskTraits.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@ template<typename scalar> class MaskedAssign
data = st_array[0];
else
data = st_array[1];

return *this;
}

Expand All @@ -116,7 +116,7 @@ template<typename scalar> class MaskedAssign
data = st_array[0]+st_array[1];
else
data = st_array[2];

return *this;
}

Expand All @@ -136,7 +136,7 @@ template<typename scalar> class MaskedAssign
data = st_array[0]-st_array[1];
else
data = st_array[2];

return *this;
}

Expand All @@ -156,7 +156,7 @@ template<typename scalar> class MaskedAssign
data = st_array[0]*st_array[1];
else
data = st_array[2];

return *this;
}

Expand All @@ -176,13 +176,30 @@ template<typename scalar> class MaskedAssign
data = st_array[0]/st_array[1];
else
data = st_array[2];

return *this;
}
};

template<typename scalar> class Mask;

namespace traits
{
template <typename>
struct is_mask : public std::false_type {};

template <typename scalar_t>
struct is_mask<Mask<scalar_t>> : public std::true_type {};

template <typename scalar_t>
struct is_mask<const Mask<scalar_t>> : public std::true_type {};

template <typename T>
inline constexpr bool is_mask_v = is_mask<T>::value;

#define CONCEPT_TRAITS_IS_MASK(M) std::enable_if_t<traits::is_mask_v<std::decay_t<M>>>
}

template<typename S> class MaskedAssign< Sacado::MP::Vector<S> >
{
typedef Sacado::MP::Vector<S> scalar;
Expand All @@ -208,7 +225,7 @@ template<typename S> class MaskedAssign< Sacado::MP::Vector<S> >
for(std::size_t i=0; i<size; ++i)
if(m.get(i))
data[i] = s[i];

return *this;
}

Expand All @@ -229,7 +246,7 @@ template<typename S> class MaskedAssign< Sacado::MP::Vector<S> >
data[i] = st_array[0][i];
else
data[i] = st_array[1][i];

return *this;
}

Expand All @@ -248,7 +265,7 @@ template<typename S> class MaskedAssign< Sacado::MP::Vector<S> >
for(std::size_t i=0; i<size; ++i)
if(m.get(i))
data[i] += s[i];

return *this;
}

Expand All @@ -269,7 +286,7 @@ template<typename S> class MaskedAssign< Sacado::MP::Vector<S> >
data[i] = st_array[0][i]+st_array[1][i];
else
data[i] = st_array[2][i];

return *this;
}

Expand All @@ -287,7 +304,7 @@ template<typename S> class MaskedAssign< Sacado::MP::Vector<S> >
for(std::size_t i=0; i<size; ++i)
if(m.get(i))
data[i] -= s[i];

return *this;
}

Expand All @@ -308,7 +325,7 @@ template<typename S> class MaskedAssign< Sacado::MP::Vector<S> >
data[i] = st_array[0][i]-st_array[1][i];
else
data[i] = st_array[2][i];

return *this;
}

Expand All @@ -326,7 +343,7 @@ template<typename S> class MaskedAssign< Sacado::MP::Vector<S> >
for(std::size_t i=0; i<size; ++i)
if(m.get(i))
data[i] *= s[i];

return *this;
}

Expand All @@ -347,7 +364,7 @@ template<typename S> class MaskedAssign< Sacado::MP::Vector<S> >
data[i] = st_array[0][i]*st_array[1][i];
else
data[i] = st_array[2][i];

return *this;
}

Expand All @@ -365,7 +382,7 @@ template<typename S> class MaskedAssign< Sacado::MP::Vector<S> >
for(std::size_t i=0; i<size; ++i)
if(m.get(i))
data[i] /= s[i];

return *this;
}

Expand All @@ -386,7 +403,7 @@ template<typename S> class MaskedAssign< Sacado::MP::Vector<S> >
data[i] = st_array[0][i]/st_array[1][i];
else
data[i] = st_array[2][i];

return *this;
}
};
Expand All @@ -406,19 +423,19 @@ template <typename ordinal_t, typename value_t, typename device_t> class MaskedA
{
if(m.get(0))
data = s;

return *this;
}

KOKKOS_INLINE_FUNCTION MaskedAssign<scalar>& operator = (const std::initializer_list<scalar> & KOKKOS_RESTRICT st)
{
auto st_array = st.begin();

if(m.get(0))
data = st_array[0];
else
data = st_array[1];

return *this;
}

Expand All @@ -427,7 +444,7 @@ template <typename ordinal_t, typename value_t, typename device_t> class MaskedA
{
if(m.get(0))
data += s;

return *this;
}

Expand All @@ -439,15 +456,15 @@ template <typename ordinal_t, typename value_t, typename device_t> class MaskedA
data = st_array[0]+st_array[1];
else
data = st_array[2];

return *this;
}

KOKKOS_INLINE_FUNCTION MaskedAssign<scalar>& operator -= (const scalar & KOKKOS_RESTRICT s)
{
if(m.get(0))
data -= s;

return *this;
}

Expand All @@ -459,15 +476,15 @@ template <typename ordinal_t, typename value_t, typename device_t> class MaskedA
data = st_array[0]-st_array[1];
else
data = st_array[2];

return *this;
}

KOKKOS_INLINE_FUNCTION MaskedAssign<scalar>& operator *= (const scalar & KOKKOS_RESTRICT s)
{
if(m.get(0))
data *= s;

return *this;
}

Expand All @@ -479,15 +496,15 @@ template <typename ordinal_t, typename value_t, typename device_t> class MaskedA
data = st_array[0]*st_array[1];
else
data = st_array[2];

return *this;
}

KOKKOS_INLINE_FUNCTION MaskedAssign<scalar>& operator /= (const scalar & KOKKOS_RESTRICT s)
{
if(m.get(0))
data /= s;

return *this;
}

Expand All @@ -499,7 +516,7 @@ template <typename ordinal_t, typename value_t, typename device_t> class MaskedA
data = st_array[0]/st_array[1];
else
data = st_array[2];

return *this;
}
};
Expand Down Expand Up @@ -603,15 +620,15 @@ template<typename scalar> class Mask
{
return !(this==m2);
}

KOKKOS_INLINE_FUNCTION Mask<scalar> operator&& (const Mask<scalar> &m2)
{
Mask<scalar> m3;
for(std::size_t i=0; i<size; ++i)
m3.set(i,(this->get(i) && m2.get(i)));
return m3;
}

KOKKOS_INLINE_FUNCTION Mask<scalar> operator|| (const Mask<scalar> &m2)
{
Mask<scalar> m3;
Expand Down Expand Up @@ -665,7 +682,7 @@ template<typename scalar> class Mask

return v2;
}

KOKKOS_INLINE_FUNCTION bool get (int i) const
{
return this->data[i];
Expand All @@ -675,15 +692,15 @@ template<typename scalar> class Mask
{
this->data[i] = b;
}

KOKKOS_INLINE_FUNCTION Mask<scalar> operator! ()
{
Mask<scalar> m2;
for(std::size_t i=0; i<size; ++i)
m2.set(i,!(this->get(i)));
return m2;
}

KOKKOS_INLINE_FUNCTION operator bool() const
{
return this->get(0);
Expand Down Expand Up @@ -794,9 +811,9 @@ namespace Sacado {
template <typename S> KOKKOS_INLINE_FUNCTION Vector<S> copysign(const Vector<S> &a1, const Vector<S> &a2)
{
typedef EnsembleTraits_m< Vector<S> > ET;

Vector<S> a_out;

using std::copysign;
#ifdef STOKHOS_HAVE_PRAGMA_IVDEP
#pragma ivdep
Expand All @@ -810,7 +827,7 @@ namespace Sacado {
for(std::size_t i=0; i<ET::size; ++i){
ET::coeff(a_out,i) = copysign(ET::coeff(a1,i),ET::coeff(a2,i));
}

return a_out;
}
}
Expand All @@ -821,7 +838,7 @@ template<typename S> KOKKOS_INLINE_FUNCTION Mask<Sacado::MP::Vector<S> > signbit
{
typedef EnsembleTraits_m<Sacado::MP::Vector<S> > ET;
using std::signbit;

Mask<Sacado::MP::Vector<S> > mask;
#ifdef STOKHOS_HAVE_PRAGMA_IVDEP
#pragma ivdep
Expand Down Expand Up @@ -935,32 +952,37 @@ template<typename S> KOKKOS_INLINE_FUNCTION Mask<Sacado::MP::Vector<S> > signbit

#endif

namespace MaskLogic{

template<typename T> KOKKOS_INLINE_FUNCTION bool OR(Mask<T> m){
return (((double) m)!=0.);
namespace MaskLogic
{
template <typename M, typename = CONCEPT_TRAITS_IS_MASK(M)>
KOKKOS_INLINE_FUNCTION
bool OR(M&& m) {
return m.operator double() != 0.;
}

KOKKOS_INLINE_FUNCTION bool OR(bool m){
return m;
}

template<typename T> KOKKOS_INLINE_FUNCTION bool XOR(Mask<T> m){
return (((double) m)==1./m.getSize());
template <typename M, typename = CONCEPT_TRAITS_IS_MASK(M)>
KOKKOS_INLINE_FUNCTION
bool XOR(M&& m) {
return m.operator double() == 1. / m.getSize();
}

KOKKOS_INLINE_FUNCTION bool XOR(bool m){
return m;
}

template<typename T> KOKKOS_INLINE_FUNCTION bool AND(Mask<T> m){
return (((double) m)==1.);
template <typename M, typename = CONCEPT_TRAITS_IS_MASK(M)>
KOKKOS_INLINE_FUNCTION
bool AND(M&& m) {
return m.operator double() == 1.;
}

KOKKOS_INLINE_FUNCTION bool AND(bool m){
return m;
}

}

#endif // STOKHOS_MP_VECTOR_MASKTRAITS_HPP

0 comments on commit d8a6bd8

Please sign in to comment.