various changes/fixes to address comments

Author: gberg617

include/RAJA/util/SubView.hpp

@@ -27,15 +27,42 @@ namespace RAJA
 {
 
 // Slice descriptors
-struct RangeSlice { long start, end; };
-struct FixedSlice { long idx; };
-struct NoSlice { };
+
+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;
+    }
+};
+
+struct NoSlice { 
+    static constexpr bool reduces_dimension = false;
+
+    template<typename IndexType = Index_type>
+    RAJA_INLINE RAJA_HOST_DEVICE constexpr IndexType map_index(IndexType& idx) const {
+        return idx;
+    }
+};
 
 // Helper to count non-fixed dimensions at compile time
 template <typename... Slices>
 RAJA_INLINE RAJA_HOST_DEVICE constexpr size_t count_nonfixed_dims() {
-    return (0 + ... + ((std::is_same_v<Slices, RangeSlice> || 
-                        std::is_same_v<Slices, NoSlice>) ? 1 : 0));
+    return (!Slices::reduces_dimension + ...);
 }
 
 template <typename T, size_t N, RAJA::Index_type... Is>
@@ -48,45 +75,42 @@ RAJA_INLINE RAJA_HOST_DEVICE constexpr auto array_to_tuple(const camp::array<T,
     return array_to_tuple_impl(arr, camp::make_idx_seq_t<N>{});
 }
 
-template <typename ViewType, typename... Slices>
+template <typename ViewType, typename IndexType = Index_type, typename... Slices>
 class SubView {
     ViewType view_;
     camp::tuple<Slices...> slices_;
+    std::array<IndexType, sizeof...(Slices)> map_;
 
-    template <typename IndexType, IndexType... Is>
-    RAJA_INLINE RAJA_HOST_DEVICE constexpr auto map_indices(IndexType* idxs, camp::idx_seq<Is...>) const {
-        camp::array<RAJA::Index_type, sizeof...(Is)> parent_indices{};
-        RAJA::Index_type idx = 0;
+    RAJA_INLINE RAJA_HOST_DEVICE constexpr void make_subview_index_map() {
+        size_t sub_idx = 0;
+        size_t i = 0;
+        ((map_[i++] = (Slices::reduces_dimension ? -1 : sub_idx++)), ...);
+    }
 
+    template<IndexType I>
+    RAJA_INLINE RAJA_HOST_DEVICE constexpr auto map_subview_idx_to_parent(IndexType* idxs) const {
+        return camp::get<I>(slices_).map_index(idxs[map_[I]]);
+    }
+
+    template <IndexType... Is>
+    RAJA_INLINE RAJA_HOST_DEVICE constexpr auto map_indices(IndexType* idxs, camp::idx_seq<Is...>) const {
         // For each slice, map subview index to parent index
-        (
-            (
-                parent_indices[Is] = [&] {
-                    const auto& s = camp::get<Is>(slices_);
-                    if constexpr (std::is_same_v<std::decay_t<decltype(s)>, RangeSlice>) {
-                        return s.start + idxs[idx++];
-                    } else if constexpr (std::is_same_v<std::decay_t<decltype(s)>, FixedSlice>) {
-                        return s.idx;
-                    } else {
-                        return idxs[idx++];
-                    }
-                }()
-            ), ...
-        );
-
-        return parent_indices;
+        return camp::array{(map_subview_idx_to_parent<Is>(idxs))...};
     }
 
 public:
-    RAJA_INLINE RAJA_HOST_DEVICE SubView(ViewType view, Slices... slices)
-        : view_(view), slices_(slices...) {}
+
+    RAJA_INLINE RAJA_HOST_DEVICE constexpr SubView(ViewType view, Slices... slices)
+        : view_(view), slices_(slices...) { make_subview_index_map(); }
 
     template <typename... Idxs>
-    RAJA_INLINE RAJA_HOST_DEVICE constexpr auto operator()(Idxs... idxs) const {
+    RAJA_INLINE RAJA_HOST_DEVICE constexpr IndexType operator()(Idxs... idxs) const {
         constexpr size_t nidx = count_nonfixed_dims<Slices...>();
         static_assert(sizeof...(idxs) == nidx, "Wrong number of indices for subview");
+
         camp::array<RAJA::Index_type, nidx> arr{idxs...};
         auto parent_indices = map_indices(arr.data(), camp::make_idx_seq_t<sizeof...(Slices)>());
+
         return camp::apply(view_, array_to_tuple(parent_indices));
     }
 };

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

@@ -7,6 +7,7 @@
 
 #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"
@@ -21,7 +22,7 @@ TEST(SubView, RangeSubView1D)
     View<Index_type, Layout<1>> view(&a[0], Layout<1>(5));
 
     // sv = View[1:3]
-    auto sv = SubView(view, RangeSlice{1,3});
+    auto sv = SubView(view, RangeSlice<>{1,3});
 
     EXPECT_EQ(sv(0), 2);
     EXPECT_EQ(sv(1), 3);
@@ -37,7 +38,7 @@ TEST(SubView, RangeSubView2D)
     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});
+    auto sv = SubView(view, RangeSlice<>{1,2}, RangeSlice<>{1,2});
 
     EXPECT_EQ(sv(0,0), 5);
     EXPECT_EQ(sv(0,1), 6);
@@ -54,7 +55,7 @@ TEST(SubView, RangeFixedSubView2D)
     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});
+    auto sv = SubView(view, RangeSlice<>{1,2}, FixedSlice<>{1});
 
     EXPECT_EQ(sv(0), 5);
     EXPECT_EQ(sv(1), 8);
@@ -69,7 +70,7 @@ TEST(SubView, FixedFirstDimSubView2D)
     View<Index_type, Layout<2>> view(&a[0][0], Layout<2>(3,3));
 
     // sv = View[1,:]
-    auto sv = SubView(view, FixedSlice{1}, NoSlice{});
+    auto sv = SubView(view, FixedSlice<>{1}, NoSlice{});
 
     EXPECT_EQ(sv(0), 4);
     EXPECT_EQ(sv(1), 5);
@@ -85,7 +86,7 @@ TEST(SubView, RangeFirstDimSubView2D)
     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{});
+    auto sv = SubView(view, RangeSlice<>{1,2}, NoSlice{});
 
     EXPECT_EQ(sv(0,0), 4);
     EXPECT_EQ(sv(0,1), 5);
@@ -105,7 +106,7 @@ void test_subviewGPU() {
         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{});
+        auto sv = SubView(view, RangeSlice<>{1,2}, NoSlice{});
 
         //printf("sv(0,0): %ld\n", sv(0,0));