Add SubView class

include/RAJA/RAJA.hpp

@@ -118,7 +118,7 @@
 #include "RAJA/util/StaticLayout.hpp"
 #include "RAJA/util/IndexLayout.hpp"
 #include "RAJA/util/View.hpp"
-
+#include "RAJA/util/SubView.hpp"
 
 //
 // View for sequences of objects

include/RAJA/util/SubView.hpp

@@ -0,0 +1,144 @@
+/*!
+ ******************************************************************************
+ *
+ * \file
+ *
+ * \brief   RAJA header file defining the SubView class
+ *
+ ******************************************************************************
+ */
+
+//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
+// Copyright (c) 2016-25, Lawrence Livermore National Security, LLC
+// and RAJA project contributors. See the RAJA/LICENSE file for details.
+//
+// SPDX-License-Identifier: (BSD-3-Clause)
+//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
+
+#ifndef RAJA_SUBVIEW_HPP
+#define RAJA_SUBVIEW_HPP
+
+#include "RAJA/util/for_each.hpp"
+#include "RAJA/util/types.hpp"
+#include "camp/number.hpp"
+#include "camp/tuple.hpp"
+#include "camp/array.hpp"
+
+namespace RAJA
+{
+
+// Slice descriptors
+
+// template<typename SliceType>
+// RAJA_INLINE RAJA_HOST_DEVICE constexpr IndexType map_index(IndexType& idx) const {
+//     return start + idx;
+// }
+
+template<typename IndexType = Index_type>
+struct RangeSlice {
+    IndexType start, end; 
+
+    static constexpr bool reduces_dimension = false;
+
+    RAJA_INLINE RAJA_HOST_DEVICE constexpr IndexType map_index(IndexType& idx) const {
+        return start + idx;
+    }
+};
+
+template<typename IndexType = Index_type>
+struct FixedSlice { 
+    IndexType idx; 
+
+    static constexpr bool reduces_dimension = true;
+
+    RAJA_INLINE RAJA_HOST_DEVICE constexpr IndexType map_index(IndexType&) const {
+        return idx;
+    }
+};
+
+template<typename IndexType = Index_type>
+struct NoSlice { 
+    static constexpr bool reduces_dimension = false;
+
+    RAJA_INLINE RAJA_HOST_DEVICE constexpr IndexType map_index(IndexType& idx) const {
+        return idx;
+    }
+};
+
+template <typename T, size_t N, RAJA::Index_type... Is>
+RAJA_INLINE RAJA_HOST_DEVICE constexpr auto array_to_tuple_impl(const camp::array<T, N>& arr, camp::idx_seq<Is...>) {
+    return camp::make_tuple(arr[Is]...);
+}
+
+template <typename T, size_t N>
+RAJA_INLINE RAJA_HOST_DEVICE constexpr auto array_to_tuple(const camp::array<T, N>& arr) {
+    return array_to_tuple_impl(arr, camp::make_idx_seq_t<N>{});
+}
+
+template <typename ViewType, typename SliceTypes, typename IndexType = Index_type>
+class SubView;
+
+template <typename ViewType, typename IndexType, typename... Slices>
+class SubView<ViewType, camp::list<Slices...>, IndexType> {
+    ViewType view_;
+    camp::tuple<Slices...> slices_;
+    std::array<IndexType, sizeof...(Slices)> map_;
+
+    RAJA_INLINE RAJA_HOST_DEVICE constexpr auto make_subview_index_map() {
+        size_t sub_idx = 0;
+        std::array<IndexType, sizeof...(Slices)> map;
+
+        for_each_tuple_index( slices_, 
+            [&](auto slice, auto index) { 
+                map[index] = decltype(slice)::reduces_dimension ? -1 : sub_idx++;
+            });
+
+        return map;
+    }
+
+public:
+
+    RAJA_INLINE RAJA_HOST_DEVICE constexpr SubView(ViewType view, Slices... slices)
+        : view_(view), slices_(slices...), map_(make_subview_index_map()) { }
+
+    RAJA_INLINE RAJA_HOST_DEVICE constexpr void set_slices(Slices... slices) {
+        slices_ = camp::tuple<Slices...>(slices...);
+    }
+
+    template<IndexType Index, typename Slice>
+    RAJA_INLINE RAJA_HOST_DEVICE constexpr void set_slice(Slice slice) {
+        camp::get<Index>(slices_) = slice;
+    }
+
+    RAJA_INLINE RAJA_HOST_DEVICE constexpr auto& get_slices() {
+        return slices_;
+    }
+
+    template<IndexType Index>
+    RAJA_INLINE RAJA_HOST_DEVICE constexpr auto& get_slice() {
+        return camp::get<Index>(slices_);
+    }
+
+    template <typename... Idxs>
+    RAJA_INLINE RAJA_HOST_DEVICE constexpr IndexType operator()(Idxs... idxs) const {
+        constexpr size_t nidx = ((Slices::reduces_dimension == false ? 1 : 0) + ...);
+        static_assert(sizeof...(idxs) == nidx, "Wrong number of indices for subview");
+
+        camp::array<Index_type, nidx> arr{idxs...};
+        camp::array<Index_type, sizeof...(Slices)> parent_indices;
+
+        for_each_tuple_index( slices_,
+            [&](auto slice, auto index) { 
+                parent_indices[index] = slice.map_index(arr[map_[index]]); 
+            });
+
+        return camp::apply(view_, array_to_tuple(parent_indices));
+    }
+};
+
+template <typename ViewType, typename... Slices>
+SubView(ViewType, Slices...) -> SubView<ViewType, camp::list<Slices...>>;
+
+}  // namespace RAJA
+
+#endif

