375431bd5265919cd3cbb2ee0a9cda2e33378104
[namd.git] / src / ComputePme.C
1 /**
2 ***  Copyright (c) 1995, 1996, 1997, 1998, 1999, 2000 by
3 ***  The Board of Trustees of the University of Illinois.
4 ***  All rights reserved.
5 **/
6
7 #ifdef NAMD_FFTW
8 //#define MANUAL_DEBUG_FFTW3 1
9 #ifdef NAMD_FFTW_3
10 #include <fftw3.h>
11 #else
12 // fftw2 doesn't have these defined
13 #define fftwf_malloc fftw_malloc
14 #define fftwf_free fftw_free
15 #ifdef NAMD_FFTW_NO_TYPE_PREFIX
16 #include <fftw.h>
17 #include <rfftw.h>
18 #else
19 #include <sfftw.h>
20 #include <srfftw.h>
21 #endif
22 #endif
23 #endif
24
25 #include <vector>
26 #include <algorithm>
27 #include <deque>
28 using namespace std;
29
30 #include "InfoStream.h"
31 #include "Node.h"
32 #include "PatchMap.h"
33 #include "PatchMap.inl"
34 #include "AtomMap.h"
35 #include "ComputePme.h"
36 #include "ComputePmeMgr.decl.h"
37 #include "PmeBase.inl"
38 #include "PmeRealSpace.h"
39 #include "PmeKSpace.h"
40 #include "ComputeNonbondedUtil.h"
41 #include "PatchMgr.h"
42 #include "Molecule.h"
43 #include "ReductionMgr.h"
44 #include "ComputeMgr.h"
45 #include "ComputeMgr.decl.h"
46 // #define DEBUGM
47 #define MIN_DEBUG_LEVEL 3
48 #include "Debug.h"
49 #include "SimParameters.h"
50 #include "WorkDistrib.h"
51 #include "varsizemsg.h"
52 #include "Random.h"
53 #include "ckhashtable.h"
54 #include "Priorities.h"
55
56 #include "ComputeMoa.h"
57 #include "ComputeMoaMgr.decl.h" 
58
59 //#define     USE_RANDOM_TOPO         1
60
61 //#define USE_TOPO_SFC                    1
62 //#define     USE_CKLOOP                1
63 //#include "TopoManager.h"
64
65 #include "DeviceCUDA.h"
66 #ifdef NAMD_CUDA
67 #include <cuda_runtime.h>
68 #include <cuda.h>
69 void cuda_errcheck(const char *msg);
70 #ifdef WIN32
71 #define __thread __declspec(thread)
72 #endif
73 extern __thread DeviceCUDA *deviceCUDA;
74 #endif
75
76 #include "ComputePmeCUDAKernel.h"
77
78 #ifndef SQRT_PI
79 #define SQRT_PI 1.7724538509055160273 /* mathematica 15 digits*/
80 #endif
81
82 #if CMK_PERSISTENT_COMM 
83 #define USE_PERSISTENT      1
84 #endif
85
86 #if USE_PERSISTENT
87 #define Z_PERSIST 1
88 #define Y_PERSIST 1
89 #define X_PERSIST 1
90 #endif
91
92 #if defined(NAMD_CUDA) && defined(MEM_OPT_VERSION)
93 #define USE_NODE_PAR_RECEIVE    1
94 #endif
95
96 int ComputePmeUtil::numGrids;
97 Bool ComputePmeUtil::alchOn;
98 Bool ComputePmeUtil::alchFepOn;
99 Bool ComputePmeUtil::alchThermIntOn;
100 Bool ComputePmeUtil::alchDecouple;
101 BigReal ComputePmeUtil::alchElecLambdaStart;
102 Bool ComputePmeUtil::lesOn;
103 int ComputePmeUtil::lesFactor;
104 Bool ComputePmeUtil::pairOn;
105 Bool ComputePmeUtil::selfOn;
106
107 char *pencilPMEProcessors;
108
109 class PmeAckMsg : public CMessage_PmeAckMsg {
110 };
111
112 class PmeGridMsg : public CMessage_PmeGridMsg {
113 public:
114
115   int sourceNode;
116   int sequence;
117   int hasData;
118   Lattice lattice;
119   int start;
120   int len;
121   int zlistlen;
122   int *zlist;
123   char *fgrid;
124   float *qgrid;
125   CkArrayIndex3D destElem;
126 };
127
128 class PmeTransMsg : public CMessage_PmeTransMsg {
129 public:
130
131   int sourceNode;
132   int sequence;
133   int hasData;
134   Lattice lattice;
135   int x_start;
136   int nx;
137   float *qgrid;
138   CkArrayIndex3D destElem;
139 };
140
141 class PmeSharedTransMsg : public CMessage_PmeSharedTransMsg {
142 public:
143   PmeTransMsg *msg;
144   int *count;
145   CmiNodeLock lock;
146 };
147
148 class PmeUntransMsg : public CMessage_PmeUntransMsg {
149 public:
150
151   int sourceNode;
152   int y_start;
153   int ny;
154   float *qgrid;
155   CkArrayIndex3D destElem;
156 };
157
158 class PmeSharedUntransMsg : public CMessage_PmeSharedUntransMsg {
159 public:
160   PmeUntransMsg *msg;
161   int *count;
162   CmiNodeLock lock;
163 };
164
165 class PmeEvirMsg : public CMessage_PmeEvirMsg {
166 public:
167   PmeReduction *evir;
168 };
169
170 class PmePencilMap : public CBase_PmePencilMap {
171 public:
172   PmePencilMap(int i_a, int i_b, int n_b, int n, int *d)
173     : ia(i_a), ib(i_b), nb(n_b),
174       size(n), data(newcopyint(n,d)) {
175   }
176   virtual int registerArray(CkArrayIndexMax&, CkArrayID) {
177     //Return an ``arrayHdl'', given some information about the array
178     return 0;
179   }
180   virtual int procNum(int, const CkArrayIndex &i) {
181     //Return the home processor number for this element of this array
182     return data[ i.data()[ia] * nb + i.data()[ib] ];
183   }
184   virtual void populateInitial(int, CkArrayIndexMax &, void *msg, CkArrMgr *mgr) {
185     int mype = CkMyPe();
186     for ( int i=0; i < size; ++i ) {
187       if ( data[i] == mype ) {
188         CkArrayIndex3D ai(0,0,0);
189         ai.data()[ia] = i / nb;
190         ai.data()[ib] = i % nb;
191         if ( procNum(0,ai) != mype ) NAMD_bug("PmePencilMap is inconsistent");
192         if ( ! msg ) NAMD_bug("PmePencilMap multiple pencils on a pe?");
193         mgr->insertInitial(ai,msg);
194         msg = 0;
195       }
196     }
197     mgr->doneInserting();
198     if ( msg ) CkFreeMsg(msg);
199   }
200 private:
201   const int ia, ib, nb, size;
202   const int* const data;
203   static int* newcopyint(int n, int *d) {
204     int *newd = new int[n];
205     memcpy(newd, d, n*sizeof(int));
206     return newd;
207   }
208 };
209
210 // use this idiom since messages don't have copy constructors
211 struct PmePencilInitMsgData {
212   PmeGrid grid;
213   int xBlocks, yBlocks, zBlocks;
214   CProxy_PmeXPencil xPencil;
215   CProxy_PmeYPencil yPencil;
216   CProxy_PmeZPencil zPencil;
217   CProxy_ComputePmeMgr pmeProxy;
218   CProxy_NodePmeMgr pmeNodeProxy;
219   CProxy_PmePencilMap xm;
220   CProxy_PmePencilMap ym;
221   CProxy_PmePencilMap zm;
222 };
223
224 class PmePencilInitMsg : public CMessage_PmePencilInitMsg {
225 public:
226    PmePencilInitMsg(PmePencilInitMsgData &d) { data = d; }
227    PmePencilInitMsgData data;
228 };
229
230
231 struct LocalPmeInfo {
232   int nx, x_start;
233   int ny_after_transpose, y_start_after_transpose;
234 };
235
236 struct NodePmeInfo {
237   int npe, pe_start, real_node;
238 };
239
240
241 static int findRecipEvirPe() {
242     PatchMap *patchMap = PatchMap::Object();
243     {
244       int mype = CkMyPe();
245       if ( patchMap->numPatchesOnNode(mype) ) {
246         return mype; 
247       }
248     }
249     {
250       int node = CmiMyNode();
251       int firstpe = CmiNodeFirst(node);
252       int nodeSize = CmiNodeSize(node);
253       int myrank = CkMyRank();
254       for ( int i=0; i<nodeSize; ++i ) {
255         int pe = firstpe + (myrank+i)%nodeSize;
256         if ( patchMap->numPatchesOnNode(pe) ) {
257           return pe;
258         }
259       }
260     }
261     {
262       int *pelist;
263       int nodeSize;
264       CmiGetPesOnPhysicalNode(CmiPhysicalNodeID(CkMyPe()), &pelist, &nodeSize);
265       int myrank;
266       for ( int i=0; i<nodeSize; ++i ) {
267         if ( pelist[i] == CkMyPe() ) myrank = i;
268       }
269       for ( int i=0; i<nodeSize; ++i ) {
270         int pe = pelist[(myrank+i)%nodeSize];
271         if ( patchMap->numPatchesOnNode(pe) ) {
272           return pe;
273         }
274       }
275     }
276     {
277       int mype = CkMyPe();
278       int npes = CkNumPes();
279       for ( int i=0; i<npes; ++i ) {
280         int pe = (mype+i)%npes;
281         if ( patchMap->numPatchesOnNode(pe) ) {
282           return pe;
283         }
284       }
285     }
286     NAMD_bug("findRecipEvirPe() failed!");
287     return -999;  // should never happen
288 }
289
290
291 //Assigns gridPeMap and transPeMap to different set of processors.
292 void generatePmePeList2(int *gridPeMap, int numGridPes, int *transPeMap, int numTransPes){
293   int ncpus = CkNumPes();
294   
295   for ( int i=0; i<numGridPes; ++i ) {
296     gridPeMap[i] = WorkDistrib::peDiffuseOrdering[ncpus - numGridPes + i];
297   }
298   std::sort(gridPeMap,gridPeMap+numGridPes);
299   int firstTransPe = ncpus - numGridPes - numTransPes;
300   if ( firstTransPe < 0 ) {
301     firstTransPe = 0;
302     // 0 should be first in list, skip if possible
303     if ( ncpus > numTransPes ) firstTransPe = 1;
304   }
305   for ( int i=0; i<numTransPes; ++i ) {
306     transPeMap[i] = WorkDistrib::peDiffuseOrdering[firstTransPe + i];
307   }
308   std::sort(transPeMap,transPeMap+numTransPes);
309 }
310
311 #if USE_TOPOMAP 
312 //Topology aware PME allocation
313 bool generateBGLORBPmePeList(int *pemap, int numPes, int *block_pes=0, 
314                              int nbpes=0);
315 #endif
316
317
318 int compare_bit_reversed(int a, int b) {
319   int d = a ^ b;
320   int c = 1;
321   if ( d ) while ( ! (d & c) ) {
322     c = c << 1;
323   }
324   return (a & c) - (b & c);
325 }
326
327 inline bool less_than_bit_reversed(int a, int b) {
328   int d = a ^ b;
329   int c = 1;
330   if ( d ) while ( ! (d & c) ) {
331     c = c << 1;
332   }
333   return d && (b & c);
334 }
335
336 struct sortop_bit_reversed {
337   inline bool operator() (int a, int b) const {
338     return less_than_bit_reversed(a,b);
339   }
340 };
341
342 struct ijpair {
343   int i,j;
344   ijpair() {;}
345   ijpair(int I, int J) : i(I), j(J) {;}
346 };
347
348 struct ijpair_sortop_bit_reversed {
349   inline bool operator() (const ijpair &a, const ijpair &b) const {
350     return ( less_than_bit_reversed(a.i,b.i)
351              || ( (a.i == b.i) && less_than_bit_reversed(a.j,b.j) ) );
352   }
353 };
354
355 class ComputePmeMgr : public CBase_ComputePmeMgr, public ComputePmeUtil {
356 public:
357   friend class ComputePme;
358   friend class NodePmeMgr;
359   ComputePmeMgr();
360   ~ComputePmeMgr();
361
362   void initialize(CkQdMsg*);
363   void initialize_pencils(CkQdMsg*);
364   void activate_pencils(CkQdMsg*);
365   void recvArrays(CProxy_PmeXPencil, CProxy_PmeYPencil, CProxy_PmeZPencil);
366   void initialize_computes();
367
368   void sendData(Lattice &, int sequence);
369   void sendDataPart(int first, int last, Lattice &, int sequence, int sourcepe, int errors);
370   Lattice *sendDataHelper_lattice;
371   int sendDataHelper_sequence;
372   int sendDataHelper_sourcepe;
373   int sendDataHelper_errors;
374   void sendPencils(Lattice &, int sequence);
375   void sendPencilsPart(int first, int last, Lattice &, int sequence, int sourcepe);
376   void recvGrid(PmeGridMsg *);
377   void gridCalc1(void);
378   void sendTransBarrier(void);
379   void sendTransSubset(int first, int last);
380   void sendTrans(void);
381   void fwdSharedTrans(PmeTransMsg *);
382   void recvSharedTrans(PmeSharedTransMsg *);
383   void sendDataHelper(int);
384   void sendPencilsHelper(int);
385   void recvTrans(PmeTransMsg *);
386   void procTrans(PmeTransMsg *);
387   void gridCalc2(void);
388   #ifdef OPENATOM_VERSION
389   void gridCalc2Moa(void);
390   #endif // OPENATOM_VERSION
391   void gridCalc2R(void);
392   void fwdSharedUntrans(PmeUntransMsg *);
393   void recvSharedUntrans(PmeSharedUntransMsg *);
394   void sendUntrans(void);
395   void sendUntransSubset(int first, int last);
396   void recvUntrans(PmeUntransMsg *);
397   void procUntrans(PmeUntransMsg *);
398   void gridCalc3(void);
399   void sendUngrid(void);
400   void sendUngridSubset(int first, int last);
401   void recvUngrid(PmeGridMsg *);
402   void recvAck(PmeAckMsg *);
403   void copyResults(PmeGridMsg *);
404   void copyPencils(PmeGridMsg *);
405   void ungridCalc(void);
406   void recvRecipEvir(PmeEvirMsg *);
407   void addRecipEvirClient(void);
408   void submitReductions();
409
410 #if 0 && USE_PERSISTENT
411   void setup_recvgrid_persistent();
412 #endif
413
414   static CmiNodeLock fftw_plan_lock;
415   CmiNodeLock pmemgr_lock;  // for accessing this object from other threads
416
417 #ifdef NAMD_CUDA
418   float *a_data_host;
419   float *a_data_dev;
420   float *f_data_host;
421   float *f_data_dev;
422   int cuda_atoms_count;
423   int cuda_atoms_alloc;
424   static CmiNodeLock cuda_lock;
425   void chargeGridSubmitted(Lattice &lattice, int sequence);
426   cudaEvent_t end_charges;
427   cudaEvent_t *end_forces;
428   int forces_count;
429   int forces_done_count;
430   double charges_time;
431   double forces_time;
432   int check_charges_count;
433   int check_forces_count;
434   int master_pe;
435   int this_pe;
436
437   void cuda_submit_charges(Lattice &lattice, int sequence);
438   struct cuda_submit_charges_args {
439     ComputePmeMgr *mgr; Lattice *lattice; int sequence;
440   };
441   static std::deque<cuda_submit_charges_args> cuda_submit_charges_deque;
442   static bool cuda_busy;
443
444   int chargeGridSubmittedCount;
445   void sendChargeGridReady();
446 #endif
447   Lattice *saved_lattice;  // saved by chargeGridSubmitted
448   int saved_sequence;      // saved by chargeGridSubmitted
449   void pollChargeGridReady();
450   void pollForcesReady();
451   void recvChargeGridReady();
452   void chargeGridReady(Lattice &lattice, int sequence);
453
454   ResizeArray<ComputePme*> pmeComputes;
455
456 private:
457
458 #if 0 && USE_PERSISTENT
459   PersistentHandle   *recvGrid_handle;
460 #endif
461
462   CProxy_ComputePmeMgr pmeProxy;
463   CProxy_ComputePmeMgr pmeProxyDir;
464   CProxy_NodePmeMgr pmeNodeProxy;
465   NodePmeMgr *nodePmeMgr;
466   ComputePmeMgr *masterPmeMgr;
467   
468   void addCompute(ComputePme *c) {
469     if ( ! pmeComputes.size() ) initialize_computes();
470     pmeComputes.add(c);
471     c->setMgr(this);
472   }
473
474   ResizeArray<ComputePme*> heldComputes;
475   PmeGrid myGrid;
476   Lattice lattice;
477   PmeKSpace *myKSpace;
478   float *qgrid;
479   float *kgrid;
480
481 #ifdef NAMD_FFTW
482 #ifdef NAMD_FFTW_3
483   fftwf_plan *forward_plan_x, *backward_plan_x;
484   fftwf_plan *forward_plan_yz, *backward_plan_yz;
485   fftwf_complex *work;
486 #else
487   fftw_plan forward_plan_x, backward_plan_x;
488   rfftwnd_plan forward_plan_yz, backward_plan_yz;
489   fftw_complex *work;
490 #endif
491 #else
492   float *work;
493 #endif
494
495   int qsize, fsize, bsize;
496   int offload;
497   BigReal alchLambda;  // set on each step in ComputePme::ungridForces()
498   BigReal alchLambda2; // set on each step in ComputePme::ungridForces()
499
500   float **q_arr;
501   // q_list and q_count not used for offload
502   float **q_list;
503   int q_count;
504   char *f_arr;
505   char *fz_arr;
506   PmeReduction evir[PME_MAX_EVALS];
507   SubmitReduction *reduction;
508
509   int noWorkCount;
510   int doWorkCount;
511   int ungridForcesCount;
512
513 #ifdef NAMD_CUDA
514 #define NUM_STREAMS 1
515   cudaStream_t streams[NUM_STREAMS];
516   int stream;
517
518   float **q_arr_dev;
519   float **v_arr_dev;
520   float *q_data_host;
521   float *q_data_dev;
522   float *v_data_dev;
523   int *ffz_host;
524   int *ffz_dev;
525   int q_data_size;
526   int ffz_size;
527
528   int f_data_mgr_alloc;
529   float *f_data_mgr_host;
530   float *f_data_mgr_dev;
531   float **afn_host;
532   float **afn_dev;
533
534   float *bspline_coeffs_dev;
535   float *bspline_dcoeffs_dev;
536 #endif
537   int recipEvirCount;   // used in compute only
538   int recipEvirClients; // used in compute only
539   int recipEvirPe;      // used in trans only
540   
541   LocalPmeInfo *localInfo;
542   NodePmeInfo *gridNodeInfo;
543   NodePmeInfo *transNodeInfo;
544   int qgrid_size;
545   int qgrid_start;
546   int qgrid_len;
547   int fgrid_start;
548   int fgrid_len;
549
550   int numSources;
551   int numGridPes;
552   int numTransPes;
553   int numGridNodes;
554   int numTransNodes;
555   int numDestRecipPes;
556   int myGridPe, myGridNode;
557   int myTransPe, myTransNode;
558   int *gridPeMap;
559   int *transPeMap;
560   int *recipPeDest;
561   int *gridPeOrder;
562   int *gridNodeOrder;
563   int *transNodeOrder;
564   int grid_count;
565   int trans_count;
566   int untrans_count;
567   int ungrid_count;
568   PmeGridMsg **gridmsg_reuse;
569   PmeReduction recip_evir2[PME_MAX_EVALS];
570
571   int compute_sequence;  // set from patch computes, used for priorities
572   int grid_sequence;  // set from grid messages, used for priorities
573   int useBarrier;
574   int sendTransBarrier_received;
575
576   int usePencils;
577   int xBlocks, yBlocks, zBlocks;
578   CProxy_PmeXPencil xPencil;
579   CProxy_PmeYPencil yPencil;
580   CProxy_PmeZPencil zPencil;
581   char *pencilActive;
582   ijpair *activePencils;
583   int numPencilsActive;
584   int strayChargeErrors;
585 };
586
587 ResizeArray<ComputePme*>& getComputes(ComputePmeMgr *mgr) {
588     return mgr->pmeComputes ;
589 }
590
591   CmiNodeLock ComputePmeMgr::fftw_plan_lock;
592 #ifdef NAMD_CUDA
593   CmiNodeLock ComputePmeMgr::cuda_lock;
594   std::deque<ComputePmeMgr::cuda_submit_charges_args> ComputePmeMgr::cuda_submit_charges_deque;
595   bool ComputePmeMgr::cuda_busy;
596 #endif
597
598 int isPmeProcessor(int p){ 
599   SimParameters *simParams = Node::Object()->simParameters;
600   if (simParams->usePMECUDA) {
601     return 0;
602   } else {
603     return pencilPMEProcessors[p];
604   }
605 }
606
607 class NodePmeMgr : public CBase_NodePmeMgr {
608 public:
609   friend class ComputePmeMgr;
610   friend class ComputePme;
611   NodePmeMgr();
612   ~NodePmeMgr();
613   void initialize();
614   void sendDataHelper(int);
615   void sendPencilsHelper(int);
616   void recvTrans(PmeTransMsg *);
617   void recvUntrans(PmeUntransMsg *);
618   void registerXPencil(CkArrayIndex3D, PmeXPencil *);
619   void registerYPencil(CkArrayIndex3D, PmeYPencil *);
620   void registerZPencil(CkArrayIndex3D, PmeZPencil *);
621   void recvXTrans(PmeTransMsg *);
622   void recvYTrans(PmeTransMsg *);
623   void recvYUntrans(PmeUntransMsg *);
624   void recvZGrid(PmeGridMsg *);
625   void recvZUntrans(PmeUntransMsg *);
626
627   void recvUngrid(PmeGridMsg *);
628
629   void recvPencilMapProxies(CProxy_PmePencilMap _xm, CProxy_PmePencilMap _ym, CProxy_PmePencilMap _zm){
630       xm=_xm; ym=_ym; zm=_zm;
631   }
632   CProxy_PmePencilMap xm;
633   CProxy_PmePencilMap ym;
634   CProxy_PmePencilMap zm;
635
636 private:
637   CProxy_ComputePmeMgr mgrProxy;
638   ComputePmeMgr *mgrObject;
639   ComputePmeMgr **mgrObjects;
640 #ifdef NAMD_CUDA
641   ComputePmeMgr *masterPmeMgr;
642   int master_pe;
643 #endif
644   CProxy_PmeXPencil xPencil;
645   CProxy_PmeYPencil yPencil;
646   CProxy_PmeZPencil zPencil;
647   CkHashtableT<CkArrayIndex3D,PmeXPencil*> xPencilObj;
648   CkHashtableT<CkArrayIndex3D,PmeYPencil*> yPencilObj;
649   CkHashtableT<CkArrayIndex3D,PmeZPencil*> zPencilObj;  
650
651 #ifdef NAMD_CUDA
652   cudaEvent_t end_charge_memset;
653   cudaEvent_t end_all_pme_kernels;
654   cudaEvent_t end_potential_memcpy;
655 #endif
656 };
657
658 NodePmeMgr::NodePmeMgr() {
659   mgrObjects = new ComputePmeMgr*[CkMyNodeSize()];
660 }
661
662 NodePmeMgr::~NodePmeMgr() {
663   delete [] mgrObjects;
664 }
665
666 void NodePmeMgr::initialize() {
667   CProxy_ComputePmeMgr proxy = CkpvAccess(BOCclass_group).computePmeMgr;
668   mgrObjects[CkMyRank()] = proxy.ckLocalBranch();
669   if ( CkMyRank() == 0 ) {
670     mgrProxy = proxy;
671     mgrObject = proxy.ckLocalBranch();
672   }
673 }
674
675 void NodePmeMgr::recvTrans(PmeTransMsg *msg) {
676   mgrObject->fwdSharedTrans(msg);
677 }
678
679 void NodePmeMgr::recvUntrans(PmeUntransMsg *msg) {
680   mgrObject->fwdSharedUntrans(msg);
681 }
682
683 void NodePmeMgr::recvUngrid(PmeGridMsg *msg) {
684 #ifdef NAMD_CUDA
685   masterPmeMgr->recvUngrid(msg);
686 #else
687   NAMD_bug("NodePmeMgr::recvUngrid called in non-CUDA build.");
688 #endif
689 }
690
691 void NodePmeMgr::registerXPencil(CkArrayIndex3D idx, PmeXPencil *obj)
692 {
693   CmiLock(ComputePmeMgr::fftw_plan_lock);
694   xPencilObj.put(idx)=obj;
695   CmiUnlock(ComputePmeMgr::fftw_plan_lock);
696 }
697 void NodePmeMgr::registerYPencil(CkArrayIndex3D idx, PmeYPencil *obj)
698 {
699   CmiLock(ComputePmeMgr::fftw_plan_lock);
700   yPencilObj.put(idx)=obj;
701   CmiUnlock(ComputePmeMgr::fftw_plan_lock);
702 }
703 void NodePmeMgr::registerZPencil(CkArrayIndex3D idx, PmeZPencil *obj)
704 {
705   CmiLock(ComputePmeMgr::fftw_plan_lock);
706   zPencilObj.put(idx)=obj;
707   CmiUnlock(ComputePmeMgr::fftw_plan_lock);
708 }
709
710 ComputePmeMgr::ComputePmeMgr() : pmeProxy(thisgroup), 
711                                  pmeProxyDir(thisgroup) {
712
713   CkpvAccess(BOCclass_group).computePmeMgr = thisgroup;
714   pmeNodeProxy = CkpvAccess(BOCclass_group).nodePmeMgr;
715   nodePmeMgr = pmeNodeProxy[CkMyNode()].ckLocalBranch();
716
717   pmeNodeProxy.ckLocalBranch()->initialize();
718
719   if ( CmiMyRank() == 0 ) {
720     fftw_plan_lock = CmiCreateLock();
721   }
722   pmemgr_lock = CmiCreateLock();
723
724   myKSpace = 0;
725   kgrid = 0;
726   work = 0;
727   grid_count = 0;
728   trans_count = 0;
729   untrans_count = 0;
730   ungrid_count = 0;
731   gridmsg_reuse= new PmeGridMsg*[CkNumPes()];
732   useBarrier = 0;
733   sendTransBarrier_received = 0;
734   usePencils = 0;
735
736 #ifdef NAMD_CUDA
737  // offload has not been set so this happens on every run
738   if ( CmiMyRank() == 0 ) {
739     cuda_lock = CmiCreateLock();
740   }
741
742 #if CUDA_VERSION >= 5050
743   int leastPriority, greatestPriority;
744   cudaDeviceGetStreamPriorityRange(&leastPriority, &greatestPriority);
745   cuda_errcheck("in cudaDeviceGetStreamPriorityRange");
746   //if ( CkMyNode() == 0 ) {
747   //  CkPrintf("Pe %d PME CUDA stream priority range %d %d\n", CkMyPe(), leastPriority, greatestPriority);
748   //}
749 #define CUDA_STREAM_CREATE(X) cudaStreamCreateWithPriority(X,cudaStreamDefault,greatestPriority)
750 #else
751 #define CUDA_STREAM_CREATE(X) cudaStreamCreate(X)
752 #endif
753
754   stream = 0;
755   for ( int i=0; i<NUM_STREAMS; ++i ) {
756 #if 1
757     CUDA_STREAM_CREATE(&streams[i]);
758     cuda_errcheck("cudaStreamCreate");
759 #else
760   streams[i] = 0;  // XXXX Testing!!!
761 #endif
762   }
763
764   this_pe = CkMyPe();
765  
766   cudaEventCreateWithFlags(&end_charges,cudaEventDisableTiming);
767   end_forces = 0;
768   check_charges_count = 0;
769   check_forces_count = 0;
770   chargeGridSubmittedCount = 0;
771
772   cuda_atoms_count = 0;
773   cuda_atoms_alloc = 0;
774
775   f_data_mgr_alloc = 0;
776   f_data_mgr_host = 0;
777   f_data_mgr_dev = 0;
778   afn_host = 0;
779   afn_dev = 0;
780
781 #define CUDA_EVENT_ID_PME_CHARGES 80
782 #define CUDA_EVENT_ID_PME_FORCES 81
783 #define CUDA_EVENT_ID_PME_TICK 82
784 #define CUDA_EVENT_ID_PME_COPY 83
785 #define CUDA_EVENT_ID_PME_KERNEL 84
786   if ( 0 == CkMyPe() ) {
787     traceRegisterUserEvent("CUDA PME charges", CUDA_EVENT_ID_PME_CHARGES);
788     traceRegisterUserEvent("CUDA PME forces", CUDA_EVENT_ID_PME_FORCES);
789     traceRegisterUserEvent("CUDA PME tick", CUDA_EVENT_ID_PME_TICK);
790     traceRegisterUserEvent("CUDA PME memcpy", CUDA_EVENT_ID_PME_COPY);
791     traceRegisterUserEvent("CUDA PME kernel", CUDA_EVENT_ID_PME_KERNEL);
792   }
793 #endif
794   recipEvirCount = 0;
795   recipEvirClients = 0;
796   recipEvirPe = -999;
797 }
798
799
800 void ComputePmeMgr::recvArrays(
801         CProxy_PmeXPencil x, CProxy_PmeYPencil y, CProxy_PmeZPencil z) {
802   xPencil = x;  yPencil = y;  zPencil = z;
803   
804     if(CmiMyRank()==0)
805     {
806       pmeNodeProxy.ckLocalBranch()->xPencil=x;
807       pmeNodeProxy.ckLocalBranch()->yPencil=y;
808       pmeNodeProxy.ckLocalBranch()->zPencil=z;
809     }
810 }
811
812 #if USE_TOPO_SFC
813  struct Coord
814   {
815     int x, y, z;
816     Coord(): x(0), y(0), z(0) {}
817     Coord(int a, int b, int c): x(a), y(b), z(c) {}
818   };
819   extern void SFC_grid(int xdim, int ydim, int zdim, int xdim1, int ydim1, int zdim1, vector<Coord> &result);
820
821   void sort_sfc(SortableResizeArray<int> &procs, TopoManager &tmgr, vector<Coord> &result)
822   {
823      SortableResizeArray<int> newprocs(procs.size());
824      int num = 0;
825      for (int i=0; i<result.size(); i++) {
826        Coord &c = result[i];
827        for (int j=0; j<procs.size(); j++) {
828          int pe = procs[j];
829          int x,y,z,t;
830          tmgr.rankToCoordinates(pe, x, y, z, t);    
831          if (x==c.x && y==c.y && z==c.z)
832            newprocs[num++] = pe;
833        }
834      } 
835      CmiAssert(newprocs.size() == procs.size());
836      procs = newprocs;
837   }
838
839   int find_level_grid(int x) 
840   {
841      int a = sqrt(x);
842      int b;
843      for (; a>0; a--) {
844        if (x%a == 0) break;
845      }
846      if (a==1) a = x;
847      b = x/a;
848      //return a>b?a:b;
849      return b;
850   }
851   CmiNodeLock tmgr_lock;
852 #endif
853
854 void Pme_init()
855 {
856 #if USE_TOPO_SFC
857   if (CkMyRank() == 0) 
858     tmgr_lock = CmiCreateLock();
859 #endif
860 }
861
862 void ComputePmeMgr::initialize(CkQdMsg *msg) {
863   delete msg;
864
865   localInfo = new LocalPmeInfo[CkNumPes()];
866   gridNodeInfo = new NodePmeInfo[CkNumNodes()];
867   transNodeInfo = new NodePmeInfo[CkNumNodes()];
868   gridPeMap = new int[CkNumPes()];
869   transPeMap = new int[CkNumPes()];
870   recipPeDest = new int[CkNumPes()];
871   gridPeOrder = new int[CkNumPes()];
872   gridNodeOrder = new int[CkNumNodes()];
873   transNodeOrder = new int[CkNumNodes()];
874
875   if (CkMyRank() == 0) {
876     pencilPMEProcessors = new char [CkNumPes()];
877     memset (pencilPMEProcessors, 0, sizeof(char) * CkNumPes());
878   }
879
880   SimParameters *simParams = Node::Object()->simParameters;
881   PatchMap *patchMap = PatchMap::Object();
882
883   offload = simParams->PMEOffload;
884 #ifdef NAMD_CUDA
885   if ( offload && ! deviceCUDA->one_device_per_node() ) {
886     NAMD_die("PME offload requires exactly one CUDA device per process.  Use \"PMEOffload no\".");
887   }
888   if ( offload ) {
889     int dev;
890     cudaGetDevice(&dev);
891     cuda_errcheck("in cudaGetDevice");
892     if ( dev != deviceCUDA->getDeviceID() ) NAMD_bug("ComputePmeMgr::initialize dev != deviceCUDA->getDeviceID()");
893     cudaDeviceProp deviceProp;
894     cudaGetDeviceProperties(&deviceProp, dev);
895     cuda_errcheck("in cudaGetDeviceProperties");
896     if ( deviceProp.major < 2 )
897       NAMD_die("PME offload requires CUDA device of compute capability 2.0 or higher.  Use \"PMEOffload no\".");
898   }
899 #endif
900
901   alchLambda = -1.;  // illegal value to catch if not updated
902   alchLambda2 = -1.;
903   useBarrier = simParams->PMEBarrier;
904
905   if ( numGrids != 1 || simParams->PMEPencils == 0 ) usePencils = 0;
906   else if ( simParams->PMEPencils > 0 ) usePencils = 1;
907   else {
908     int nrps = simParams->PMEProcessors;
909     if ( nrps <= 0 ) nrps = CkNumPes();
910     if ( nrps > CkNumPes() ) nrps = CkNumPes();
911     int dimx = simParams->PMEGridSizeX;
912     int dimy = simParams->PMEGridSizeY;
913     int maxslabs = 1 + (dimx - 1) / simParams->PMEMinSlices;
914     if ( maxslabs > nrps ) maxslabs = nrps;
915     int maxpencils = ( simParams->PMEGridSizeX * (int64) simParams->PMEGridSizeY
916                 * simParams->PMEGridSizeZ ) / simParams->PMEMinPoints;
917     if ( maxpencils > nrps ) maxpencils = nrps;
918     if ( maxpencils > 3 * maxslabs ) usePencils = 1;
919     else usePencils = 0;
920   }
921
922   if ( usePencils ) {
923     int nrps = simParams->PMEProcessors;
924     if ( nrps <= 0 ) nrps = CkNumPes();
925     if ( nrps > CkNumPes() ) nrps = CkNumPes();
926     if ( simParams->PMEPencils > 1 &&
927          simParams->PMEPencils * simParams->PMEPencils <= nrps ) {
928       xBlocks = yBlocks = zBlocks = simParams->PMEPencils;
929     } else {
930       int nb2 = ( simParams->PMEGridSizeX * (int64) simParams->PMEGridSizeY
931                 * simParams->PMEGridSizeZ ) / simParams->PMEMinPoints;
932       if ( nb2 > nrps ) nb2 = nrps;
933       if ( nb2 < 1 ) nb2 = 1;
934       int nb = (int) sqrt((float)nb2);
935       if ( nb < 1 ) nb = 1;
936       xBlocks = zBlocks = nb;
937       yBlocks = nb2 / nb;
938     }
939
940     if ( simParams->PMEPencilsX > 0 ) xBlocks = simParams->PMEPencilsX;
941     if ( simParams->PMEPencilsY > 0 ) yBlocks = simParams->PMEPencilsY;
942     if ( simParams->PMEPencilsZ > 0 ) zBlocks = simParams->PMEPencilsZ;
943
944     int dimx = simParams->PMEGridSizeX;
945     int bx = 1 + ( dimx - 1 ) / xBlocks;
946     xBlocks = 1 + ( dimx - 1 ) / bx;
947
948     int dimy = simParams->PMEGridSizeY;
949     int by = 1 + ( dimy - 1 ) / yBlocks;
950     yBlocks = 1 + ( dimy - 1 ) / by;
951
952     int dimz = simParams->PMEGridSizeZ / 2 + 1;  // complex
953     int bz = 1 + ( dimz - 1 ) / zBlocks;
954     zBlocks = 1 + ( dimz - 1 ) / bz;
955
956     if ( xBlocks * yBlocks > CkNumPes() ) {
957       NAMD_die("PME pencils xBlocks * yBlocks > numPes");
958     }
959     if ( xBlocks * zBlocks > CkNumPes() ) {
960       NAMD_die("PME pencils xBlocks * zBlocks > numPes");
961     }
962     if ( yBlocks * zBlocks > CkNumPes() ) {
963       NAMD_die("PME pencils yBlocks * zBlocks > numPes");
964     }
965
966     if ( ! CkMyPe() ) {
967       iout << iINFO << "PME using " << xBlocks << " x " <<
968         yBlocks << " x " << zBlocks <<
969         " pencil grid for FFT and reciprocal sum.\n" << endi;
970     }
971   } else { // usePencils
972
973   {  // decide how many pes to use for reciprocal sum
974
975     // rules based on work available
976     int minslices = simParams->PMEMinSlices;
977     int dimx = simParams->PMEGridSizeX;
978     int nrpx = ( dimx + minslices - 1 ) / minslices;
979     int dimy = simParams->PMEGridSizeY;
980     int nrpy = ( dimy + minslices - 1 ) / minslices;
981
982     // rules based on processors available
983     int nrpp = CkNumPes();
984     // if ( nrpp > 32 ) nrpp = 32;  // cap to limit messages
985     if ( nrpp < nrpx ) nrpx = nrpp;
986     if ( nrpp < nrpy ) nrpy = nrpp;
987
988     // user override
989     int nrps = simParams->PMEProcessors;
990     if ( nrps > CkNumPes() ) nrps = CkNumPes();
991     if ( nrps > 0 ) nrpx = nrps;
992     if ( nrps > 0 ) nrpy = nrps;
993
994     // make sure there aren't any totally empty processors
995     int bx = ( dimx + nrpx - 1 ) / nrpx;
996     nrpx = ( dimx + bx - 1 ) / bx;
997     int by = ( dimy + nrpy - 1 ) / nrpy;
998     nrpy = ( dimy + by - 1 ) / by;
999     if ( bx != ( dimx + nrpx - 1 ) / nrpx )
1000       NAMD_bug("Error in selecting number of PME processors.");
1001     if ( by != ( dimy + nrpy - 1 ) / nrpy )
1002       NAMD_bug("Error in selecting number of PME processors.");
1003
1004     numGridPes = nrpx;
1005     numTransPes = nrpy;
1006   }
1007   if ( ! CkMyPe() ) {
1008     iout << iINFO << "PME using " << numGridPes << " and " << numTransPes <<
1009       " processors for FFT and reciprocal sum.\n" << endi;
1010   }
1011
1012   int sum_npes = numTransPes + numGridPes;
1013   int max_npes = (numTransPes > numGridPes)?numTransPes:numGridPes;
1014
1015 #if 0 // USE_TOPOMAP
1016   /* This code is being disabled permanently for slab PME on Blue Gene machines */
1017   PatchMap * pmap = PatchMap::Object();
1018   
1019   int patch_pes = pmap->numNodesWithPatches();
1020   TopoManager tmgr;
1021   if(tmgr.hasMultipleProcsPerNode())
1022     patch_pes *= 2;
1023
1024   bool done = false;
1025   if(CkNumPes() > 2*sum_npes + patch_pes) {    
1026     done = generateBGLORBPmePeList(transPeMap, numTransPes);
1027     done &= generateBGLORBPmePeList(gridPeMap, numGridPes, transPeMap, numTransPes);    
1028   }
1029   else 
1030     if(CkNumPes() > 2 *max_npes + patch_pes) {
1031       done = generateBGLORBPmePeList(transPeMap, max_npes);
1032       gridPeMap = transPeMap;
1033     }
1034
1035   if (!done)
1036 #endif
1037     {
1038       //generatePmePeList(transPeMap, max_npes);
1039       //gridPeMap = transPeMap;
1040       generatePmePeList2(gridPeMap, numGridPes, transPeMap, numTransPes);
1041     }
1042   
1043   if ( ! CkMyPe() ) {
1044     iout << iINFO << "PME GRID LOCATIONS:";
1045     int i;
1046     for ( i=0; i<numGridPes && i<10; ++i ) {
1047       iout << " " << gridPeMap[i];
1048     }
1049     if ( i < numGridPes ) iout << " ...";
1050     iout << "\n" << endi;
1051     iout << iINFO << "PME TRANS LOCATIONS:";
1052     for ( i=0; i<numTransPes && i<10; ++i ) {
1053       iout << " " << transPeMap[i];
1054     }
1055     if ( i < numTransPes ) iout << " ...";
1056     iout << "\n" << endi;
1057   }
1058
1059   // sort based on nodes and physical nodes
1060   std::sort(gridPeMap,gridPeMap+numGridPes,WorkDistrib::pe_sortop_compact());
1061
1062   myGridPe = -1;
1063   myGridNode = -1;
1064   int i = 0;
1065   int node = -1;
1066   int real_node = -1;
1067   for ( i=0; i<numGridPes; ++i ) {
1068     if ( gridPeMap[i] == CkMyPe() ) myGridPe = i;
1069     if (CkMyRank() == 0) pencilPMEProcessors[gridPeMap[i]] |= 1;
1070     int real_node_i = CkNodeOf(gridPeMap[i]);
1071     if ( real_node_i == real_node ) {
1072       gridNodeInfo[node].npe += 1;
1073     } else {
1074       real_node = real_node_i;
1075       ++node;
1076       gridNodeInfo[node].real_node = real_node;
1077       gridNodeInfo[node].pe_start = i;
1078       gridNodeInfo[node].npe = 1;
1079     }
1080     if ( CkMyNode() == real_node_i ) myGridNode = node;
1081   }
1082   numGridNodes = node + 1;
1083   myTransPe = -1;
1084   myTransNode = -1;
1085   node = -1;
1086   real_node = -1;
1087   for ( i=0; i<numTransPes; ++i ) {
1088     if ( transPeMap[i] == CkMyPe() ) myTransPe = i;
1089     if (CkMyRank() == 0) pencilPMEProcessors[transPeMap[i]] |= 2;
1090     int real_node_i = CkNodeOf(transPeMap[i]);
1091     if ( real_node_i == real_node ) {
1092       transNodeInfo[node].npe += 1;
1093     } else {
1094       real_node = real_node_i;
1095       ++node;
1096       transNodeInfo[node].real_node = real_node;
1097       transNodeInfo[node].pe_start = i;
1098       transNodeInfo[node].npe = 1;
1099     }
1100     if ( CkMyNode() == real_node_i ) myTransNode = node;
1101   }
1102   numTransNodes = node + 1;
1103
1104   if ( ! CkMyPe() ) {
1105     iout << iINFO << "PME USING " << numGridNodes << " GRID NODES AND "
1106          << numTransNodes << " TRANS NODES\n" << endi;
1107   }
1108
1109   { // generate random orderings for grid and trans messages
1110     int i;
1111     for ( i = 0; i < numGridPes; ++i ) {
1112       gridPeOrder[i] = i;
1113     }
1114     Random rand(CkMyPe());
1115     if ( myGridPe < 0 ) {
1116       rand.reorder(gridPeOrder,numGridPes);
1117     } else {  // self last
1118       gridPeOrder[myGridPe] = numGridPes-1;
1119       gridPeOrder[numGridPes-1] = myGridPe;
1120       rand.reorder(gridPeOrder,numGridPes-1);
1121     } 
1122     for ( i = 0; i < numGridNodes; ++i ) {
1123       gridNodeOrder[i] = i;
1124     }
1125     if ( myGridNode < 0 ) {
1126       rand.reorder(gridNodeOrder,numGridNodes);
1127     } else {  // self last
1128       gridNodeOrder[myGridNode] = numGridNodes-1;
1129       gridNodeOrder[numGridNodes-1] = myGridNode;
1130       rand.reorder(gridNodeOrder,numGridNodes-1);
1131     }
1132     for ( i = 0; i < numTransNodes; ++i ) {
1133       transNodeOrder[i] = i;
1134     }
1135     if ( myTransNode < 0 ) {
1136       rand.reorder(transNodeOrder,numTransNodes);
1137     } else {  // self last
1138       transNodeOrder[myTransNode] = numTransNodes-1;
1139       transNodeOrder[numTransNodes-1] = myTransNode;
1140       rand.reorder(transNodeOrder,numTransNodes-1);
1141     }
1142   }
1143   
1144   } // ! usePencils
1145
1146   myGrid.K1 = simParams->PMEGridSizeX;
1147   myGrid.K2 = simParams->PMEGridSizeY;
1148   myGrid.K3 = simParams->PMEGridSizeZ;
1149   myGrid.order = simParams->PMEInterpOrder;
1150   myGrid.dim2 = myGrid.K2;
1151   myGrid.dim3 = 2 * (myGrid.K3/2 + 1);
1152
1153   if ( ! usePencils ) {
1154     myGrid.block1 = ( myGrid.K1 + numGridPes - 1 ) / numGridPes;
1155     myGrid.block2 = ( myGrid.K2 + numTransPes - 1 ) / numTransPes;
1156     myGrid.block3 = myGrid.dim3 / 2;  // complex
1157   }
1158
1159   if ( usePencils ) {
1160     myGrid.block1 = ( myGrid.K1 + xBlocks - 1 ) / xBlocks;
1161     myGrid.block2 = ( myGrid.K2 + yBlocks - 1 ) / yBlocks;
1162     myGrid.block3 = ( myGrid.K3/2 + 1 + zBlocks - 1 ) / zBlocks;  // complex
1163
1164
1165       int pe = 0;
1166       int x,y,z;
1167
1168                 SortableResizeArray<int> zprocs(xBlocks*yBlocks);
1169                 SortableResizeArray<int> yprocs(xBlocks*zBlocks);
1170                 SortableResizeArray<int> xprocs(yBlocks*zBlocks);
1171       
1172                 // decide which pes to use by bit reversal and patch use
1173                 int i;
1174                 int ncpus = CkNumPes();
1175                 SortableResizeArray<int> patches, nopatches, pmeprocs;
1176                 PatchMap *pmap = PatchMap::Object();
1177                 for ( int icpu=0; icpu<ncpus; ++icpu ) {
1178                         int ri = WorkDistrib::peDiffuseOrdering[icpu];
1179                         if ( ri ) { // keep 0 for special case
1180                                 // pretend pe 1 has patches to avoid placing extra PME load on node
1181                                 if ( ri == 1 || pmap->numPatchesOnNode(ri) ) patches.add(ri);
1182                                 else nopatches.add(ri);
1183                         }
1184                 }
1185
1186 #if USE_RANDOM_TOPO
1187             Random rand(CkMyPe());
1188             int *tmp = new int[patches.size()];
1189             int nn = patches.size();
1190             for (i=0;i<nn;i++)  tmp[i] = patches[i];
1191             rand.reorder(tmp, nn);
1192             patches.resize(0);
1193             for (i=0;i<nn;i++)  patches.add(tmp[i]);
1194             delete [] tmp;
1195             tmp = new int[nopatches.size()];
1196             nn = nopatches.size();
1197             for (i=0;i<nn;i++)  tmp[i] = nopatches[i];
1198             rand.reorder(tmp, nn);
1199             nopatches.resize(0);
1200             for (i=0;i<nn;i++)  nopatches.add(tmp[i]);
1201             delete [] tmp;
1202 #endif
1203
1204                 // only use zero if it eliminates overloading or has patches
1205                 int useZero = 0;
1206                 int npens = xBlocks*yBlocks;
1207                 if ( npens % ncpus == 0 ) useZero = 1;
1208                 if ( npens == nopatches.size() + 1 ) useZero = 1;
1209                 npens += xBlocks*zBlocks;
1210                 if ( npens % ncpus == 0 ) useZero = 1;
1211                 if ( npens == nopatches.size() + 1 ) useZero = 1;
1212                 npens += yBlocks*zBlocks;
1213                 if ( npens % ncpus == 0 ) useZero = 1;
1214                 if ( npens == nopatches.size() + 1 ) useZero = 1;
1215
1216                 // add nopatches then patches in reversed order
1217                 for ( i=nopatches.size()-1; i>=0; --i ) pmeprocs.add(nopatches[i]);
1218                 if ( useZero && ! pmap->numPatchesOnNode(0) ) pmeprocs.add(0);
1219                 for ( i=patches.size()-1; i>=0; --i ) pmeprocs.add(patches[i]);
1220                 if ( pmap->numPatchesOnNode(0) ) pmeprocs.add(0);
1221   
1222                 int npes = pmeprocs.size();
1223                 for ( i=0; i<xBlocks*yBlocks; ++i, ++pe ) zprocs[i] = pmeprocs[pe%npes];
1224                 if ( i>1 && zprocs[0] == zprocs[i-1] ) zprocs[0] = 0;
1225 #if !USE_RANDOM_TOPO
1226                 zprocs.sort();
1227 #endif
1228                 for ( i=0; i<xBlocks*zBlocks; ++i, ++pe ) yprocs[i] = pmeprocs[pe%npes];
1229                 if ( i>1 && yprocs[0] == yprocs[i-1] ) yprocs[0] = 0;
1230 #if !USE_RANDOM_TOPO
1231                 yprocs.sort();
1232 #endif
1233       for ( i=0; i<yBlocks*zBlocks; ++i, ++pe ) xprocs[i] = pmeprocs[pe%npes];
1234       if ( i>1 && xprocs[0] == xprocs[i-1] ) xprocs[0] = 0;
1235 #if !USE_RANDOM_TOPO
1236       xprocs.sort();
1237 #endif
1238
1239 #if USE_TOPO_SFC
1240   CmiLock(tmgr_lock);
1241   //{
1242   TopoManager tmgr;
1243   int xdim = tmgr.getDimNX();
1244   int ydim = tmgr.getDimNY();
1245   int zdim = tmgr.getDimNZ();
1246   int xdim1 = find_level_grid(xdim);
1247   int ydim1 = find_level_grid(ydim);
1248   int zdim1 = find_level_grid(zdim);
1249   if(CkMyPe() == 0)
1250       printf("xdim: %d %d %d, %d %d %d\n", xdim, ydim, zdim, xdim1, ydim1, zdim1);
1251
1252   vector<Coord> result;
1253   SFC_grid(xdim, ydim, zdim, xdim1, ydim1, zdim1, result);
1254   sort_sfc(xprocs, tmgr, result);
1255   sort_sfc(yprocs, tmgr, result);
1256   sort_sfc(zprocs, tmgr, result);
1257   //}
1258   CmiUnlock(tmgr_lock);
1259 #endif
1260
1261
1262                 if(CkMyPe() == 0){  
1263               iout << iINFO << "PME Z PENCIL LOCATIONS:";
1264           for ( i=0; i<zprocs.size() && i<10; ++i ) {
1265 #if USE_TOPO_SFC
1266               int x,y,z,t;
1267               tmgr.rankToCoordinates(zprocs[i], x,y, z, t);
1268               iout << " " << zprocs[i] << "(" << x << " " << y << " " << z << ")";
1269 #else
1270               iout << " " << zprocs[i];
1271 #endif
1272           }
1273           if ( i < zprocs.size() ) iout << " ...";
1274               iout << "\n" << endi;
1275                 }
1276
1277     if (CkMyRank() == 0) {
1278       for (pe=0, x = 0; x < xBlocks; ++x)
1279         for (y = 0; y < yBlocks; ++y, ++pe ) {
1280           pencilPMEProcessors[zprocs[pe]] = 1;
1281         }
1282     }
1283      
1284                 if(CkMyPe() == 0){  
1285               iout << iINFO << "PME Y PENCIL LOCATIONS:";
1286           for ( i=0; i<yprocs.size() && i<10; ++i ) {
1287 #if USE_TOPO_SFC
1288               int x,y,z,t;
1289               tmgr.rankToCoordinates(yprocs[i], x,y, z, t);
1290               iout << " " << yprocs[i] << "(" << x << " " << y << " " << z << ")";
1291 #else
1292               iout << " " << yprocs[i];
1293 #endif
1294           }
1295           if ( i < yprocs.size() ) iout << " ...";
1296               iout << "\n" << endi;
1297                 }
1298
1299     if (CkMyRank() == 0) {
1300       for (pe=0, z = 0; z < zBlocks; ++z )
1301         for (x = 0; x < xBlocks; ++x, ++pe ) {
1302           pencilPMEProcessors[yprocs[pe]] = 1;
1303         }
1304     }
1305     
1306                 if(CkMyPe() == 0){  
1307                 iout << iINFO << "PME X PENCIL LOCATIONS:";
1308                     for ( i=0; i<xprocs.size() && i<10; ++i ) {
1309 #if USE_TOPO_SFC
1310                 int x,y,z,t;
1311                 tmgr.rankToCoordinates(xprocs[i], x,y, z, t);
1312                 iout << " " << xprocs[i] << "(" << x << "  " << y << " " << z << ")";
1313 #else
1314                 iout << " " << xprocs[i];
1315 #endif
1316             }
1317                 if ( i < xprocs.size() ) iout << " ...";
1318                 iout << "\n" << endi;
1319                 }
1320
1321     if (CkMyRank() == 0) {
1322       for (pe=0, y = 0; y < yBlocks; ++y )      
1323         for (z = 0; z < zBlocks; ++z, ++pe ) {
1324           pencilPMEProcessors[xprocs[pe]] = 1;
1325         }
1326     }
1327         
1328
1329         // creating the pencil arrays
1330         if ( CkMyPe() == 0 ){
1331 #if !USE_RANDOM_TOPO
1332         // std::sort(zprocs.begin(),zprocs.end(),WorkDistrib::pe_sortop_compact());
1333         WorkDistrib::sortPmePes(zprocs.begin(),xBlocks,yBlocks);
1334         std::sort(yprocs.begin(),yprocs.end(),WorkDistrib::pe_sortop_compact());
1335         std::sort(xprocs.begin(),xprocs.end(),WorkDistrib::pe_sortop_compact());
1336 #endif
1337 #if 1
1338         CProxy_PmePencilMap zm = CProxy_PmePencilMap::ckNew(0,1,yBlocks,xBlocks*yBlocks,zprocs.begin());
1339         CProxy_PmePencilMap ym;
1340         if ( simParams->PMEPencilsYLayout )
1341           ym = CProxy_PmePencilMap::ckNew(0,2,zBlocks,zBlocks*xBlocks,yprocs.begin()); // new
1342         else
1343           ym = CProxy_PmePencilMap::ckNew(2,0,xBlocks,zBlocks*xBlocks,yprocs.begin()); // old
1344         CProxy_PmePencilMap xm;
1345         if ( simParams->PMEPencilsXLayout )
1346           xm = CProxy_PmePencilMap::ckNew(2,1,yBlocks,yBlocks*zBlocks,xprocs.begin()); // new
1347         else
1348           xm = CProxy_PmePencilMap::ckNew(1,2,zBlocks,yBlocks*zBlocks,xprocs.begin()); // old
1349         pmeNodeProxy.recvPencilMapProxies(xm,ym,zm);
1350         CkArrayOptions zo(xBlocks,yBlocks,1);  zo.setMap(zm);
1351         CkArrayOptions yo(xBlocks,1,zBlocks);  yo.setMap(ym);
1352         CkArrayOptions xo(1,yBlocks,zBlocks);  xo.setMap(xm);
1353         zo.setAnytimeMigration(false);  zo.setStaticInsertion(true);
1354         yo.setAnytimeMigration(false);  yo.setStaticInsertion(true);
1355         xo.setAnytimeMigration(false);  xo.setStaticInsertion(true);
1356         zPencil = CProxy_PmeZPencil::ckNew(zo);  // (xBlocks,yBlocks,1);
1357         yPencil = CProxy_PmeYPencil::ckNew(yo);  // (xBlocks,1,zBlocks);
1358         xPencil = CProxy_PmeXPencil::ckNew(xo);  // (1,yBlocks,zBlocks);
1359 #else
1360         zPencil = CProxy_PmeZPencil::ckNew();  // (xBlocks,yBlocks,1);
1361         yPencil = CProxy_PmeYPencil::ckNew();  // (xBlocks,1,zBlocks);
1362         xPencil = CProxy_PmeXPencil::ckNew();  // (1,yBlocks,zBlocks);
1363
1364                 for (pe=0, x = 0; x < xBlocks; ++x)
1365                         for (y = 0; y < yBlocks; ++y, ++pe ) {
1366                                 zPencil(x,y,0).insert(zprocs[pe]);
1367                         }
1368         zPencil.doneInserting();
1369
1370                 for (pe=0, x = 0; x < xBlocks; ++x)
1371                         for (z = 0; z < zBlocks; ++z, ++pe ) {
1372                                 yPencil(x,0,z).insert(yprocs[pe]);
1373                         }
1374         yPencil.doneInserting();
1375
1376
1377                 for (pe=0, y = 0; y < yBlocks; ++y )    
1378                         for (z = 0; z < zBlocks; ++z, ++pe ) {
1379                                 xPencil(0,y,z).insert(xprocs[pe]);
1380                         }
1381                 xPencil.doneInserting();     
1382 #endif
1383
1384                 pmeProxy.recvArrays(xPencil,yPencil,zPencil);
1385                 PmePencilInitMsgData msgdata;
1386                 msgdata.grid = myGrid;
1387                 msgdata.xBlocks = xBlocks;
1388                 msgdata.yBlocks = yBlocks;
1389                 msgdata.zBlocks = zBlocks;
1390                 msgdata.xPencil = xPencil;
1391                 msgdata.yPencil = yPencil;
1392                 msgdata.zPencil = zPencil;
1393                 msgdata.pmeProxy = pmeProxyDir;
1394         msgdata.pmeNodeProxy = pmeNodeProxy;
1395         msgdata.xm = xm;
1396         msgdata.ym = ym;
1397         msgdata.zm = zm;
1398                 xPencil.init(new PmePencilInitMsg(msgdata));
1399                 yPencil.init(new PmePencilInitMsg(msgdata));
1400                 zPencil.init(new PmePencilInitMsg(msgdata));
1401         }
1402
1403     return;  // continue in initialize_pencils() at next startup stage
1404   }
1405
1406
1407   int pe;
1408   int nx = 0;
1409   for ( pe = 0; pe < numGridPes; ++pe ) {
1410     localInfo[pe].x_start = nx;
1411     nx += myGrid.block1;
1412     if ( nx > myGrid.K1 ) nx = myGrid.K1;
1413     localInfo[pe].nx = nx - localInfo[pe].x_start;
1414   }
1415   int ny = 0;
1416   for ( pe = 0; pe < numTransPes; ++pe ) {
1417     localInfo[pe].y_start_after_transpose = ny;
1418     ny += myGrid.block2;
1419     if ( ny > myGrid.K2 ) ny = myGrid.K2;
1420     localInfo[pe].ny_after_transpose =
1421                         ny - localInfo[pe].y_start_after_transpose;
1422   }
1423
1424   {  // decide how many pes this node exchanges charges with
1425
1426   PatchMap *patchMap = PatchMap::Object();
1427   Lattice lattice = simParams->lattice;
1428   BigReal sysdima = lattice.a_r().unit() * lattice.a();
1429   BigReal cutoff = simParams->cutoff;
1430   BigReal patchdim = simParams->patchDimension;
1431   int numPatches = patchMap->numPatches();
1432   int numNodes = CkNumPes();
1433   int *source_flags = new int[numNodes];
1434   int node;
1435   for ( node=0; node<numNodes; ++node ) {
1436     source_flags[node] = 0;
1437     recipPeDest[node] = 0;
1438   }
1439
1440   // // make sure that we don't get ahead of ourselves on this node
1441   // if ( CkMyPe() < numPatches && myRecipPe >= 0 ) {
1442   //   source_flags[CkMyPe()] = 1;
1443   //   recipPeDest[myRecipPe] = 1;
1444   // }
1445
1446   for ( int pid=0; pid < numPatches; ++pid ) {
1447     int pnode = patchMap->node(pid);
1448 #ifdef NAMD_CUDA
1449     if ( offload ) pnode = CkNodeFirst(CkNodeOf(pnode));
1450 #endif
1451     int shift1 = (myGrid.K1 + myGrid.order - 1)/2;
1452     BigReal minx = patchMap->min_a(pid);
1453     BigReal maxx = patchMap->max_a(pid);
1454     BigReal margina = 0.5 * ( patchdim - cutoff ) / sysdima;
1455     // min1 (max1) is smallest (largest) grid line for this patch
1456     int min1 = ((int) floor(myGrid.K1 * (minx - margina))) + shift1 - myGrid.order + 1;
1457     int max1 = ((int) floor(myGrid.K1 * (maxx + margina))) + shift1;
1458     for ( int i=min1; i<=max1; ++i ) {
1459       int ix = i;
1460       while ( ix >= myGrid.K1 ) ix -= myGrid.K1;
1461       while ( ix < 0 ) ix += myGrid.K1;
1462       // set source_flags[pnode] if this patch sends to our node
1463       if ( myGridPe >= 0 && ix >= localInfo[myGridPe].x_start &&
1464            ix < localInfo[myGridPe].x_start + localInfo[myGridPe].nx ) {
1465         source_flags[pnode] = 1;
1466       }
1467       // set dest_flags[] for node that our patch sends to
1468 #ifdef NAMD_CUDA
1469       if ( offload ) {
1470         if ( pnode == CkNodeFirst(CkMyNode()) ) {
1471           recipPeDest[ix / myGrid.block1] = 1;
1472         }
1473       } else
1474 #endif
1475       if ( pnode == CkMyPe() ) {
1476         recipPeDest[ix / myGrid.block1] = 1;
1477       }
1478     }
1479   }
1480
1481   int numSourcesSamePhysicalNode = 0;
1482   numSources = 0;
1483   numDestRecipPes = 0;
1484   for ( node=0; node<numNodes; ++node ) {
1485     if ( source_flags[node] ) ++numSources;
1486     if ( recipPeDest[node] ) ++numDestRecipPes;
1487     if ( source_flags[node] && CmiPeOnSamePhysicalNode(node,CkMyPe()) ) ++numSourcesSamePhysicalNode;
1488   }
1489
1490 #if 0
1491   if ( numSources ) {
1492     CkPrintf("pe %5d pme %5d of %5d on same physical node\n",
1493             CkMyPe(), numSourcesSamePhysicalNode, numSources);
1494     iout << iINFO << "PME " << CkMyPe() << " sources:";
1495     for ( node=0; node<numNodes; ++node ) {
1496       if ( source_flags[node] ) iout << " " << node;
1497     }
1498     iout << "\n" << endi;
1499   }
1500 #endif
1501
1502   delete [] source_flags;
1503
1504   // CkPrintf("PME on node %d has %d sources and %d destinations\n",
1505   //           CkMyPe(), numSources, numDestRecipPes);
1506
1507   }  // decide how many pes this node exchanges charges with (end)
1508
1509   ungrid_count = numDestRecipPes;
1510
1511   sendTransBarrier_received = 0;
1512
1513   if ( myGridPe < 0 && myTransPe < 0 ) return;
1514   // the following only for nodes doing reciprocal sum
1515
1516   if ( myTransPe >= 0 ) {
1517     recipEvirPe = findRecipEvirPe();
1518     pmeProxy[recipEvirPe].addRecipEvirClient();
1519   }
1520
1521   if ( myTransPe >= 0 ) {
1522       int k2_start = localInfo[myTransPe].y_start_after_transpose;
1523       int k2_end = k2_start + localInfo[myTransPe].ny_after_transpose;
1524       #ifdef OPENATOM_VERSION
1525       if ( simParams->openatomOn ) { 
1526         CProxy_ComputeMoaMgr moaProxy(CkpvAccess(BOCclass_group).computeMoaMgr);
1527         myKSpace = new PmeKSpace(myGrid, k2_start, k2_end, 0, myGrid.dim3/2, moaProxy);
1528       } else {
1529         myKSpace = new PmeKSpace(myGrid, k2_start, k2_end, 0, myGrid.dim3/2);
1530       }
1531       #else  // OPENATOM_VERSION
1532       myKSpace = new PmeKSpace(myGrid, k2_start, k2_end, 0, myGrid.dim3/2);
1533       #endif // OPENATOM_VERSION
1534   }
1535
1536   int local_size = myGrid.block1 * myGrid.K2 * myGrid.dim3;
1537   int local_size_2 = myGrid.block2 * myGrid.K1 * myGrid.dim3;
1538   if ( local_size < local_size_2 ) local_size = local_size_2;
1539   qgrid = new float[local_size*numGrids];
1540   if ( numGridPes > 1 || numTransPes > 1 ) {
1541     kgrid = new float[local_size*numGrids];
1542   } else {
1543     kgrid = qgrid;
1544   }
1545   qgrid_size = local_size;
1546
1547   if ( myGridPe >= 0 ) {
1548   qgrid_start = localInfo[myGridPe].x_start * myGrid.K2 * myGrid.dim3;
1549   qgrid_len = localInfo[myGridPe].nx * myGrid.K2 * myGrid.dim3;
1550   fgrid_start = localInfo[myGridPe].x_start * myGrid.K2;
1551   fgrid_len = localInfo[myGridPe].nx * myGrid.K2;
1552   }
1553
1554   int n[3]; n[0] = myGrid.K1; n[1] = myGrid.K2; n[2] = myGrid.K3;
1555 #ifdef NAMD_FFTW
1556   CmiLock(fftw_plan_lock);
1557 #ifdef NAMD_FFTW_3
1558   work = new fftwf_complex[n[0]];
1559   int fftwFlags = simParams->FFTWPatient ? FFTW_PATIENT  : simParams->FFTWEstimate ? FFTW_ESTIMATE  : FFTW_MEASURE ;
1560   if ( myGridPe >= 0 ) {
1561     forward_plan_yz=new fftwf_plan[numGrids];
1562     backward_plan_yz=new fftwf_plan[numGrids];
1563   }
1564   if ( myTransPe >= 0 ) {
1565     forward_plan_x=new fftwf_plan[numGrids];
1566     backward_plan_x=new fftwf_plan[numGrids];
1567   }
1568   /* need one plan per grid */
1569   if ( ! CkMyPe() ) iout << iINFO << "Optimizing 4 FFT steps.  1..." << endi;
1570   if ( myGridPe >= 0 ) {
1571     for( int g=0; g<numGrids; g++)
1572       {
1573         forward_plan_yz[g] = fftwf_plan_many_dft_r2c(2, n+1, 
1574                                                      localInfo[myGridPe].nx,
1575                                                      qgrid + qgrid_size * g,
1576                                                      NULL,
1577                                                      1,
1578                                                      myGrid.dim2 * myGrid.dim3,
1579                                                      (fftwf_complex *) 
1580                                                      (qgrid + qgrid_size * g),
1581                                                      NULL,
1582                                                      1,
1583                                                      myGrid.dim2 * (myGrid.dim3/2),
1584                                                      fftwFlags);
1585       }
1586   }
1587   int zdim = myGrid.dim3;
1588   int xStride=localInfo[myTransPe].ny_after_transpose *( myGrid.dim3 / 2);
1589   if ( ! CkMyPe() ) iout << " 2..." << endi;
1590   if ( myTransPe >= 0 ) {
1591     for( int g=0; g<numGrids; g++)
1592       {
1593
1594         forward_plan_x[g] = fftwf_plan_many_dft(1, n, xStride,
1595                                                 (fftwf_complex *)
1596                                                 (kgrid+qgrid_size*g),
1597                                                 NULL,
1598                                                 xStride,
1599                                                 1,
1600                                                 (fftwf_complex *)
1601                                                 (kgrid+qgrid_size*g),
1602                                                 NULL,
1603                                                 xStride,
1604                                                 1,
1605                                                 FFTW_FORWARD,fftwFlags);
1606         
1607       }
1608   }
1609   if ( ! CkMyPe() ) iout << " 3..." << endi;
1610   if ( myTransPe >= 0 ) {
1611     for( int g=0; g<numGrids; g++)
1612       {
1613         backward_plan_x[g] = fftwf_plan_many_dft(1, n, xStride,
1614                                                  (fftwf_complex *)
1615                                                  (kgrid+qgrid_size*g),
1616                                                  NULL,
1617                                                  xStride,
1618                                                  1,
1619                                                  (fftwf_complex *)
1620                                                  (kgrid+qgrid_size*g),
1621                                                  NULL,
1622                                                  xStride,
1623                                                  1,
1624                                                  FFTW_BACKWARD, fftwFlags);
1625
1626       }
1627   }
1628   if ( ! CkMyPe() ) iout << " 4..." << endi;
1629   if ( myGridPe >= 0 ) {
1630     for( int g=0; g<numGrids; g++)
1631       {
1632         backward_plan_yz[g] = fftwf_plan_many_dft_c2r(2, n+1, 
1633                                                       localInfo[myGridPe].nx,
1634                                                       (fftwf_complex *)
1635                                                       (qgrid + qgrid_size * g),
1636                                                       NULL,
1637                                                       1,
1638                                                       myGrid.dim2*(myGrid.dim3/2),
1639                                                       qgrid + qgrid_size * g,
1640                                                       NULL,
1641                                                       1,
1642                                                       myGrid.dim2 * myGrid.dim3,
1643                                                       fftwFlags);
1644       }
1645   }
1646   if ( ! CkMyPe() ) iout << "   Done.\n" << endi;
1647
1648 #else
1649   work = new fftw_complex[n[0]];
1650
1651   if ( ! CkMyPe() ) iout << iINFO << "Optimizing 4 FFT steps.  1..." << endi;
1652   if ( myGridPe >= 0 ) {
1653   forward_plan_yz = rfftwnd_create_plan_specific(2, n+1, FFTW_REAL_TO_COMPLEX,
1654         ( simParams->FFTWEstimate ? FFTW_ESTIMATE : FFTW_MEASURE )
1655         | FFTW_IN_PLACE | FFTW_USE_WISDOM, qgrid, 1, 0, 0);
1656   }
1657   if ( ! CkMyPe() ) iout << " 2..." << endi;
1658   if ( myTransPe >= 0 ) {
1659       forward_plan_x = fftw_create_plan_specific(n[0], FFTW_FORWARD,
1660         ( simParams->FFTWEstimate ? FFTW_ESTIMATE : FFTW_MEASURE )
1661         | FFTW_IN_PLACE | FFTW_USE_WISDOM, (fftw_complex *) kgrid,
1662         localInfo[myTransPe].ny_after_transpose * myGrid.dim3 / 2, work, 1);
1663   }
1664   if ( ! CkMyPe() ) iout << " 3..." << endi;
1665   if ( myTransPe >= 0 ) {
1666   backward_plan_x = fftw_create_plan_specific(n[0], FFTW_BACKWARD,
1667         ( simParams->FFTWEstimate ? FFTW_ESTIMATE : FFTW_MEASURE )
1668         | FFTW_IN_PLACE | FFTW_USE_WISDOM, (fftw_complex *) kgrid,
1669         localInfo[myTransPe].ny_after_transpose * myGrid.dim3 / 2, work, 1);
1670   }
1671   if ( ! CkMyPe() ) iout << " 4..." << endi;
1672   if ( myGridPe >= 0 ) {
1673   backward_plan_yz = rfftwnd_create_plan_specific(2, n+1, FFTW_COMPLEX_TO_REAL,
1674         ( simParams->FFTWEstimate ? FFTW_ESTIMATE : FFTW_MEASURE )
1675         | FFTW_IN_PLACE | FFTW_USE_WISDOM, qgrid, 1, 0, 0);
1676   }
1677   if ( ! CkMyPe() ) iout << "   Done.\n" << endi;
1678 #endif
1679   CmiUnlock(fftw_plan_lock);
1680 #else
1681   NAMD_die("Sorry, FFTW must be compiled in to use PME.");
1682 #endif
1683
1684   if ( myGridPe >= 0 && numSources == 0 )
1685                 NAMD_bug("PME grid elements exist without sources.");
1686   grid_count = numSources;
1687   memset( (void*) qgrid, 0, qgrid_size * numGrids * sizeof(float) );
1688   trans_count = numGridPes;
1689 }
1690
1691
1692
1693 void ComputePmeMgr::initialize_pencils(CkQdMsg *msg) {
1694   delete msg;
1695   if ( ! usePencils ) return;
1696
1697   SimParameters *simParams = Node::Object()->simParameters;
1698
1699   PatchMap *patchMap = PatchMap::Object();
1700   Lattice lattice = simParams->lattice;
1701   BigReal sysdima = lattice.a_r().unit() * lattice.a();
1702   BigReal sysdimb = lattice.b_r().unit() * lattice.b();
1703   BigReal cutoff = simParams->cutoff;
1704   BigReal patchdim = simParams->patchDimension;
1705   int numPatches = patchMap->numPatches();
1706
1707   pencilActive = new char[xBlocks*yBlocks];
1708   for ( int i=0; i<xBlocks; ++i ) {
1709     for ( int j=0; j<yBlocks; ++j ) {
1710       pencilActive[i*yBlocks+j] = 0;
1711     }
1712   }
1713
1714   for ( int pid=0; pid < numPatches; ++pid ) {
1715     int pnode = patchMap->node(pid);
1716 #ifdef NAMD_CUDA
1717     if ( offload ) {
1718       if ( CkNodeOf(pnode) != CkMyNode() ) continue;
1719     } else
1720 #endif
1721     if ( pnode != CkMyPe() ) continue;
1722
1723     int shift1 = (myGrid.K1 + myGrid.order - 1)/2;
1724     int shift2 = (myGrid.K2 + myGrid.order - 1)/2;
1725
1726     BigReal minx = patchMap->min_a(pid);
1727     BigReal maxx = patchMap->max_a(pid);
1728     BigReal margina = 0.5 * ( patchdim - cutoff ) / sysdima;
1729     // min1 (max1) is smallest (largest) grid line for this patch
1730     int min1 = ((int) floor(myGrid.K1 * (minx - margina))) + shift1 - myGrid.order + 1;
1731     int max1 = ((int) floor(myGrid.K1 * (maxx + margina))) + shift1;
1732
1733     BigReal miny = patchMap->min_b(pid);
1734     BigReal maxy = patchMap->max_b(pid);
1735     BigReal marginb = 0.5 * ( patchdim - cutoff ) / sysdimb;
1736     // min2 (max2) is smallest (largest) grid line for this patch
1737     int min2 = ((int) floor(myGrid.K2 * (miny - marginb))) + shift2 - myGrid.order + 1;
1738     int max2 = ((int) floor(myGrid.K2 * (maxy + marginb))) + shift2;
1739
1740     for ( int i=min1; i<=max1; ++i ) {
1741       int ix = i;
1742       while ( ix >= myGrid.K1 ) ix -= myGrid.K1;
1743       while ( ix < 0 ) ix += myGrid.K1;
1744       for ( int j=min2; j<=max2; ++j ) {
1745         int jy = j;
1746         while ( jy >= myGrid.K2 ) jy -= myGrid.K2;
1747         while ( jy < 0 ) jy += myGrid.K2;
1748         pencilActive[(ix / myGrid.block1)*yBlocks + (jy / myGrid.block2)] = 1;
1749       }
1750     }
1751   }
1752
1753   numPencilsActive = 0;
1754   for ( int i=0; i<xBlocks; ++i ) {
1755     for ( int j=0; j<yBlocks; ++j ) {
1756       if ( pencilActive[i*yBlocks+j] ) {
1757         ++numPencilsActive;
1758 #ifdef NAMD_CUDA
1759         if ( CkMyPe() == deviceCUDA->getMasterPe() || ! offload )
1760 #endif
1761         zPencil(i,j,0).dummyRecvGrid(CkMyPe(),0);
1762       }
1763     }
1764   }
1765   activePencils = new ijpair[numPencilsActive];
1766   numPencilsActive = 0;
1767   for ( int i=0; i<xBlocks; ++i ) {
1768     for ( int j=0; j<yBlocks; ++j ) {
1769       if ( pencilActive[i*yBlocks+j] ) {
1770         activePencils[numPencilsActive++] = ijpair(i,j);
1771       }
1772     }
1773   }
1774   if ( simParams->PMESendOrder ) {
1775     std::sort(activePencils,activePencils+numPencilsActive,ijpair_sortop_bit_reversed());
1776   } else {
1777     Random rand(CkMyPe());
1778     rand.reorder(activePencils,numPencilsActive);
1779   }
1780   //if ( numPencilsActive ) {
1781   //  CkPrintf("node %d sending to %d pencils\n", CkMyPe(), numPencilsActive);
1782   //}
1783
1784   ungrid_count = numPencilsActive;
1785 }
1786
1787
1788 void ComputePmeMgr::activate_pencils(CkQdMsg *msg) {
1789   if ( ! usePencils ) return;
1790   if ( CkMyPe() == 0 ) zPencil.dummyRecvGrid(CkMyPe(),1);
1791 }
1792
1793
1794 ComputePmeMgr::~ComputePmeMgr() {
1795
1796   if ( CmiMyRank() == 0 ) {
1797     CmiDestroyLock(fftw_plan_lock);
1798   }
1799   CmiDestroyLock(pmemgr_lock);
1800
1801   delete myKSpace;
1802   delete [] localInfo;
1803   delete [] gridNodeInfo;
1804   delete [] transNodeInfo;
1805   delete [] gridPeMap;
1806   delete [] transPeMap;
1807   delete [] recipPeDest;
1808   delete [] gridPeOrder;
1809   delete [] gridNodeOrder;
1810   delete [] transNodeOrder;
1811   delete [] qgrid;
1812   if ( kgrid != qgrid ) delete [] kgrid;
1813   delete [] work;
1814   delete [] gridmsg_reuse;
1815
1816  if ( ! offload ) {
1817   for (int i=0; i<q_count; ++i) {
1818     delete [] q_list[i];
1819   }
1820   delete [] q_list;
1821   delete [] fz_arr;
1822  }
1823   delete [] f_arr;
1824   delete [] q_arr;
1825 }
1826
1827 void ComputePmeMgr::recvGrid(PmeGridMsg *msg) {
1828   // CkPrintf("recvGrid from %d on Pe(%d)\n",msg->sourceNode,CkMyPe());
1829   if ( grid_count == 0 ) {
1830     NAMD_bug("Message order failure in ComputePmeMgr::recvGrid\n");
1831   }
1832   if ( grid_count == numSources ) {
1833     lattice = msg->lattice;
1834     grid_sequence = msg->sequence;
1835   }
1836
1837   int zdim = myGrid.dim3;
1838   int zlistlen = msg->zlistlen;
1839   int *zlist = msg->zlist;
1840   float *qmsg = msg->qgrid;
1841   for ( int g=0; g<numGrids; ++g ) {
1842     char *f = msg->fgrid + fgrid_len * g;
1843     float *q = qgrid + qgrid_size * g;
1844     for ( int i=0; i<fgrid_len; ++i ) {
1845       if ( f[i] ) {
1846         for ( int k=0; k<zlistlen; ++k ) {
1847           q[zlist[k]] += *(qmsg++);
1848         }
1849       }
1850       q += zdim;
1851     }
1852   }
1853
1854   gridmsg_reuse[numSources-grid_count] = msg;
1855   --grid_count;
1856
1857   if ( grid_count == 0 ) {
1858     pmeProxyDir[CkMyPe()].gridCalc1();
1859     if ( useBarrier ) pmeProxyDir[0].sendTransBarrier();
1860   }
1861 }
1862 #ifdef MANUAL_DEBUG_FFTW3
1863
1864 /* utility functions for manual debugging */
1865 void dumpMatrixFloat(const char *infilename, float *matrix, int xdim, int ydim, int zdim,int pe)
1866 {
1867
1868   char fmt[1000];
1869   char filename[1000];
1870   strncpy(fmt,infilename,999);
1871   strncat(fmt,"_%d.out",999);
1872   sprintf(filename,fmt, pe);
1873   FILE *loutfile = fopen(filename, "w");
1874 #ifdef PAIRCALC_TEST_DUMP
1875   fprintf(loutfile,"%d\n",ydim);
1876 #endif
1877   fprintf(loutfile,"%d %d %d\n",xdim,ydim, zdim);
1878   for(int i=0;i<xdim;i++)
1879     for(int j=0;j<ydim;j++)
1880       for(int k=0;k<zdim;k++)
1881         fprintf(loutfile,"%d %d %d %.8f\n",i,j,k,matrix[i*zdim*ydim+j*zdim +k]);
1882   fclose(loutfile);
1883
1884 }
1885
1886 void dumpMatrixFloat3(const char *infilename, float *matrix, int xdim, int ydim, int zdim,int x, int y, int z)
1887 {
1888   char fmt[1000];
1889   char filename[1000];
1890   strncpy(fmt,infilename,999);
1891   strncat(fmt,"_%d_%d_%d.out",999);
1892   sprintf(filename,fmt, x,y,z);
1893   FILE *loutfile = fopen(filename, "w");
1894   CkAssert(loutfile!=NULL);
1895   CkPrintf("opened %s for dump\n",filename);
1896   fprintf(loutfile,"%d %d %d\n",xdim,ydim, zdim);
1897   for(int i=0;i<xdim;i++)
1898     for(int j=0;j<ydim;j++)
1899       for(int k=0;k<zdim;k++)
1900         fprintf(loutfile,"%d %d %d %.8f\n",i,j,k,matrix[i*zdim*ydim+j*zdim +k]);
1901   fclose(loutfile);
1902 }
1903
1904 #endif
1905
1906 void ComputePmeMgr::gridCalc1(void) {
1907   // CkPrintf("gridCalc1 on Pe(%d)\n",CkMyPe());
1908
1909 #ifdef NAMD_FFTW
1910   for ( int g=0; g<numGrids; ++g ) {
1911 #ifdef NAMD_FFTW_3
1912     fftwf_execute(forward_plan_yz[g]);
1913 #else
1914     rfftwnd_real_to_complex(forward_plan_yz, localInfo[myGridPe].nx,
1915         qgrid + qgrid_size * g, 1, myGrid.dim2 * myGrid.dim3, 0, 0, 0);
1916 #endif
1917
1918   }
1919 #endif
1920
1921   if ( ! useBarrier ) pmeProxyDir[CkMyPe()].sendTrans();
1922 }
1923
1924 void ComputePmeMgr::sendTransBarrier(void) {
1925   sendTransBarrier_received += 1;
1926   // CkPrintf("sendTransBarrier on %d %d\n",myGridPe,numGridPes-sendTransBarrier_received);
1927   if ( sendTransBarrier_received < numGridPes ) return;
1928   sendTransBarrier_received = 0;
1929   for ( int i=0; i<numGridPes; ++i ) {
1930     pmeProxyDir[gridPeMap[i]].sendTrans();
1931   }
1932 }
1933
1934 static inline void PmeSlabSendTrans(int first, int last, void *result, int paraNum, void *param) {
1935   ComputePmeMgr *mgr = (ComputePmeMgr *)param;
1936   mgr->sendTransSubset(first, last);
1937 }
1938
1939 void ComputePmeMgr::sendTrans(void) {
1940
1941   untrans_count = numTransPes;
1942
1943 #if     CMK_SMP && USE_CKLOOP
1944   int useCkLoop = Node::Object()->simParameters->useCkLoop;
1945   if ( useCkLoop >= CKLOOP_CTRL_PME_SENDTRANS && CkNumPes() >= 2 * numGridPes) {
1946     CkLoop_Parallelize(PmeSlabSendTrans, 1, (void *)this, CkMyNodeSize(), 0, numTransNodes-1, 0); // no sync
1947   } else
1948 #endif
1949   {
1950     sendTransSubset(0, numTransNodes-1);
1951   }
1952
1953 }
1954
1955 void ComputePmeMgr::sendTransSubset(int first, int last) {
1956   // CkPrintf("sendTrans on Pe(%d)\n",CkMyPe());
1957
1958   // send data for transpose
1959   int zdim = myGrid.dim3;
1960   int nx = localInfo[myGridPe].nx;
1961   int x_start = localInfo[myGridPe].x_start;
1962   int slicelen = myGrid.K2 * zdim;
1963
1964   ComputePmeMgr **mgrObjects = pmeNodeProxy.ckLocalBranch()->mgrObjects;
1965
1966 #if CMK_BLUEGENEL
1967   CmiNetworkProgressAfter (0);
1968 #endif
1969
1970   for (int j=first; j<=last; j++) {
1971     int node = transNodeOrder[j];  // different order on each node
1972     int pe = transNodeInfo[node].pe_start;
1973     int npe = transNodeInfo[node].npe;
1974     int totlen = 0;
1975     if ( node != myTransNode ) for (int i=0; i<npe; ++i, ++pe) {
1976       LocalPmeInfo &li = localInfo[pe];
1977       int cpylen = li.ny_after_transpose * zdim;
1978       totlen += cpylen;
1979     }
1980     PmeTransMsg *newmsg = new (nx * totlen * numGrids,
1981                                 PRIORITY_SIZE) PmeTransMsg;
1982     newmsg->sourceNode = myGridPe;
1983     newmsg->lattice = lattice;
1984     newmsg->x_start = x_start;
1985     newmsg->nx = nx;
1986     for ( int g=0; g<numGrids; ++g ) {
1987       float *qmsg = newmsg->qgrid + nx * totlen * g;
1988       pe = transNodeInfo[node].pe_start;
1989       for (int i=0; i<npe; ++i, ++pe) {
1990         LocalPmeInfo &li = localInfo[pe];
1991         int cpylen = li.ny_after_transpose * zdim;
1992         if ( node == myTransNode ) {
1993           ComputePmeMgr *m = mgrObjects[CkRankOf(transPeMap[pe])];
1994           qmsg = m->kgrid + m->qgrid_size * g + x_start*cpylen;
1995         }
1996         float *q = qgrid + qgrid_size * g + li.y_start_after_transpose * zdim;
1997         for ( int x = 0; x < nx; ++x ) {
1998           CmiMemcpy((void*)qmsg, (void*)q, cpylen*sizeof(float));
1999           q += slicelen;
2000           qmsg += cpylen;
2001         }
2002       }
2003     }
2004     newmsg->sequence = grid_sequence;
2005     SET_PRIORITY(newmsg,grid_sequence,PME_TRANS_PRIORITY)
2006     if ( node == myTransNode ) newmsg->nx = 0;
2007     if ( npe > 1 ) {
2008       if ( node == myTransNode ) fwdSharedTrans(newmsg);
2009       else pmeNodeProxy[transNodeInfo[node].real_node].recvTrans(newmsg);
2010     } else pmeProxy[transPeMap[transNodeInfo[node].pe_start]].recvTrans(newmsg);
2011   }
2012 }
2013
2014 void ComputePmeMgr::fwdSharedTrans(PmeTransMsg *msg) {
2015   // CkPrintf("fwdSharedTrans on Pe(%d)\n",CkMyPe());
2016   int pe = transNodeInfo[myTransNode].pe_start;
2017   int npe = transNodeInfo[myTransNode].npe;
2018   CmiNodeLock lock = CmiCreateLock();
2019   int *count = new int; *count = npe;
2020   for (int i=0; i<npe; ++i, ++pe) {
2021     PmeSharedTransMsg *shmsg = new (PRIORITY_SIZE) PmeSharedTransMsg;
2022     SET_PRIORITY(shmsg,msg->sequence,PME_TRANS_PRIORITY)
2023     shmsg->msg = msg;
2024     shmsg->count = count;
2025     shmsg->lock = lock;
2026     pmeProxy[transPeMap[pe]].recvSharedTrans(shmsg);
2027   }
2028 }
2029
2030 void ComputePmeMgr::recvSharedTrans(PmeSharedTransMsg *msg) {
2031   procTrans(msg->msg);
2032   CmiLock(msg->lock);
2033   int count = --(*msg->count);
2034   CmiUnlock(msg->lock);
2035   if ( count == 0 ) {
2036     CmiDestroyLock(msg->lock);
2037     delete msg->count;
2038     delete msg->msg;
2039   }
2040   delete msg;
2041 }
2042
2043 void ComputePmeMgr::recvTrans(PmeTransMsg *msg) {
2044   procTrans(msg);
2045   delete msg;
2046 }
2047
2048 void ComputePmeMgr::procTrans(PmeTransMsg *msg) {
2049   // CkPrintf("procTrans on Pe(%d)\n",CkMyPe());
2050   if ( trans_count == numGridPes ) {
2051     lattice = msg->lattice;
2052     grid_sequence = msg->sequence;
2053   }
2054
2055  if ( msg->nx ) {
2056   int zdim = myGrid.dim3;
2057   NodePmeInfo &nodeInfo(transNodeInfo[myTransNode]);
2058   int first_pe = nodeInfo.pe_start;
2059   int last_pe = first_pe+nodeInfo.npe-1;
2060   int y_skip = localInfo[myTransPe].y_start_after_transpose
2061              - localInfo[first_pe].y_start_after_transpose;
2062   int ny_msg = localInfo[last_pe].y_start_after_transpose
2063              + localInfo[last_pe].ny_after_transpose
2064              - localInfo[first_pe].y_start_after_transpose;
2065   int ny = localInfo[myTransPe].ny_after_transpose;
2066   int x_start = msg->x_start;
2067   int nx = msg->nx;
2068   for ( int g=0; g<numGrids; ++g ) {
2069     CmiMemcpy((void*)(kgrid + qgrid_size * g + x_start*ny*zdim),
2070         (void*)(msg->qgrid + nx*(ny_msg*g+y_skip)*zdim),
2071         nx*ny*zdim*sizeof(float));
2072   }
2073  }
2074
2075   --trans_count;
2076
2077   if ( trans_count == 0 ) {
2078     pmeProxyDir[CkMyPe()].gridCalc2();
2079   }
2080 }
2081
2082 void ComputePmeMgr::gridCalc2(void) {
2083   // CkPrintf("gridCalc2 on Pe(%d)\n",CkMyPe());
2084
2085 #if CMK_BLUEGENEL
2086   CmiNetworkProgressAfter (0);
2087 #endif
2088
2089   int zdim = myGrid.dim3;
2090   // int y_start = localInfo[myTransPe].y_start_after_transpose;
2091   int ny = localInfo[myTransPe].ny_after_transpose;
2092
2093   for ( int g=0; g<numGrids; ++g ) {
2094     // finish forward FFT (x dimension)
2095 #ifdef NAMD_FFTW
2096 #ifdef NAMD_FFTW_3
2097     fftwf_execute(forward_plan_x[g]);
2098 #else
2099     fftw(forward_plan_x, ny * zdim / 2, (fftw_complex *)(kgrid+qgrid_size*g),
2100         ny * zdim / 2, 1, work, 1, 0);
2101 #endif
2102 #endif
2103   }
2104
2105 #ifdef OPENATOM_VERSION
2106     if ( ! simParams -> openatomOn ) { 
2107 #endif // OPENATOM_VERSION
2108       gridCalc2R();
2109 #ifdef OPENATOM_VERSION
2110     } else {
2111       gridCalc2Moa();
2112     }
2113 #endif // OPENATOM_VERSION
2114 }
2115
2116 #ifdef OPENATOM_VERSION
2117 void ComputePmeMgr::gridCalc2Moa(void) {
2118
2119   int zdim = myGrid.dim3;
2120   // int y_start = localInfo[myTransPe].y_start_after_transpose;
2121   int ny = localInfo[myTransPe].ny_after_transpose;
2122
2123   SimParameters *simParams = Node::Object()->simParameters;
2124
2125   CProxy_ComputeMoaMgr moaProxy(CkpvAccess(BOCclass_group).computeMoaMgr);
2126
2127   for ( int g=0; g<numGrids; ++g ) {
2128     #ifdef OPENATOM_VERSION_DEBUG 
2129     CkPrintf("Sending recQ on processor %d \n", CkMyPe());
2130     for ( int i=0; i<=(ny * zdim / 2); ++i) 
2131     {
2132       CkPrintf("PE, g,fftw_q,k*q*g, kgrid, qgrid_size value %d pre-send = %d, %d, %f %f, %d, \n", i, CkMyPe(), g, (kgrid+qgrid_size*g)[i], kgrid[i], qgrid_size);
2133     }
2134     #endif // OPENATOM_VERSION_DEBUG
2135 //     mqcpProxy[CkMyPe()].recvQ((ny * zdim / 2),((fftw_complex *)(kgrid+qgrid_size*g)));
2136     CkCallback resumePme(CkIndex_ComputePmeMgr::gridCalc2R(), thishandle);
2137     moaProxy[CkMyPe()].recvQ(g,numGrids,(ny * zdim / 2),(kgrid+qgrid_size*g), resumePme);
2138   }
2139 }
2140 #endif // OPENATOM_VERSION
2141
2142 void ComputePmeMgr::gridCalc2R(void) {
2143
2144   int useCkLoop = 0;
2145 #if CMK_SMP && USE_CKLOOP
2146   if ( Node::Object()->simParameters->useCkLoop >= CKLOOP_CTRL_PME_KSPACE
2147        && CkNumPes() >= 2 * numTransPes ) {
2148     useCkLoop = 1;
2149   }
2150 #endif
2151
2152   int zdim = myGrid.dim3;
2153   // int y_start = localInfo[myTransPe].y_start_after_transpose;
2154   int ny = localInfo[myTransPe].ny_after_transpose;
2155
2156   for ( int g=0; g<numGrids; ++g ) {
2157     // reciprocal space portion of PME
2158     BigReal ewaldcof = ComputeNonbondedUtil::ewaldcof;
2159     recip_evir2[g][0] = myKSpace->compute_energy(kgrid+qgrid_size*g,
2160                         lattice, ewaldcof, &(recip_evir2[g][1]), useCkLoop);
2161     // CkPrintf("Ewald reciprocal energy = %f\n", recip_evir2[g][0]);
2162
2163     // start backward FFT (x dimension)
2164
2165 #ifdef NAMD_FFTW
2166 #ifdef NAMD_FFTW_3
2167     fftwf_execute(backward_plan_x[g]);
2168 #else
2169     fftw(backward_plan_x, ny * zdim / 2, (fftw_complex *)(kgrid+qgrid_size*g),
2170         ny * zdim / 2, 1, work, 1, 0);
2171 #endif
2172 #endif
2173   }
2174   
2175   pmeProxyDir[CkMyPe()].sendUntrans();
2176 }
2177
2178 static inline void PmeSlabSendUntrans(int first, int last, void *result, int paraNum, void *param) {
2179   ComputePmeMgr *mgr = (ComputePmeMgr *)param;
2180   mgr->sendUntransSubset(first, last);
2181 }
2182
2183 void ComputePmeMgr::sendUntrans(void) {
2184
2185   trans_count = numGridPes;
2186
2187   { // send energy and virial
2188     PmeEvirMsg *newmsg = new (numGrids, PRIORITY_SIZE) PmeEvirMsg;
2189     for ( int g=0; g<numGrids; ++g ) {
2190       newmsg->evir[g] = recip_evir2[g];
2191     }
2192     SET_PRIORITY(newmsg,grid_sequence,PME_UNGRID_PRIORITY)
2193     CmiEnableUrgentSend(1);
2194     pmeProxy[recipEvirPe].recvRecipEvir(newmsg);
2195     CmiEnableUrgentSend(0);
2196   }
2197
2198 #if     CMK_SMP && USE_CKLOOP
2199   int useCkLoop = Node::Object()->simParameters->useCkLoop;
2200   if ( useCkLoop >= CKLOOP_CTRL_PME_SENDUNTRANS && CkNumPes() >= 2 * numTransPes) {
2201     CkLoop_Parallelize(PmeSlabSendUntrans, 1, (void *)this, CkMyNodeSize(), 0, numGridNodes-1, 0); // no sync
2202   } else
2203 #endif
2204   {
2205     sendUntransSubset(0, numGridNodes-1);
2206   }
2207
2208 }
2209
2210 void ComputePmeMgr::sendUntransSubset(int first, int last) {
2211
2212   int zdim = myGrid.dim3;
2213   int y_start = localInfo[myTransPe].y_start_after_transpose;
2214   int ny = localInfo[myTransPe].ny_after_transpose;
2215   int slicelen = myGrid.K2 * zdim;
2216
2217   ComputePmeMgr **mgrObjects = pmeNodeProxy.ckLocalBranch()->mgrObjects;
2218
2219 #if CMK_BLUEGENEL
2220   CmiNetworkProgressAfter (0);
2221 #endif
2222
2223   // send data for reverse transpose
2224   for (int j=first; j<=last; j++) {
2225     int node = gridNodeOrder[j];  // different order on each node
2226     int pe = gridNodeInfo[node].pe_start;
2227     int npe = gridNodeInfo[node].npe;
2228     int totlen = 0;
2229     if ( node != myGridNode ) for (int i=0; i<npe; ++i, ++pe) {
2230       LocalPmeInfo &li = localInfo[pe];
2231       int cpylen = li.nx * zdim;
2232       totlen += cpylen;
2233     }
2234     PmeUntransMsg *newmsg = new (ny * totlen * numGrids, PRIORITY_SIZE) PmeUntransMsg;
2235     newmsg->sourceNode = myTransPe;
2236     newmsg->y_start = y_start;
2237     newmsg->ny = ny;
2238     for ( int g=0; g<numGrids; ++g ) {
2239       float *qmsg = newmsg->qgrid + ny * totlen * g;
2240       pe = gridNodeInfo[node].pe_start;
2241       for (int i=0; i<npe; ++i, ++pe) {
2242         LocalPmeInfo &li = localInfo[pe];
2243         if ( node == myGridNode ) {
2244           ComputePmeMgr *m = mgrObjects[CkRankOf(gridPeMap[pe])];
2245           qmsg = m->qgrid + m->qgrid_size * g + y_start * zdim;
2246           float *q = kgrid + qgrid_size*g + li.x_start*ny*zdim;
2247           int cpylen = ny * zdim;
2248           for ( int x = 0; x < li.nx; ++x ) {
2249             CmiMemcpy((void*)qmsg, (void*)q, cpylen*sizeof(float));
2250             q += cpylen;
2251             qmsg += slicelen;
2252           }
2253         } else {
2254           CmiMemcpy((void*)qmsg,
2255                 (void*)(kgrid + qgrid_size*g + li.x_start*ny*zdim),
2256                 li.nx*ny*zdim*sizeof(float));
2257           qmsg += li.nx*ny*zdim;
2258         }
2259       }
2260     }
2261     SET_PRIORITY(newmsg,grid_sequence,PME_UNTRANS_PRIORITY)
2262     if ( node == myGridNode ) newmsg->ny = 0;
2263     if ( npe > 1 ) {
2264       if ( node == myGridNode ) fwdSharedUntrans(newmsg);
2265       else pmeNodeProxy[gridNodeInfo[node].real_node].recvUntrans(newmsg);
2266     } else pmeProxy[gridPeMap[gridNodeInfo[node].pe_start]].recvUntrans(newmsg);
2267   }
2268 }
2269
2270 void ComputePmeMgr::fwdSharedUntrans(PmeUntransMsg *msg) {
2271   int pe = gridNodeInfo[myGridNode].pe_start;
2272   int npe = gridNodeInfo[myGridNode].npe;
2273   CmiNodeLock lock = CmiCreateLock();
2274   int *count = new int; *count = npe;
2275   for (int i=0; i<npe; ++i, ++pe) {
2276     PmeSharedUntransMsg *shmsg = new PmeSharedUntransMsg;
2277     shmsg->msg = msg;
2278     shmsg->count = count;
2279     shmsg->lock = lock;
2280     pmeProxy[gridPeMap[pe]].recvSharedUntrans(shmsg);
2281   }
2282 }
2283
2284 void ComputePmeMgr::recvSharedUntrans(PmeSharedUntransMsg *msg) {
2285   procUntrans(msg->msg);
2286   CmiLock(msg->lock);
2287   int count = --(*msg->count);
2288   CmiUnlock(msg->lock);
2289   if ( count == 0 ) {
2290     CmiDestroyLock(msg->lock);
2291     delete msg->count;
2292     delete msg->msg;
2293   }
2294   delete msg;
2295 }
2296
2297 void ComputePmeMgr::recvUntrans(PmeUntransMsg *msg) {
2298   procUntrans(msg);
2299   delete msg;
2300 }
2301
2302 void ComputePmeMgr::procUntrans(PmeUntransMsg *msg) {
2303   // CkPrintf("recvUntrans on Pe(%d)\n",CkMyPe());
2304
2305 #if CMK_BLUEGENEL
2306   CmiNetworkProgressAfter (0);
2307 #endif
2308
2309   NodePmeInfo &nodeInfo(gridNodeInfo[myGridNode]);
2310   int first_pe = nodeInfo.pe_start;
2311   int g;
2312
2313  if ( msg->ny ) {
2314   int zdim = myGrid.dim3;
2315   int last_pe = first_pe+nodeInfo.npe-1;
2316   int x_skip = localInfo[myGridPe].x_start
2317              - localInfo[first_pe].x_start;
2318   int nx_msg = localInfo[last_pe].x_start
2319              + localInfo[last_pe].nx
2320              - localInfo[first_pe].x_start;
2321   int nx = localInfo[myGridPe].nx;
2322   int y_start = msg->y_start;
2323   int ny = msg->ny;
2324   int slicelen = myGrid.K2 * zdim;
2325   int cpylen = ny * zdim;
2326   for ( g=0; g<numGrids; ++g ) {
2327     float *q = qgrid + qgrid_size * g + y_start * zdim;
2328     float *qmsg = msg->qgrid + (nx_msg*g+x_skip) * cpylen;
2329     for ( int x = 0; x < nx; ++x ) {
2330       CmiMemcpy((void*)q, (void*)qmsg, cpylen*sizeof(float));
2331       q += slicelen;
2332       qmsg += cpylen;
2333     }
2334   }
2335  }
2336
2337   --untrans_count;
2338
2339   if ( untrans_count == 0 ) {
2340     pmeProxyDir[CkMyPe()].gridCalc3();
2341   }
2342 }
2343
2344 void ComputePmeMgr::gridCalc3(void) {
2345   // CkPrintf("gridCalc3 on Pe(%d)\n",CkMyPe());
2346
2347   // finish backward FFT
2348 #ifdef NAMD_FFTW
2349   for ( int g=0; g<numGrids; ++g ) {
2350 #ifdef NAMD_FFTW_3
2351     fftwf_execute(backward_plan_yz[g]);
2352 #else
2353     rfftwnd_complex_to_real(backward_plan_yz, localInfo[myGridPe].nx,
2354         (fftw_complex *) (qgrid + qgrid_size * g),
2355         1, myGrid.dim2 * myGrid.dim3 / 2, 0, 0, 0);
2356 #endif
2357   }
2358
2359 #endif
2360
2361   pmeProxyDir[CkMyPe()].sendUngrid();
2362 }
2363
2364 static inline void PmeSlabSendUngrid(int first, int last, void *result, int paraNum, void *param) {
2365   ComputePmeMgr *mgr = (ComputePmeMgr *)param;
2366   mgr->sendUngridSubset(first, last);
2367 }
2368
2369 void ComputePmeMgr::sendUngrid(void) {
2370
2371 #if     CMK_SMP && USE_CKLOOP
2372   int useCkLoop = Node::Object()->simParameters->useCkLoop;
2373   if ( useCkLoop >= CKLOOP_CTRL_PME_SENDUNTRANS && CkNumPes() >= 2 * numGridPes) {
2374     CkLoop_Parallelize(PmeSlabSendUngrid, 1, (void *)this, CkMyNodeSize(), 0, numSources-1, 1); // sync
2375   } else
2376 #endif
2377   {
2378     sendUngridSubset(0, numSources-1);
2379   }
2380
2381   grid_count = numSources;
2382   memset( (void*) qgrid, 0, qgrid_size * numGrids * sizeof(float) );
2383 }
2384
2385 void ComputePmeMgr::sendUngridSubset(int first, int last) {
2386
2387 #ifdef NAMD_CUDA
2388   const int UNGRID_PRIORITY = ( offload ? PME_OFFLOAD_UNGRID_PRIORITY : PME_UNGRID_PRIORITY );
2389 #else
2390   const int UNGRID_PRIORITY = PME_UNGRID_PRIORITY ;
2391 #endif
2392
2393   for ( int j=first; j<=last; ++j ) {
2394     // int msglen = qgrid_len;
2395     PmeGridMsg *newmsg = gridmsg_reuse[j];
2396     int pe = newmsg->sourceNode;
2397     int zdim = myGrid.dim3;
2398     int flen = newmsg->len;
2399     int fstart = newmsg->start;
2400     int zlistlen = newmsg->zlistlen;
2401     int *zlist = newmsg->zlist;
2402     float *qmsg = newmsg->qgrid;
2403     for ( int g=0; g<numGrids; ++g ) {
2404       char *f = newmsg->fgrid + fgrid_len * g;
2405       float *q = qgrid + qgrid_size * g + (fstart-fgrid_start) * zdim;
2406       for ( int i=0; i<flen; ++i ) {
2407         if ( f[i] ) {
2408           for ( int k=0; k<zlistlen; ++k ) {
2409             *(qmsg++) = q[zlist[k]];
2410           }
2411         }
2412         q += zdim;
2413       }
2414     }
2415     newmsg->sourceNode = myGridPe;
2416
2417     SET_PRIORITY(newmsg,grid_sequence,UNGRID_PRIORITY)
2418     CmiEnableUrgentSend(1);
2419 #ifdef NAMD_CUDA
2420     if ( offload ) {
2421       pmeNodeProxy[CkNodeOf(pe)].recvUngrid(newmsg);
2422     } else
2423 #endif
2424     pmeProxyDir[pe].recvUngrid(newmsg);
2425     CmiEnableUrgentSend(0);
2426   }
2427 }
2428
2429 void ComputePmeMgr::recvUngrid(PmeGridMsg *msg) {
2430   // CkPrintf("recvUngrid on Pe(%d)\n",CkMyPe());
2431 #ifdef NAMD_CUDA
2432   if ( ! offload )  // would need lock
2433 #endif
2434   if ( ungrid_count == 0 ) {
2435     NAMD_bug("Message order failure in ComputePmeMgr::recvUngrid\n");
2436   }
2437
2438   if ( usePencils ) copyPencils(msg);
2439   else copyResults(msg);
2440   delete msg;
2441   recvAck(0);
2442 }
2443
2444 void ComputePmeMgr::recvAck(PmeAckMsg *msg) {
2445   if ( msg ) delete msg;
2446 #ifdef NAMD_CUDA
2447   if ( offload ) {
2448     CmiLock(cuda_lock);
2449     if ( ungrid_count == 0 ) {
2450       NAMD_bug("Message order failure in ComputePmeMgr::recvUngrid\n");
2451     }
2452     int uc = --ungrid_count;
2453     CmiUnlock(cuda_lock);
2454
2455     if ( uc == 0 ) {
2456       pmeProxyDir[master_pe].ungridCalc();
2457     }
2458     return;
2459   }
2460 #endif
2461   --ungrid_count;
2462
2463   if ( ungrid_count == 0 ) {
2464     pmeProxyDir[CkMyPe()].ungridCalc();
2465   }
2466 }
2467
2468 #ifdef NAMD_CUDA
2469 #define count_limit 1000000
2470 #define CUDA_POLL(FN,ARG) CcdCallFnAfter(FN,ARG,0.1)
2471 #define EVENT_STRIDE 10
2472
2473 extern "C" void CcdCallBacksReset(void *ignored,double curWallTime);  // fix Charm++
2474
2475 void cuda_check_pme_forces(void *arg, double walltime) {
2476   ComputePmeMgr *argp = (ComputePmeMgr *) arg;
2477
2478  while ( 1 ) { // process multiple events per call
2479   cudaError_t err = cudaEventQuery(argp->end_forces[argp->forces_done_count/EVENT_STRIDE]);
2480   if ( err == cudaSuccess ) {
2481     argp->check_forces_count = 0;
2482     for ( int i=0; i<EVENT_STRIDE; ++i ) {
2483       WorkDistrib::messageEnqueueWork(argp->pmeComputes[argp->forces_done_count]);
2484       if ( ++(argp->forces_done_count) == argp->forces_count ) break;
2485     }
2486     if ( argp->forces_done_count == argp->forces_count ) { // last event
2487       traceUserBracketEvent(CUDA_EVENT_ID_PME_FORCES,argp->forces_time,walltime);
2488       argp->forces_time = walltime - argp->forces_time;
2489       //CkPrintf("cuda_check_pme_forces forces_time == %f\n", argp->forces_time);
2490       return;
2491     } else { // more events
2492       continue; // check next event
2493     }
2494   } else if ( err != cudaErrorNotReady ) {
2495     cuda_errcheck("in cuda_check_pme_forces");
2496     NAMD_bug("cuda_errcheck missed error in cuda_check_pme_forces");
2497   } else if ( ++(argp->check_forces_count) >= count_limit ) {
2498     char errmsg[256];
2499     sprintf(errmsg,"cuda_check_pme_forces polled %d times over %f s on seq %d",
2500             argp->check_forces_count, walltime - argp->forces_time,
2501             argp->saved_sequence);
2502     cuda_errcheck(errmsg);
2503     NAMD_die(errmsg);
2504   } else {
2505     break; // call again
2506   }
2507  } // while ( 1 )
2508  CcdCallBacksReset(0,walltime);  // fix Charm++
2509  CUDA_POLL(cuda_check_pme_forces, arg);
2510 }
2511 #endif // NAMD_CUDA
2512
2513 void ComputePmeMgr::ungridCalc(void) {
2514   // CkPrintf("ungridCalc on Pe(%d)\n",CkMyPe());
2515
2516   ungridForcesCount = pmeComputes.size();
2517
2518 #ifdef NAMD_CUDA
2519  if ( offload ) {
2520   //CmiLock(cuda_lock);
2521   cudaSetDevice(deviceCUDA->getDeviceID());
2522
2523   if ( this == masterPmeMgr ) {
2524     double before = CmiWallTimer();
2525     cudaMemcpyAsync(v_data_dev, q_data_host, q_data_size, cudaMemcpyHostToDevice, 0 /*streams[stream]*/);
2526     cudaEventRecord(nodePmeMgr->end_potential_memcpy, 0 /*streams[stream]*/);
2527     traceUserBracketEvent(CUDA_EVENT_ID_PME_COPY,before,CmiWallTimer());
2528
2529     const int myrank = CkMyRank();
2530     for ( int i=0; i<CkMyNodeSize(); ++i ) {
2531       if ( myrank != i && nodePmeMgr->mgrObjects[i]->pmeComputes.size() ) {
2532         nodePmeMgr->mgrObjects[i]->ungridCalc();
2533       }
2534     }
2535     if ( ! pmeComputes.size() ) return;
2536   }
2537
2538   if ( ! end_forces ) {
2539     int n=(pmeComputes.size()-1)/EVENT_STRIDE+1;
2540     end_forces = new cudaEvent_t[n];
2541     for ( int i=0; i<n; ++i ) {
2542       cudaEventCreateWithFlags(&end_forces[i],cudaEventDisableTiming);
2543     }
2544   }
2545
2546   const int pcsz = pmeComputes.size();
2547   if ( ! afn_host ) {
2548     cudaMallocHost((void**) &afn_host, 3*pcsz*sizeof(float*));
2549     cudaMalloc((void**) &afn_dev, 3*pcsz*sizeof(float*));
2550     cuda_errcheck("malloc params for pme");
2551   }
2552   int totn = 0;
2553   for ( int i=0; i<pcsz; ++i ) {
2554     int n = pmeComputes[i]->numGridAtoms[0];
2555     totn += n;
2556   }
2557   if ( totn > f_data_mgr_alloc ) {
2558     if ( f_data_mgr_alloc ) {
2559       CkPrintf("Expanding CUDA forces allocation because %d > %d\n", totn, f_data_mgr_alloc);
2560       cudaFree(f_data_mgr_dev);
2561       cudaFreeHost(f_data_mgr_host);
2562     }
2563     f_data_mgr_alloc = 1.2 * (totn + 100);
2564     cudaMalloc((void**) &f_data_mgr_dev, 3*f_data_mgr_alloc*sizeof(float));
2565     cudaMallocHost((void**) &f_data_mgr_host, 3*f_data_mgr_alloc*sizeof(float));
2566     cuda_errcheck("malloc forces for pme");
2567   }
2568   // CkPrintf("pe %d pcsz %d totn %d alloc %d\n", CkMyPe(), pcsz, totn, f_data_mgr_alloc);
2569   float *f_dev = f_data_mgr_dev;
2570   float *f_host = f_data_mgr_host;
2571   for ( int i=0; i<pcsz; ++i ) {
2572     int n = pmeComputes[i]->numGridAtoms[0];
2573     pmeComputes[i]->f_data_dev = f_dev;
2574     pmeComputes[i]->f_data_host = f_host;
2575     afn_host[3*i  ] = a_data_dev + 7 * pmeComputes[i]->cuda_atoms_offset;
2576     afn_host[3*i+1] = f_dev;
2577     afn_host[3*i+2] = f_dev + n;  // avoid type conversion issues
2578     f_dev += 3*n;
2579     f_host += 3*n;
2580   }
2581   //CmiLock(cuda_lock);
2582   double before = CmiWallTimer();
2583   cudaMemcpyAsync(afn_dev, afn_host, 3*pcsz*sizeof(float*), cudaMemcpyHostToDevice, streams[stream]);
2584   traceUserBracketEvent(CUDA_EVENT_ID_PME_COPY,before,CmiWallTimer());
2585   cudaStreamWaitEvent(streams[stream], nodePmeMgr->end_potential_memcpy, 0);
2586   traceUserEvent(CUDA_EVENT_ID_PME_TICK);
2587
2588   for ( int i=0; i<pcsz; ++i ) {
2589     // cudaMemsetAsync(pmeComputes[i]->f_data_dev, 0, 3*n*sizeof(float), streams[stream]);
2590     if ( i%EVENT_STRIDE == 0 ) {
2591       int dimy = pcsz - i;
2592       if ( dimy > EVENT_STRIDE ) dimy = EVENT_STRIDE;
2593       int maxn = 0;
2594       int subtotn = 0;
2595       for ( int j=0; j<dimy; ++j ) {
2596         int n = pmeComputes[i+j]->numGridAtoms[0];
2597         subtotn += n;
2598         if ( n > maxn ) maxn = n;
2599       }
2600       // CkPrintf("pe %d dimy %d maxn %d subtotn %d\n", CkMyPe(), dimy, maxn, subtotn);
2601       before = CmiWallTimer();
2602       cuda_pme_forces(
2603         bspline_coeffs_dev,
2604         v_arr_dev, afn_dev+3*i, dimy, maxn, /*
2605         pmeComputes[i]->a_data_dev,
2606         pmeComputes[i]->f_data_dev,
2607         n, */ myGrid.K1, myGrid.K2, myGrid.K3, myGrid.order,
2608         streams[stream]);
2609       traceUserBracketEvent(CUDA_EVENT_ID_PME_KERNEL,before,CmiWallTimer());
2610       before = CmiWallTimer();
2611       cudaMemcpyAsync(pmeComputes[i]->f_data_host, pmeComputes[i]->f_data_dev, 3*subtotn*sizeof(float),
2612         cudaMemcpyDeviceToHost, streams[stream]);
2613       traceUserBracketEvent(CUDA_EVENT_ID_PME_COPY,before,CmiWallTimer());
2614       cudaEventRecord(end_forces[i/EVENT_STRIDE], streams[stream]);
2615       traceUserEvent(CUDA_EVENT_ID_PME_TICK);
2616     }
2617     // CkPrintf("pe %d c %d natoms %d fdev %lld fhost %lld\n", CkMyPe(), i, (int64)afn_host[3*i+2], pmeComputes[i]->f_data_dev, pmeComputes[i]->f_data_host);
2618   }
2619   //CmiUnlock(cuda_lock);
2620  } else
2621 #endif // NAMD_CUDA
2622  {
2623   for ( int i=0; i<pmeComputes.size(); ++i ) {
2624     WorkDistrib::messageEnqueueWork(pmeComputes[i]);
2625     // pmeComputes[i]->ungridForces();
2626   }
2627  }
2628   // submitReductions();  // must follow all ungridForces()
2629
2630 #ifdef NAMD_CUDA
2631  if ( offload ) {
2632   forces_time = CmiWallTimer();
2633   forces_count = ungridForcesCount;
2634   forces_done_count = 0;
2635   pmeProxy[this_pe].pollForcesReady();
2636  }
2637 #endif
2638
2639   ungrid_count = (usePencils ? numPencilsActive : numDestRecipPes );
2640 }
2641
2642 void ComputePmeMgr::pollForcesReady() {
2643 #ifdef NAMD_CUDA
2644   CcdCallBacksReset(0,CmiWallTimer());  // fix Charm++
2645   CUDA_POLL(cuda_check_pme_forces,this);
2646 #else
2647   NAMD_bug("ComputePmeMgr::pollForcesReady() called in non-CUDA build.");
2648 #endif
2649 }
2650
2651 void ComputePme::atomUpdate() { atomsChanged = 1; }
2652
2653 ComputePme::ComputePme(ComputeID c, PatchID pid) : Compute(c), patchID(pid)
2654 {
2655   DebugM(4,"ComputePme created.\n");
2656   basePriority = PME_PRIORITY;
2657   setNumPatches(1);
2658
2659   CProxy_ComputePmeMgr::ckLocalBranch(
2660         CkpvAccess(BOCclass_group).computePmeMgr)->addCompute(this);
2661
2662   SimParameters *simParams = Node::Object()->simParameters;
2663
2664   qmForcesOn =  simParams->qmForcesOn;
2665   offload = simParams->PMEOffload;
2666
2667   numGridsMax = numGrids;
2668
2669   myGrid.K1 = simParams->PMEGridSizeX;
2670   myGrid.K2 = simParams->PMEGridSizeY;
2671   myGrid.K3 = simParams->PMEGridSizeZ;
2672   myGrid.order = simParams->PMEInterpOrder;
2673   myGrid.dim2 = myGrid.K2;
2674   myGrid.dim3 = 2 * (myGrid.K3/2 + 1);
2675
2676 #ifdef NAMD_CUDA
2677   cuda_atoms_offset = 0;
2678   f_data_host = 0;
2679   f_data_dev = 0;
2680  if ( ! offload )
2681 #endif
2682  {
2683   for ( int g=0; g<numGrids; ++g ) myRealSpace[g] = new PmeRealSpace(myGrid);
2684  }
2685
2686   atomsChanged = 0;
2687   
2688   qmLoclIndx = 0;
2689   qmLocalCharges = 0;
2690 }
2691
2692 void ComputePme::initialize() {
2693   if (!(patch = PatchMap::Object()->patch(patchID))) {
2694     NAMD_bug("ComputePme used with unknown patch.");
2695   }
2696   positionBox = patch->registerPositionPickup(this);
2697   avgPositionBox = patch->registerAvgPositionPickup(this);
2698   forceBox = patch->registerForceDeposit(this);
2699 #ifdef NAMD_CUDA
2700  if ( offload ) {
2701   myMgr->cuda_atoms_count += patch->getNumAtoms();
2702  }
2703 #endif
2704 }
2705
2706 void ComputePmeMgr::initialize_computes() {
2707
2708   noWorkCount = 0;
2709   doWorkCount = 0;
2710   ungridForcesCount = 0;
2711
2712   reduction = ReductionMgr::Object()->willSubmit(REDUCTIONS_BASIC);
2713
2714   SimParameters *simParams = Node::Object()->simParameters;
2715
2716   strayChargeErrors = 0;
2717
2718 #ifdef NAMD_CUDA
2719  PatchMap *patchMap = PatchMap::Object();
2720  int pe = master_pe = CkNodeFirst(CkMyNode());
2721  for ( int i=0; i<CkMyNodeSize(); ++i, ++pe ) {
2722     if ( ! patchMap->numPatchesOnNode(master_pe) ) master_pe = pe;
2723     if ( ! patchMap->numPatchesOnNode(pe) ) continue;
2724     if ( master_pe < 1 && pe != deviceCUDA->getMasterPe() ) master_pe = pe;
2725     if ( master_pe == deviceCUDA->getMasterPe() ) master_pe = pe;
2726     if ( WorkDistrib::pe_sortop_diffuse()(pe,master_pe)
2727         && pe != deviceCUDA->getMasterPe() ) {
2728       master_pe = pe;
2729     }
2730  }
2731  if ( ! patchMap->numPatchesOnNode(master_pe) ) {
2732    NAMD_bug("ComputePmeMgr::initialize_computes() master_pe has no patches.");
2733  }
2734
2735  masterPmeMgr = nodePmeMgr->mgrObjects[master_pe - CkNodeFirst(CkMyNode())];
2736  bool cudaFirst = 1;
2737  if ( offload ) {
2738   CmiLock(cuda_lock);
2739   cudaFirst = ! masterPmeMgr->chargeGridSubmittedCount++;
2740  }
2741
2742  if ( cudaFirst ) {
2743   nodePmeMgr->master_pe = master_pe;
2744   nodePmeMgr->masterPmeMgr = masterPmeMgr;
2745  }
2746 #endif
2747
2748   qsize = myGrid.K1 * myGrid.dim2 * myGrid.dim3;
2749   fsize = myGrid.K1 * myGrid.dim2;
2750   if ( myGrid.K2 != myGrid.dim2 ) NAMD_bug("PME myGrid.K2 != myGrid.dim2");
2751 #ifdef NAMD_CUDA
2752  if ( ! offload )
2753 #endif
2754  {
2755   q_arr = new float*[fsize*numGrids];
2756   memset( (void*) q_arr, 0, fsize*numGrids * sizeof(float*) );
2757   q_list = new float*[fsize*numGrids];
2758   memset( (void*) q_list, 0, fsize*numGrids * sizeof(float*) );
2759   q_count = 0;
2760  }
2761
2762 #ifdef NAMD_CUDA
2763  if ( cudaFirst || ! offload ) {
2764 #endif
2765   f_arr = new char[fsize*numGrids];
2766   // memset to non-zero value has race condition on BlueGene/Q
2767   // memset( (void*) f_arr, 2, fsize*numGrids * sizeof(char) );
2768   for ( int n=fsize*numGrids, i=0; i<n; ++i ) f_arr[i] = 2;
2769
2770   for ( int g=0; g<numGrids; ++g ) {
2771     char *f = f_arr + g*fsize;
2772     if ( usePencils ) {
2773       int K1 = myGrid.K1;
2774       int K2 = myGrid.K2;
2775       int block1 = ( K1 + xBlocks - 1 ) / xBlocks;
2776       int block2 = ( K2 + yBlocks - 1 ) / yBlocks;
2777       int dim2 = myGrid.dim2;
2778       for (int ap=0; ap<numPencilsActive; ++ap) {
2779         int ib = activePencils[ap].i;
2780         int jb = activePencils[ap].j;
2781         int ibegin = ib*block1;
2782         int iend = ibegin + block1;  if ( iend > K1 ) iend = K1;
2783         int jbegin = jb*block2;
2784         int jend = jbegin + block2;  if ( jend > K2 ) jend = K2;
2785         int flen = numGrids * (iend - ibegin) * (jend - jbegin);
2786         for ( int i=ibegin; i<iend; ++i ) {
2787           for ( int j=jbegin; j<jend; ++j ) {
2788             f[i*dim2+j] = 0;
2789           }
2790         }
2791       }
2792     } else {
2793       int block1 = ( myGrid.K1 + numGridPes - 1 ) / numGridPes;
2794       bsize = block1 * myGrid.dim2 * myGrid.dim3;
2795       for (int pe=0; pe<numGridPes; pe++) {
2796         if ( ! recipPeDest[pe] ) continue;
2797         int start = pe * bsize;
2798         int len = bsize;
2799         if ( start >= qsize ) { start = 0; len = 0; }
2800         if ( start + len > qsize ) { len = qsize - start; }
2801         int zdim = myGrid.dim3;
2802         int fstart = start / zdim;
2803         int flen = len / zdim;
2804         memset(f + fstart, 0, flen*sizeof(char));
2805         // CkPrintf("pe %d enabled slabs %d to %d\n", CkMyPe(), fstart/myGrid.dim2, (fstart+flen)/myGrid.dim2-1);
2806       }
2807     }
2808   }
2809 #ifdef NAMD_CUDA
2810  }
2811  if ( offload ) {
2812  cudaSetDevice(deviceCUDA->getDeviceID());
2813  if ( cudaFirst ) {
2814
2815   int f_alloc_count = 0;
2816   for ( int n=fsize, i=0; i<n; ++i ) {
2817     if ( f_arr[i] == 0 ) {
2818       ++f_alloc_count;
2819     }
2820   }
2821   // CkPrintf("pe %d f_alloc_count == %d (%d slabs)\n", CkMyPe(), f_alloc_count, f_alloc_count/myGrid.dim2);
2822
2823   q_arr = new float*[fsize*numGrids];
2824   memset( (void*) q_arr, 0, fsize*numGrids * sizeof(float*) );
2825
2826   float **q_arr_dev_host = new float*[fsize];
2827   cudaMalloc((void**) &q_arr_dev, fsize * sizeof(float*));
2828
2829   float **v_arr_dev_host = new float*[fsize];
2830   cudaMalloc((void**) &v_arr_dev, fsize * sizeof(float*));
2831
2832   int q_stride = myGrid.K3+myGrid.order-1;
2833   q_data_size = f_alloc_count * q_stride * sizeof(float);
2834   ffz_size = (fsize + q_stride) * sizeof(int);
2835
2836   // tack ffz onto end of q_data to allow merged transfer
2837   cudaMallocHost((void**) &q_data_host, q_data_size+ffz_size);
2838   ffz_host = (int*)(((char*)q_data_host) + q_data_size);
2839   cudaMalloc((void**) &q_data_dev, q_data_size+ffz_size);
2840   ffz_dev = (int*)(((char*)q_data_dev) + q_data_size);
2841   cudaMalloc((void**) &v_data_dev, q_data_size);
2842   cuda_errcheck("malloc grid data for pme");
2843   cudaMemset(q_data_dev, 0, q_data_size + ffz_size);  // for first time
2844   cudaEventCreateWithFlags(&(nodePmeMgr->end_charge_memset),cudaEventDisableTiming);
2845   cudaEventRecord(nodePmeMgr->end_charge_memset, 0);
2846   cudaEventCreateWithFlags(&(nodePmeMgr->end_all_pme_kernels),cudaEventDisableTiming);
2847   cudaEventCreateWithFlags(&(nodePmeMgr->end_potential_memcpy),cudaEventDisableTiming);
2848
2849   f_alloc_count = 0;
2850   for ( int n=fsize, i=0; i<n; ++i ) {
2851     if ( f_arr[i] == 0 ) {
2852       q_arr[i] = q_data_host + f_alloc_count * q_stride;
2853       q_arr_dev_host[i] = q_data_dev + f_alloc_count * q_stride;
2854       v_arr_dev_host[i] = v_data_dev + f_alloc_count * q_stride;
2855       ++f_alloc_count;
2856     } else {
2857       q_arr[i] = 0;
2858       q_arr_dev_host[i] = 0;
2859       v_arr_dev_host[i] = 0;
2860     }
2861   }
2862
2863   cudaMemcpy(q_arr_dev, q_arr_dev_host, fsize * sizeof(float*), cudaMemcpyHostToDevice);
2864   cudaMemcpy(v_arr_dev, v_arr_dev_host, fsize * sizeof(float*), cudaMemcpyHostToDevice);
2865   delete [] q_arr_dev_host;
2866   delete [] v_arr_dev_host;
2867   delete [] f_arr;
2868   f_arr = new char[fsize + q_stride];
2869   fz_arr = f_arr + fsize;
2870   memset(f_arr, 0, fsize + q_stride);
2871   memset(ffz_host, 0, (fsize + q_stride)*sizeof(int));
2872
2873   cuda_errcheck("initialize grid data for pme");
2874
2875   cuda_init_bspline_coeffs(&bspline_coeffs_dev, &bspline_dcoeffs_dev, myGrid.order);
2876   cuda_errcheck("initialize bspline coefficients for pme");
2877
2878 #define XCOPY(X) masterPmeMgr->X = X;
2879   XCOPY(bspline_coeffs_dev)
2880   XCOPY(bspline_dcoeffs_dev)
2881   XCOPY(q_arr)
2882   XCOPY(q_arr_dev)
2883   XCOPY(v_arr_dev)
2884   XCOPY(q_data_size)
2885   XCOPY(q_data_host)
2886   XCOPY(q_data_dev)
2887   XCOPY(v_data_dev)
2888   XCOPY(ffz_size)
2889   XCOPY(ffz_host)
2890   XCOPY(ffz_dev)
2891   XCOPY(f_arr)
2892   XCOPY(fz_arr)
2893 #undef XCOPY
2894   //CkPrintf("pe %d init first\n", CkMyPe());
2895  } else { // cudaFirst
2896   //CkPrintf("pe %d init later\n", CkMyPe());
2897 #define XCOPY(X) X = masterPmeMgr->X;
2898   XCOPY(bspline_coeffs_dev)
2899   XCOPY(bspline_dcoeffs_dev)
2900   XCOPY(q_arr)
2901   XCOPY(q_arr_dev)
2902   XCOPY(v_arr_dev)
2903   XCOPY(q_data_size)
2904   XCOPY(q_data_host)
2905   XCOPY(q_data_dev)
2906   XCOPY(v_data_dev)
2907   XCOPY(ffz_size)
2908   XCOPY(ffz_host)
2909   XCOPY(ffz_dev)
2910   XCOPY(f_arr)
2911   XCOPY(fz_arr)
2912 #undef XCOPY
2913  } // cudaFirst
2914   CmiUnlock(cuda_lock);
2915  } else // offload
2916 #endif // NAMD_CUDA
2917  {
2918   fz_arr = new char[myGrid.K3+myGrid.order-1];
2919  }
2920
2921 #if 0 && USE_PERSISTENT
2922   recvGrid_handle = NULL;
2923 #endif
2924 }
2925
2926 ComputePme::~ComputePme()
2927 {
2928 #ifdef NAMD_CUDA
2929   if ( ! offload )
2930 #endif
2931   {
2932     for ( int g=0; g<numGridsMax; ++g ) delete myRealSpace[g];
2933   }
2934 }
2935
2936 #if 0 && USE_PERSISTENT 
2937 void ComputePmeMgr::setup_recvgrid_persistent() 
2938 {
2939     int K1 = myGrid.K1;
2940     int K2 = myGrid.K2;
2941     int dim2 = myGrid.dim2;
2942     int dim3 = myGrid.dim3;
2943     int block1 = myGrid.block1;
2944     int block2 = myGrid.block2;
2945
2946     CkArray *zPencil_local = zPencil.ckLocalBranch();
2947     recvGrid_handle = (PersistentHandle*) malloc( sizeof(PersistentHandle) * numPencilsActive);
2948     for (int ap=0; ap<numPencilsActive; ++ap) {
2949         int ib = activePencils[ap].i;
2950         int jb = activePencils[ap].j;
2951         int ibegin = ib*block1;
2952         int iend = ibegin + block1;  if ( iend > K1 ) iend = K1;
2953         int jbegin = jb*block2;
2954         int jend = jbegin + block2;  if ( jend > K2 ) jend = K2;
2955         int flen = numGrids * (iend - ibegin) * (jend - jbegin);
2956         // f is changing
2957         int fcount = 0;
2958         for ( int g=0; g<numGrids; ++g ) {
2959             char *f = f_arr + g*fsize;
2960             for ( int i=ibegin; i<iend; ++i ) {
2961                 for ( int j=jbegin; j<jend; ++j ) {
2962                     fcount += f[i*dim2+j];
2963                 }
2964             }
2965         }
2966         int zlistlen = 0;
2967         for ( int i=0; i<myGrid.K3; ++i ) {
2968             if ( fz_arr[i] ) ++zlistlen;
2969         }
2970         int hd = ( fcount? 1 : 0 );  // has data?
2971         int peer = zPencil_local->homePe(CkArrayIndex3D(ib, jb, 0));
2972         int compress_start = sizeof(PmeGridMsg ) + sizeof(envelope) + sizeof(int)*hd*zlistlen + sizeof(char)*hd*flen +sizeof(PmeReduction)*hd*numGrids ;
2973         int compress_size = sizeof(float)*hd*fcount*zlistlen;
2974         int size = compress_start +  compress_size  + PRIORITY_SIZE/8+6;
2975         recvGrid_handle[ap] =  CmiCreateCompressPersistentSize(peer, size, compress_start, compress_size, CMI_FLOATING);
2976     }
2977 }
2978 #endif
2979
2980 int ComputePme::noWork() {
2981
2982   if ( patch->flags.doFullElectrostatics ) {
2983     // In QM/MM simulations, atom charges form QM regions need special treatment.
2984     if ( qmForcesOn ) {
2985         return 1;
2986     }
2987     if ( ! myMgr->ungridForcesCount && ! myMgr->recipEvirCount ) return 0;  // work to do, enqueue as usual
2988     myMgr->heldComputes.add(this);
2989     return 1;  // don't enqueue yet
2990   }
2991
2992   positionBox->skip();
2993   forceBox->skip();
2994
2995   if ( ++(myMgr->noWorkCount) == myMgr->pmeComputes.size() ) {
2996     myMgr->noWorkCount = 0;
2997     myMgr->reduction->submit();
2998   }
2999
3000   atomsChanged = 0;
3001
3002   return 1;  // no work for this step
3003 }
3004
3005 void ComputePmeMgr::addRecipEvirClient() {
3006   ++recipEvirClients;
3007 }
3008
3009 void ComputePmeMgr::recvRecipEvir(PmeEvirMsg *msg) {
3010   if ( ! pmeComputes.size() ) NAMD_bug("ComputePmeMgr::recvRecipEvir() called on pe without patches");
3011   for ( int g=0; g<numGrids; ++g ) {
3012     evir[g] += msg->evir[g];
3013   }
3014   delete msg;
3015   // CkPrintf("recvRecipEvir pe %d %d %d\n", CkMyPe(), ungridForcesCount, recipEvirCount);
3016   if ( ! --recipEvirCount && ! ungridForcesCount ) submitReductions();
3017 }
3018
3019 void ComputePme::doQMWork() {
3020     
3021 //     iout << CkMyPe() << ") ----> PME doQMWork.\n" << endi ;
3022     
3023     
3024     int numQMAtms = Node::Object()->molecule->get_numQMAtoms();
3025     const Real *qmAtmChrg = Node::Object()->molecule->get_qmAtmChrg() ;
3026     const int *qmAtmIndx = Node::Object()->molecule->get_qmAtmIndx() ;
3027     const Real *qmAtomGroup = Node::Object()->molecule->get_qmAtomGroup() ;
3028     
3029     const CompAtomExt *xExt = patch->getCompAtomExtInfo();
3030     
3031     // Determine number of qm atoms in this patch for the current step.
3032     numLocalQMAtoms = 0;
3033     for (int paIter=0; paIter<patch->getNumAtoms(); paIter++) {
3034         if ( qmAtomGroup[xExt[paIter].id] != 0 ) {
3035             numLocalQMAtoms++;
3036         }
3037     }
3038     
3039     // We prepare a charge vector with QM charges for use in the PME calculation.
3040     
3041     // Clears data from last step, if there is any.
3042     if (qmLoclIndx != 0)
3043         delete [] qmLoclIndx;
3044     if (qmLocalCharges != 0)
3045         delete [] qmLocalCharges;
3046     
3047     qmLoclIndx = new int[numLocalQMAtoms] ;
3048     qmLocalCharges = new Real[numLocalQMAtoms] ;
3049     
3050     // I am assuming there will be (in general) more QM atoms among all QM groups
3051     // than MM atoms in a patch.
3052     int procAtms = 0;
3053     
3054     for (int paIter=0; paIter<patch->getNumAtoms(); paIter++) {
3055         
3056         for (int i=0; i<numQMAtms; i++) {
3057             
3058             if (qmAtmIndx[i] == xExt[paIter].id) {
3059                 
3060                 qmLoclIndx[procAtms] = paIter ;
3061                 qmLocalCharges[procAtms] = qmAtmChrg[i];
3062                 
3063                 procAtms++;
3064                 break;
3065             }
3066             
3067         }
3068         
3069         if (procAtms == numLocalQMAtoms)
3070             break;
3071     }
3072     
3073     doWork();
3074     return ;
3075 }
3076
3077 void ComputePme::doWork()
3078 {
3079   DebugM(4,"Entering ComputePme::doWork().\n");
3080
3081   if ( basePriority >= COMPUTE_HOME_PRIORITY ) {
3082 #ifdef NAMD_CUDA
3083     basePriority = ( offload ? PME_OFFLOAD_PRIORITY : PME_PRIORITY );
3084 #else
3085     basePriority = PME_PRIORITY;
3086 #endif
3087     ungridForces();
3088     // CkPrintf("doWork 2 pe %d %d %d\n", CkMyPe(), myMgr->ungridForcesCount, myMgr->recipEvirCount);
3089     if ( ! --(myMgr->ungridForcesCount) && ! myMgr->recipEvirCount ) myMgr->submitReductions();
3090     return;
3091   }
3092   basePriority = COMPUTE_HOME_PRIORITY + PATCH_PRIORITY(patchID);
3093   // CkPrintf("doWork 1 pe %d %d %d\n", CkMyPe(), myMgr->ungridForcesCount, myMgr->recipEvirCount);
3094
3095 #ifdef TRACE_COMPUTE_OBJECTS
3096     double traceObjStartTime = CmiWallTimer();
3097 #endif
3098
3099 #ifdef NAMD_CUDA
3100   if ( offload ) cudaSetDevice(deviceCUDA->getDeviceID());
3101 #endif
3102
3103   // allocate storage
3104   numLocalAtoms = patch->getNumAtoms();
3105
3106   Lattice &lattice = patch->flags.lattice;
3107
3108   localData_alloc.resize(numLocalAtoms*(numGrids+ ((numGrids>1 || selfOn)?1:0)));
3109   localData = localData_alloc.begin();
3110   localPartition_alloc.resize(numLocalAtoms);
3111   localPartition = localPartition_alloc.begin();
3112
3113   int g;
3114   for ( g=0; g<numGrids; ++g ) {
3115     localGridData[g] = localData + numLocalAtoms*(g+1);
3116   }
3117
3118   // get positions and charges
3119   PmeParticle * data_ptr = localData;
3120   unsigned char * part_ptr = localPartition;
3121   const BigReal coulomb_sqrt = sqrt( COULOMB * ComputeNonbondedUtil::scaling
3122                                 * ComputeNonbondedUtil::dielectric_1 );
3123
3124   {
3125     CompAtom *x = positionBox->open();
3126     // CompAtomExt *xExt = patch->getCompAtomExtInfo();
3127     if ( patch->flags.doMolly ) {
3128       positionBox->close(&x);
3129       x = avgPositionBox->open();
3130     }
3131     int numAtoms = patch->getNumAtoms();
3132
3133     for(int i=0; i<numAtoms; ++i)
3134     {
3135       data_ptr->x = x[i].position.x;
3136       data_ptr->y = x[i].position.y;
3137       data_ptr->z = x[i].position.z;
3138       data_ptr->cg = coulomb_sqrt * x[i].charge;
3139       ++data_ptr;
3140       *part_ptr = x[i].partition;
3141       ++part_ptr;
3142     }
3143
3144     // QM loop to overwrite charges of QM atoms.
3145     // They are zero for NAMD, but are updated in ComputeQM.
3146     if ( qmForcesOn ) {
3147         
3148         for(int i=0; i<numLocalQMAtoms; ++i)
3149         {
3150           localData[qmLoclIndx[i]].cg = coulomb_sqrt * qmLocalCharges[i];
3151         }
3152         
3153     }
3154     
3155     if ( patch->flags.doMolly ) { avgPositionBox->close(&x); }
3156     else { positionBox->close(&x); }
3157   }
3158
3159   // copy to other grids if needed
3160   if ( (alchOn && (!alchDecouple)) || lesOn ) {
3161     for ( g=0; g<numGrids; ++g ) {
3162       PmeParticle *lgd = localGridData[g];
3163       int nga = 0;
3164       for(int i=0; i<numLocalAtoms; ++i) {
3165         if ( localPartition[i] == 0 || localPartition[i] == (g+1) ) {
3166           // for FEP/TI: grid 0 gets non-alch + partition 1;
3167           // grid 1 gets non-alch + partition 2;
3168           // grid 2 (only if called for with numGrids=3) gets only non-alch
3169           lgd[nga++] = localData[i];
3170         }
3171       }
3172       numGridAtoms[g] = nga;
3173     }
3174   } else if ( alchOn && alchDecouple) {
3175     // alchemical decoupling: four grids
3176     // g=0: partition 0 and partition 1
3177     // g=1: partition 0 and partition 2
3178     // g=2: only partition 1 atoms
3179     // g=3: only partition 2 atoms
3180     // plus one grid g=4, only partition 0, if numGrids=5
3181     for ( g=0; g<2; ++g ) {  // same as before for first 2
3182       PmeParticle *lgd = localGridData[g];
3183       int nga = 0;
3184       for(int i=0; i<numLocalAtoms; ++i) {
3185         if ( localPartition[i] == 0 || localPartition[i] == (g+1) ) {
3186           lgd[nga++] = localData[i];
3187         }
3188       }
3189       numGridAtoms[g] = nga;
3190     }
3191     for (g=2 ; g<4 ; ++g ) {  // only alchemical atoms for these 2
3192       PmeParticle *lgd = localGridData[g];
3193       int nga = 0;
3194       for(int i=0; i<numLocalAtoms; ++i) {
3195         if ( localPartition[i] == (g-1) ) {
3196           lgd[nga++] = localData[i];
3197         }
3198       }
3199       numGridAtoms[g] = nga;
3200     }
3201     for (g=4 ; g<numGrids ; ++g ) {  // only non-alchemical atoms 
3202       // numGrids=5 only if alchElecLambdaStart > 0
3203       PmeParticle *lgd = localGridData[g];
3204       int nga = 0;
3205       for(int i=0; i<numLocalAtoms; ++i) {
3206         if ( localPartition[i] == 0 ) {
3207           lgd[nga++] = localData[i];
3208         }
3209       }
3210       numGridAtoms[g] = nga;
3211     }
3212   } else if ( selfOn ) {
3213     if ( numGrids != 1 ) NAMD_bug("ComputePme::doWork assertion 1 failed");
3214     g = 0;
3215     PmeParticle *lgd = localGridData[g];
3216     int nga = 0;
3217     for(int i=0; i<numLocalAtoms; ++i) {
3218       if ( localPartition[i] == 1 ) {
3219         lgd[nga++] = localData[i];
3220       }
3221     }
3222     numGridAtoms[g] = nga;
3223   } else if ( pairOn ) {
3224     if ( numGrids != 3 ) NAMD_bug("ComputePme::doWork assertion 2 failed");
3225     g = 0;
3226     PmeParticle *lgd = localGridData[g];
3227     int nga = 0;
3228     for(int i=0; i<numLocalAtoms; ++i) {
3229       if ( localPartition[i] == 1 || localPartition[i] == 2 ) {
3230         lgd[nga++] = localData[i];
3231       }
3232     }
3233     numGridAtoms[g] = nga;
3234     for ( g=1; g<3; ++g ) {
3235       PmeParticle *lgd = localGridData[g];
3236       int nga = 0;
3237       for(int i=0; i<numLocalAtoms; ++i) {
3238         if ( localPartition[i] == g ) {
3239           lgd[nga++] = localData[i];
3240         }
3241       }
3242       numGridAtoms[g] = nga;
3243     }
3244   } else {
3245     if ( numGrids != 1 ) NAMD_bug("ComputePme::doWork assertion 3 failed");
3246     localGridData[0] = localData;
3247     numGridAtoms[0] = numLocalAtoms;
3248   }
3249
3250  if ( ! myMgr->doWorkCount ) {
3251   myMgr->doWorkCount = myMgr->pmeComputes.size();
3252
3253 #ifdef NAMD_CUDA
3254  if ( !  offload )
3255 #endif // NAMD_CUDA
3256  {
3257   memset( (void*) myMgr->fz_arr, 0, (myGrid.K3+myGrid.order-1) * sizeof(char) );
3258
3259   for (int i=0; i<myMgr->q_count; ++i) {
3260     memset( (void*) (myMgr->q_list[i]), 0, (myGrid.K3+myGrid.order-1) * sizeof(float) );
3261   }
3262  }
3263
3264   for ( g=0; g<numGrids; ++g ) {
3265     myMgr->evir[g] = 0;
3266   }
3267
3268   myMgr->strayChargeErrors = 0;
3269
3270   myMgr->compute_sequence = sequence();
3271  }
3272
3273   if ( sequence() != myMgr->compute_sequence ) NAMD_bug("ComputePme sequence mismatch in doWork()");
3274
3275   int strayChargeErrors = 0;
3276
3277   // calculate self energy
3278   BigReal ewaldcof = ComputeNonbondedUtil::ewaldcof;
3279   for ( g=0; g<numGrids; ++g ) {
3280     BigReal selfEnergy = 0;
3281     data_ptr = localGridData[g];
3282     int i;
3283     for(i=0; i<numGridAtoms[g]; ++i)
3284     {
3285       selfEnergy += data_ptr->cg * data_ptr->cg;
3286       ++data_ptr;
3287     }
3288     selfEnergy *= -1. * ewaldcof / SQRT_PI;
3289     myMgr->evir[g][0] += selfEnergy;
3290
3291     float **q = myMgr->q_arr + g*myMgr->fsize;
3292     char *f = myMgr->f_arr + g*myMgr->fsize;
3293
3294     scale_coordinates(localGridData[g], numGridAtoms[g], lattice, myGrid);
3295 #ifdef NAMD_CUDA
3296    if ( offload ) {
3297     if ( myMgr->cuda_atoms_alloc == 0 ) {  // first call
3298       int na = myMgr->cuda_atoms_alloc = 1.2 * (myMgr->cuda_atoms_count + 1000);
3299       cuda_errcheck("before malloc atom data for pme");
3300       cudaMallocHost((void**) &(myMgr->a_data_host), 7*na*sizeof(float));
3301       cudaMalloc((void**) &(myMgr->a_data_dev), 7*na*sizeof(float));
3302       cuda_errcheck("malloc atom data for pme");
3303       myMgr->cuda_atoms_count = 0;
3304     }
3305     cuda_atoms_offset = myMgr->cuda_atoms_count;
3306     int n = numGridAtoms[g];
3307     myMgr->cuda_atoms_count += n;
3308     if ( myMgr->cuda_atoms_count > myMgr->cuda_atoms_alloc ) {
3309       CkPrintf("Pe %d expanding CUDA PME atoms allocation because %d > %d\n",
3310                         CkMyPe(), myMgr->cuda_atoms_count, myMgr->cuda_atoms_alloc);
3311       cuda_errcheck("before malloc expanded atom data for pme");
3312       int na = myMgr->cuda_atoms_alloc = 1.2 * (myMgr->cuda_atoms_count + 1000);
3313       const float *a_data_host_old = myMgr->a_data_host;
3314       cudaMallocHost((void**) &(myMgr->a_data_host), 7*na*sizeof(float));
3315       cuda_errcheck("malloc expanded host atom data for pme");
3316       memcpy(myMgr->a_data_host, a_data_host_old, 7*cuda_atoms_offset*sizeof(float));
3317       cudaFreeHost((void*) a_data_host_old);
3318       cuda_errcheck("free expanded host atom data for pme");
3319       cudaFree(myMgr->a_data_dev);
3320       cuda_errcheck("free expanded dev atom data for pme");
3321       cudaMalloc((void**) &(myMgr->a_data_dev), 7*na*sizeof(float));
3322       cuda_errcheck("malloc expanded dev atom data for pme");
3323     }
3324     float *a_data_host = myMgr->a_data_host + 7 * cuda_atoms_offset;
3325     data_ptr = localGridData[g];
3326     double order_1 = myGrid.order - 1;
3327     double K1 = myGrid.K1;
3328     double K2 = myGrid.K2;
3329     double K3 = myGrid.K3;
3330     int found_negative = 0;
3331     for ( int i=0; i<n; ++i ) {
3332       if ( data_ptr[i].x < 0 || data_ptr[i].y < 0 || data_ptr[i].z < 0 ) {
3333         found_negative = 1;
3334         // CkPrintf("low coord: %f %f %f\n", data_ptr[i].x, data_ptr[i].y, data_ptr[i].z);
3335       }
3336       double x_int = (int) data_ptr[i].x;
3337       double y_int = (int) data_ptr[i].y;
3338       double z_int = (int) data_ptr[i].z;
3339       a_data_host[7*i  ] = data_ptr[i].x - x_int;  // subtract in double precision
3340       a_data_host[7*i+1] = data_ptr[i].y - y_int;
3341       a_data_host[7*i+2] = data_ptr[i].z - z_int;
3342       a_data_host[7*i+3] = data_ptr[i].cg;
3343       x_int -= order_1;  if ( x_int < 0 ) x_int += K1;
3344       y_int -= order_1;  if ( y_int < 0 ) y_int += K2;
3345       z_int -= order_1;  if ( z_int < 0 ) z_int += K3;
3346       a_data_host[7*i+4] = x_int;
3347       a_data_host[7*i+5] = y_int;
3348       a_data_host[7*i+6] = z_int;
3349     }
3350     if ( found_negative ) NAMD_bug("found negative atom coordinate in ComputePme::doWork");
3351    } else
3352 #endif // NAMD_CUDA
3353    {
3354     myRealSpace[g]->set_num_atoms(numGridAtoms[g]);
3355     myRealSpace[g]->fill_charges(q, myMgr->q_list, myMgr->q_count, strayChargeErrors, f, myMgr->fz_arr, localGridData[g]);
3356    }
3357   }
3358   myMgr->strayChargeErrors += strayChargeErrors;
3359
3360 #ifdef TRACE_COMPUTE_OBJECTS
3361     traceUserBracketEvent(TRACE_COMPOBJ_IDOFFSET+this->cid, traceObjStartTime, CmiWallTimer());
3362 #endif
3363
3364  if ( --(myMgr->doWorkCount) == 0 ) {
3365 // cudaDeviceSynchronize();  // XXXX
3366 #ifdef NAMD_CUDA
3367   if ( offload ) {
3368     ComputePmeMgr::cuda_submit_charges_args args;
3369     args.mgr = myMgr;
3370     args.lattice = &lattice;
3371     args.sequence = sequence();
3372     CmiLock(ComputePmeMgr::cuda_lock);
3373     if ( ComputePmeMgr::cuda_busy ) {
3374       ComputePmeMgr::cuda_submit_charges_deque.push_back(args);
3375     } else if ( CkMyPe() == deviceCUDA->getMasterPe() ) {
3376       // avoid adding work to nonbonded data preparation pe
3377       args.mgr->cuda_submit_charges(*args.lattice, args.sequence);
3378     } else {
3379       ComputePmeMgr::cuda_busy = true;
3380       while ( 1 ) {
3381         CmiUnlock(ComputePmeMgr::cuda_lock);
3382         args.mgr->cuda_submit_charges(*args.lattice, args.sequence);
3383         CmiLock(ComputePmeMgr::cuda_lock);
3384         if ( ComputePmeMgr::cuda_submit_charges_deque.size() ) {
3385           args = ComputePmeMgr::cuda_submit_charges_deque.front();
3386           ComputePmeMgr::cuda_submit_charges_deque.pop_front();
3387         } else {
3388           ComputePmeMgr::cuda_busy = false;
3389           break;
3390         }
3391       }
3392     }
3393     CmiUnlock(ComputePmeMgr::cuda_lock);
3394   } else
3395 #endif // NAMD_CUDA
3396   {
3397     myMgr->chargeGridReady(lattice,sequence());
3398   }
3399  }
3400  atomsChanged = 0;
3401 }
3402
3403 #ifdef NAMD_CUDA
3404
3405 void ComputePmeMgr::cuda_submit_charges(Lattice &lattice, int sequence) {
3406
3407     int n = cuda_atoms_count;
3408     //CkPrintf("pe %d cuda_atoms_count %d\n", CkMyPe(), cuda_atoms_count);
3409     cuda_atoms_count = 0;
3410
3411     const double before = CmiWallTimer();
3412     cudaMemcpyAsync(a_data_dev, a_data_host, 7*n*sizeof(float),
3413                           cudaMemcpyHostToDevice, streams[stream]);
3414     const double after = CmiWallTimer();
3415
3416     cudaStreamWaitEvent(streams[stream], nodePmeMgr->end_charge_memset, 0);
3417
3418     cuda_pme_charges(
3419       bspline_coeffs_dev,
3420       q_arr_dev, ffz_dev, ffz_dev + fsize,
3421       a_data_dev, n,
3422       myGrid.K1, myGrid.K2, myGrid.K3, myGrid.order,
3423       streams[stream]);
3424     const double after2 = CmiWallTimer();
3425
3426     chargeGridSubmitted(lattice,sequence);  // must be inside lock
3427
3428     masterPmeMgr->charges_time = before;
3429     traceUserBracketEvent(CUDA_EVENT_ID_PME_COPY,before,after);
3430     traceUserBracketEvent(CUDA_EVENT_ID_PME_KERNEL,after,after2);
3431 }
3432
3433 void cuda_check_pme_charges(void *arg, double walltime) {
3434   ComputePmeMgr *argp = (ComputePmeMgr *) arg;
3435
3436   cudaError_t err = cudaEventQuery(argp->end_charges);
3437   if ( err == cudaSuccess ) {
3438     traceUserBracketEvent(CUDA_EVENT_ID_PME_CHARGES,argp->charges_time,walltime);
3439     argp->charges_time = walltime - argp->charges_time;
3440     argp->sendChargeGridReady();
3441     argp->check_charges_count = 0;
3442   } else if ( err != cudaErrorNotReady ) {
3443     cuda_errcheck("in cuda_check_pme_charges");
3444     NAMD_bug("cuda_errcheck missed error in cuda_check_pme_charges");
3445   } else if ( ++(argp->check_charges_count) >= count_limit ) {
3446     char errmsg[256];
3447     sprintf(errmsg,"cuda_check_pme_charges polled %d times over %f s on seq %d",
3448             argp->check_charges_count, walltime - argp->charges_time,
3449             argp->saved_sequence);
3450     cuda_errcheck(errmsg);
3451     NAMD_die(errmsg);
3452   } else {
3453     CcdCallBacksReset(0,walltime);  // fix Charm++
3454     CUDA_POLL(cuda_check_pme_charges, arg);
3455   }
3456 }
3457
3458 void ComputePmeMgr::chargeGridSubmitted(Lattice &lattice, int sequence) {
3459   saved_lattice = &lattice;
3460   saved_sequence = sequence;
3461
3462   // cudaDeviceSynchronize();  //  XXXX TESTING
3463   //int q_stride = myGrid.K3+myGrid.order-1;
3464   //for (int n=fsize+q_stride, j=0; j<n; ++j) {
3465   //  if ( ffz_host[j] != 0 && ffz_host[j] != 1 ) {
3466   //    CkPrintf("pre-memcpy flag %d/%d == %d on pe %d in ComputePmeMgr::chargeGridReady\n", j, n, ffz_host[j], CkMyPe());
3467   //  }
3468   //}
3469   //CmiLock(cuda_lock);
3470
3471  if ( --(masterPmeMgr->chargeGridSubmittedCount) == 0 ) {
3472   double before = CmiWallTimer();
3473   cudaEventRecord(nodePmeMgr->end_all_pme_kernels, 0);  // when all streams complete
3474   cudaStreamWaitEvent(streams[stream], nodePmeMgr->end_all_pme_kernels, 0);
3475   cudaMemcpyAsync(q_data_host, q_data_dev, q_data_size+ffz_size,
3476                         cudaMemcpyDeviceToHost, streams[stream]);
3477   traceUserBracketEvent(CUDA_EVENT_ID_PME_COPY,before,CmiWallTimer());
3478   cudaEventRecord(masterPmeMgr->end_charges, streams[stream]);
3479   cudaMemsetAsync(q_data_dev, 0, q_data_size + ffz_size, streams[stream]);  // for next time
3480   cudaEventRecord(nodePmeMgr->end_charge_memset, streams[stream]);
3481   //CmiUnlock(cuda_lock);
3482   // cudaDeviceSynchronize();  //  XXXX TESTING
3483   // cuda_errcheck("after memcpy grid to host");
3484
3485   SimParameters *simParams = Node::Object()->simParameters;
3486   if ( ! simParams->useCUDA2 ) {
3487     CProxy_ComputeMgr cm(CkpvAccess(BOCclass_group).computeMgr);
3488     cm[deviceCUDA->getMasterPe()].recvYieldDevice(-1);
3489   }
3490
3491   pmeProxy[master_pe].pollChargeGridReady();
3492  }
3493 }
3494
3495 void ComputePmeMgr::sendChargeGridReady() {
3496   for ( int i=0; i<CkMyNodeSize(); ++i ) {
3497     ComputePmeMgr *mgr = nodePmeMgr->mgrObjects[i];
3498     int cs = mgr->pmeComputes.size();
3499     if ( cs ) {
3500       mgr->ungridForcesCount = cs;
3501       mgr->recipEvirCount = mgr->recipEvirClients;
3502       masterPmeMgr->chargeGridSubmittedCount++;
3503     }
3504   }
3505   pmeProxy[master_pe].recvChargeGridReady();
3506 }
3507 #endif // NAMD_CUDA
3508
3509 void ComputePmeMgr::pollChargeGridReady() {
3510 #ifdef NAMD_CUDA
3511   CcdCallBacksReset(0,CmiWallTimer());  // fix Charm++
3512   CUDA_POLL(cuda_check_pme_charges,this);
3513 #else
3514   NAMD_bug("ComputePmeMgr::pollChargeGridReady() called in non-CUDA build.");
3515 #endif
3516 }
3517
3518 void ComputePmeMgr::recvChargeGridReady() {
3519   chargeGridReady(*saved_lattice,saved_sequence);
3520 }
3521
3522 void ComputePmeMgr::chargeGridReady(Lattice &lattice, int sequence) {
3523
3524 #ifdef NAMD_CUDA
3525  if ( offload ) {
3526   int errcount = 0;
3527   int q_stride = myGrid.K3+myGrid.order-1;
3528   for (int n=fsize+q_stride, j=fsize; j<n; ++j) {
3529     f_arr[j] = ffz_host[j];
3530     if ( ffz_host[j] & ~1 ) ++errcount;
3531   }
3532   if ( errcount ) NAMD_bug("bad flag in ComputePmeMgr::chargeGridReady");
3533  }
3534 #endif
3535   recipEvirCount = recipEvirClients;
3536   ungridForcesCount = pmeComputes.size();
3537
3538   for (int j=0; j<myGrid.order-1; ++j) {
3539     fz_arr[j] |= fz_arr[myGrid.K3+j];
3540   }
3541
3542   if ( usePencils ) {
3543     sendPencils(lattice,sequence);
3544   } else {
3545     sendData(lattice,sequence);
3546   }
3547 }
3548
3549
3550 void ComputePmeMgr::sendPencilsPart(int first, int last, Lattice &lattice, int sequence, int sourcepe) {
3551
3552   // iout << "Sending charge grid for " << numLocalAtoms << " atoms to FFT on " << iPE << ".\n" << endi;
3553
3554 #if 0 && USE_PERSISTENT
3555     if (recvGrid_handle== NULL) setup_recvgrid_persistent();
3556 #endif
3557   int K1 = myGrid.K1;
3558   int K2 = myGrid.K2;
3559   int dim2 = myGrid.dim2;
3560   int dim3 = myGrid.dim3;
3561   int block1 = myGrid.block1;
3562   int block2 = myGrid.block2;
3563
3564   // int savedMessages = 0;
3565   NodePmeMgr *npMgr = pmeNodeProxy[CkMyNode()].ckLocalBranch();
3566
3567   for (int ap=first; ap<=last; ++ap) {
3568     int ib = activePencils[ap].i;
3569     int jb = activePencils[ap].j;
3570     int ibegin = ib*block1;
3571     int iend = ibegin + block1;  if ( iend > K1 ) iend = K1;
3572     int jbegin = jb*block2;
3573     int jend = jbegin + block2;  if ( jend > K2 ) jend = K2;
3574     int flen = numGrids * (iend - ibegin) * (jend - jbegin);
3575
3576     int fcount = 0;
3577     for ( int g=0; g<numGrids; ++g ) {
3578       char *f = f_arr + g*fsize;
3579 #ifdef NAMD_CUDA
3580      if ( offload ) {
3581       int errcount = 0;
3582       for ( int i=ibegin; i<iend; ++i ) {
3583        for ( int j=jbegin; j<jend; ++j ) {
3584         int k = i*dim2+j;
3585         f[k] = ffz_host[k];
3586         fcount += f[k];
3587         if ( ffz_host[k] & ~1 ) ++errcount;
3588        }
3589       }
3590       if ( errcount ) NAMD_bug("bad flag in ComputePmeMgr::sendPencilsPart");
3591      } else
3592 #endif
3593       for ( int i=ibegin; i<iend; ++i ) {
3594        for ( int j=jbegin; j<jend; ++j ) {
3595         fcount += f[i*dim2+j];
3596        }
3597       }
3598     }
3599
3600 #ifdef NETWORK_PROGRESS
3601     CmiNetworkProgress();
3602 #endif
3603
3604     if ( ! pencilActive[ib*yBlocks+jb] )
3605       NAMD_bug("PME activePencils list inconsistent");
3606
3607     int zlistlen = 0;
3608     for ( int i=0; i<myGrid.K3; ++i ) {
3609       if ( fz_arr[i] ) ++zlistlen;
3610     }
3611
3612     int hd = ( fcount? 1 : 0 );  // has data?
3613     // if ( ! hd ) ++savedMessages;
3614
3615     
3616     PmeGridMsg *msg = new ( hd*zlistlen, hd*flen,
3617         hd*fcount*zlistlen, PRIORITY_SIZE) PmeGridMsg;
3618     msg->sourceNode = sourcepe;
3619     msg->hasData = hd;
3620     msg->lattice = lattice;
3621    if ( hd ) {
3622 #if 0
3623     msg->start = fstart;
3624     msg->len = flen;
3625 #else
3626     msg->start = -1;   // obsolete?
3627     msg->len = -1;   // obsolete?
3628 #endif
3629     msg->zlistlen = zlistlen;
3630     int *zlist = msg->zlist;
3631     zlistlen = 0;
3632     for ( int i=0; i<myGrid.K3; ++i ) {
3633       if ( fz_arr[i] ) zlist[zlistlen++] = i;
3634     }
3635     char *fmsg = msg->fgrid;
3636     float *qmsg = msg->qgrid;
3637     for ( int g=0; g<numGrids; ++g ) {
3638       char *f = f_arr + g*fsize;
3639       float **q = q_arr + g*fsize;
3640       for ( int i=ibegin; i<iend; ++i ) {
3641        for ( int j=jbegin; j<jend; ++j ) {
3642         *(fmsg++) = f[i*dim2+j];
3643         if( f[i*dim2+j] ) {
3644           for (int h=0; h<myGrid.order-1; ++h) {
3645             q[i*dim2+j][h] += q[i*dim2+j][myGrid.K3+h];
3646           }
3647           for ( int k=0; k<zlistlen; ++k ) {
3648             *(qmsg++) = q[i*dim2+j][zlist[k]];
3649           }
3650         }
3651        }
3652       }
3653     }
3654    }
3655
3656     msg->sequence = compute_sequence;
3657     SET_PRIORITY(msg,compute_sequence,PME_GRID_PRIORITY)
3658     CmiEnableUrgentSend(1);
3659 #if USE_NODE_PAR_RECEIVE
3660     msg->destElem=CkArrayIndex3D(ib,jb,0);
3661     CProxy_PmePencilMap lzm = npMgr->zm;
3662     int destproc = lzm.ckLocalBranch()->procNum(0, msg->destElem);
3663     int destnode = CmiNodeOf(destproc);
3664     
3665 #if  0 
3666     CmiUsePersistentHandle(&recvGrid_handle[ap], 1);
3667 #endif
3668     pmeNodeProxy[destnode].recvZGrid(msg);
3669 #if 0 
3670     CmiUsePersistentHandle(NULL, 0);
3671 #endif
3672 #else
3673 #if 0 
3674     CmiUsePersistentHandle(&recvGrid_handle[ap], 1);
3675 #endif
3676     zPencil(ib,jb,0).recvGrid(msg);
3677 #if 0 
3678     CmiUsePersistentHandle(NULL, 0);
3679 #endif
3680 #endif
3681     CmiEnableUrgentSend(0);
3682   }
3683
3684
3685   // if ( savedMessages ) {
3686   //   CkPrintf("Pe %d eliminated %d PME messages\n",CkMyPe(),savedMessages);
3687   // }
3688
3689 }
3690
3691
3692 void ComputePmeMgr::sendPencilsHelper(int iter) {
3693   nodePmeMgr->sendPencilsHelper(iter);
3694 }
3695
3696 void NodePmeMgr::sendPencilsHelper(int iter) {
3697 #ifdef NAMD_CUDA
3698   ComputePmeMgr *obj = masterPmeMgr;
3699   obj->sendPencilsPart(iter, iter, *obj->sendDataHelper_lattice, obj->sendDataHelper_sequence, obj->sendDataHelper_sourcepe);
3700 #else
3701   NAMD_bug("NodePmeMgr::sendPencilsHelper called in non-CUDA build");
3702 #endif
3703 }
3704
3705 void ComputePmeMgr::sendPencils(Lattice &lattice, int sequence) {
3706
3707   sendDataHelper_lattice = &lattice;
3708   sendDataHelper_sequence = sequence;
3709   sendDataHelper_sourcepe = CkMyPe();
3710
3711 #ifdef NAMD_CUDA
3712   if ( offload ) {
3713     for ( int ap=0; ap < numPencilsActive; ++ap ) {
3714 #if CMK_MULTICORE
3715       // nodegroup messages on multicore are delivered to sending pe, or pe 0 if expedited
3716       int ib = activePencils[ap].i;
3717       int jb = activePencils[ap].j;
3718       int destproc = nodePmeMgr->zm.ckLocalBranch()->procNum(0, CkArrayIndex3D(ib,jb,0));
3719       pmeProxy[destproc].sendPencilsHelper(ap);
3720 #else
3721       pmeNodeProxy[CkMyNode()].sendPencilsHelper(ap);
3722 #endif
3723     }
3724   } else
3725 #endif
3726   {
3727     sendPencilsPart(0,numPencilsActive-1,lattice,sequence,CkMyPe());
3728   }
3729
3730   if ( strayChargeErrors ) {
3731    strayChargeErrors = 0;
3732    iout << iERROR << "Stray PME grid charges detected: "
3733         << CkMyPe() << " sending to (x,y)";
3734    int K1 = myGrid.K1;
3735    int K2 = myGrid.K2;
3736    int dim2 = myGrid.dim2;
3737    int block1 = myGrid.block1;
3738    int block2 = myGrid.block2;
3739    for (int ib=0; ib<xBlocks; ++ib) {
3740     for (int jb=0; jb<yBlocks; ++jb) {
3741      int ibegin = ib*block1;
3742      int iend = ibegin + block1;  if ( iend > K1 ) iend = K1;
3743      int jbegin = jb*block2;
3744      int jend = jbegin + block2;  if ( jend > K2 ) jend = K2;
3745      int flen = numGrids * (iend - ibegin) * (jend - jbegin);
3746
3747      for ( int g=0; g<numGrids; ++g ) {
3748        char *f = f_arr + g*fsize;
3749        if ( ! pencilActive[ib*yBlocks+jb] ) {
3750            for ( int i=ibegin; i<iend; ++i ) {
3751             for ( int j=jbegin; j<jend; ++j ) {
3752              if ( f[i*dim2+j] == 3 ) {
3753                f[i*dim2+j] = 2;
3754                iout << " (" << i << "," << j << ")";
3755              }
3756             }
3757            }
3758        }
3759      }
3760     }
3761    }
3762    iout << "\n" << endi;
3763   }
3764  
3765 }
3766
3767
3768 void ComputePmeMgr::copyPencils(PmeGridMsg *msg) {
3769
3770   int K1 = myGrid.K1;
3771   int K2 = myGrid.K2;
3772   int dim2 = myGrid.dim2;
3773   int dim3 = myGrid.dim3;
3774   int block1 = myGrid.block1;
3775   int block2 = myGrid.block2;
3776
3777   // msg->sourceNode = thisIndex.x * initdata.yBlocks + thisIndex.y;
3778   int ib = msg->sourceNode / yBlocks;
3779   int jb = msg->sourceNode % yBlocks;
3780
3781   int ibegin = ib*block1;
3782   int iend = ibegin + block1;  if ( iend > K1 ) iend = K1;
3783   int jbegin = jb*block2;
3784   int jend = jbegin + block2;  if ( jend > K2 ) jend = K2;
3785
3786   int zlistlen = msg->zlistlen;
3787   int *zlist = msg->zlist;
3788   float *qmsg = msg->qgrid;
3789   int g;
3790   for ( g=0; g<numGrids; ++g ) {
3791     char *f = f_arr + g*fsize;
3792     float **q = q_arr + g*fsize;
3793     for ( int i=ibegin; i<iend; ++i ) {
3794      for ( int j=jbegin; j<jend; ++j ) {
3795       if( f[i*dim2+j] ) {
3796         f[i*dim2+j] = 0;
3797         for ( int k=0; k<zlistlen; ++k ) {
3798           q[i*dim2+j][zlist[k]] = *(qmsg++);
3799         }
3800         for (int h=0; h<myGrid.order-1; ++h) {
3801           q[i*dim2+j][myGrid.K3+h] = q[i*dim2+j][h];
3802         }
3803       }
3804      }
3805     }
3806   }
3807 }
3808
3809
3810 void ComputePmeMgr::sendDataPart(int first, int last, Lattice &lattice, int sequence, int sourcepe, int errors) {
3811
3812   // iout << "Sending charge grid for " << numLocalAtoms << " atoms to FFT on " << iPE << ".\n" << endi;
3813
3814   bsize = myGrid.block1 * myGrid.dim2 * myGrid.dim3;
3815
3816   CProxy_ComputePmeMgr pmeProxy(CkpvAccess(BOCclass_group).computePmeMgr);
3817   for (int j=first; j<=last; j++) {
3818     int pe = gridPeOrder[j];  // different order
3819     if ( ! recipPeDest[pe] && ! errors ) continue;
3820     int start = pe * bsize;
3821     int len = bsize;
3822     if ( start >= qsize ) { start = 0; len = 0; }
3823     if ( start + len > qsize ) { len = qsize - start; }
3824     int zdim = myGrid.dim3;
3825     int fstart = start / zdim;
3826     int flen = len / zdim;
3827     int fcount = 0;
3828     int i;
3829
3830     int g;
3831     for ( g=0; g<numGrids; ++g ) {
3832       char *f = f_arr + fstart + g*fsize;
3833 #ifdef NAMD_CUDA
3834      if ( offload ) {
3835       int errcount = 0;
3836       for ( i=0; i<flen; ++i ) {
3837         f[i] = ffz_host[fstart+i];
3838         fcount += f[i];
3839         if ( ffz_host[fstart+i] & ~1 ) ++errcount;
3840       }
3841       if ( errcount ) NAMD_bug("bad flag in ComputePmeMgr::sendDataPart");
3842      } else
3843 #endif
3844       for ( i=0; i<flen; ++i ) {
3845         fcount += f[i];
3846       }
3847       if ( ! recipPeDest[pe] ) {
3848         int errfound = 0;
3849         for ( i=0; i<flen; ++i ) {
3850           if ( f[i] == 3 ) {
3851             errfound = 1;
3852             break;
3853           }
3854         }
3855         if ( errfound ) {
3856           iout << iERROR << "Stray PME grid charges detected: "
3857                 << sourcepe << " sending to " << gridPeMap[pe] << " for planes";
3858           int iz = -1;
3859           for ( i=0; i<flen; ++i ) {
3860             if ( f[i] == 3 ) {
3861               f[i] = 2;
3862               int jz = (i+fstart)/myGrid.K2;
3863               if ( iz != jz ) { iout << " " << jz;  iz = jz; }
3864             }
3865           }
3866           iout << "\n" << endi;
3867         }
3868       }
3869     }
3870
3871 #ifdef NETWORK_PROGRESS
3872     CmiNetworkProgress();
3873 #endif
3874
3875     if ( ! recipPeDest[pe] ) continue;
3876
3877     int zlistlen = 0;
3878     for ( i=0; i<myGrid.K3; ++i ) {
3879       if ( fz_arr[i] ) ++zlistlen;
3880     }
3881
3882     PmeGridMsg *msg = new (zlistlen, flen*numGrids,
3883                                 fcount*zlistlen, PRIORITY_SIZE) PmeGridMsg;
3884
3885     msg->sourceNode = sourcepe;
3886     msg->lattice = lattice;
3887     msg->start = fstart;
3888     msg->len = flen;
3889     msg->zlistlen = zlistlen;
3890     int *zlist = msg->zlist;
3891     zlistlen = 0;
3892     for ( i=0; i<myGrid.K3; ++i ) {
3893       if ( fz_arr[i] ) zlist[zlistlen++] = i;
3894     }
3895     float *qmsg = msg->qgrid;
3896     for ( g=0; g<numGrids; ++g ) {
3897       char *f = f_arr + fstart + g*fsize;
3898       CmiMemcpy((void*)(msg->fgrid+g*flen),(void*)f,flen*sizeof(char));
3899       float **q = q_arr + fstart + g*fsize;
3900       for ( i=0; i<flen; ++i ) {
3901         if ( f[i] ) {
3902           for (int h=0; h<myGrid.order-1; ++h) {
3903             q[i][h] += q[i][myGrid.K3+h];
3904           }
3905           for ( int k=0; k<zlistlen; ++k ) {
3906             *(qmsg++) = q[i][zlist[k]];
3907           }
3908         }
3909       }
3910     }
3911
3912     msg->sequence = compute_sequence;
3913     SET_PRIORITY(msg,compute_sequence,PME_GRID_PRIORITY)
3914     pmeProxy[gridPeMap[pe]].recvGrid(msg);
3915   }
3916
3917 }
3918
3919 void ComputePmeMgr::sendDataHelper(int iter) {
3920   nodePmeMgr->sendDataHelper(iter);
3921 }
3922
3923 void NodePmeMgr::sendDataHelper(int iter) {
3924 #ifdef NAMD_CUDA
3925   ComputePmeMgr *obj = masterPmeMgr;
3926   obj->sendDataPart(iter, iter, *obj->sendDataHelper_lattice, obj->sendDataHelper_sequence, obj->sendDataHelper_sourcepe, obj->sendDataHelper_errors);
3927 #else
3928   NAMD_bug("NodePmeMgr::sendDataHelper called in non-CUDA build");
3929 #endif
3930 }
3931
3932 void ComputePmeMgr::sendData(Lattice &lattice, int sequence) {
3933
3934   sendDataHelper_lattice = &lattice;
3935   sendDataHelper_sequence = sequence;
3936   sendDataHelper_sourcepe = CkMyPe();
3937   sendDataHelper_errors = strayChargeErrors;
3938   strayChargeErrors = 0;
3939
3940 #ifdef NAMD_CUDA
3941   if ( offload ) {
3942     for ( int i=0; i < numGridPes; ++i ) {
3943       int pe = gridPeOrder[i];  // different order
3944       if ( ! recipPeDest[pe] && ! sendDataHelper_errors ) continue;
3945 #if CMK_MULTICORE
3946       // nodegroup messages on multicore are delivered to sending pe, or pe 0 if expedited
3947       pmeProxy[gridPeMap[pe]].sendDataHelper(i);
3948 #else
3949       pmeNodeProxy[CkMyNode()].sendDataHelper(i);
3950 #endif
3951     }
3952   } else
3953 #endif
3954   {
3955     sendDataPart(0,numGridPes-1,lattice,sequence,CkMyPe(),sendDataHelper_errors);
3956   }
3957  
3958 }
3959
3960 void ComputePmeMgr::copyResults(PmeGridMsg *msg) {
3961
3962   int zdim = myGrid.dim3;
3963   int flen = msg->len;
3964   int fstart = msg->start;
3965   int zlistlen = msg->zlistlen;
3966   int *zlist = msg->zlist;
3967   float *qmsg = msg->qgrid;
3968   int g;
3969   for ( g=0; g<numGrids; ++g ) {
3970     char *f = msg->fgrid + g*flen;
3971     float **q = q_arr + fstart + g*fsize;
3972     for ( int i=0; i<flen; ++i ) {
3973       if ( f[i] ) {
3974         f[i] = 0;
3975         for ( int k=0; k<zlistlen; ++k ) {
3976           q[i][zlist[k]] = *(qmsg++);
3977         }
3978         for (int h=0; h<myGrid.order-1; ++h) {
3979           q[i][myGrid.K3+h] = q[i][h];
3980         }
3981       }
3982     }
3983   }
3984 }
3985
3986 void ComputePme::ungridForces() {
3987
3988     if ( sequence() != myMgr->compute_sequence ) NAMD_bug("ComputePme sequence mismatch in ungridForces()");
3989  
3990     SimParameters *simParams = Node::Object()->simParameters;
3991
3992     localResults_alloc.resize(numLocalAtoms* ((numGrids>1 || selfOn)?2:1));
3993     Vector *localResults = localResults_alloc.begin();
3994     Vector *gridResults;
3995
3996     if ( alchOn || lesOn || selfOn || pairOn ) {
3997       for(int i=0; i<numLocalAtoms; ++i) { localResults[i] = 0.; }
3998       gridResults = localResults + numLocalAtoms;
3999     } else {
4000       gridResults = localResults;
4001     }
4002
4003     Vector pairForce = 0.;
4004     Lattice &lattice = patch->flags.lattice;
4005     int g = 0;
4006     if(!simParams->commOnly) {
4007     for ( g=0; g<numGrids; ++g ) {
4008 #ifdef NETWORK_PROGRESS
4009       CmiNetworkProgress();
4010 #endif
4011
4012 #ifdef NAMD_CUDA
4013       if ( offload ) {
4014         int errfound = 0;
4015         for ( int n=numGridAtoms[g], i=0; i<n; ++i ) {
4016           // Neither isnan() nor x != x worked when testing on Cray; this does.
4017           if ( ((int*)f_data_host)[3*i] == 0x7fffffff ) { errfound = 1; }  // CUDA NaN
4018           gridResults[i].x = f_data_host[3*i];
4019           gridResults[i].y = f_data_host[3*i+1];
4020           gridResults[i].z = f_data_host[3*i+2];
4021         }
4022         if ( errfound ) {
4023           int errcount = 0;
4024           for ( int n=numGridAtoms[g], i=0; i<n; ++i ) {
4025             float f = f_data_host[3*i];
4026             if ( ((int*)f_data_host)[3*i] == 0x7fffffff ) {  // CUDA NaN
4027               ++errcount;
4028               gridResults[i] = 0.;
4029             }
4030           }
4031           iout << iERROR << "Stray PME grid charges detected: "
4032                 << errcount << " atoms on pe " << CkMyPe() << "\n" << endi;
4033         }
4034       } else
4035 #endif // NAMD_CUDA
4036         {
4037           myRealSpace[g]->compute_forces(myMgr->q_arr+g*myMgr->fsize, localGridData[g], gridResults);
4038         }
4039       scale_forces(gridResults, numGridAtoms[g], lattice);
4040       
4041       if (alchOn) {
4042         float scale = 1.;
4043         BigReal elecLambdaUp, elecLambdaDown;
4044         if ( simParams->alchFepWhamOn ) {
4045           if ( simParams->alchFepElecOn ) {
4046             elecLambdaUp = simParams->alchElecLambda;
4047             elecLambdaDown = 1.0 - simParams->alchElecLambda;
4048           }
4049           else {
4050             elecLambdaUp = 0.0;
4051             elecLambdaDown = 1.0;
4052           }
4053         }
4054         else {
4055           BigReal alchLambda = simParams->getCurrentLambda(patch->flags.step);
4056           myMgr->alchLambda = alchLambda;
4057           BigReal alchLambda2 = simParams->getCurrentLambda2(patch->flags.step);
4058           myMgr->alchLambda2 = alchLambda2;
4059           elecLambdaUp = simParams->getElecLambda(alchLambda);
4060           elecLambdaDown = simParams->getElecLambda(1. - alchLambda);
4061         }
4062         
4063         if ( g == 0 ) scale = elecLambdaUp;
4064         else if ( g == 1 ) scale = elecLambdaDown;
4065         else if ( g == 2 ) scale = (elecLambdaUp + elecLambdaDown - 1)*(-1);
4066
4067         if (alchDecouple) {
4068           if ( g == 2 ) scale = 1 - elecLambdaUp;
4069           else if ( g == 3 ) scale = 1 - elecLambdaDown;
4070           else if ( g == 4 ) scale = (elecLambdaUp + elecLambdaDown - 1)*(-1);
4071         }
4072         int nga = 0;
4073         if (!alchDecouple) {
4074           for(int i=0; i<numLocalAtoms; ++i) {
4075             if ( localPartition[i] == 0 || localPartition[i] == (g+1) ) {
4076               // (g=2: only partition 0)
4077               localResults[i] += gridResults[nga++] * scale;
4078             }
4079           }
4080         }
4081         else {  // alchDecouple
4082           if ( g < 2 ) {
4083             for(int i=0; i<numLocalAtoms; ++i) {
4084               if ( localPartition[i] == 0 || localPartition[i] == (g+1) ) {
4085                 // g = 0: partition 0 or partition 1
4086                 // g = 1: partition 0 or partition 2
4087                 localResults[i] += gridResults[nga++] * scale;
4088               }
4089             }
4090           }
4091           else {
4092             for(int i=0; i<numLocalAtoms; ++i) {
4093               if ( localPartition[i] == (g-1) || localPartition[i] == (g-4)) {
4094                 // g = 2: partition 1 only
4095                 // g = 3: partition 2 only
4096                 // g = 4: partition 0 only
4097                 localResults[i] += gridResults[nga++] * scale;
4098               }
4099             }
4100           }
4101         }
4102       } else if ( lesOn ) {
4103         float scale = 1.;
4104         if ( alchFepOn ) {
4105           if(simParams->alchFepWhamOn) {
4106             if(simParams->alchFepElecOn) {
4107               if ( g == 0 ) scale = simParams->alchElecLambda;
4108               else if ( g == 1 ) scale = 1. - simParams->alchElecLambda;
4109             }
4110             else {
4111               if ( g == 0 ) scale = 0.0;
4112               else if ( g == 1 ) scale = 1.0;
4113             }
4114           }
4115           else {
4116             BigReal alchLambda = simParams->getCurrentLambda(patch->flags.step);
4117             myMgr->alchLambda = alchLambda;
4118             BigReal alchLambda2 = simParams->getCurrentLambda2(patch->flags.step);
4119             myMgr->alchLambda2 = alchLambda2;
4120             if ( g == 0 ) scale = alchLambda;
4121             else if ( g == 1 ) scale = 1. - alchLambda;
4122           }
4123         } else if ( lesOn ) {
4124           scale = 1.0 / (float)lesFactor;
4125         }
4126         int nga = 0;
4127         for(int i=0; i<numLocalAtoms; ++i) {
4128           if ( localPartition[i] == 0 || localPartition[i] == (g+1) ) {
4129             localResults[i] += gridResults[nga++] * scale;
4130           }
4131         }
4132       } else if ( selfOn ) {
4133         PmeParticle *lgd = localGridData[g];
4134         int nga = 0;
4135         for(int i=0; i<numLocalAtoms; ++i) {
4136           if ( localPartition[i] == 1 ) {
4137             pairForce += gridResults[nga];  // should add up to almost zero
4138             localResults[i] += gridResults[nga++];
4139           }
4140         }
4141       } else if ( pairOn ) {
4142         if ( g == 0 ) {
4143           int nga = 0;
4144           for(int i=0; i<numLocalAtoms; ++i) {
4145             if ( localPartition[i] == 1 ) {
4146               pairForce += gridResults[nga];
4147             }
4148             if ( localPartition[i] == 1 || localPartition[i] == 2 ) {
4149               localResults[i] += gridResults[nga++];
4150             }
4151           }
4152         } else if ( g == 1 ) {
4153           int nga = 0;
4154           for(int i=0; i<numLocalAtoms; ++i) {
4155             if ( localPartition[i] == g ) {
4156               pairForce -= gridResults[nga];  // should add up to almost zero
4157               localResults[i] -= gridResults[nga++];
4158             }
4159           }
4160         } else {
4161           int nga = 0;
4162           for(int i=0; i<numLocalAtoms; ++i) {
4163             if ( localPartition[i] == g ) {
4164               localResults[i] -= gridResults[nga++];
4165             }
4166          }
4167         }
4168       }
4169     }
4170     }
4171
4172     Vector *results_ptr = localResults;
4173     
4174     // add in forces
4175     {
4176       Results *r = forceBox->open();
4177       Force *f = r->f[Results::slow];
4178       int numAtoms = patch->getNumAtoms();
4179
4180       if ( ! myMgr->strayChargeErrors && ! simParams->commOnly ) {
4181         for(int i=0; i<numAtoms; ++i) {
4182           f[i].x += results_ptr->x;
4183           f[i].y += results_ptr->y;
4184           f[i].z += results_ptr->z;
4185           ++results_ptr;
4186         }
4187       }
4188       forceBox->close(&r);
4189     }
4190
4191     if ( pairOn || selfOn ) {
4192         ADD_VECTOR_OBJECT(myMgr->reduction,REDUCTION_PAIR_ELECT_FORCE,pairForce);
4193     }
4194
4195 }
4196
4197 void ComputePmeMgr::submitReductions() {
4198
4199     SimParameters *simParams = Node::Object()->simParameters;
4200
4201     for ( int g=0; g<numGrids; ++g ) {
4202       float scale = 1.;
4203       if (alchOn) {
4204         BigReal elecLambdaUp, elecLambdaDown;
4205         if( simParams->alchFepWhamOn ) {
4206           if( simParams->alchFepElecOn ) {
4207             elecLambdaUp = simParams->alchElecLambda;
4208             elecLambdaDown = 1.0 - simParams->alchElecLambda;
4209           }
4210           else {
4211             elecLambdaUp = 0.0;
4212             elecLambdaDown = 1.0;
4213           }
4214         }
4215         else {
4216           // alchLambda set on each step in ComputePme::ungridForces()
4217           if ( alchLambda < 0 || alchLambda > 1 ) {
4218             NAMD_bug("ComputePmeMgr::submitReductions alchLambda out of range");
4219           }
4220           elecLambdaUp = simParams->getElecLambda(alchLambda);
4221           elecLambdaDown = simParams->getElecLambda(1-alchLambda);
4222         }
4223         if ( g == 0 ) scale = elecLambdaUp;
4224         else if ( g == 1 ) scale = elecLambdaDown;
4225         else if ( g == 2 ) scale = (elecLambdaUp + elecLambdaDown - 1)*(-1);
4226         if (alchDecouple) {
4227           if ( g == 2 ) scale = 1-elecLambdaUp;
4228           else if ( g == 3 ) scale = 1-elecLambdaDown;
4229           else if ( g == 4 ) scale = (elecLambdaUp + elecLambdaDown - 1)*(-1);
4230         }
4231       } else if ( lesOn ) {
4232         scale = 1.0 / lesFactor;
4233       } else if ( pairOn ) {
4234         scale = ( g == 0 ? 1. : -1. );
4235       }
4236       reduction->item(REDUCTION_ELECT_ENERGY_SLOW) += evir[g][0] * scale;
4237       reduction->item(REDUCTION_VIRIAL_SLOW_XX) += evir[g][1] * scale;
4238       reduction->item(REDUCTION_VIRIAL_SLOW_XY) += evir[g][2] * scale;
4239       reduction->item(REDUCTION_VIRIAL_SLOW_XZ) += evir[g][3] * scale;
4240       reduction->item(REDUCTION_VIRIAL_SLOW_YX) += evir[g][2] * scale;
4241       reduction->item(REDUCTION_VIRIAL_SLOW_YY) += evir[g][4] * scale;
4242       reduction->item(REDUCTION_VIRIAL_SLOW_YZ) += evir[g][5] * scale;
4243       reduction->item(REDUCTION_VIRIAL_SLOW_ZX) += evir[g][3] * scale;
4244       reduction->item(REDUCTION_VIRIAL_SLOW_ZY) += evir[g][5] * scale;
4245       reduction->item(REDUCTION_VIRIAL_SLOW_ZZ) += evir[g][6] * scale;
4246
4247       float scale2 = 0.;
4248
4249       // why is this declared/defined again here?
4250       SimParameters *simParams = Node::Object()->simParameters;
4251
4252       if (alchFepOn) {
4253         BigReal elecLambda2Up=0.0, elecLambda2Down=0.0;
4254         if(simParams->alchFepWhamOn) {
4255           if(simParams->alchFepElecOn) {
4256             elecLambda2Up = simParams->alchElecLambda;
4257             elecLambda2Down =  1.0 - simParams->alchElecLambda;
4258           }
4259           else {
4260             elecLambda2Up = 0.0;
4261             elecLambda2Down =  1.0;
4262           }
4263         }
4264         else {
4265           elecLambda2Up = simParams->getElecLambda(alchLambda2);
4266           elecLambda2Down = simParams->getElecLambda(1.-alchLambda2);
4267         }
4268         
4269         if ( g == 0 ) scale2 = elecLambda2Up;
4270         else if ( g == 1 ) scale2 = elecLambda2Down;
4271         else if ( g == 2 ) scale2 = (elecLambda2Up + elecLambda2Down - 1)*(-1);
4272         if (alchDecouple && g == 2 ) scale2 = 1 - elecLambda2Up;
4273         else if (alchDecouple && g == 3 ) scale2 = 1 - elecLambda2Down;
4274         else if (alchDecouple && g == 4 ) scale2 = (elecLambda2Up + elecLambda2Down - 1)*(-1);
4275       }
4276       if(simParams->alchFepWhamOn && simParams->alchFepElecOn)  {       // FEP with wham post-process
4277         if( g==0 )      scale2 = scale + 1.0;
4278         else if( g==1 ) scale2 = scale - 1.0;
4279         else if( g==2 ) scale2 = scale - 1.0;
4280         else if( g==3 ) scale2 = scale + 1.0;
4281       }
4282       reduction->item(REDUCTION_ELECT_ENERGY_SLOW_F) += evir[g][0] * scale2;
4283       
4284       if (alchThermIntOn) {
4285         
4286         // no decoupling:
4287         // part. 1 <-> all of system except partition 2: g[0] - g[2] 
4288         // (interactions between all atoms [partition 0 OR partition 1], 
4289         // minus all [within partition 0])
4290         // U = elecLambdaUp * (U[0] - U[2])
4291         // dU/dl = U[0] - U[2];
4292         
4293         // part. 2 <-> all of system except partition 1: g[1] - g[2] 
4294         // (interactions between all atoms [partition 0 OR partition 2], 
4295         // minus all [within partition 0])
4296         // U = elecLambdaDown * (U[1] - U[2])
4297         // dU/dl = U[1] - U[2];
4298
4299         // alchDecouple:
4300         // part. 1 <-> part. 0: g[0] - g[2] - g[4] 
4301         // (interactions between all atoms [partition 0 OR partition 1]
4302         // minus all [within partition 1] minus all [within partition 0]
4303         // U = elecLambdaUp * (U[0] - U[4]) + (1-elecLambdaUp)* U[2]
4304         // dU/dl = U[0] - U[2] - U[4];
4305
4306         // part. 2 <-> part. 0: g[1] - g[3] - g[4] 
4307         // (interactions between all atoms [partition 0 OR partition 2]
4308         // minus all [within partition 2] minus all [within partition 0]
4309         // U = elecLambdaDown * (U[1] - U[4]) + (1-elecLambdaDown)* U[3]
4310         // dU/dl = U[1] - U[3] - U[4];
4311         
4312         
4313         if ( g == 0 ) reduction->item(REDUCTION_ELECT_ENERGY_PME_TI_1) += evir[g][0];
4314         if ( g == 1 ) reduction->item(REDUCTION_ELECT_ENERGY_PME_TI_2) += evir[g][0];
4315         if (!alchDecouple) {
4316           if ( g == 2 ) reduction->item(REDUCTION_ELECT_ENERGY_PME_TI_1) -= evir[g][0];
4317           if ( g == 2 ) reduction->item(REDUCTION_ELECT_ENERGY_PME_TI_2) -= evir[g][0];
4318         }
4319         else {  // alchDecouple
4320           if ( g == 2 ) reduction->item(REDUCTION_ELECT_ENERGY_PME_TI_1) -= evir[g][0];
4321           if ( g == 3 ) reduction->item(REDUCTION_ELECT_ENERGY_PME_TI_2) -= evir[g][0];
4322           if ( g == 4 ) reduction->item(REDUCTION_ELECT_ENERGY_PME_TI_1) -= evir[g][0];
4323           if ( g == 4 ) reduction->item(REDUCTION_ELECT_ENERGY_PME_TI_2) -= evir[g][0];
4324         }
4325       }
4326     }
4327
4328     alchLambda = -1.;  // illegal value to catch if not updated
4329
4330     reduction->item(REDUCTION_STRAY_CHARGE_ERRORS) += strayChargeErrors;
4331     reduction->submit();
4332
4333   for ( int i=0; i<heldComputes.size(); ++i ) {
4334     WorkDistrib::messageEnqueueWork(heldComputes[i]);
4335   }
4336   heldComputes.resize(0);
4337 }
4338
4339 #if USE_TOPOMAP 
4340
4341 #define NPRIMES 8
4342 const static unsigned int NAMDPrimes[] = {
4343   3,
4344   5,
4345   7,
4346   11,
4347   13,
4348   17,
4349   19,
4350   23,  
4351   29,
4352   31,
4353   37,
4354   59,
4355   73,
4356   93,
4357   113,
4358   157,
4359   307,
4360   617,
4361   1217                  //This should b enough for 64K nodes of BGL. 
4362 };
4363
4364 #include "RecBisection.h"
4365
4366 /***-----------------------------------------------------**********
4367     The Orthogonal Recursive Bisection strategy, which allocates PME
4368     objects close to the patches they communicate, and at the same
4369     time spreads them around the grid 
4370 ****----------------------------------------------------------****/
4371
4372 bool generateBGLORBPmePeList(int *pemap, int numPes, 
4373                              int *block_pes, int nbpes) {
4374
4375   PatchMap *pmap = PatchMap::Object();
4376   int *pmemap = new int [CkNumPes()];
4377
4378   if (pemap == NULL)
4379     return false;
4380
4381   TopoManager tmgr;
4382
4383   memset(pmemap, 0, sizeof(int) * CkNumPes());
4384
4385   for(int count = 0; count < CkNumPes(); count++) {
4386     if(count < nbpes)
4387       pmemap[block_pes[count]] = 1;
4388     
4389     if(pmap->numPatchesOnNode(count)) {
4390       pmemap[count] = 1;
4391       
4392       //Assumes an XYZT mapping !!
4393       if(tmgr.hasMultipleProcsPerNode()) {
4394         pmemap[(count + CkNumPes()/2)% CkNumPes()] = 1;
4395       }
4396     }
4397   }
4398
4399   if(numPes + nbpes + pmap->numNodesWithPatches() > CkNumPes())
4400     //NAMD_bug("PME ORB Allocator: Processors Unavailable\n");
4401     return false;
4402
4403   CProxy_Node nd(CkpvAccess(BOCclass_group).node);
4404   Node *node = nd.ckLocalBranch();
4405   SimParameters *simParams = node->simParameters;
4406
4407   //first split PME processors into patch groups
4408
4409   int xsize = 0, ysize = 0, zsize = 0;
4410
4411   xsize = tmgr.getDimNX();
4412   ysize = tmgr.getDimNY();
4413   zsize = tmgr.getDimNZ();
4414   
4415   int nx = xsize, ny = ysize, nz = zsize;
4416   DimensionMap dm;
4417   
4418   dm.x = 0;
4419   dm.y = 1;
4420   dm.z = 2;
4421   
4422   findOptimalDimensions(xsize, ysize, zsize, nx, ny, nz, dm);
4423
4424   //group size processors have to be allocated to each YZ plane
4425   int group_size = numPes/nx;
4426   if(numPes % nx)
4427     group_size ++;
4428
4429   int my_prime = NAMDPrimes[0];
4430   int density = (ny * nz)/group_size + 1;