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: /* Complex DFTs of rank == 1 via six-step algorithm. */ cannam@127: cannam@127: #include "mpi-dft.h" cannam@127: #include "mpi-transpose.h" cannam@127: #include "dft.h" cannam@127: cannam@127: typedef struct { cannam@127: solver super; cannam@127: rdftapply apply; /* apply_ddft_first or apply_ddft_last */ 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_dft super; cannam@127: cannam@127: triggen *t; cannam@127: plan *cldt, *cld_ddft, *cld_dft; cannam@127: INT roff, ioff; cannam@127: int preserve_input; cannam@127: INT vn, xmin, xmax, xs, m, r; cannam@127: } P; cannam@127: cannam@127: static void do_twiddle(triggen *t, INT ir, INT m, INT vn, R *xr, R *xi) cannam@127: { cannam@127: void (*rotate)(triggen *, INT, R, R, R *) = t->rotate; cannam@127: INT im, iv; cannam@127: for (im = 0; im < m; ++im) cannam@127: for (iv = 0; iv < vn; ++iv) { cannam@127: /* TODO: modify/inline rotate function cannam@127: so that it can do whole vn vector at once? */ cannam@127: R c[2]; cannam@127: rotate(t, ir * im, *xr, *xi, c); cannam@127: *xr = c[0]; *xi = c[1]; cannam@127: xr += 2; xi += 2; cannam@127: } cannam@127: } cannam@127: cannam@127: /* radix-r DFT of size r*m. This is equivalent to an m x r 2d DFT, cannam@127: plus twiddle factors between the size-m and size-r 1d DFTs, where cannam@127: the m dimension is initially distributed. The output is transposed cannam@127: to r x m where the r dimension is distributed. cannam@127: cannam@127: This algorithm follows the general sequence: cannam@127: global transpose (m x r -> r x m) cannam@127: DFTs of size m cannam@127: multiply by twiddles + global transpose (r x m -> m x r) cannam@127: DFTs of size r cannam@127: global transpose (m x r -> r x m) cannam@127: where the multiplication by twiddles can come before or after cannam@127: the middle transpose. The first/last transposes are omitted cannam@127: for SCRAMBLED_IN/OUT formats, respectively. cannam@127: cannam@127: However, we wish to exploit our dft-rank1-bigvec solver, which cannam@127: solves a vector of distributed DFTs via transpose+dft+transpose. cannam@127: Therefore, we can group *either* the DFTs of size m *or* the cannam@127: DFTs of size r with their surrounding transposes as a single cannam@127: distributed-DFT (ddft) plan. These two variations correspond to cannam@127: apply_ddft_first or apply_ddft_last, respectively. cannam@127: */ cannam@127: cannam@127: static void apply_ddft_first(const plan *ego_, R *I, R *O) cannam@127: { cannam@127: const P *ego = (const P *) ego_; cannam@127: plan_dft *cld_dft; cannam@127: plan_rdft *cldt, *cld_ddft; cannam@127: INT roff, ioff, im, mmax, ms, r, vn; cannam@127: triggen *t; cannam@127: R *dI, *dO; cannam@127: cannam@127: /* distributed size-m DFTs, with output in m x r format */ cannam@127: cld_ddft = (plan_rdft *) ego->cld_ddft; cannam@127: cld_ddft->apply(ego->cld_ddft, I, O); cannam@127: cannam@127: cldt = (plan_rdft *) ego->cldt; cannam@127: if (ego->preserve_input || !cldt) I = O; cannam@127: cannam@127: /* twiddle multiplications, followed by 1d DFTs of size-r */ cannam@127: cld_dft = (plan_dft *) ego->cld_dft; cannam@127: roff = ego->roff; ioff = ego->ioff; cannam@127: mmax = ego->xmax; ms = ego->xs; cannam@127: t = ego->t; r = ego->r; vn = ego->vn; cannam@127: dI = O; dO = I; cannam@127: for (im = ego->xmin; im <= mmax; ++im) { cannam@127: do_twiddle(t, im, r, vn, dI+roff, dI+ioff); cannam@127: cld_dft->apply((plan *) cld_dft, dI+roff, dI+ioff, dO+roff, dO+ioff); cannam@127: dI += ms; dO += ms; cannam@127: } cannam@127: cannam@127: /* final global transpose (m x r -> r x m), if not SCRAMBLED_OUT */ cannam@127: if (cldt) cannam@127: cldt->apply((plan *) cldt, I, O); cannam@127: } cannam@127: cannam@127: static void apply_ddft_last(const plan *ego_, R *I, R *O) cannam@127: { cannam@127: const P *ego = (const P *) ego_; cannam@127: plan_dft *cld_dft; cannam@127: plan_rdft *cldt, *cld_ddft; cannam@127: INT roff, ioff, ir, rmax, rs, m, vn; cannam@127: triggen *t; cannam@127: R *dI, *dO0, *dO; cannam@127: cannam@127: /* initial global transpose (m x r -> r x m), if not SCRAMBLED_IN */ cannam@127: cldt = (plan_rdft *) ego->cldt; cannam@127: if (cldt) { cannam@127: cldt->apply((plan *) cldt, I, O); cannam@127: dI = O; cannam@127: } cannam@127: else cannam@127: dI = I; cannam@127: if (ego->preserve_input) dO = O; else dO = I; cannam@127: dO0 = dO; cannam@127: cannam@127: /* 1d DFTs of size m, followed by twiddle multiplications */ cannam@127: cld_dft = (plan_dft *) ego->cld_dft; cannam@127: roff = ego->roff; ioff = ego->ioff; cannam@127: rmax = ego->xmax; rs = ego->xs; cannam@127: t = ego->t; m = ego->m; vn = ego->vn; cannam@127: for (ir = ego->xmin; ir <= rmax; ++ir) { cannam@127: cld_dft->apply((plan *) cld_dft, dI+roff, dI+ioff, dO+roff, dO+ioff); cannam@127: do_twiddle(t, ir, m, vn, dO+roff, dO+ioff); cannam@127: dI += rs; dO += rs; cannam@127: } cannam@127: cannam@127: /* distributed size-r DFTs, with output in r x m format */ cannam@127: cld_ddft = (plan_rdft *) ego->cld_ddft; cannam@127: cld_ddft->apply(ego->cld_ddft, dO0, O); cannam@127: } cannam@127: cannam@127: static int applicable(const S *ego, const problem *p_, cannam@127: const planner *plnr, cannam@127: INT *r, INT rblock[2], INT mblock[2]) cannam@127: { cannam@127: const problem_mpi_dft *p = (const problem_mpi_dft *) p_; cannam@127: int n_pes; cannam@127: MPI_Comm_size(p->comm, &n_pes); cannam@127: return (1 cannam@127: && p->sz->rnk == 1 cannam@127: cannam@127: && ONLY_SCRAMBLEDP(p->flags) cannam@127: cannam@127: && (!ego->preserve_input || (!NO_DESTROY_INPUTP(plnr) cannam@127: && p->I != p->O)) cannam@127: cannam@127: && (!(p->flags & SCRAMBLED_IN) || ego->apply == apply_ddft_last) cannam@127: && (!(p->flags & SCRAMBLED_OUT) || ego->apply == apply_ddft_first) cannam@127: cannam@127: && (!NO_SLOWP(plnr) /* slow if dft-serial is applicable */ cannam@127: || !XM(dft_serial_applicable)(p)) cannam@127: cannam@127: /* disallow if dft-rank1-bigvec is applicable since the cannam@127: data distribution may be slightly different (ugh!) */ cannam@127: && (p->vn < n_pes || p->flags) cannam@127: cannam@127: && (*r = XM(choose_radix)(p->sz->dims[0], n_pes, cannam@127: p->flags, p->sign, cannam@127: rblock, mblock)) cannam@127: cannam@127: /* ddft_first or last has substantial advantages in the cannam@127: bigvec transpositions for the common case where cannam@127: n_pes == n/r or r, respectively */ cannam@127: && (!NO_UGLYP(plnr) cannam@127: || !(*r == n_pes && ego->apply == apply_ddft_first) cannam@127: || !(p->sz->dims[0].n / *r == n_pes cannam@127: && ego->apply == apply_ddft_last)) 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->cldt, wakefulness); cannam@127: X(plan_awake)(ego->cld_dft, wakefulness); cannam@127: X(plan_awake)(ego->cld_ddft, wakefulness); cannam@127: cannam@127: switch (wakefulness) { cannam@127: case SLEEPY: cannam@127: X(triggen_destroy)(ego->t); ego->t = 0; cannam@127: break; cannam@127: default: cannam@127: ego->t = X(mktriggen)(AWAKE_SQRTN_TABLE, ego->r * ego->m); cannam@127: break; cannam@127: } 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->cldt); cannam@127: X(plan_destroy_internal)(ego->cld_dft); cannam@127: X(plan_destroy_internal)(ego->cld_ddft); 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-dft-rank1/%D%s%s%(%p%)%(%p%)%(%p%))", cannam@127: ego->r, cannam@127: ego->super.apply == apply_ddft_first ? "/first" : "/last", cannam@127: ego->preserve_input==2 ?"/p":"", cannam@127: ego->cld_ddft, ego->cld_dft, ego->cldt); 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_dft *p; cannam@127: P *pln; cannam@127: plan *cld_dft = 0, *cld_ddft = 0, *cldt = 0; cannam@127: R *ri, *ii, *ro, *io, *I, *O; cannam@127: INT r, rblock[2], m, mblock[2], rp, mp, mpblock[2], mpb; cannam@127: int my_pe, n_pes, preserve_input, ddft_first; cannam@127: dtensor *sz; cannam@127: static const plan_adt padt = { cannam@127: XM(dft_solve), awake, print, destroy cannam@127: }; cannam@127: cannam@127: UNUSED(ego); cannam@127: cannam@127: if (!applicable(ego, p_, plnr, &r, rblock, mblock)) cannam@127: return (plan *) 0; cannam@127: cannam@127: p = (const problem_mpi_dft *) p_; cannam@127: cannam@127: MPI_Comm_rank(p->comm, &my_pe); cannam@127: MPI_Comm_size(p->comm, &n_pes); cannam@127: cannam@127: m = p->sz->dims[0].n / r; cannam@127: cannam@127: /* some hackery so that we can plan both ddft_first and ddft_last cannam@127: as if they were ddft_first */ cannam@127: if ((ddft_first = (ego->apply == apply_ddft_first))) { cannam@127: rp = r; mp = m; cannam@127: mpblock[IB] = mblock[IB]; mpblock[OB] = mblock[OB]; cannam@127: mpb = XM(block)(mp, mpblock[OB], my_pe); cannam@127: } cannam@127: else { cannam@127: rp = m; mp = r; cannam@127: mpblock[IB] = rblock[IB]; mpblock[OB] = rblock[OB]; cannam@127: mpb = XM(block)(mp, mpblock[IB], my_pe); cannam@127: } cannam@127: cannam@127: preserve_input = ego->preserve_input ? 2 : NO_DESTROY_INPUTP(plnr); cannam@127: cannam@127: sz = XM(mkdtensor)(1); cannam@127: sz->dims[0].n = mp; cannam@127: sz->dims[0].b[IB] = mpblock[IB]; cannam@127: sz->dims[0].b[OB] = mpblock[OB]; cannam@127: I = (ddft_first || !preserve_input) ? p->I : p->O; cannam@127: O = p->O; cannam@127: cld_ddft = X(mkplan_d)(plnr, XM(mkproblem_dft_d)(sz, rp * p->vn, cannam@127: I, O, p->comm, p->sign, cannam@127: RANK1_BIGVEC_ONLY)); cannam@127: if (XM(any_true)(!cld_ddft, p->comm)) goto nada; cannam@127: cannam@127: I = TAINT((ddft_first || !p->flags) ? p->O : p->I, rp * p->vn * 2); cannam@127: O = TAINT((preserve_input || (ddft_first && p->flags)) ? p->O : p->I, cannam@127: rp * p->vn * 2); cannam@127: X(extract_reim)(p->sign, I, &ri, &ii); cannam@127: X(extract_reim)(p->sign, O, &ro, &io); cannam@127: cld_dft = X(mkplan_d)(plnr, cannam@127: X(mkproblem_dft_d)(X(mktensor_1d)(rp, p->vn*2,p->vn*2), cannam@127: X(mktensor_1d)(p->vn, 2, 2), cannam@127: ri, ii, ro, io)); cannam@127: if (XM(any_true)(!cld_dft, p->comm)) goto nada; cannam@127: cannam@127: if (!p->flags) { /* !(SCRAMBLED_IN or SCRAMBLED_OUT) */ cannam@127: I = (ddft_first && preserve_input) ? p->O : p->I; cannam@127: O = p->O; cannam@127: cldt = X(mkplan_d)(plnr, cannam@127: XM(mkproblem_transpose)( cannam@127: m, r, p->vn * 2, cannam@127: I, O, cannam@127: ddft_first ? mblock[OB] : mblock[IB], cannam@127: ddft_first ? rblock[OB] : rblock[IB], cannam@127: p->comm, 0)); cannam@127: if (XM(any_true)(!cldt, p->comm)) goto nada; cannam@127: } cannam@127: cannam@127: pln = MKPLAN_MPI_DFT(P, &padt, ego->apply); cannam@127: cannam@127: pln->cld_ddft = cld_ddft; cannam@127: pln->cld_dft = cld_dft; cannam@127: pln->cldt = cldt; cannam@127: pln->preserve_input = preserve_input; cannam@127: X(extract_reim)(p->sign, p->O, &ro, &io); cannam@127: pln->roff = ro - p->O; cannam@127: pln->ioff = io - p->O; cannam@127: pln->vn = p->vn; cannam@127: pln->m = m; cannam@127: pln->r = r; cannam@127: pln->xmin = (ddft_first ? mblock[OB] : rblock[IB]) * my_pe; cannam@127: pln->xmax = pln->xmin + mpb - 1; cannam@127: pln->xs = rp * p->vn * 2; cannam@127: pln->t = 0; cannam@127: cannam@127: X(ops_add)(&cld_ddft->ops, &cld_dft->ops, &pln->super.super.ops); cannam@127: if (cldt) X(ops_add2)(&cldt->ops, &pln->super.super.ops); cannam@127: { cannam@127: double n0 = (1 + pln->xmax - pln->xmin) * (mp - 1) * pln->vn; cannam@127: pln->super.super.ops.mul += 8 * n0; cannam@127: pln->super.super.ops.add += 4 * n0; cannam@127: pln->super.super.ops.other += 8 * n0; cannam@127: } cannam@127: cannam@127: return &(pln->super.super); cannam@127: cannam@127: nada: cannam@127: X(plan_destroy_internal)(cldt); cannam@127: X(plan_destroy_internal)(cld_dft); cannam@127: X(plan_destroy_internal)(cld_ddft); cannam@127: return (plan *) 0; cannam@127: } cannam@127: cannam@127: static solver *mksolver(rdftapply apply, int preserve_input) cannam@127: { cannam@127: static const solver_adt sadt = { PROBLEM_MPI_DFT, mkplan, 0 }; cannam@127: S *slv = MKSOLVER(S, &sadt); cannam@127: slv->apply = apply; cannam@127: slv->preserve_input = preserve_input; cannam@127: return &(slv->super); cannam@127: } cannam@127: cannam@127: void XM(dft_rank1_register)(planner *p) cannam@127: { cannam@127: rdftapply apply[] = { apply_ddft_first, apply_ddft_last }; cannam@127: unsigned int iapply; cannam@127: int preserve_input; cannam@127: for (iapply = 0; iapply < sizeof(apply) / sizeof(apply[0]); ++iapply) cannam@127: for (preserve_input = 0; preserve_input <= 1; ++preserve_input) cannam@127: REGISTER_SOLVER(p, mksolver(apply[iapply], preserve_input)); cannam@127: }