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