Python like Range

examples/CMakeLists.txt

@@ -19,6 +19,10 @@ raja_add_executable(
   NAME resource-forall
   SOURCES resource-forall.cpp)
 
+raja_add_executable(
+  NAME raja-ranges
+  SOURCES raja-ranges.cpp)
+
 raja_add_executable(
   NAME dynamic-forall
   SOURCES dynamic-forall.cpp)

examples/raja-ranges.cpp

@@ -0,0 +1,86 @@
+//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
+// Copyright (c) Lawrence Livermore National Security, LLC and other
+// RAJA Project Developers. See top-level LICENSE and COPYRIGHT
+// files for dates and other details. No copyright assignment is required
+// to contribute to RAJA.
+//
+// SPDX-License-Identifier: (BSD-3-Clause)
+//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
+
+#include <iostream>
+
+#include "RAJA/RAJA.hpp"
+
+/*
+ * RAJA ranges example
+ *
+ * Demonstrates Python-like range helpers backed by RAJA segments:
+ *   1. RAJA::Range(N) for [0, N)
+ *   2. RAJA::Range<T>(N) for a TypedRangeSegment<T> over [0, N)
+ *   3. RAJA::Range(begin, end) for [begin, end)
+ *   4. RAJA::Range(begin, end, step) for a strided half-open interval
+ */
+
+RAJA_INDEX_VALUE(CellIndex, "CellIndex");
+
+int main(int RAJA_UNUSED_ARG(argc), char** RAJA_UNUSED_ARG(argv))
+{
+  constexpr RAJA::Index_type N = 8;
+
+  int values[N] = {};
+  int typed_values[N] = {};
+  int subrange_values[N] = {};
+  int odd_values[N] = {};
+
+  // Equivalent to Python range(N): [0, N)
+  RAJA::forall<RAJA::seq_exec>(RAJA::Range(N), [&](RAJA::Index_type i) {
+    values[i] = static_cast<int>(i * i);
+  });
+
+  // Equivalent to Python range(N), but preserving a strong index type.
+  RAJA::forall<RAJA::seq_exec>(RAJA::Range<CellIndex>(N), [&](CellIndex i) {
+    typed_values[*i] = static_cast<int>(*i + 10);
+  });
+
+  // Equivalent to Python range(2, 6): [2, 6)
+  RAJA::forall<RAJA::seq_exec>(RAJA::Range(2, 6), [&](int i) {
+    subrange_values[i] = i;
+  });
+
+  // Equivalent to Python range(1, N, 2): odd indices in [1, N)
+  RAJA::forall<RAJA::seq_exec>(RAJA::Range(1, N, 2), [&](int i) {
+    odd_values[i] = i;
+  });
+
+  std::cout << "Range(N):";
+  for (auto i : RAJA::Range(N)) {
+    std::cout << ' ' << values[i];
+  }
+  std::cout << '\n';
+
+  std::cout << "Range<CellIndex>(N):";
+  for (auto i : RAJA::Range<CellIndex>(N)) {
+    std::cout << ' ' << typed_values[*i];
+  }
+  std::cout << '\n';
+
+  std::cout << "Range(2, 6):";
+  for (auto i : RAJA::Range(2, 6)) {
+    std::cout << ' ' << subrange_values[i];
+  }
+  std::cout << '\n';
+
+  std::cout << "Range(1, N, 2):";
+  for (auto i : RAJA::Range(1, N, 2)) {
+    std::cout << ' ' << odd_values[i];
+  }
+  std::cout << '\n';
+
+  std::cout << "Range(N - 1, -1, -2):";
+  for (auto i : RAJA::Range(N - 1, -1, -2)) {
+    std::cout << ' ' << i;
+  }
+  std::cout << '\n';
+
+  return 0;
+}

include/RAJA/index/RangeSegment.hpp

@@ -530,6 +530,33 @@ using RangeSegment = TypedRangeSegment<Index_type>;
 //! Alias for TypedRangeStrideSegment<Index_type>
 using RangeStrideSegment = TypedRangeStrideSegment<Index_type>;
 
+/*!
+ * \brief Function to make a RangeSegment for the interval [0, end)
+ *
+ *  \return a newly constructed RangeSegment over the half-open interval
+ *          starting at zero and ending at @end.
+ */
+RAJA_HOST_DEVICE RAJA_INLINE constexpr RangeSegment Range(
+    Index_type end) noexcept
+{
+  return {0, end};
+}
+
+/*!
+ * \brief Function to make a TypedRangeSegment for the interval [0, end)
+ *
+ *  \return a newly constructed TypedRangeSegment over the half-open interval
+ *          starting at zero and ending at @end.
+ */
+template<typename StorageT>
+RAJA_HOST_DEVICE RAJA_INLINE constexpr TypedRangeSegment<StorageT> Range(
+    strip_index_type_t<StorageT> end) noexcept
+{
+  static_assert(!std::is_floating_point<strip_index_type_t<StorageT>>::value,
+                "Range requires a non-floating point index type.");
+  return {strip_index_type_t<StorageT> {0}, end};
+}
+
 namespace detail
 {
 
@@ -549,6 +576,56 @@ using common_type_t = typename common_type<Ts...>::type;
 
 }  // namespace detail
 
