comm.hpp

Go to the documentation of this file.

// Copyright (c) 2014-2023, Lawrence Livermore National Security, LLC.
// Produced at the Lawrence Livermore National Laboratory.
// Written by the LBANN Research Team (B. Van Essen, et al.) listed in
// the CONTRIBUTORS file. <lbann-dev@llnl.gov>
//
// LLNL-CODE-697807.
// All rights reserved.
//
// This file is part of LBANN: Livermore Big Artificial Neural Network
// Toolkit. For details, see http://software.llnl.gov/LBANN or
// https://github.com/LLNL/LBANN.
//
// Licensed under the Apache License, Version 2.0 (the "Licensee"); 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
//
// 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.

#ifndef LBANN_COMM_HPP_INCLUDED
#define LBANN_COMM_HPP_INCLUDED

#include "base.hpp"

#ifdef LBANN_HAS_CUDA
#include <cuda_runtime.h>
#endif // LBANN_HAS_CUDA
#ifdef LBANN_HAS_ALUMINUM
#include <Al.hpp>
#endif // LBANN_HAS_ALUMINUM

#include "lbann/comm_nb_request.hpp"

#include "detect_El_mpi.hpp"

#include <map>
#include <typeindex>
#include <vector>

namespace lbann {

#ifdef LBANN_HAS_ALUMINUM

::Al::ReductionOperator mpi_op_to_al_op(El::mpi::Op op);
#endif

enum class GridType
{
  NO_GRID = 0,
  PRIMARY_GRID = 1,
  SECONDARY_GRID = 2
};

/* Notes on Synchronization
 *
 * The updated interface exposes a synchronization handle/device
 * tagging mechanism used by Hydrogen: El::SyncInfo<D>, where D is an
 * El::Device. When operating on Matrix objects, this should be
 * handled automagically, assuming the Matrix is setup properly. Users
 * must be aware of this when making MPI calls through Hydrogen or
 * through lbann_comm with raw data buffers (T[]).
 *
 * When dealing with El::Matrix objects, users should be aware of the
 * following. There is no synchronization for CPU objects
 * (El::SyncInfo<El::Device::CPU> is an empty struct), but GPU Matrix
 * objects now have an associated stream and event. These are
 * GPUManager::Stream() and GPUManager::Event() by default, resp., but
 * can be overriden by a user. Note: the Matrix never owns these; it
 * will not free these resources at destruction. There are many
 * methods in which multiple El::Matrix objects might interact. This
 * should work properly; otherwise, report bugs to benson31.
 *
 * When dealing with raw data (T[]), users should be aware of the
 * following. In the near future, all El::mpi functions will have an
 * El::SyncInfo object as their last parameter, and it will be a
 * required parameter. In lbann_comm, this means that when the call
 * trickles down to an El::mpi function, an appropriate El::SyncInfo
 * must be available. Since many of LBANN's uses of this interface are
 * for communicating CPU buffers, there is "shortcut" API that assumes
 * the data is CPU memory, thus providing the default
 * El::SyncInfo<El::Device::CPU> object to El::mpi. If a user wishes
 * to communicate GPU data, they must use the "full" API, which adds a
 * final El::SyncInfo parameter to the function. This ensures the
 * appropriate synchronization semantics, especially when working with
 * Aluminum as the communication frontend.
 */

class lbann_comm
{
public:
  lbann_comm(int procs_per_trainer = 0,
             El::mpi::Comm world = El::mpi::COMM_WORLD.GetMPIComm());
  lbann_comm(const lbann_comm&) = delete;
  lbann_comm& operator=(const lbann_comm&) = delete;
  ~lbann_comm();

  void split_trainers(int procs_per_trainer = -1, int trainer_grid_height = -1);

  void split_trainer_grid(int num_process_primary_grid = 0,
                          bool create_two_models = false,
                          bool enable_async_comm = false,
                          bool enable_topo_aware = false);

