1 /*=========================================================================
2
3 Program: Visualization Toolkit
4 Module: TestPartitionBalancer.cxx
5
6 Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
7 All rights reserved.
8 See Copyright.txt or http://www.kitware.com/Copyright.htm for details.
9
10 This software is distributed WITHOUT ANY WARRANTY; without even
11 the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
12 PURPOSE. See the above copyright notice for more information.
13 =========================================================================*/
14
15 #include "vtkPartitionBalancer.h"
16
17 #include "vtkFieldData.h"
18 #include "vtkImageData.h"
19 #include "vtkLogger.h"
20 #include "vtkMPIController.h"
21 #include "vtkMultiProcessController.h"
22 #include "vtkNew.h"
23 #include "vtkPartitionedDataSet.h"
24 #include "vtkPartitionedDataSetCollection.h"
25 #include "vtkStringArray.h"
26
27 namespace
28 {
29 constexpr const char* names[2][4] = { { "r0_PD0_DS0", "r0_PD0_DS1", "r0_PD0_DS2", "r0_PD1_DS0" },
30 { "r1_PD0_DS0", "r2_PD0_DS1", nullptr, nullptr } };
31
32 //----------------------------------------------------------------------------
GenerateDataSet(int rank,int id)33 vtkNew<vtkImageData> GenerateDataSet(int rank, int id)
34 {
35 vtkNew<vtkStringArray> array;
36 array->SetName(names[rank][id]);
37 vtkNew<vtkImageData> ds;
38 ds->GetFieldData()->AddArray(array);
39 return ds;
40 }
41
42 //----------------------------------------------------------------------------
TestExpandPDS0(vtkPartitionedDataSet * outPDS0,int rank)43 bool TestExpandPDS0(vtkPartitionedDataSet* outPDS0, int rank)
44 {
45 bool retVal = true;
46
47 if (outPDS0->GetNumberOfPartitions() != 5)
48 {
49 vtkLog(ERROR,
50 "Wrong number of generated partitions in PD0 in rank "
51 << rank << ". There are " << outPDS0->GetNumberOfPartitions() << " instead of 5.");
52 retVal = false;
53 }
54
55 vtkDataSet* outDS0 = outPDS0->GetPartition(0);
56 vtkDataSet* outDS1 = outPDS0->GetPartition(1);
57 vtkDataSet* outDS2 = outPDS0->GetPartition(2);
58 vtkDataSet* outDS3 = outPDS0->GetPartition(3);
59 vtkDataSet* outDS4 = outPDS0->GetPartition(4);
60
61 if (rank == 0)
62 {
63 if (!outDS0 || !outDS1 || !outDS2)
64 {
65 vtkLog(ERROR,
66 "Output partitioned data set r0 - PD0 has nullptr at wrong places."
67 << " All those pointers should be non nullptr: DS0 == " << outDS0 << ", DS1 == " << outDS1
68 << ", DS2 == " << outDS2);
69 retVal = false;
70 }
71 if (retVal != EXIT_FAILURE &&
72 (!outDS0->GetFieldData()->GetAbstractArray(names[0][0]) ||
73 !outDS1->GetFieldData()->GetAbstractArray(names[0][1]) ||
74 !outDS2->GetFieldData()->GetAbstractArray(names[0][2])))
75 {
76 vtkLog(ERROR, "Output partitioned data set r0 - PD0 have been wrongly copied.");
77 retVal = false;
78 }
79 if (outDS3 || outDS4)
80 {
81 vtkLog(ERROR,
82 "Output partitioned data set r0 - PD0"
83 << " should have nullptr at partition 3 and 4");
84 retVal = false;
85 }
86 }
87 if (rank == 1)
88 {
89 if (!outDS3 || !outDS4)
90 {
91 vtkLog(ERROR,
92 "Output partitioned data set r1 - PD0 has nullptr at wrong places."
93 << " All those pointers should be non nullptr: DS3 == " << outDS3
94 << ", DS4 == " << outDS4);
95 retVal = false;
96 }
97 if (retVal != EXIT_FAILURE &&
98 (!outDS3->GetFieldData()->GetAbstractArray(names[1][0]) ||
99 !outDS4->GetFieldData()->GetAbstractArray(names[1][1])))
100 {
101 vtkLog(ERROR, "Output partitioned data set r1 - PD0 have been wrongly copied.");
102 retVal = false;
103 }
104 if (outDS0 || outDS1 || outDS2)
105 {
106 vtkLog(ERROR,
107 "Output partitioned data set r1 - PD0"
108 << " should have nullptr at partition 3 and 4");
109 retVal = false;
110 }
111 }
112
113 return retVal;
114 }
115
116 //----------------------------------------------------------------------------
TestExpandPDS1(vtkPartitionedDataSet * outPDS1,int rank)117 bool TestExpandPDS1(vtkPartitionedDataSet* outPDS1, int rank)
118 {
119 bool retVal = true;
120
121 if (outPDS1->GetNumberOfPartitions() != 1)
122 {
123 vtkLog(ERROR,
124 "Wrong number of generated partitions in PD0 in rank "
125 << rank << ". There are " << outPDS1->GetNumberOfPartitions() << " instead of 1.");
126 retVal = false;
127 }
128
129 vtkDataSet* outDS100 = outPDS1->GetPartition(0);
130
131 if (rank == 0)
132 {
133 if (!outDS100)
134 {
135 vtkLog(ERROR, "Output partitioned data set r0 - PD1 has a nullptr partition.") retVal = false;
136 }
137 else if (!outDS100->GetFieldData()->GetAbstractArray(names[0][3]))
138 {
139 vtkLog(ERROR, "DS100 has been wrongly copied in rank 0.");
140 }
141 }
142 if (rank == 1)
143 {
144 if (outDS100)
145 {
146 vtkLog(ERROR, "Output partitioned data set r1 - PD1 should have a nullptr partition.")
147 retVal = false;
148 }
149 }
150
151 return retVal;
152 }
153
154 //----------------------------------------------------------------------------
TestSquashPDS0(vtkPartitionedDataSet * outPDS0,int rank)155 bool TestSquashPDS0(vtkPartitionedDataSet* outPDS0, int rank)
156 {
157 bool retVal = true;
158
159 if (outPDS0->GetNumberOfPartitions() != 3)
160 {
161 vtkLog(ERROR,
162 "Wrong number of generated partitions in PD0 in rank "
163 << rank << ". There are " << outPDS0->GetNumberOfPartitions() << " instead of 3.");
164 retVal = false;
165 }
166
167 if (rank == 0)
168 {
169 vtkDataSet* outDS0 = outPDS0->GetPartition(0);
170 vtkDataSet* outDS1 = outPDS0->GetPartition(1);
171 vtkDataSet* outDS2 = outPDS0->GetPartition(2);
172
173 if (!outDS0 || !outDS1 || !outDS2)
174 {
175 vtkLog(ERROR,
176 "Output partitioned data set r0 - PD0 has nullptr at wrong places."
177 << " All those pointers should be non nullptr: DS0 == " << outDS0 << ", DS1 == " << outDS1
178 << ", DS2 == " << outDS2);
179 retVal = false;
180 }
181 if (retVal != false &&
182 (!outDS0->GetFieldData()->GetAbstractArray(names[0][0]) ||
183 !outDS1->GetFieldData()->GetAbstractArray(names[0][1]) ||
184 !outDS2->GetFieldData()->GetAbstractArray(names[0][2])))
185 {
186 vtkLog(ERROR, "Output partitioned data set r0 - PD0 have been wrongly copied.");
187 retVal = false;
188 }
189 }
190 if (rank == 1)
191 {
192 vtkDataSet* outDS3 = outPDS0->GetPartition(0);
193 vtkDataSet* outDS4 = outPDS0->GetPartition(1);
194 vtkDataSet* outDS5 = outPDS0->GetPartition(2);
195
196 if (!outDS3 || !outDS4)
197 {
198 vtkLog(ERROR,
199 "Output partitioned data set r1 - PD0 has nullptr at wrong places."
200 << " All those pointers should be non nullptr: DS3 == " << outDS3
201 << ", DS4 == " << outDS4);
202 retVal = false;
203 }
204 if (retVal != false &&
205 (!outDS3->GetFieldData()->GetAbstractArray(names[1][0]) ||
206 !outDS4->GetFieldData()->GetAbstractArray(names[1][1])))
207 {
208 vtkLog(ERROR, "Output partitioned data set r1 - PD0 have been wrongly copied.");
209 retVal = false;
210 }
211 if (outDS5)
212 {
213 vtkLog(ERROR,
214 "Output partitioned data set r1 - PD0"
215 << " should have nullptr at partition 2");
216 retVal = false;
217 }
218 }
219
220 return retVal;
221 }
222
223 //----------------------------------------------------------------------------
TestSquashPDS1(vtkPartitionedDataSet * outPDS1,int rank)224 bool TestSquashPDS1(vtkPartitionedDataSet* outPDS1, int rank)
225 {
226 bool retVal = true;
227
228 if (outPDS1->GetNumberOfPartitions() != 1)
229 {
230 vtkLog(ERROR,
231 "Wrong number of generated partitions in PD1 in rank "
232 << rank << ". There are " << outPDS1->GetNumberOfPartitions() << " instead of 1.");
233 retVal = false;
234 }
235
236 vtkDataSet* outDS100 = outPDS1->GetPartition(0);
237
238 if (rank == 0)
239 {
240 if (!outDS100)
241 {
242 vtkLog(ERROR, "Output partitioned data set r0 - PD1 has a nullptr partition.") retVal = false;
243 }
244 else if (!outDS100->GetFieldData()->GetAbstractArray(names[0][3]))
245 {
246 vtkLog(ERROR, "DS100 has been wrongly copied in rank 0.");
247 }
248 }
249 if (rank == 1)
250 {
251 if (outDS100)
252 {
253 vtkLog(ERROR, "Output partitioned data set r1 - PD1 should have a nullptr partition.")
254 retVal = false;
255 }
256 }
257
258 return retVal;
259 }
260 } // anonmymous namespace
261
262 //----------------------------------------------------------------------------
TestPartitionBalancer(int argc,char * argv[])263 int TestPartitionBalancer(int argc, char* argv[])
264 {
265 int retVal = EXIT_SUCCESS;
266
267 vtkNew<vtkMPIController> controller;
268
269 controller->Initialize(&argc, &argv);
270 vtkMultiProcessController::SetGlobalController(controller);
271
272 int rank = controller->GetLocalProcessId();
273
274 vtkNew<vtkPartitionedDataSetCollection> pdsc;
275 pdsc->SetNumberOfPartitionedDataSets(2);
276
277 // rank 0: PDC [ PD (DS0, DS1, DS2) ] [PD (nullptr, DS100) ]
278 // rank 1: PDC [ PD (DS3, nullptr, DS4) ] [PD () ]
279
280 if (rank == 0)
281 {
282 vtkNew<vtkPartitionedDataSet> pds0;
283 pds0->SetNumberOfPartitions(3);
284 pds0->SetPartition(0, GenerateDataSet(0, 0));
285 pds0->SetPartition(1, GenerateDataSet(0, 1));
286 pds0->SetPartition(2, GenerateDataSet(0, 2));
287 pdsc->SetPartitionedDataSet(0, pds0);
288
289 vtkNew<vtkPartitionedDataSet> pds1;
290 pds1->SetNumberOfPartitions(2);
291 pds1->SetPartition(0, nullptr);
292 pds1->SetPartition(0, GenerateDataSet(0, 3));
293 pdsc->SetPartitionedDataSet(1, pds1);
294 }
295 else if (rank == 1)
296 {
297 vtkNew<vtkPartitionedDataSet> pds0;
298 pds0->SetNumberOfPartitions(3);
299 pds0->SetPartition(0, GenerateDataSet(1, 0));
300 pds0->SetPartition(1, nullptr);
301 pds0->SetPartition(2, GenerateDataSet(1, 1));
302 pdsc->SetPartitionedDataSet(0, pds0);
303
304 pdsc->SetPartitionedDataSet(1, vtkNew<vtkPartitionedDataSet>());
305 }
306
307 vtkNew<vtkPartitionBalancer> balancer;
308 balancer->SetInputDataObject(pdsc);
309 balancer->SetController(controller);
310
311 if (rank == 0)
312 {
313 vtkLog(INFO, "Testing vtkPartitionBalancer for vtkPartitionedDataSetCollection input");
314 }
315
316 if (rank == 0)
317 {
318 vtkLog(INFO, "*** Expand mode");
319 }
320
321 balancer->SetModeToExpand();
322 balancer->Update();
323
324 {
325 // Looking if output looks like that:
326 // rank 0: PDC [ PD (DS0, DS1, DS2, nullptr, nullptr) ] [PD (DS100) ]
327 // rank 1: PDC [ PD (nullptr, nullptr, nullptr, DS3, DS4) ] [PD (nullptr) ]
328 auto outPDSC = vtkPartitionedDataSetCollection::SafeDownCast(balancer->GetOutputDataObject(0));
329
330 if (outPDSC->GetNumberOfPartitionedDataSets() != 2)
331 {
332 vtkLog(ERROR,
333 "Wrong number of generated partitioned data sets in rank "
334 << rank << ". There are " << outPDSC->GetNumberOfPartitionedDataSets()
335 << " instead of 2");
336 retVal = EXIT_FAILURE;
337 }
338
339 retVal = !TestExpandPDS0(outPDSC->GetPartitionedDataSet(0), rank) ||
340 !TestExpandPDS1(outPDSC->GetPartitionedDataSet(1), rank)
341 ? EXIT_FAILURE
342 : retVal;
343 }
344
345 if (rank == 0)
346 {
347 vtkLog(INFO, "*** Squash mode");
348 }
349
350 balancer->SetModeToSquash();
351 balancer->Update();
352
353 {
354 // Looking if output looks like that:
355 // rank 0: PDC [ PD (DS0, DS1, DS2) ] [PD (DS100) ]
356 // rank 1: PDC [ PD (DS3, DS4, nullptr) ] [PD (nullptr) ]
357 auto outPDSC = vtkPartitionedDataSetCollection::SafeDownCast(balancer->GetOutputDataObject(0));
358
359 if (outPDSC->GetNumberOfPartitionedDataSets() != 2)
360 {
361 vtkLog(ERROR,
362 "Wrong number of generated partitioned data sets in rank "
363 << rank << ". There are " << outPDSC->GetNumberOfPartitionedDataSets()
364 << " instead of 2");
365 retVal = EXIT_FAILURE;
366 }
367
368 retVal = !TestSquashPDS0(outPDSC->GetPartitionedDataSet(0), rank) ||
369 !TestSquashPDS1(outPDSC->GetPartitionedDataSet(1), rank)
370 ? EXIT_FAILURE
371 : retVal;
372 }
373
374 if (rank == 0)
375 {
376 vtkLog(INFO, "Testing vtkPartitionBalancer for vtkPartitionedDataSet input");
377 }
378
379 // Here, we will test the same input as before, but directly feed in partitioned data sets.
380
381 if (rank == 0)
382 {
383 vtkLog(INFO, "*** Expand mode");
384 }
385
386 balancer->SetModeToExpand();
387 balancer->SetInputDataObject(pdsc->GetPartitionedDataSet(0));
388 balancer->Update();
389
390 retVal =
391 !TestExpandPDS0(vtkPartitionedDataSet::SafeDownCast(balancer->GetOutputDataObject(0)), rank)
392 ? EXIT_FAILURE
393 : retVal;
394
395 balancer->SetInputDataObject(pdsc->GetPartitionedDataSet(1));
396 balancer->Update();
397
398 retVal =
399 !TestExpandPDS1(vtkPartitionedDataSet::SafeDownCast(balancer->GetOutputDataObject(0)), rank)
400 ? EXIT_FAILURE
401 : retVal;
402
403 if (rank == 0)
404 {
405 vtkLog(INFO, "*** Squash mode");
406 }
407
408 balancer->SetModeToSquash();
409 balancer->SetInputDataObject(pdsc->GetPartitionedDataSet(0));
410 balancer->Update();
411
412 retVal =
413 !TestSquashPDS0(vtkPartitionedDataSet::SafeDownCast(balancer->GetOutputDataObject(0)), rank)
414 ? EXIT_FAILURE
415 : retVal;
416
417 balancer->SetInputDataObject(pdsc->GetPartitionedDataSet(1));
418 balancer->Update();
419
420 retVal =
421 !TestSquashPDS1(vtkPartitionedDataSet::SafeDownCast(balancer->GetOutputDataObject(0)), rank)
422 ? EXIT_FAILURE
423 : retVal;
424
425 controller->Finalize();
426
427 return retVal;
428 }
429