src/smpi/colls/allreduce-rab-rdb.c

   1 #include "colls.h"
   2
   3 int smpi_coll_tuned_allreduce_rab_rdb(void *sbuff, void *rbuff, int count,
   4                                       MPI_Datatype dtype, MPI_Op op,
   5                                       MPI_Comm comm)
   6 {
   7   int nprocs, rank, type_size, tag = 543;
   8   int mask, dst, pof2, newrank, rem, newdst, i,
   9       send_idx, recv_idx, last_idx, send_cnt, recv_cnt, *cnts, *disps;
  10   MPI_Aint lb, extent;
  11   MPI_Status status;
  12   void *tmp_buf = NULL;
  13
  14 #ifdef MPICH2_REDUCTION
  15   MPI_User_function *uop = MPIR_Op_table[op % 16 - 1];
  16 #else
  17   MPI_User_function *uop;
  18   struct MPIR_OP *op_ptr;
  19   op_ptr = (MPI_User_function *) MPIR_ToPointer(op);
  20   uop = op_ptr->op;
  21 #endif
  22
  23   MPI_Comm_size(comm, &nprocs);
  24   MPI_Comm_rank(comm, &rank);
  25
  26   MPI_Type_extent(dtype, &extent);
  27   tmp_buf = (void *) malloc(count * extent);
  28   if (!tmp_buf) {
  29     printf("Could not allocate memory for tmp_buf\n");
  30     return 1;
  31   }
  32
  33   MPIR_Localcopy(sbuff, count, dtype, rbuff, count, dtype);
  34
  35   MPI_Type_size(dtype, &type_size);
  36
  37   // find nearest power-of-two less than or equal to comm_size
  38   pof2 = 1;
  39   while (pof2 <= nprocs)
  40     pof2 <<= 1;
  41   pof2 >>= 1;
  42
  43   rem = nprocs - pof2;
  44
  45   // In the non-power-of-two case, all even-numbered
  46   // processes of rank < 2*rem send their data to
  47   // (rank+1). These even-numbered processes no longer
  48   // participate in the algorithm until the very end. The
  49   // remaining processes form a nice power-of-two.
  50
  51   if (rank < 2 * rem) {
  52     // even
  53     if (rank % 2 == 0) {
  54
  55       MPI_Send(rbuff, count, dtype, rank + 1, tag, comm);
  56
  57       // temporarily set the rank to -1 so that this
  58       // process does not pariticipate in recursive
  59       // doubling
  60       newrank = -1;
  61     } else                      // odd
  62     {
  63       MPI_Recv(tmp_buf, count, dtype, rank - 1, tag, comm, &status);
  64       // do the reduction on received data. since the
  65       // ordering is right, it doesn't matter whether
  66       // the operation is commutative or not.
  67       (*uop) (tmp_buf, rbuff, &count, &dtype);
  68
  69       // change the rank
  70       newrank = rank / 2;
  71     }
  72   }
  73
  74   else                          // rank >= 2 * rem
  75     newrank = rank - rem;
  76
  77   // If op is user-defined or count is less than pof2, use
  78   // recursive doubling algorithm. Otherwise do a reduce-scatter
  79   // followed by allgather. (If op is user-defined,
  80   // derived datatypes are allowed and the user could pass basic
  81   // datatypes on one process and derived on another as long as
  82   // the type maps are the same. Breaking up derived
  83   // datatypes to do the reduce-scatter is tricky, therefore
  84   // using recursive doubling in that case.)
  85
  86   if (newrank != -1) {
  87     // do a reduce-scatter followed by allgather. for the
  88     // reduce-scatter, calculate the count that each process receives
  89     // and the displacement within the buffer
  90
  91     cnts = (int *) malloc(pof2 * sizeof(int));
  92     disps = (int *) malloc(pof2 * sizeof(int));
  93
  94     for (i = 0; i < (pof2 - 1); i++)
  95       cnts[i] = count / pof2;
  96     cnts[pof2 - 1] = count - (count / pof2) * (pof2 - 1);
  97
  98     disps[0] = 0;
  99     for (i = 1; i < pof2; i++)
 100       disps[i] = disps[i - 1] + cnts[i - 1];
 101
 102     mask = 0x1;
 103     send_idx = recv_idx = 0;
 104     last_idx = pof2;
 105     while (mask < pof2) {
 106       newdst = newrank ^ mask;
 107       // find real rank of dest
 108       dst = (newdst < rem) ? newdst * 2 + 1 : newdst + rem;
 109
 110       send_cnt = recv_cnt = 0;
 111       if (newrank < newdst) {
 112         send_idx = recv_idx + pof2 / (mask * 2);
 113         for (i = send_idx; i < last_idx; i++)
 114           send_cnt += cnts[i];
 115         for (i = recv_idx; i < send_idx; i++)
 116           recv_cnt += cnts[i];
 117       } else {
 118         recv_idx = send_idx + pof2 / (mask * 2);
 119         for (i = send_idx; i < recv_idx; i++)
 120           send_cnt += cnts[i];
 121         for (i = recv_idx; i < last_idx; i++)
 122           recv_cnt += cnts[i];
 123       }
 124
 125       // Send data from recvbuf. Recv into tmp_buf
 126       MPI_Sendrecv((char *) rbuff + disps[send_idx] * extent, send_cnt,
 127                    dtype, dst, tag,
 128                    (char *) tmp_buf + disps[recv_idx] * extent, recv_cnt,
 129                    dtype, dst, tag, comm, &status);
 130
 131       // tmp_buf contains data received in this step.
 132       // recvbuf contains data accumulated so far
 133
 134       // This algorithm is used only for predefined ops
 135       // and predefined ops are always commutative.
 136       (*uop) ((char *) tmp_buf + disps[recv_idx] * extent,
 137               (char *) rbuff + disps[recv_idx] * extent, &recv_cnt, &dtype);
 138
 139       // update send_idx for next iteration
 140       send_idx = recv_idx;
 141       mask <<= 1;
 142
 143       // update last_idx, but not in last iteration because the value
 144       // is needed in the allgather step below.
 145       if (mask < pof2)
 146         last_idx = recv_idx + pof2 / mask;
 147     }
 148
 149     // now do the allgather
 150
 151     mask >>= 1;
 152     while (mask > 0) {
 153       newdst = newrank ^ mask;
 154       // find real rank of dest
 155       dst = (newdst < rem) ? newdst * 2 + 1 : newdst + rem;
 156
 157       send_cnt = recv_cnt = 0;
 158       if (newrank < newdst) {
 159         // update last_idx except on first iteration
 160         if (mask != pof2 / 2)
 161           last_idx = last_idx + pof2 / (mask * 2);
 162
 163         recv_idx = send_idx + pof2 / (mask * 2);
 164         for (i = send_idx; i < recv_idx; i++)
 165           send_cnt += cnts[i];
 166         for (i = recv_idx; i < last_idx; i++)
 167           recv_cnt += cnts[i];
 168       } else {
 169         recv_idx = send_idx - pof2 / (mask * 2);
 170         for (i = send_idx; i < last_idx; i++)
 171           send_cnt += cnts[i];
 172         for (i = recv_idx; i < send_idx; i++)
 173           recv_cnt += cnts[i];
 174       }
 175
 176       MPI_Sendrecv((char *) rbuff + disps[send_idx] * extent, send_cnt,
 177                    dtype, dst, tag,
 178                    (char *) rbuff + disps[recv_idx] * extent, recv_cnt,
 179                    dtype, dst, tag, comm, &status);
 180
 181       if (newrank > newdst)
 182         send_idx = recv_idx;
 183
 184       mask >>= 1;
 185     }
 186
 187     free(cnts);
 188     free(disps);
 189
 190   }
 191   // In the non-power-of-two case, all odd-numbered processes of
 192   // rank < 2 * rem send the result to (rank-1), the ranks who didn't
 193   // participate above.
 194
 195   if (rank < 2 * rem) {
 196     if (rank % 2)               // odd
 197       MPI_Send(rbuff, count, dtype, rank - 1, tag, comm);
 198     else                        // even
 199       MPI_Recv(rbuff, count, dtype, rank + 1, tag, comm, &status);
 200   }
 201
 202   free(tmp_buf);
 203   return 0;
 204 }