  inline GridType get_grid_type() const noexcept { return m_grid_type; }

  inline int get_trainer_rank() const noexcept { return m_trainer_rank; }
  inline int get_rank_in_trainer() const noexcept { return m_rank_in_trainer; }
  inline int get_rank_in_world() const
  {
    return El::mpi::Rank(get_world_comm());
  }
  inline int get_world_rank(int trainer, int rank) const noexcept
  {
    if (m_secondary_grid_ranks.size() == 0) {
      return m_procs_per_trainer * trainer + rank;
    }
    else {
      return (m_secondary_grid_ranks.size() + m_primary_grid_ranks.size()) *
               trainer +
             rank;
    }
  }
  inline int map_world_rank_to_trainer_rank(int world_rank) const noexcept
  {
    return (world_rank / m_procs_per_trainer);
  }
  inline int map_world_rank_to_rank_in_trainer(int world_rank) const noexcept
  {
    return (world_rank % m_procs_per_trainer);
  }
  inline int get_trainer_master() const noexcept { return 0; }
  inline int get_intertrainer_master() const noexcept { return 0; }
  inline int get_world_master() const noexcept { return 0; }
  inline bool am_trainer_master() const noexcept
  {
    return get_rank_in_trainer() == get_trainer_master();
  }
  inline bool am_world_master() const noexcept
  {
    return get_rank_in_world() == get_world_master();
  }
  inline El::Grid& get_trainer_grid() { return *m_grid; }
  inline const El::Grid& get_trainer_grid() const { return *m_grid; }
  inline El::Grid& get_secondary_grid() { return *m_secondary_grid; }
  inline const El::Grid& get_secondary_grid() const
  {
    return *m_secondary_grid;
  }
  inline El::Grid& get_subset_grid() { return *m_subset_grid; }
  inline const El::Grid& get_subset_grid() const { return *m_subset_grid; }
  inline int get_num_trainers() const noexcept { return m_num_trainers; }
  /* Return the number of processes in a trainer. */
  inline int get_procs_per_trainer() const noexcept
  {
    return m_procs_per_trainer;
  }
  inline int get_procs_per_node() const noexcept { return m_procs_per_node; }
  inline int get_procs_in_world() const
  {
    return El::mpi::Size(get_world_comm());
  }
  inline int get_rank_in_node() const noexcept { return m_rank_in_node; }
  inline bool is_world_rank_on_node(int rank) const
  {
    return std::find(m_world_ranks_on_node.begin(),
                     m_world_ranks_on_node.end(),
                     rank) != m_world_ranks_on_node.end();
  }

  inline int get_default_threads_per_proc() const noexcept
  {
    return m_threads_per_proc;
  }

  void reset_threads() const noexcept;

  void intertrainer_sum_matrix(AbsMat& mat) const;
  void intertrainer_sum_matrix(AbsDistMat& mat) const;
  void intertrainer_broadcast_matrix(AbsMat& mat, int root) const;
  void intertrainer_broadcast_matrix(AbsDistMat& mat, int root) const;

  template <typename T, bool S = is_instantiated_El_mpi_type<T>::value>
  void broadcast(int root, T& val, const El::mpi::Comm& c) const;

  template <typename T>
  void broadcast_custom(int root, T& val, const El::mpi::Comm& c) const;
  template <typename T>
  void broadcast_native(int root, T& val, const El::mpi::Comm& c) const;

  template <typename T>
  void world_broadcast(int root, T& val) const;
  template <typename T>
  void intertrainer_broadcast(int root, T& val) const;
  template <typename T>
  void trainer_broadcast(int root, T& val) const;

  // Default to cpu memory
  template <typename T>
  void broadcast(const int root,
                 T* data,
                 const int count,
                 const El::mpi::Comm& c) const;

  template <typename T,
            El::Device D,
            bool S = is_instantiated_El_mpi_type<T>::value>
  void broadcast(const int root,
                 T* data,
                 const int count,
                 const El::mpi::Comm& c,
                 El::SyncInfo<D> const& syncInfo) const;

