ComputeLonepairsCUDAKernel.cu

Go to the documentation of this file.

#include "ComputeLonepairsCUDA.h"
#include "ComputeLonepairsCUDAKernel.h"
#include "Vector.h"
#include <string>

#ifdef NAMD_HIP
#include <hipcub/hipcub.hpp>
#define cub hipcub
#endif //  NAMD_HIP

#ifdef NAMD_CUDA
#include <cuda.h>
#if __CUDACC_VER_MAJOR__ >= 11
#include <cub/cub.cuh>
#else
#include <namd_cub/cub.cuh>
#endif
#endif  // NAMD_CUDA

#if defined(NAMD_CUDA) || defined(NAMD_HIP)

// See https://github.com/HanatoK/lone_pair_force/blob/main/lone_pair_relative.ipynb for the maths
__global__ void repositionRelativeKernel(
  double* __restrict__ d_pos_x,
  double* __restrict__ d_pos_y,
  double* __restrict__ d_pos_z,
  const ComputeLonepairsCUDA::LonepairRelative* __restrict__ d_lprelative_list,
  size_t lprelative_list_size) {
  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
  if (tid < lprelative_list_size) {
    const int i = d_lprelative_list[tid].i_soaid;
    const int j = d_lprelative_list[tid].j_soaid;
    const int k = d_lprelative_list[tid].k_soaid;
    const int l = d_lprelative_list[tid].l_soaid;
    // Bernard Brooks: Do all LP math in double precision.  Cost is not excessive.
    const double dcosa = d_lprelative_list[tid].dcosa;
    const double dsinacost = d_lprelative_list[tid].dsinacost;
    const double dsinasint = d_lprelative_list[tid].dsinasint;
    const Vector rj{d_pos_x[j], d_pos_y[j], d_pos_z[j]};
    const Vector rk{d_pos_x[k], d_pos_y[k], d_pos_z[k]};
    const Vector rl{d_pos_x[l], d_pos_y[l], d_pos_z[l]};
    const Vector v = rk;
    Vector a = v - rj;
    Vector b = rl - v;
    BigReal invlen = rnorm3d(a.x, a.y, a.z);
    a *= invlen;
    Vector c = cross(b, a);
    invlen = rnorm3d(c.x, c.y, c.z);
    c *= invlen;
    b = cross(a, c);
    const Vector ri = rj + dcosa * a + dsinacost * b + dsinasint * c;
    d_pos_x[i] = ri.x;
    d_pos_y[i] = ri.y;
    d_pos_z[i] = ri.z;
  }
}

void repositionRelative(
  double* d_pos_x,
  double* d_pos_y,
  double* d_pos_z,
  const ComputeLonepairsCUDA::LonepairRelative* d_lprelative_list,
  size_t lprelative_list_size,
  cudaStream_t stream) {
  const int block_size = 128;
  const int grid = (lprelative_list_size + block_size - 1) / block_size;
  repositionRelativeKernel<<<grid, block_size, 0, stream>>>(
    d_pos_x, d_pos_y, d_pos_z, d_lprelative_list, lprelative_list_size);
}

__global__ void repositionBisectorKernel(
  double* __restrict__ d_pos_x,
  double* __restrict__ d_pos_y,
  double* __restrict__ d_pos_z,
  const ComputeLonepairsCUDA::LonepairBisector* __restrict__ d_lpbisector_list,
  size_t lpbisector_list_size) {
  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
  if (tid < lpbisector_list_size) {
    const int i = d_lpbisector_list[tid].i_soaid;
    const int j = d_lpbisector_list[tid].j_soaid;
    const int k = d_lpbisector_list[tid].k_soaid;
    const int l = d_lpbisector_list[tid].l_soaid;
    // Bernard Brooks: Do all LP math in double precision.  Cost is not excessive.
    const double dcosa = d_lpbisector_list[tid].dcosa;
    const double dsinacost = d_lpbisector_list[tid].dsinacost;
    const double dsinasint = d_lpbisector_list[tid].dsinasint;
    const Vector rj{d_pos_x[j], d_pos_y[j], d_pos_z[j]};
    const Vector rk{d_pos_x[k], d_pos_y[k], d_pos_z[k]};
    const Vector rl{d_pos_x[l], d_pos_y[l], d_pos_z[l]};
    // The middle point of l and k
    const Vector v = 0.5 * (rl + rk);
    Vector a = v - rj;
    Vector b = rl - v;
    BigReal invlen = rnorm3d(a.x, a.y, a.z);
    a *= invlen;
    Vector c = cross(b, a);
    invlen = rnorm3d(c.x, c.y, c.z);
    c *= invlen;
    b = cross(a, c);
    const Vector ri = rj + dcosa * a + dsinacost * b + dsinasint * c;
    d_pos_x[i] = ri.x;
    d_pos_y[i] = ri.y;
    d_pos_z[i] = ri.z;
  }
}

void repositionBisector(
  double* d_pos_x,
  double* d_pos_y,
  double* d_pos_z,
  const ComputeLonepairsCUDA::LonepairBisector* d_lpbisector_list,
  size_t lpbisector_list_size,
  cudaStream_t stream) {
  const int block_size = 128;
  const int grid = (lpbisector_list_size + block_size - 1) / block_size;
  repositionBisectorKernel<<<grid, block_size, 0, stream>>>(
    d_pos_x, d_pos_y, d_pos_z, d_lpbisector_list, lpbisector_list_size);
}

