Use pika's transform_mpi and polling support

CMakeLists.txt

@@ -163,7 +163,7 @@ if(DLAF_WITH_SCALAPACK)
 endif()
 
 # ----- pika
-find_package(pika 0.19.1 REQUIRED)
+find_package(pika 0.30.0 REQUIRED)
 
 # ----- BLASPP/LAPACKPP
 find_package(blaspp REQUIRED)

ci/common-ci.yml

@@ -34,7 +34,7 @@ stages:
     reports:
       dotenv: build.env
   variables:
-    SPACK_SHA: develop-2024-10-06
+    SPACK_SHA: develop-2024-11-10
     SPACK_DLAF_REPO: ./spack
     DOCKER_BUILD_ARGS: '[
         "BASE_IMAGE",

ci/cuda/gcc11_codecov.yml

@@ -22,6 +22,8 @@ cuda gcc11 codecov build:
 
 cuda gcc11 codecov test:
   extends: .run_common
+  variables:
+    PIKA_MPI_ENABLE_POOL: 1
   needs:
     - cuda gcc11 codecov build
   trigger:

ci/cuda/gcc11_debug_scalapack.yml

@@ -20,6 +20,8 @@ cuda gcc11 debug scalapack build:
 
 cuda gcc11 debug scalapack test:
   extends: .run_common
+  variables:
+    PIKA_MPI_ENABLE_POOL: 1
   needs:
     - cuda gcc11 debug scalapack build
   trigger:

ci/cuda/gcc11_release.yml

@@ -21,6 +21,8 @@ cuda gcc11 release build:
 
 cuda gcc11 release test:
   extends: .run_common
+  variables:
+    PIKA_MPI_ENABLE_POOL: 1
   needs:
     - cuda gcc11 release build
   trigger:

ci/cuda/gcc11_release_scalapack.yml

@@ -21,6 +21,8 @@ cuda gcc11 release scalapack build:
 
 cuda gcc11 release scalapack test:
   extends: .run_common
+  variables:
+    PIKA_MPI_ENABLE_POOL: 1
   needs:
     - cuda gcc11 release scalapack build
   trigger:

ci/docker/debug-cuda-scalapack.yaml

@@ -35,6 +35,6 @@ spack:
         - '+fortran'
     pika:
       require:
-        - '@0.19.1'
+        - '@0.30.0'
         - 'build_type=Debug'
         - 'malloc=system'

include/dlaf/init.h

@@ -50,7 +50,6 @@ struct configuration {
   double umpire_device_memory_pool_coalescing_reallocation_ratio = 1.0;
   std::size_t num_gpu_blas_handles = 16;
   std::size_t num_gpu_lapack_handles = 16;
-  std::string mpi_pool = "mpi";
 };
 
 std::ostream& operator<<(std::ostream& os, const configuration& cfg);
@@ -107,10 +106,4 @@ struct [[nodiscard]] ScopedInitializer {
   ScopedInitializer& operator=(ScopedInitializer&&) = delete;
   ScopedInitializer& operator=(const ScopedInitializer&) = delete;
 };
-
-/// Initialize the MPI pool.
-///
-///
-void initResourcePartitionerHandler(pika::resource::partitioner& rp,
-                                    const pika::program_options::variables_map& vm);
 }

include/dlaf/schedulers.h

@@ -47,9 +47,4 @@ auto getBackendScheduler(
   }
 #endif
 }
-
-inline auto getMPIScheduler() {
-  return pika::execution::experimental::thread_pool_scheduler{
-      &pika::resource::get_thread_pool(getConfiguration().mpi_pool)};
-}
 }  // namespace dlaf::internal

include/dlaf/sender/transform_mpi.h

@@ -13,6 +13,7 @@
 #include <utility>
 
 #include <pika/execution.hpp>
+#include <pika/mpi.hpp>
 
 #include <dlaf/common/consume_rvalues.h>
 #include <dlaf/common/unwrap.h>
@@ -30,53 +31,35 @@ inline void consumeCommunicatorWrapper(CommunicatorPipelineExclusiveWrapper& com
   [[maybe_unused]] auto comm_wrapper_local = std::move(comm_wrapper);
 }
 
-/// \overload consumeCommunicatorWrapper
+/// \overload consumeCommunicatorWrapper (for non communicator types)
 template <typename T>
 void consumeCommunicatorWrapper(T&) {}
 
 /// Helper type for wrapping MPI calls.
 ///
