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