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,52 @@
 
 namespace dr::mp {
 
-/// Collective for_each on distributed range
+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;
+  }
+
+  assert(aligned(dr));
+
+  for (auto &s : dr::ranges::segments(dr)) {
+    if (!s.is_compute()) {
+      s.swap_state();
+      continue;
+    }
+
+    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);
+    }
+
+    s.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)) {

include/dr/mp/containers/distributed_vector.hpp

@@ -254,7 +254,7 @@ template <typename T, class BackendT = MpiBackend> class distributed_vector {
     if (!finalized()) {
       fence();
       if (data_ != nullptr) {
-        backend.deallocate(data_, data_size_ * sizeof(value_type));
+        backend_.deallocate(data_, data_size_ * sizeof(value_type));
       }
 
       delete halo_;
@@ -274,7 +274,14 @@ template <typename T, class BackendT = MpiBackend> class distributed_vector {
 
   auto segments() const { return rng::views::all(segments_); }
 
-  void fence() { backend.fence(); }
+  void fence() { backend_.fence(); }
+
+  backend_type& backend(const std::size_t segment_index) { return backend_; }
+  const backend_type& backend(const std::size_t segment_index) const { 
+    return backend_;
+  }
+
+  T *data(const std::size_t segment_index) { return data_; }
 
 private:
   void init(auto size, auto dist) {
@@ -295,7 +302,7 @@ template <typename T, class BackendT = MpiBackend> class distributed_vector {
     data_size_ = segment_size_ + hb.prev + hb.next;
 
     if (size_ > 0) {
-      data_ = static_cast<T *>(backend.allocate(data_size_ * sizeof(T)));
+      data_ = static_cast<T *>(backend_.allocate(data_size_ * sizeof(T)));
     }
 
     halo_ = new span_halo<T>(default_comm(), data_, data_size_, hb);
@@ -319,7 +326,7 @@ template <typename T, class BackendT = MpiBackend> class distributed_vector {
   distribution distribution_;
   std::size_t size_;
   std::vector<dv_segment<distributed_vector>> segments_;
-  BackendT backend;
+  BackendT backend_;
 };
 
 template <typename T, typename B>