start moving single loop kernels to separate files

Author: zhi-wang

ext/ext/ck3.py

@@ -307,7 +307,7 @@ class Variable:
     def __init__(self, iork:str, defstr:str) -> None:
         vs = defstr.split()
         self.iork = iork
-        self.location = vs[0] # shared / register
+        self.location = vs[0]  # shared / register
         self.type = vs[1]
         self.name = vs[2]
         self.readfrom, self.addto, self.onlyif = None, None, None
@@ -323,7 +323,6 @@ def __init__(self, iork:str, defstr:str) -> None:
                 s3 = s2.replace('onlyif:', '')
                 self.onlyif = s3
 
-
     @staticmethod
     def _get_src(src, index) -> str:
         if ',' in src:
@@ -332,18 +331,16 @@ def _get_src(src, index) -> str:
         else:
             return '{}[{}]'.format(src, index)
 
-
     def zero(self) -> str:
         v1 = ''
         if self.location == 'shared':
             v1 = '{}[threadIdx.x] = 0;'.format(self.name)
         else:
             v1 = '{} = 0;'.format(self.name)
-        if self.onlyif != None:
-            v1 =  'if CONSTEXPR ({}) {}'.format(self.onlyif, v1)
+        if self.onlyif is not None:
+            v1 = 'if CONSTEXPR ({}) {}'.format(self.onlyif, v1)
         return v1
 
-
     def save(self) -> str:
         dst = self.addto
         v1, suffix = '', ''
@@ -354,11 +351,10 @@ def save(self) -> str:
         else:
             vs = dst.split(',')
             v1 = 'atomic_add({}{}, &{}[{}][{}]);'.format(self.name, suffix, vs[0], self.iork, vs[1])
-        if self.onlyif != None:
+        if self.onlyif is not None:
             v1 = 'if CONSTEXPR ({}) {}'.format(self.onlyif, v1)
         return v1
 
-
     def init_exclude(self) -> str:
         rhs = self._get_src(self.readfrom, self.iork)
         if self.location == 'shared':
@@ -369,7 +365,6 @@ def init_exclude(self) -> str:
         else:
             return '{} = {};'.format(self.name, rhs)
 
-
     def init_block(self) -> str:
         if self.readfrom in ['x', 'y', 'z']:
             if self.location == 'shared':
@@ -382,14 +377,12 @@ def init_block(self) -> str:
             else:
                 return '{} = {};'.format(self.name, self._get_src(self.readfrom, self.iork))
 
-
     def shuffle(self) -> str:
         if self.location == 'register':
             return '{0:} = __shfl_sync(ALL_LANES, {0:}, ilane + 1);'.format(self.name)
         else:
             raise ValueError('Cannot shuffle variables in the shared memory.')
 
-
     def ikreplace(self, code:str) -> str:
         old_name = '@{}@'.format(self.name)
         new_name = self.name
@@ -402,7 +395,6 @@ def ikreplace(self, code:str) -> str:
         code = code.replace(old_name, new_name)
         return code
 
-
     def iterreplace(self, code:str) -> str:
         old_name = '@{}@'.format('i')
         new_name = self.name
@@ -425,7 +417,6 @@ def __init__(self, iork:str, lst:list) -> None:
             else:
                 d[v.type] = [v]
 
-
     def declare(self) -> str:
         s = ''
         for t in self.shared.keys():
@@ -441,7 +432,6 @@ def declare(self) -> str:
         s = s.replace(',;', ';')
         return s
 
-
     def zero(self) -> str:
         s = ''
         for t in self.shared.keys():
@@ -452,7 +442,6 @@ def zero(self) -> str:
                 s = s + v.zero()
         return s
 
-
     def save(self) -> str:
         s = ''
         for t in self.shared.keys():
@@ -463,7 +452,6 @@ def save(self) -> str:
                 s = s + v.save()
         return s
 
-
     def init_exclude(self) -> str:
         s = ''
         for t in self.shared.keys():
@@ -474,7 +462,6 @@ def init_exclude(self) -> str:
                 s = s + v.init_exclude()
         return s
 
-
     def init_block(self) -> str:
         s = ''
         for t in self.shared.keys():
@@ -485,15 +472,13 @@ def init_block(self) -> str:
                 s = s + v.init_block()
         return s
 
-
     def shuffle(self) -> str:
         s = ''
         for t in self.register.keys():
             for v in self.register[t]:
                 s = s + v.shuffle()
         return s
 
-
     def ikreplace(self, code:str) -> str:
         for t in self.shared.keys():
             for v in self.shared[t]:
@@ -531,7 +516,6 @@ def _func_param(ptype:str, pname:str) -> str:
         else:
             raise ValueError('Do not know how to parse type: {}'.format(ptype))
 
-
     @staticmethod
     def _load_scale_param(ptype:str, stem:str, input:str, separate_scaled_pairwise:bool) -> str:
         if ptype == 'real_const_array':
