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 {
 
  class communicator {
    // Wrapped `MPI_Comm` object.
    MPI_Comm _com = MPI_COMM_WORLD;
 
    public:
    communicator() = default;
 
    communicator(MPI_Comm c) : _com(c) {}
 
    [[nodiscard]] MPI_Comm get() const noexcept { return _com; }
 
    [[nodiscard]] int rank() const {
      if (has_env) {
        int num = 0;
        check_mpi_call(MPI_Comm_rank(_com, &num), "MPI_Comm_rank");
        return num;
      } else
        return 0;
    }
 
    [[nodiscard]] int size() const {
      if (has_env) {
        int num = 0;
        check_mpi_call(MPI_Comm_size(_com, &num), "MPI_Comm_size");
        return num;
      } else
        return 1;
    }
 
    [[nodiscard]] communicator split(int color, int key = 0) const {
      if (has_env) {
        communicator c;
        check_mpi_call(MPI_Comm_split(_com, color, key, &c._com), "MPI_Comm_split");
        return c;
      } else
        return {};
    }
 
    [[nodiscard]] communicator duplicate() const {
      if (has_env) {
        communicator c;
        check_mpi_call(MPI_Comm_dup(_com, &c._com), "MPI_Comm_dup");
        return c;
      } else
        return {};
    }
 
    void free() {
      if (has_env) { check_mpi_call(MPI_Comm_free(&_com), "MPI_Comm_free"); }
    }
 
    void abort(int error_code) const {
      if (has_env)
        check_mpi_call(MPI_Abort(_com, 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(_com, boost::mpi::comm_duplicate); }
    inline communicator(boost::mpi::communicator c) : _com(c) {}
#endif // BOOST_MPI_HPP
 
    void barrier(long poll_msec = 1) const {
      if (has_env) {
        if (poll_msec == 0) {
          check_mpi_call(MPI_Barrier(_com), "MPI_Barrier");
        } else {
          MPI_Request req{};
          int flag = 0;
          check_mpi_call(MPI_Ibarrier(_com, &req), "MPI_Ibarrier");
          while (!flag) {
            check_mpi_call(MPI_Test(&req, &flag, MPI_STATUS_IGNORE), "MPI_Test");
            usleep(poll_msec * 1000);
          }
        }
      }
    }
  };
 
} // namespace mpi