__global__ void repositionColinearKernel(
  double* __restrict__ d_pos_x,
  double* __restrict__ d_pos_y,
  double* __restrict__ d_pos_z,
  const ComputeLonepairsCUDA::LonepairColinear* __restrict__ d_lpcolinear_list,
  size_t lpcolinear_list_size) {
  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
  if (tid < lpcolinear_list_size) {
    const int i = d_lpcolinear_list[tid].i_soaid;
    const int j = d_lpcolinear_list[tid].j_soaid;
    const int k = d_lpcolinear_list[tid].k_soaid;
    const double distance = d_lpcolinear_list[tid].distance;
    const double scale_factor = d_lpcolinear_list[tid].scale_factor;
    const Vector rj{d_pos_x[j], d_pos_y[j], d_pos_z[j]};
    const Vector rk{d_pos_x[k], d_pos_y[k], d_pos_z[k]};
    const Vector rkj = rj - rk;
    const double inv_r = rnorm3d(rkj.x, rkj.y, rkj.z); // rkj.rlength()
    const Vector ri = rj + (scale_factor + distance * inv_r) * rkj;
    d_pos_x[i] = ri.x;
    d_pos_y[i] = ri.y;
    d_pos_z[i] = ri.z;
  }
}

void repositionColinear(
  double* __restrict__ d_pos_x,
  double* __restrict__ d_pos_y,
  double* __restrict__ d_pos_z,
  const ComputeLonepairsCUDA::LonepairColinear* __restrict__ d_lpcolinear_list,
  size_t lpcolinear_list_size,
  cudaStream_t stream) {
  const int block_size = 128;
  const int grid = (lpcolinear_list_size + block_size - 1) / block_size;
  repositionColinearKernel<<<grid, block_size, 0, stream>>>(
    d_pos_x, d_pos_y, d_pos_z, d_lpcolinear_list, lpcolinear_list_size);
}

// See https://github.com/HanatoK/lone_pair_force/blob/main/lone_pair_relative.ipynb for the maths
// I try to project multiple forces at once for MTS.
template <int numForces>
__inline__ __device__ void projectRelativeForces(
  const Vector& ri,
  const Vector& rj,
  const Vector& rk,
  const Vector& rl,
  const Vector (&fi)[numForces],
  Vector (&fj)[numForces],
  Vector (&fk)[numForces],
  Vector (&fl)[numForces]) {
  const Vector rji = ri - rj;
  const Vector rjk = rk - rj;
  const Vector rkl = rl - rk;
  const Vector rki = ri - rk;
  const double inv_rjk_norm2 = 1.0 / rjk.length2();
  const Vector plane_jkl = cross(rjk, rkl);
  // Project the force on rij to keep the distance unchanged
  const double rji_norm2 = rji.length2();
  const double inv_rij_norm2 = 1.0 / rji_norm2;
  if (rji_norm2 > 0) {
    #pragma unroll
    for (int m = 0; m < numForces; ++m) {
      fj[m] += inv_rij_norm2 * fi[m].dot(rji) * rji;
    }
  }
  // Project the force on the dihedral angle to keep it unchanged
  Vector fp[numForces];
  #pragma unroll
  for (int m = 0; m < numForces; ++m) {
    fp[m] = 0.0;
  }
  // Get the normal vector of plane ikj
  Vector normal_ikj = cross(rji, rjk);
  const double normal_ikj_norm = normal_ikj.length();
  // The arm of the force on plane ikj (v2 + v3)
  const Vector hijk_ri = rji - (rji.dot(rjk) * inv_rjk_norm2 * rjk);
  // Only project the force to plane ijk if the ijk are not colinear
  // If rji and rjk is colinear then rji×rjk is zero
  if (normal_ikj_norm > 1e-6) {
    normal_ikj /= normal_ikj_norm;
    // Force on plane ikj
    #pragma unroll
    for (int m = 0; m < numForces; ++m) {
      fp[m] = fi[m].dot(normal_ikj) * normal_ikj;
    }
  }
  // Torque of the force on l after projecting the force on plane ijk to l
  const Vector dir_l_to_jk = cross(plane_jkl, rjk).unit();
  const Vector h_l = rkl.dot(dir_l_to_jk) * dir_l_to_jk;
  const double inv_h_l_norm2 = 1.0 / h_l.length2();
  #pragma unroll
  for (int m = 0; m < numForces; ++m) {
    // Torque of fp
    const Vector torque_p = cross(hijk_ri, fp[m]);
    // The force on l
    const Vector fpl = inv_h_l_norm2 * cross(torque_p, h_l);
    fl[m] += fpl;
    // (a) The remaining force on j and k after subtracting the force on l from the torsional force
    const Vector fj_prime = inv_rjk_norm2 * cross(cross(rki, fp[m]) - cross(rkl, fpl), -rjk);
    const Vector fk_prime = fp[m] - fpl - fj_prime;
    // (b) The remaining angular force
    const Vector fai = fi[m] - fp[m] - fj[m];
    // Sum the torsional force (a) on j and k
    fj[m] += fj_prime;
    fk[m] += fk_prime;
    // Torque of the angular force (b) on k
    const Vector torque_k = cross(rji, fai);
    const Vector fak = cross(torque_k, rjk) * inv_rjk_norm2;
    fk[m] += fak;
    fj[m] += fai - fak;
  }
}

enum class LonepairThreeHostsType {Relative, Bisector};

