Skip to content

Commit

Permalink
Refactor create_mirror for View of MP Vector
Browse files Browse the repository at this point in the history
  • Loading branch information
@maartenarnst
maartenarnst committed Sep 9, 2024
1 parent 49d2b72 commit e3dadea
Show file tree
Hide file tree
Showing 2 changed files with 134 additions and 186 deletions.
206 changes: 81 additions & 125 deletions packages/stokhos/src/sacado/kokkos/vector/KokkosExp_View_MP_Vector_Contiguous.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,150 +92,106 @@ struct FlatArrayType< View<D,P...>,
typedef View<flat_data_type,P...> type;
};

template< class T , class ... P >
inline
typename std::enable_if<
std::is_same< typename ViewTraits<T,P...>::specialize ,
Kokkos::Experimental::Impl::ViewMPVectorContiguous >::value &&
!std::is_same< typename Kokkos::ViewTraits<T,P...>::array_layout,
Kokkos::LayoutStride >::value,
typename Kokkos::View<T,P...>::HostMirror>::type
create_mirror(const Kokkos::View<T,P...> & src)
template <class T, class... P, class... ViewCtorArgs>
inline auto create_mirror(
const Impl::ViewCtorProp<ViewCtorArgs...>& arg_prop,
const View<T, P...>& src,
typename std::enable_if_t<
std::is_same_v<typename ViewTraits<T, P...>::specialize,
Experimental::Impl::ViewMPVectorContiguous>>*
)
{
typedef View<T,P...> src_type ;
typedef typename src_type::HostMirror dst_type ;
static_assert(std::is_same_v<typename ViewTraits<T, P...>::array_layout, LayoutLeft> ||
std::is_same_v<typename ViewTraits<T, P...>::array_layout, LayoutRight> ||
std::is_same_v<typename ViewTraits<T, P...>::array_layout, LayoutStride>);

typename src_type::array_layout layout = src.layout();
layout.dimension[src_type::rank] = Kokkos::dimension_scalar(src);
using src_type = View<T, P...>;

return dst_type(std::string(src.label()).append("_mirror"), layout);
}
auto layout = [&] () {
if constexpr ( ! std::is_same_v<typename ViewTraits<T, P...>::array_layout, LayoutStride>) {
return src.layout();
} else {
LayoutStride layout;

template< class T , class ... P >
inline
typename std::enable_if<
std::is_same< typename ViewTraits<T,P...>::specialize ,
Kokkos::Experimental::Impl::ViewMPVectorContiguous >::value &&
std::is_same< typename Kokkos::ViewTraits<T,P...>::array_layout,
Kokkos::LayoutStride >::value,
typename Kokkos::View<T,P...>::HostMirror>::type
create_mirror(const Kokkos::View<T,P...> & src)
{
typedef View<T,P...> src_type ;
typedef typename src_type::HostMirror dst_type ;

Kokkos::LayoutStride layout ;

layout.dimension[0] = src.extent(0);
layout.dimension[1] = src.extent(1);
layout.dimension[2] = src.extent(2);
layout.dimension[3] = src.extent(3);
layout.dimension[4] = src.extent(4);
layout.dimension[5] = src.extent(5);
layout.dimension[6] = src.extent(6);
layout.dimension[7] = src.extent(7);

layout.stride[0] = src.stride_0();
layout.stride[1] = src.stride_1();
layout.stride[2] = src.stride_2();
layout.stride[3] = src.stride_3();
layout.stride[4] = src.stride_4();
layout.stride[5] = src.stride_5();
layout.stride[6] = src.stride_6();
layout.stride[7] = src.stride_7();
for (int idx = 0; idx <= 7; ++idx) {
layout.dimension[idx] = src.extent(idx);
layout.stride [idx] = src.stride(idx);
}

layout.dimension[src_type::rank] = Kokkos::dimension_scalar(src);
return layout;
}
}();

layout.dimension[src_type::rank] = dimension_scalar(src);

const auto prop_copy = Impl::with_properties_if_unset(
arg_prop, std::string(src.label()).append("_mirror"));

return dst_type(std::string(src.label()).append("_mirror"), layout);
if constexpr (Impl::ViewCtorProp<ViewCtorArgs...>::has_memory_space){
return typename Impl::MirrorType<typename Impl::ViewCtorProp<ViewCtorArgs...>::memory_space, T, P ...>::view_type(prop_copy, layout);
} else {
return typename View<T, P...>::HostMirror(prop_copy, layout);
}
}

