From 9fbb1ce57215d5fb9e028c449774813e50d37d10 Mon Sep 17 00:00:00 2001
From: Ansgar Esztermann <aeszter@gwdg.de>
Date: Wed, 28 Apr 2010 16:36:41 +0200
Subject: [PATCH] Implement a Morton curve in gather_f_bsplines() to improve cache locality

In gather_f_bsplines(), copy the original row-major FFT grid to a different
array organized along a space filling (fractal) curve. This improves locality,
thereby reducing the number of L2 cache misses.
---
 src/mdlib/pme.c |   94 ++++++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 files changed, 93 insertions(+), 1 deletions(-)

diff --git a/src/mdlib/pme.c b/src/mdlib/pme.c
index cd7aca4..5e6009c 100644
--- a/src/mdlib/pme.c
+++ b/src/mdlib/pme.c
@@ -97,6 +97,8 @@
 /* #define TAKETIME (step > 1 && timesteps < 10) */
 #define TAKETIME FALSE
 
+#define PME_GRID_MORTON
+
 #ifdef GMX_DOUBLE
 #define mpi_type MPI_DOUBLE
 #else
@@ -174,6 +176,10 @@ typedef struct gmx_pme {
     t_fftgrid *gridA,*gridB;
     int  *nnx,*nny,*nnz;
     
+#ifdef PME_GRID_MORTON
+    real *morton;
+#endif
+
     pme_atomcomm_t atc[2];
     matrix    recipbox;
     splinevec bsp_mod;
@@ -1157,6 +1163,31 @@ real solve_pme(gmx_pme_t pme,t_fftgrid *grid,
     return(0.5*energy);
 }
 
+#ifdef PME_GRID_MORTON
+#define DO_FSPLINE(order)                      \
+for(ithx=0; (ithx<order); ithx++)              \
+{           								   \
+    tx      = thx[ithx];                       \
+    dx      = dthx[ithx];                      \
+											   \
+    for(ithy=0; (ithy<order); ithy++)          \
+	{										   \
+        ty       = thy[ithy];                  \
+        dy       = dthy[ithy];                 \
+        fxy1     = fz1 = 0;                    \
+                                               \
+        for(ithz=0; (ithz<order); ithz++)      \
+		{     								   \
+		gval = morton[morton_index(i0[ithx], j0[ithy], k0[ithz])];   \
+		fxy1 += thz[ithz]*gval;                \
+		fz1  += dthz[ithz]*gval;               \
+		}                                      \
+		fx += dx*ty*fxy1;                      \
+		fy += tx*dy*fxy1;                      \
+		fz += tx*ty*fz1;                       \
+	}                                          \
+}
+#else
 #define DO_FSPLINE(order)                      \
 for(ithx=0; (ithx<order); ithx++)              \
 {           								   \
@@ -1182,7 +1213,36 @@ for(ithx=0; (ithx<order); ithx++)              \
 		fz += tx*ty*fz1;                       \
 	}                                          \
 }
+#endif
+
+
+#ifdef PME_GRID_MORTON
+static int morton_index(int x, int y, int z)
+{
+    static const unsigned int s[] = {1, 3, 9};
+    static const unsigned int m[] = {011111111111, 07007007007, 037000000777};
+    x = (x | (x << s[2])) & m[2];
+    x = (x | (x << s[1])) & m[1];
+    x = (x | (x << s[0])) & m[0];
+    y = (y | (y << s[2])) & m[2];
+    y = (y | (y << s[1])) & m[1];
+    y = (y | (y << s[0])) & m[0];
+    z = (z | (z << s[2])) & m[2];
+    z = (z | (z << s[1])) & m[1];
+    z = (z | (z << s[0])) & m[0];
+
+    return x | (y << 1) | (z << 2);
+}
 
+void to_morton (real *morton, real *rows, int nx, int ny, int nz, int la12, int la2)
+{
+    int x, y, z;
+    for (x=0; x<nx; ++x)
+        for (y=0; y<ny; ++y)
+            for (z=0; z<nz; ++z)
+                morton[morton_index(x, y, z)] = rows[INDEX(x, y, z)];
+}
+#endif
 
 void gather_f_bsplines(gmx_pme_t pme,t_fftgrid *grid,
                        bool bClearF,pme_atomcomm_t *atc,real scale)
@@ -1198,8 +1258,18 @@ void gather_f_bsplines(gmx_pme_t pme,t_fftgrid *grid,
     int     norder,*idxptr;
     real    rxx,ryx,ryy,rzx,rzy,rzz;
     int     order;
+
+#ifdef PME_GRID_MORTON
+    real   *morton;
+#endif
     
     unpack_fftgrid(grid,&nx,&ny,&nz,&nx2,&ny2,&nz2,&la2,&la12,TRUE,&ptr);
+
+#ifdef PME_GRID_MORTON
+    morton = pme->morton;
+    to_morton(morton, ptr, nx, ny, nz, la12, la2);
+#endif
+
     order = pme->pme_order;
     thx   = atc->theta[XX];
     thy   = atc->theta[YY];
@@ -1678,11 +1748,17 @@ int gmx_pme_init(gmx_pme_t *pmedata,t_commrec *cr,int nnodes_major,
     pme_atomcomm_t *atc;
     int nminor,b,d,i,lbnd,rbnd,maxlr;
 
+#ifdef PME_GRID_MORTON
+    int morton_size;
+#endif
     
     if (debug)
         fprintf(debug,"Creating PME data structures.\n");
     snew(pme,1);
-    
+
+#ifdef PME_GRID_MORTON
+    pme->morton = NULL;
+#endif
    
     pme->redist_init   = FALSE;
     pme->sum_qgrid_tmp = NULL;
@@ -1807,6 +1883,22 @@ int gmx_pme_init(gmx_pme_t *pmedata,t_commrec *cr,int nnodes_major,
     snew(pme->bsp_mod[YY],pme->nky);
     snew(pme->bsp_mod[ZZ],pme->nkz);
     
+#ifdef PME_GRID_MORTON
+    morton_size = max(pme->nkx, pme->nky);
+    morton_size = max(morton_size, pme->nkz);
+
+    /* round up to next power of two */
+    morton_size--;
+    morton_size |= morton_size >> 1;
+    morton_size |= morton_size >> 2;
+    morton_size |= morton_size >> 4;
+    morton_size |= morton_size >> 8;
+    morton_size |= morton_size >> 16;
+    morton_size++;
+
+    snew(pme->morton, morton_size * morton_size * morton_size);
+#endif
+
     pme->gridA = mk_fftgrid(pme->nkx,pme->nky,pme->nkz,
                             NULL,pme->overlap[0].s2g,cr,
                             bReproducible);
-- 
1.6.3.3