communicator.hpp 4.71 KB
Newer Older
Kitware Robot's avatar
Kitware Robot committed
1
2
3
4
5
6
7
8
9
10
namespace diy
{
namespace mpi
{

  //! \ingroup MPI
  //! Simple wrapper around `MPI_Comm`.
  class communicator
  {
    public:
Kitware Robot's avatar
Kitware Robot committed
11
                inline
Kitware Robot's avatar
Kitware Robot committed
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
                communicator(MPI_Comm comm = MPI_COMM_WORLD, bool owner = false);

                inline
                ~communicator();

                communicator(const communicator& other):
                    comm_(other.comm_),
                    rank_(other.rank_),
                    size_(other.size_),
                    owner_(false)                   {}

                communicator(communicator&& other):
                    comm_(other.comm_),
                    rank_(other.rank_),
                    size_(other.size_),
                    owner_(other.owner_)                    { other.owner_ = false; }

    communicator&
                operator=(const communicator& other)        { comm_ = other.comm_; rank_ = other.rank_; size_ = other.size_; owner_ = false; return *this; }
    communicator&
                operator=(communicator&& other)             { comm_ = other.comm_; rank_ = other.rank_; size_ = other.size_; owner_ = other.owner_; other.owner_ = false; return *this; }
Kitware Robot's avatar
Kitware Robot committed
33
34
35
36
37
38
39
40
41
42

      int       rank() const                        { return rank_; }
      int       size() const                        { return size_; }

      //! Send `x` to processor `dest` using `tag` (blocking).
      template<class T>
      void      send(int dest, int tag, const T& x) const   { detail::send<T>()(comm_, dest, tag, x); }

      //! Receive `x` from `dest` using `tag` (blocking).
      //! If `T` is an `std::vector<...>`, `recv` will resize it to fit exactly the sent number of values.
Kitware Robot's avatar
Kitware Robot committed
43
44
45
46
47
48
49
50
51
52
      template <class T>
      status recv(int source, int tag, T &x) const
      {
#if defined(DIY_NO_MPI) && defined(__CUDACC_VER_MAJOR__) && __CUDACC_VER_MAJOR__ < 8 // CUDA 7.5 workaround
        (void) source; (void)tag; (void)x;
        DIY_UNSUPPORTED_MPI_CALL(MPI_Recv);
#else
        return detail::recv<T>{}(comm_, source, tag, x);
#endif
      }
Kitware Robot's avatar
Kitware Robot committed
53
54

      //! Non-blocking version of `send()`.
Kitware Robot's avatar
Kitware Robot committed
55
56
57
58
59
60
61
62
63
64
      template <class T>
      request isend(int dest, int tag, const T &x) const
      {
#if defined(DIY_NO_MPI) && defined(__CUDACC_VER_MAJOR__) && __CUDACC_VER_MAJOR__ < 8 // CUDA 7.5 workaround
        (void) dest; (void)tag; (void)x;
        DIY_UNSUPPORTED_MPI_CALL(MPI_Send);
#else
        return detail::isend<T>{}(comm_, dest, tag, x);
#endif
      }
Kitware Robot's avatar
Kitware Robot committed
65
66

      //! Non-blocking version of `recv()`.
Kitware Robot's avatar
Kitware Robot committed
67
      //! If `T` is an `std::vector<...>`, its size must be big enough to accommodate the sent values.
Kitware Robot's avatar
Kitware Robot committed
68
69
70
71
72
73
74
75
76
77
      template <class T>
      request irecv(int source, int tag, T &x) const
      {
#if defined(DIY_NO_MPI) && defined(__CUDACC_VER_MAJOR__) && __CUDACC_VER_MAJOR__ < 8 // CUDA 7.5 workaround
        (void)source; (void)tag; (void)x;
        DIY_UNSUPPORTED_MPI_CALL(MPI_Irecv);
#else
        return detail::irecv<T>()(comm_, source, tag, x);
#endif
      }
Kitware Robot's avatar
Kitware Robot committed
78
79

