28da4eff96995ab83b02eb1d7944256e0eac4ba8
[charm.git] / examples / charm++ / load_balancing / stencil3d / stencil3d.C
1 /** \file stencil3d.C
2  *  Author: Abhinav S Bhatele
3  *  Date Created: December 28th, 2010
4  *
5  *  This example is written to be used with periodic measurement-based load
6  *  balancers at sync.
7  *
8  *
9  *
10  *            *****************
11  *         *               *  *
12  *   ^  *****************     *
13  *   |  *               *     *
14  *   |  *               *     *
15  *   |  *               *     *
16  *   Y  *               *     *
17  *   |  *               *     *
18  *   |  *               *     *
19  *   |  *               *  * 
20  *   ~  *****************    Z
21  *      <------ X ------> 
22  *
23  *   X: left, right --> wrap_x
24  *   Y: top, bottom --> wrap_y
25  *   Z: front, back --> wrap_z
26  */
27
28 #include "stencil3d.decl.h"
29 #include "TopoManager.h"
30
31 /*readonly*/ CProxy_Main mainProxy;
32 /*readonly*/ int arrayDimX;
33 /*readonly*/ int arrayDimY;
34 /*readonly*/ int arrayDimZ;
35 /*readonly*/ int blockDimX;
36 /*readonly*/ int blockDimY;
37 /*readonly*/ int blockDimZ;
38
39 // specify the number of worker chares in each dimension
40 /*readonly*/ int num_chare_x;
41 /*readonly*/ int num_chare_y;
42 /*readonly*/ int num_chare_z;
43
44 static unsigned long next = 1;
45
46 int myrand(int numpes) {
47   next = next * 1103515245 + 12345;
48   return((unsigned)(next/65536) % numpes);
49 }
50
51 // We want to wrap entries around, and because mod operator % 
52 // sometimes misbehaves on negative values. -1 maps to the highest value.
53 #define wrap_x(a)       (((a)+num_chare_x)%num_chare_x)
54 #define wrap_y(a)       (((a)+num_chare_y)%num_chare_y)
55 #define wrap_z(a)       (((a)+num_chare_z)%num_chare_z)
56
57 #define index(a,b,c)    ((a)+(b)*(blockDimX+2)+(c)*(blockDimX+2)*(blockDimY+2))
58
59 #define MAX_ITER        26
60 #define LBPERIOD        5
61 #define LEFT            1
62 #define RIGHT           2
63 #define TOP             3
64 #define BOTTOM          4
65 #define FRONT           5
66 #define BACK            6
67 #define DIVIDEBY7       0.14285714285714285714
68
69 double startTime;
70 double endTime;
71
72 /** \class Main
73  *
74  */
75 class Main : public CBase_Main {
76   public:
77     CProxy_Stencil array;
78     int iterations;
79
80     Main(CkArgMsg* m) {
81       if ( (m->argc != 3) && (m->argc != 7) ) {
82         CkPrintf("%s [array_size] [block_size]\n", m->argv[0]);
83         CkPrintf("OR %s [array_size_X] [array_size_Y] [array_size_Z] [block_size_X] [block_size_Y] [block_size_Z]\n", m->argv[0]);
84         CkAbort("Abort");
85       }
86
87       // set iteration counter to zero
88       iterations = 0;
89
90       // store the main proxy
91       mainProxy = thisProxy;
92         
93       if(m->argc == 3) {
94         arrayDimX = arrayDimY = arrayDimZ = atoi(m->argv[1]);
95         blockDimX = blockDimY = blockDimZ = atoi(m->argv[2]); 
96       }
97       else if (m->argc == 7) {
98         arrayDimX = atoi(m->argv[1]);
99         arrayDimY = atoi(m->argv[2]);
100         arrayDimZ = atoi(m->argv[3]);
101         blockDimX = atoi(m->argv[4]); 
102         blockDimY = atoi(m->argv[5]); 
103         blockDimZ = atoi(m->argv[6]);
104       }
105
106       if (arrayDimX < blockDimX || arrayDimX % blockDimX != 0)
107         CkAbort("array_size_X % block_size_X != 0!");
108       if (arrayDimY < blockDimY || arrayDimY % blockDimY != 0)
109         CkAbort("array_size_Y % block_size_Y != 0!");
110       if (arrayDimZ < blockDimZ || arrayDimZ % blockDimZ != 0)
111         CkAbort("array_size_Z % block_size_Z != 0!");
112
113       num_chare_x = arrayDimX / blockDimX;
114       num_chare_y = arrayDimY / blockDimY;
115       num_chare_z = arrayDimZ / blockDimZ;
116
117       // print info
118       CkPrintf("\nSTENCIL COMPUTATION WITH BARRIERS\n");
119       CkPrintf("Running Stencil on %d processors with (%d, %d, %d) chares\n", CkNumPes(), num_chare_x, num_chare_y, num_chare_z);
120       CkPrintf("Array Dimensions: %d %d %d\n", arrayDimX, arrayDimY, arrayDimZ);
121       CkPrintf("Block Dimensions: %d %d %d\n", blockDimX, blockDimY, blockDimZ);
122
123       // Create new array of worker chares
124       array = CProxy_Stencil::ckNew(num_chare_x, num_chare_y, num_chare_z);
125
126       //Start the computation
127       startTime = CmiWallTimer();
128       array.doStep();
129     }
130
131     // Each worker reports back to here when it completes an iteration
132     void report() {
133       iterations++;
134       endTime = CmiWallTimer();
135       CkPrintf("[%d] Time per iteration: %f %f\n", iterations, (endTime - startTime), endTime);
136
137       if(iterations % LBPERIOD == 0)
138         array.startBalancing();
139       else {
140         if(iterations == MAX_ITER) {
141           CkExit();
142         } else {
143           startTime = CmiWallTimer();
144           array.doStep();
145         }
146       }
147     }
148 };
149
150 /** \class Stencil
151  *
152  */
153
154 class Stencil: public CBase_Stencil {
155   Stencil_SDAG_CODE
156
157   public:
158     int iterations;
159     int imsg;
160
161     double *temperature;
162     double *new_temperature;
163
164     // Constructor, initialize values
165     Stencil() {
166       __sdag_init();
167       usesAtSync=CmiTrue;
168
169       int i, j, k;
170       // allocate a three dimensional array
171       temperature = new double[(blockDimX+2) * (blockDimY+2) * (blockDimZ+2)];
172       new_temperature = new double[(blockDimX+2) * (blockDimY+2) * (blockDimZ+2)];
173
174       for(k=0; k<blockDimZ+2; ++k)
175         for(j=0; j<blockDimY+2; ++j)
176           for(i=0; i<blockDimX+2; ++i)
177             temperature[index(i, j, k)] = 0.0;
178
179       iterations = 0;
180       imsg = 0;
181       constrainBC();
182     }
183
184     void pup(PUP::er &p)
185     {
186       CBase_Stencil::pup(p);
187       __sdag_pup(p);
188       p|iterations;
189       p|imsg;
190
191       size_t size = (blockDimX+2) * (blockDimY+2) * (blockDimZ+2);
192       if (p.isUnpacking()) {
193         temperature = new double[size];
194         new_temperature = new double[size];
195       }
196       p(temperature, size);
197       p(new_temperature, size);
198     }
199
200     Stencil(CkMigrateMessage* m) { __sdag_init(); }
201
202     ~Stencil() { 
203       delete [] temperature; 
204       delete [] new_temperature; 
205     }
206
207     // Send ghost faces to the six neighbors
208     void begin_iteration(void) {
209       iterations++;
210
211       // Copy different faces into messages
212       double *leftGhost =  new double[blockDimY*blockDimZ];
213       double *rightGhost =  new double[blockDimY*blockDimZ];
214       double *topGhost =  new double[blockDimX*blockDimZ];
215       double *bottomGhost =  new double[blockDimX*blockDimZ];
216       double *frontGhost =  new double[blockDimX*blockDimY];
217       double *backGhost =  new double[blockDimX*blockDimY];
218
219       for(int k=0; k<blockDimZ; ++k)
220         for(int j=0; j<blockDimY; ++j) {
221           leftGhost[k*blockDimY+j] = temperature[index(1, j+1, k+1)];
222           rightGhost[k*blockDimY+j] = temperature[index(blockDimX, j+1, k+1)];
223         }
224
225       for(int k=0; k<blockDimZ; ++k)
226         for(int i=0; i<blockDimX; ++i) {
227           topGhost[k*blockDimX+i] = temperature[index(i+1, 1, k+1)];
228           bottomGhost[k*blockDimX+i] = temperature[index(i+1, blockDimY, k+1)];
229         }
230
231       for(int j=0; j<blockDimY; ++j)
232         for(int i=0; i<blockDimX; ++i) {
233           frontGhost[j*blockDimX+i] = temperature[index(i+1, j+1, 1)];
234           backGhost[j*blockDimX+i] = temperature[index(i+1, j+1, blockDimZ)];
235         }
236
237       // Send my left face
238       thisProxy(wrap_x(thisIndex.x-1), thisIndex.y, thisIndex.z)
239           .receiveGhosts(iterations, RIGHT, blockDimY, blockDimZ, leftGhost);
240       // Send my right face
241       thisProxy(wrap_x(thisIndex.x+1), thisIndex.y, thisIndex.z)
242           .receiveGhosts(iterations, LEFT, blockDimY, blockDimZ, rightGhost);
243       // Send my bottom face
244       thisProxy(thisIndex.x, wrap_y(thisIndex.y-1), thisIndex.z)
245           .receiveGhosts(iterations, TOP, blockDimX, blockDimZ, bottomGhost);
246       // Send my top face
247       thisProxy(thisIndex.x, wrap_y(thisIndex.y+1), thisIndex.z)
248           .receiveGhosts(iterations, BOTTOM, blockDimX, blockDimZ, topGhost);
249       // Send my front face
250       thisProxy(thisIndex.x, thisIndex.y, wrap_z(thisIndex.z-1))
251           .receiveGhosts(iterations, BACK, blockDimX, blockDimY, frontGhost);
252       // Send my back face
253       thisProxy(thisIndex.x, thisIndex.y, wrap_z(thisIndex.z+1))
254           .receiveGhosts(iterations, FRONT, blockDimX, blockDimY, backGhost);
255
256       delete[] leftGhost;
257       delete[] rightGhost;
258       delete[] bottomGhost;
259       delete[] topGhost;
260       delete[] frontGhost;
261       delete[] backGhost;
262     }
263
264     void processGhosts(int dir, int height, int width, double gh[]) {
265       switch(dir) {
266         case LEFT:
267           for(int k=0; k<width; ++k)
268             for(int j=0; j<height; ++j) {
269               temperature[index(0, j+1, k+1)] = gh[k*height+j];
270             }
271           break;
272         case RIGHT:
273           for(int k=0; k<width; ++k)
274             for(int j=0; j<height; ++j) {
275               temperature[index(blockDimX+1, j+1, k+1)] = gh[k*height+j];
276             }
277           break;
278         case BOTTOM:
279           for(int k=0; k<width; ++k)
280             for(int i=0; i<height; ++i) {
281               temperature[index(i+1, 0, k+1)] = gh[k*height+i];
282             }
283           break;
284         case TOP:
285           for(int k=0; k<width; ++k)
286             for(int i=0; i<height; ++i) {
287               temperature[index(i+1, blockDimY+1, k+1)] = gh[k*height+i];
288             }
289           break;
290         case FRONT:
291           for(int j=0; j<width; ++j)
292             for(int i=0; i<height; ++i) {
293               temperature[index(i+1, j+1, 0)] = gh[j*height+i];
294             }
295           break;
296         case BACK:
297           for(int j=0; j<width; ++j)
298             for(int i=0; i<height; ++i) {
299               temperature[index(i+1, j+1, blockDimZ+1)] = gh[j*height+i];
300             }
301           break;
302         default:
303           CkAbort("ERROR\n");
304       }
305     }
306
307
308     void check_and_compute() {
309       compute_kernel();
310
311       // calculate error
312       // not being done right now since we are doing a fixed no. of iterations
313
314       double *tmp;
315       tmp = temperature;
316       temperature = new_temperature;
317       new_temperature = tmp;
318
319       constrainBC();
320
321       contribute(0, 0, CkReduction::concat, CkCallback(CkIndex_Main::report(), mainProxy));
322     }
323
324     // Check to see if we have received all neighbor values yet
325     // If all neighbor values have been received, we update our values and proceed
326     void compute_kernel() {
327 #pragma unroll    
328       for(int k=1; k<blockDimZ+1; ++k)
329         for(int j=1; j<blockDimY+1; ++j)
330           for(int i=1; i<blockDimX+1; ++i) {
331             // update my value based on the surrounding values
332             new_temperature[index(i, j, k)] = (temperature[index(i-1, j, k)] 
333                                             +  temperature[index(i+1, j, k)]
334                                             +  temperature[index(i, j-1, k)]
335                                             +  temperature[index(i, j+1, k)]
336                                             +  temperature[index(i, j, k-1)]
337                                             +  temperature[index(i, j, k+1)]
338                                             +  temperature[index(i, j, k)] ) * DIVIDEBY7;
339           } // end for
340     }
341
342     // Enforce some boundary conditions
343     void constrainBC() {
344       // Heat left, top and front faces of each chare's block
345       for(int k=1; k<blockDimZ+1; ++k)
346         for(int i=1; i<blockDimX+1; ++i)
347           temperature[index(i, 1, k)] = 255.0;
348       for(int k=1; k<blockDimZ+1; ++k)
349         for(int j=1; j<blockDimY+1; ++j)
350           temperature[index(1, j, k)] = 255.0;
351       for(int j=1; j<blockDimY+1; ++j)
352         for(int i=1; i<blockDimX+1; ++i)
353           temperature[index(i, j, 1)] = 255.0;
354     }
355
356     void startBalancing() {
357       AtSync();
358     }
359
360     void ResumeFromSync() {
361       doStep();
362     }
363 };
364
365 #include "stencil3d.def.h"