  template <typename T>
  void world_broadcast(const int root, T* data, const int count) const;

  template <typename T, El::Device D>
  void world_broadcast(const int root,
                       T* data,
                       const int count,
                       El::SyncInfo<D> const& syncInfo) const;
  template <typename T>
  void intertrainer_broadcast(const int root, T* data, const int count) const;
  template <typename T, El::Device D>
  void intertrainer_broadcast(const int root,
                              T* data,
                              const int count,
                              El::SyncInfo<D> const& syncInfo) const;
  template <typename T>
  void trainer_broadcast(const int root, T* data, const int count) const;

  template <typename T, El::Device D>
  void trainer_broadcast(const int root,
                         T* data,
                         const int count,
                         El::SyncInfo<D> const& syncInfo) const;

  template <typename T>
  size_t
  resize(const int root, std::vector<T>& data, const El::mpi::Comm& c) const;

  template <typename T>
  void
  broadcast(const int root, std::vector<T>& data, const El::mpi::Comm& c) const;
  template <typename T>
  void world_broadcast(int root, std::vector<T>& data) const;
  template <typename T>
  void intertrainer_broadcast(int root, std::vector<T>& data) const;
  template <typename T>
  void trainer_broadcast(int root, std::vector<T>& data) const;

  template <typename T>
  void all_gather(const T* src,
                  int src_count,
                  T* rcv,
                  int rcv_count,
                  const El::mpi::Comm& c) const;
  template <typename T, El::Device D>
  void all_gather(const T* src,
                  int src_count,
                  T* rcv,
                  int rcv_count,
                  const El::mpi::Comm& c,
                  El::SyncInfo<D> const& syncInfo) const;

  template <typename T>
  void all_gather(std::vector<T> const& src,
                  std::vector<T>& rcs,
                  std::vector<int> const& rcv_counts,
                  std::vector<int> const& rcv_disp,
                  const El::mpi::Comm& c) const;
  template <typename T>
  void trainer_all_gather(std::vector<T> const& src,
                          std::vector<T>& rcs,
                          std::vector<int> const& rcv_counts,
                          std::vector<int> const& rcv_disp) const;
  template <typename T>
  void
  all_gather(T const& src, std::vector<T>& data, const El::mpi::Comm& c) const;
  template <typename T>
  void world_all_gather(T const& src, std::vector<T>& data) const;
  template <typename T>
  void trainer_all_gather(T const& src, std::vector<T>& data) const;

