text: ////////////////////////////////////////////////////////////////////////////////
  // Copyright (c) 2014-2022, 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 "detect_El_mpi.hpp"
  
  #include <map>
  #include <typeindex>
  #include <vector>
  
  namespace lbann {
  
  #ifdef LBANN_HAS_ALUMINUM
  /** Convert an MPI_Op to an Aluminum reduction operator. */
  ::Al::ReductionOperator mpi_op_to_al_op(El::mpi::Op op);
  #endif
  
  /** Grid types in sub-grid parallelism (2nd order) */
  enum class GridType 
  {   
    NO_GRID = 0, 
    PRIMARY_GRID = 1, 
    SECONDARY_GRID = 2
  };
  
  
  namespace Al {
  
  /** Dummy Aluminum backend. */
  class dummy_backend
  {
  public:
    using req_type = int;
    static constexpr req_type null_req = 0;
  };
  
  // Define aliases for Aluminum backends
  #ifdef LBANN_HAS_ALUMINUM
  using mpi_backend = ::Al::MPIBackend;
  #else
  using mpi_backend = lbann::Al::dummy_backend;
  #endif // LBANN_HAS_ALUMINUM
  using mpi_req_type = mpi_backend::req_type;
  static const mpi_req_type mpi_null_req = mpi_backend::null_req;
  /// @todo MPI-CUDA backend
  #if defined(LBANN_HAS_ALUMINUM) && defined(AL_HAS_NCCL)
  using nccl_backend = ::Al::NCCLBackend;
  // LBANN does its own synchronization on this.
  #else
  using nccl_backend = lbann::Al::dummy_backend;
  #endif // defined(LBANN_HAS_ALUMINUM) && defined(AL_HAS_NCCL)
  using nccl_req_type = nccl_backend::req_type;
  static const nccl_req_type nccl_null_req = nccl_backend::null_req;
  #if defined(LBANN_HAS_ALUMINUM) && defined(AL_HAS_MPI_CUDA)
  using mpicuda_backend = ::Al::MPICUDABackend;
  #else
  using mpicuda_backend = lbann::Al::dummy_backend;
  #endif // defined(LBANN_HAS_ALUMINUM) && defined(AL_HAS_MPI_CUDA)
  #if defined(LBANN_HAS_ALUMINUM) && defined(AL_HAS_HOST_TRANSFER)
  using hosttransfer_backend = ::Al::HostTransferBackend;
  #else
  using hosttransfer_backend = lbann::Al::dummy_backend;
  #endif // defined(LBANN_HAS_ALUMINUM) && defined(AL_HAS_HOST_TRANSFER)
  using mpicuda_req_type = mpicuda_backend::req_type;
  static const mpicuda_req_type mpicuda_null_req = mpicuda_backend::null_req;
  
  /** Wrapper for Aluminum non-blocking routine requests. */
  struct request
  {
    mpi_req_type mpi_req = mpi_null_req;
    nccl_req_type nccl_req = nccl_null_req;
    mpicuda_req_type mpicuda_req = mpicuda_null_req;
    MPI_Request raw_mpi_req = MPI_REQUEST_NULL;
  };
  } // namespace Al
  
  /* 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.
   */
  
  /**
   * Manage communication.
   * This supports separate trainers, each of which are split over potentially
   * several processes. Every trainer is split over the same number of processes.
   * The corresponding processes between trainers are on the "inter-trainer
   * communicator".
   * You can also do point-to-point or broadcast communication to arbitrary sets
   * of processes.
   */
  class lbann_comm
  {
  public:
    /**
     * Init communicators for trainers each with procs_per_trainer processes,
     * defaulting to every process in one trainer.
     */
    lbann_comm(int procs_per_trainer = 0,
               El::mpi::Comm world = El::mpi::COMM_WORLD.GetMPIComm());
    /** Don't allow copying; it doesn't make sense for the communicator. */
    lbann_comm(const lbann_comm&) = delete;
    /** Don't allow assignment; it doesn't make sense for the communicator. */
    lbann_comm& operator=(const lbann_comm&) = delete;
    ~lbann_comm();
  
    /** @brief Construct communicators for trainers
     *
     *  Invalidates any existing trainer communicators.
     *
     *  @param procs_per_trainer Number of MPI ranks in a trainer.
     *  Default is size of world communicator.
     *  @param trainer_grid_height Height of 2D process grid for each
     *  trainer. Must divide @c procs_per_trainer. Default grid is
     *  approximately square.
     */
    void split_trainers(int procs_per_trainer=-1, int trainer_grid_height=-1);
  
    /** Split the commicator for the given trainer into primary and seconday*/
    void split_trainer_grid(int num_process_primary_grid=0, bool create_two_models=false);
  
    /** Get trainer grid number (0: no primary/secondary grid, 1: part of primary grid, 2: part of secondary grid). */
    inline GridType get_grid_type() const noexcept { return m_grid_type; }
  
    /** Get which trainer this process is in. */
    inline int get_trainer_rank() const noexcept { return m_trainer_rank; }
    /** Get the rank of this process in its trainer. */
    inline int get_rank_in_trainer() const noexcept { return m_rank_in_trainer; }
    /** Get my rank in COMM_WORLD. */
    inline int get_rank_in_world() const
    {
      return El::mpi::Rank(get_world_comm());
    }
    /** Return the COMM_WORLD rank of the rank'th processor in trainer. */
    inline int get_world_rank(int trainer, int rank) const noexcept
    {
      return m_procs_per_trainer * trainer + rank;
    }
    /** Return the "rank" of the trainer that this rank is in */
    inline int map_world_rank_to_trainer_rank(int world_rank) const noexcept
    {
      return (world_rank / m_procs_per_trainer);
    }
    /** Return the "rank" within the trainer that this rank is in */
    inline int map_world_rank_to_rank_in_trainer(int world_rank) const noexcept
    {
      return (world_rank % m_procs_per_trainer);
    }
    /** Return the rank of the master process in this trainer. */
    inline int get_trainer_master() const noexcept { return 0; }
    /** Return the rank of the inter-trainer master process. */
    inline int get_intertrainer_master() const noexcept { return 0; }
    /** Return the rank of the world master process. */
    inline int get_world_master() const noexcept { return 0; }
    /** Return true if this process is the master process in its trainer. */
    inline bool am_trainer_master() const noexcept
    {
      return get_rank_in_trainer() == get_trainer_master();
    }
    /** Return true if this process is the world master process. */
    inline bool am_world_master() const noexcept
    {
      return get_rank_in_world() == get_world_master();
    }
    /** Return a grid to use for this trainer. */
    inline El::Grid& get_trainer_grid() { return *m_grid; }
    /** Return a read-only grid to use for this trainer. */
    inline const El::Grid& get_trainer_grid() const { return *m_grid; }
    /** Return secondary grid to use for this trainer. */
    inline El::Grid& get_secondary_grid() { return *m_secondary_grid; }
    /** Return read-only secondary grid to use for this trainer. */
    inline const El::Grid& get_secondary_grid() const { return *m_secondary_grid; }
    /** Return the total number of trainers. */
    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;
    }
    /** Return the number of processes in a compute node. */
    inline int get_procs_per_node() const noexcept { return m_procs_per_node; }
    /** Return the total number of ranks. */
    inline int get_procs_in_world() const
    {
      return El::mpi::Size(get_world_comm());
    }
    /** Return the rank of this process within its compute node. */
    inline int get_rank_in_node() const noexcept { return m_rank_in_node; }
    /** Return true if rank (in COMM_WORLD) is on this compute 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();
    }
  
    /** Get default number of threads per process.
     *  This is the number of OpenMP threads to use for parallel
     *  regions, provided omp_set_num_threads has not been called or the
     *  num_threads directive has not been provided.
     */
    inline int get_default_threads_per_proc() const noexcept
    {
      return m_threads_per_proc;
    }
  
    /** Reset the number of threads per process to the default. */
    void reset_threads() const noexcept;
  
    /** Perform a sum reduction of mat over the inter-trainer communicator. */
    void intertrainer_sum_matrix(AbsMat& mat) const;
    void intertrainer_sum_matrix(AbsDistMat& mat) const;
    /** Broadcast mat over the inter-trainer communicator starting from root. */
    void intertrainer_broadcast_matrix(AbsMat& mat, int root) const;
    void intertrainer_broadcast_matrix(AbsDistMat& mat, int root) const;
  
    /// Broadcast a scalar value over an arbitrary communicator
    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;
  
    /// World broadcast of a scalar.
    template <typename T> void world_broadcast(int root, T& val) const;
    /// Inter-trainer broadcast of a scalar.
    template <typename T> void intertrainer_broadcast(int root, T& val) const;
    /// Within-trainer broadcast of a scalar.
    template <typename T> void trainer_broadcast(int root, T& val) const;
  
    /**
     * Broadcast a buffer over an arbitrary communicator assuming that
     * the buffer space is already allocated.
     */
  
    // 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;
  
    /// World broadcast of a buffer.
    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;
    /// Inter-trainer broadcast of a buffer.
    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;
    /// Within-trainer broadcast of a buffer.
    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;
  
    /**
     * Resize vector<> over an arbitrary communicator to match the one on root.
     */
    template <typename T>
    size_t
    resize(const int root, std::vector<T>& data, const El::mpi::Comm& c) const;
  
    /**
     * Broadcast vector<> over an arbitrary communicator;
     * vector<> for non-root processes will be resized as needed.
     */
    template <typename T>
    void
    broadcast(const int root, std::vector<T>& data, const El::mpi::Comm& c) const;
    /// Broadcast vector<> to world.
    template <typename T>
    void world_broadcast(int root, std::vector<T>& data) const;
    /**
     * Broadcast vector<> within trainer;
     * vector<> for non-root processes will be resized as needed.
     */
    /// Broadcast vector<> across trainers.
    template <typename T>
    void intertrainer_broadcast(int root, std::vector<T>& data) const;
    /// Broadcast vector<> within trainer.
    template <typename T>
    void trainer_broadcast(int root, std::vector<T>& data) const;
  
    /** Allgather over an arbitrary communicator */
    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;
  
    /**
     * Allgatherv over an arbitrary communicator;
     * all vectors must be correctly sized prior to entry.
     */
    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;
    /**
     * Allgatherv over a trainer communicator;
     * all vectors must be correctly sized prior to entry.
     */
    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;
    /**
     * Allgather for a single element over an arbitrary communicator;
     * std::vector<T> &data must be correctly sized prior to entry.
     */
    template <typename T>
    void
    all_gather(T const& src, std::vector<T>& data, const El::mpi::Comm& c) const;
    /**
     * Allgather for a single element over the world communicator;
     * std::vector<T> &data must be correctly sized prior to entry.
     */
    template <typename T>
    void world_all_gather(T const& src, std::vector<T>& data) const;
    /**
     * Allgather for a single element over the trainer communicator;
     * std::vector<T> &data must be correctly sized prior to entry.
     */
    template <typename T>
    void trainer_all_gather(T const& src, std::vector<T>& data) const;
  
    /** Within-trainer scalar gather (for non-root processes). */
    template <typename T> void trainer_gather(T snd, int root) const;
    /** Within-trainer scalar gather (for root processes). */
    template <typename T> void trainer_gather(T snd, T* rcv) const;
    /** Within-trainer scalar-array gather (for non-root processes). */
    template <typename T>
    void trainer_gather(T const* snd, int count, int root) const;
    /** Within-trainer scalar-array gather (for root processes). */
    template <typename T>
    void trainer_gather(T const* snd, int count, T* rcv) const;
    /** Within-trainer variable-length-array gather (for non-root processes). */
    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;
    /** Inter-trainer gather (for non-root processes). */
    template <typename T> void intertrainer_gather(T snd, int root) const;
    /** Inter-trainer gather (for root processes). */
    template <typename T>
    void intertrainer_gather(T snd, std::vector<T>& rcv) const;
    /** Inter-trainer scalar-array gather (for non-root processes). */
    template <typename T>
    void intertrainer_gather(T const* snd, int count, int root) const;
    /** Inter-trainer scalar-array gather (for root processes). */
    template <typename T>
    void intertrainer_gather(T const* snd, int count, T* rcv) const;
    /** Scalar gather (for non-root processes). */
    template <typename T>
    void gather(T snd, int root, const El::mpi::Comm& c) const;
    /** Scalar gather (for root processes). */
    template <typename T>
    void gather(T snd, T* rcv, const El::mpi::Comm& c) const;
    /** Scalar gather (for root processes). */
    template <typename T>
    void gather(T snd, std::vector<T>& rcv, const El::mpi::Comm& c) const;
    /** Scalar-array gather (for non-root processes). */
    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;
    /** Scalar-array gather (for root processes). */
    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;
    /** Scalar scatter (for non-root processes). */
    template <typename T> T scatter(int root, const El::mpi::Comm& c) const;
    /** Scalar scatter (for root processes). */
    template <typename T> T scatter(T const* snd, const El::mpi::Comm& c) const;
    /** Inter-trainer reduce (for non-root processes). */
    template <typename T>
    void
    intertrainer_reduce(T snd, int root, El::mpi::Op op = El::mpi::SUM) const;
    /** Inter-trainer reduce (for root processes). */
    template <typename T>
    T intertrainer_reduce(T snd, El::mpi::Op op = El::mpi::SUM) const;
    /** Within-trainer reduce (for non-root processes). */
    template <typename T>
    void trainer_reduce(T snd, int root, El::mpi::Op op = El::mpi::SUM) const;
    /** Within-trainer reduce (for root processes). */
    template <typename T>
    T trainer_reduce(T snd, El::mpi::Op op = El::mpi::SUM) const;
    /** Within-trainer scalar array reduce (for non-root processes). */
    template <typename T>
    void trainer_reduce(T const* snd,
                        int count,
                        int root,
                        El::mpi::Op op = El::mpi::SUM) const;
    /** Within-trainer scalar array reduce (for root processes). */
    template <typename T>
    void trainer_reduce(T const* snd,
                        int count,
                        T* rcv,
                        El::mpi::Op op = El::mpi::SUM) const;
    /** Scalar reduce (for non-root processes). */
    template <typename T>
    void reduce(T snd,
                int root,
                const El::mpi::Comm& c,
                El::mpi::Op op = El::mpi::SUM) const;
    /** Scalar reduce (for root processes). */
    template <typename T>
    T reduce(T snd, const El::mpi::Comm& c, El::mpi::Op op = El::mpi::SUM) const;
  
    /** Scalar-array reduce (for non-root processes). */
    // 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;
    /** Scalar-array reduce (for root processes). */
    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;
    /** Inter-trainer all-reduce. */
    template <typename T>
    T intertrainer_allreduce(T snd, El::mpi::Op op = El::mpi::SUM) const;
    /** Within-trainer all-reduce. */
    template <typename T>
    T trainer_allreduce(T snd, El::mpi::Op op = El::mpi::SUM) const;
    /** Scalar array within-trainer all-reduce. */
    template <typename T>
    void trainer_allreduce(T const* snd,
                           int count,
                           T* rcv,
                           El::mpi::Op op = El::mpi::SUM) const;
    /** Scalar allreduce. */
    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.
    /** Scalar-array allreduce. */
    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;
    /** In-place scalar-array allreduce. */
    template <typename T>
    void allreduce(T* data,
                   int count,
                   const El::mpi::Comm& c,
                   El::mpi::Op op = El::mpi::SUM) const;
    /** Matrix allreduce. */
    template <typename TensorDataType>
    void allreduce(El::AbstractMatrix<TensorDataType>& m,
                   const El::mpi::Comm& c,
                   El::mpi::Op op = El::mpi::SUM) const;
    /** Matrix allreduce. */
    template <typename TensorDataType>
    void allreduce(El::AbstractDistMatrix<TensorDataType>& m,
                   const El::mpi::Comm& c,
                   El::mpi::Op op = El::mpi::SUM) const;
    /** Non-blocking matrix allreduce.
     *  If LBANN has not been built with Aluminum, then this calls a
     *  blocking matrix allreduce.
     */
    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;
    /** Non-blocking matrix allreduce.
     *  If LBANN has not been built with Aluminum, then this calls a
     *  blocking matrix allreduce.
     */
    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;
    /** Non-blocking in-place scalar-array allreduce.
     *  If LBANN has not been built with Aluminum, then this calls a blocking
     *  allreduce.
     *  This currently only supports host pointers (i.e. the MPI backend).
     */
    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;
  
    /** Wait for a all non-blocking requests to complete. */
    template <typename T>
    void wait_all(std::vector<El::mpi::Request<T>>& req) const;
  
    /** Wait for a non-blocking request to complete. */
    template <typename T> void wait(El::mpi::Request<T>& req) const;
  
    /** Wait for a non-blocking request to complete. */
    void wait(Al::request& req) const;
    /** Test whether a non-blocking request has completed; true if it has. */
    bool test(Al::request& req) const;
  
    /** Barrier among the inter-trainer processes. */
    void intertrainer_barrier() const;
    /** Barrier among processes in this trainer. */
    void trainer_barrier() const;
    /** Barrier among all processes. */
    void global_barrier() const;
    /** Barrier on an arbitrary communicator. */
    void barrier(const El::mpi::Comm& c) const;
  
    /** Send a buffer to rank in trainer. */
    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);
    }
  
    /** Corresponding non-blocking sends. */
    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);
    }
  
    /** Corresponding receive to send. */
    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);
    }
    /** As above, but receive from anyone. */
    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;
  
    /** Corresponding non-blocking receives. */
    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;
  
    /** Send/recv to/from ranks. */
    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;
  
    /** Determine the size (count) of an incoming message. */
    template <typename T> int get_count(int trainer, int rank) const;
    template <typename T> int get_count(int trainer) const;
  
    // Statistics methods.
    /** Return the number of trainer barriers performed. */
    inline size_t get_num_trainer_barriers() const noexcept
    {
      return m_num_trainer_barriers;
    }
    /** Return the number of inter-trainer barriers performed. */
    inline size_t get_num_intertrainer_barriers() const noexcept
    {
      return m_num_intertrainer_barriers;
    }
    /** Return the number of global barriers performed. */
    inline size_t get_num_global_barriers() const noexcept
    {
      return m_num_global_barriers;
    }
    /** Return the number of bytes sent. */
    inline size_t get_bytes_sent() const noexcept { return m_bytes_sent; }
    /** Return the number of bytes received. */
    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;
    }
  
    /** Return true if mat can be transmitted. */
    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();
    }
    /** Return true if the local portion of dist_mat can be transmitted. */
    static inline bool is_sendable(const AbsDistMat& dist_mat) noexcept
    {
      return is_sendable(dist_mat.LockedMatrix());
    }
  
    /** Developer's note: to get the raw MPI communicator, which may be needed
     *  when working with external libraries, by example:
     *     comm->get_intertrainer_comm().GetMPIComm()
     */
  
    /** Return the intertrainer communicator. */
    const El::mpi::Comm& get_intertrainer_comm() const noexcept
    {
      return m_intertrainer_comm;
    }
  
    /** Return the trainer communicator. */
    const El::mpi::Comm& get_trainer_comm() const noexcept
    {
      return m_trainer_comm;
    }
  
    /** Return the combined grid communicator for a trainer. */
    const El::mpi::Comm& get_combined_grid_comm() const noexcept
    {
      return m_combined_grid_comm;
    }
  
    /** Return the world communicator. */
    const El::mpi::Comm& get_world_comm() const noexcept { return m_world_comm; }
  
    /** Return the communicator for this node. */
    const El::mpi::Comm& get_node_comm() const noexcept { return m_node_comm; }
  
    /** Return the communicator for this grid in sub-grid parallelism. */
    const El::mpi::Comm& get_KFAC_comm() const noexcept { return m_trainer_comm;  }
  
    /** Return the ranks of primary grid in the trainer */
    std::vector<int> get_primary_grid_ranks(){ return m_primary_grid_ranks; }
  
    /** Return the ranks of secondary grid in the trainer */
    std::vector<int> get_secondary_grid_ranks(){ return m_secondary_grid_ranks; }
  
    bool get_KFAC_subgrid_create_two_models(){ return m_create_two_models; }
  
    /**
     * Return a communicator containing num_per_group processors.
     *
     * This will attempt to pack processes so that the processes in each group
     * are physically close together on the system.
     *
     * num_per_group must evenly divide the number of processors in the world.
     */
    const El::mpi::Comm& get_packed_group_comm(int num_per_group) const;
  
    /** Return true if rank (in comm) is on the local node. */
    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);
    }
  
    /** throws an lbann_exception **/
    void lbann_comm_abort(std::string msg) const;
  
  private:
    /** World communicator. */
    const El::mpi::Comm m_world_comm;
    /** Communicator for every process in this trainer. */
    El::mpi::Comm m_trainer_comm;
    /** Communicator for every process with the same trainer rank. */
    El::mpi::Comm m_intertrainer_comm;
    /** Communicator for every process in the same compute node. */
    El::mpi::Comm m_node_comm;
    /** Communicator for primary grid in each trainer */
    El::mpi::Comm m_primary_grid_comm;
    /** Communicator for secondary grid in each trainer */
    El::mpi::Comm m_secondary_grid_comm;
    /** Combined communicator for primary and secondary grid in each trainer */
    El::mpi::Comm m_combined_grid_comm;
    /** Packed group communicators. */
    mutable std::unordered_map<int, El::mpi::Comm> m_group_communicators;
    /** Grid for this trainer. */
    std::unique_ptr<El::Grid> m_grid;
    /** Number of trainers. */
    int m_num_trainers;
    /** Number of processors per trainer. */
    int m_procs_per_trainer;
    /** Rank of the trainer this process is in. */
    int m_trainer_rank;
    /** Rank of this process within its trainer. */
    int m_rank_in_trainer;
    /** Number of processers per compute node. */
    int m_procs_per_node;
    /** Rank of this process within its compute node. */
    int m_rank_in_node;
    /** The list of world ranks that are on this compute node. */
    std::vector<int> m_world_ranks_on_node;
    /** Default number of threads per process.
     *  This is the number of OpenMP threads to use for parallel
     *  regions, provided omp_set_num_threads has not been called or the
     *  num_threads directive has not been provided.
     */
    int m_threads_per_proc;
  
    /** Grid typer for current process when sub-grid parallelism is enabled */
    GridType m_grid_type = GridType::NO_GRID;
  
    bool m_create_two_models=false;
  
    std::unique_ptr<El::Grid> m_secondary_grid;
  
    /**
    Ranks in primary and secondary grids
    */
    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;
  
    /** Setup communicator for processes in the same compute node. */
    void setup_node_comm();
  
    /** Initialize the default number of threads per process.
     *  This is the number of OpenMP threads to use for parallel
     *  regions, provided omp_set_num_threads has not been called or the
     *  num_threads directive has not been provided. If the environment
     *  variable OMP_NUM_THREADS is defined, it's value is used for the
     *  default. Otherwise, then the default is the number of hardware
     *  cores per node divided by the number of processes per node.
     */
    void setup_threads();
  
    /**
     * Keep track of the number of broadcast bytes transmitted and received
     */
    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
  
  /** Get the current rank within MPI_COMM_WORLD.
   *  This function is safe to call even if MPI has not initialized or
   *  has been finalized. In either case it returns a negative value.
   */
  int get_rank_in_world();
  
  } // namespace lbann
  
  #endif // LBANN_COMM_HPP_INCLUDED
  