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