ParMetisPartitioner.cc 15.6 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
#include "ParMetisPartitioner.h"
#include "Mesh.h"
#include "Traverse.h"
#include "ElInfo.h"
#include "Element.h"
#include "FixVec.h"
#include "PartitionElementData.h"
#include "DOFVector.h"
#include "mpi.h"

#include <queue>

namespace AMDiS {

15
  ParMetisMesh::ParMetisMesh(Mesh *mesh, MPI::Comm *comm)
16
    : nElements(0),
17
      mpiComm(comm)
18
19
  {
    FUNCNAME("ParMetisMesh::ParMetisMesh()");
20
21

    int mpiSize = mpiComm->Get_size();
22
23
24
25
26
27
28
29
30
31

    int nodeCounter = 0;
    int elementCounter = 0;

    dim_ = mesh->getDim();
    int dow = Global::getGeo(WORLD);

    TraverseStack stack;
    ElInfo *elInfo = stack.traverseFirst(mesh, -1, 
					 Mesh::CALL_EVERY_EL_PREORDER);
32
    while (elInfo) {
33
34
35
36
37
38
      Element *element = elInfo->getElement();

      // get partition data
      PartitionElementData *partitionData = dynamic_cast<PartitionElementData*>
	(element->getElementData(PARTITION_ED));

39
40
41
      if (partitionData && 
	  partitionData->getPartitionStatus() == IN &&
	  partitionData->getLevel() == 0) {
42
43
44
45
46
47
	elementCounter++;
      }

      elInfo = stack.traverseNext(elInfo);
    }

48
    nElements = elementCounter;
49

50
    TEST_EXIT(nElements > 0)("no elements in ParMETIS mesh\n");
51
52

    // allocate memory
53
54
55
    eptr_ = GET_MEMORY(int, nElements + 1);
    eind_ = GET_MEMORY(int, nElements * (dim_ + 1));
    elmdist = GET_MEMORY(int, mpiSize + 1);
56

57
    elem_p2a_ = GET_MEMORY(int, nElements);
58

59
    if (dim_ == dow) {
60
      xyz_ = GET_MEMORY(float, nElements * dim_);
61
62
63
64
65
66
67
68
69
70
71
    } else {
      xyz_ = NULL;
    }

    eptr_[0] = 0;

    int *ptr_eptr = eptr_ + 1;
    int *ptr_eind = eind_;
    float *ptr_xyz = xyz_;
    
    // gather element numbers and create elmdist
72
    mpiComm->Allgather(&nElements, 1, MPI_INT, elmdist + 1, 1, MPI_INT);
73

74
    elmdist[0] = 0;
75
    for (int i = 2; i < mpiSize + 1; i++) {
76
      elmdist[i] += elmdist[i - 1];
77
78
79
80
81
82
83
84
85
86
87
    }

    // traverse mesh and fill distributed ParMETIS data
    DimVec<double> bary(dim_, DEFAULT_VALUE, 1.0 / (dim_ + 1));
    WorldVector<double> world;

    elementCounter = 0;

    elInfo = stack.traverseFirst(mesh, -1, 
				 Mesh::CALL_EVERY_EL_PREORDER | 
				 Mesh::FILL_COORDS);
88
    while (elInfo) {
89
90
91
92
93
94
95
96
      Element *element = elInfo->getElement();
      int index = element->getIndex();

      // get partition data
      PartitionElementData *partitionData = dynamic_cast<PartitionElementData*>
	(element->getElementData(PARTITION_ED));

      // if element in partition
97
98
99
      if (partitionData && 
	  partitionData->getPartitionStatus() == IN &&
	  partitionData->getLevel() == 0) {
100
101
102
103
104
105
106
107
108
109
	// remember index
	setParMetisIndex(index, elementCounter);
	setAMDiSIndex(elementCounter, index);

	// write eptr entry
	nodeCounter += dim_ + 1;
	*ptr_eptr = nodeCounter;
	ptr_eptr++;

	// write eind entries (element nodes)
110
	for (int i = 0; i < dim_ + 1; i++) {
111
112
113
114
115
	  *ptr_eind = element->getDOF(i, 0);
	  ptr_eind++;
	}

	// write xyz element coordinates
116
	if (ptr_xyz) {
117
	  elInfo->coordToWorld(bary, world);
118
	  for (int i = 0; i < dim_; i++) {
119
120
121
122
123
124
125
126
127
128
129
130
131
	    *ptr_xyz = static_cast<float>(world[i]); 
	    ptr_xyz++;
	  }
	}

	elementCounter++;
      }
      elInfo = stack.traverseNext(elInfo);
    }
  }