template<class Space, class T, class ... P, typename Enabled>
typename std::enable_if<
std::is_same< typename ViewTraits<T,P...>::specialize ,
Kokkos::Experimental::Impl::ViewMPVectorContiguous >::value,
typename Impl::MirrorType<Space,T,P ...>::view_type>::type
create_mirror(const Space& , const Kokkos::View<T,P...> & src)
template <class T, class... P>
inline auto create_mirror(
const View<T, P...>& src,
typename std::enable_if_t<
std::is_same_v<typename ViewTraits<T, P...>::specialize,
Experimental::Impl::ViewMPVectorContiguous>>*
)
{
typedef View<T,P...> src_type ;
typename src_type::array_layout layout = src.layout();
layout.dimension[src_type::rank] = Kokkos::dimension_scalar(src);
return typename Impl::MirrorType<Space,T,P ...>::view_type(src.label(),layout);
return create_mirror(view_alloc(), src);
}

template< class T , class ... P >
inline
typename std::enable_if<
std::is_same< typename ViewTraits<T,P...>::specialize ,
Kokkos::Experimental::Impl::ViewMPVectorContiguous >::value &&
!std::is_same< typename Kokkos::ViewTraits<T,P...>::array_layout,
Kokkos::LayoutStride >::value,
typename Kokkos::View<T,P...>::HostMirror>::type
create_mirror(Kokkos::Impl::WithoutInitializing_t wi,
const Kokkos::View<T,P...> & src)
template <class Space, class T, class... P, typename Enable>
inline auto create_mirror(
const Space& space,
const View<T, P...>& src,
typename std::enable_if_t<
std::is_same_v<typename ViewTraits<T, P...>::specialize,
Experimental::Impl::ViewMPVectorContiguous>>*
)
{
typedef View<T,P...> src_type ;
typedef typename src_type::HostMirror dst_type ;

typename src_type::array_layout layout = src.layout();
layout.dimension[src_type::rank] = Kokkos::dimension_scalar(src);

return dst_type(
Kokkos::view_alloc(std::string(src.label()).append("_mirror"), wi), layout);
return create_mirror(view_alloc(space), src);
}

template< class T , class ... P >
inline
typename std::enable_if<
std::is_same< typename ViewTraits<T,P...>::specialize ,
Kokkos::Experimental::Impl::ViewMPVectorContiguous >::value &&
std::is_same< typename Kokkos::ViewTraits<T,P...>::array_layout,
Kokkos::LayoutStride >::value,
typename Kokkos::View<T,P...>::HostMirror>::type
create_mirror(Kokkos::Impl::WithoutInitializing_t wi,
const Kokkos::View<T,P...> & src)
template <class T, class... P>
inline auto create_mirror(
Impl::WithoutInitializing_t wi,
const View<T, P...>& src,
typename std::enable_if_t<
std::is_same_v<typename ViewTraits<T, P...>::specialize,
Experimental::Impl::ViewMPVectorContiguous>>*
)
{
typedef View<T,P...> src_type ;
typedef typename src_type::HostMirror dst_type ;

Kokkos::LayoutStride layout ;

layout.dimension[0] = src.extent(0);
layout.dimension[1] = src.extent(1);
layout.dimension[2] = src.extent(2);
layout.dimension[3] = src.extent(3);
layout.dimension[4] = src.extent(4);
layout.dimension[5] = src.extent(5);
layout.dimension[6] = src.extent(6);
layout.dimension[7] = src.extent(7);

layout.stride[0] = src.stride_0();
layout.stride[1] = src.stride_1();
layout.stride[2] = src.stride_2();
layout.stride[3] = src.stride_3();
layout.stride[4] = src.stride_4();
layout.stride[5] = src.stride_5();
layout.stride[6] = src.stride_6();
layout.stride[7] = src.stride_7();

layout.dimension[src_type::rank] = Kokkos::dimension_scalar(src);
return create_mirror(view_alloc(wi), src);
}

