Fix pencil PME hang with oversized periodic cell
[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        }