grid.hh

// This file is part of fml which is released under the Boost Software
// License, Version 1.0. See accompanying file LICENSE or copy at
// https://www.boost.org/LICENSE_1_0.txt
 
#ifndef FML_MPI_GRID_H
#define FML_MPI_GRID_H
#pragma once
 
 
#define OMPI_SKIP_MPICXX 1
#include <mpi.h>
 
#include <cmath>
#include <stdexcept>
 
#include "../_internals/print.hh"
 
#include "internals/blacs_prototypes.h"
 
 
namespace fml
{
  enum gridshape
  {
    PROC_GRID_SQUARE,
    PROC_GRID_WIDE,
    PROC_GRID_TALL
  };
  
  enum blacsops
  {
    BLACS_SUM,
    BLACS_MAX,
    BLACS_MIN
  };
}
 
 
 
namespace fml
{
  class grid
  {
    public:
      // constructors/destructor and comm management
      grid();
      grid(const gridshape gridtype);
      void set(const int blacs_context);
      void exit();
      void finalize(const bool mpi_continue=false);
      
      // utilities
      void printf(const int row, const int col, const char *fmt, ...) const;
      void info() const;
      bool rank0() const;
      bool ingrid() const;
      
      // send/recv
      void send(const int m, const int n, const int *x,    const int rdest=0, const int cdest=0) const;
      void send(const int m, const int n, const float *x,  const int rdest=0, const int cdest=0) const;
      void send(const int m, const int n, const double *x, const int rdest=0, const int cdest=0) const;
      void send(const int m, const int n, const int ldx, const int *x,    const int rdest=0, const int cdest=0) const;
      void send(const int m, const int n, const int ldx, const float *x,  const int rdest=0, const int cdest=0) const;
      void send(const int m, const int n, const int ldx, const double *x, const int rdest=0, const int cdest=0) const;
      
      void recv(const int m, const int n, int *x,    const int rsrc=0, const int csrc=0) const;
      void recv(const int m, const int n, float *x,  const int rsrc=0, const int csrc=0) const;
      void recv(const int m, const int n, double *x, const int rsrc=0, const int csrc=0) const;
      void recv(const int m, const int n, const int ldx, int *x,    const int rsrc=0, const int csrc=0) const;
      void recv(const int m, const int n, const int ldx, float *x,  const int rsrc=0, const int csrc=0) const;
      void recv(const int m, const int n, const int ldx, double *x, const int rsrc=0, const int csrc=0) const;
      
      // collectives
      void barrier(const char scope='A') const;
      
      void allreduce(const int m, const int n, int *x,    const char scope='A', const blacsops op=BLACS_SUM) const;
      void allreduce(const int m, const int n, float *x,  const char scope='A', const blacsops op=BLACS_SUM) const;
      void allreduce(const int m, const int n, double *x, const char scope='A', const blacsops op=BLACS_SUM) const;
      
      void reduce(const int m, const int n, int *x,    const char scope='A', const blacsops op=BLACS_SUM, const int rdest=0, const int cdest=0) const;
      void reduce(const int m, const int n, float *x,  const char scope='A', const blacsops op=BLACS_SUM, const int rdest=0, const int cdest=0) const;
      void reduce(const int m, const int n, double *x, const char scope='A', const blacsops op=BLACS_SUM, const int rdest=0, const int cdest=0) const;
      
      void bcast(const int m, const int n, int *x,    const char scope='A', const int rsrc=0, const int csrc=0) const;
      void bcast(const int m, const int n, float *x,  const char scope='A', const int rsrc=0, const int csrc=0) const;
      void bcast(const int m, const int n, double *x, const char scope='A', const int rsrc=0, const int csrc=0) const;
      
      
      int ictxt() const {return _ictxt;};
      int nprocs() const {return _nprocs;};
      int nprow() const {return _nprow;};
      int npcol() const {return _npcol;};
      int myrow() const {return _myrow;};
      int mycol() const {return _mycol;};
      
      bool valid_grid() const {return (_ictxt!=UNINITIALIZED_GRID && _ictxt!=EXITED_GRID);};
    
    protected:
      int _ictxt;
      int _nprocs;
      int _nprow;
      int _npcol;
      int _myrow;
      int _mycol;
    
    private:
      static const int UNINITIALIZED_GRID = -1;
      static const int EXITED_GRID = -11;
      
      void squarish(int *nr, int *nc) const;
  };
}
 
 
 
// -----------------------------------------------------------------------------
// public
// -----------------------------------------------------------------------------
 
// constructors/destructor and grid management
 
inline fml::grid::grid()
{
  _ictxt = UNINITIALIZED_GRID;
  _nprocs = _nprow = _npcol = _myrow = _mycol = -1;
}
 
 
 
