expol: simplify PCG codes

Author: zhi-wang

include/ff/cuinduce.h

@@ -0,0 +1,30 @@
+#pragma once
+#include "ff/precision.h"
+
+namespace tinker {
+// udir = polarity * field
+
+__global__
+void pcgUdirV1(int n, const real* polarity, //
+   real (*udir)[3], const real (*field)[3]);
+
+__global__
+void pcgUdirV2(int n, const real* polarity, //
+   real (*udir)[3], real (*udirp)[3], const real (*field)[3], const real (*fieldp)[3]);
+
+// r(0) = E - (1/polarity + Tu) u(0) = (udir - u(0))/polarity + mutual field
+
+__global__
+void pcgRsd0V1(int n, const real* polarity_inv, real (*rsd)[3], //
+   const real (*udir)[3], const real (*uind)[3], const real (*field)[3]);
+
+__global__
+void pcgRsd0V2(int n, const real* polarity_inv, real (*rsd)[3], real (*rsp)[3], //
+   const real (*udir)[3], const real (*udip)[3], const real (*uind)[3], const real (*uinp)[3],
+   const real (*field)[3], const real (*fielp)[3]);
+
+__global__
+void pcgRsd0V3(int n, const real* polarity_inv, real (*rsd)[3], //
+   const real (*udir)[3], const real (*uind)[3], const real (*field)[3],
+   const real (*polscale)[3][3]);
+}

src/acc/hippo/expol.cpp

@@ -142,8 +142,8 @@ void alterpol_acc(real (*polscale)[3][3], real (*polinv)[3][3])
    }
 }
 