  template <typename T>
  void trainer_gather(T snd, int root) const;
  template <typename T>
  void trainer_gather(T snd, T* rcv) const;
  template <typename T>
  void trainer_gather(T const* snd, int count, int root) const;
  template <typename T>
  void trainer_gather(T const* snd, int count, T* rcv) const;
  template <typename T>
  void trainer_gatherv(T const* snd, int count, int root) const;
  template <typename T>
  void trainer_gatherv(T const* snd,
                       int count,
                       T* rcv,
                       int const* rcv_counts,
                       int const* rcv_displacements) const;
  template <typename T>
  void intertrainer_gather(T snd, int root) const;
  template <typename T>
  void intertrainer_gather(T snd, std::vector<T>& rcv) const;
  template <typename T>
  void intertrainer_gather(T const* snd, int count, int root) const;
  template <typename T>
  void intertrainer_gather(T const* snd, int count, T* rcv) const;
  template <typename T>
  void gather(T snd, int root, const El::mpi::Comm& c) const;
  template <typename T>
  void gather(T snd, T* rcv, const El::mpi::Comm& c) const;
  template <typename T>
  void gather(T snd, std::vector<T>& rcv, const El::mpi::Comm& c) const;
  template <typename T>
  void gather(T const* snd, int count, int root, const El::mpi::Comm& c) const;
  template <typename T, El::Device D>
  void gather(T const* snd,
              int count,
              int root,
              const El::mpi::Comm& c,
              El::SyncInfo<D> const& syncInfo) const;
  template <typename T>
  void gather(T const* snd, int count, T* rcv, const El::mpi::Comm& c) const;
  template <typename T, El::Device D>
  void gather(T const* snd,
              int count,
              T* rcv,
              const El::mpi::Comm& c,
              El::SyncInfo<D> const& syncInfo) const;
  template <typename T>
  T scatter(int root, const El::mpi::Comm& c) const;
  template <typename T>
  T scatter(T const* snd, const El::mpi::Comm& c) const;
  template <typename T>
  void
  intertrainer_reduce(T snd, int root, El::mpi::Op op = El::mpi::SUM) const;
  template <typename T>
  T intertrainer_reduce(T snd, El::mpi::Op op = El::mpi::SUM) const;
  template <typename T>
  void trainer_reduce(T snd, int root, El::mpi::Op op = El::mpi::SUM) const;
  template <typename T>
  T trainer_reduce(T snd, El::mpi::Op op = El::mpi::SUM) const;
  template <typename T>
  void trainer_reduce(T const* snd,
                      int count,
                      int root,
                      El::mpi::Op op = El::mpi::SUM) const;
  template <typename T>
  void trainer_reduce(T const* snd,
                      int count,
                      T* rcv,
                      El::mpi::Op op = El::mpi::SUM) const;
  template <typename T>
  void reduce(T snd,
              int root,
              const El::mpi::Comm& c,
              El::mpi::Op op = El::mpi::SUM) const;
  template <typename T>
  T reduce(T snd, const El::mpi::Comm& c, El::mpi::Op op = El::mpi::SUM) const;

  // Op is "SUM"
  template <typename T>
  void reduce(T const* snd, int count, int root, const El::mpi::Comm& c) const;
  template <typename T, El::Device D>
  void reduce(T const* snd,
              int count,
              int root,
              const El::mpi::Comm& c,
              El::SyncInfo<D> const& syncInfo) const;

  template <typename T>
  void reduce(T const* snd,
              int count,
              int root,
              const El::mpi::Comm& c,
              El::mpi::Op op) const;
  template <typename T, El::Device D>
  void reduce(T const* snd,
              int count,
              int root,
              const El::mpi::Comm& c,
              El::mpi::Op op,
              El::SyncInfo<D> const& syncInfo) const;
  template <typename T, El::Device D>
  void reduce(T const* snd,
              int count,
              T* rcv,
              const El::mpi::Comm& c,
              El::SyncInfo<D> const& syncInfo) const;
  template <typename T>
  void reduce(T const* snd, int count, T* rcv, const El::mpi::Comm& c) const;

  template <typename T>
  void reduce(T const* snd,
              int count,
              T* rcv,
              const El::mpi::Comm& c,
              El::mpi::Op op) const;
  template <typename T, El::Device D>
  void reduce(T const* snd,
              int count,
              T* rcv,
              const El::mpi::Comm& c,
              El::mpi::Op op,
              El::SyncInfo<D> const& syncInfo) const;
  template <typename T>
  T intertrainer_allreduce(T snd, El::mpi::Op op = El::mpi::SUM) const;
  template <typename T>
  T trainer_allreduce(T snd, El::mpi::Op op = El::mpi::SUM) const;
  template <typename T>
  void trainer_allreduce(T const* snd,
                         int count,
                         T* rcv,
                         El::mpi::Op op = El::mpi::SUM) const;
  template <typename T>
  T allreduce(T snd,
              const El::mpi::Comm& c,
              El::mpi::Op op = El::mpi::SUM) const;

