Wrap communicators in RoundRobin and Pipeline in CommunicatorGrid

include/dlaf/auxiliary/norm.h

@@ -32,7 +32,7 @@ namespace dlaf::auxiliary {
 /// @pre @p A has tilesize (NB x NB)
 /// @return the max norm of the Matrix @p A or 0 if `A.size().isEmpty()`
 template <Backend backend, Device device, class T>
-dlaf::BaseType<T> max_norm(comm::CommunicatorGrid grid, comm::Index2D rank, blas::Uplo uplo,
+dlaf::BaseType<T> max_norm(comm::CommunicatorGrid& grid, comm::Index2D rank, blas::Uplo uplo,
                            Matrix<const T, device>& A) {
   using dlaf::matrix::equal_process_grid;
   using dlaf::matrix::single_tile_per_block;

include/dlaf/auxiliary/norm/api.h

@@ -20,10 +20,10 @@ struct Norm {};
 
 template <class T>
 struct Norm<Backend::MC, Device::CPU, T> {
-  static dlaf::BaseType<T> max_L(comm::CommunicatorGrid comm_grid, comm::Index2D rank,
+  static dlaf::BaseType<T> max_L(comm::CommunicatorGrid& comm_grid, comm::Index2D rank,
                                  Matrix<const T, Device::CPU>& matrix);
 
-  static dlaf::BaseType<T> max_G(comm::CommunicatorGrid comm_grid, comm::Index2D rank,
+  static dlaf::BaseType<T> max_G(comm::CommunicatorGrid& comm_grid, comm::Index2D rank,
                                  Matrix<const T, Device::CPU>& matrix);
 };
 

include/dlaf/auxiliary/norm/mc.h

@@ -29,7 +29,7 @@ namespace dlaf::auxiliary::internal {
 // - sy/he lower
 // - tr lower non-unit
 template <class T>
-dlaf::BaseType<T> Norm<Backend::MC, Device::CPU, T>::max_L(comm::CommunicatorGrid comm_grid,
+dlaf::BaseType<T> Norm<Backend::MC, Device::CPU, T>::max_L(comm::CommunicatorGrid& comm_grid,
                                                            comm::Index2D rank,
                                                            Matrix<const T, Device::CPU>& matrix) {
   using namespace dlaf::matrix;
@@ -92,7 +92,7 @@ dlaf::BaseType<T> Norm<Backend::MC, Device::CPU, T>::max_L(comm::CommunicatorGri
 }
 
 template <class T>
-dlaf::BaseType<T> Norm<Backend::MC, Device::CPU, T>::max_G(comm::CommunicatorGrid comm_grid,
+dlaf::BaseType<T> Norm<Backend::MC, Device::CPU, T>::max_G(comm::CommunicatorGrid& comm_grid,
                                                            comm::Index2D rank,
                                                            Matrix<const T, Device::CPU>& matrix) {
   using namespace dlaf::matrix;

include/dlaf/common/pipeline.h

@@ -12,6 +12,8 @@
 
 /// @file
 
+#include <utility>
+
 #include <pika/async_rw_mutex.hpp>
 #include <pika/execution.hpp>
 
@@ -32,25 +34,78 @@ class Pipeline {
   using AsyncRwMutex = pika::execution::experimental::async_rw_mutex<T>;
 
 public:
-  using Wrapper = typename AsyncRwMutex::readwrite_access_type;
-  using Sender = pika::execution::experimental::unique_any_sender<Wrapper>;
+  using ReadOnlyWrapper = typename AsyncRwMutex::read_access_type;
+  using ReadWriteWrapper = typename AsyncRwMutex::readwrite_access_type;
+  using ReadOnlySender = pika::execution::experimental::any_sender<ReadOnlyWrapper>;
+  using ReadWriteSender = pika::execution::experimental::unique_any_sender<ReadWriteWrapper>;
+
+  /// Create an invalid Pipeline.
+  Pipeline() = default;
 
   /// Create a Pipeline by moving in the resource (it takes the ownership).
   explicit Pipeline(T object) : pipeline(std::move(object)) {}
-  Pipeline(Pipeline&&) = default;
-  Pipeline& operator=(Pipeline&&) = default;
+
+  Pipeline(Pipeline&& other) noexcept
+      : pipeline(std::exchange(other.pipeline, std::nullopt)),
+        nested_sender(std::exchange(other.nested_sender, std::nullopt)) {}
+
+  Pipeline& operator=(Pipeline&& other) noexcept {
+    if (this != &other) {
+      pipeline = std::exchange(other.pipeline, std::nullopt);
+      nested_sender = std::exchange(other.nested_sender, std::nullopt);
+    }
+
+    return *this;
+  };
+
   Pipeline(const Pipeline&) = delete;
   Pipeline& operator=(const Pipeline&) = delete;
 
-  /// Enqueue for the resource.
+  ~Pipeline() {
+    release_parent_pipeline();
+  }
+
+  /// Enqueue for exclusive read-write access to the resource.
   ///
   /// @return a sender that will become ready as soon as the previous user releases the resource.
   /// @pre valid()
-  Sender operator()() {
+  ReadWriteSender readwrite() {
     DLAF_ASSERT(valid(), "");
     return pipeline->readwrite();
   }
 
+  /// Enqueue for shared read-only access to the resource.
+  ///
+  /// @return a sender that will become ready as soon as the previous user releases the resource.
+  /// @pre valid()
+  ReadOnlySender read() {
+    DLAF_ASSERT(valid(), "");
+    return pipeline->read();
+  }
+
+  /// Create a sub pipeline to the value contained in the current Pipeline
+  ///
+  /// All accesses to the sub pipeline are sequenced after previous accesses and before later accesses to
+  /// the original pipeline, independently of when values are accessed in the sub pipeline.
+  Pipeline sub_pipeline() {
+    namespace ex = pika::execution::experimental;
+
+    // Move value from pipeline into sub pipeline, then store a sender of the wrapper of the pipeline in
+    // a sender which we will release when the sub pipeline is released. This ensures that all accesses
+    // to the sub pipeline happen after previous accesses and before later accesses to the pipeline.
+    Pipeline sub_pipeline(T{});
+    sub_pipeline.nested_sender =
+        ex::when_all(sub_pipeline.pipeline->readwrite(), this->pipeline->readwrite()) |
+        ex::then([](auto sub_wrapper, auto wrapper) {
+          sub_wrapper.get() = std::move(wrapper.get());
+
+          return wrapper;
+        }) |
+        ex::ensure_started();
+
+    return sub_pipeline;
+  }
+
   /// Check if the pipeline is valid.
   ///
   /// @return true if the pipeline hasn't been reset, otherwise false.
@@ -62,10 +117,27 @@ class Pipeline {
   ///
   /// @post !valid()
   void reset() noexcept {
+    release_parent_pipeline();
     pipeline.reset();
   }
 
 private:
+  void release_parent_pipeline() {
+    namespace ex = pika::execution::experimental;
+
+    if (nested_sender) {
+      DLAF_ASSERT(valid(), "");
+
+      auto s =
+          ex::when_all(pipeline->readwrite(), std::move(nested_sender.value())) |
+          ex::then([](auto sub_wrapper, auto wrapper) { wrapper.get() = std::move(sub_wrapper.get()); });
+      ex::start_detached(std::move(s));
+      nested_sender.reset();
+    }
+  }
+
   std::optional<AsyncRwMutex> pipeline = std::nullopt;
+  std::optional<pika::execution::experimental::unique_any_sender<ReadWriteWrapper>> nested_sender =
+      std::nullopt;
 };
 }

include/dlaf/common/round_robin.h

@@ -13,28 +13,46 @@
 
 #include <vector>
 
+#include <dlaf/common/assert.h>
+
 namespace dlaf {
 namespace common {
 
 template <class T>
-struct RoundRobin {
+class RoundRobin {
+public:
+  RoundRobin() = default;
+
   template <class... Args>
-  RoundRobin(std::size_t n, Args... args) : curr_index_(0) {
+  RoundRobin(std::size_t n, Args... args) {
     pool_.reserve(n);
     for (std::size_t i = 0; i < n; ++i)
       pool_.emplace_back(args...);
   }
 
+  RoundRobin(RoundRobin&&) = default;
+  RoundRobin(const RoundRobin&) = delete;
+  RoundRobin& operator=(RoundRobin&&) = default;
+  RoundRobin& operator=(const RoundRobin&) = delete;
+
   T& currentResource() {
+    DLAF_ASSERT(curr_index_ < pool_.size(), curr_index_, pool_.size());
     return pool_[curr_index_];
   }
 
   T& nextResource() {
+    DLAF_ASSERT(!pool_.empty(), "");
     curr_index_ = (curr_index_ + 1) % pool_.size();
+    DLAF_ASSERT_HEAVY(curr_index_ < pool_.size(), curr_index_, pool_.size());
     return pool_[curr_index_];
   }
 
-  std::size_t curr_index_;
+  std::size_t size() const noexcept {
+    return pool_.size();
+  }
+
+private:
+  std::size_t curr_index_ = 0;
   std::vector<T> pool_;
 };
 

include/dlaf/common/with_result_of.h

@@ -0,0 +1,41 @@
+//
+// Distributed Linear Algebra with Future (DLAF)
+//
+// Copyright (c) 2018-2023, ETH Zurich
+// All rights reserved.
+//
+// Please, refer to the LICENSE file in the root directory.
+// SPDX-License-Identifier: BSD-3-Clause
+//
+
+#pragma once
+
+/// @file
+
+#include <type_traits>
+#include <utility>
+
+namespace dlaf::internal {
+// Based on https://quuxplusone.github.io/blog/2018/05/17/super-elider-round-2/ and
+// https://akrzemi1.wordpress.com/2018/05/16/rvalues-redefined/.
+//
+// Because of the conversion operator and guaranteed copy-elision, useful for
+// emplacing immovable types into e.g. variants and optionals. Can also be used
+// to construct new instances for each element in a vector, where the element
+// type has reference semantics and regular copy construction is not what is
+// wanted.
+template <typename F>
+class WithResultOf {
+  F&& f;
+
+public:
+  using ResultType = std::invoke_result_t<F&&>;
+  explicit WithResultOf(F&& f) : f(std::forward<F>(f)) {}
+  operator ResultType() {
+    return std::forward<F>(f)();
+  }
+};
+
+template <typename F>
+WithResultOf(F&&) -> WithResultOf<F&&>;
+}

include/dlaf/communication/broadcast_panel.h

@@ -17,6 +17,8 @@
 #include <pika/execution.hpp>
 
 #include <dlaf/common/index2d.h>
+#include <dlaf/communication/communicator_pipeline.h>
+#include <dlaf/communication/index.h>
 #include <dlaf/communication/kernels/broadcast.h>
 #include <dlaf/communication/message.h>
 #include <dlaf/matrix/copy_tile.h>
@@ -56,7 +58,7 @@ std::pair<SizeType, comm::IndexT_MPI> transposedOwner(const matrix::Distribution
 template <class T, Device D, Coord axis, matrix::StoreTransposed storage,
           class = std::enable_if_t<!std::is_const_v<T>>>
 void broadcast(comm::IndexT_MPI rank_root, matrix::Panel<axis, T, D, storage>& panel,
-               common::Pipeline<comm::Communicator>& serial_comm) {
+               comm::CommunicatorPipeline<coord_to_communicator_type(orthogonal(axis))>& serial_comm) {
   constexpr auto comm_coord = axis;
 
   // do not schedule communication tasks if there is no reason to do so...
@@ -68,12 +70,25 @@ void broadcast(comm::IndexT_MPI rank_root, matrix::Panel<axis, T, D, storage>& p
   namespace ex = pika::execution::experimental;
   for (const auto& index : panel.iteratorLocal()) {
     if (rank == rank_root)
-      ex::start_detached(scheduleSendBcast(serial_comm(), panel.read(index)));
+      ex::start_detached(scheduleSendBcast(serial_comm.exclusive(), panel.read(index)));
     else
-      ex::start_detached(scheduleRecvBcast(serial_comm(), rank_root, panel.readwrite(index)));
+      ex::start_detached(scheduleRecvBcast(serial_comm.exclusive(), rank_root, panel.readwrite(index)));
   }
 }
 
+namespace internal {
+template <Coord C>
+auto& get_taskchain(comm::CommunicatorPipeline<comm::CommunicatorType::Row>& row_task_chain,
+                    comm::CommunicatorPipeline<comm::CommunicatorType::Col>& col_task_chain) {
+  if constexpr (C == Coord::Row) {
+    return row_task_chain;
+  }
+  else {
+    return col_task_chain;
+  }
+}
+}  // namespace internal
+
 /// Broadcast
 ///
 /// Given a source panel on a rank, this communication pattern makes every rank access tiles of both:
@@ -107,17 +122,13 @@ template <class T, Device D, Coord axis, matrix::StoreTransposed storage,
           matrix::StoreTransposed storageT, class = std::enable_if_t<!std::is_const_v<T>>>
 void broadcast(comm::IndexT_MPI rank_root, matrix::Panel<axis, T, D, storage>& panel,
                matrix::Panel<orthogonal(axis), T, D, storageT>& panelT,
-               common::Pipeline<comm::Communicator>& row_task_chain,
-               common::Pipeline<comm::Communicator>& col_task_chain) {
+               comm::CommunicatorPipeline<comm::CommunicatorType::Row>& row_task_chain,
+               comm::CommunicatorPipeline<comm::CommunicatorType::Col>& col_task_chain) {
   constexpr Coord axisT = orthogonal(axis);
 
   constexpr Coord coord = std::decay_t<decltype(panel)>::coord;
   constexpr Coord coordT = std::decay_t<decltype(panelT)>::coord;
 
-  auto get_taskchain = [&](Coord comm_dir) -> auto& {
-    return comm_dir == Coord::Row ? row_task_chain : col_task_chain;
-  };
-
   // Note:
   // Given a source panel, this communication pattern makes every rank access tiles of both the
   // source panel and it's tranposed variant (just tile coordinates, data is not transposed).
@@ -160,15 +171,15 @@ void broadcast(comm::IndexT_MPI rank_root, matrix::Panel<axis, T, D, storage>& p
 
   // STEP 1
   constexpr auto comm_dir_step1 = orthogonal(axis);
-  auto& chain_step1 = get_taskchain(comm_dir_step1);
+  auto& chain_step1 = internal::get_taskchain<comm_dir_step1>(row_task_chain, col_task_chain);
 
   broadcast(rank_root, panel, chain_step1);
 
   // STEP 2
   constexpr auto comm_dir_step2 = orthogonal(axisT);
   constexpr auto comm_coord_step2 = axisT;
 
-  auto& chain_step2 = get_taskchain(comm_dir_step2);
+  auto& chain_step2 = internal::get_taskchain<comm_dir_step2>(row_task_chain, col_task_chain);
 
   const SizeType last_tile = std::max(panelT.rangeStart(), panelT.rangeEnd() - 1);
   const auto owner = dist.template rankGlobalTile<coordT>(last_tile);
@@ -186,11 +197,12 @@ void broadcast(comm::IndexT_MPI rank_root, matrix::Panel<axis, T, D, storage>& p
       panelT.setTile(indexT, panel.read({coord, index_diag_local}));
 
       if (dist.commGridSize().get(comm_coord_step2) > 1)
-        ex::start_detached(scheduleSendBcast(chain_step2(), panelT.read(indexT)));
+        ex::start_detached(scheduleSendBcast(chain_step2.exclusive(), panelT.read(indexT)));
     }
     else {
       if (dist.commGridSize().get(comm_coord_step2) > 1)
-        ex::start_detached(scheduleRecvBcast(chain_step2(), owner_diag, panelT.readwrite(indexT)));
+        ex::start_detached(scheduleRecvBcast(chain_step2.exclusive(), owner_diag,
+                                             panelT.readwrite(indexT)));
     }
   }
 }

include/dlaf/communication/communicator.h

@@ -15,13 +15,11 @@
 #include <memory>
 
 #include <dlaf/communication/error.h>
+#include <dlaf/communication/index.h>
 
 namespace dlaf {
 namespace comm {
 
-/// Type used for indexes in MPI API.
-using IndexT_MPI = int;
-
 class CommunicatorImpl;
 
 /// MPI-compatible wrapper for the MPI_Comm.