fix buffer problem in PMEMimic
[charm.git] / examples / charm++ / PMEMimic / PMEMimic.C
1 #include <stdio.h>
2 #include "PMEMimic.decl.h"
3
4 /*readonly*/ CProxy_Main mainProxy;
5 /*readonly*/ 
6
7 int     N;
8 int     grid_x;
9 int     grid_y;
10 int     grid_z;
11 int     max_iter;
12 int     pes_per_node;
13
14 CProxy_PMEPencil pme_x;
15 CProxy_PMEPencil pme_y;
16 CProxy_PMEPencil pme_z;
17
18 class DataMsg : public CMessage_DataMsg
19 {
20 public:
21     int phrase;
22     char data[2048];
23
24 };
25
26 class PMEMap : public CkArrayMap 
27
28     int offset;
29 public: 
30     PMEMap(int off) { offset = off;} 
31     PMEMap(CkMigrateMessage *m){} 
32     int registerArray(CkArrayIndex& numElements,CkArrayID aid) { 
33         return 0; 
34     } 
35     int procNum(int /*arrayHdl*/,const CkArrayIndex &idx) { 
36         int penum;
37         int *index =  (int *)idx.data();
38         int obj_index =  index[0]*grid_x + index[1];
39         penum = obj_index * pes_per_node + offset;
40         return penum; 
41     } 
42 }; 
43
44
45 /*mainchare*/
46 class Main : public CBase_Main
47 {
48     double startTimer;
49     int done_pme, iteration;
50 public:
51
52     Main(CkArgMsg* m)
53     {
54         //Process command-line arguments
55  
56         grid_x = grid_y = grid_z = 10;
57         max_iter = 100;
58         pes_per_node = 3;
59         if(m->argc > 1)
60         {
61             pes_per_node = atoi(m->argv[1]);;
62             grid_x = grid_y = grid_z = atoi(m->argv[2]);
63             max_iter = atoi(m->argv[3]);
64         }
65         delete m;
66
67     //Start the computation
68       CkPrintf("Running PMEMimic on %d processors for %d elements\n",
69           CkNumPes(), grid_x);
70       mainProxy = thisProxy;
71
72       CProxy_PMEMap myMap_x=CProxy_PMEMap::ckNew(0); 
73       CkArrayOptions opts_x(grid_y, grid_z); 
74       opts_x.setMap(myMap_x);
75
76       CProxy_PMEMap myMap_y=CProxy_PMEMap::ckNew(1); 
77       CkArrayOptions opts_y(grid_x, grid_z); 
78       opts_y.setMap(myMap_y);
79
80       CProxy_PMEMap myMap_z=CProxy_PMEMap::ckNew(2); 
81       CkArrayOptions opts_z(grid_x, grid_y); 
82       opts_z.setMap(myMap_z);
83
84       pme_x = CProxy_PMEPencil::ckNew(0, opts_x);
85       pme_y = CProxy_PMEPencil::ckNew(1, opts_y);
86       pme_z = CProxy_PMEPencil::ckNew(2, opts_z);
87
88       done_pme=0;
89       startTimer = CmiWallTimer();
90       pme_x.start();
91       
92     };
93
94     void done()
95     {
96         done_pme++;
97         if(done_pme == grid_x)
98         {
99             done_pme = 0;
100
101             CkPrintf("PME(%d, %d, %d) on %d PEs, %d iteration, avg time:%f\n", grid_x, grid_y, grid_z, CkNumPes(), max_iter, CmiWallTimer()-startTimer);
102             CkExit();
103         }
104     }
105 };
106
107 /*array [1D]*/
108 class PMEPencil : public CBase_PMEPencil
109 {
110     int PME_index;
111     int buffered_num, buffered_phrase;
112     int recv_nums, iteration;
113 public:
114   PMEPencil(int i)
115   {
116       PME_index = i;
117       recv_nums = 0;
118       iteration = 0;
119       buffered_num = 0;
120   }
121   PMEPencil(CkMigrateMessage *m) {}
122
123   void start()
124   {
125    //thisindex.x thisindex.y
126     // x (yz), y(x, z)
127     for(int x=0; x<grid_x; x++)
128     {
129       DataMsg *msg= new DataMsg;
130       msg->phrase = 1;
131       pme_y(x, thisIndex.y).recvTrans(msg);  
132     }
133   }
134   void recvTrans(DataMsg *msg_recv)
135   {
136     int expect_num, index;
137     expect_num = grid_x;
138     index = msg_recv->phrase;
139
140     if(msg_recv->phrase != PME_index)
141     {
142         buffered_num++;
143         buffered_phrase = msg_recv->phrase;
144         delete msg_recv;
145         return;
146     }
147     recv_nums++;
148     if(recv_nums == expect_num)
149     {
150         //CkPrintf("[%d, %d] phrase %d, iter=%d\n", thisIndex.x, thisIndex.y, msg_recv->phrase, iteration);
151         if(index == 0  ) //x (y,z) to y(x,z)
152         {
153             iteration++;
154             if(iteration == max_iter)
155             {
156                 mainProxy.done();
157                 return;
158             }
159             for(int x=0; x<grid_x; x++)
160             {
161                 DataMsg *msg= new DataMsg;
162                 msg->phrase = msg_recv->phrase+1;
163                 pme_y(x, thisIndex.y).recvTrans(msg);  
164             }
165         }else if(index == 1) //y(x,z) send to z(x,y)
166         {
167             for(int y=0; y<grid_y; y++)
168             {
169                 DataMsg *msg= new DataMsg;
170                 msg->phrase = msg_recv->phrase+1;
171                 pme_z(thisIndex.x, y).recvTrans(msg); 
172             }
173             PME_index = 3;
174             recv_nums = buffered_num;
175         }else if(index == 2) //Z(x,y) send to y(x,z)
176         {
177             for(int z=0; z<grid_z; z++)
178             {
179                 DataMsg *msg= new DataMsg;
180                 msg->phrase = msg_recv->phrase+1;
181                 pme_y(thisIndex.x, z).recvTrans(msg); 
182             }
183         } else if(index == 3) //y(x,z) to x(y,z)
184         {
185             for(int y=0; y<grid_y; y++)
186             {
187                 DataMsg *msg= new DataMsg;
188                 msg->phrase = 0;
189                 pme_x(y, thisIndex.y).recvTrans(msg); 
190             }
191             PME_index = 1;
192             recv_nums = buffered_num;
193         }
194         recv_nums = 0;
195     }
196     delete msg_recv;
197   }
198 };
199
200
201 #include "PMEMimic.def.h"