include/RAJA/util/for_each.hpp

@@ -24,6 +24,9 @@
 #include <type_traits>
 
 #include "camp/list.hpp"
+#include "camp/concepts.hpp"
+#include "camp/number.hpp"
+#include "camp/tuple.hpp"
 
 #include "RAJA/pattern/detail/algorithm.hpp"
 
@@ -36,6 +39,20 @@ namespace RAJA
 namespace detail
 {
 
+// compile time expansion applying func to each of the indices
+RAJA_SUPPRESS_HD_WARN
+template<typename UnaryFunc, typename IndexType, IndexType... Is>
+RAJA_HOST_DEVICE RAJA_INLINE UnaryFunc for_each_index(camp::idx_seq<Is...>,
+                                                      UnaryFunc func)
+{
+  // braced init lists are evaluated in order
+  // create integral_constant type to allow UnaryFunc arguments to be used in compile-time context
+  int seq_unused_array[] = {0, (func(std::integral_constant<std::size_t, Is>{}), 0)...};
+  RAJA_UNUSED_VAR(seq_unused_array);
+
+  return func;
+}
+
 // runtime loop applying func to each element in the range in order
 RAJA_SUPPRESS_HD_WARN
 template<typename Iter, typename UnaryFunc>
@@ -79,8 +96,31 @@ RAJA_HOST_DEVICE RAJA_INLINE UnaryFunc for_each_tuple(Tuple&& t,
   return func;
 }
 
+// compile time expansion applying func to a each type in the tuple in order
+RAJA_SUPPRESS_HD_WARN
+template<typename Tuple, typename BinaryFunc, camp::idx_t... Is>
+RAJA_HOST_DEVICE RAJA_INLINE BinaryFunc for_each_tuple_index(Tuple&& t,
+                                                             BinaryFunc func,
+                                                             camp::idx_seq<Is...>)
+{
+  using camp::get;
+  // braced init lists are evaluated in order
+  int seq_unused_array[] = {0, (func(get<Is>(std::forward<Tuple>(t)), Is), 0)...};
+  RAJA_UNUSED_VAR(seq_unused_array);
+
+  return func;
+}
+
 }  // namespace detail
 
+RAJA_SUPPRESS_HD_WARN
+template<size_t N, typename UnaryFunc>
+RAJA_HOST_DEVICE RAJA_INLINE
+UnaryFunc for_each_index(UnaryFunc func)
+{
+  return detail::for_each_index(camp::make_idx_seq_t<N>(), std::move(func));
+}
+
 /*!
   \brief Apply func to all the elements in the given range in order
   using a sequential for loop in O(N) operations and O(1) extra memory
@@ -89,7 +129,7 @@ RAJA_HOST_DEVICE RAJA_INLINE UnaryFunc for_each_tuple(Tuple&& t,
 RAJA_SUPPRESS_HD_WARN
 template<typename Container, typename UnaryFunc>
 RAJA_HOST_DEVICE RAJA_INLINE
-    concepts::enable_if_t<UnaryFunc, type_traits::is_range<Container>>
+    camp::concepts::enable_if_t<UnaryFunc, camp::type_traits::is_range<Container>>
     for_each(Container&& c, UnaryFunc func)
 {
   using std::begin;
@@ -123,6 +163,20 @@ RAJA_HOST_DEVICE RAJA_INLINE UnaryFunc for_each_tuple(Tuple&& t, UnaryFunc func)
       camp::make_idx_seq_t<std::tuple_size<camp::decay<Tuple>>::value> {});
 }
 
+/*!
+  \brief Apply func to each object in the given tuple or tuple like type as well as 
+  the index of the tuple element in order using a compile-time expansion in O(N) 
+  operations and O(1) extra memory
+*/
+RAJA_SUPPRESS_HD_WARN
+template<typename Tuple, typename BinaryFunc>
+RAJA_HOST_DEVICE RAJA_INLINE BinaryFunc for_each_tuple_index(Tuple&& t, BinaryFunc func)
+{
+  return detail::for_each_tuple_index(
+      std::forward<Tuple>(t), std::move(func),
+      camp::make_idx_seq_t<std::tuple_size<camp::decay<Tuple>>::value> {});
+}
+
 }  // namespace RAJA
 
 #endif

test/unit/view-layout/CMakeLists.txt

