1 /*
2 * Copyright (c) 2003, 2007-14 Matteo Frigo
3 * Copyright (c) 2003, 2007-14 Massachusetts Institute of Technology
4 *
5 * This program is free software; you can redistribute it and/or modify
6 * it under the terms of the GNU General Public License as published by
7 * the Free Software Foundation; either version 2 of the License, or
8 * (at your option) any later version.
9 *
10 * This program is distributed in the hope that it will be useful,
11 * but WITHOUT ANY WARRANTY; without even the implied warranty of
12 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13 * GNU General Public License for more details.
14 *
15 * You should have received a copy of the GNU General Public License
16 * along with this program; if not, write to the Free Software
17 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
18 *
19 */
20
21 /* Real-input (r2c) DFTs of rank >= 2, for the case where we are distributed
22 across the first dimension only, and the output is transposed both
23 in data distribution and in ordering (for the first 2 dimensions).
24
25 Conversely, real-output (c2r) DFTs where the input is transposed.
26
27 We don't currently support transposed-input r2c or transposed-output
28 c2r transforms. */
29
30 #include "mpi-rdft2.h"
31 #include "mpi-transpose.h"
32 #include "rdft/rdft.h"
33 #include "dft/dft.h"
34
35 typedef struct {
36 solver super;
37 int preserve_input; /* preserve input even if DESTROY_INPUT was passed */
38 } S;
39
40 typedef struct {
41 plan_mpi_rdft2 super;
42
43 plan *cld1, *cldt, *cld2;
44 INT vn;
45 int preserve_input;
46 } P;
47
apply_r2c(const plan * ego_,R * I,R * O)48 static void apply_r2c(const plan *ego_, R *I, R *O)
49 {
50 const P *ego = (const P *) ego_;
51 plan_rdft2 *cld1;
52 plan_dft *cld2;
53 plan_rdft *cldt;
54
55 /* RDFT2 local dimensions */
56 cld1 = (plan_rdft2 *) ego->cld1;
57 if (ego->preserve_input) {
58 cld1->apply(ego->cld1, I, I+ego->vn, O, O+1);
59 I = O;
60 }
61 else
62 cld1->apply(ego->cld1, I, I+ego->vn, I, I+1);
63
64 /* global transpose */
65 cldt = (plan_rdft *) ego->cldt;
66 cldt->apply(ego->cldt, I, O);
67
68 /* DFT final local dimension */
69 cld2 = (plan_dft *) ego->cld2;
70 cld2->apply(ego->cld2, O, O+1, O, O+1);
71 }
72
apply_c2r(const plan * ego_,R * I,R * O)73 static void apply_c2r(const plan *ego_, R *I, R *O)
74 {
75 const P *ego = (const P *) ego_;
76 plan_rdft2 *cld1;
77 plan_dft *cld2;
78 plan_rdft *cldt;
79
80 /* IDFT local dimensions */
81 cld2 = (plan_dft *) ego->cld2;
82 if (ego->preserve_input) {
83 cld2->apply(ego->cld2, I+1, I, O+1, O);
84 I = O;
85 }
86 else
87 cld2->apply(ego->cld2, I+1, I, I+1, I);
88
89 /* global transpose */
90 cldt = (plan_rdft *) ego->cldt;
91 cldt->apply(ego->cldt, I, O);
92
93 /* RDFT2 final local dimension */
94 cld1 = (plan_rdft2 *) ego->cld1;
95 cld1->apply(ego->cld1, O, O+ego->vn, O, O+1);
96 }
97
applicable(const S * ego,const problem * p_,const planner * plnr)98 static int applicable(const S *ego, const problem *p_,
99 const planner *plnr)
100 {
101 const problem_mpi_rdft2 *p = (const problem_mpi_rdft2 *) p_;
102 return (1
103 && p->sz->rnk > 1
104 && (!ego->preserve_input || (!NO_DESTROY_INPUTP(plnr)
105 && p->I != p->O))
106 && ((p->flags == TRANSPOSED_OUT && p->kind == R2HC
107 && XM(is_local_after)(1, p->sz, IB)
108 && XM(is_local_after)(2, p->sz, OB)
109 && XM(num_blocks)(p->sz->dims[0].n,
110 p->sz->dims[0].b[OB]) == 1)
111 ||
112 (p->flags == TRANSPOSED_IN && p->kind == HC2R
113 && XM(is_local_after)(1, p->sz, OB)
114 && XM(is_local_after)(2, p->sz, IB)
115 && XM(num_blocks)(p->sz->dims[0].n,
116 p->sz->dims[0].b[IB]) == 1))
117 && (!NO_SLOWP(plnr) /* slow if rdft2-serial is applicable */
118 || !XM(rdft2_serial_applicable)(p))
119 );
120 }
121
awake(plan * ego_,enum wakefulness wakefulness)122 static void awake(plan *ego_, enum wakefulness wakefulness)
123 {
124 P *ego = (P *) ego_;
125 X(plan_awake)(ego->cld1, wakefulness);
126 X(plan_awake)(ego->cldt, wakefulness);
127 X(plan_awake)(ego->cld2, wakefulness);
128 }
129
destroy(plan * ego_)130 static void destroy(plan *ego_)
131 {
132 P *ego = (P *) ego_;
133 X(plan_destroy_internal)(ego->cld2);
134 X(plan_destroy_internal)(ego->cldt);
135 X(plan_destroy_internal)(ego->cld1);
136 }
137
print(const plan * ego_,printer * p)138 static void print(const plan *ego_, printer *p)
139 {
140 const P *ego = (const P *) ego_;
141 p->print(p, "(mpi-rdft2-rank-geq2-transposed%s%(%p%)%(%p%)%(%p%))",
142 ego->preserve_input==2 ?"/p":"",
143 ego->cld1, ego->cldt, ego->cld2);
144 }
145
mkplan(const solver * ego_,const problem * p_,planner * plnr)146 static plan *mkplan(const solver *ego_, const problem *p_, planner *plnr)
147 {
148 const S *ego = (const S *) ego_;
149 const problem_mpi_rdft2 *p;
150 P *pln;
151 plan *cld1 = 0, *cldt = 0, *cld2 = 0;
152 R *r0, *r1, *cr, *ci, *ri, *ii, *ro, *io, *I, *O;
153 tensor *sz;
154 int i, my_pe, n_pes;
155 INT nrest, n1, b1;
156 static const plan_adt padt = {
157 XM(rdft2_solve), awake, print, destroy
158 };
159 block_kind k1, k2;
160
161 UNUSED(ego);
162
163 if (!applicable(ego, p_, plnr))
164 return (plan *) 0;
165
166 p = (const problem_mpi_rdft2 *) p_;
167
168 I = p->I; O = p->O;
169 if (p->kind == R2HC) {
170 k1 = IB; k2 = OB;
171 r1 = (r0 = I) + p->vn;
172 if (ego->preserve_input || NO_DESTROY_INPUTP(plnr)) {
173 ci = (cr = O) + 1;
174 I = O;
175 }
176 else
177 ci = (cr = I) + 1;
178 io = ii = (ro = ri = O) + 1;
179 }
180 else {
181 k1 = OB; k2 = IB;
182 r1 = (r0 = O) + p->vn;
183 ci = (cr = O) + 1;
184 if (ego->preserve_input || NO_DESTROY_INPUTP(plnr)) {
185 ri = (ii = I) + 1;
186 ro = (io = O) + 1;
187 I = O;
188 }
189 else
190 ro = ri = (io = ii = I) + 1;
191 }
192
193 MPI_Comm_rank(p->comm, &my_pe);
194 MPI_Comm_size(p->comm, &n_pes);
195
196 sz = X(mktensor)(p->sz->rnk - 1); /* tensor of last rnk-1 dimensions */
197 i = p->sz->rnk - 2; A(i >= 0);
198 sz->dims[i].n = p->sz->dims[i+1].n / 2 + 1;
199 sz->dims[i].is = sz->dims[i].os = 2 * p->vn;
200 for (--i; i >= 0; --i) {
201 sz->dims[i].n = p->sz->dims[i+1].n;
202 sz->dims[i].is = sz->dims[i].os = sz->dims[i+1].n * sz->dims[i+1].is;
203 }
204 nrest = 1; for (i = 1; i < sz->rnk; ++i) nrest *= sz->dims[i].n;
205 {
206 INT ivs = 1 + (p->kind == HC2R), ovs = 1 + (p->kind == R2HC);
207 INT is = sz->dims[0].n * sz->dims[0].is;
208 INT b = XM(block)(p->sz->dims[0].n, p->sz->dims[0].b[k1], my_pe);
209 sz->dims[p->sz->rnk - 2].n = p->sz->dims[p->sz->rnk - 1].n;
210 cld1 = X(mkplan_d)(plnr,
211 X(mkproblem_rdft2_d)(sz,
212 X(mktensor_2d)(b, is, is,
213 p->vn,ivs,ovs),
214 r0, r1, cr, ci, p->kind));
215 if (XM(any_true)(!cld1, p->comm)) goto nada;
216 }
217
218 nrest *= p->vn;
219 n1 = p->sz->dims[1].n;
220 b1 = p->sz->dims[1].b[k2];
221 if (p->sz->rnk == 2) { /* n1 dimension is cut in ~half */
222 n1 = n1 / 2 + 1;
223 b1 = b1 == p->sz->dims[1].n ? n1 : b1;
224 }
225
226 if (p->kind == R2HC)
227 cldt = X(mkplan_d)(plnr,
228 XM(mkproblem_transpose)(
229 p->sz->dims[0].n, n1, nrest * 2,
230 I, O,
231 p->sz->dims[0].b[IB], b1,
232 p->comm, 0));
233 else
234 cldt = X(mkplan_d)(plnr,
235 XM(mkproblem_transpose)(
236 n1, p->sz->dims[0].n, nrest * 2,
237 I, O,
238 b1, p->sz->dims[0].b[OB],
239 p->comm, 0));
240 if (XM(any_true)(!cldt, p->comm)) goto nada;
241
242 {
243 INT is = p->sz->dims[0].n * nrest * 2;
244 INT b = XM(block)(n1, b1, my_pe);
245 cld2 = X(mkplan_d)(plnr,
246 X(mkproblem_dft_d)(X(mktensor_1d)(
247 p->sz->dims[0].n,
248 nrest * 2, nrest * 2),
249 X(mktensor_2d)(b, is, is,
250 nrest, 2, 2),
251 ri, ii, ro, io));
252 if (XM(any_true)(!cld2, p->comm)) goto nada;
253 }
254
255 pln = MKPLAN_MPI_RDFT2(P, &padt, p->kind == R2HC ? apply_r2c : apply_c2r);
256 pln->cld1 = cld1;
257 pln->cldt = cldt;
258 pln->cld2 = cld2;
259 pln->preserve_input = ego->preserve_input ? 2 : NO_DESTROY_INPUTP(plnr);
260 pln->vn = p->vn;
261
262 X(ops_add)(&cld1->ops, &cld2->ops, &pln->super.super.ops);
263 X(ops_add2)(&cldt->ops, &pln->super.super.ops);
264
265 return &(pln->super.super);
266
267 nada:
268 X(plan_destroy_internal)(cld2);
269 X(plan_destroy_internal)(cldt);
270 X(plan_destroy_internal)(cld1);
271 return (plan *) 0;
272 }
273
mksolver(int preserve_input)274 static solver *mksolver(int preserve_input)
275 {
276 static const solver_adt sadt = { PROBLEM_MPI_RDFT2, mkplan, 0 };
277 S *slv = MKSOLVER(S, &sadt);
278 slv->preserve_input = preserve_input;
279 return &(slv->super);
280 }
281
XM(rdft2_rank_geq2_transposed_register)282 void XM(rdft2_rank_geq2_transposed_register)(planner *p)
283 {
284 int preserve_input;
285 for (preserve_input = 0; preserve_input <= 1; ++preserve_input)
286 REGISTER_SOLVER(p, mksolver(preserve_input));
287 }
288