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