Mercurial > hg > sv-dependency-builds
diff src/fftw-3.3.3/mpi/mpi-bench.c @ 10:37bf6b4a2645
Add FFTW3
| author | Chris Cannam |
|---|---|
| date | Wed, 20 Mar 2013 15:35:50 +0000 |
| parents | |
| children |
line wrap: on
line diff
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/fftw-3.3.3/mpi/mpi-bench.c Wed Mar 20 15:35:50 2013 +0000 @@ -0,0 +1,844 @@ +/**************************************************************************/ +/* NOTE to users: this is the FFTW-MPI self-test and benchmark program. + It is probably NOT a good place to learn FFTW usage, since it has a + lot of added complexity in order to exercise and test the full API, + etcetera. We suggest reading the manual. */ +/**************************************************************************/ + +#include <math.h> +#include <stdio.h> +#include <string.h> +#include "fftw3-mpi.h" +#include "fftw-bench.h" + +#if defined(BENCHFFT_SINGLE) +# define BENCH_MPI_TYPE MPI_FLOAT +#elif defined(BENCHFFT_LDOUBLE) +# define BENCH_MPI_TYPE MPI_LONG_DOUBLE +#elif defined(BENCHFFT_QUAD) +# error MPI quad-precision type is unknown +#else +# define BENCH_MPI_TYPE MPI_DOUBLE +#endif + +#if SIZEOF_PTRDIFF_T == SIZEOF_INT +# define FFTW_MPI_PTRDIFF_T MPI_INT +#elif SIZEOF_PTRDIFF_T == SIZEOF_LONG +# define FFTW_MPI_PTRDIFF_T MPI_LONG +#elif SIZEOF_PTRDIFF_T == SIZEOF_LONG_LONG +# define FFTW_MPI_PTRDIFF_T MPI_LONG_LONG +#else +# error MPI type for ptrdiff_t is unknown +# define FFTW_MPI_PTRDIFF_T MPI_LONG +#endif + +static const char *mkversion(void) { return FFTW(version); } +static const char *mkcc(void) { return FFTW(cc); } +static const char *mkcodelet_optim(void) { return FFTW(codelet_optim); } +static const char *mknproc(void) { + static char buf[32]; + int ncpus; + MPI_Comm_size(MPI_COMM_WORLD, &ncpus); +#ifdef HAVE_SNPRINTF + snprintf(buf, 32, "%d", ncpus); +#else + sprintf(buf, "%d", ncpus); +#endif + return buf; +} + +BEGIN_BENCH_DOC +BENCH_DOC("name", "fftw3_mpi") +BENCH_DOCF("version", mkversion) +BENCH_DOCF("cc", mkcc) +BENCH_DOCF("codelet-optim", mkcodelet_optim) +BENCH_DOCF("nproc", mknproc) +END_BENCH_DOC + +static int n_pes = 1, my_pe = 0; + +/* global variables describing the shape of the data and its distribution */ +static int rnk; +static ptrdiff_t vn, iNtot, oNtot; +static ptrdiff_t *local_ni=0, *local_starti=0; +static ptrdiff_t *local_no=0, *local_starto=0; +static ptrdiff_t *all_local_ni=0, *all_local_starti=0; /* n_pes x rnk arrays */ +static ptrdiff_t *all_local_no=0, *all_local_starto=0; /* n_pes x rnk arrays */ +static ptrdiff_t *istrides = 0, *ostrides = 0; +static ptrdiff_t *total_ni=0, *total_no=0; +static int *isend_cnt = 0, *isend_off = 0; /* for MPI_Scatterv */ +static int *orecv_cnt = 0, *orecv_off = 0; /* for MPI_Gatherv */ + +static bench_real *local_in = 0, *local_out = 0; +static bench_real *all_local_in = 0, *all_local_out = 0; +static int all_local_in_alloc = 0, all_local_out_alloc = 0; +static FFTW(plan) plan_scramble_in = 0, plan_unscramble_out = 0; + +static void alloc_rnk(int rnk_) { + rnk = rnk_; + bench_free(local_ni); + if (rnk == 0) + local_ni = 0; + else + local_ni = (ptrdiff_t *) bench_malloc(sizeof(ptrdiff_t) * rnk + * (8 + n_pes * 4)); + + local_starti = local_ni + rnk; + local_no = local_ni + 2 * rnk; + local_starto = local_ni + 3 * rnk; + istrides = local_ni + 4 * rnk; + ostrides = local_ni + 5 * rnk; + total_ni = local_ni + 6 * rnk; + total_no = local_ni + 7 * rnk; + all_local_ni = local_ni + 8 * rnk; + all_local_starti = local_ni + (8 + n_pes) * rnk; + all_local_no = local_ni + (8 + 2 * n_pes) * rnk; + all_local_starto = local_ni + (8 + 3 * n_pes) * rnk; +} + +static void setup_gather_scatter(void) +{ + int i, j; + ptrdiff_t off; + + MPI_Gather(local_ni, rnk, FFTW_MPI_PTRDIFF_T, + all_local_ni, rnk, FFTW_MPI_PTRDIFF_T, + 0, MPI_COMM_WORLD); + MPI_Bcast(all_local_ni, rnk*n_pes, FFTW_MPI_PTRDIFF_T, 0, MPI_COMM_WORLD); + MPI_Gather(local_starti, rnk, FFTW_MPI_PTRDIFF_T, + all_local_starti, rnk, FFTW_MPI_PTRDIFF_T, + 0, MPI_COMM_WORLD); + MPI_Bcast(all_local_starti, rnk*n_pes, FFTW_MPI_PTRDIFF_T, 0, MPI_COMM_WORLD); + + MPI_Gather(local_no, rnk, FFTW_MPI_PTRDIFF_T, + all_local_no, rnk, FFTW_MPI_PTRDIFF_T, + 0, MPI_COMM_WORLD); + MPI_Bcast(all_local_no, rnk*n_pes, FFTW_MPI_PTRDIFF_T, 0, MPI_COMM_WORLD); + MPI_Gather(local_starto, rnk, FFTW_MPI_PTRDIFF_T, + all_local_starto, rnk, FFTW_MPI_PTRDIFF_T, + 0, MPI_COMM_WORLD); + MPI_Bcast(all_local_starto, rnk*n_pes, FFTW_MPI_PTRDIFF_T, 0, MPI_COMM_WORLD); + + off = 0; + for (i = 0; i < n_pes; ++i) { + ptrdiff_t N = vn; + for (j = 0; j < rnk; ++j) + N *= all_local_ni[i * rnk + j]; + isend_cnt[i] = N; + isend_off[i] = off; + off += N; + } + iNtot = off; + all_local_in_alloc = 1; + + istrides[rnk - 1] = vn; + for (j = rnk - 2; j >= 0; --j) + istrides[j] = total_ni[j + 1] * istrides[j + 1]; + + off = 0; + for (i = 0; i < n_pes; ++i) { + ptrdiff_t N = vn; + for (j = 0; j < rnk; ++j) + N *= all_local_no[i * rnk + j]; + orecv_cnt[i] = N; + orecv_off[i] = off; + off += N; + } + oNtot = off; + all_local_out_alloc = 1; + + ostrides[rnk - 1] = vn; + for (j = rnk - 2; j >= 0; --j) + ostrides[j] = total_no[j + 1] * ostrides[j + 1]; +} + +static void copy_block_out(const bench_real *in, + int rnk, ptrdiff_t *n, ptrdiff_t *start, + ptrdiff_t is, ptrdiff_t *os, ptrdiff_t vn, + bench_real *out) +{ + ptrdiff_t i; + if (rnk == 0) { + for (i = 0; i < vn; ++i) + out[i] = in[i]; + } + else if (rnk == 1) { /* this case is just an optimization */ + ptrdiff_t j; + out += start[0] * os[0]; + for (j = 0; j < n[0]; ++j) { + for (i = 0; i < vn; ++i) + out[i] = in[i]; + in += is; + out += os[0]; + } + } + else { + /* we should do n[0] for locality, but this way is simpler to code */ + for (i = 0; i < n[rnk - 1]; ++i) + copy_block_out(in + i * is, + rnk - 1, n, start, is * n[rnk - 1], os, vn, + out + (start[rnk - 1] + i) * os[rnk - 1]); + } +} + +static void copy_block_in(bench_real *in, + int rnk, ptrdiff_t *n, ptrdiff_t *start, + ptrdiff_t is, ptrdiff_t *os, ptrdiff_t vn, + const bench_real *out) +{ + ptrdiff_t i; + if (rnk == 0) { + for (i = 0; i < vn; ++i) + in[i] = out[i]; + } + else if (rnk == 1) { /* this case is just an optimization */ + ptrdiff_t j; + out += start[0] * os[0]; + for (j = 0; j < n[0]; ++j) { + for (i = 0; i < vn; ++i) + in[i] = out[i]; + in += is; + out += os[0]; + } + } + else { + /* we should do n[0] for locality, but this way is simpler to code */ + for (i = 0; i < n[rnk - 1]; ++i) + copy_block_in(in + i * is, + rnk - 1, n, start, is * n[rnk - 1], os, vn, + out + (start[rnk - 1] + i) * os[rnk - 1]); + } +} + +static void do_scatter_in(bench_real *in) +{ + bench_real *ali; + int i; + if (all_local_in_alloc) { + bench_free(all_local_in); + all_local_in = (bench_real*) bench_malloc(iNtot*sizeof(bench_real)); + all_local_in_alloc = 0; + } + ali = all_local_in; + for (i = 0; i < n_pes; ++i) { + copy_block_in(ali, + rnk, all_local_ni + i * rnk, + all_local_starti + i * rnk, + vn, istrides, vn, + in); + ali += isend_cnt[i]; + } + MPI_Scatterv(all_local_in, isend_cnt, isend_off, BENCH_MPI_TYPE, + local_in, isend_cnt[my_pe], BENCH_MPI_TYPE, + 0, MPI_COMM_WORLD); +} + +static void do_gather_out(bench_real *out) +{ + bench_real *alo; + int i; + + if (all_local_out_alloc) { + bench_free(all_local_out); + all_local_out = (bench_real*) bench_malloc(oNtot*sizeof(bench_real)); + all_local_out_alloc = 0; + } + MPI_Gatherv(local_out, orecv_cnt[my_pe], BENCH_MPI_TYPE, + all_local_out, orecv_cnt, orecv_off, BENCH_MPI_TYPE, + 0, MPI_COMM_WORLD); + MPI_Bcast(all_local_out, oNtot, BENCH_MPI_TYPE, 0, MPI_COMM_WORLD); + alo = all_local_out; + for (i = 0; i < n_pes; ++i) { + copy_block_out(alo, + rnk, all_local_no + i * rnk, + all_local_starto + i * rnk, + vn, ostrides, vn, + out); + alo += orecv_cnt[i]; + } +} + +static void alloc_local(ptrdiff_t nreal, int inplace) +{ + bench_free(local_in); + if (local_out != local_in) bench_free(local_out); + local_in = local_out = 0; + if (nreal > 0) { + ptrdiff_t i; + local_in = (bench_real*) bench_malloc(nreal * sizeof(bench_real)); + if (inplace) + local_out = local_in; + else + local_out = (bench_real*) bench_malloc(nreal * sizeof(bench_real)); + for (i = 0; i < nreal; ++i) local_in[i] = local_out[i] = 0.0; + } +} + +void after_problem_rcopy_from(bench_problem *p, bench_real *ri) +{ + UNUSED(p); + do_scatter_in(ri); + if (plan_scramble_in) FFTW(execute)(plan_scramble_in); +} + +void after_problem_rcopy_to(bench_problem *p, bench_real *ro) +{ + UNUSED(p); + if (plan_unscramble_out) FFTW(execute)(plan_unscramble_out); + do_gather_out(ro); +} + +void after_problem_ccopy_from(bench_problem *p, bench_real *ri, bench_real *ii) +{ + UNUSED(ii); + after_problem_rcopy_from(p, ri); +} + +void after_problem_ccopy_to(bench_problem *p, bench_real *ro, bench_real *io) +{ + UNUSED(io); + after_problem_rcopy_to(p, ro); +} + +void after_problem_hccopy_from(bench_problem *p, bench_real *ri, bench_real *ii) +{ + UNUSED(ii); + after_problem_rcopy_from(p, ri); +} + +void after_problem_hccopy_to(bench_problem *p, bench_real *ro, bench_real *io) +{ + UNUSED(io); + after_problem_rcopy_to(p, ro); +} + +static FFTW(plan) mkplan_transpose_local(ptrdiff_t nx, ptrdiff_t ny, + ptrdiff_t vn, + bench_real *in, bench_real *out) +{ + FFTW(iodim64) hdims[3]; + FFTW(r2r_kind) k[3]; + FFTW(plan) pln; + + hdims[0].n = nx; + hdims[0].is = ny * vn; + hdims[0].os = vn; + hdims[1].n = ny; + hdims[1].is = vn; + hdims[1].os = nx * vn; + hdims[2].n = vn; + hdims[2].is = 1; + hdims[2].os = 1; + k[0] = k[1] = k[2] = FFTW_R2HC; + pln = FFTW(plan_guru64_r2r)(0, 0, 3, hdims, in, out, k, FFTW_ESTIMATE); + BENCH_ASSERT(pln != 0); + return pln; +} + +static int tensor_rowmajor_transposedp(bench_tensor *t) +{ + bench_iodim *d; + int i; + + BENCH_ASSERT(FINITE_RNK(t->rnk)); + if (t->rnk < 2) + return 0; + + d = t->dims; + if (d[0].is != d[1].is * d[1].n + || d[0].os != d[1].is + || d[1].os != d[0].os * d[0].n) + return 0; + if (t->rnk > 2 && d[1].is != d[2].is * d[2].n) + return 0; + for (i = 2; i + 1 < t->rnk; ++i) { + d = t->dims + i; + if (d[0].is != d[1].is * d[1].n + || d[0].os != d[1].os * d[1].n) + return 0; + } + + if (t->rnk > 2 && t->dims[t->rnk-1].is != t->dims[t->rnk-1].os) + return 0; + return 1; +} + +static int tensor_contiguousp(bench_tensor *t, int s) +{ + return (t->dims[t->rnk-1].is == s + && ((tensor_rowmajorp(t) && + t->dims[t->rnk-1].is == t->dims[t->rnk-1].os) + || tensor_rowmajor_transposedp(t))); +} + +static FFTW(plan) mkplan_complex(bench_problem *p, unsigned flags) +{ + FFTW(plan) pln = 0; + int i; + ptrdiff_t ntot; + + vn = p->vecsz->rnk == 1 ? p->vecsz->dims[0].n : 1; + + if (p->sz->rnk < 1 + || p->split + || !tensor_contiguousp(p->sz, vn) + || tensor_rowmajor_transposedp(p->sz) + || p->vecsz->rnk > 1 + || (p->vecsz->rnk == 1 && (p->vecsz->dims[0].is != 1 + || p->vecsz->dims[0].os != 1))) + return 0; + + alloc_rnk(p->sz->rnk); + for (i = 0; i < rnk; ++i) { + total_ni[i] = total_no[i] = p->sz->dims[i].n; + local_ni[i] = local_no[i] = total_ni[i]; + local_starti[i] = local_starto[i] = 0; + } + if (rnk > 1) { + ptrdiff_t n, start, nT, startT; + ntot = FFTW(mpi_local_size_many_transposed) + (p->sz->rnk, total_ni, vn, + FFTW_MPI_DEFAULT_BLOCK, FFTW_MPI_DEFAULT_BLOCK, + MPI_COMM_WORLD, + &n, &start, &nT, &startT); + if (flags & FFTW_MPI_TRANSPOSED_IN) { + local_ni[1] = nT; + local_starti[1] = startT; + } + else { + local_ni[0] = n; + local_starti[0] = start; + } + if (flags & FFTW_MPI_TRANSPOSED_OUT) { + local_no[1] = nT; + local_starto[1] = startT; + } + else { + local_no[0] = n; + local_starto[0] = start; + } + } + else if (rnk == 1) { + ntot = FFTW(mpi_local_size_many_1d) + (total_ni[0], vn, MPI_COMM_WORLD, p->sign, flags, + local_ni, local_starti, local_no, local_starto); + } + alloc_local(ntot * 2, p->in == p->out); + + pln = FFTW(mpi_plan_many_dft)(p->sz->rnk, total_ni, vn, + FFTW_MPI_DEFAULT_BLOCK, + FFTW_MPI_DEFAULT_BLOCK, + (FFTW(complex) *) local_in, + (FFTW(complex) *) local_out, + MPI_COMM_WORLD, p->sign, flags); + + vn *= 2; + + if (rnk > 1) { + ptrdiff_t nrest = 1; + for (i = 2; i < rnk; ++i) nrest *= p->sz->dims[i].n; + if (flags & FFTW_MPI_TRANSPOSED_IN) + plan_scramble_in = mkplan_transpose_local( + p->sz->dims[0].n, local_ni[1], vn * nrest, + local_in, local_in); + if (flags & FFTW_MPI_TRANSPOSED_OUT) + plan_unscramble_out = mkplan_transpose_local( + local_no[1], p->sz->dims[0].n, vn * nrest, + local_out, local_out); + } + + return pln; +} + +static int tensor_real_contiguousp(bench_tensor *t, int sign, int s) +{ + return (t->dims[t->rnk-1].is == s + && ((tensor_real_rowmajorp(t, sign, 1) && + t->dims[t->rnk-1].is == t->dims[t->rnk-1].os))); +} + +static FFTW(plan) mkplan_real(bench_problem *p, unsigned flags) +{ + FFTW(plan) pln = 0; + int i; + ptrdiff_t ntot; + + vn = p->vecsz->rnk == 1 ? p->vecsz->dims[0].n : 1; + + if (p->sz->rnk < 2 + || p->split + || !tensor_real_contiguousp(p->sz, p->sign, vn) + || tensor_rowmajor_transposedp(p->sz) + || p->vecsz->rnk > 1 + || (p->vecsz->rnk == 1 && (p->vecsz->dims[0].is != 1 + || p->vecsz->dims[0].os != 1))) + return 0; + + alloc_rnk(p->sz->rnk); + for (i = 0; i < rnk; ++i) { + total_ni[i] = total_no[i] = p->sz->dims[i].n; + local_ni[i] = local_no[i] = total_ni[i]; + local_starti[i] = local_starto[i] = 0; + } + local_ni[rnk-1] = local_no[rnk-1] = total_ni[rnk-1] = total_no[rnk-1] + = p->sz->dims[rnk-1].n / 2 + 1; + { + ptrdiff_t n, start, nT, startT; + ntot = FFTW(mpi_local_size_many_transposed) + (p->sz->rnk, total_ni, vn, + FFTW_MPI_DEFAULT_BLOCK, FFTW_MPI_DEFAULT_BLOCK, + MPI_COMM_WORLD, + &n, &start, &nT, &startT); + if (flags & FFTW_MPI_TRANSPOSED_IN) { + local_ni[1] = nT; + local_starti[1] = startT; + } + else { + local_ni[0] = n; + local_starti[0] = start; + } + if (flags & FFTW_MPI_TRANSPOSED_OUT) { + local_no[1] = nT; + local_starto[1] = startT; + } + else { + local_no[0] = n; + local_starto[0] = start; + } + } + alloc_local(ntot * 2, p->in == p->out); + + total_ni[rnk - 1] = p->sz->dims[rnk - 1].n; + if (p->sign < 0) + pln = FFTW(mpi_plan_many_dft_r2c)(p->sz->rnk, total_ni, vn, + FFTW_MPI_DEFAULT_BLOCK, + FFTW_MPI_DEFAULT_BLOCK, + local_in, + (FFTW(complex) *) local_out, + MPI_COMM_WORLD, flags); + else + pln = FFTW(mpi_plan_many_dft_c2r)(p->sz->rnk, total_ni, vn, + FFTW_MPI_DEFAULT_BLOCK, + FFTW_MPI_DEFAULT_BLOCK, + (FFTW(complex) *) local_in, + local_out, + MPI_COMM_WORLD, flags); + + total_ni[rnk - 1] = p->sz->dims[rnk - 1].n / 2 + 1; + vn *= 2; + + { + ptrdiff_t nrest = 1; + for (i = 2; i < rnk; ++i) nrest *= total_ni[i]; + if (flags & FFTW_MPI_TRANSPOSED_IN) + plan_scramble_in = mkplan_transpose_local( + total_ni[0], local_ni[1], vn * nrest, + local_in, local_in); + if (flags & FFTW_MPI_TRANSPOSED_OUT) + plan_unscramble_out = mkplan_transpose_local( + local_no[1], total_ni[0], vn * nrest, + local_out, local_out); + } + + return pln; +} + +static FFTW(plan) mkplan_transpose(bench_problem *p, unsigned flags) +{ + ptrdiff_t ntot, nx, ny; + int ix=0, iy=1, i; + const bench_iodim *d = p->vecsz->dims; + FFTW(plan) pln; + + if (p->vecsz->rnk == 3) { + for (i = 0; i < 3; ++i) + if (d[i].is == 1 && d[i].os == 1) { + vn = d[i].n; + ix = (i + 1) % 3; + iy = (i + 2) % 3; + break; + } + if (i == 3) return 0; + } + else { + vn = 1; + ix = 0; + iy = 1; + } + + if (d[ix].is == d[iy].n * vn && d[ix].os == vn + && d[iy].os == d[ix].n * vn && d[iy].is == vn) { + nx = d[ix].n; + ny = d[iy].n; + } + else if (d[iy].is == d[ix].n * vn && d[iy].os == vn + && d[ix].os == d[iy].n * vn && d[ix].is == vn) { + nx = d[iy].n; + ny = d[ix].n; + } + else + return 0; + + alloc_rnk(2); + ntot = vn * FFTW(mpi_local_size_2d_transposed)(nx, ny, MPI_COMM_WORLD, + &local_ni[0], + &local_starti[0], + &local_no[0], + &local_starto[0]); + local_ni[1] = ny; + local_starti[1] = 0; + local_no[1] = nx; + local_starto[1] = 0; + total_ni[0] = nx; total_ni[1] = ny; + total_no[1] = nx; total_no[0] = ny; + alloc_local(ntot, p->in == p->out); + + pln = FFTW(mpi_plan_many_transpose)(nx, ny, vn, + FFTW_MPI_DEFAULT_BLOCK, + FFTW_MPI_DEFAULT_BLOCK, + local_in, local_out, + MPI_COMM_WORLD, flags); + + if (flags & FFTW_MPI_TRANSPOSED_IN) + plan_scramble_in = mkplan_transpose_local(local_ni[0], ny, vn, + local_in, local_in); + if (flags & FFTW_MPI_TRANSPOSED_OUT) + plan_unscramble_out = mkplan_transpose_local + (nx, local_no[0], vn, local_out, local_out); + +#if 0 + if (pln && vn == 1) { + int i, j; + bench_real *ri = (bench_real *) p->in; + bench_real *ro = (bench_real *) p->out; + if (!ri || !ro) return pln; + setup_gather_scatter(); + for (i = 0; i < nx * ny; ++i) + ri[i] = i; + after_problem_rcopy_from(p, ri); + FFTW(execute)(pln); + after_problem_rcopy_to(p, ro); + if (my_pe == 0) { + for (i = 0; i < nx; ++i) { + for (j = 0; j < ny; ++j) + printf(" %3g", ro[j * nx + i]); + printf("\n"); + } + } + } +#endif + + return pln; +} + +static FFTW(plan) mkplan_r2r(bench_problem *p, unsigned flags) +{ + FFTW(plan) pln = 0; + int i; + ptrdiff_t ntot; + FFTW(r2r_kind) *k; + + if ((p->sz->rnk == 0 || (p->sz->rnk == 1 && p->sz->dims[0].n == 1)) + && p->vecsz->rnk >= 2 && p->vecsz->rnk <= 3) + return mkplan_transpose(p, flags); + + vn = p->vecsz->rnk == 1 ? p->vecsz->dims[0].n : 1; + + if (p->sz->rnk < 1 + || p->split + || !tensor_contiguousp(p->sz, vn) + || tensor_rowmajor_transposedp(p->sz) + || p->vecsz->rnk > 1 + || (p->vecsz->rnk == 1 && (p->vecsz->dims[0].is != 1 + || p->vecsz->dims[0].os != 1))) + return 0; + + alloc_rnk(p->sz->rnk); + for (i = 0; i < rnk; ++i) { + total_ni[i] = total_no[i] = p->sz->dims[i].n; + local_ni[i] = local_no[i] = total_ni[i]; + local_starti[i] = local_starto[i] = 0; + } + if (rnk > 1) { + ptrdiff_t n, start, nT, startT; + ntot = FFTW(mpi_local_size_many_transposed) + (p->sz->rnk, total_ni, vn, + FFTW_MPI_DEFAULT_BLOCK, FFTW_MPI_DEFAULT_BLOCK, + MPI_COMM_WORLD, + &n, &start, &nT, &startT); + if (flags & FFTW_MPI_TRANSPOSED_IN) { + local_ni[1] = nT; + local_starti[1] = startT; + } + else { + local_ni[0] = n; + local_starti[0] = start; + } + if (flags & FFTW_MPI_TRANSPOSED_OUT) { + local_no[1] = nT; + local_starto[1] = startT; + } + else { + local_no[0] = n; + local_starto[0] = start; + } + } + else if (rnk == 1) { + ntot = FFTW(mpi_local_size_many_1d) + (total_ni[0], vn, MPI_COMM_WORLD, p->sign, flags, + local_ni, local_starti, local_no, local_starto); + } + alloc_local(ntot, p->in == p->out); + + k = (FFTW(r2r_kind) *) bench_malloc(sizeof(FFTW(r2r_kind)) * p->sz->rnk); + for (i = 0; i < p->sz->rnk; ++i) + switch (p->k[i]) { + case R2R_R2HC: k[i] = FFTW_R2HC; break; + case R2R_HC2R: k[i] = FFTW_HC2R; break; + case R2R_DHT: k[i] = FFTW_DHT; break; + case R2R_REDFT00: k[i] = FFTW_REDFT00; break; + case R2R_REDFT01: k[i] = FFTW_REDFT01; break; + case R2R_REDFT10: k[i] = FFTW_REDFT10; break; + case R2R_REDFT11: k[i] = FFTW_REDFT11; break; + case R2R_RODFT00: k[i] = FFTW_RODFT00; break; + case R2R_RODFT01: k[i] = FFTW_RODFT01; break; + case R2R_RODFT10: k[i] = FFTW_RODFT10; break; + case R2R_RODFT11: k[i] = FFTW_RODFT11; break; + default: BENCH_ASSERT(0); + } + + pln = FFTW(mpi_plan_many_r2r)(p->sz->rnk, total_ni, vn, + FFTW_MPI_DEFAULT_BLOCK, + FFTW_MPI_DEFAULT_BLOCK, + local_in, local_out, + MPI_COMM_WORLD, k, flags); + bench_free(k); + + if (rnk > 1) { + ptrdiff_t nrest = 1; + for (i = 2; i < rnk; ++i) nrest *= p->sz->dims[i].n; + if (flags & FFTW_MPI_TRANSPOSED_IN) + plan_scramble_in = mkplan_transpose_local( + p->sz->dims[0].n, local_ni[1], vn * nrest, + local_in, local_in); + if (flags & FFTW_MPI_TRANSPOSED_OUT) + plan_unscramble_out = mkplan_transpose_local( + local_no[1], p->sz->dims[0].n, vn * nrest, + local_out, local_out); + } + + return pln; +} + +FFTW(plan) mkplan(bench_problem *p, unsigned flags) +{ + FFTW(plan) pln = 0; + FFTW(destroy_plan)(plan_scramble_in); plan_scramble_in = 0; + FFTW(destroy_plan)(plan_unscramble_out); plan_unscramble_out = 0; + if (p->scrambled_in) { + if (p->sz->rnk == 1 && p->sz->dims[0].n != 1) + flags |= FFTW_MPI_SCRAMBLED_IN; + else + flags |= FFTW_MPI_TRANSPOSED_IN; + } + if (p->scrambled_out) { + if (p->sz->rnk == 1 && p->sz->dims[0].n != 1) + flags |= FFTW_MPI_SCRAMBLED_OUT; + else + flags |= FFTW_MPI_TRANSPOSED_OUT; + } + switch (p->kind) { + case PROBLEM_COMPLEX: + pln =mkplan_complex(p, flags); + break; + case PROBLEM_REAL: + pln = mkplan_real(p, flags); + break; + case PROBLEM_R2R: + pln = mkplan_r2r(p, flags); + break; + default: BENCH_ASSERT(0); + } + if (pln) setup_gather_scatter(); + return pln; +} + +void main_init(int *argc, char ***argv) +{ +#ifdef HAVE_SMP +# if MPI_VERSION >= 2 /* for MPI_Init_thread */ + int provided; + MPI_Init_thread(argc, argv, MPI_THREAD_FUNNELED, &provided); + threads_ok = provided >= MPI_THREAD_FUNNELED; +# else + MPI_Init(argc, argv); + threads_ok = 0; +# endif +#else + MPI_Init(argc, argv); +#endif + MPI_Comm_rank(MPI_COMM_WORLD, &my_pe); + MPI_Comm_size(MPI_COMM_WORLD, &n_pes); + if (my_pe != 0) verbose = -999; + no_speed_allocation = 1; /* so we can benchmark transforms > memory */ + always_pad_real = 1; /* out-of-place real transforms are padded */ + isend_cnt = (int *) bench_malloc(sizeof(int) * n_pes); + isend_off = (int *) bench_malloc(sizeof(int) * n_pes); + orecv_cnt = (int *) bench_malloc(sizeof(int) * n_pes); + orecv_off = (int *) bench_malloc(sizeof(int) * n_pes); + + /* init_threads must be called before any other FFTW function, + including mpi_init, because it has to register the threads hooks + before the planner is initalized */ +#ifdef HAVE_SMP + if (threads_ok) { BENCH_ASSERT(FFTW(init_threads)()); } +#endif + FFTW(mpi_init)(); +} + +void initial_cleanup(void) +{ + alloc_rnk(0); + alloc_local(0, 0); + bench_free(all_local_in); all_local_in = 0; + bench_free(all_local_out); all_local_out = 0; + bench_free(isend_off); isend_off = 0; + bench_free(isend_cnt); isend_cnt = 0; + bench_free(orecv_off); orecv_off = 0; + bench_free(orecv_cnt); orecv_cnt = 0; + FFTW(destroy_plan)(plan_scramble_in); plan_scramble_in = 0; + FFTW(destroy_plan)(plan_unscramble_out); plan_unscramble_out = 0; +} + +void final_cleanup(void) +{ + MPI_Finalize(); +} + +void bench_exit(int status) +{ + MPI_Abort(MPI_COMM_WORLD, status); +} + +double bench_cost_postprocess(double cost) +{ + double cost_max; + MPI_Allreduce(&cost, &cost_max, 1, MPI_DOUBLE, MPI_MAX, MPI_COMM_WORLD); + return cost_max; +} + + +int import_wisdom(FILE *f) +{ + int success = 1, sall; + if (my_pe == 0) success = FFTW(import_wisdom_from_file)(f); + FFTW(mpi_broadcast_wisdom)(MPI_COMM_WORLD); + MPI_Allreduce(&success, &sall, 1, MPI_INT, MPI_LAND, MPI_COMM_WORLD); + return sall; +} + +void export_wisdom(FILE *f) +{ + FFTW(mpi_gather_wisdom)(MPI_COMM_WORLD); + if (my_pe == 0) FFTW(export_wisdom_to_file)(f); +}