Add for_each_n algorithm to shp and mhp (oneapi-src#786)

include/dr/mhp/algorithms/for_each.hpp

@@ -53,4 +53,13 @@ void for_each(DI first, DI last, auto op) {
   mhp::for_each(rng::subrange(first, last), op);
 }
 
+/// Collective for_each on iterator/sentinel for a distributed range
+template <dr::distributed_iterator DI, std::integral I>
+DI for_each_n(DI first, I n, auto op) {
+  auto last = first;
+  rng::advance(last, n);
+  mhp::for_each(first, last, op);
+  return last;
+}
+
 } // namespace dr::mhp

include/dr/shp/algorithms/for_each.hpp

@@ -54,4 +54,19 @@ void for_each(Iter begin, Iter end, Fn &&fn) {
   for_each(dr::shp::par_unseq, begin, end, std::forward<Fn>(fn));
 }
 
+template <typename ExecutionPolicy, dr::distributed_iterator Iter,
+          std::integral I, typename Fn>
+Iter for_each_n(ExecutionPolicy &&policy, Iter begin, I n, Fn fn) {
+  auto end = begin;
+  rng::advance(end, n);
+  for_each(std::forward<ExecutionPolicy>(policy), begin, end,
+           std::forward<Fn>(fn));
+  return end;
+}
+
+template <dr::distributed_iterator Iter, std::integral I, typename Fn>
+Iter for_each_n(Iter &&r, I n, Fn fn) {
+  return for_each_n(dr::shp::par_unseq, std::forward<Iter>(r), n, fn);
+}
+
 } // namespace dr::shp

test/gtest/common/for_each.cpp

@@ -9,6 +9,22 @@ template <typename T> class ForEach : public testing::Test {
 public:
 };
 
+template <typename T>
+void test_foreach_n(std::vector<T> v, int n, int initial_skip, auto func) {
+  auto size = v.size();
+  xhp::distributed_vector<T> d_v(size);
+
+  for (std::size_t idx = 0; idx < size; idx++) {
+    d_v[idx] = v[idx];
+  }
+  barrier();
+
+  xhp::for_each_n(d_v.begin() + initial_skip, n, func);
+  rng::for_each_n(v.begin() + initial_skip, n, func);
+
+  EXPECT_TRUE(equal(v, d_v));
+}
+
 TYPED_TEST_SUITE(ForEach, AllTypes);
 
 TYPED_TEST(ForEach, Range) {
@@ -45,6 +61,33 @@ TYPED_TEST(ForEach, RangeAlignedZip) {
   EXPECT_EQ(local, dist);
 }
 
+TYPED_TEST(ForEach, ForEachN) {
+  auto negate = [](auto &&v) { v = -v; };
+  test_foreach_n<int>({1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, 5, 0, negate);
+}
+
+TYPED_TEST(ForEach, ForEachNLongerThanSize) {
+  auto negate = [](auto &&v) { v = -v; };
+  test_foreach_n<int>({1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, 10, 0, negate);
+}
+
+TYPED_TEST(ForEach, ForEachNPartial) {
+  auto negate = [](auto &&v) { v = -v; };
+  test_foreach_n<int>({1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, 5, 2, negate);
+}
+
+TYPED_TEST(ForEach, ForEachNWholeLength) {
+  Ops1<TypeParam> ops(10);
+
+  auto negate = [](auto &&v) { v = -v; };
+  auto input = ops.vec;
+
+  xhp::for_each_n(ops.dist_vec.begin(), 10, negate);
+  rng::for_each(ops.vec.begin(), ops.vec.end(), negate);
+
+  EXPECT_TRUE(check_unary_op(input, ops.vec, ops.dist_vec));
+}
+
 // Disabled. Not sure how to support this properly for MHP. We need to
 // copy the values local so we can operate on them. Read-only data
 // seems doable but writing misaligned data is harder. We should