template <typename LplistT, int maxForceNumer, bool doVirial, int blockSize, LonepairThreeHostsType lptype>
__global__ void redistributeForceThreeHostsKernel(
  double* __restrict__ d_f_normal_x,
  double* __restrict__ d_f_normal_y,
  double* __restrict__ d_f_normal_z,
  double* __restrict__ d_f_nbond_x,
  double* __restrict__ d_f_nbond_y,
  double* __restrict__ d_f_nbond_z,
  double* __restrict__ d_f_slow_x,
  double* __restrict__ d_f_slow_y,
  double* __restrict__ d_f_slow_z,
  cudaTensor* __restrict__ d_virial_normal,
  cudaTensor* __restrict__ d_virial_nbond,
  cudaTensor* __restrict__ d_virial_slow,
  const double* __restrict__ d_pos_x,
  const double* __restrict__ d_pos_y,
  const double* __restrict__ d_pos_z,
  const LplistT* __restrict__ d_lp_list,
  size_t lp_list_size) {
  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
  cudaTensor vir_slow, vir_nbond, vir_normal;
  if (doVirial) {
    switch (maxForceNumer) {
      // Intentionally fall-through
      case 2: {
        vir_slow.xx = 0;
        vir_slow.xy = 0;
        vir_slow.xz = 0;
        vir_slow.yx = 0;
        vir_slow.yy = 0;
        vir_slow.yz = 0;
        vir_slow.zx = 0;
        vir_slow.zy = 0;
        vir_slow.zz = 0;
      }
      case 1: {
        vir_nbond.xx = 0;
        vir_nbond.xy = 0;
        vir_nbond.xz = 0;
        vir_nbond.yx = 0;
        vir_nbond.yy = 0;
        vir_nbond.yz = 0;
        vir_nbond.zx = 0;
        vir_nbond.zy = 0;
        vir_nbond.zz = 0;
      }
      case 0: {
        vir_normal.xx = 0;
        vir_normal.xy = 0;
        vir_normal.xz = 0;
        vir_normal.yx = 0;
        vir_normal.yy = 0;
        vir_normal.yz = 0;
        vir_normal.zx = 0;
        vir_normal.zy = 0;
        vir_normal.zz = 0;
      }
    }
  }
  if (tid < lp_list_size) {
    const int i = d_lp_list[tid].i_soaid;
    const int j = d_lp_list[tid].j_soaid;
    const int k = d_lp_list[tid].k_soaid;
    const int l = d_lp_list[tid].l_soaid;
    // Bernard Brooks: Do all LP math in double precision.  Cost is not excessive.
    const Vector ri{d_pos_x[i], d_pos_y[i], d_pos_z[i]};
    const Vector rj{d_pos_x[j], d_pos_y[j], d_pos_z[j]};
    const Vector rk{d_pos_x[k], d_pos_y[k], d_pos_z[k]};
    const Vector rl{d_pos_x[l], d_pos_y[l], d_pos_z[l]};
    Vector fi[maxForceNumer+1];
    Vector fj[maxForceNumer+1];
    Vector fk[maxForceNumer+1];
    Vector fl[maxForceNumer+1];
    #pragma unroll
    for (int m = 0; m < maxForceNumer + 1; ++m) {
      fj[m] = 0.0;
      fk[m] = 0.0;
      fl[m] = 0.0;
    }
    switch (maxForceNumer) {
      // Intentionally fall-through
      case 2: fi[2] = Vector{d_f_slow_x[i], d_f_slow_y[i], d_f_slow_z[i]};
      case 1: fi[1] = Vector{d_f_nbond_x[i], d_f_nbond_y[i], d_f_nbond_z[i]};
      case 0: fi[0] = Vector{d_f_normal_x[i], d_f_normal_y[i], d_f_normal_z[i]};
    }
    switch (lptype) {
      case LonepairThreeHostsType::Relative: {
        projectRelativeForces<maxForceNumer+1>(ri, rj, rk, rl, fi, fj, fk, fl);
        break;
      }
      case LonepairThreeHostsType::Bisector: {
        // The difference between the bisector and the relative cases is
        // that the order of the vertices of the former is changed to ijmk
        // from ijkl, so we can just reuse the function.
        const Vector rm = 0.5 * (rj + rk);
        // Store the force of m to fl, so I rotate the buffers fl and fk
        projectRelativeForces<maxForceNumer+1>(ri, rj, rm, rk, fi, fj, fl, fk);
        // fk and fl actually store the forces on m and k, respectively.
        // The real force on k is 0.5 * fl + fk
        // The real force on l is 0.5 * fl
        #pragma unroll
        for (int m = 0; m < maxForceNumer + 1; ++m) {
          fl[m] *= 0.5;
          fk[m] += fl[m];
        }
        break;
      }
    }
    // Add the forces back to jkl
    switch (maxForceNumer) {
      // Intentionally fall-through
      case 2: {
        atomicAdd(&(d_f_slow_x[j]), fj[2].x);
        atomicAdd(&(d_f_slow_y[j]), fj[2].y);
        atomicAdd(&(d_f_slow_z[j]), fj[2].z);
        atomicAdd(&(d_f_slow_x[k]), fk[2].x);
        atomicAdd(&(d_f_slow_y[k]), fk[2].y);
        atomicAdd(&(d_f_slow_z[k]), fk[2].z);
        atomicAdd(&(d_f_slow_x[l]), fl[2].x);
        atomicAdd(&(d_f_slow_y[l]), fl[2].y);
        atomicAdd(&(d_f_slow_z[l]), fl[2].z);
      }
      case 1: {
        atomicAdd(&(d_f_nbond_x[j]), fj[1].x);
        atomicAdd(&(d_f_nbond_y[j]), fj[1].y);
        atomicAdd(&(d_f_nbond_z[j]), fj[1].z);
        atomicAdd(&(d_f_nbond_x[k]), fk[1].x);
        atomicAdd(&(d_f_nbond_y[k]), fk[1].y);
        atomicAdd(&(d_f_nbond_z[k]), fk[1].z);
        atomicAdd(&(d_f_nbond_x[l]), fl[1].x);
        atomicAdd(&(d_f_nbond_y[l]), fl[1].y);
        atomicAdd(&(d_f_nbond_z[l]), fl[1].z);
      }
      case 0: {
        atomicAdd(&(d_f_normal_x[j]), fj[0].x);
        atomicAdd(&(d_f_normal_y[j]), fj[0].y);
        atomicAdd(&(d_f_normal_z[j]), fj[0].z);
        atomicAdd(&(d_f_normal_x[k]), fk[0].x);
        atomicAdd(&(d_f_normal_y[k]), fk[0].y);
        atomicAdd(&(d_f_normal_z[k]), fk[0].z);
        atomicAdd(&(d_f_normal_x[l]), fl[0].x);
        atomicAdd(&(d_f_normal_y[l]), fl[0].y);
        atomicAdd(&(d_f_normal_z[l]), fl[0].z);
      }
    }
    // Compute the virial contribution after redistributing the forces
    if (doVirial) {
      switch (maxForceNumer) {
        // Intentionally fall-through
        case 2: {
          vir_slow.xx = fj[2].x * rj.x + fk[2].x * rk.x + fl[2].x * rl.x - fi[2].x * ri.x;
          vir_slow.xy = fj[2].x * rj.y + fk[2].x * rk.y + fl[2].x * rl.y - fi[2].x * ri.y;
          vir_slow.xz = fj[2].x * rj.z + fk[2].x * rk.z + fl[2].x * rl.z - fi[2].x * ri.z;
          vir_slow.yx = fj[2].y * rj.x + fk[2].y * rk.x + fl[2].y * rl.x - fi[2].y * ri.x;
          vir_slow.yy = fj[2].y * rj.y + fk[2].y * rk.y + fl[2].y * rl.y - fi[2].y * ri.y;
          vir_slow.yz = fj[2].y * rj.z + fk[2].y * rk.z + fl[2].y * rl.z - fi[2].y * ri.z;
          vir_slow.zx = fj[2].z * rj.x + fk[2].z * rk.x + fl[2].z * rl.x - fi[2].z * ri.x;
          vir_slow.zy = fj[2].z * rj.y + fk[2].z * rk.y + fl[2].z * rl.y - fi[2].z * ri.y;
          vir_slow.zz = fj[2].z * rj.z + fk[2].z * rk.z + fl[2].z * rl.z - fi[2].z * ri.z;
        }
        case 1: {
          vir_nbond.xx = fj[1].x * rj.x + fk[1].x * rk.x + fl[1].x * rl.x - fi[1].x * ri.x;
          vir_nbond.xy = fj[1].x * rj.y + fk[1].x * rk.y + fl[1].x * rl.y - fi[1].x * ri.y;
          vir_nbond.xz = fj[1].x * rj.z + fk[1].x * rk.z + fl[1].x * rl.z - fi[1].x * ri.z;
          vir_nbond.yx = fj[1].y * rj.x + fk[1].y * rk.x + fl[1].y * rl.x - fi[1].y * ri.x;
          vir_nbond.yy = fj[1].y * rj.y + fk[1].y * rk.y + fl[1].y * rl.y - fi[1].y * ri.y;
          vir_nbond.yz = fj[1].y * rj.z + fk[1].y * rk.z + fl[1].y * rl.z - fi[1].y * ri.z;
          vir_nbond.zx = fj[1].z * rj.x + fk[1].z * rk.x + fl[1].z * rl.x - fi[1].z * ri.x;
          vir_nbond.zy = fj[1].z * rj.y + fk[1].z * rk.y + fl[1].z * rl.y - fi[1].z * ri.y;
          vir_nbond.zz = fj[1].z * rj.z + fk[1].z * rk.z + fl[1].z * rl.z - fi[1].z * ri.z;
        }
        case 0: {
          vir_normal.xx = fj[0].x * rj.x + fk[0].x * rk.x + fl[0].x * rl.x - fi[0].x * ri.x;
          vir_normal.xy = fj[0].x * rj.y + fk[0].x * rk.y + fl[0].x * rl.y - fi[0].x * ri.y;
          vir_normal.xz = fj[0].x * rj.z + fk[0].x * rk.z + fl[0].x * rl.z - fi[0].x * ri.z;
          vir_normal.yx = fj[0].y * rj.x + fk[0].y * rk.x + fl[0].y * rl.x - fi[0].y * ri.x;
          vir_normal.yy = fj[0].y * rj.y + fk[0].y * rk.y + fl[0].y * rl.y - fi[0].y * ri.y;
          vir_normal.yz = fj[0].y * rj.z + fk[0].y * rk.z + fl[0].y * rl.z - fi[0].y * ri.z;
          vir_normal.zx = fj[0].z * rj.x + fk[0].z * rk.x + fl[0].z * rl.x - fi[0].z * ri.x;
          vir_normal.zy = fj[0].z * rj.y + fk[0].z * rk.y + fl[0].z * rl.y - fi[0].z * ri.y;
          vir_normal.zz = fj[0].z * rj.z + fk[0].z * rk.z + fl[0].z * rl.z - fi[0].z * ri.z;
        }
      }
    }
    // Clean the forces on the lone pair
    switch (maxForceNumer) {
      // Intentionally fall-through
      case 2: {
        d_f_slow_x[i] = 0;
        d_f_slow_y[i] = 0;
        d_f_slow_z[i] = 0;
      }
      case 1: {
        d_f_nbond_x[i] = 0;
        d_f_nbond_y[i] = 0;
        d_f_nbond_z[i] = 0;
      }
      case 0: {
        d_f_normal_x[i] = 0;
        d_f_normal_y[i] = 0;
        d_f_normal_z[i] = 0;
      }
    }
  }
  typedef cub::BlockReduce<cudaTensor, blockSize> BlockReduce;
  __shared__ typename BlockReduce::TempStorage temp_storage;
  if (doVirial) {
    switch (maxForceNumer) {
      // Intentionally fall-through
      case 2: {
        vir_slow = BlockReduce(temp_storage).Reduce(vir_slow, [](const cudaTensor& a, const cudaTensor& b){return a+b;});
        __syncthreads();
      }
      case 1: {
        vir_nbond = BlockReduce(temp_storage).Reduce(vir_nbond, [](const cudaTensor& a, const cudaTensor& b){return a+b;});
        __syncthreads();
      }
      case 0: {
        vir_normal = BlockReduce(temp_storage).Reduce(vir_normal, [](const cudaTensor& a, const cudaTensor& b){return a+b;});
        __syncthreads();
      }
    }
    if (threadIdx.x == 0) {
      switch (maxForceNumer) {
        // Intentionally fall-through
        case 2: {
          atomicAdd(&(d_virial_slow->xx), vir_slow.xx);
          atomicAdd(&(d_virial_slow->xy), vir_slow.xy);
          atomicAdd(&(d_virial_slow->xz), vir_slow.xz);
          atomicAdd(&(d_virial_slow->yx), vir_slow.yx);
          atomicAdd(&(d_virial_slow->yy), vir_slow.yy);
          atomicAdd(&(d_virial_slow->yz), vir_slow.yz);
          atomicAdd(&(d_virial_slow->zx), vir_slow.zx);
          atomicAdd(&(d_virial_slow->zy), vir_slow.zy);
          atomicAdd(&(d_virial_slow->zz), vir_slow.zz);
        }
        case 1: {
          atomicAdd(&(d_virial_nbond->xx), vir_nbond.xx);
          atomicAdd(&(d_virial_nbond->xy), vir_nbond.xy);
          atomicAdd(&(d_virial_nbond->xz), vir_nbond.xz);
          atomicAdd(&(d_virial_nbond->yx), vir_nbond.yx);
          atomicAdd(&(d_virial_nbond->yy), vir_nbond.yy);
          atomicAdd(&(d_virial_nbond->yz), vir_nbond.yz);
          atomicAdd(&(d_virial_nbond->zx), vir_nbond.zx);
          atomicAdd(&(d_virial_nbond->zy), vir_nbond.zy);
          atomicAdd(&(d_virial_nbond->zz), vir_nbond.zz);
        }
        case 0: {
          atomicAdd(&(d_virial_normal->xx), vir_normal.xx);
          atomicAdd(&(d_virial_normal->xy), vir_normal.xy);
          atomicAdd(&(d_virial_normal->xz), vir_normal.xz);
          atomicAdd(&(d_virial_normal->yx), vir_normal.yx);
          atomicAdd(&(d_virial_normal->yy), vir_normal.yy);
          atomicAdd(&(d_virial_normal->yz), vir_normal.yz);
          atomicAdd(&(d_virial_normal->zx), vir_normal.zx);
          atomicAdd(&(d_virial_normal->zy), vir_normal.zy);
          atomicAdd(&(d_virial_normal->zz), vir_normal.zz);
        }
      }
    }
  }
}

