2c143f00aec69ca5017d7d0cae362852efc28288
[charm.git] / examples / ampi / Cjacobi3D / jacobi.C
1 #include <stdio.h>
2 #include <stdlib.h>
3 #include "mpi.h"
4
5 #if CMK_BLUEGENE_CHARM
6 extern void BgPrintf(char *);
7 #define BGPRINTF(x)  if (thisIndex == 0) BgPrintf(x);
8 #else
9 #define BGPRINTF(x)
10 #endif
11
12 #ifndef DIMX
13 #define DIMX 100
14 #define DIMY 100
15 #define DIMZ 100
16 #endif
17
18 int NX, NY, NZ;
19
20 class chunk {
21   public:
22     double t[DIMX+2][DIMY+2][DIMZ+2];
23     int xidx, yidx, zidx;
24     int xm, xp, ym, yp, zm, zp;
25     double sbxm[DIMY*DIMZ];
26     double sbxp[DIMY*DIMZ];
27     double sbym[DIMX*DIMZ];
28     double sbyp[DIMX*DIMZ];
29     double sbzm[DIMX*DIMY];
30     double sbzp[DIMX*DIMY];
31     double rbxm[DIMY*DIMZ];
32     double rbxp[DIMY*DIMZ];
33     double rbym[DIMX*DIMZ];
34     double rbyp[DIMX*DIMZ];
35     double rbzm[DIMX*DIMY];
36     double rbzp[DIMX*DIMY];
37 };
38
39 #ifdef AMPI_MIG
40 void chunk_pup(pup_er p, void *d)
41 {
42   chunk **cpp = (chunk **) d;
43   if(pup_isUnpacking(p))
44     *cpp = new chunk;
45   chunk *cp = *cpp;
46   pup_doubles(p, &cp->t[0][0][0], (DIMX+2)*(DIMY+2)*(DIMZ+2));
47   pup_int(p, &cp->xidx);
48   pup_int(p, &cp->yidx);
49   pup_int(p, &cp->zidx);
50   pup_int(p, &cp->xp);
51   pup_int(p, &cp->xm);
52   pup_int(p, &cp->yp);
53   pup_int(p, &cp->ym);
54   pup_int(p, &cp->zp);
55   pup_int(p, &cp->zm);
56   pup_doubles(p, cp->sbxm, (DIMY*DIMZ));
57   pup_doubles(p, cp->sbxp, (DIMY*DIMZ));
58   pup_doubles(p, cp->rbxm, (DIMY*DIMZ));
59   pup_doubles(p, cp->rbxp, (DIMY*DIMZ));
60   pup_doubles(p, cp->sbym, (DIMX*DIMZ));
61   pup_doubles(p, cp->sbyp, (DIMX*DIMZ));
62   pup_doubles(p, cp->rbym, (DIMX*DIMZ));
63   pup_doubles(p, cp->rbyp, (DIMX*DIMZ));
64   pup_doubles(p, cp->sbzm, (DIMX*DIMY));
65   pup_doubles(p, cp->sbzp, (DIMX*DIMY));
66   pup_doubles(p, cp->rbzm, (DIMX*DIMY));
67   pup_doubles(p, cp->rbzp, (DIMX*DIMY));
68   if(pup_isDeleting(p))
69     delete cp;
70 }
71 #endif
72
73 #define abs(x) ((x)<0.0 ? -(x) : (x))
74
75 int index1d(int ix, int iy, int iz)
76 {
77   return NY*NZ*ix + NZ*iy + iz;
78 }
79
80 void index3d(int index, int& ix, int& iy, int& iz)
81 {
82   ix = index/(NY*NZ);
83   iy = (index%(NY*NZ))/NZ;
84   iz = index%NZ;
85 }
86
87 static void copyout(double *d, double t[DIMX+2][DIMY+2][DIMZ+2],
88                     int sx, int ex, int sy, int ey, int sz, int ez)
89 {
90   int i, j, k;
91   int l = 0;
92   for(i=sx; i<=ex; i++)
93     for(j=sy; j<=ey; j++)
94       for(k=sz; k<=ez; k++, l++)
95         d[l] = t[i][j][k];
96 }
97
98 static void copyin(double *d, double t[DIMX+2][DIMY+2][DIMZ+2],
99                     int sx, int ex, int sy, int ey, int sz, int ez)
100 {
101   int i, j, k;
102   int l = 0;
103   for(i=sx; i<=ex; i++)
104     for(j=sy; j<=ey; j++)
105       for(k=sz; k<=ez; k++, l++)
106         t[i][j][k] = d[l];
107 }
108
109 int main(int ac, char** av)
110 {
111   int i,j,k,m,cidx;
112   int iter, niter;
113   MPI_Status status;
114   double error, tval, maxerr, tmpmaxerr, starttime, endtime, itertime;
115   chunk *cp;
116   int thisIndex, ierr, nblocks;
117   double *times;
118   double avgtime = 0.0;
119   
120   MPI_Init(&ac, &av);
121   MPI_Comm_rank(MPI_COMM_WORLD, &thisIndex);
122   MPI_Comm_size(MPI_COMM_WORLD, &nblocks);
123
124   if (ac < 2) {
125     if (thisIndex == 0)
126       printf("Usage: jacobi DIM [nIter].\n");
127     MPI_Finalize();
128   }
129   NX = NY = NZ = atoi(av[1]);
130   if (NX*NY*NZ != nblocks) {
131     if (thisIndex == 0) 
132       printf("%d x %d x %d != %d\n", NX,NY,NZ, nblocks);
133     MPI_Finalize();
134   }
135   if (ac == 3)
136     niter = atoi(av[2]);
137   else
138     niter = 100;
139   times = new double [niter];
140
141   MPI_Bcast(&niter, 1, MPI_INT, 0, MPI_COMM_WORLD);
142
143   cp = new chunk;
144 #if ! defined(NO_PUP) && defined(AMPI_MIG)
145   MPI_Register((void*)&cp, (MPI_PupFn) chunk_pup);
146 #endif
147
148   index3d(thisIndex, cp->xidx, cp->yidx, cp->zidx);
149   cp->xp = index1d((cp->xidx+1)%NX,cp->yidx,cp->zidx);
150   cp->xm = index1d((cp->xidx+NX-1)%NX,cp->yidx,cp->zidx);
151   cp->yp = index1d(cp->xidx,(cp->yidx+1)%NY,cp->zidx);
152   cp->ym = index1d(cp->xidx,(cp->yidx+NY-1)%NY,cp->zidx);
153   cp->zp = index1d(cp->xidx,cp->yidx,(cp->zidx+1)%NZ);
154   cp->zm = index1d(cp->xidx,cp->yidx,(cp->zidx+NZ-1)%NZ);
155   for(i=1; i<=DIMZ; i++)
156     for(j=1; j<=DIMY; j++)
157       for(k=1; k<=DIMX; k++)
158         cp->t[k][j][i] = DIMY*DIMX*(i-1) + DIMX*(j-2) + (k-1);
159
160   MPI_Barrier(MPI_COMM_WORLD);
161   starttime = MPI_Wtime();
162
163   maxerr = 0.0;
164   for(iter=1; iter<=niter; iter++) {
165     BGPRINTF("interation starts at %f\n");
166     maxerr = 0.0;
167     copyout(cp->sbxm, cp->t, 1, 1, 1, DIMY, 1, DIMZ);
168     copyout(cp->sbxp, cp->t, DIMX, DIMX, 1, DIMY, 1, DIMZ);
169     copyout(cp->sbym, cp->t, 1, DIMX, 1, 1, 1, DIMZ);
170     copyout(cp->sbyp, cp->t, 1, DIMX, DIMY, DIMY, 1, DIMZ);
171     copyout(cp->sbzm, cp->t, 1, DIMX, 1, DIMY, 1, 1);
172     copyout(cp->sbzp, cp->t, 1, DIMX, 1, DIMY, DIMZ, DIMZ);
173
174     MPI_Request rreq[6];
175     MPI_Status rsts[6];
176
177     MPI_Irecv(cp->rbxp, DIMY*DIMZ, MPI_DOUBLE, cp->xp, 0, MPI_COMM_WORLD, &rreq[0]);
178     MPI_Irecv(cp->rbxm, DIMY*DIMZ, MPI_DOUBLE, cp->xm, 1, MPI_COMM_WORLD, &rreq[1]);
179     MPI_Irecv(cp->rbyp, DIMX*DIMZ, MPI_DOUBLE, cp->yp, 2, MPI_COMM_WORLD, &rreq[2]);
180     MPI_Irecv(cp->rbym, DIMX*DIMZ, MPI_DOUBLE, cp->ym, 3, MPI_COMM_WORLD, &rreq[3]);
181     MPI_Irecv(cp->rbzm, DIMX*DIMY, MPI_DOUBLE, cp->zm, 5, MPI_COMM_WORLD, &rreq[4]);
182     MPI_Irecv(cp->rbzp, DIMX*DIMY, MPI_DOUBLE, cp->zp, 4, MPI_COMM_WORLD, &rreq[5]);
183
184     MPI_Send(cp->sbxm, DIMY*DIMZ, MPI_DOUBLE, cp->xm, 0, MPI_COMM_WORLD);
185     MPI_Send(cp->sbxp, DIMY*DIMZ, MPI_DOUBLE, cp->xp, 1, MPI_COMM_WORLD);
186     MPI_Send(cp->sbym, DIMX*DIMZ, MPI_DOUBLE, cp->ym, 2, MPI_COMM_WORLD);
187     MPI_Send(cp->sbyp, DIMX*DIMZ, MPI_DOUBLE, cp->yp, 3, MPI_COMM_WORLD);
188     MPI_Send(cp->sbzm, DIMX*DIMY, MPI_DOUBLE, cp->zm, 4, MPI_COMM_WORLD);
189     MPI_Send(cp->sbzp, DIMX*DIMY, MPI_DOUBLE, cp->zp, 5, MPI_COMM_WORLD);
190
191     MPI_Waitall(6, rreq, rsts);
192
193     copyin(cp->sbxm, cp->t, 0, 0, 1, DIMY, 1, DIMZ);
194     copyin(cp->sbxp, cp->t, DIMX+1, DIMX+1, 1, DIMY, 1, DIMZ);
195     copyin(cp->sbym, cp->t, 1, DIMX, 0, 0, 1, DIMZ);
196     copyin(cp->sbyp, cp->t, 1, DIMX, DIMY+1, DIMY+1, 1, DIMZ);
197     copyin(cp->sbzm, cp->t, 1, DIMX, 1, DIMY, 0, 0);
198     copyin(cp->sbzp, cp->t, 1, DIMX, 1, DIMY, DIMZ+1, DIMZ+1);
199     if(iter > 25 &&  iter < 85 && thisIndex == 35)
200       m = 9;
201     else
202       m = 1;
203     for(; m>0; m--)
204       for(i=1; i<=DIMZ; i++)
205         for(j=1; j<=DIMY; j++)
206           for(k=1; k<=DIMX; k++) {
207             tval = (cp->t[k][j][i] + cp->t[k][j][i+1] +
208                  cp->t[k][j][i-1] + cp->t[k][j+1][i]+ 
209                  cp->t[k][j-1][i] + cp->t[k+1][j][i] + cp->t[k-1][j][i])/7.0;
210             error = abs(tval-cp->t[k][j][i]);
211             cp->t[k][j][i] = tval;
212             if (error > maxerr) maxerr = error;
213           }
214     endtime = MPI_Wtime();
215     itertime = endtime - starttime;
216     double  it;
217     MPI_Allreduce(&itertime, &it, 1, MPI_DOUBLE, MPI_SUM,
218                    MPI_COMM_WORLD);
219     itertime = it/nblocks;
220     starttime = MPI_Wtime();
221   }
222
223   for(iter=0;iter<niter;iter++){
224     avgtime += times[iter];
225   }
226   if (thisIndex == 0){
227     printf("Average time %lf\n", avgtime/niter);
228   }
229   delete [] times;
230   MPI_Finalize();
231   return 0;
232 }