-/// Wrapper type around calls to MPI functions. Provides a call operator that
-/// creates an MPI request and passes it as the last argument to the provided
-/// callable. The wrapper then waits for the request to complete with
-/// yield_while.
+/// The wrapper explicitly releases any dla communicator objects when the pika::transform_mpi
+/// function returns (e.g. a message has been sent/posted) to prevent blocking access to many
+/// queued mpi operations.
+/// The mpi operations can complete asynchronously later, but the commmunicator is
+/// released/made available once the mpi task has been safely initiated
 ///
 /// This could in theory be a lambda inside transformMPI.  However, clang at
 /// least until version 12 fails with an internal compiler error with a trailing
 /// decltype for SFINAE. GCC has no problems with a lambda.
 template <typename F>
 struct MPICallHelper {
   std::decay_t<F> f;
-  template <typename... Ts>
-  auto operator()(Ts&&... ts) -> decltype(std::move(f)(dlaf::common::internal::unwrap(ts)...,
-                                                       std::declval<MPI_Request*>())) {
-    MPI_Request req;
-    auto is_request_completed = [&req] {
-      int flag;
-      MPI_Test(&req, &flag, MPI_STATUS_IGNORE);
-      return flag == 0;
-    };
 
-    // Note:
-    // Callables passed to transformMPI have their arguments passed by reference, but doing so
-    // with PromiseGuard would keep the guard alive until the completion of the MPI operation,
-    // whereas we are only looking to guard the submission of the MPI operation. We therefore
-    // explicitly release CommunicatorPipelineExclusiveWrapper after submitting the MPI operation
-    // with consumeCommunicatorWrapper.
-    //
-    // We also use unwrap various types passed to the MPI operation, including PromiseGuards of
-    // any type, to allow the MPI operation not to care whether a Communicator was wrapped in a
-    // PromiseGuard or not.
-    using result_type = decltype(std::move(f)(dlaf::common::internal::unwrap(ts)..., &req));
+  template <typename... Ts>
+  auto operator()(Ts&&... ts) -> decltype(std::move(f)(dlaf::common::internal::unwrap(ts)...)) {
+    using result_type = decltype(std::move(f)(dlaf::common::internal::unwrap(ts)...));
     if constexpr (std::is_void_v<result_type>) {
-      std::move(f)(dlaf::common::internal::unwrap(ts)..., &req);
+      std::move(f)(dlaf::common::internal::unwrap(ts)...);
       (internal::consumeCommunicatorWrapper(ts), ...);
-      pika::util::yield_while(is_request_completed);
     }
     else {
-      auto r = std::move(f)(dlaf::common::internal::unwrap(ts)..., &req);
+      auto r = std::move(f)(dlaf::common::internal::unwrap(ts)...);
       (internal::consumeCommunicatorWrapper(ts), ...);
-      pika::util::yield_while(is_request_completed);
       return r;
     }
   }
@@ -85,16 +68,16 @@ struct MPICallHelper {
 template <typename F>
 MPICallHelper(F&&) -> MPICallHelper<std::decay_t<F>>;
 
-/// Lazy transformMPI. This does not submit the work and returns a sender.
+/// Lazy transformMPI. Returns a sender that will submit the work passed in
 template <typename F, typename Sender,
           typename = std::enable_if_t<pika::execution::experimental::is_sender_v<Sender>>>
 [[nodiscard]] decltype(auto) transformMPI(F&& f, Sender&& sender) {
-  using dlaf::internal::continues_on;
-  namespace ex = pika::execution::experimental;
-
-  return continues_on(std::forward<Sender>(sender), dlaf::internal::getMPIScheduler()) |
-         ex::then(dlaf::common::internal::ConsumeRvalues{MPICallHelper{std::forward<F>(f)}}) |
-         ex::drop_operation_state();
+  using dlaf::common::internal::ConsumeRvalues;
+  using pika::execution::experimental::drop_operation_state;
+  using pika::mpi::experimental::transform_mpi;
+  return std::forward<Sender>(sender)                                        //
+         | transform_mpi(ConsumeRvalues{MPICallHelper{std::forward<F>(f)}})  //
+         | drop_operation_state();
 }
 
 template <typename F>

miniapp/miniapp_band_to_tridiag.cpp

@@ -209,6 +209,5 @@ int main(int argc, char** argv) {
 
   pika::init_params p;
   p.desc_cmdline = desc_commandline;
-  p.rp_callback = dlaf::initResourcePartitionerHandler;
   return pika::init(pika_main, argc, argv, p);
 }

miniapp/miniapp_bt_band_to_tridiag.cpp

@@ -219,6 +219,5 @@ int main(int argc, char** argv) {
 
   pika::init_params p;
   p.desc_cmdline = desc_commandline;
-  p.rp_callback = dlaf::initResourcePartitionerHandler;
   return pika::init(pika_main, argc, argv, p);
 }

miniapp/miniapp_bt_reduction_to_band.cpp

@@ -238,6 +238,5 @@ int main(int argc, char** argv) {
 
   pika::init_params p;
   p.desc_cmdline = desc_commandline;
-  p.rp_callback = dlaf::initResourcePartitionerHandler;
   return pika::init(pika_main, argc, argv, p);
 }

miniapp/miniapp_cholesky.cpp

@@ -228,7 +228,6 @@ int main(int argc, char** argv) {
 
   pika::init_params p;
   p.desc_cmdline = desc_commandline;
-  p.rp_callback = dlaf::initResourcePartitionerHandler;
   return pika::init(pika_main, argc, argv, p);
 }
 

miniapp/miniapp_communication.cpp

@@ -606,6 +606,5 @@ int main(int argc, char** argv) {
 
   pika::init_params p;
   p.desc_cmdline = desc_commandline;
-  p.rp_callback = dlaf::initResourcePartitionerHandler;
   return pika::init(pika_main, argc, argv, p);
 }

miniapp/miniapp_eigensolver.cpp

@@ -257,7 +257,6 @@ int main(int argc, char** argv) {
 
   pika::init_params p;
   p.desc_cmdline = desc_commandline;
-  p.rp_callback = dlaf::initResourcePartitionerHandler;
   return pika::init(pika_main, argc, argv, p);
 }
 

miniapp/miniapp_gen_eigensolver.cpp

@@ -286,7 +286,6 @@ int main(int argc, char** argv) {
 
   pika::init_params p;
   p.desc_cmdline = desc_commandline;
-  p.rp_callback = dlaf::initResourcePartitionerHandler;
   return pika::init(pika_main, argc, argv, p);
 }
 

miniapp/miniapp_gen_to_std.cpp

@@ -216,6 +216,5 @@ int main(int argc, char** argv) {
 
   pika::init_params p;
   p.desc_cmdline = desc_commandline;
-  p.rp_callback = dlaf::initResourcePartitionerHandler;
   return pika::init(pika_main, argc, argv, p);
 }

miniapp/miniapp_reduction_to_band.cpp

@@ -257,6 +257,5 @@ int main(int argc, char** argv) {
 
   pika::init_params p;
   p.desc_cmdline = desc_commandline;
-  p.rp_callback = dlaf::initResourcePartitionerHandler;
   return pika::init(pika_main, argc, argv, p);
 }

miniapp/miniapp_triangular_multiplication.cpp

@@ -225,6 +225,5 @@ int main(int argc, char** argv) {
 
   pika::init_params p;
   p.desc_cmdline = desc_commandline;
-  p.rp_callback = dlaf::initResourcePartitionerHandler;
   return pika::init(pika_main, argc, argv, p);
 }

miniapp/miniapp_triangular_solver.cpp

@@ -244,6 +244,5 @@ int main(int argc, char** argv) {
 
   pika::init_params p;
   p.desc_cmdline = desc_commandline;
-  p.rp_callback = dlaf::initResourcePartitionerHandler;
   return pika::init(pika_main, argc, argv, p);
 }

miniapp/miniapp_tridiag_solver.cpp

@@ -225,6 +225,5 @@ int main(int argc, char** argv) {
 
   pika::init_params p;
   p.desc_cmdline = desc_commandline;
-  p.rp_callback = dlaf::initResourcePartitionerHandler;
   return pika::init(pika_main, argc, argv, p);
 }

spack/packages/dla-future/package.py

@@ -93,6 +93,7 @@ class DlaFuture(CMakePackage, CudaPackage, ROCmPackage):
     depends_on("[email protected]:", when="@0.3")
     depends_on("[email protected]:", when="@0.4.0:")
     conflicts("^[email protected]:", when="@:0.6")
+    depends_on("[email protected]:", when="@0.7.0:")
     depends_on("[email protected]:", when="@:0.2")
     depends_on("pika +mpi")
     depends_on("pika +cuda", when="+cuda")

src/c_api/init.cpp

@@ -26,7 +26,6 @@ void dlaf_initialize(int argc_pika, const char** argv_pika, int argc_dlaf,
 
     // pika initialization
     pika::init_params params;
-    params.rp_callback = dlaf::initResourcePartitionerHandler;
     params.desc_cmdline = desc;
     // After pika 0.21.0 pika::start reports errors only by exception and returns void
 #if PIKA_VERSION_FULL >= 0x001500