void redistributeForceRelative(
  double*        d_f_normal_x,
  double*        d_f_normal_y,
  double*        d_f_normal_z,
  double*        d_f_nbond_x,
  double*        d_f_nbond_y,
  double*        d_f_nbond_z,
  double*        d_f_slow_x,
  double*        d_f_slow_y,
  double*        d_f_slow_z,
  cudaTensor*    d_virial_normal,
  cudaTensor*    d_virial_nbond,
  cudaTensor*    d_virial_slow,
  const double*  d_pos_x,
  const double*  d_pos_y,
  const double*  d_pos_z,
  const int      maxForceNumber,
  const int      doVirial,
  const ComputeLonepairsCUDA::LonepairRelative* d_lprelative_list,
  size_t lprelative_list_size,
  cudaStream_t stream) {
  const int block_size = 128;
  const int grid = (lprelative_list_size + block_size - 1) / block_size;
#define CALL(MAXFORCENUMBER, DOVIRIAL) \
  redistributeForceThreeHostsKernel<ComputeLonepairsCUDA::LonepairRelative, MAXFORCENUMBER, DOVIRIAL, block_size, LonepairThreeHostsType::Relative> \
  <<<grid, block_size, 0, stream>>>( \
    d_f_normal_x, d_f_normal_y, d_f_normal_z, \
    d_f_nbond_x, d_f_nbond_y, d_f_nbond_z, \
    d_f_slow_x, d_f_slow_y, d_f_slow_z, \
    d_virial_normal, d_virial_nbond, d_virial_slow, \
    d_pos_x, d_pos_y, d_pos_z, \
    d_lprelative_list, lprelative_list_size \
  );
  const int option = (maxForceNumber << 0) + (doVirial << 2);
  switch (option) {
    case ((0 << 0) + (0 << 2)): CALL(0, 0); break;
    case ((1 << 0) + (0 << 2)): CALL(1, 0); break;
    case ((2 << 0) + (0 << 2)): CALL(2, 0); break;
    case ((0 << 0) + (1 << 2)): CALL(0, 1); break;
    case ((1 << 0) + (1 << 2)): CALL(1, 1); break;
    case ((2 << 0) + (1 << 2)): CALL(2, 1); break;
    default: {
      const std::string error =
        "redistributeForceRelative: no kernel called (maxForceNumber = " +
        std::to_string(maxForceNumber) + ", doVirial = " + std::to_string(doVirial);
      NAMD_bug(error.c_str());
    }
  }
#undef CALL
}