return dst_type(
Kokkos::view_alloc(std::string(src.label()).append("_mirror"), wi), layout);
template <class Space, class T, class... P, typename Enable>
inline auto create_mirror(
Impl::WithoutInitializing_t wi,
const Space& space,
const View<T, P...>& src,
typename std::enable_if_t<
std::is_same_v<typename ViewTraits<T, P...>::specialize,
Experimental::Impl::ViewMPVectorContiguous>>*
)
{
return create_mirror(view_alloc(wi, space), src);
}

template<class Space, class T, class ... P, typename Enable>
typename std::enable_if<
std::is_same< typename ViewTraits<T,P...>::specialize ,
Kokkos::Experimental::Impl::ViewMPVectorContiguous >::value,
typename Impl::MirrorType<Space,T,P ...>::view_type>::type
create_mirror(Kokkos::Impl::WithoutInitializing_t wi,
const Space&, const Kokkos::View<T,P...> & src)
template <class T, class... P, class... ViewCtorArgs>
inline auto create_mirror_view(
const Impl::ViewCtorProp<ViewCtorArgs...>& arg_prop,
const Kokkos::View<T, P...>& src,
typename std::enable_if_t<
std::is_same_v<typename ViewTraits<T, P...>::specialize,
Experimental::Impl::ViewMPVectorContiguous>>*
)
{
typedef View<T,P...> src_type ;
typename src_type::array_layout layout = src.layout();
layout.dimension[src_type::rank] = Kokkos::dimension_scalar(src);
return typename Impl::MirrorType<Space,T,P ...>::view_type(
Kokkos::view_alloc(src.label(), wi), layout);
return Impl::create_mirror_view(src, arg_prop);
}

