worker.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2025 VTT Technical Research Centre of Finland Ltd
// SPDX-License-Identifier: AGPL-3.0-or-later
 
#ifndef PFC_MPI_WORKER_HPP
#define PFC_MPI_WORKER_HPP
 
#include <mpi.h>
 
#include <iostream>
 
namespace pfc {
 
class MPI_Worker {
  MPI_Comm m_comm; 
  int m_rank;      
  int m_num_procs; 
  bool m_owns_mpi; 
 
public:
  MPI_Worker(int argc, char *argv[], MPI_Comm comm = MPI_COMM_WORLD,
             bool verbose = true)
      : m_comm(comm) {
    int initialized = 0;
    MPI_Initialized(&initialized);
    if (!initialized) {
      MPI_Init(&argc, &argv);
      m_owns_mpi = true;
    } else {
      m_owns_mpi = false;
    }
    MPI_Comm_rank(m_comm, &m_rank);
    MPI_Comm_size(m_comm, &m_num_procs);
    if (m_owns_mpi && verbose) {
      std::cout << "MPI_Init(): initialized " << m_num_procs << " processes"
                << std::endl;
    }
    if (m_rank != 0) {
      mute();
    }
  }
  ~MPI_Worker() {
    if (m_owns_mpi) {
      int finalized = 0;
      MPI_Finalized(&finalized);
      if (!finalized) {
        MPI_Finalize();
      }
    }
  }
 
  int get_rank() const { return m_rank; }
 
  int get_num_ranks() const { return m_num_procs; }
 
  void mute() { std::cout.setstate(std::ios::failbit); }
 
  void unmute() { std::cout.clear(); }
};
 
} // namespace pfc
 
#endif // PFC_MPI_WORKER_HPP