  // FIXME (trb): Based on the backend choice of "MPIBackend", I'm
  // assuming this is intended as a CPU-only call.
  template <typename T>
  void allreduce(T const* snd,
                 int count,
                 T* rcv,
                 const El::mpi::Comm& c,
                 El::mpi::Op op = El::mpi::SUM) const;
  template <typename T>
  void allreduce(T* data,
                 int count,
                 const El::mpi::Comm& c,
                 El::mpi::Op op = El::mpi::SUM) const;
  template <typename TensorDataType>
  void allreduce(El::AbstractMatrix<TensorDataType>& m,
                 const El::mpi::Comm& c,
                 El::mpi::Op op = El::mpi::SUM) const;
  template <typename TensorDataType>
  void allreduce(El::AbstractDistMatrix<TensorDataType>& m,
                 const El::mpi::Comm& c,
                 El::mpi::Op op = El::mpi::SUM) const;
  template <typename TensorDataType>
  void nb_allreduce(El::AbstractMatrix<TensorDataType>& m,
                    const El::mpi::Comm& c,
                    Al::request& req,
                    El::mpi::Op op = El::mpi::SUM) const;
  template <typename TensorDataType>
  void nb_allreduce(El::AbstractDistMatrix<TensorDataType>& m,
                    const El::mpi::Comm& c,
                    Al::request& req,
                    El::mpi::Op op = El::mpi::SUM) const;
  template <typename T>
  void nb_allreduce(T* data,
                    int count,
                    const El::mpi::Comm& c,
                    Al::request& req,
                    El::mpi::Op op = El::mpi::SUM) const;

  template <typename T>
  void wait_all(std::vector<El::mpi::Request<T>>& req) const;

  template <typename T>
  void wait(El::mpi::Request<T>& req) const;

  void wait(Al::request& req) const;
  bool test(Al::request& req) const;

  void intertrainer_barrier() const;
  void trainer_barrier() const;
  void global_barrier() const;
  void barrier(const El::mpi::Comm& c) const;

  template <typename T>
  void send(const T* data, int count, int trainer, int rank) const;
  template <typename T, El::Device D>
  void send(const T* data,
            int count,
            int trainer,
            int rank,
            El::SyncInfo<D> const& syncInfo) const;
  template <typename T, El::Device D>
  void send(const T* data,
            int count,
            int trainer,
            El::SyncInfo<D> const& syncInfo) const;
  void send(const AbsMat& mat, int trainer, int rank) const;
  void send(const DistMat& mat, int trainer, int rank) const;
  void send(const AbsMat& mat, int trainer) const
  {
    send(mat, trainer, m_rank_in_trainer);
  }
  void send(const DistMat& mat, int trainer) const
  {
    send(mat, trainer, m_rank_in_trainer);
  }

  template <typename T>
  void nb_send(const T* data,
               int count,
               int trainer,
               int rank,
               El::mpi::Request<T>& req) const;
  template <typename T>
  void nb_tagged_send(const T* data,
                      int count,
                      int rank,
                      int tag,
                      El::mpi::Request<T>& req,
                      const El::mpi::Comm& c) const;
  template <typename T>
  void nb_send(const T* data,
               int count,
               int trainer,
               El::mpi::Request<T>& req) const;
  void nb_send(const AbsMat& mat,
               int trainer,
               int rank,
               El::mpi::Request<DataType>& req) const;
  void nb_send(const DistMat& mat,
               int trainer,
               int rank,
               El::mpi::Request<DataType>& req) const;
  void
  nb_send(const AbsMat& mat, int trainer, El::mpi::Request<DataType>& req) const
  {
    nb_send(mat, trainer, m_rank_in_trainer, req);
  }
  void nb_send(const DistMat& mat,
               int trainer,
               El::mpi::Request<DataType>& req) const
  {
    nb_send(mat, trainer, m_rank_in_trainer, req);
  }

