Interleaved computation with communication in halo exchange

include/dr/mp.hpp

@@ -79,3 +79,4 @@
 #include <dr/mp/algorithms/transpose.hpp>
 #include <dr/mp/containers/distributed_vector.hpp>
 #include <dr/mp/containers/distributed_mdarray.hpp>
+#include <dr/mp/containers/dual_distributed_vector.hpp>

include/dr/mp/algorithms/for_each.hpp

@@ -18,7 +18,58 @@
 
 namespace dr::mp {
 
-/// Collective for_each on distributed range
+// the concept doesn't work yet... for some reason
+template <typename R>
+concept dual_vector_range =
+    dr::distributed_range<R> && requires(R &r) { dr::ranges::segments(r)[0].is_compute(); };
+
+void for_each(dual_vector_range auto &&dr, auto op) {
+  partial_for_each(dr, op);
+  partial_for_each(dr, op);
+}
+
+void partial_for_each(dual_vector_range auto &&dr, auto op) {
+  dr::drlog.debug(dr::logger::for_each, "partial_for_each: parallel execution\n");
+  if (rng::empty(dr)) {
+    return;
+  }
+
+  auto is_local = [](const auto &segment) {
+    return dr::ranges::rank(segment) == default_comm().rank();
+  };
+
+  for (auto &seg : dr::ranges::segments(dr) | rng::views::filter(is_local)) {
+    if (!seg.is_compute()) {
+      seg.swap_state();
+      continue;
+    }
+
+    auto b = dr::ranges::local(rng::begin(seg));
+    auto s = rng::subrange(b, b + rng::distance(seg));
+
+    if (mp::use_sycl()) {
+      dr::drlog.debug("  using sycl\n");
+
+      assert(rng::distance(s) > 0);
+#ifdef SYCL_LANGUAGE_VERSION
+      dr::__detail::parallel_for(
+          dr::mp::sycl_queue(), sycl::range<1>(rng::distance(s)),
+          [first = rng::begin(s), op](auto idx) { op(first[idx]); })
+          .wait();
+#else
+      assert(false);
+#endif
+    } else {
+      dr::drlog.debug("  using cpu\n");
+      rng::for_each(s, op);
+    }
+
+    seg.swap_state();
+  }
+  barrier();
+}
+
+// Collective for_each on distributed range
 void for_each(dr::distributed_range auto &&dr, auto op) {
   dr::drlog.debug(dr::logger::for_each, "for_each: parallel execution\n");
   if (rng::empty(dr)) {