Merge pull request #20 from uliegecsm/slice

view(slice): take a subview without specifying all trailing Kokkos::ALL

CMakeLists.txt

@@ -65,6 +65,7 @@ target_sources(
         include/kokkos-utils/impl/type_traits.hpp
 
         include/kokkos-utils/view/extents.hpp
+        include/kokkos-utils/view/slice.hpp
 )
 
 target_include_directories(

include/kokkos-utils/concepts/View.hpp

@@ -16,7 +16,7 @@ concept View = Kokkos::is_view_v<T>;
 
 //! Specify that a type is a @c Kokkos::View of given rank @p Rank.
 template <typename T, std::size_t Rank>
-concept ViewOfRank = View<T> && T::rank == Rank;
+concept ViewOfRank = View<T> && T::rank() == Rank;
 
 //! Specify that a type is a modifiable @c Kokkos::View.
 template <typename T>

include/kokkos-utils/view/extents.hpp

@@ -18,7 +18,7 @@ namespace impl
 
 //! Implementation of @ref Kokkos::utils::view::extents.
 template <concepts::View ViewType, size_t... Ints>
-KOKKOS_INLINE_FUNCTION
+KOKKOS_FUNCTION
 constexpr auto extents(const ViewType& view, std::index_sequence<Ints...>)
 {
     Kokkos::Array<size_t, sizeof...(Ints)> extents {};
@@ -30,10 +30,10 @@ constexpr auto extents(const ViewType& view, std::index_sequence<Ints...>)
 
 //! Get all extents of @p view.
 template <concepts::View ViewType>
-KOKKOS_INLINE_FUNCTION
+KOKKOS_FUNCTION
 constexpr auto extents(const ViewType& view)
 {
-    constexpr auto rank = ViewType::rank;
+    constexpr auto rank = ViewType::rank();
     constexpr auto dims = std::make_index_sequence<rank>{};
     return impl::extents(view, dims);
 }

include/kokkos-utils/view/slice.hpp

@@ -0,0 +1,78 @@
+#ifndef KOKKOS_UTILS_VIEW_SLICE_HPP
+#define KOKKOS_UTILS_VIEW_SLICE_HPP
+
+#include "kokkos-utils/concepts/View.hpp"
+
+/**
+ * @file
+ *
+ * This file is a collection of helper functions similar to what can be found in
+ * the standard @c mdspan header, related to retrieving slices of a @c Kokkos::View.
+ *
+ * References:
+ *  - https://en.cppreference.com/w/cpp/container/mdspan
+ *  - https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2023/p2630r4.html
+ */
+
+namespace Kokkos::utils::view
+{
+
+namespace impl
+{
+
+template <typename ViewType, typename... Indices, size_t... AllsIndices>
+KOKKOS_FUNCTION
+constexpr auto slice(ViewType&& view, Indices&&... indices, std::index_sequence<AllsIndices...>)
+{
+    return Kokkos::subview(
+        std::forward<ViewType>(view),
+        std::forward<Indices>(indices)...,
+        std::get<AllsIndices>(std::array<Kokkos::ALL_t, sizeof...(AllsIndices)>{})...
+    );
+}
+
+} // namespace impl
+
+/**
+ * @brief Get a subview, given the first @p indices. The rest is filled with @c Kokkos::ALL.
+ *
+ * @note Though similar to @c std::submdspan, we think it it OK not to provide the remaining
+ *       slice specifiers if they are all @c Kokkos::ALL iif the user specifies what is the expected rank
+ *       of the input view. This ensures that the user will not face unexpected behaviors.
+ *       Therefore, we allow the following:
+ *       @code
+ *       auto sliced = slice<4>(view_of_rank_4, 1, 2);
+ *       @endcode
+ *       which is strictly equivalent to:
+ *       @code
+ *       auto sliced = subview(view_of_rank_4, 1, 2, Kokkos::ALL, Kokkos::ALL);
+ *       @endcode
+ */
+template <size_t Rank, typename ViewType, typename... Indices>
+requires concepts::ViewOfRank<std::remove_reference_t<ViewType>, Rank>
+KOKKOS_FUNCTION
+constexpr auto slice(ViewType&& view, Indices&&... indices)
+{
+    constexpr auto size = Rank - sizeof...(Indices);
+    return impl::slice<ViewType, Indices...>(
+        std::forward<ViewType>(view),
+        std::forward<Indices>(indices)...,
+        std::make_index_sequence<size>{}
+    );
+}
+
+//! @overload
+template <typename ViewType, typename... Indices>
+requires concepts::ViewOfRank<std::remove_reference_t<ViewType>, sizeof...(Indices)>
+KOKKOS_FUNCTION
+constexpr auto slice(ViewType&& view, Indices&&... indices)
+{
+    return slice<std::remove_reference_t<ViewType>::rank()>(
+        std::forward<ViewType>(view),
+        std::forward<Indices>(indices)...
+    );
+}
+
+} // namespace Kokkos::utils::view
+
+#endif // KOKKOS_UTILS_VIEW_SLICE_HPP

tests/atomics/test_InsertOp.cpp

@@ -12,7 +12,8 @@ using execution_space = Kokkos::DefaultExecutionSpace;
  *
  * @addtogroup unittests
  *
- * **Atomic insertion**
+ * Atomic insertion
+ * ----------------
  *
  * This group of tests check the behavior of our atomic insert operators.
  */

tests/concepts/test_DualView.cpp

@@ -10,7 +10,8 @@ using execution_space = Kokkos::DefaultExecutionSpace;
  *
  * @addtogroup unittests
  *
- * **Concepts related to @c Kokkos::DualView**
+ * Concepts related to @c Kokkos::DualView
+ * ---------------------------------------
  *
  * This group of tests check the behavior of our concepts related to @c Kokkos::DualView.
  */

tests/concepts/test_ExecutionSpace.cpp

@@ -11,7 +11,8 @@ using execution_space = Kokkos::DefaultExecutionSpace;
  *
  * @addtogroup unittests
  *
- * **Concepts related to @ref Kokkos execution space**
+ * Concepts related to @ref Kokkos execution space
+ * -----------------------------------------------
  *
  * This group of tests check the behavior of our concepts related to @ref Kokkos execution space.
  */

tests/concepts/test_MemorySpace.cpp

@@ -11,7 +11,8 @@ using execution_space = Kokkos::DefaultExecutionSpace;
  *
  * @addtogroup unittests
  *
- * **Concepts related to @ref Kokkos memory space**
+ * Concepts related to @ref Kokkos memory space
+ * --------------------------------------------
  *
  * This group of tests check the behavior of our concepts related to @ref Kokkos memory space.
  */

tests/concepts/test_Space.cpp

@@ -11,7 +11,8 @@ using execution_space = Kokkos::DefaultExecutionSpace;
  *
  * @addtogroup unittests
  *
- * **Concepts related to @ref Kokkos space**
+ * Concepts related to @ref Kokkos space
+ * -------------------------------------
  *
  * This group of tests check the behavior of our concepts related to @ref Kokkos space.
  */

tests/concepts/test_View.cpp

@@ -13,7 +13,8 @@ using execution_space = Kokkos::DefaultExecutionSpace;
  *
  * @addtogroup unittests
  *
- * **Concepts related to @c Kokkos::View**
+ * Concepts related to @c Kokkos::View
+ * -----------------------------------
  *
  * This group of tests check the behavior of our concepts related to @c Kokkos::View.
  */

tests/impl/test_type_list.cpp

@@ -5,7 +5,8 @@
 /**
  * @addtogroup unittests
  *
- * **Type list**
+ * Type list
+ * ---------
  *
  * This group of tests check the behavior of our utilities extending @c Kokkos::Impl::type_list,
  * which are found in @ref kokkos-utils/impl/type_list.hpp.

tests/impl/test_type_traits.cpp

@@ -8,7 +8,8 @@
 /**
  * @addtogroup unittests
  *
- * **Type traits**
+ * Type traits
+ * -----------
  *
  * This group of tests check the behavior of our utilities extending the standard @c type_traits.
  */

tests/view/CMakeLists.txt

@@ -1,2 +1,5 @@
 ### TEST : extents ###
 add_one_test(NAME extents)
+
+### TEST : slice ###
+add_one_test(NAME slice)

tests/view/test_extents.cpp

@@ -12,7 +12,8 @@ using execution_space = Kokkos::DefaultExecutionSpace;
  *
  * @addtogroup unittests
  *
- * **Get extents of a @c Kokkos::View**
+ * Get extents of a @c Kokkos::View
+ * --------------------------------
  *
  * This group of tests check the behavior of our helpers related to @c Kokkos::View that can be
  * found in @ref extents.hpp.
@@ -23,6 +24,14 @@ namespace Kokkos::utils::tests::view
 
 constexpr size_t dim_1 = 5, dim_2 = 3, dim_3 = 6, dim_4 = 9;
 
+//! @test Check that @ref view::extents works as expected for a rank-0 @c Kokkos::View.
+TEST(view, extents_rank_0)
+{
+    const Kokkos::View<double, execution_space> view_0("rank-0 view");
+
+    static_assert(utils::view::extents(view_0) == Kokkos::Array<size_t, 0>{});
+}
+
 //! @test Check that @ref view::extents works as expected for a rank-1 @c Kokkos::View.
 TEST(view, extents_rank_1)
 {

tests/view/test_slice.cpp

@@ -0,0 +1,119 @@
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+#include "Kokkos_Core.hpp"
+
+#include "kokkos-utils/view/extents.hpp"
+#include "kokkos-utils/view/slice.hpp"
+
+using execution_space = Kokkos::DefaultExecutionSpace;
+
+/**
+ * @file
+ *
+ * @addtogroup unittests
+ *
+ * Get slices of a @c Kokkos::View
+ * -------------------------------
+ *
+ * This group of tests check the behavior of our helpers related to @c Kokkos::View that can be
+ * found in @ref slice.hpp.
+ */
+
+namespace Kokkos::utils::tests::view
+{
+
+#define CHECK_THAT(__slice__, __rank__, __extents__)       \
+    {                                                      \
+        const auto tmp = __slice__;                        \
+        ASSERT_EQ(tmp.rank(), __rank__);                   \
+        ASSERT_EQ(utils::view::extents(tmp), __extents__); \
+    }
+
+constexpr size_t dim_1 = 5, dim_2 = 3, dim_3 = 6, dim_4 = 9;
+
+//! @test Check that @ref view::slice works as expected for a rank-1 @c Kokkos::View.
+TEST(view, slice_rank_1)
+{
+    const Kokkos::View<double[dim_1], execution_space> view_1_static("rank-1 view with static extent");
+    const Kokkos::View<double*      , execution_space> view_1_dynami("rank-1 view with dynamic extent", dim_1);
+
+    CHECK_THAT(utils::view::slice<1>(view_1_static), 1, Kokkos::Array{dim_1});
+    CHECK_THAT(utils::view::slice<1>(view_1_dynami), 1, Kokkos::Array{dim_1});
+
+    CHECK_THAT(utils::view::slice(view_1_static, 0), 0, (Kokkos::Array<size_t, 0>{}));
+    CHECK_THAT(utils::view::slice(view_1_dynami, 0), 0, (Kokkos::Array<size_t, 0>{}));
+}
+
+//! @test Check that @ref view::slice works as expected for a rank-2 @c Kokkos::View.
+TEST(view, slice_rank_2)
+{
+    const Kokkos::View<double[dim_1][dim_2], execution_space> view_2_static("rank-2 view with static extents");
+    const Kokkos::View<double**            , execution_space> view_2_dynami("rank-2 view with dynamic extents", dim_1, dim_2);
+    const Kokkos::View<double*[dim_2]      , execution_space> view_2_mixed ("rank-2 view with mixed extents"  , dim_1);
+
+    CHECK_THAT(utils::view::slice<2>(view_2_static, 0), 1, Kokkos::Array{dim_2});
+    CHECK_THAT(utils::view::slice<2>(view_2_dynami, 0), 1, Kokkos::Array{dim_2});
+    CHECK_THAT(utils::view::slice<2>(view_2_mixed , 0), 1, Kokkos::Array{dim_2});
+}
+
+//! @test Check that @ref view::slice works as expected for a rank-3 @c Kokkos::View.
+TEST(view, slice_rank_3)
+{
+    const Kokkos::View<double[dim_1][dim_2][dim_3], execution_space> view_3_static("rank-3 view of static extents");
+    const Kokkos::View<double***                  , execution_space> view_3_dynami("rank-3 view of dynamic extents", dim_1, dim_2, dim_3);
+    const Kokkos::View<double*[dim_2][dim_3]      , execution_space> view_3_mixed ("rank-3 view of mixed extents"  , dim_1);
+
+    CHECK_THAT(utils::view::slice<3>(view_3_static, 0), 2, (Kokkos::Array{dim_2, dim_3}));
+    CHECK_THAT(utils::view::slice<3>(view_3_dynami, 0), 2, (Kokkos::Array{dim_2, dim_3}));
+    CHECK_THAT(utils::view::slice<3>(view_3_mixed , 0), 2, (Kokkos::Array{dim_2, dim_3}));
+
+    CHECK_THAT(utils::view::slice<3>(view_3_mixed , 0, Kokkos::ALL), 2, (Kokkos::Array{dim_2, dim_3}));
+}
+
+//! @test Check that @ref view::slice works as expected for a rank-4 @c Kokkos::View.
+TEST(view, slice_rank_4)
+{
+    const Kokkos::View<double[dim_1][dim_2][dim_3][dim_4], execution_space> view_4_static("rank-4 view of static extents");
+
+    CHECK_THAT(utils::view::slice<4>(view_4_static, 0, Kokkos::ALL), 3, (Kokkos::Array{dim_2, dim_3, dim_4}));
+
+    CHECK_THAT(utils::view::slice(view_4_static, 0, 0, 0, 0), 0, (Kokkos::Array<size_t, 0>{}));
+
+    CHECK_THAT(utils::view::slice<4>(view_4_static, Kokkos::make_pair(0, 2)), 4, (Kokkos::Array<size_t, 4>{2, dim_2, dim_3, dim_4}));
+}
+
+template <concepts::ViewOfRank<4> ViewType>
+bool test_slice_on_device(const ViewType& view, const size_t expected_size)
+{
+    bool result = false;
+
+    Kokkos::parallel_reduce(
+        Kokkos::RangePolicy<execution_space>(0, 1),
+        KOKKOS_LAMBDA(const int, bool& result)
+        {
+            const auto subview = Kokkos::subview      (view, 0, Kokkos::ALL, Kokkos::ALL, Kokkos::ALL);
+            const auto slice   = utils::view::slice<4>(view, 0);
+            result = (
+                subview.size() == expected_size && subview.rank() == 3 &&
+                slice  .size() == expected_size && slice  .rank() == 3 &&
+                slice.extent(0) == dim_2 &&
+                slice.extent(1) == dim_3 &&
+                slice.extent(2) == dim_4
+            );
+        },
+        Kokkos::LAnd<bool>(result)
+    );
+
+    return result;
+}
+
+//! @test Check that @ref view::slice works on device.
+TEST(view, slice_on_device)
+{
+    const Kokkos::View<double[dim_1][dim_2][dim_3][dim_4], execution_space> view_4_static("rank-4 view of static extents");
+
+    ASSERT_TRUE(test_slice_on_device(view_4_static, dim_2 * dim_3 * dim_4));
+}
+
+} // namespace Kokkos::utils::tests::view