Chris@10: /*
Chris@10:  * Copyright (c) 2003, 2007-11 Matteo Frigo
Chris@10:  * Copyright (c) 2003, 2007-11 Massachusetts Institute of Technology
Chris@10:  *
Chris@10:  * This program is free software; you can redistribute it and/or modify
Chris@10:  * it under the terms of the GNU General Public License as published by
Chris@10:  * the Free Software Foundation; either version 2 of the License, or
Chris@10:  * (at your option) any later version.
Chris@10:  *
Chris@10:  * This program is distributed in the hope that it will be useful,
Chris@10:  * but WITHOUT ANY WARRANTY; without even the implied warranty of
Chris@10:  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
Chris@10:  * GNU General Public License for more details.
Chris@10:  *
Chris@10:  * You should have received a copy of the GNU General Public License
Chris@10:  * along with this program; if not, write to the Free Software
Chris@10:  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
Chris@10:  *
Chris@10:  */
Chris@10: 
Chris@10: /* Recursive "radix-r" distributed transpose, which breaks a transpose
Chris@10:    over p processes into p/r transposes over r processes plus r
Chris@10:    transposes over p/r processes.  If performed recursively, this
Chris@10:    produces a total of O(p log p) messages vs. O(p^2) messages for a
Chris@10:    direct approach.
Chris@10: 
Chris@10:    However, this is not necessarily an improvement.  The total size of
Chris@10:    all the messages is actually increased from O(N) to O(N log p)
Chris@10:    where N is the total data size.  Also, the amount of local data
Chris@10:    rearrangement is increased.  So, it's not clear, a priori, what the
Chris@10:    best algorithm will be, and we'll leave it to the planner.  (In
Chris@10:    theory and practice, it looks like this becomes advantageous for
Chris@10:    large p, in the limit where the message sizes are small and
Chris@10:    latency-dominated.)
Chris@10: */
Chris@10: 
Chris@10: #include "mpi-transpose.h"
Chris@10: #include <string.h>
Chris@10: 
Chris@10: typedef struct {
Chris@10:      solver super;
Chris@10:      int (*radix)(int np);
Chris@10:      const char *nam;
Chris@10:      int preserve_input; /* preserve input even if DESTROY_INPUT was passed */
Chris@10: } S;
Chris@10: 
Chris@10: typedef struct {
Chris@10:      plan_mpi_transpose super;
Chris@10: 
Chris@10:      plan *cld1, *cldtr, *cldtm;
Chris@10:      int preserve_input;
Chris@10: 
Chris@10:      int r; /* "radix" */
Chris@10:      const char *nam;
Chris@10: } P;
Chris@10: 
Chris@10: static void apply(const plan *ego_, R *I, R *O)
Chris@10: {
Chris@10:      const P *ego = (const P *) ego_;
Chris@10:      plan_rdft *cld1, *cldtr, *cldtm;
Chris@10: 
Chris@10:      cld1 = (plan_rdft *) ego->cld1;
Chris@10:      if (cld1) cld1->apply((plan *) cld1, I, O);
Chris@10: 
Chris@10:      if (ego->preserve_input) I = O;
Chris@10: 
Chris@10:      cldtr = (plan_rdft *) ego->cldtr;
Chris@10:      if (cldtr) cldtr->apply((plan *) cldtr, O, I);
Chris@10: 
Chris@10:      cldtm = (plan_rdft *) ego->cldtm;
Chris@10:      if (cldtm) cldtm->apply((plan *) cldtm, I, O);
Chris@10: }
Chris@10: 
Chris@10: static int radix_sqrt(int np)
Chris@10: {
Chris@10:      int r;
Chris@10:      for (r = (int) (X(isqrt)(np)); np % r != 0; ++r)
Chris@10: 	  ;
Chris@10:      return r;
Chris@10: }
Chris@10: 
Chris@10: static int radix_first(int np)
Chris@10: {
Chris@10:      int r = (int) (X(first_divisor)(np));
Chris@10:      return (r >= (int) (X(isqrt)(np)) ? 0 : r);
Chris@10: }
Chris@10: 
Chris@10: /* the local allocated space on process pe required for the given transpose
Chris@10:    dimensions and block sizes */
Chris@10: static INT transpose_space(INT nx, INT ny, INT block, INT tblock, int pe)
Chris@10: {
Chris@10:      return X(imax)(XM(block)(nx, block, pe) * ny,
Chris@10: 		    nx * XM(block)(ny, tblock, pe));
Chris@10: }
Chris@10: 
Chris@10: /* check whether the recursive transposes fit within the space
Chris@10:    that must have been allocated on each process for this transpose;
Chris@10:    this must be modified if the subdivision in mkplan is changed! */
Chris@10: static int enough_space(INT nx, INT ny, INT block, INT tblock,
Chris@10: 			int r, int n_pes)
Chris@10: {
Chris@10:      int pe;
Chris@10:      int m = n_pes / r;
Chris@10:      for (pe = 0; pe < n_pes; ++pe) {
Chris@10: 	  INT space = transpose_space(nx, ny, block, tblock, pe);
Chris@10: 	  INT b1 = XM(block)(nx, r * block, pe / r);
Chris@10: 	  INT b2 = XM(block)(ny, m * tblock, pe % r);
Chris@10: 	  if (transpose_space(b1, ny, block, m*tblock, pe % r) > space
Chris@10: 	      || transpose_space(nx, b2, r*block, tblock, pe / r) > space)
Chris@10: 	       return 0;
Chris@10:      }
Chris@10:      return 1;
Chris@10: }
Chris@10: 
Chris@10: /* In theory, transpose-recurse becomes advantageous for message sizes
Chris@10:    below some minimum, assuming that the time is dominated by
Chris@10:    communications.  In practice, we want to constrain the minimum
Chris@10:    message size for transpose-recurse to keep the planning time down.
Chris@10:    I've set this conservatively according to some simple experiments
Chris@10:    on a Cray XT3 where the crossover message size was 128, although on
Chris@10:    a larger-latency machine the crossover will be larger. */
Chris@10: #define SMALL_MESSAGE 2048
Chris@10: 
Chris@10: static int applicable(const S *ego, const problem *p_,
Chris@10: 		      const planner *plnr, int *r)
Chris@10: {
Chris@10:      const problem_mpi_transpose *p = (const problem_mpi_transpose *) p_;
Chris@10:      int n_pes;
Chris@10:      MPI_Comm_size(p->comm, &n_pes);
Chris@10:      return (1
Chris@10: 	     && p->tblock * n_pes == p->ny
Chris@10: 	     && (!ego->preserve_input || (!NO_DESTROY_INPUTP(plnr)
Chris@10:                                           && p->I != p->O))
Chris@10: 	     && (*r = ego->radix(n_pes)) && *r < n_pes && *r > 1
Chris@10: 	     && enough_space(p->nx, p->ny, p->block, p->tblock, *r, n_pes)
Chris@10: 	     && (!CONSERVE_MEMORYP(plnr) || *r > 8
Chris@10: 		 || !X(toobig)((p->nx * (p->ny / n_pes) * p->vn) / *r))
Chris@10: 	     && (!NO_SLOWP(plnr) || 
Chris@10: 		 (p->nx * (p->ny / n_pes) * p->vn) / n_pes <= SMALL_MESSAGE)
Chris@10: 	     && ONLY_TRANSPOSEDP(p->flags)
Chris@10: 	  );
Chris@10: }
Chris@10: 
Chris@10: static void awake(plan *ego_, enum wakefulness wakefulness)
Chris@10: {
Chris@10:      P *ego = (P *) ego_;
Chris@10:      X(plan_awake)(ego->cld1, wakefulness);
Chris@10:      X(plan_awake)(ego->cldtr, wakefulness);
Chris@10:      X(plan_awake)(ego->cldtm, wakefulness);
Chris@10: }
Chris@10: 
Chris@10: static void destroy(plan *ego_)
Chris@10: {
Chris@10:      P *ego = (P *) ego_;
Chris@10:      X(plan_destroy_internal)(ego->cldtm);
Chris@10:      X(plan_destroy_internal)(ego->cldtr);
Chris@10:      X(plan_destroy_internal)(ego->cld1);
Chris@10: }
Chris@10: 
Chris@10: static void print(const plan *ego_, printer *p)
Chris@10: {
Chris@10:      const P *ego = (const P *) ego_;
Chris@10:      p->print(p, "(mpi-transpose-recurse/%s/%d%s%(%p%)%(%p%)%(%p%))",
Chris@10: 	      ego->nam, ego->r, ego->preserve_input==2 ?"/p":"",
Chris@10: 	      ego->cld1, ego->cldtr, ego->cldtm);
Chris@10: }
Chris@10: 
Chris@10: static plan *mkplan(const solver *ego_, const problem *p_, planner *plnr)
Chris@10: {
Chris@10:      const S *ego = (const S *) ego_;
Chris@10:      const problem_mpi_transpose *p;
Chris@10:      P *pln;
Chris@10:      plan *cld1 = 0, *cldtr = 0, *cldtm = 0;
Chris@10:      R *I, *O;
Chris@10:      int me, np, r, m;
Chris@10:      INT b;
Chris@10:      MPI_Comm comm2;
Chris@10:      static const plan_adt padt = {
Chris@10:           XM(transpose_solve), awake, print, destroy
Chris@10:      };
Chris@10: 
Chris@10:      UNUSED(ego);
Chris@10: 
Chris@10:      if (!applicable(ego, p_, plnr, &r))
Chris@10:           return (plan *) 0;
Chris@10: 
Chris@10:      p = (const problem_mpi_transpose *) p_;
Chris@10: 
Chris@10:      MPI_Comm_size(p->comm, &np);
Chris@10:      MPI_Comm_rank(p->comm, &me);
Chris@10:      m = np / r;
Chris@10:      A(r * m == np);
Chris@10: 
Chris@10:      I = p->I; O = p->O;
Chris@10: 
Chris@10:      b = XM(block)(p->nx, p->block, me);
Chris@10:      A(p->tblock * np == p->ny); /* this is currently required for cld1 */
Chris@10:      if (p->flags & TRANSPOSED_IN) { 
Chris@10:           /* m x r x (bt x b x vn) -> r x m x (bt x b x vn) */
Chris@10: 	  INT vn = p->vn * b * p->tblock;
Chris@10: 	  cld1 = X(mkplan_f_d)(plnr,
Chris@10:                                X(mkproblem_rdft_0_d)(X(mktensor_3d)
Chris@10: 						     (m, r*vn, vn,
Chris@10: 						      r, vn, m*vn,
Chris@10: 						      vn, 1, 1),
Chris@10:                                                      I, O),
Chris@10:                                0, 0, NO_SLOW);
Chris@10:      }
Chris@10:      else if (I != O) { /* combine cld1 with TRANSPOSED_IN permutation */
Chris@10:           /* b x m x r x bt x vn -> r x m x bt x b x vn */
Chris@10: 	  INT vn = p->vn;
Chris@10: 	  INT bt = p->tblock;
Chris@10: 	  cld1 = X(mkplan_f_d)(plnr,
Chris@10:                                X(mkproblem_rdft_0_d)(X(mktensor_5d)
Chris@10: 						     (b, m*r*bt*vn, vn,
Chris@10: 						      m, r*bt*vn, bt*b*vn,
Chris@10: 						      r, bt*vn, m*bt*b*vn,
Chris@10: 						      bt, vn, b*vn,
Chris@10: 						      vn, 1, 1),
Chris@10:                                                      I, O),
Chris@10:                                0, 0, NO_SLOW);
Chris@10:      }
Chris@10:      else { /* TRANSPOSED_IN permutation must be separate for in-place */
Chris@10: 	  /* b x (m x r) x bt x vn -> b x (r x m) x bt x vn */
Chris@10: 	  INT vn = p->vn * p->tblock;
Chris@10: 	  cld1 = X(mkplan_f_d)(plnr,
Chris@10:                                X(mkproblem_rdft_0_d)(X(mktensor_4d)
Chris@10: 						     (m, r*vn, vn,
Chris@10: 						      r, vn, m*vn,
Chris@10: 						      vn, 1, 1,
Chris@10: 						      b, np*vn, np*vn),
Chris@10:                                                      I, O),
Chris@10:                                0, 0, NO_SLOW);
Chris@10:      }
Chris@10:      if (XM(any_true)(!cld1, p->comm)) goto nada;
Chris@10: 
Chris@10:      if (ego->preserve_input || NO_DESTROY_INPUTP(plnr)) I = O;
Chris@10: 
Chris@10:      b = XM(block)(p->nx, r * p->block, me / r);
Chris@10:      MPI_Comm_split(p->comm, me / r, me, &comm2);
Chris@10:      if (b)
Chris@10: 	  cldtr = X(mkplan_d)(plnr, XM(mkproblem_transpose)
Chris@10: 			      (b, p->ny, p->vn,
Chris@10: 			       O, I, p->block, m * p->tblock, comm2, 
Chris@10: 			       p->I != p->O
Chris@10: 			       ? TRANSPOSED_IN : (p->flags & TRANSPOSED_IN)));
Chris@10:      MPI_Comm_free(&comm2);
Chris@10:      if (XM(any_true)(b && !cldtr, p->comm)) goto nada;
Chris@10:      
Chris@10:      b = XM(block)(p->ny, m * p->tblock, me % r);
Chris@10:      MPI_Comm_split(p->comm, me % r, me, &comm2);
Chris@10:      if (b)
Chris@10: 	  cldtm = X(mkplan_d)(plnr, XM(mkproblem_transpose)
Chris@10: 			      (p->nx, b, p->vn,
Chris@10: 			       I, O, r * p->block, p->tblock, comm2, 
Chris@10: 			       TRANSPOSED_IN | (p->flags & TRANSPOSED_OUT)));
Chris@10:      MPI_Comm_free(&comm2);
Chris@10:      if (XM(any_true)(b && !cldtm, p->comm)) goto nada;
Chris@10: 
Chris@10:      pln = MKPLAN_MPI_TRANSPOSE(P, &padt, apply);
Chris@10: 
Chris@10:      pln->cld1 = cld1;
Chris@10:      pln->cldtr = cldtr;
Chris@10:      pln->cldtm = cldtm;
Chris@10:      pln->preserve_input = ego->preserve_input ? 2 : NO_DESTROY_INPUTP(plnr);
Chris@10:      pln->r = r;
Chris@10:      pln->nam = ego->nam;
Chris@10: 
Chris@10:      pln->super.super.ops = cld1->ops;
Chris@10:      if (cldtr) X(ops_add2)(&cldtr->ops, &pln->super.super.ops);
Chris@10:      if (cldtm) X(ops_add2)(&cldtm->ops, &pln->super.super.ops);
Chris@10: 
Chris@10:      return &(pln->super.super);
Chris@10: 
Chris@10:  nada:
Chris@10:      X(plan_destroy_internal)(cldtm);
Chris@10:      X(plan_destroy_internal)(cldtr);
Chris@10:      X(plan_destroy_internal)(cld1);
Chris@10:      return (plan *) 0;
Chris@10: }
Chris@10: 
Chris@10: static solver *mksolver(int preserve_input,
Chris@10: 			int (*radix)(int np), const char *nam)
Chris@10: {
Chris@10:      static const solver_adt sadt = { PROBLEM_MPI_TRANSPOSE, mkplan, 0 };
Chris@10:      S *slv = MKSOLVER(S, &sadt);
Chris@10:      slv->preserve_input = preserve_input;
Chris@10:      slv->radix = radix;
Chris@10:      slv->nam = nam;
Chris@10:      return &(slv->super);
Chris@10: }
Chris@10: 
Chris@10: void XM(transpose_recurse_register)(planner *p)
Chris@10: {
Chris@10:      int preserve_input;
Chris@10:      for (preserve_input = 0; preserve_input <= 1; ++preserve_input) {
Chris@10: 	  REGISTER_SOLVER(p, mksolver(preserve_input, radix_sqrt, "sqrt"));
Chris@10: 	  REGISTER_SOLVER(p, mksolver(preserve_input, radix_first, "first"));
Chris@10:      }
Chris@10: }