  ParMetisMesh::~ParMetisMesh()
  {
132
    if (eptr_)
133
      FREE_MEMORY(eptr_, int, nElements + 1);
134
135
    
    if (eind_)     
136
      FREE_MEMORY(eind_, int, nElements * (dim_ + 1));
137
    
138
139
    if (elmdist)
      FREE_MEMORY(elmdist, int, mpiComm->Get_size() + 1);
140
141
    
    if (xyz_)
142
      FREE_MEMORY(xyz_, float, nElements * dim_);
143
144
    
    if (elem_p2a_) 
145
      FREE_MEMORY(elem_p2a_, int, nElements);
146
147
148
  }

  ParMetisGraph::ParMetisGraph(ParMetisMesh *parMetisMesh,
149
			       MPI::Comm *comm,
150
151
152
153
154
			       int ncommonnodes)
    : parMetisMesh_(parMetisMesh)
  {
    int numflag = 0;

155
156
    if (ncommonnodes == -1) 
      ncommonnodes = parMetisMesh->getDim();
157

158
    MPI_Comm tmpComm = MPI_Comm(*comm);
159
160
161
162
163
164
165
166

    ParMETIS_V3_Mesh2Dual(parMetisMesh_->getElementDist(),
			  parMetisMesh_->getElementPtr(),
			  parMetisMesh_->getElementInd(),
			  &numflag,
			  &ncommonnodes,
			  &xadj_,
			  &adjncy_,
167
			  &tmpComm);
168
169
170
171
172
173
174
175
  }

  ParMetisGraph::~ParMetisGraph()
  {
    free(xadj_);
    free(adjncy_);
  }

176

177
178
179
180
181
  void ParMetisPartitioner::deletePartitionData() 
  {
    TraverseStack stack;
    ElInfo *elInfo;
    elInfo = stack.traverseFirst(mesh_, -1, Mesh::CALL_EVERY_EL_PREORDER);
182
    while (elInfo) {
183
184
185
186
187
188
189
      Element *element = elInfo->getElement();
      element->deleteElementData(PARTITION_ED);
      elInfo = stack.traverseNext(elInfo);
    }
  }

  void ParMetisPartitioner::createPartitionData() {
190
191
    int mpiRank = mpiComm->Get_rank();
    int mpiSize = mpiComm->Get_size();
192
193
194
195
196
197
198

    TraverseStack stack;
    ElInfo *elInfo;

    // === create initial partitioning on AMDiS mesh ===
    int totalElementCounter = 0;
    elInfo = stack.traverseFirst(mesh_, -1, Mesh::CALL_LEAF_EL);
199
    while (elInfo) {
200
201
202
203
204
205
206
207
208
      Element *element = elInfo->getElement();

      TEST_EXIT(element->getElementData(PARTITION_ED) == NULL)
	("mesh already partitioned\n");

      PartitionElementData *elData = 
	NEW PartitionElementData(element->getElementData());
      element->setElementData(elData);

209
      if (totalElementCounter % mpiSize == mpiRank) {
210
211
212
213
214
215
216
217
218
219
220
221
222
223
	elData->setPartitionStatus(IN);
      } else {
	elData->setPartitionStatus(UNDEFINED);
      }
      totalElementCounter++;

      elInfo = stack.traverseNext(elInfo);
    }
  }