void redistributeForceBisector(
  double*        d_f_normal_x,
  double*        d_f_normal_y,
  double*        d_f_normal_z,
  double*        d_f_nbond_x,
  double*        d_f_nbond_y,
  double*        d_f_nbond_z,
  double*        d_f_slow_x,
  double*        d_f_slow_y,
  double*        d_f_slow_z,
  cudaTensor*    d_virial_normal,
  cudaTensor*    d_virial_nbond,
  cudaTensor*    d_virial_slow,
  const double*  d_pos_x,
  const double*  d_pos_y,
  const double*  d_pos_z,
  const int      maxForceNumber,
  const int      doVirial,
  const ComputeLonepairsCUDA::LonepairBisector* d_lpbisector_list,
  size_t lpbisector_list_size,
  cudaStream_t stream) {
  const int block_size = 128;
  const int grid = (lpbisector_list_size + block_size - 1) / block_size;
#define CALL(MAXFORCENUMBER, DOVIRIAL) \
  redistributeForceThreeHostsKernel<ComputeLonepairsCUDA::LonepairBisector, MAXFORCENUMBER, DOVIRIAL, block_size, LonepairThreeHostsType::Bisector> \
  <<<grid, block_size, 0, stream>>>( \
    d_f_normal_x, d_f_normal_y, d_f_normal_z, \
    d_f_nbond_x, d_f_nbond_y, d_f_nbond_z, \
    d_f_slow_x, d_f_slow_y, d_f_slow_z, \
    d_virial_normal, d_virial_nbond, d_virial_slow, \
    d_pos_x, d_pos_y, d_pos_z, \
    d_lpbisector_list, lpbisector_list_size \
  );
  const int option = (maxForceNumber << 0) + (doVirial << 2);
  switch (option) {
    case ((0 << 0) + (0 << 2)): CALL(0, 0); break;
    case ((1 << 0) + (0 << 2)): CALL(1, 0); break;
    case ((2 << 0) + (0 << 2)): CALL(2, 0); break;
    case ((0 << 0) + (1 << 2)): CALL(0, 1); break;
    case ((1 << 0) + (1 << 2)): CALL(1, 1); break;
    case ((2 << 0) + (1 << 2)): CALL(2, 1); break;
    default: {
      const std::string error =
        "redistributeForceBisector: no kernel called (maxForceNumber = " +
        std::to_string(maxForceNumber) + ", doVirial = " + std::to_string(doVirial);
      NAMD_bug(error.c_str());
    }
  }
#undef CALL
}