inline fml::grid::grid(const fml::gridshape gridtype)
{
  char order = 'R';
  
  int mypnum;
  Cblacs_pinfo(&mypnum, &_nprocs);
  
  Cblacs_get(-1, 0, &_ictxt);
  
  if (gridtype == PROC_GRID_SQUARE)
  {
    int nr, nc;
    squarish(&nr, &nc);
    
    Cblacs_gridinit(&_ictxt, &order, nr, nc);
  }
  else if (gridtype == PROC_GRID_TALL)
    Cblacs_gridinit(&_ictxt, &order, _nprocs, 1);
  else if (gridtype == PROC_GRID_WIDE)
    Cblacs_gridinit(&_ictxt, &order, 1, _nprocs);
  else
    throw std::runtime_error("Process grid should be one of PROC_GRID_SQUARE, PROC_GRID_TALL, or PROC_GRID_WIDE");
  
  Cblacs_gridinfo(_ictxt, &_nprow, &_npcol, &_myrow, &_mycol);
}
 
 
 
inline void fml::grid::set(const int blacs_context)
{
  _ictxt = blacs_context;
  Cblacs_gridinfo(_ictxt, &_nprow, &_npcol, &_myrow, &_mycol);
  
  if (_nprow == -1)
    throw std::runtime_error("context handle does not point at a valid context");
  
  _nprocs = _nprow * _npcol;
}
 
 
 
inline void fml::grid::exit()
{
  if (_ictxt != EXITED_GRID && _ictxt != UNINITIALIZED_GRID)
    Cblacs_gridexit(_ictxt);
  
  _ictxt = EXITED_GRID;
  _nprocs = _nprow = _npcol = _myrow = _mycol = -1;
}
 
 
 
inline void fml::grid::finalize(const bool mpi_continue)
{
  exit();
  
  int cont = (int) mpi_continue;
  Cblacs_exit(cont);
}
 
 
 
// utilities
 
inline void fml::grid::printf(const int row, const int col, const char *fmt, ...) const
{
  if (_myrow == row && _mycol == col)
  {
    va_list args;
    
    va_start(args, fmt);
    fml::print::vprintf(fmt, args);
    va_end(args);
  }
}
 
 
 
inline void fml::grid::info() const
{
  printf(0, 0, "## Grid %d %dx%d\n\n", _ictxt, _nprow, _npcol);
}
 
 
 
inline bool fml::grid::rank0() const
{
  return (_myrow==0 && _mycol==0);
}
 
 
 
inline bool fml::grid::ingrid() const
{
  return !(_myrow==-1 && _mycol==-1);
}
 
 
 
// send/recv
 
inline void fml::grid::send(const int m, const int n, const int *x, const int rdest, const int cdest) const
{
  Cigesd2d(_ictxt, m, n, x, m, rdest, cdest);
}
 
inline void fml::grid::send(const int m, const int n, const float *x, const int rdest, const int cdest) const
{
  Csgesd2d(_ictxt, m, n, x, m, rdest, cdest);
}
 
inline void fml::grid::send(const int m, const int n, const double *x, const int rdest, const int cdest) const
{
  Cdgesd2d(_ictxt, m, n, x, m, rdest, cdest);
}
 
inline void fml::grid::send(const int m, const int n, const int ldx, const int *x, const int rdest, const int cdest) const
{
  Cigesd2d(_ictxt, m, n, x, ldx, rdest, cdest);
}
 
inline void fml::grid::send(const int m, const int n, const int ldx, const float *x, const int rdest, const int cdest) const
{
  Csgesd2d(_ictxt, m, n, x, ldx, rdest, cdest);
}
 
inline void fml::grid::send(const int m, const int n, const int ldx, const double *x, const int rdest, const int cdest) const
{
  Cdgesd2d(_ictxt, m, n, x, ldx, rdest, cdest);
}
 
 
 
inline void fml::grid::recv(const int m, const int n, int *x, const int rsrc, const int csrc) const
{
  Cigerv2d(_ictxt, m, n, x, m, rsrc, csrc);
}
 
inline void fml::grid::recv(const int m, const int n, float *x, const int rsrc, const int csrc) const
{
  Csgerv2d(_ictxt, m, n, x, m, rsrc, csrc);
}
 
inline void fml::grid::recv(const int m, const int n, double *x, const int rsrc, const int csrc) const
{
  Cdgerv2d(_ictxt, m, n, x, m, rsrc, csrc);
}
 
inline void fml::grid::recv(const int m, const int n, const int ldx, int *x, const int rsrc, const int csrc) const
{
  Cigerv2d(_ictxt, m, n, x, ldx, rsrc, csrc);
}
 