  void ParMetisPartitioner::partition(std::map<int, double> *elemWeights,
				      PartitionMode mode,
				      float itr) 
  {
224
    int mpiSize = mpiComm->Get_size();
225
226
227
228
229

    TraverseStack stack;
    ElInfo *elInfo;

    // === create parmetis mesh ===
230
231
232
    if (parMetisMesh_) 
      DELETE parMetisMesh_;
    parMetisMesh_ = NEW ParMetisMesh(mesh_, mpiComm);
233
234
235
236
237
238
239
240
241
242
243

    int numElements = parMetisMesh_->getNumElements();

    // === create weight array ===
    int *wgts = elemWeights ? GET_MEMORY(int, numElements) : NULL;
    float *floatWgts = elemWeights ? GET_MEMORY(float, numElements) : NULL;
    float maxWgt = 0.0;

    float *ptr_floatWgts = floatWgts;

    elInfo = stack.traverseFirst(mesh_, -1, Mesh::CALL_EVERY_EL_PREORDER);
244
    while (elInfo) {
245
246
247
248
249
250
      Element *element = elInfo->getElement();

      // get partition data
      PartitionElementData *partitionData = dynamic_cast<PartitionElementData*>
	(element->getElementData(PARTITION_ED));

251
252
253
      if (partitionData &&
	  partitionData->getPartitionStatus() == IN &&
	  partitionData->getLevel() == 0) {
254
255
256
257
258
259
260
261
262
263
264
265
266
267
	int index = element->getIndex();

	// get weight 
	float wgt = static_cast<float>((*elemWeights)[index]);
	maxWgt = max(wgt, maxWgt);

	// write float weight
	*ptr_floatWgts = wgt;
	ptr_floatWgts++;
      }
      elInfo = stack.traverseNext(elInfo);
    }

    float tmp;
268
    mpiComm->Allreduce(&maxWgt, &tmp, 1, MPI_FLOAT, MPI_MAX);
269
270
271
    maxWgt = tmp;

    // === create dual graph ===
272
    ParMetisGraph parMetisGraph(parMetisMesh_, mpiComm);
273
274
275
276
277
278
279
280
281
282
283
284

    // === partitioning of dual graph ===
    int wgtflag = elemWeights ? 2 : 0; // weights at vertices only!
    int numflag = 0; // c numbering style!
    int ncon = elemWeights ? 1 : 0; // one weight at each vertex!
    int nparts = mpiSize; // number of partitions
    float *tpwgts = elemWeights ? GET_MEMORY(float, mpiSize) : NULL;
    float ubvec = 1.05;
    int options[3] = {0, 0, 0}; // default options
    int edgecut = -1;
    int *part = GET_MEMORY(int, numElements);
    
285
286
    if (elemWeights) {
      for (int i = 0; i < mpiSize; i++) {
287
	// set tpwgts
288
	tpwgts[i] = 1.0 / nparts;
289
290
291
      }      

      float scale = 10000 / maxWgt;
292
      for (int i = 0; i < numElements; i++) {
293
294
295
296
297
	// scale wgts
	wgts[i] = static_cast<int>(floatWgts[i] * scale);
      }
    }

298
299
    MPI_Comm tmpComm = MPI_Comm(*mpiComm);

300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
    switch(mode) {
    case INITIAL:
      ParMETIS_V3_PartKway(parMetisMesh_->getElementDist(),
			   parMetisGraph.getXAdj(),
			   parMetisGraph.getAdjncy(),
			   wgts,
			   NULL,
			   &wgtflag,
			   &numflag,
			   &ncon,
			   &nparts,
			   tpwgts,
			   &ubvec,
			   options,
			   &edgecut,
			   part,
316
			   &tmpComm);
317
318
319
320
      break;
    case ADAPTIVE_REPART:
      {
	int *vsize = GET_MEMORY(int, numElements);
321
	for (int i = 0; i < numElements; i++) {
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
	  vsize[i] = 1;
	}
	ParMETIS_V3_AdaptiveRepart(parMetisMesh_->getElementDist(),
				   parMetisGraph.getXAdj(),
				   parMetisGraph.getAdjncy(),
				   wgts,
				   NULL,
				   vsize,
				   &wgtflag,
				   &numflag,
				   &ncon,
				   &nparts,
				   tpwgts,
				   &ubvec,
				   &itr,
				   options,
				   &edgecut,
				   part,
340
				   &tmpComm);
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
	FREE_MEMORY(vsize, int, numElements);
      }
      break;
    case REFINE_PART:
      ParMETIS_V3_RefineKway(parMetisMesh_->getElementDist(),
			     parMetisGraph.getXAdj(),
			     parMetisGraph.getAdjncy(),
			     wgts,
			     NULL,
			     &wgtflag,
			     &numflag,
			     &ncon,
			     &nparts,
			     tpwgts,
			     &ubvec,
			     options,
			     &edgecut,
			     part,
359
			     &tmpComm);
360
361
362
363
364
365
366
367
      break;
    default: 
      ERROR_EXIT("unknown partitioning mode\n");
    }

    // === distribute new partition data ===
    distributePartitioning(part);

368
369
370
371
372
373
    if (floatWgts) 
      FREE_MEMORY(floatWgts, float, numElements);
    if (wgts) 
      FREE_MEMORY(wgts, int, numElements);    
    if (tpwgts) 
      FREE_MEMORY(tpwgts, float, mpiSize);
374
375
376
377
378
379
380
381
382
383
    FREE_MEMORY(part, int, numElements);
  }