  template <typename T>
  void recv(T* data, int count, int trainer, int rank) const;
  template <typename T>
  void recv(T* data, int count, int trainer) const;
  template <typename T>
  void recv(T* data, int count) const;
  template <typename T, El::Device D>
  void recv(T* data,
            int count,
            int trainer,
            int rank,
            El::SyncInfo<D> const& syncInfo) const;
  template <typename T, El::Device D>
  void
  recv(T* data, int count, int trainer, El::SyncInfo<D> const& syncInfo) const;
  void recv(AbsMat& mat, int trainer, int rank) const;
  void recv(DistMat& mat, int trainer, int rank) const;
  void recv(AbsMat& mat, int trainer) const
  {
    recv(mat, trainer, m_rank_in_trainer);
  }
  void recv(DistMat& mat, int trainer) const
  {
    recv(mat, trainer, m_rank_in_trainer);
  }
  template <typename T, El::Device D>
  void recv(T* data, int count, El::SyncInfo<D> const& syncInfo) const;
  void recv(AbsMat& mat) const;
  void recv(DistMat& mat) const;

  template <typename T>
  void nb_recv(T* data,
               int count,
               int trainer,
               int rank,
               El::mpi::Request<T>& req) const;
  template <typename T>
  void nb_tagged_recv(T* data,
                      int count,
                      int rank,
                      int tag,
                      El::mpi::Request<T>& req,
                      const El::mpi::Comm& c) const;

  template <typename T>
  void nb_recv(T* data, int count, int trainer, El::mpi::Request<T>& req) const;
  void nb_recv(AbsMat& mat,
               int trainer,
               int rank,
               El::mpi::Request<DataType>& req) const;
  void nb_recv(DistMat& mat,
               int trainer,
               int rank,
               El::mpi::Request<DataType>& req) const;
  void nb_recv(AbsMat& mat, int trainer, El::mpi::Request<DataType>& req) const
  {
    nb_recv(mat, trainer, m_rank_in_trainer, req);
  }
  void nb_recv(DistMat& mat, int trainer, El::mpi::Request<DataType>& req) const
  {
    nb_recv(mat, trainer, m_rank_in_trainer, req);
  }
  template <typename T>
  void nb_recv(T* data, int count, El::mpi::Request<T>& req) const;
  void nb_recv(AbsMat& mat, El::mpi::Request<DataType>& req) const;
  void nb_recv(DistMat& mat, El::mpi::Request<DataType>& req) const;

  template <typename T, El::Device D>
  void sendrecv(const T* snd,
                int send_count,
                int send_trainer,
                int send_rank,
                T* rcv,
                int recv_count,
                int recv_trainer,
                int recv_rank) const;
  template <typename T, El::Device D>
  void sendrecv(const T* snd,
                int send_count,
                int send_trainer,
                T* rcv,
                int recv_count,
                int recv_trainer) const;

  template <typename T, El::Device D>
  void sendrecv(const T* snd,
                int send_count,
                int send_trainer,
                int send_rank,
                T* rcv,
                int recv_count,
                int recv_trainer,
                int recv_rank,
                El::SyncInfo<D> const& syncInfo) const;
  template <typename T, El::Device D>
  void sendrecv(const T* snd,
                int send_count,
                int send_trainer,
                T* rcv,
                int recv_count,
                int recv_trainer,
                El::SyncInfo<D> const& syncInfo) const;

  template <typename T>
  int get_count(int trainer, int rank) const;
  template <typename T>
  int get_count(int trainer) const;

  // Statistics methods.
  inline size_t get_num_trainer_barriers() const noexcept
  {
    return m_num_trainer_barriers;
  }
  inline size_t get_num_intertrainer_barriers() const noexcept
  {
    return m_num_intertrainer_barriers;
  }
  inline size_t get_num_global_barriers() const noexcept
  {
    return m_num_global_barriers;
  }
  inline size_t get_bytes_sent() const noexcept { return m_bytes_sent; }
  inline size_t get_bytes_received() const noexcept { return m_bytes_received; }