template <class Space, class T, class... P>
Expand Down
114 changes: 53 additions & 61 deletions packages/stokhos/src/sacado/kokkos/vector/Kokkos_View_MP_Vector_Fwd.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,67 +60,59 @@ namespace Kokkos {
// Declare overloads of create_mirror() so they are in scope
// Kokkos_Core.hpp is included below

template< class T , class ... P >
inline
typename std::enable_if<
std::is_same< typename ViewTraits<T,P...>::specialize ,
Kokkos::Experimental::Impl::ViewMPVectorContiguous >::value &&
!std::is_same< typename Kokkos::ViewTraits<T,P...>::array_layout,
Kokkos::LayoutStride >::value,
typename Kokkos::View<T,P...>::HostMirror>::type
create_mirror(const Kokkos::View<T,P...> & src);

template< class T , class ... P >
inline
typename std::enable_if<
std::is_same< typename ViewTraits<T,P...>::specialize ,
Kokkos::Experimental::Impl::ViewMPVectorContiguous >::value &&
std::is_same< typename Kokkos::ViewTraits<T,P...>::array_layout,
Kokkos::LayoutStride >::value,
typename Kokkos::View<T,P...>::HostMirror>::type
create_mirror(const Kokkos::View<T,P...> & src);

template<class Space, class T, class ... P,
typename Enable = std::enable_if_t<Kokkos::is_space<Space>::value>>
typename std::enable_if<
std::is_same< typename ViewTraits<T,P...>::specialize ,
Kokkos::Experimental::Impl::ViewMPVectorContiguous >::value,
typename Impl::MirrorType<Space,T,P ...>::view_type>::type
create_mirror(const Space&,
const Kokkos::View<T,P...> & src);

template< class T , class ... P >
inline
typename std::enable_if<
std::is_same< typename ViewTraits<T,P...>::specialize ,
Kokkos::Experimental::Impl::ViewMPVectorContiguous >::value &&
!std::is_same< typename Kokkos::ViewTraits<T,P...>::array_layout,
Kokkos::LayoutStride >::value,
typename Kokkos::View<T,P...>::HostMirror>::type
create_mirror(Kokkos::Impl::WithoutInitializing_t wi,
const Kokkos::View<T,P...> & src);

template< class T , class ... P >
inline
typename std::enable_if<
std::is_same< typename ViewTraits<T,P...>::specialize ,
Kokkos::Experimental::Impl::ViewMPVectorContiguous >::value &&
std::is_same< typename Kokkos::ViewTraits<T,P...>::array_layout,
Kokkos::LayoutStride >::value,
typename Kokkos::View<T,P...>::HostMirror>::type
create_mirror(Kokkos::Impl::WithoutInitializing_t wi,
const Kokkos::View<T,P...> & src);

template<class Space, class T, class ... P,
typename Enable = std::enable_if_t<Kokkos::is_space<Space>::value>>
typename std::enable_if<
std::is_same< typename ViewTraits<T,P...>::specialize ,
Kokkos::Experimental::Impl::ViewMPVectorContiguous >::value,
typename Impl::MirrorType<Space,T,P ...>::view_type>::type
create_mirror(
Kokkos::Impl::WithoutInitializing_t wi,
const Space&,
const Kokkos::View<T,P...> & src);
template <class T, class... P, class... ViewCtorArgs>
inline auto create_mirror(
const Impl::ViewCtorProp<ViewCtorArgs...>& arg_prop,
const View<T, P...>& src,
typename std::enable_if_t<
std::is_same_v<typename ViewTraits<T, P...>::specialize,
Experimental::Impl::ViewMPVectorContiguous>>* = nullptr
);

template <class T, class... P>
inline auto create_mirror(
const View<T, P...>& src,
typename std::enable_if_t<
std::is_same_v<typename ViewTraits<T, P...>::specialize,
Experimental::Impl::ViewMPVectorContiguous>>* = nullptr
);

template <class Space, class T, class... P, typename Enable = std::enable_if_t<Kokkos::is_space<Space>::value>>
inline auto create_mirror(
const Space& space,
const View<T, P...>& src,
typename std::enable_if_t<
std::is_same_v<typename ViewTraits<T, P...>::specialize,
Experimental::Impl::ViewMPVectorContiguous>>* = nullptr
);

template <class T, class... P>
inline auto create_mirror(
Impl::WithoutInitializing_t wi,
const View<T, P...>& src,
typename std::enable_if_t<
std::is_same_v<typename ViewTraits<T, P...>::specialize,
Experimental::Impl::ViewMPVectorContiguous>>* = nullptr
);

template <class Space, class T, class... P, typename Enable = std::enable_if_t<is_space<Space>::value>>
inline auto create_mirror(
Impl::WithoutInitializing_t wi,
const Space& space,
const View<T, P...>& src,
typename std::enable_if_t<
std::is_same_v<typename ViewTraits<T, P...>::specialize,
Experimental::Impl::ViewMPVectorContiguous>>* = nullptr
);

template <class T, class... P, class... ViewCtorArgs>
inline auto create_mirror_view(
const Impl::ViewCtorProp<ViewCtorArgs...>& arg_prop,
const Kokkos::View<T, P...>& src,
typename std::enable_if_t<
std::is_same_v<typename ViewTraits<T, P...>::specialize,
Experimental::Impl::ViewMPVectorContiguous>>* = nullptr
);

template <class Space, class T, class... P>
typename Impl::MirrorViewType<Space, T, P...>::view_type
Expand Down

0 comments on commit e3dadea

Please sign in to comment.