  void ParMetisPartitioner::fillCoarsePartitionVec(std::map<int, int> *partitionVec)
  {
    TEST_EXIT(partitionVec)("no partition vector\n");

    partitionVec->clear();

    // update ParMETIS mesh to new partitioning
384
385
    if (!parMetisMesh_) 
      parMetisMesh_ = NEW ParMetisMesh(mesh_, mpiComm);
386

387
388
    int mpiRank = mpiComm->Get_rank();
    int mpiSize = mpiComm->Get_size();
389
390
391
392

    int *numPartitionElements = GET_MEMORY(int, mpiSize);

    int *elmdist = parMetisMesh_->getElementDist();
393
    for (int i = 0;  i < mpiSize; i++) {
394
395
396
397
398
399
      numPartitionElements[i] = elmdist[i+1] - elmdist[i];
    }

    // === count number of elements ===
    int numElements = 0;
    int localElements = parMetisMesh_->getNumElements();
400
    mpiComm->Allreduce(&localElements, &numElements, 1, MPI_INT, MPI_SUM);
401
402
403
404

    int *partitionElements = GET_MEMORY(int, numElements);

    // distribute partition elements
405
406
407
408
409
410
411
    mpiComm->Allgatherv(parMetisMesh_->getAMDiSIndices(),
			numPartitionElements[mpiRank], 
			MPI_INT, 
			partitionElements, 
			numPartitionElements, 
			elmdist, 
			MPI_INT);
412
413

    // fill partitionVec
414
415
    for (int i = 0; i < mpiSize; i++) {
      for (int j = 0; j < numPartitionElements[i]; j++) {
416
417
418
419
420
421
422
423
424
425
	(*partitionVec)[partitionElements[elmdist[i] + j]] = i;
      }
    }

    FREE_MEMORY(partitionElements, int, numElements);
    FREE_MEMORY(numPartitionElements, int, mpiSize);
  }