-void dexpol_acc(const int vers, const real (*uind)[3], grad_prec* depx, grad_prec* depy,
-   grad_prec* depz, VirialBuffer restrict vir_ep)
+void dexpol_acc(int vers, const real (*uind)[3], grad_prec* depx, grad_prec* depy, grad_prec* depz,
+   VirialBuffer vir_ep)
 {
    auto do_v = vers & calc::virial;
    real cut = switchCut(Switch::REPULS);

src/cu/amoeba/pcg.cu

@@ -1,5 +1,6 @@
 #include "ff/amoeba/induce.h"
 #include "ff/amoebamod.h"
+#include "ff/cuinduce.h"
 #include "ff/switch.h"
 #include "seq/launch.h"
 #include "tool/error.h"
@@ -10,36 +11,6 @@
 #include <tinker/detail/units.hh>
 
 namespace tinker {
-__global__
-void pcgUdir(int n, const real* restrict polarity, real (*restrict udir)[3],
-   real (*restrict udirp)[3], const real (*restrict field)[3], const real (*restrict fieldp)[3])
-{
-   for (int i = ITHREAD; i < n; i += STRIDE) {
-      real poli = polarity[i];
-      #pragma unroll
-      for (int j = 0; j < 3; ++j) {
-         udir[i][j] = poli * field[i][j];
-         udirp[i][j] = poli * fieldp[i][j];
-      }
-   }
-}
-
-__global__
-void pcgRsd(int n, const real* restrict polarity_inv, //
-   real (*restrict rsd)[3], real (*restrict rsp)[3],  //
-   const real (*restrict udir)[3], const real (*restrict udip)[3], const real (*restrict uind)[3],
-   const real (*restrict uinp)[3], const real (*restrict field)[3], const real (*restrict fielp)[3])
-{
-   for (int i = ITHREAD; i < n; i += STRIDE) {
-      real poli_inv = polarity_inv[i];
-      #pragma unroll
-      for (int j = 0; j < 3; ++j) {
-         rsd[i][j] = (udir[i][j] - uind[i][j]) * poli_inv + field[i][j];
-         rsp[i][j] = (udip[i][j] - uinp[i][j]) * poli_inv + fielp[i][j];
-      }
-   }
-}
-
 __global__
 void pcgRsd0(
    int n, const real* restrict polarity, real (*restrict rsd)[3], real (*restrict rsdp)[3])
@@ -81,10 +52,8 @@ void pcgP2(int n, const real* restrict polarity,      //
 {
    real kaval = *ka, kapval = *kap;
    real a = *ksum / kaval, ap = *ksump / kapval;
-   if (kaval == 0)
-      a = 0;
-   if (kapval == 0)
-      ap = 0;
+   if (kaval == 0) a = 0;
+   if (kapval == 0) ap = 0;
    for (int i = ITHREAD; i < n; i += STRIDE) {
       #pragma unroll
       for (int j = 0; j < 3; ++j) {
@@ -111,10 +80,8 @@ void pcgP3(int n, const real* restrict ksum, const real* restrict ksump, const r
 {
    real kaval = *ksum, kapval = *ksump;
    real b = *ksum1 / kaval, bp = *ksump1 / kapval;
-   if (kaval == 0)
-      b = 0;
-   if (kapval == 0)
-      bp = 0;
+   if (kaval == 0) b = 0;
+   if (kapval == 0) bp = 0;
    for (int i = ITHREAD; i < n; i += STRIDE) {
       #pragma unroll
       for (int j = 0; j < 3; ++j) {
@@ -162,7 +129,7 @@ void induceMutualPcg1_cu(real (*uind)[3], real (*uinp)[3])
    // get the electrostatic field due to permanent multipoles
    dfield(field, fieldp);
    // direct induced dipoles
-   launch_k1s(g::s0, n, pcgUdir, n, polarity, udir, udirp, field, fieldp);
+   launch_k1s(g::s0, n, pcgUdirV2, n, polarity, udir, udirp, field, fieldp);
 
    // initial induced dipole
    if (predict) {
@@ -188,7 +155,7 @@ void induceMutualPcg1_cu(real (*uind)[3], real (*uinp)[3])
    if (predict) {
       ufield(uind, uinp, field, fieldp);
       launch_k1s(
-         g::s0, n, pcgRsd, n, polarity_inv, rsd, rsdp, udir, udirp, uind, uinp, field, fieldp);
+         g::s0, n, pcgRsd0V2, n, polarity_inv, rsd, rsdp, udir, udirp, uind, uinp, field, fieldp);
    } else if (dirguess) {
       ufield(udir, udirp, rsd, rsdp);
    } else {
@@ -285,17 +252,12 @@ void induceMutualPcg1_cu(real (*uind)[3], real (*uinp)[3])
          print(stdout, " %8d       %-16.10f\n", iter, eps);
       }
 
-      if (eps < poleps)
-         done = true;
-      if (eps > epsold)
-         done = true;
-      if (iter >= politer)
-         done = true;
+      if (eps < poleps) done = true;
+      if (eps > epsold) done = true;
+      if (iter >= politer) done = true;
 
       // apply a "peek" iteration to the mutual induced dipoles
-      if (done) {
-         launch_k1s(g::s0, n, pcgPeek, n, pcgpeek, polarity, uind, uinp, rsd, rsdp);
-      }
+      if (done) launch_k1s(g::s0, n, pcgPeek, n, pcgpeek, polarity, uind, uinp, rsd, rsdp);
    }
 
    // print the results from the conjugate gradient iteration

src/cu/aplus/pcg.cu

@@ -1,6 +1,7 @@
 #include "ff/amoeba/induce.h"
 #include "ff/amoebamod.h"
 #include "ff/aplus/induce.h"
+#include "ff/cuinduce.h"
 #include "ff/switch.h"
 #include "seq/launch.h"
 #include "tool/error.h"
@@ -11,32 +12,6 @@
 #include <tinker/detail/units.hh>
 
 namespace tinker {
-__global__
-void pcgUdirAplus(
-   int n, const real* restrict polarity, real (*restrict udir)[3], const real (*restrict field)[3])
-{
-   for (int i = ITHREAD; i < n; i += STRIDE) {
-      real poli = polarity[i];
-      #pragma unroll
-      for (int j = 0; j < 3; ++j) {
-         udir[i][j] = poli * field[i][j];
-      }
-   }
-}
-
-__global__
-void pcgRsd4(int n, const real* restrict polarity_inv, //
-   real (*restrict rsd)[3],                            //
-   const real (*restrict udir)[3], const real (*restrict uind)[3], const real (*restrict field)[3])
-{
-   for (int i = ITHREAD; i < n; i += STRIDE) {
-      real poli_inv = polarity_inv[i];
-      #pragma unroll
-      for (int j = 0; j < 3; ++j)
-         rsd[i][j] = (udir[i][j] - uind[i][j]) * poli_inv + field[i][j];
-   }
-}
-
 __global__
 void pcgRsd3(int n, const real* restrict polarity, real (*restrict rsd)[3])
 {
@@ -69,8 +44,7 @@ void pcgP8(int n, const real* restrict polarity, //
 {
    real kaval = *ka;
    real a = *ksum / kaval;
-   if (kaval == 0)
-      a = 0;
+   if (kaval == 0) a = 0;
    for (int i = ITHREAD; i < n; i += STRIDE) {
       #pragma unroll
       for (int j = 0; j < 3; ++j) {
@@ -91,8 +65,7 @@ void pcgP9(int n, const real* restrict ksum, const real* restrict ksum1, real (*
 {
    real ksumval = *ksum;
    real b = *ksum1 / ksumval;
-   if (ksumval == 0)
-      b = 0;
+   if (ksumval == 0) b = 0;
    for (int i = ITHREAD; i < n; i += STRIDE) {
       #pragma unroll
       for (int j = 0; j < 3; ++j)
@@ -131,7 +104,7 @@ void induceMutualPcg3_cu(real (*uind)[3])
    // get the electrostatic field due to permanent multipoles
    dfieldAplus(field);
    // direct induced dipoles
-   launch_k1s(g::s0, n, pcgUdirAplus, n, polarity, udir, field);
+   launch_k1s(g::s0, n, pcgUdirV1, n, polarity, udir, field);
 
    // initial induced dipole
    if (predict) {
@@ -155,7 +128,7 @@ void induceMutualPcg3_cu(real (*uind)[3])
    // if do not use pcgguess, r(0) = E - T Zero = E
    if (predict) {
       ufieldAplus(uind, field);
-      launch_k1s(g::s0, n, pcgRsd4, n, polarity_inv, rsd, udir, uind, field);
+      launch_k1s(g::s0, n, pcgRsd0V1, n, polarity_inv, rsd, udir, uind, field);
    } else if (dirguess) {
       ufieldAplus(udir, rsd);
    } else {
@@ -241,16 +214,12 @@ void induceMutualPcg3_cu(real (*uind)[3])
          print(stdout, " %8d       %-16.10f\n", iter, eps);
       }
 
-      if (eps < poleps)
-         done = true;
-      if (eps > epsold)
-         done = true;
-      if (iter >= politer)
-         done = true;
+      if (eps < poleps) done = true;
+      if (eps > epsold) done = true;
+      if (iter >= politer) done = true;
 
       // apply a "peek" iteration to the mutual induced dipoles
-      if (done)
-         launch_k1s(g::s0, n, pcgPeek2, n, pcgpeek, polarity, uind, rsd);
+      if (done) launch_k1s(g::s0, n, pcgPeek2, n, pcgpeek, polarity, uind, rsd);
    }
 
    // print the results from the conjugate gradient iteration

src/cu/cmakesrc.txt

@@ -33,6 +33,7 @@ hippo/expol.cu
 hippo/field.cu
 hippo/pcg.cu
 hippo/precond.cu
+induce.cu
 mathparallel.cu
 mathzero.cu
 mdhc.cu

src/cu/hippo/expol.cu

@@ -735,7 +735,7 @@ void dexpol_cu1(int n, TINKER_IMAGE_PARAMS, VirialBuffer restrict vep, grad_prec
 }
 
 void dexpol_cu(int vers, const real (*uind)[3], grad_prec* depx, grad_prec* depy, grad_prec* depz,
-   VirialBuffer restrict vir_ep)
+   VirialBuffer vir_ep)
 {
    const auto& st = *mspatial_v2_unit;
    real cut = switchCut(Switch::REPULS);
@@ -766,6 +766,7 @@ void dexpol_cu(int vers, const real (*uind)[3], grad_prec* depx, grad_prec* depy
 #include "ff/amoeba/induce.h"
 #include "ff/amoebamod.h"
 #include "ff/atom.h"
+#include "ff/cuinduce.h"
 #include "ff/hippo/induce.h"
 #include "ff/hippomod.h"
 #include "ff/switch.h"
@@ -780,20 +781,7 @@ void dexpol_cu(int vers, const real (*uind)[3], grad_prec* depx, grad_prec* depy
 
 namespace tinker {
 __global__
-void eppcgUdirDonly(
-   int n, const real* restrict polarity, real (*restrict udir)[3], const real (*restrict field)[3])
-{
-   for (int i = ITHREAD; i < n; i += STRIDE) {
-      real poli = polarity[i];
-      #pragma unroll
-      for (int j = 0; j < 3; ++j) {
-         udir[i][j] = poli * field[i][j];
-      }
-   }
-}
-
-__global__
-void eppcgUdirGuess(int n, const real* restrict polarity, real (*restrict uind)[3],
+static void eppcgUdirGuess(int n, const real* restrict polarity, real (*restrict uind)[3],
    const real (*restrict field)[3], const real (*restrict polinv)[3][3])
 {
    for (int i = ITHREAD; i < n; i += STRIDE) {
@@ -807,24 +795,6 @@ void eppcgUdirGuess(int n, const real* restrict polarity, real (*restrict uind)[
    }
 }
 
-__global__
-void eppcgRsd2(int n, const real* restrict polarity_inv, //
-   real (*restrict rsd)[3],                              //
-   const real (*restrict udir)[3], const real (*restrict uind)[3], const real (*restrict field)[3],
-   const real (*restrict polscale)[3][3])
-{
-   for (int i = ITHREAD; i < n; i += STRIDE) {
-      real poli_inv = polarity_inv[i];
-      #pragma unroll
-      for (int j = 0; j < 3; ++j) {
-         rsd[i][j] = (udir[i][j] - uind[i][0] * polscale[i][0][j] - uind[i][1] * polscale[i][1][j] -
-                        uind[i][2] * polscale[i][2][j]) *
-               poli_inv +
-            field[i][j];
-      }
-   }
-}
-
 __global__
 void eppcgRsd1(int n, const real* restrict polarity, real (*restrict rsd)[3])
 {
@@ -921,7 +891,7 @@ void induceMutualPcg4_cu(real (*uind)[3])
    // get the electrostatic field due to permanent multipoles
    dfieldChgpen(field);
    // direct induced dipoles
-   launch_k1s(g::s0, n, eppcgUdirDonly, n, polarity, udir, field);
+   launch_k1s(g::s0, n, pcgUdirV1, n, polarity, udir, field);
 
    alterpol(polscale, polinv);
 
@@ -936,7 +906,7 @@ void induceMutualPcg4_cu(real (*uind)[3])
 
    if (predict) {
       ufieldChgpen(uind, field);
-      launch_k1s(g::s0, n, eppcgRsd2, n, polarity_inv, rsd, udir, uind, field, polscale);
+      launch_k1s(g::s0, n, pcgRsd0V3, n, polarity_inv, rsd, udir, uind, field, polscale);
    } else if (dirguess) {
       // uind is used here instead of udir since without exchange polarization udir = uind
       // but with exchange polarization udir != uind (for dirguess).