Skip to content

Commit

Permalink
Merge pull request #5878 from masterleinad/aligned_subview
Browse files Browse the repository at this point in the history 
Remove Aligned memory trait when creating subviews and deprecate subview overload
  • Loading branch information
@crtrott
crtrott committed Feb 18, 2023
2 parents 3d2dc6a + 4ca0340 commit af89aa7
Show file tree
Hide file tree
Showing 5 changed files with 103 additions and 62 deletions.
24 changes: 13 additions & 11 deletions core/src/Kokkos_MemoryTraits.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,17 +48,19 @@ template <unsigned T>
struct MemoryTraits {
//! Tag this class as a kokkos memory traits:
using memory_traits = MemoryTraits<T>;
enum : bool {
is_unmanaged = (unsigned(0) != (T & unsigned(Kokkos::Unmanaged)))
};
enum : bool {
is_random_access = (unsigned(0) != (T & unsigned(Kokkos::RandomAccess)))
};
enum : bool { is_atomic = (unsigned(0) != (T & unsigned(Kokkos::Atomic))) };
enum : bool {
is_restrict = (unsigned(0) != (T & unsigned(Kokkos::Restrict)))
};
enum : bool { is_aligned = (unsigned(0) != (T & unsigned(Kokkos::Aligned))) };

static constexpr unsigned impl_value = T;

static constexpr bool is_unmanaged =
(unsigned(0) != (T & unsigned(Kokkos::Unmanaged)));
static constexpr bool is_random_access =
(unsigned(0) != (T & unsigned(Kokkos::RandomAccess)));
static constexpr bool is_atomic =
(unsigned(0) != (T & unsigned(Kokkos::Atomic)));
static constexpr bool is_restrict =
(unsigned(0) != (T & unsigned(Kokkos::Restrict)));
static constexpr bool is_aligned =
(unsigned(0) != (T & unsigned(Kokkos::Aligned)));
};

} // namespace Kokkos
Expand Down
63 changes: 43 additions & 20 deletions core/src/Kokkos_View.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ static_assert(false,
#include <View/Hooks/Kokkos_ViewHooks.hpp>

#include <impl/Kokkos_Tools.hpp>
#include <impl/Kokkos_Utilities.hpp>

#ifdef KOKKOS_ENABLE_IMPL_MDSPAN
#include <View/MDSpan/Kokkos_MDSpan_Extents.hpp>
Expand Down Expand Up @@ -1745,44 +1746,66 @@ void apply_to_view_of_static_rank(Function&& f, View<Args...> a) {
//----------------------------------------------------------------------------
//----------------------------------------------------------------------------

template <class V, class... Args>
using Subview =
typename Kokkos::Impl::ViewMapping<void /* deduce subview type from source
view traits */
,
typename V::traits, Args...>::type;
namespace Impl {
template <class ValueType, class TypeList>
struct TypeListToViewTraits;

template <class ValueType, class... Properties>
struct TypeListToViewTraits<ValueType, Kokkos::Impl::type_list<Properties...>> {
using type = ViewTraits<ValueType, Properties...>;
};

// It is not safe to assume that subviews of views with the Aligned memory trait
// are also aligned. Hence, just remove that attribute for subviews.
template <class D, class... P>
struct RemoveAlignedMemoryTrait {
private:
using type_list_in = Kokkos::Impl::type_list<P...>;
using memory_traits = typename ViewTraits<D, P...>::memory_traits;
using type_list_in_wo_memory_traits =
typename Kokkos::Impl::type_list_remove_first<memory_traits,
type_list_in>::type;
using new_memory_traits =
Kokkos::MemoryTraits<memory_traits::impl_value & ~Kokkos::Aligned>;
using new_type_list = typename Kokkos::Impl::concat_type_list<
type_list_in_wo_memory_traits,
Kokkos::Impl::type_list<new_memory_traits>>::type;

public:
using type = typename TypeListToViewTraits<D, new_type_list>::type;
};
} // namespace Impl

template <class D, class... P, class... Args>
KOKKOS_INLINE_FUNCTION
typename Kokkos::Impl::ViewMapping<void /* deduce subview type from source
view traits */
,
ViewTraits<D, P...>, Args...>::type
subview(const View<D, P...>& src, Args... args) {
KOKKOS_INLINE_FUNCTION auto subview(const View<D, P...>& src, Args... args) {
static_assert(View<D, P...>::rank == sizeof...(Args),
"subview requires one argument for each source View rank");

return typename Kokkos::Impl::ViewMapping<
void /* deduce subview type from source view traits */
,
ViewTraits<D, P...>, Args...>::type(src, args...);
typename Impl::RemoveAlignedMemoryTrait<D, P...>::type,
Args...>::type(src, args...);
}

#ifdef KOKKOS_ENABLE_DEPRECATED_CODE_4
template <class MemoryTraits, class D, class... P, class... Args>
KOKKOS_INLINE_FUNCTION typename Kokkos::Impl::ViewMapping<
void /* deduce subview type from source view traits */
,
ViewTraits<D, P...>, Args...>::template apply<MemoryTraits>::type
subview(const View<D, P...>& src, Args... args) {
KOKKOS_DEPRECATED KOKKOS_INLINE_FUNCTION auto subview(const View<D, P...>& src,
Args... args) {
static_assert(View<D, P...>::rank == sizeof...(Args),
"subview requires one argument for each source View rank");
static_assert(Kokkos::is_memory_traits<MemoryTraits>::value);

return typename Kokkos::Impl::ViewMapping<
void /* deduce subview type from source view traits */
,
ViewTraits<D, P...>,
Args...>::template apply<MemoryTraits>::type(src, args...);
typename Impl::RemoveAlignedMemoryTrait<D, P..., MemoryTraits>::type,
Args...>::type(src, args...);
}
#endif

template <class V, class... Args>
using Subview = decltype(subview(std::declval<V>(), std::declval<Args>()...));

} /* namespace Kokkos */

Expand Down
30 changes: 6 additions & 24 deletions core/src/impl/Kokkos_ViewMapping.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2712,14 +2712,8 @@ struct ViewDataHandle<
Traits,
std::enable_if_t<(std::is_void<typename Traits::specialize>::value &&
(!Traits::memory_traits::is_aligned) &&
Traits::memory_traits::is_restrict
#ifdef KOKKOS_ENABLE_CUDA
&& (!(std::is_same<typename Traits::memory_space,
Kokkos::CudaSpace>::value ||
std::is_same<typename Traits::memory_space,
Kokkos::CudaUVMSpace>::value))
#endif
&& (!Traits::memory_traits::is_atomic))>> {
Traits::memory_traits::is_restrict &&
(!Traits::memory_traits::is_atomic))>> {
using value_type = typename Traits::value_type;
using handle_type = typename Traits::value_type* KOKKOS_RESTRICT;
using return_type = typename Traits::value_type& KOKKOS_RESTRICT;
Expand All @@ -2742,14 +2736,8 @@ struct ViewDataHandle<
Traits,
std::enable_if_t<(std::is_void<typename Traits::specialize>::value &&
Traits::memory_traits::is_aligned &&
(!Traits::memory_traits::is_restrict)
#ifdef KOKKOS_ENABLE_CUDA
&& (!(std::is_same<typename Traits::memory_space,
Kokkos::CudaSpace>::value ||
std::is_same<typename Traits::memory_space,
Kokkos::CudaUVMSpace>::value))
#endif
&& (!Traits::memory_traits::is_atomic))>> {
(!Traits::memory_traits::is_restrict) &&
(!Traits::memory_traits::is_atomic))>> {
using value_type = typename Traits::value_type;
// typedef work-around for intel compilers error #3186: expected typedef
// declaration
Expand Down Expand Up @@ -2787,14 +2775,8 @@ struct ViewDataHandle<
Traits,
std::enable_if_t<(std::is_void<typename Traits::specialize>::value &&
Traits::memory_traits::is_aligned &&
Traits::memory_traits::is_restrict
#ifdef KOKKOS_ENABLE_CUDA
&& (!(std::is_same<typename Traits::memory_space,
Kokkos::CudaSpace>::value ||
std::is_same<typename Traits::memory_space,
Kokkos::CudaUVMSpace>::value))
#endif
&& (!Traits::memory_traits::is_atomic))>> {
Traits::memory_traits::is_restrict &&
(!Traits::memory_traits::is_atomic))>> {
using value_type = typename Traits::value_type;
// typedef work-around for intel compilers error #3186: expected typedef
// declaration
Expand Down
3 changes: 1 addition & 2 deletions core/unit_test/TestViewMapping_subview.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,8 +100,7 @@ struct TestViewMappingSubview {

KOKKOS_INLINE_FUNCTION
void operator()(const int, long& error_count) const {
auto Ad = Kokkos::subview<Kokkos::MemoryUnmanaged>(
Aa, Kokkos::pair<int, int>(1, AN - 1));
auto Ad = Kokkos::subview(Aa, Kokkos::pair<int, int>(1, AN - 1));

for (int i = 1; i < AN - 1; ++i)
if (&Aa[i] != &Ab[i - 1]) ++error_count;
Expand Down
45 changes: 40 additions & 5 deletions core/unit_test/TestViewSubview.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2110,15 +2110,33 @@ void test_unmanaged_subview_reset() {
template <std::underlying_type_t<Kokkos::MemoryTraitsFlags> MTF>
struct TestSubviewMemoryTraitsConstruction {
void operator()() const noexcept {
using view_type = Kokkos::View<double*, Kokkos::HostSpace>;
using size_type = view_type::size_type;
using memory_traits_type = Kokkos::MemoryTraits<MTF>;

view_type v("v", 7);
using view_type =
Kokkos::View<double*, Kokkos::HostSpace, memory_traits_type>;
using size_type = typename view_type::size_type;

// Create a managed View first and then apply the desired memory traits to
// an unmanaged version of it since a managed View can't use the Unmanaged
// trait.
Kokkos::View<double*, Kokkos::HostSpace> v_original("v", 7);
view_type v(v_original.data(), v_original.size());
for (size_type i = 0; i != v.size(); ++i) v[i] = static_cast<double>(i);

std::pair<int, int> range(3, 5);
auto sv = Kokkos::subview<memory_traits_type>(v, range);
auto sv = Kokkos::subview(v, range);

// check that the subview memory traits are the same as the original view
// (with the Aligned trait stripped).
using view_memory_traits = typename decltype(v)::memory_traits;
using subview_memory_traits = typename decltype(sv)::memory_traits;
static_assert(view_memory_traits::impl_value ==
memory_traits_type::impl_value);
if constexpr (memory_traits_type::is_aligned)
static_assert(subview_memory_traits::impl_value + Kokkos::Aligned ==
memory_traits_type::impl_value);
else
static_assert(subview_memory_traits::impl_value ==
memory_traits_type::impl_value);

ASSERT_EQ(2u, sv.size());
EXPECT_EQ(3., sv[0]);
Expand All @@ -2132,6 +2150,7 @@ inline void test_subview_memory_traits_construction() {
// RandomAccess (2)
// Atomic (4)
// Restricted (8)
// Aligned (16)
TestSubviewMemoryTraitsConstruction<0>()();
TestSubviewMemoryTraitsConstruction<1>()();
TestSubviewMemoryTraitsConstruction<2>()();
Expand All @@ -2148,6 +2167,22 @@ inline void test_subview_memory_traits_construction() {
TestSubviewMemoryTraitsConstruction<13>()();
TestSubviewMemoryTraitsConstruction<14>()();
TestSubviewMemoryTraitsConstruction<15>()();
TestSubviewMemoryTraitsConstruction<16>()();
TestSubviewMemoryTraitsConstruction<17>()();
TestSubviewMemoryTraitsConstruction<18>()();
TestSubviewMemoryTraitsConstruction<19>()();
TestSubviewMemoryTraitsConstruction<20>()();
TestSubviewMemoryTraitsConstruction<21>()();
TestSubviewMemoryTraitsConstruction<22>()();
TestSubviewMemoryTraitsConstruction<23>()();
TestSubviewMemoryTraitsConstruction<24>()();
TestSubviewMemoryTraitsConstruction<25>()();
TestSubviewMemoryTraitsConstruction<26>()();
TestSubviewMemoryTraitsConstruction<27>()();
TestSubviewMemoryTraitsConstruction<28>()();
TestSubviewMemoryTraitsConstruction<29>()();
TestSubviewMemoryTraitsConstruction<30>()();
TestSubviewMemoryTraitsConstruction<31>()();
}

//----------------------------------------------------------------------------
Expand Down

0 comments on commit af89aa7

Please sign in to comment.