  void ParMetisPartitioner::distributePartitioning(int *part) 
  {
426
427
    int mpiSize = mpiComm->Get_size();
    int mpiRank = mpiComm->Get_rank();
428
429
430
431
    int numElements = parMetisMesh_->getNumElements();

    // count elements per partition in this rank
    int *numPartitionElements = GET_MEMORY(int, mpiSize);
432
    for (int i = 0; i < mpiSize; i++) 
433
      numPartitionElements[i] = 0;
434
    for (int i = 0; i < numElements; i++) {
435
436
437
438
439
      numPartitionElements[part[i]]++;
    }

    // collect number of partition elements from all ranks for this rank
    int *numRankElements = GET_MEMORY(int, mpiSize);
440
441
    mpiComm->Alltoall(numPartitionElements, 1, MPI_INT,
		      numRankElements, 1, MPI_INT);
442
443
444
445

    // sum up partition elements over all ranks
    int *sumPartitionElements = GET_MEMORY(int, mpiSize);

446
447
    mpiComm->Allreduce(numPartitionElements, sumPartitionElements, mpiSize,
		       MPI_INT, MPI_SUM);
448
449
450
451
452

    
    // prepare distribution (fill partitionElements with AMDiS indices)
    int *bufferOffset = GET_MEMORY(int, mpiSize);
    bufferOffset[0] = 0;
453
    for (int i = 1; i < mpiSize; i++) {
454
455
456
457
458
459
      bufferOffset[i] = bufferOffset[i - 1] + numPartitionElements[i - 1];
    }

    int *partitionElements = GET_MEMORY(int, numElements);
    int **partitionPtr = GET_MEMORY(int*, mpiSize);

460
    for (int i = 0; i < mpiSize; i++) {
461
462
463
      partitionPtr[i] = partitionElements + bufferOffset[i];
    }

464
    for (int i = 0; i < numElements; i++) {
465
466
467
468
469
470
471
472
473
474
      int partition = part[i];
      int amdisIndex = parMetisMesh_->getAMDiSIndex(i);
      *(partitionPtr[partition]) = amdisIndex;
      ++(partitionPtr[partition]);
    }

    // all to all: partition elements to rank elements
    int *rankElements = GET_MEMORY(int, sumPartitionElements[mpiRank]);
    int *recvBufferOffset = GET_MEMORY(int, mpiSize);
    recvBufferOffset[0] = 0;
475
    for (int i = 1; i < mpiSize; i++) {
476
477
478
      recvBufferOffset[i] = recvBufferOffset[i - 1] + numRankElements[i - 1];
    }

479
480
481
482
483
484
485
486
    mpiComm->Alltoallv(partitionElements, 
		       numPartitionElements,
		       bufferOffset,
		       MPI_INT,
		       rankElements,
		       numRankElements,
		       recvBufferOffset,
		       MPI_INT);
487
488
489
490
491
492
    

    // === partition AMDiS mesh ===

    // write data in stl map
    std::map<int, bool> elementInPartition;
493
    for (int i = 0; i < mpiSize; i++) {
494
495
      int *rankStart = rankElements + recvBufferOffset[i];
      int *rankEnd = rankStart + numRankElements[i];
496
      for (int *rankPtr = rankStart; rankPtr < rankEnd; ++rankPtr) {
497
498
499
500
501
502
	elementInPartition[*rankPtr] = true;
      }
    }

    TraverseStack stack;
    ElInfo *elInfo = stack.traverseFirst(mesh_, -1, Mesh::CALL_EVERY_EL_PREORDER);
503
    while (elInfo) {
504
505
506
507
508
509
      Element *element = elInfo->getElement();

      // get partition data
      PartitionElementData *partitionData = dynamic_cast<PartitionElementData*>
	(element->getElementData(PARTITION_ED));

510
      if (partitionData && partitionData->getLevel() == 0) {
511
	int amdisIndex = element->getIndex();
512
	if (elementInPartition[amdisIndex]) {
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
	  partitionData->setPartitionStatus(IN);
	} else {
	  partitionData->setPartitionStatus(OUT);
	}
	descendPartitionData(element);
      }

      elInfo = stack.traverseNext(elInfo);
    }

    DELETE parMetisMesh_;
    parMetisMesh_ = NULL;

    FREE_MEMORY(rankElements, int, sumPartitionElements[mpiRank]);
    FREE_MEMORY(numPartitionElements, int, mpiSize);
    FREE_MEMORY(numRankElements, int, mpiSize);
    FREE_MEMORY(sumPartitionElements, int, mpiSize);
    FREE_MEMORY(partitionElements, int, numElements);
    FREE_MEMORY(partitionPtr, int*, mpiSize);
    FREE_MEMORY(bufferOffset, int, mpiSize);
    FREE_MEMORY(recvBufferOffset, int, mpiSize);
  }

  void ParMetisPartitioner::descendPartitionData(Element *element) 
  {
538
    if (!element->isLeaf()) {
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
      Element *child0 = element->getChild(0);
      Element *child1 = element->getChild(1);

      // get partition data
      PartitionElementData *parentData = dynamic_cast<PartitionElementData*>
	(element->getElementData(PARTITION_ED));
      PartitionElementData *child0Data = dynamic_cast<PartitionElementData*>
	(child0->getElementData(PARTITION_ED));
      PartitionElementData *child1Data = dynamic_cast<PartitionElementData*>
	(child1->getElementData(PARTITION_ED));
      
      TEST_EXIT(parentData && child0Data && child1Data)("no partition data\n");

      child0Data->setPartitionStatus(parentData->getPartitionStatus());
      child1Data->setPartitionStatus(parentData->getPartitionStatus());

      descendPartitionData(child0);
      descendPartitionData(child1);
    }
  }


  void ParMetisPartitioner::fillLeafPartitionVec(std::map<int, int> *coarseVec,
						 std::map<int, int> *fineVec)
  {
    int partition = -1;
    TraverseStack stack;
    ElInfo *elInfo = stack.traverseFirst(mesh_, -1, Mesh::CALL_EVERY_EL_PREORDER);
567
    while (elInfo) {
568
569
570
      Element *element = elInfo->getElement();
      PartitionElementData *partitionData = dynamic_cast<PartitionElementData*>
	(element->getElementData(PARTITION_ED));
571
572
      if (partitionData) {
	if (partitionData->getLevel() == 0) {
573
574
575
576
577
578
579
580
581
582
	  partition = (*(coarseVec))[element->getIndex()];
	}
	if(element->isLeaf()) {
	  (*(fineVec))[element->getIndex()] = partition;
	}
      }
      elInfo = stack.traverseNext(elInfo);
    }
  }
}