template <int maxForceNumer, bool doVirial, int blockSize>
__global__ void redistributeForceColinearKernel(
  double* __restrict__ d_f_normal_x,
  double* __restrict__ d_f_normal_y,
  double* __restrict__ d_f_normal_z,
  double* __restrict__ d_f_nbond_x,
  double* __restrict__ d_f_nbond_y,
  double* __restrict__ d_f_nbond_z,
  double* __restrict__ d_f_slow_x,
  double* __restrict__ d_f_slow_y,
  double* __restrict__ d_f_slow_z,
  cudaTensor* __restrict__ d_virial_normal,
  cudaTensor* __restrict__ d_virial_nbond,
  cudaTensor* __restrict__ d_virial_slow,
  const double* __restrict__ d_pos_x,
  const double* __restrict__ d_pos_y,
  const double* __restrict__ d_pos_z,
  const ComputeLonepairsCUDA::LonepairColinear* __restrict__ d_lp_list,
  size_t lp_list_size) {
  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
  cudaTensor vir_slow, vir_nbond, vir_normal;
  if (doVirial) {
    switch (maxForceNumer) {
      // Intentionally fall-through
      case 2: {
        vir_slow.xx = 0;
        vir_slow.xy = 0;
        vir_slow.xz = 0;
        vir_slow.yx = 0;
        vir_slow.yy = 0;
        vir_slow.yz = 0;
        vir_slow.zx = 0;
        vir_slow.zy = 0;
        vir_slow.zz = 0;
      }
      case 1: {
        vir_nbond.xx = 0;
        vir_nbond.xy = 0;
        vir_nbond.xz = 0;
        vir_nbond.yx = 0;
        vir_nbond.yy = 0;
        vir_nbond.yz = 0;
        vir_nbond.zx = 0;
        vir_nbond.zy = 0;
        vir_nbond.zz = 0;
      }
      case 0: {
        vir_normal.xx = 0;
        vir_normal.xy = 0;
        vir_normal.xz = 0;
        vir_normal.yx = 0;
        vir_normal.yy = 0;
        vir_normal.yz = 0;
        vir_normal.zx = 0;
        vir_normal.zy = 0;
        vir_normal.zz = 0;
      }
    }
  }
  if (tid < lp_list_size) {
    const int i = d_lp_list[tid].i_soaid;
    const int j = d_lp_list[tid].j_soaid;
    const int k = d_lp_list[tid].k_soaid;
    double distance = d_lp_list[tid].distance;
    const double scale_factor = d_lp_list[tid].scale_factor;
    const Vector rj{d_pos_x[j], d_pos_y[j], d_pos_z[j]};
    const Vector rk{d_pos_x[k], d_pos_y[k], d_pos_z[k]};
    const Vector rkj = rj - rk;
    const double inv_rkj_norm = rnorm3d(rkj.x, rkj.y, rkj.z); // rkj.rlength()
    distance *= inv_rkj_norm;
    // Prepare the force buffers
    Vector fi[maxForceNumer+1];
    Vector fj[maxForceNumer+1];
    Vector fk[maxForceNumer+1];
    #pragma unroll
    for (int m = 0; m < maxForceNumer + 1; ++m) {
      fj[m] = 0.0;
      fk[m] = 0.0;
    }
    switch (maxForceNumer) {
      // Intentionally fall-through
      case 2: fi[2] = Vector{d_f_slow_x[i], d_f_slow_y[i], d_f_slow_z[i]};
      case 1: fi[1] = Vector{d_f_nbond_x[i], d_f_nbond_y[i], d_f_nbond_z[i]};
      case 0: fi[0] = Vector{d_f_normal_x[i], d_f_normal_y[i], d_f_normal_z[i]};
    }
    // Project the forces
    #pragma unroll
    for (int m = 0; m < maxForceNumer + 1; ++m) {
      const double fdot = distance * (fi[m].dot(rkj)) * inv_rkj_norm * inv_rkj_norm;
      fj[m] = (1.0 + scale_factor + distance) * fi[m] - fdot * rkj;
      fk[m] = fi[m] - fj[m];
    }
    // Add the forces back
    switch (maxForceNumer) {
      // Intentionally fall-through
      case 2: {
        atomicAdd(&(d_f_slow_x[j]), fj[2].x);
        atomicAdd(&(d_f_slow_y[j]), fj[2].y);
        atomicAdd(&(d_f_slow_z[j]), fj[2].z);
        atomicAdd(&(d_f_slow_x[k]), fk[2].x);
        atomicAdd(&(d_f_slow_y[k]), fk[2].y);
        atomicAdd(&(d_f_slow_z[k]), fk[2].z);
      }
      case 1: {
        atomicAdd(&(d_f_nbond_x[j]), fj[1].x);
        atomicAdd(&(d_f_nbond_y[j]), fj[1].y);
        atomicAdd(&(d_f_nbond_z[j]), fj[1].z);
        atomicAdd(&(d_f_nbond_x[k]), fk[1].x);
        atomicAdd(&(d_f_nbond_y[k]), fk[1].y);
        atomicAdd(&(d_f_nbond_z[k]), fk[1].z);
      }
      case 0: {
        atomicAdd(&(d_f_normal_x[j]), fj[0].x);
        atomicAdd(&(d_f_normal_y[j]), fj[0].y);
        atomicAdd(&(d_f_normal_z[j]), fj[0].z);
        atomicAdd(&(d_f_normal_x[k]), fk[0].x);
        atomicAdd(&(d_f_normal_y[k]), fk[0].y);
        atomicAdd(&(d_f_normal_z[k]), fk[0].z);
      }
    }
    if (doVirial) {
      const Vector ri{d_pos_x[i], d_pos_y[i], d_pos_z[i]};
      switch (maxForceNumer) {
        // Intentionally fall-through
        case 2: {
          vir_slow.xx = fj[2].x * rj.x + fk[2].x * rk.x - fi[2].x * ri.x;
          vir_slow.xy = fj[2].x * rj.y + fk[2].x * rk.y - fi[2].x * ri.y;
          vir_slow.xz = fj[2].x * rj.z + fk[2].x * rk.z - fi[2].x * ri.z;
          vir_slow.yx = fj[2].y * rj.x + fk[2].y * rk.x - fi[2].y * ri.x;
          vir_slow.yy = fj[2].y * rj.y + fk[2].y * rk.y - fi[2].y * ri.y;
          vir_slow.yz = fj[2].y * rj.z + fk[2].y * rk.z - fi[2].y * ri.z;
          vir_slow.zx = fj[2].z * rj.x + fk[2].z * rk.x - fi[2].z * ri.x;
          vir_slow.zy = fj[2].z * rj.y + fk[2].z * rk.y - fi[2].z * ri.y;
          vir_slow.zz = fj[2].z * rj.z + fk[2].z * rk.z - fi[2].z * ri.z;
        }
        case 1: {
          vir_nbond.xx = fj[1].x * rj.x + fk[1].x * rk.x - fi[1].x * ri.x;
          vir_nbond.xy = fj[1].x * rj.y + fk[1].x * rk.y - fi[1].x * ri.y;
          vir_nbond.xz = fj[1].x * rj.z + fk[1].x * rk.z - fi[1].x * ri.z;
          vir_nbond.yx = fj[1].y * rj.x + fk[1].y * rk.x - fi[1].y * ri.x;
          vir_nbond.yy = fj[1].y * rj.y + fk[1].y * rk.y - fi[1].y * ri.y;
          vir_nbond.yz = fj[1].y * rj.z + fk[1].y * rk.z - fi[1].y * ri.z;
          vir_nbond.zx = fj[1].z * rj.x + fk[1].z * rk.x - fi[1].z * ri.x;
          vir_nbond.zy = fj[1].z * rj.y + fk[1].z * rk.y - fi[1].z * ri.y;
          vir_nbond.zz = fj[1].z * rj.z + fk[1].z * rk.z - fi[1].z * ri.z;
        }
        case 0: {
          vir_normal.xx = fj[0].x * rj.x + fk[0].x * rk.x - fi[0].x * ri.x;
          vir_normal.xy = fj[0].x * rj.y + fk[0].x * rk.y - fi[0].x * ri.y;
          vir_normal.xz = fj[0].x * rj.z + fk[0].x * rk.z - fi[0].x * ri.z;
          vir_normal.yx = fj[0].y * rj.x + fk[0].y * rk.x - fi[0].y * ri.x;
          vir_normal.yy = fj[0].y * rj.y + fk[0].y * rk.y - fi[0].y * ri.y;
          vir_normal.yz = fj[0].y * rj.z + fk[0].y * rk.z - fi[0].y * ri.z;
          vir_normal.zx = fj[0].z * rj.x + fk[0].z * rk.x - fi[0].z * ri.x;
          vir_normal.zy = fj[0].z * rj.y + fk[0].z * rk.y - fi[0].z * ri.y;
          vir_normal.zz = fj[0].z * rj.z + fk[0].z * rk.z - fi[0].z * ri.z;
        }
      }
    }
    // Clean the forces on the lone pair
    switch (maxForceNumer) {
      // Intentionally fall-through
      case 2: {
        d_f_slow_x[i] = 0;
        d_f_slow_y[i] = 0;
        d_f_slow_z[i] = 0;
      }
      case 1: {
        d_f_nbond_x[i] = 0;
        d_f_nbond_y[i] = 0;
        d_f_nbond_z[i] = 0;
      }
      case 0: {
        d_f_normal_x[i] = 0;
        d_f_normal_y[i] = 0;
        d_f_normal_z[i] = 0;
      }
    }
  }
  typedef cub::BlockReduce<cudaTensor, blockSize> BlockReduce;
  __shared__ typename BlockReduce::TempStorage temp_storage;
  if (doVirial) {
    switch (maxForceNumer) {
      // Intentionally fall-through
      case 2: {
        vir_slow = BlockReduce(temp_storage).Reduce(vir_slow, [](const cudaTensor& a, const cudaTensor& b){return a+b;});
        __syncthreads();
      }
      case 1: {
        vir_nbond = BlockReduce(temp_storage).Reduce(vir_nbond, [](const cudaTensor& a, const cudaTensor& b){return a+b;});
        __syncthreads();
      }
      case 0: {
        vir_normal = BlockReduce(temp_storage).Reduce(vir_normal, [](const cudaTensor& a, const cudaTensor& b){return a+b;});
        __syncthreads();
      }
    }
    if (threadIdx.x == 0) {
      switch (maxForceNumer) {
        // Intentionally fall-through
        case 2: {
          atomicAdd(&(d_virial_slow->xx), vir_slow.xx);
          atomicAdd(&(d_virial_slow->xy), vir_slow.xy);
          atomicAdd(&(d_virial_slow->xz), vir_slow.xz);
          atomicAdd(&(d_virial_slow->yx), vir_slow.yx);
          atomicAdd(&(d_virial_slow->yy), vir_slow.yy);
          atomicAdd(&(d_virial_slow->yz), vir_slow.yz);
          atomicAdd(&(d_virial_slow->zx), vir_slow.zx);
          atomicAdd(&(d_virial_slow->zy), vir_slow.zy);
          atomicAdd(&(d_virial_slow->zz), vir_slow.zz);
        }
        case 1: {
          atomicAdd(&(d_virial_nbond->xx), vir_nbond.xx);
          atomicAdd(&(d_virial_nbond->xy), vir_nbond.xy);
          atomicAdd(&(d_virial_nbond->xz), vir_nbond.xz);
          atomicAdd(&(d_virial_nbond->yx), vir_nbond.yx);
          atomicAdd(&(d_virial_nbond->yy), vir_nbond.yy);
          atomicAdd(&(d_virial_nbond->yz), vir_nbond.yz);
          atomicAdd(&(d_virial_nbond->zx), vir_nbond.zx);
          atomicAdd(&(d_virial_nbond->zy), vir_nbond.zy);
          atomicAdd(&(d_virial_nbond->zz), vir_nbond.zz);
        }
        case 0: {
          atomicAdd(&(d_virial_normal->xx), vir_normal.xx);
          atomicAdd(&(d_virial_normal->xy), vir_normal.xy);
          atomicAdd(&(d_virial_normal->xz), vir_normal.xz);
          atomicAdd(&(d_virial_normal->yx), vir_normal.yx);
          atomicAdd(&(d_virial_normal->yy), vir_normal.yy);
          atomicAdd(&(d_virial_normal->yz), vir_normal.yz);
          atomicAdd(&(d_virial_normal->zx), vir_normal.zx);
          atomicAdd(&(d_virial_normal->zy), vir_normal.zy);
          atomicAdd(&(d_virial_normal->zz), vir_normal.zz);
        }
      }
    }
  }
}