  inline void reset_stats_counters() noexcept
  {
    m_num_trainer_barriers = 0;
    m_num_intertrainer_barriers = 0;
    m_num_global_barriers = 0;
    m_bytes_sent = 0;
    m_bytes_received = 0;
  }

  static inline bool is_sendable(const AbsMat& mat) noexcept
  {
    // This assumes we do not transmit mat with a datatype smaller than
    // DataType.
    // MPI uses "int" as its count type; do calculations with larger ints.
    size_t count = (size_t)mat.Height() * (size_t)mat.Width();
    return count <= (size_t)std::numeric_limits<int>::max();
  }
  static inline bool is_sendable(const AbsDistMat& dist_mat) noexcept
  {
    return is_sendable(dist_mat.LockedMatrix());
  }

  const El::mpi::Comm& get_intertrainer_comm() const noexcept
  {
    return m_intertrainer_comm;
  }

  const El::mpi::Comm& get_trainer_comm() const noexcept
  {
    return m_trainer_comm;
  }

  const El::mpi::Comm& get_combined_grid_comm() const noexcept
  {
    return m_combined_grid_comm;
  }

  const El::mpi::Comm& get_world_comm() const noexcept { return m_world_comm; }

  const El::mpi::Comm& get_node_comm() const noexcept { return m_node_comm; }

  const El::mpi::Comm& get_KFAC_comm() const noexcept { return m_trainer_comm; }

  std::vector<int> get_primary_grid_ranks() { return m_primary_grid_ranks; }

  std::vector<int> get_secondary_grid_ranks() { return m_secondary_grid_ranks; }

  bool get_KFAC_subgrid_create_two_models() { return m_create_two_models; }

  bool enable_subgrid_async_communication() { return m_subgrid_async_progress; }

  const El::mpi::Comm& get_packed_group_comm(int num_per_group) const;

  bool is_rank_node_local(int rank, const El::mpi::Comm& comm) const
  {
    // Translating to COMM_WORLD is typically constant time.
    int world_rank = El::mpi::Translate(comm, rank, get_world_comm());
    return is_world_rank_on_node(world_rank);
  }

  void lbann_comm_abort(std::string msg) const;

private:
  const El::mpi::Comm m_world_comm;
  El::mpi::Comm m_trainer_comm;
  El::mpi::Comm m_intertrainer_comm;
  El::mpi::Comm m_node_comm;
  El::mpi::Comm m_primary_grid_comm;
  El::mpi::Comm m_secondary_grid_comm;
  El::mpi::Comm m_combined_grid_comm;
  mutable std::unordered_map<int, El::mpi::Comm> m_group_communicators;
  std::unique_ptr<El::Grid> m_grid;
  int m_num_trainers;
  int m_procs_per_trainer;
  int m_trainer_rank;
  int m_rank_in_trainer;
  int m_procs_per_node;
  int m_rank_in_node;
  std::vector<int> m_world_ranks_on_node;
  int m_threads_per_proc;

  GridType m_grid_type = GridType::NO_GRID;

  bool m_create_two_models = false, m_subgrid_async_progress = false;

  std::unique_ptr<El::Grid> m_secondary_grid, m_subset_grid;

  std::vector<int> m_primary_grid_ranks;
  std::vector<int> m_secondary_grid_ranks;

  // Various statistics counters.
  mutable size_t m_num_trainer_barriers;
  mutable size_t m_num_intertrainer_barriers;
  mutable size_t m_num_global_barriers;
  mutable size_t m_bytes_sent;
  mutable size_t m_bytes_received;

  void setup_node_comm();

  void setup_threads();

  void count_bytes_broadcast(const size_t bytes,
                             const int rank,
                             const int root) const noexcept
  {
    if (rank == root) {
      m_bytes_sent += bytes;
    }
    else {
      m_bytes_received += bytes;
    }
  }
}; // class lbann_comm

int get_rank_in_world();

} // namespace lbann

#endif // LBANN_COMM_HPP_INCLUDED