inline void fml::grid::recv(const int m, const int n, const int ldx, float *x, const int rsrc, const int csrc) const
{
  Csgerv2d(_ictxt, m, n, x, ldx, rsrc, csrc);
}
 
inline void fml::grid::recv(const int m, const int n, const int ldx, double *x, const int rsrc, const int csrc) const
{
  Cdgerv2d(_ictxt, m, n, x, ldx, rsrc, csrc);
}
 
 
 
// collectives
 
inline void fml::grid::barrier(const char scope) const
{
  Cblacs_barrier(_ictxt, &scope);
}
 
 
 
inline void fml::grid::allreduce(const int m, const int n, int *x, const char scope, const blacsops op) const
{
  reduce(m, n, x, scope, op, -1, -1);
}
 
inline void fml::grid::allreduce(const int m, const int n, float *x, const char scope, const blacsops op) const
{
  reduce(m, n, x, scope, op, -1, -1);
}
 
inline void fml::grid::allreduce(const int m, const int n, double *x, const char scope, const blacsops op) const
{
  reduce(m, n, x, scope, op, -1, -1);
}
 
 
 
inline void fml::grid::reduce(const int m, const int n, int *x, const char scope, const blacsops op, const int rdest, const int cdest) const
{
  char top = ' ';
  
  if (op == BLACS_SUM)
    Cigsum2d(_ictxt, &scope, &top, m, n, x, m, rdest, cdest);
  else if (op == BLACS_MAX)
    Cigamx2d(_ictxt, &scope, &top, m, n, x, m, NULL, NULL, -1, rdest, cdest);
  else if (op == BLACS_MIN)
    Cigamn2d(_ictxt, &scope, &top, m, n, x, m, NULL, NULL, -1, rdest, cdest);
}
 
inline void fml::grid::reduce(const int m, const int n, float *x, const char scope, const blacsops op, const int rdest, const int cdest) const
{
  char top = ' ';
  
  if (op == BLACS_SUM)
    Csgsum2d(_ictxt, &scope, &top, m, n, x, m, rdest, cdest);
  else if (op == BLACS_MAX)
    Csgamx2d(_ictxt, &scope, &top, m, n, x, m, NULL, NULL, -1, rdest, cdest);
  else if (op == BLACS_MIN)
    Csgamn2d(_ictxt, &scope, &top, m, n, x, m, NULL, NULL, -1, rdest, cdest);
}
 
inline void fml::grid::reduce(const int m, const int n, double *x, const char scope, const blacsops op, const int rdest, const int cdest) const
{
  char top = ' ';
  
  if (op == BLACS_SUM)
    Cdgsum2d(_ictxt, &scope, &top, m, n, x, m, rdest, cdest);
  else if (op == BLACS_MAX)
    Cdgamx2d(_ictxt, &scope, &top, m, n, x, m, NULL, NULL, -1, rdest, cdest);
  else if (op == BLACS_MIN)
    Cdgamn2d(_ictxt, &scope, &top, m, n, x, m, NULL, NULL, -1, rdest, cdest);
}
 
 
 
inline void fml::grid::bcast(const int m, const int n, int *x, const char scope, const int rsrc, const int csrc) const
{
  char top = ' ';
  if (rsrc == _myrow && csrc == _mycol)
    Cigebs2d(_ictxt, &scope, &top, m, n, x, m);
  else
    Cigebr2d(_ictxt, &scope, &top, m, n, x, m, rsrc, csrc);
}
 
inline void fml::grid::bcast(const int m, const int n, float *x, const char scope, const int rsrc, const int csrc) const
{
  char top = ' ';
  if (rsrc == _myrow && csrc == _mycol)
    Csgebs2d(_ictxt, &scope, &top, m, n, x, m);
  else
    Csgebr2d(_ictxt, &scope, &top, m, n, x, m, rsrc, csrc);
}
 
inline void fml::grid::bcast(const int m, const int n, double *x, const char scope, const int rsrc, const int csrc) const
{
  char top = ' ';
  if (rsrc == _myrow && csrc == _mycol)
    Cdgebs2d(_ictxt, &scope, &top, m, n, x, m);
  else
    Cdgebr2d(_ictxt, &scope, &top, m, n, x, m, rsrc, csrc);
}
 
 
 
// -----------------------------------------------------------------------------
// private
// -----------------------------------------------------------------------------
 
inline void fml::grid::squarish(int *nr, int *nc) const
{
  int n = (int) sqrt((double) _nprocs);
  n = (n<1)?1:n; // suppresses bogus compiler warning
  
  for (int i=0; i<n; i++)
  {
    (*nc) = n - i;
    (*nr) = _nprocs % (*nc);
    if ((*nr) == 0)
      break;
  }
  
  (*nr) = _nprocs / (*nc);
}
 
 
#endif