@@ -28,3 +28,7 @@ raja_add_test(
 raja_add_test(
   NAME test-indexlayout
   SOURCES test-indexlayout.cpp)
+
+raja_add_test(
+  NAME test-subview
+  SOURCES test-subview.cpp)

test/unit/view-layout/test-subview.cpp

@@ -0,0 +1,132 @@
+//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
+// Copyright (c) 2016-25, Lawrence Livermore National Security, LLC
+// and RAJA project contributors. See the RAJA/LICENSE file for details.
+//
+// SPDX-License-Identifier: (BSD-3-Clause)
+//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
+
+#include <gtest/gtest.h>
+#include "RAJA/policy/PolicyBase.hpp"
+#include "RAJA/util/SubView.hpp"
+#include "RAJA/util/types.hpp"
+#include "RAJA_test-base.hpp"
+#include "RAJA_unit-test-forone.hpp"
+
+using namespace RAJA;
+
+TEST(SubView, RangeSubView1D)
+{
+
+    Index_type a[] = {1,2,3,4,5};
+
+    View<Index_type, Layout<1>> view(&a[0], Layout<1>(5));
+
+    // sv = View[1:3]
+    auto sv = SubView(view, RangeSlice<>{1,3});
+
+    EXPECT_EQ(sv(0), 2);
+    EXPECT_EQ(sv(1), 3);
+    EXPECT_EQ(sv(2), 4);
+
+}
+
+TEST(SubView, RangeSubView2D)
+{
+
+    Index_type a[3][3] = {{1,2,3},{4,5,6},{7,8,9}};
+
+    View<Index_type, Layout<2>> view(&a[0][0], Layout<2>(3,3));
+
+    // sv = View[1:2,1:2]
+    auto sv = SubView(view, RangeSlice<>{1,2}, RangeSlice<>{1,2});
+
+    EXPECT_EQ(sv(0,0), 5);
+    EXPECT_EQ(sv(0,1), 6);
+    EXPECT_EQ(sv(1,0), 8);
+    EXPECT_EQ(sv(1,1), 9);
+
+}
+
+TEST(SubView, RangeFixedSubView2D)
+{
+
+    Index_type a[3][3] = {{1,2,3},{4,5,6},{7,8,9}};
+
+    View<Index_type, Layout<2>> view(&a[0][0], Layout<2>(3,3));
+
+    // sv = View[1:2,1]
+    auto sv = SubView(view, RangeSlice<>{1,2}, FixedSlice<>{1});
+
+    EXPECT_EQ(sv(0), 5);
+    EXPECT_EQ(sv(1), 8);
+
+}
+
+TEST(SubView, FixedFirstDimSubView2D)
+{
+
+    Index_type a[3][3] = {{1,2,3},{4,5,6},{7,8,9}};
+
+    View<Index_type, Layout<2>> view(&a[0][0], Layout<2>(3,3));
+
+    // sv = View[1,:]
+    auto sv = SubView(view, FixedSlice<>{1}, NoSlice{});
+
+    EXPECT_EQ(sv(0), 4);
+    EXPECT_EQ(sv(1), 5);
+    EXPECT_EQ(sv(2), 6);
+
+}
+
+TEST(SubView, RangeFirstDimSubView2D)
+{
+
+    Index_type a[3][3] = {{1,2,3},{4,5,6},{7,8,9}};
+
+    View<Index_type, Layout<2>> view(&a[0][0], Layout<2>(3,3));
+
+    // sv = View[1:2,:]
+    auto sv = SubView(view, RangeSlice<>{1,2}, NoSlice{});
+
+    EXPECT_EQ(sv(0,0), 4);
+    EXPECT_EQ(sv(0,1), 5);
+    EXPECT_EQ(sv(0,2), 6);
+
+    EXPECT_EQ(sv(1,0), 7);
+    EXPECT_EQ(sv(1,1), 8);
+    EXPECT_EQ(sv(1,2), 9);
+
+    sv.set_slice<0>(RangeSlice<>{0,1});
+
+    EXPECT_EQ(sv(0,0), 1);
+    EXPECT_EQ(sv(0,1), 2);
+    EXPECT_EQ(sv(0,2), 3);
+
+    EXPECT_EQ(sv(1,0), 4);
+    EXPECT_EQ(sv(1,1), 5);
+    EXPECT_EQ(sv(1,2), 6);
+
+}
+
+// To Do tests: subview of subview, "sliding winow" modify slice entries, GPU tests.
+
+// void test_subviewGPU() {
+// #if defined(RAJA_ENABLE_HIP)
+//     forone<test_hip>([=] __host__ __device__ () {
+//         Index_type a[3][3] = {{1,2,3},{4,5,6},{7,8,9}};
+
+//         View<Index_type, Layout<2>> view(&a[0][0], Layout<2>(3,3));
+
+//         // sv = View[1:2,:]
+//         auto sv = SubView(view, RangeSlice<>{1,2}, NoSlice{});
+
+//         //printf("sv(0,0): %ld\n", sv(0,0));
+
+//     });
+// #endif
+// }
+
+// TEST(SubView, RangeFirstDimSubView2DGPU)
+// {
+//     test_subviewGPU();
+// }