      //! probe
Kitware Robot's avatar
Kitware Robot committed
80
81
      inline
      status    probe(int source, int tag) const;
Kitware Robot's avatar
Kitware Robot committed
82
83
84
85
86
87
88

      //! iprobe
      inline
      optional<status>
                iprobe(int source, int tag) const;

      //! barrier
Kitware Robot's avatar
Kitware Robot committed
89
90
      inline
      void      barrier() const;
Kitware Robot's avatar
Kitware Robot committed
91
92
93

                operator MPI_Comm() const                   { return comm_; }

Kitware Robot's avatar
Kitware Robot committed
94
95
96
97
98
99
      //! split
      //! When keys are the same, the ties are broken by the rank in the original comm.
      inline
      communicator
                split(int color, int key = 0) const;

Kitware Robot's avatar
Kitware Robot committed
100
101
102
103
    private:
      MPI_Comm  comm_;
      int       rank_;
      int       size_;
Kitware Robot's avatar
Kitware Robot committed
104
      bool      owner_;
Kitware Robot's avatar
Kitware Robot committed
105
106
107
108
  };
}
}

Kitware Robot's avatar
Kitware Robot committed
109
diy::mpi::communicator::
Kitware Robot's avatar
Kitware Robot committed
110
111
communicator(MPI_Comm comm, bool owner):
    comm_(comm), rank_(0), size_(1), owner_(owner)
Kitware Robot's avatar
Kitware Robot committed
112
113
114
115
116
117
118
119
120
121
{
#ifndef DIY_NO_MPI
  if (comm != MPI_COMM_NULL)
  {
    MPI_Comm_rank(comm_, &rank_);
    MPI_Comm_size(comm_, &size_);
  }
#endif
}

Kitware Robot's avatar
Kitware Robot committed
122
123
124
125
126
127
128
129
130
diy::mpi::communicator::
~communicator()
{
#ifndef DIY_NO_MPI
    if (owner_)
        MPI_Comm_free(&comm_);
#endif
}

Kitware Robot's avatar
Kitware Robot committed
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
diy::mpi::status
diy::mpi::communicator::
probe(int source, int tag) const
{
  (void) source;
  (void) tag;

#ifndef DIY_NO_MPI
  status s;
  MPI_Probe(source, tag, comm_, &s.s);
  return s;
#else
  DIY_UNSUPPORTED_MPI_CALL(MPI_Probe);
#endif
}

Kitware Robot's avatar
Kitware Robot committed
147
148
149
150
diy::mpi::optional<diy::mpi::status>
diy::mpi::communicator::
iprobe(int source, int tag) const
{
Kitware Robot's avatar
Kitware Robot committed
151
152
153
  (void) source;
  (void) tag;
#ifndef DIY_NO_MPI
Kitware Robot's avatar
Kitware Robot committed
154
155
156
157
158
  status s;
  int flag;
  MPI_Iprobe(source, tag, comm_, &flag, &s.s);
  if (flag)
    return s;
Kitware Robot's avatar
Kitware Robot committed
159
#endif
Kitware Robot's avatar
Kitware Robot committed
160
161
162
  return optional<status>();
}

Kitware Robot's avatar
Kitware Robot committed
163
164
165
166
167
168
169
170
void
diy::mpi::communicator::
barrier() const
{
#ifndef DIY_NO_MPI
  MPI_Barrier(comm_);
#endif
}
Kitware Robot's avatar
Kitware Robot committed
171
172
173
174
175
176
177
178
179
180
181
182
183

diy::mpi::communicator
diy::mpi::communicator::
split(int color, int key) const
{
#ifndef DIY_NO_MPI
    MPI_Comm newcomm;
    MPI_Comm_split(comm_, color, key, &newcomm);
    return communicator(newcomm, true);
#else
    return communicator();
#endif
}