Use Hilbert curve for reordering water molecules

platforms/cuda/CMakeLists.txt

@@ -73,7 +73,7 @@ SET(SOURCE_FILES) # empty
 SET(SOURCE_INCLUDE_FILES)
 
 FOREACH(subdir ${OPENMM_SOURCE_SUBDIRS})
-    FILE(GLOB src_files  ${CMAKE_CURRENT_SOURCE_DIR}/${subdir}/src/*.cpp ${CMAKE_CURRENT_SOURCE_DIR}/${subdir}/src/*/*.cpp)
+    FILE(GLOB_RECURSE src_files  ${CMAKE_CURRENT_SOURCE_DIR}/${subdir}/src/*.cpp ${CMAKE_CURRENT_SOURCE_DIR}/${subdir}/src/*.c)
     FILE(GLOB incl_files ${CMAKE_CURRENT_SOURCE_DIR}/${subdir}/src/*.h)
     SET(SOURCE_FILES         ${SOURCE_FILES}         ${src_files})   #append
     SET(SOURCE_INCLUDE_FILES ${SOURCE_INCLUDE_FILES} ${incl_files})

platforms/cuda/src/CudaKernels.cpp

@@ -47,7 +47,7 @@ using namespace std;
 
 static void calcForces(OpenMMContextImpl& context, CudaPlatform::PlatformData& data) {
     _gpuContext* gpu = data.gpu;
-    if (data.stepCount%200 == 0)
+    if (data.stepCount%100 == 0)
         gpuReorderAtoms(gpu);
     data.stepCount++;
     kClearForces(gpu);

platforms/cuda/src/kernels/gpu.cpp

@@ -50,6 +50,7 @@ using namespace std;
 
 #include "gputypes.h"
 #include "cudaKernels.h"
+#include "hilbert.h"
 #include "OpenMMException.h"
 
 using OpenMM::OpenMMException;
@@ -2873,7 +2874,7 @@ static void findMoleculeGroups(gpuContext gpu)
         if (atom3 != -1)
             constraints.push_back(Constraint(atom1, atom3, distance2));
         if (atom4 != -1)
-            constraints.push_back(Constraint(atom1, atom3, distance2));
+            constraints.push_back(Constraint(atom1, atom4, distance2));
     }
     for (int i = 0; i < gpu->sim.settleConstraints; i++)
     {
@@ -3082,10 +3083,6 @@ void gpuReorderAtoms(gpuContext gpu)
             maxz = max(maxz, posq[i].z);
         }
     }
-    float binWidth = 0.2*sqrt(gpu->sim.nonbondedCutoffSqr);
-    int xbins = 1 + (int) ((maxx-minx)/binWidth);
-    int ybins = 1 + (int) ((maxy-miny)/binWidth);
-    int zbins = 1 + (int) ((maxz-minz)/binWidth);
 
     // Loop over each group of identical molecules and reorder them.
 
@@ -3130,17 +3127,37 @@ void gpuReorderAtoms(gpuContext gpu)
 
         // Select a bin for each molecule, then sort them by bin.
 
+        bool useHilbert = (numMolecules > 5000); // For small systems, a simple zigzag curve works better than a Hilbert curve.
+        float binWidth;
+        if (useHilbert)
+        binWidth = max(max(maxx-minx, maxy-miny), maxz-minz)/255.0;
+        else
+            binWidth = 0.2*sqrt(gpu->sim.nonbondedCutoffSqr);
+        int xbins = 1 + (int) ((maxx-minx)/binWidth);
+        int ybins = 1 + (int) ((maxy-miny)/binWidth);
         vector<pair<int, int> > molBins(numMolecules);
+        bitmask_t coords[3];
         for (int i = 0; i < numMolecules; i++)
         {
             int x = (int) ((molPos[i].x-minx)/binWidth);
             int y = (int) ((molPos[i].y-miny)/binWidth);
             int z = (int) ((molPos[i].z-minz)/binWidth);
-            int yodd = y&1;
-            int zodd = z&1;
-            int bin = z*xbins*ybins;
-            bin += (zodd ? ybins-y : y)*xbins;
-            bin += (yodd ? xbins-x : x);
+            int bin;
+            if (useHilbert)
+            {
+                coords[0] = x;
+                coords[1] = y;
+                coords[2] = z;
+                bin = (int) hilbert_c2i(3, 8, coords);
+            }
+            else
+            {
+                int yodd = y&1;
+                int zodd = z&1;
+                bin = z*xbins*ybins;
+                bin += (zodd ? ybins-y : y)*xbins;
+                bin += (yodd ? xbins-x : x);
+            }
             molBins[i] = pair<int, int>(bin, i);
         }
         sort(molBins.begin(), molBins.end());

platforms/cuda/src/kernels/hilbert.h

+/* C header file for Hilbert curve functions */
+#if !defined(_hilbert_h_)
+#define _hilbert_h_
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/* define the bitmask_t type as an integer of sufficient size */
+typedef unsigned long long bitmask_t;
+/* define the halfmask_t type as an integer of 1/2 the size of bitmask_t */
+typedef unsigned long halfmask_t;
+
+/*****************************************************************
+ * hilbert_i2c
+ *
+ * Convert an index into a Hilbert curve to a set of coordinates.
+ * Inputs:
+ *  nDims:      Number of coordinate axes.
+ *  nBits:      Number of bits per axis.
+ *  index:      The index, contains nDims*nBits bits (so nDims*nBits must be <= 8*sizeof(bitmask_t)).
+ * Outputs:
+ *  coord:      The list of nDims coordinates, each with nBits bits.
+ * Assumptions:
+ *      nDims*nBits <= (sizeof index) * (bits_per_byte)
+ */
+
+void hilbert_i2c(unsigned nDims, unsigned nBits, bitmask_t index, bitmask_t coord[]);
+
+/*****************************************************************
+ * hilbert_c2i
+ *
+ * Convert coordinates of a point on a Hilbert curve to its index.
+ * Inputs:
+ *  nDims:      Number of coordinates.
+ *  nBits:      Number of bits/coordinate.
+ *  coord:      Array of n nBits-bit coordinates.
+ * Outputs:
+ *  index:      Output index value.  nDims*nBits bits.
+ * Assumptions:
+ *      nDims*nBits <= (sizeof bitmask_t) * (bits_per_byte)
+ */
+
+bitmask_t hilbert_c2i(unsigned nDims, unsigned nBits, bitmask_t const coord[]);
+
+/*****************************************************************
+ * hilbert_cmp, hilbert_ieee_cmp
+ *
+ * Determine which of two points lies further along the Hilbert curve
+ * Inputs:
+ *  nDims:      Number of coordinates.
+ *  nBytes:     Number of bytes of storage/coordinate (hilbert_cmp only)
+ *  nBits:      Number of bits/coordinate. (hilbert_cmp only)
+ *  coord1:     Array of nDims nBytes-byte coordinates (or doubles for ieee_cmp).
+ *  coord2:     Array of nDims nBytes-byte coordinates (or doubles for ieee_cmp).
+ * Return value:
+ *      -1, 0, or 1 according to whether
+           coord1<coord2, coord1==coord2, coord1>coord2
+ * Assumptions:
+ *      nBits <= (sizeof bitmask_t) * (bits_per_byte)
+ */
+
+int hilbert_cmp(unsigned nDims, unsigned nBytes, unsigned nBits, void const* coord1, void const* coord2);
+int hilbert_ieee_cmp(unsigned nDims, double const* coord1, double const* coord2);
+
+/*****************************************************************
+ * hilbert_box_vtx
+ *
+ * Determine the first or last vertex of a box to lie on a Hilbert curve
+ * Inputs:
+ *  nDims:      Number of coordinates.
+ *  nBytes:     Number of bytes/coordinate.
+ *  nBits:      Number of bits/coordinate. (hilbert_cmp only)
+ *  findMin:    Is it the least vertex sought?
+ *  coord1:     Array of nDims nBytes-byte coordinates - one corner of box
+ *  coord2:     Array of nDims nBytes-byte coordinates - opposite corner
+ * Output:
+ *      c1 and c2 modified to refer to selected corner
+ *      value returned is log2 of size of largest power-of-two-aligned box that
+ *      contains the selected corner and no other corners
+ * Assumptions:
+ *      nBits <= (sizeof bitmask_t) * (bits_per_byte)
+ */
+unsigned
+hilbert_box_vtx(unsigned nDims, unsigned nBytes, unsigned nBits,
+		int findMin, void* c1, void* c2);
+unsigned
+hilbert_ieee_box_vtx(unsigned nDims,
+		     int findMin, double* c1, double* c2);
+
+/*****************************************************************
+ * hilbert_box_pt
+ *
+ * Determine the first or last point of a box to lie on a Hilbert curve
+ * Inputs:
+ *  nDims:      Number of coordinates.
+ *  nBytes:     Number of bytes/coordinate.
+ *  nBits:      Number of bits/coordinate.
+ *  findMin:    Is it the least vertex sought?
+ *  coord1:     Array of nDims nBytes-byte coordinates - one corner of box
+ *  coord2:     Array of nDims nBytes-byte coordinates - opposite corner
+ * Output:
+ *      c1 and c2 modified to refer to least point
+ * Assumptions:
+ *      nBits <= (sizeof bitmask_t) * (bits_per_byte)
+ */
+unsigned
+hilbert_box_pt(unsigned nDims, unsigned nBytes, unsigned nBits,
+	       int findMin, void* coord1, void* coord2);
+unsigned
+hilbert_ieee_box_pt(unsigned nDims,
+		    int findMin, double* c1, double* c2);
+
+/*****************************************************************
+ * hilbert_nextinbox
+ *
+ * Determine the first point of a box after a given point to lie on a Hilbert curve
+ * Inputs:
+ *  nDims:      Number of coordinates.
+ *  nBytes:     Number of bytes/coordinate.
+ *  nBits:      Number of bits/coordinate.
+ *  findPrev:   Is the previous point sought?
+ *  coord1:     Array of nDims nBytes-byte coordinates - one corner of box
+ *  coord2:     Array of nDims nBytes-byte coordinates - opposite corner
+ *  point:      Array of nDims nBytes-byte coordinates - lower bound on point returned
+ *
+ * Output:
+      if returns 1:
+ *      c1 and c2 modified to refer to least point after "point" in box
+      else returns 0:
+        arguments unchanged; "point" is beyond the last point of the box
+ * Assumptions:
+ *      nBits <= (sizeof bitmask_t) * (bits_per_byte)
+ */
+int
+hilbert_nextinbox(unsigned nDims, unsigned nBytes, unsigned nBits,
+		  int findPrev, void* coord1, void* coord2,
+		  void const* point);
+
+/*****************************************************************
+ * hilbert_incr
+ *
+ * Advance from one point to its successor on a Hilbert curve
+ * Inputs:
+ *  nDims:      Number of coordinates.
+ *  nBits:      Number of bits/coordinate.
+ *  coord:      Array of nDims nBits-bit coordinates.
+ * Output:
+ *  coord:      Next point on Hilbert curve
+ * Assumptions:
+ *      nBits <= (sizeof bitmask_t) * (bits_per_byte)
+ */
+
+void
+hilbert_incr(unsigned nDims, unsigned nBits, bitmask_t coord[]);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* _hilbert_h_ */