ZoltanLB which uses zoltan hyper-graph partitioner. This is a multicast aware load...
[charm.git] / src / ck-ldb / ZoltanLB.C
1 /** \file ZoltanLB.C
2  *
3  *  Updated by Abhinav Bhatele, 2010-11-26 to use ckgraph
4  */
5
6 /**
7  * \addtogroup CkLdb
8 */
9 /*@{*/
10
11
12 #include "ZoltanLB.h"
13 #include "mpi.h"
14 #include "ckgraph.h"
15 #include "zoltan.h"
16
17 typedef struct {
18   int numMyVertices;
19   // Ids of the vertices (chares)
20   ZOLTAN_ID_TYPE *vtxGID;
21   // Weight/load of each chare
22   int *vtxWgt;
23
24   int numMyHEdges;
25   // Ids of the edges
26   ZOLTAN_ID_TYPE *edgeGID;
27   int *edgWgt;
28
29   // For compressed hyperedge storage. nborIndex indexes into nborGID.
30   int *nborIndex;
31   // nborGID contains ids of the chares that constitute an edge
32   ZOLTAN_ID_TYPE *nborGID;
33
34   int numAllNbors;
35
36   ProcArray *parr;
37   ObjGraph *ogr;
38 } HGRAPH_DATA;
39
40 static int get_number_of_vertices(void *data, int *ierr);
41 static void get_vertex_list(void *data, int sizeGID, int sizeLID,
42             ZOLTAN_ID_PTR globalID, ZOLTAN_ID_PTR localID,
43                   int wgt_dim, float *obj_wgts, int *ierr);
44 static void get_hypergraph_size(void *data, int *num_lists, int *num_nonzeroes,
45                                 int *format, int *ierr);
46 static void get_hypergraph(void *data, int sizeGID, int num_edges, int num_nonzeroes,
47                            int format, ZOLTAN_ID_PTR edgeGID, int *vtxPtr,
48                            ZOLTAN_ID_PTR vtxGID, int *ierr);
49 static void get_hypergraph_edge_size(void *data, int *num_edges, int *ierr);
50
51 static void get_hypergraph_edge_wgts(void *data, int numGID, int numLID, int num_edges,
52                                      int edge_weight_dim, ZOLTAN_ID_PTR edgeGID, ZOLTAN_ID_PTR edgeLID,
53                                      float *edge_wgts, int *ierr);
54
55 CreateLBFunc_Def(ZoltanLB, "Use Zoltan(tm) to partition object graph")
56
57 ZoltanLB::ZoltanLB(const CkLBOptions &opt): CentralLB(opt)
58 {
59   lbname = "ZoltanLB";
60   if (CkMyPe() == 0)
61     CkPrintf("[%d] ZoltanLB created\n",CkMyPe());
62 }
63
64 void ZoltanLB::work(LDStats* stats)
65 {
66   /** ========================== INITIALIZATION ============================= */
67   ProcArray *parr = new ProcArray(stats);
68   ObjGraph *ogr = new ObjGraph(stats);
69
70   /** ============================= STRATEGY ================================ */
71   if (_lb_args.debug() >= 2) {
72     CkPrintf("[%d] In ZoltanLB Strategy...\n", CkMyPe());
73   }
74
75   int rc;
76   float ver;
77   struct Zoltan_Struct *zz;
78   int changes, numGidEntries, numLidEntries, numImport, numExport;
79   int myRank, numProcs;
80   ZOLTAN_ID_PTR importGlobalGids, importLocalGids, exportGlobalGids, exportLocalGids;
81   int *importProcs, *importToPart, *exportProcs, *exportToPart;
82   int *parts;
83
84   HGRAPH_DATA hg;
85
86   hg.parr = parr;
87   hg.ogr = ogr;
88
89   int numPes = parr->procs.size();
90   // convert ObjGraph to the adjacency structure
91   int numVertices = ogr->vertices.size();       // number of vertices
92   int numEdges = 0;                             // number of edges
93   int numAllNbors = 0;
94   hg.numMyVertices = numVertices;
95   hg.vtxGID = (ZOLTAN_ID_TYPE *)malloc(sizeof(ZOLTAN_ID_TYPE) * numVertices);
96   hg.vtxWgt = (int *)malloc(sizeof(int) * numVertices);
97
98   double maxLoad = 0.0;
99   double maxBytes = 0.0;
100   int i, j, k, vert;
101
102   /** remove duplicate edges from recvFrom */
103   for(i = 0; i < numVertices; i++) {
104     hg.vtxGID[i] = (ZOLTAN_ID_TYPE) i;
105
106     for(j = 0; j < ogr->vertices[i].sendToList.size(); j++) {
107       vert = ogr->vertices[i].sendToList[j].getNeighborId();
108       for(k = 0; k < ogr->vertices[i].recvFromList.size(); k++) {
109         if(ogr->vertices[i].recvFromList[k].getNeighborId() == vert) {
110           ogr->vertices[i].sendToList[j].setNumBytes(ogr->vertices[i].sendToList[j].getNumBytes() + ogr->vertices[i].recvFromList[k].getNumBytes());
111           ogr->vertices[i].recvFromList.erase(ogr->vertices[i].recvFromList.begin() + k);
112         }
113       }
114     }
115   }
116
117   /** the object load is normalized to an integer between 0 and 256 */
118   for(i = 0; i < numVertices; i++) {
119     if(ogr->vertices[i].getVertexLoad() > maxLoad)
120       maxLoad = ogr->vertices[i].getVertexLoad();
121     numEdges += ogr->vertices[i].sendToList.size();
122     numEdges += ogr->vertices[i].mcastToList.size();
123
124     numAllNbors += 2 * ogr->vertices[i].sendToList.size(); // For each point to point edge, add 2 numAllNbors 
125     for (k = 0; k < ogr->vertices[i].mcastToList.size(); k++) {
126       McastSrc mcast_src = ogr->vertices[i].mcastToList[k];
127       numAllNbors += mcast_src.destList.size();
128       numAllNbors++;
129     }
130   }
131
132   for(i = 0; i < numVertices; i++) {
133     for(j = 0; j < ogr->vertices[i].sendToList.size(); j++) {
134       if (ogr->vertices[i].sendToList[j].getNumBytes() > maxBytes) {
135         maxBytes = ogr->vertices[i].sendToList[j].getNumBytes();
136       }
137     }
138   }
139
140   hg.numMyHEdges = numEdges;
141   hg.numAllNbors = numAllNbors;
142   hg.edgeGID = (ZOLTAN_ID_TYPE *)malloc(sizeof(ZOLTAN_ID_TYPE) * numEdges);
143   hg.edgWgt = (int *)malloc(sizeof(int) * numEdges);
144   hg.nborIndex = (int *)malloc(sizeof(int) * numEdges);
145   hg.nborGID = (ZOLTAN_ID_TYPE *)malloc(sizeof(ZOLTAN_ID_TYPE) * numAllNbors);
146
147   double ratio = 256.0/maxLoad;
148   double byteRatio = 1024.0 / maxBytes;
149   int edgeNum = 0;
150   int index = 0;
151
152   for (i = 0; i < numVertices; i++) {
153     hg.vtxWgt[i] = (int)ceil(ogr->vertices[i].getVertexLoad() * ratio);
154
155     for (j = 0; j < ogr->vertices[i].sendToList.size(); j++) {
156         hg.edgeGID[edgeNum] = edgeNum;
157         hg.edgWgt[edgeNum] = (int) ceil(ogr->vertices[i].sendToList[j].getNumBytes() * byteRatio);
158         hg.nborIndex[edgeNum++] = index;
159         hg.nborGID[index++] = i; 
160         hg.nborGID[index++] = ogr->vertices[i].sendToList[j].getNeighborId(); 
161     }
162
163     for (j = 0; j < ogr->vertices[i].mcastToList.size(); j++) {
164       hg.edgeGID[edgeNum] = edgeNum;
165       hg.edgWgt[edgeNum] = (int) ceil(ogr->vertices[i].mcastToList[j].getNumBytes() * byteRatio);
166       hg.nborIndex[edgeNum++] = index;
167       McastSrc mcast_src = ogr->vertices[i].mcastToList[j];
168
169       hg.nborGID[index++] = i; 
170       for (k = 0; k < mcast_src.destList.size(); k++) {
171         // For all the vertices in the multicast edge, add it
172         hg.nborGID[index++] = mcast_src.destList[k]; 
173       }
174     }
175   }
176
177   CkAssert(edgeNum == numEdges);
178
179   rc = Zoltan_Initialize(0, NULL, &ver);
180   zz = Zoltan_Create(MPI_COMM_WORLD);
181   char global_parts[10];
182   sprintf(global_parts, "%d", numPes);
183
184   Zoltan_Set_Param(zz, "DEBUG_LEVEL", "0");
185   Zoltan_Set_Param(zz, "LB_METHOD", "HYPERGRAPH");   /* partitioning method */
186   Zoltan_Set_Param(zz, "HYPERGRAPH_PACKAGE", "PHG"); /* version of method */
187   Zoltan_Set_Param(zz, "NUM_GID_ENTRIES", "1");/* global IDs are integers */
188   Zoltan_Set_Param(zz, "NUM_LID_ENTRIES", "1");/* local IDs are integers */
189   Zoltan_Set_Param(zz, "RETURN_LISTS", "PART"); /* export AND import lists */
190   Zoltan_Set_Param(zz, "OBJ_WEIGHT_DIM", "1"); /* use Zoltan default vertex weights */
191   Zoltan_Set_Param(zz, "EDGE_WEIGHT_DIM", "1");/* use Zoltan default hyperedge weights */
192   Zoltan_Set_Param(zz, "NUM_GLOBAL_PARTS", global_parts);
193   Zoltan_Set_Param(zz, "LB_APPROACH", "PARTITION");
194
195   /* Application defined query functions */
196
197   Zoltan_Set_Num_Obj_Fn(zz, get_number_of_vertices, &hg);
198   Zoltan_Set_Obj_List_Fn(zz, get_vertex_list, &hg);
199   Zoltan_Set_HG_Size_CS_Fn(zz, get_hypergraph_size, &hg);
200   Zoltan_Set_HG_CS_Fn(zz, get_hypergraph, &hg);
201   Zoltan_Set_HG_Size_Edge_Wts_Fn(zz, get_hypergraph_edge_size, &hg);
202   Zoltan_Set_HG_Edge_Wts_Fn(zz, get_hypergraph_edge_wgts, &hg);
203
204
205   rc = Zoltan_LB_Partition(zz, /* input (all remaining fields are output) */
206         &changes,        /* 1 if partitioning was changed, 0 otherwise */ 
207         &numGidEntries,  /* Number of integers used for a global ID */
208         &numLidEntries,  /* Number of integers used for a local ID */
209         &numImport,      /* Number of vertices to be sent to me */
210         &importGlobalGids,  /* Global IDs of vertices to be sent to me */
211         &importLocalGids,   /* Local IDs of vertices to be sent to me */
212         &importProcs,    /* Process rank for source of each incoming vertex */
213         &importToPart,   /* New partition for each incoming vertex */
214         &numExport,      /* Number of vertices I must send to other processes*/
215         &exportGlobalGids,  /* Global IDs of the vertices I must send */
216         &exportLocalGids,   /* Local IDs of the vertices I must send */
217         &exportProcs,    /* Process to which I send each of the vertices */
218         &exportToPart);  /* Partition to which each vertex will belong */
219
220   if (rc != ZOLTAN_OK){
221     CkPrintf("Zoltan exiting\n");
222     Zoltan_Destroy(&zz);
223     exit(0);
224   }
225
226
227   for(i = 0; i < numVertices; i++) {
228     if(exportToPart[i] != ogr->vertices[i].getCurrentPe())
229       ogr->vertices[i].setNewPe(exportToPart[i]);
230   }
231
232   /******************************************************************
233   ** Free the arrays allocated by Zoltan_LB_Partition, and free
234   ** the storage allocated for the Zoltan structure.
235   ******************************************************************/
236
237   Zoltan_LB_Free_Part(&importGlobalGids, &importLocalGids, 
238                       &importProcs, &importToPart);
239   Zoltan_LB_Free_Part(&exportGlobalGids, &exportLocalGids, 
240                       &exportProcs, &exportToPart);
241
242   Zoltan_Destroy(&zz);
243
244   /**********************
245   ** all done ***********
246   **********************/
247
248   if (hg.numMyVertices > 0){
249     free(hg.vtxGID);
250   }
251   if (hg.numMyHEdges > 0){
252     free(hg.edgeGID);
253     free(hg.nborIndex);
254     if (hg.numAllNbors > 0){
255       free(hg.nborGID);
256     }
257   }
258
259   ogr->convertDecisions(stats);
260
261   delete parr;
262   delete ogr;
263 }
264
265 /* Application defined query functions */
266
267 static int get_number_of_vertices(void *data, int *ierr)
268 {
269   HGRAPH_DATA *hg = (HGRAPH_DATA *)data;
270   *ierr = ZOLTAN_OK;
271   return hg->numMyVertices;
272 }
273
274 static void get_vertex_list(void *data, int sizeGID, int sizeLID,
275             ZOLTAN_ID_PTR globalID, ZOLTAN_ID_PTR localID,
276                   int wgt_dim, float *obj_wgts, int *ierr)
277 {
278 int i;
279
280   HGRAPH_DATA *hg= (HGRAPH_DATA *)data;
281   *ierr = ZOLTAN_OK;
282
283   /* In this example, return the IDs of our vertices, but no weights.
284    * Zoltan will assume equally weighted vertices.
285    */
286
287   for (i=0; i<hg->numMyVertices; i++){
288     globalID[i] = hg->vtxGID[i];
289     localID[i] = i;
290     obj_wgts[i] = hg->vtxWgt[i];
291   }
292 }
293
294 static void get_hypergraph_size(void *data, int *num_lists, int *num_nonzeroes,
295                                 int *format, int *ierr)
296 {
297   HGRAPH_DATA *hg = (HGRAPH_DATA *)data;
298   *ierr = ZOLTAN_OK;
299
300   *num_lists = hg->numMyHEdges;
301   *num_nonzeroes = hg->numAllNbors;
302
303   /* We will provide compressed hyperedge (row) format.  The alternative is
304    * is compressed vertex (column) format: ZOLTAN_COMPRESSED_VERTEX.
305    */
306
307   *format = ZOLTAN_COMPRESSED_EDGE;
308
309   return;
310 }
311
312 static void get_hypergraph(void *data, int sizeGID, int num_edges, int num_nonzeroes,
313                            int format, ZOLTAN_ID_PTR edgeGID, int *vtxPtr,
314                            ZOLTAN_ID_PTR vtxGID, int *ierr)
315 {
316 int i;
317
318   HGRAPH_DATA *hg = (HGRAPH_DATA *)data;
319   *ierr = ZOLTAN_OK;
320
321   if ( (num_edges != hg->numMyHEdges) || (num_nonzeroes != hg->numAllNbors) ||
322        (format != ZOLTAN_COMPRESSED_EDGE)) {
323     *ierr = ZOLTAN_FATAL;
324     return;
325   }
326
327   for (i=0; i < num_edges; i++){
328     edgeGID[i] = hg->edgeGID[i];
329     vtxPtr[i] = hg->nborIndex[i];
330   }
331
332   for (i=0; i < num_nonzeroes; i++){
333     vtxGID[i] = hg->nborGID[i];
334   }
335
336   return;
337 }
338
339 static void get_hypergraph_edge_size(void *data, int *num_edges, int *ierr) {
340   HGRAPH_DATA *hg = (HGRAPH_DATA *) data;
341   *ierr = ZOLTAN_OK;
342   *num_edges = hg->numMyHEdges;
343 }
344
345 static void get_hypergraph_edge_wgts(void *data, int numGID, int numLID, int num_edges,
346                                      int edge_weight_dim, ZOLTAN_ID_PTR edgeGID, ZOLTAN_ID_PTR edgeLID,
347                                      float *edge_wgts, int *ierr) {
348   int i;
349   HGRAPH_DATA *hg = (HGRAPH_DATA *) data;
350   *ierr = ZOLTAN_OK;
351   for (i = 0; i < num_edges; i++) {
352     edgeGID[i] = hg->edgeGID[i];
353     edgeLID[i] = hg->edgeGID[i];
354     edge_wgts[i] = hg->edgWgt[i];
355   }
356 }
357
358 #include "ZoltanLB.def.h"
359
360 /*@}*/