forked from kokkos/kokkos
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
kokkos#6805: add initial converting constructor from mdspan to view
- Loading branch information
Showing
6 changed files
with
273 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,93 @@ | ||
| //@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_IMPL_PUBLIC_INCLUDE | ||
| static_assert(false, | ||
| "Including non-public Kokkos header files is not allowed."); | ||
| #endif | ||
|
|
||
| #ifndef KOKKOS_EXPERIMENTAL_MDSPAN_LAYOUT_HPP | ||
| #define KOKKOS_EXPERIMENTAL_MDSPAN_LAYOUT_HPP | ||
|
|
||
| #include "Kokkos_MDSpan_Extents.hpp" | ||
| #include <Kokkos_Core_fwd.hpp> | ||
| #include <impl/Kokkos_ViewDataAnalysis.hpp> | ||
|
|
||
| namespace Kokkos::Experimental::Impl { | ||
| template <class Layout> | ||
| struct ArrayLayoutFromLayout; | ||
|
|
||
| template <std::size_t padding_value> | ||
| struct ArrayLayoutFromLayout<Experimental::layout_left_padded<padding_value>> { | ||
| using type = Kokkos::LayoutLeft; | ||
| static constexpr std::integral_constant<unsigned, | ||
| static_cast<unsigned>(padding_value)> | ||
| padding = {}; | ||
| }; | ||
|
|
||
| template <std::size_t padding_value> | ||
| struct ArrayLayoutFromLayout<Experimental::layout_right_padded<padding_value>> { | ||
| using type = Kokkos::LayoutRight; | ||
| static constexpr std::integral_constant<unsigned, | ||
| static_cast<unsigned>(padding_value)> | ||
| padding = {}; | ||
| }; | ||
|
|
||
| template <class T, class Extents, class Layout> | ||
| struct ViewOffsetFromExtents { | ||
| using value_type = T; | ||
| using data_type = typename DataTypeFromExtents<value_type, Extents>::type; | ||
| using array_layout = typename ArrayLayoutFromLayout<Layout>::type; | ||
| using data_analysis = | ||
| Kokkos::Impl::ViewDataAnalysis<data_type, array_layout, value_type>; | ||
| using type = | ||
| Kokkos::Impl::ViewOffset<typename data_analysis::dimension, array_layout>; | ||
| }; | ||
|
|
||
| template <class ElementType, class Extents, class LayoutPolicy, | ||
| class AccessorPolicy> | ||
| auto array_layout_from_mdspan( | ||
| const mdspan<ElementType, Extents, LayoutPolicy, AccessorPolicy> &mds) { | ||
| using layout_type = typename ArrayLayoutFromLayout<LayoutPolicy>::type; | ||
| const auto &ext = mds.extents(); | ||
|
|
||
| static constexpr auto rank = Extents::rank(); | ||
|
|
||
| static_assert(rank <= ARRAY_LAYOUT_MAX_RANK, | ||
| "Unsupported rank for mdspan (must be <= 8)"); | ||
| return layout_type{ | ||
| rank > 0 ? dimension_from_extent(ext, 0) : KOKKOS_IMPL_CTOR_DEFAULT_ARG, | ||
| rank > 1 ? dimension_from_extent(ext, 1) : KOKKOS_IMPL_CTOR_DEFAULT_ARG, | ||
| rank > 2 ? dimension_from_extent(ext, 2) : KOKKOS_IMPL_CTOR_DEFAULT_ARG, | ||
| rank > 3 ? dimension_from_extent(ext, 3) : KOKKOS_IMPL_CTOR_DEFAULT_ARG, | ||
| rank > 4 ? dimension_from_extent(ext, 4) : KOKKOS_IMPL_CTOR_DEFAULT_ARG, | ||
| rank > 5 ? dimension_from_extent(ext, 5) : KOKKOS_IMPL_CTOR_DEFAULT_ARG, | ||
| rank > 6 ? dimension_from_extent(ext, 6) : KOKKOS_IMPL_CTOR_DEFAULT_ARG, | ||
| rank > 7 ? dimension_from_extent(ext, 7) : KOKKOS_IMPL_CTOR_DEFAULT_ARG}; | ||
| } | ||
|
|
||
| template <class ElementType, class Extents, class LayoutPolicy, | ||
| class AccessorPolicy> | ||
| auto view_offset_from_mdspan( | ||
| const mdspan<ElementType, Extents, LayoutPolicy, AccessorPolicy> &mds) { | ||
| using offset_type = | ||
| typename ViewOffsetFromExtents<ElementType, Extents, LayoutPolicy>::type; | ||
| static constexpr auto padding = ArrayLayoutFromLayout<LayoutPolicy>::padding; | ||
| return offset_type(padding, array_layout_from_mdspan(mds)); | ||
| }; | ||
| } // namespace Kokkos::Experimental::Impl | ||
|
|
||
| #endif // KOKKOS_EXPERIMENTAL_MDSPAN_LAYOUT_HPP |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,116 @@ | ||
| //@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 | ||
|
|
||
| #include <gtest/gtest.h> | ||
|
|
||
| #define KOKKOS_IMPL_TEST_ACCESS | ||
|
|
||
| #include <Kokkos_Core.hpp> | ||
|
|
||
| template <class T, class ExecutionSpace> | ||
| struct TestViewMDSpanConversion { | ||
| using value_type = T; | ||
|
|
||
| template <class MDSpanLayout, class KokkosLayout, class DataType, | ||
| class MDSpanExtents> | ||
| static void test_conversion_from_mdspan(Kokkos::View<DataType> ref, | ||
| const MDSpanExtents &exts) { | ||
| using view_type = Kokkos::View<DataType, KokkosLayout, | ||
| Kokkos::MemoryTraits<Kokkos::Unmanaged>>; | ||
| using mdspan_type = Kokkos::mdspan<value_type, MDSpanExtents, MDSpanLayout>; | ||
|
|
||
| using converted_layout_type = | ||
| typename Kokkos::Experimental::Impl::ArrayLayoutFromLayout< | ||
| MDSpanLayout>::type; | ||
| static_assert( | ||
| std::is_constructible_v<typename view_type::traits::array_layout, | ||
| const converted_layout_type &>); | ||
| static_assert( | ||
| view_type::template mdspan_conversion_constraints< | ||
| MDSpanExtents, MDSpanLayout, Kokkos::default_accessor<value_type>>); | ||
|
|
||
| // Manually create an mdspan from ref so we have a valid pointer to play | ||
| // with | ||
| auto mds = mdspan_type{ref.data(), exts}; | ||
|
|
||
| auto test_view = view_type(mds); | ||
|
|
||
| ASSERT_EQ(test_view.data(), ref.data()); | ||
| ASSERT_EQ(test_view.data(), mds.data_handle()); | ||
| for (std::size_t r = 0; r < mdspan_type::rank(); ++r) { | ||
| ASSERT_EQ(test_view.extent(r), ref.extent(r)); | ||
| ASSERT_EQ(test_view.extent(r), exts.extent(r)); | ||
| } | ||
| } | ||
|
|
||
| static void run_test() { | ||
| static_assert(std::is_same_v< | ||
| typename Kokkos::Experimental::Impl::ArrayLayoutFromLayout< | ||
| Kokkos::Experimental::layout_left_padded<sizeof( | ||
| value_type)>>::type, | ||
| Kokkos::LayoutLeft>); | ||
| static_assert(std::is_same_v< | ||
| typename Kokkos::Experimental::Impl::ArrayLayoutFromLayout< | ||
| Kokkos::Experimental::layout_right_padded<sizeof( | ||
| value_type)>>::type, | ||
| Kokkos::LayoutRight>); | ||
|
|
||
| test_conversion_from_mdspan< | ||
| Kokkos::Experimental::layout_left_padded<sizeof(value_type)>, | ||
| Kokkos::LayoutLeft>( | ||
| Kokkos::View<double *>("ref", 7), | ||
| Kokkos::extents<std::size_t, Kokkos::dynamic_extent>(7)); | ||
| test_conversion_from_mdspan< | ||
| Kokkos::Experimental::layout_left_padded<sizeof(value_type)>, | ||
| Kokkos::LayoutLeft>(Kokkos::View<double[7]>("ref"), | ||
| Kokkos::extents<std::size_t, 7>()); | ||
| test_conversion_from_mdspan< | ||
| Kokkos::Experimental::layout_left_padded<sizeof(value_type)>, | ||
| Kokkos::LayoutLeft>( | ||
| Kokkos::View<double[7]>("ref"), | ||
| Kokkos::extents<std::size_t, Kokkos::dynamic_extent>(7)); | ||
| test_conversion_from_mdspan< | ||
| Kokkos::Experimental::layout_left_padded<sizeof(value_type)>, | ||
| Kokkos::LayoutLeft>(Kokkos::View<double *>("ref", 7), | ||
| Kokkos::extents<std::size_t, 7>()); | ||
|
|
||
| test_conversion_from_mdspan< | ||
| Kokkos::Experimental::layout_left_padded<sizeof(value_type)>, | ||
| Kokkos::LayoutLeft>(Kokkos::View<double **>("ref", 7, 3), | ||
| Kokkos::dextents<std::size_t, 2>(7, 3)); | ||
| test_conversion_from_mdspan< | ||
| Kokkos::Experimental::layout_left_padded<sizeof(value_type)>, | ||
| Kokkos::LayoutLeft>(Kokkos::View<double[7][3]>("ref"), | ||
| Kokkos::extents<std::size_t, 7, 3>()); | ||
| test_conversion_from_mdspan< | ||
| Kokkos::Experimental::layout_left_padded<sizeof(value_type)>, | ||
| Kokkos::LayoutLeft>(Kokkos::View<double[7][3]>("ref"), | ||
| Kokkos::extents<std::size_t, Kokkos::dynamic_extent, | ||
| Kokkos::dynamic_extent>(7, 3)); | ||
| test_conversion_from_mdspan< | ||
| Kokkos::Experimental::layout_left_padded<sizeof(value_type)>, | ||
| Kokkos::LayoutLeft>(Kokkos::View<double **>("ref", 7, 3), | ||
| Kokkos::extents<std::size_t, 7, 3>()); | ||
| } | ||
| }; | ||
|
|
||
| namespace Test { | ||
|
|
||
| TEST(TEST_CATEGORY, view_mdspan_conversion) { | ||
| TestViewMDSpanConversion<double, TEST_EXECSPACE>::run_test(); | ||
| } | ||
|
|
||
| } // namespace Test |