@@ -551,7 +535,7 @@ def _load_scale_param(ptype:str, stem:str, input:str, separate_scaled_pairwise:b
                 # dim = match.group(2)
                 ss = ptype.split(',')
                 v = ''
-                for i in range(1,len(ss)):
+                for i in range(1, len(ss)):
                     idx = ss[i]
                     al = rc_alphabets[idx]
                     if input is None:
@@ -561,7 +545,6 @@ def _load_scale_param(ptype:str, stem:str, input:str, separate_scaled_pairwise:b
                         v = v + '{} {}{} = {}[ii][{}];'.format(t, stem, al, input, idx)
                 return v
 
-
     def __init__(self, config) -> None:
         self.config = config
 
@@ -598,7 +581,6 @@ def _kv(self, k:str):
         else:
             return ''
 
-
     def cudaReplaceDict(self) -> dict:
         d = {}
         config = self.config
@@ -679,10 +661,10 @@ def cudaReplaceDict(self) -> dict:
         if kcfg in keys:
             vcfg, decl, zero, total = config[kcfg], '', '', ''
             for t in vcfg:
-               v1 = v1 + ', CountBuffer restrict {}'.format(t)
-               decl = decl + 'int {}tl;'.format(t)
-               zero = zero + '{}tl = 0;'.format(t)
-               total = total + 'atomic_add({}tl, {}, ithread);'.format(t, t)
+                v1 = v1 + ', CountBuffer restrict {}'.format(t)
+                decl = decl + 'int {}tl;'.format(t)
+                zero = zero + '{}tl = 0;'.format(t)
+                total = total + 'atomic_add({}tl, {}, ithread);'.format(t, t)
             v2 = '%s if CONSTEXPR (do_a) {%s}' % (decl, zero)
             v3 = 'if CONSTEXPR (do_a) {%s}' % (total)
         d[k1], d[k2], d[k3] = v1, v2, v3
@@ -809,7 +791,7 @@ def cudaReplaceDict(self) -> dict:
             v1 = kvars.ikreplace(v1)
             v1 = ifrcs.ikreplace(v1)
             v1 = kfrcs.ikreplace(v1)
-            v2 = v1 # in case no scaled pairwise interaction is given
+            v2 = v1  # in case no scaled pairwise interaction is given
         kcfg = self.yk_scaled_pairwise
         if kcfg in keys:
             v2 = config[kcfg]
@@ -840,27 +822,24 @@ def cudaReplaceDict(self) -> dict:
 
         return d
 
-
     @staticmethod
     def version() -> str:
         return '3.1.0'
 
-
     @staticmethod
     def _replace(s:str, d:dict) -> str:
         output = s
         for k in d.keys():
             v = d[k]
-            if v == None:
+            if v is None:
                 v = ''
             output = output.replace(k, v)
         return output
 
-
     def write(self, output) -> None:
         d = self.cudaReplaceDict()
         outstr = '// ck.py Version {}'.format(self.version())
-        kernel_num = 21 # default
+        kernel_num = 21  # default
         if self.yk_kernel_version_number in self.config.keys():
             kernel_num = self.config[self.yk_kernel_version_number]
         if kernel_num == 11:
@@ -896,7 +875,6 @@ def show_command(argv):
     d2 = os.path.join(d, '../..')
     d = os.path.abspath(d2)
 
-
     yaml_file = argv[1]
     with open(yaml_file) as input_file:
         config = yaml.full_load(input_file)

ext/ext/y3/mdPos_cu1.yaml renamed to ext/ext/y31/mdPos_cu1.yaml

@@ -1,6 +1,7 @@
 KERNEL_VERSION_NUMBER: 11
-
 KERNEL_IS_STATIC: True
+
+OUTPUT_DIR: src/cu
 KERNEL_NAME: mdPos_cu1
 SINGLE_LOOP_LIMIT: int n
 SINGLE_LOOP_ITER: int i

src/cu/mdPos_cu1.cc

@@ -0,0 +1,11 @@
+// ck.py Version 3.1.0
+__global__
+static void mdPos_cu1(int n, time_prec dt, pos_prec* restrict qx, pos_prec* restrict qy, pos_prec* restrict qz,
+   const vel_prec* restrict vlx, const vel_prec* restrict vly, const vel_prec* restrict vlz)
+{
+   for (int i = ITHREAD; i < n; i += STRIDE) {
+      qx[i] += dt * vlx[i];
+      qy[i] += dt * vly[i];
+      qz[i] += dt * vlz[i];
+   }
+}

src/cu/mdpq.cu

@@ -4,29 +4,18 @@
 #include <tinker/detail/units.hh>
 
 namespace tinker {
-__global__
-static void mdPos_cu1(int n, time_prec dt,                              //
-   pos_prec* restrict qx, pos_prec* restrict qy, pos_prec* restrict qz, //
-   const vel_prec* restrict vlx, const vel_prec* restrict vly, const vel_prec* restrict vlz)
-{
-   for (int i = ITHREAD; i < n; i += STRIDE) {
-      qx[i] += dt * vlx[i];
-      qy[i] += dt * vly[i];
-      qz[i] += dt * vlz[i];
-   }
-}
-
-void mdPos_cu(time_prec dt, pos_prec* qx, pos_prec* qy, pos_prec* qz, const vel_prec* vlx,
-   const vel_prec* vly, const vel_prec* vlz)
+#include "mdPos_cu1.cc"
+void mdPos_cu(time_prec dt, pos_prec* qx, pos_prec* qy, pos_prec* qz, const vel_prec* vlx, const vel_prec* vly,
+   const vel_prec* vlz)
 {
    launch_k1s(g::s0, n, mdPos_cu1, //
       n, dt, qx, qy, qz, vlx, vly, vlz);
 }
 
 __global__
 static void mdPosAxbv_cu1(int n, pos_prec sa, pos_prec sb, //
-   pos_prec* restrict xpos, pos_prec* restrict ypos, pos_prec* restrict zpos,
-   const vel_prec* restrict vx, const vel_prec* restrict vy, const vel_prec* restrict vz)
+   pos_prec* restrict xpos, pos_prec* restrict ypos, pos_prec* restrict zpos, const vel_prec* restrict vx,
+   const vel_prec* restrict vy, const vel_prec* restrict vz)
 {
    for (int i = ITHREAD; i < n; i += STRIDE) {
       xpos[i] = sa * xpos[i] + sb * vx[i];
@@ -45,8 +34,8 @@ __global__
 static void mdPosAxbvAn_cu1(int n,     //
    double3 a0, double3 a1, double3 a2, //
    double3 b0, double3 b1, double3 b2, //
-   pos_prec* restrict xpos, pos_prec* restrict ypos, pos_prec* restrict zpos,
-   const vel_prec* restrict vx, const vel_prec* restrict vy, const vel_prec* restrict vz)
+   pos_prec* restrict xpos, pos_prec* restrict ypos, pos_prec* restrict zpos, const vel_prec* restrict vx,
+   const vel_prec* restrict vy, const vel_prec* restrict vz)
 {
    pos_prec a00 = a0.x, a01 = a0.y, a02 = a0.z;
    pos_prec a10 = a1.x, a11 = a1.y, a12 = a1.z;
@@ -94,20 +83,19 @@ static void mdVel_cu1(int n, time_prec dt, const double* restrict massinv, //
    }
 }
 
-void mdVel_cu(time_prec dt, vel_prec* vlx, vel_prec* vly, vel_prec* vlz, const grad_prec* grx,
-   const grad_prec* gry, const grad_prec* grz)
+void mdVel_cu(time_prec dt, vel_prec* vlx, vel_prec* vly, vel_prec* vlz, const grad_prec* grx, const grad_prec* gry,
+   const grad_prec* grz)
 {
    launch_k1s(g::s0, n, mdVel_cu1, //
       n, dt, massinv, vlx, vly, vlz, grx, gry, grz);
 }
 
 __global__
-static void mdVel2_cu1(int n, const double* restrict massinv, vel_prec* restrict vlx,
-   vel_prec* restrict vly, vel_prec* restrict vlz, //
+static void mdVel2_cu1(int n, const double* restrict massinv, vel_prec* restrict vlx, vel_prec* restrict vly,
+   vel_prec* restrict vlz, //
    time_prec dt1, const grad_prec* restrict grx1, const grad_prec* restrict gry1,
    const grad_prec* restrict grz1, //
-   time_prec dt2, const grad_prec* restrict grx2, const grad_prec* restrict gry2,
-   const grad_prec* restrict grz2)
+   time_prec dt2, const grad_prec* restrict grx2, const grad_prec* restrict gry2, const grad_prec* restrict grz2)
 {
    const vel_prec ekcal = units::ekcal;
    for (int i = ITHREAD; i < n; i += STRIDE) {
@@ -124,8 +112,8 @@ static void mdVel2_cu1(int n, const double* restrict massinv, vel_prec* restrict
    }
 }
 
-void mdVel2_cu(time_prec dt, const grad_prec* grx, const grad_prec* gry, const grad_prec* grz,
-   time_prec dt2, const grad_prec* grx2, const grad_prec* gry2, const grad_prec* grz2)
+void mdVel2_cu(time_prec dt, const grad_prec* grx, const grad_prec* gry, const grad_prec* grz, time_prec dt2,
+   const grad_prec* grx2, const grad_prec* gry2, const grad_prec* grz2)
 {
    launch_k1s(g::s0, n, mdVel2_cu1, //
       n, massinv, vx, vy, vz,       //