+/*!
+ * \brief Function to make a TypedRangeSegment for the interval [begin, end)
+ *
+ *  \return a newly constructed TypedRangeSegment where the
+ *          value_type is equivilent to the common type of
+ *          @begin and @end. If there is no common type, then
+ *          a compiler error will be produced.
+ */
+template<typename BeginT,
+         typename EndT,
+         typename Common = detail::common_type_t<BeginT, EndT>>
+RAJA_HOST_DEVICE RAJA_INLINE constexpr TypedRangeSegment<Common> Range(
+    BeginT&& begin,
+    EndT&& end) noexcept
+{
+  return {begin, end};
+}
+
+/*!
+ * \brief Function to make a TypedRangeStrideSegment for the interval
+ *        [begin, end) with given stride
+ *
+ *  \return a newly constructed TypedRangeStrideSegment where the
+ *          value_type is equivilent to the common type of
+ *          @begin and @end. If stride is zero, execution aborts or throws.
+ */
+template<typename BeginT,
+         typename EndT,
+         typename StrideT,
+         typename Common = detail::common_type_t<BeginT, EndT>,
+         typename DiffT = make_signed_t<strip_index_type_t<Common>>>
+RAJA_HOST_DEVICE RAJA_INLINE TypedRangeStrideSegment<Common> Range(
+    BeginT&& begin,
+    EndT&& end,
+    StrideT&& stride)
+{
+  static_assert(std::is_signed<strip_index_type_t<StrideT>>::value,
+                "Range requires a signed stride type.");
+  static_assert(!std::is_floating_point<strip_index_type_t<StrideT>>::value,
+                "Range requires a non-floating point stride type.");
+
+  DiffT const typed_stride = static_cast<DiffT>(stride);
+  if (typed_stride == DiffT {0})
+  {
+    RAJA_ABORT_OR_THROW("RAJA::Range requires a non-zero stride.");
+  }
+
+  return {begin, end, typed_stride};
+}
+
 /*!
  * \brief Function to make a TypedRangeSegment for the interval [begin, end)
  *

test/unit/index/test-rangesegment.cpp

@@ -15,6 +15,8 @@
 
 #include "RAJA_unit-test-types.hpp"
 
+RAJA_INDEX_VALUE(RangeStrongIndex, "RangeStrongIndex");
+
 template<typename T>
 class RangeSegmentUnitTest : public ::testing::Test {};
 
@@ -166,3 +168,47 @@ TYPED_TEST(RangeSegmentUnitTest, Equality)
 
   ASSERT_NE(r1, r3);
 }
+
+TEST(RangeSegmentUnitTest, RangeEnd)
+{
+  auto r = RAJA::Range(RAJA::Index_type(17));
+
+  ASSERT_EQ(RAJA::RangeSegment(RAJA::Index_type(0), RAJA::Index_type(17)), r);
+  ASSERT_EQ(RAJA::Index_type(0), *r.begin());
+  ASSERT_EQ(RAJA::Index_type(17), r.size());
+}
+
+TEST(RangeSegmentUnitTest, TypedRangeEnd)
+{
+  auto r = RAJA::Range<RangeStrongIndex>(17);
+
+  ASSERT_EQ((RAJA::TypedRangeSegment<RangeStrongIndex>(0, 17)), r);
+  ASSERT_EQ(RangeStrongIndex(0), *r.begin());
+  ASSERT_EQ(RAJA::Index_type(17), r.size());
+}
+
+TEST(RangeSegmentUnitTest, RangeBeginEnd)
+{
+  auto r = RAJA::Range(3, 17);
+
+  ASSERT_EQ((RAJA::TypedRangeSegment<int>(3, 17)), r);
+  ASSERT_EQ(14, r.size());
+}
+
+TEST(RangeSegmentUnitTest, RangeBeginEndStridePositive)
+{
+  auto r = RAJA::Range(2, 11, 3);
+
+  ASSERT_EQ((RAJA::TypedRangeStrideSegment<int>(2, 11, 3)), r);
+  ASSERT_EQ(4, r.size());
+  ASSERT_EQ(2, *r.begin());
+}
+
+TEST(RangeSegmentUnitTest, RangeBeginEndStrideNegative)
+{
+  auto r = RAJA::Range(10, -1, -2);
+
+  ASSERT_EQ((RAJA::TypedRangeStrideSegment<int>(10, -1, -2)), r);
+  ASSERT_EQ(6, r.size());
+  ASSERT_EQ(10, *r.begin());
+}