communicator.hpp

Go to the documentation of this file.

// Copyright (c) 2024 Simons Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0.txt
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors: Thomas Hahn, Alexander Hampel, Olivier Parcollet, Nils Wentzell

 
#pragma once
 
#include "./environment.hpp"
#include "./utils.hpp"
 
#include <mpi.h>
 
#include <cstdlib>
#include <unistd.h>
 
namespace mpi {
 
  // Forward declaration.
  class shared_communicator;

  class communicator {
    public:
    communicator() = default;

    communicator(MPI_Comm c) : comm_(c) {}

    [[nodiscard]] MPI_Comm get() const noexcept { return comm_; }

    [[nodiscard]] bool is_null() const noexcept { return comm_ == MPI_COMM_NULL; }

    [[nodiscard]] int rank() const {
      int r = 0;
      if (has_env) check_mpi_call(MPI_Comm_rank(comm_, &r), "MPI_Comm_rank");
      return r;
    }

    [[nodiscard]] int size() const {
      int s = 1;
      if (has_env) check_mpi_call(MPI_Comm_size(comm_, &s), "MPI_Comm_size");
      return s;
    }

    [[nodiscard]] communicator split(int color, int key = 0) const {
      communicator c{};
      if (has_env) check_mpi_call(MPI_Comm_split(comm_, color, key, &c.comm_), "MPI_Comm_split");
      return c;
    }

    [[nodiscard]] shared_communicator split_shared(int split_type = MPI_COMM_TYPE_SHARED, int key = 0) const;

    [[nodiscard]] communicator duplicate() const {
      communicator c{};
      if (has_env) check_mpi_call(MPI_Comm_dup(comm_, &c.comm_), "MPI_Comm_dup");
      return c;
    }

    void free() {
      if (has_env && !is_null()) check_mpi_call(MPI_Comm_free(&comm_), "MPI_Comm_free");
    }

    void abort(int error_code) const {
      if (has_env) {
        check_mpi_call(MPI_Abort(comm_, error_code), "MPI_Abort");
      } else {
        std::abort();
      }
    }
 
#ifdef BOOST_MPI_HPP
    // Conversion to and from boost communicator, Keep for backward compatibility
    inline operator boost::mpi::communicator() const { return boost::mpi::communicator(comm_, boost::mpi::comm_duplicate); }
    inline communicator(boost::mpi::communicator c) : comm_(c) {}
#endif // BOOST_MPI_HPP

    void barrier(long poll_msec = 1) const {
      if (has_env) {
        if (poll_msec == 0) {
          check_mpi_call(MPI_Barrier(comm_), "MPI_Barrier");
        } else {
          MPI_Request req{};
          int flag = 0;
          check_mpi_call(MPI_Ibarrier(comm_, &req), "MPI_Ibarrier");
          while (!flag) {
            check_mpi_call(MPI_Test(&req, &flag, MPI_STATUS_IGNORE), "MPI_Test");
            usleep(poll_msec * 1000);
          }
        }
      }
    }
 
    private:
    MPI_Comm comm_ = MPI_COMM_WORLD;
  };

  class shared_communicator : public communicator {
    public:
    // Make the constructors of mpi::communicator accessible.
    using communicator::communicator;

    shared_communicator() : communicator(MPI_COMM_NULL) {}
  };
 
  [[nodiscard]] inline shared_communicator communicator::split_shared(int split_type, int key) const {
    shared_communicator c{};
    if (has_env) check_mpi_call(MPI_Comm_split_type(comm_, split_type, key, MPI_INFO_NULL, &c.comm_), "MPI_Comm_split_type");
    return c;
  }
 
} // namespace mpi