void redistributeForceColinear(
  double*        d_f_normal_x,
  double*        d_f_normal_y,
  double*        d_f_normal_z,
  double*        d_f_nbond_x,
  double*        d_f_nbond_y,
  double*        d_f_nbond_z,
  double*        d_f_slow_x,
  double*        d_f_slow_y,
  double*        d_f_slow_z,
  cudaTensor*    d_virial_normal,
  cudaTensor*    d_virial_nbond,
  cudaTensor*    d_virial_slow,
  const double*  d_pos_x,
  const double*  d_pos_y,
  const double*  d_pos_z,
  const int      maxForceNumber,
  const int      doVirial,
  const ComputeLonepairsCUDA::LonepairColinear* d_lpcolinear_list,
  size_t lpcolinear_list_size,
  cudaStream_t stream) {
  const int block_size = 128;
  const int grid = (lpcolinear_list_size + block_size - 1) / block_size;
#define CALL(MAXFORCENUMBER, DOVIRIAL) \
  redistributeForceColinearKernel<MAXFORCENUMBER, DOVIRIAL, block_size> \
  <<<grid, block_size, 0, stream>>>( \
    d_f_normal_x, d_f_normal_y, d_f_normal_z, \
    d_f_nbond_x, d_f_nbond_y, d_f_nbond_z, \
    d_f_slow_x, d_f_slow_y, d_f_slow_z, \
    d_virial_normal, d_virial_nbond, d_virial_slow, \
    d_pos_x, d_pos_y, d_pos_z, \
    d_lpcolinear_list, lpcolinear_list_size);
  const int option = (maxForceNumber << 0) + (doVirial << 2);
  switch (option) {
    case ((0 << 0) + (0 << 2)): CALL(0, 0); break;
    case ((1 << 0) + (0 << 2)): CALL(1, 0); break;
    case ((2 << 0) + (0 << 2)): CALL(2, 0); break;
    case ((0 << 0) + (1 << 2)): CALL(0, 1); break;
    case ((1 << 0) + (1 << 2)): CALL(1, 1); break;
    case ((2 << 0) + (1 << 2)): CALL(2, 1); break;
    default: {
      const std::string error =
        "redistributeForceColinear: no kernel called (maxForceNumber = " +
        std::to_string(maxForceNumber) + ", doVirial = " + std::to_string(doVirial);
      NAMD_bug(error.c_str());
    }
  }
#undef CALL
}

#endif // NAMD_CUDA