remove xkl - xk

bytedance · Dec 14, 2024 · 538f78e · 538f78e
1 parent 39b1de7
commit 538f78e
Show file tree

Hide file tree

Showing 9 changed files with 38 additions and 91 deletions.
diff --git a/gpu4pyscf/lib/gint/g2e.cu b/gpu4pyscf/lib/gint/g2e.cu
@@ -479,15 +479,9 @@ static void GINTg0_int3c2e(GINTEnvVars envs, double* __restrict__ g,
     const double xi = bas_x[ish];
     const double yi = bas_y[ish];
     const double zi = bas_z[ish];
-    const double xk = bas_x[ksh];
-    const double yk = bas_y[ksh];
-    const double zk = bas_z[ksh];
     const double xijxi = xij - xi;
     const double yijyi = yij - yi;
     const double zijzi = zij - zi;
-    const double xklxk = xkl - xk;
-    const double yklyk = ykl - yk;
-    const double zklzk = zkl - zk;
 
     int nmax = envs.li_ceil + envs.lj_ceil;
     int mmax = envs.lk_ceil + envs.ll_ceil;
@@ -563,9 +557,9 @@ static void GINTg0_int3c2e(GINTEnvVars envs, double* __restrict__ g,
             //}
             const double tmp3 = tmp1 * aij;
             const double b01 = b00 + tmp4 * aij;
-            const double c0px = xklxk + tmp3 * xijxkl;
-            const double c0py = yklyk + tmp3 * yijykl;
-            const double c0pz = zklzk + tmp3 * zijzkl;
+            const double c0px = tmp3 * xijxkl;
+            const double c0py = tmp3 * yijykl;
+            const double c0pz = tmp3 * zijzkl;
             double s0x = gx[i];
             double s0y = gy[i];
             double s0z = gz[i];

diff --git a/gpu4pyscf/lib/gint/g3c2e.cu b/gpu4pyscf/lib/gint/g3c2e.cu
@@ -151,18 +151,11 @@ static void GINTfill_int3c2e_kernel0010(GINTEnvVars envs, ERITensor eri, BasisPr
     double* __restrict__ x12 = c_bpcache.x12;
     double* __restrict__ y12 = c_bpcache.y12;
     double* __restrict__ z12 = c_bpcache.z12;
-
-    const int nbas = c_bpcache.nbas;
-    double* __restrict__ bas_x = c_bpcache.bas_coords;
-    double* __restrict__ bas_y = bas_x + nbas;
-    double* __restrict__ bas_z = bas_y + nbas;
 
     double gout0 = 0;
     double gout1 = 0;
     double gout2 = 0;
-    const double xk = bas_x[ksh];
-    const double yk = bas_y[ksh];
-    const double zk = bas_z[ksh];
+
     const int prim_ij0 = prim_ij;
     const int prim_ij1 = prim_ij + nprim_ij;
     const int prim_kl0 = prim_kl;
@@ -208,9 +201,9 @@ static void GINTfill_int3c2e_kernel0010(GINTEnvVars envs, ERITensor eri, BasisPr
         //const double b00 = u2 * tmp4;
         //const double tmp1 = aij * u2 / (u2 * aijkl + a1);;
         const double tmp3 = aij * u2 / (u2 * aijkl + a1);;
-        const double c0px = xkl - xk + tmp3 * xijxkl;
-        const double c0py = ykl - yk + tmp3 * yijykl;
-        const double c0pz = zkl - zk + tmp3 * zijzkl;
+        const double c0px = tmp3 * xijxkl;
+        const double c0py = tmp3 * yijykl;
+        const double c0pz = tmp3 * zijzkl;
         const double g_0 = 1;
         const double g_1 = c0px;
         const double g_2 = 1;

diff --git a/gpu4pyscf/lib/gint/g3c2e_ip1ip2.cu b/gpu4pyscf/lib/gint/g3c2e_ip1ip2.cu
@@ -165,12 +165,11 @@ static void GINTfill_int3c2e_ip1ip2_kernel000(GINTEnvVars envs, ERITensor eri, B
     double gzx = 0;
     double gzy = 0;
     double gzz = 0;
+
     const double xi = bas_x[ish];
     const double yi = bas_y[ish];
     const double zi = bas_z[ish];
-    const double xk = bas_x[ksh];
-    const double yk = bas_y[ksh];
-    const double zk = bas_z[ksh];
+
     const int prim_ij0 = prim_ij;
     const int prim_ij1 = prim_ij + nprim_ij;
     const int prim_kl0 = prim_kl;
@@ -217,9 +216,9 @@ static void GINTfill_int3c2e_ip1ip2_kernel000(GINTEnvVars envs, ERITensor eri, B
             const double c00y = yij - yi - tmp2 * yijykl;
             const double c00z = zij - zi - tmp2 * zijzkl;
             const double tmp3 = tmp1 * aij;
-            const double c0px = xkl - xk + tmp3 * xijxkl;
-            const double c0py = ykl - yk + tmp3 * yijykl;
-            const double c0pz = zkl - zk + tmp3 * zijzkl;
+            const double c0px = tmp3 * xijxkl;
+            const double c0py = tmp3 * yijykl;
+            const double c0pz = tmp3 * zijzkl;
             const double g_0 = 1;
             const double g_1 = c00x;
             const double g_2 = c0px;

diff --git a/gpu4pyscf/lib/gint/g3c2e_ip2.cu b/gpu4pyscf/lib/gint/g3c2e_ip2.cu
@@ -138,17 +138,10 @@ static void GINTfill_int3c2e_ip2_kernel000(GINTEnvVars envs, ERITensor eri, Basi
     double* __restrict__ z12 = c_bpcache.z12;
     double* __restrict__ a1 = c_bpcache.a1;
 
-    const int nbas = c_bpcache.nbas;
-    double* __restrict__ bas_x = c_bpcache.bas_coords;
-    double* __restrict__ bas_y = bas_x + nbas;
-    double* __restrict__ bas_z = bas_y + nbas;
-
     double gout0 = 0;
     double gout1 = 0;
     double gout2 = 0;
-    const double xk = bas_x[ksh];
-    const double yk = bas_y[ksh];
-    const double zk = bas_z[ksh];
+
     const int prim_ij0 = prim_ij;
     const int prim_ij1 = prim_ij + nprim_ij;
     const int prim_kl0 = prim_kl;
@@ -191,13 +184,11 @@ static void GINTfill_int3c2e_ip2_kernel000(GINTEnvVars envs, ERITensor eri, Basi
         }
         root0 /= root0 + 1 - root0 * theta;
         const double u2 = a0 * root0;
-        const double tmp4 = .5 / (u2 * aijkl + a1);
-        const double b00 = u2 * tmp4;
-        const double tmp1 = 2 * b00;
+        const double tmp1 = u2 / (u2 * aijkl + a1);
         const double tmp3 = tmp1 * aij;
-        const double c0px = xkl - xk + tmp3 * xijxkl;
-        const double c0py = ykl - yk + tmp3 * yijykl;
-        const double c0pz = zkl - zk + tmp3 * zijzkl;
+        const double c0px = tmp3 * xijxkl;
+        const double c0py = tmp3 * yijykl;
+        const double c0pz = tmp3 * zijzkl;
         const double g_0 = 1;
         const double g_1 = c0px;
         const double g_2 = 1;

diff --git a/gpu4pyscf/lib/gint/g3c2e_ipip2.cu b/gpu4pyscf/lib/gint/g3c2e_ipip2.cu
@@ -147,11 +147,6 @@ static void GINTfill_int3c2e_ipip2_kernel000(GINTEnvVars envs, ERITensor eri, Ba
     double* __restrict__ x12 = c_bpcache.x12;
     double* __restrict__ y12 = c_bpcache.y12;
     double* __restrict__ z12 = c_bpcache.z12;
-
-    const int nbas = c_bpcache.nbas;
-    double* __restrict__ bas_x = c_bpcache.bas_coords;
-    double* __restrict__ bas_y = bas_x + nbas;
-    double* __restrict__ bas_z = bas_y + nbas;
 
     double gxx = 0;
     double gxy = 0;
@@ -162,9 +157,7 @@ static void GINTfill_int3c2e_ipip2_kernel000(GINTEnvVars envs, ERITensor eri, Ba
     double gzx = 0;
     double gzy = 0;
     double gzz = 0;
-    double xk = bas_x[ksh];
-    double yk = bas_y[ksh];
-    double zk = bas_z[ksh];
+
     const int prim_ij0 = prim_ij;
     const int prim_ij1 = prim_ij + nprim_ij;
     const int prim_kl0 = prim_kl;
@@ -208,9 +201,9 @@ static void GINTfill_int3c2e_ipip2_kernel000(GINTEnvVars envs, ERITensor eri, Ba
             const double tmp1 = 2 * b00;
             const double tmp3 = tmp1 * aij;
             const double b01 = b00 + tmp4 * aij;
-            const double c0px = xkl - xk + tmp3 * xijxkl;
-            const double c0py = ykl - yk + tmp3 * yijykl;
-            const double c0pz = zkl - zk + tmp3 * zijzkl;
+            const double c0px = tmp3 * xijxkl;
+            const double c0py = tmp3 * yijykl;
+            const double c0pz = tmp3 * zijzkl;
             const double g_0 = 1;
             const double g_1 = c0px;
             const double g_2 = c0px * c0px + b01;

diff --git a/gpu4pyscf/lib/gvhf/g3c2e_ip2.cu b/gpu4pyscf/lib/gvhf/g3c2e_ip2.cu
@@ -163,17 +163,11 @@ static void GINTint3c2e_ip2_jk_kernel001(GINTEnvVars envs, JKMatrix jk, BasisPro
     double* __restrict__ a1 = c_bpcache.a1;
     int ij, kl;
     int prim_ij0, prim_ij1, prim_kl0, prim_kl1;
-    int nbas = c_bpcache.nbas;
-    double* __restrict__ bas_x = c_bpcache.bas_coords;
-    double* __restrict__ bas_y = bas_x + nbas;
-    double* __restrict__ bas_z = bas_y + nbas;
 
     double gout0 = 0;
     double gout1 = 0;
     double gout2 = 0;
-    double xk = bas_x[ksh];
-    double yk = bas_y[ksh];
-    double zk = bas_z[ksh];
+
     prim_ij0 = prim_ij;
     prim_ij1 = prim_ij + nprim_ij;
     prim_kl0 = prim_kl;
@@ -216,13 +210,11 @@ static void GINTint3c2e_ip2_jk_kernel001(GINTEnvVars envs, JKMatrix jk, BasisPro
         }
         root0 /= root0 + 1 - root0 * theta;
         double u2 = a0 * root0;
-        double tmp4 = .5 / (u2 * aijkl + a1);
-        double b00 = u2 * tmp4;
-        double tmp1 = 2 * b00;
+        double tmp1 = u2 / (u2 * aijkl + a1);
         double tmp3 = tmp1 * aij;
-        double c0px = xkl - xk + tmp3 * xijxkl;
-        double c0py = ykl - yk + tmp3 * yijykl;
-        double c0pz = zkl - zk + tmp3 * zijzkl;
+        double c0px = tmp3 * xijxkl;
+        double c0py = tmp3 * yijykl;
+        double c0pz = tmp3 * zijzkl;
         double g_0 = 1;
         double g_1 = c0px;
         double g_2 = 1;

diff --git a/gpu4pyscf/lib/gvhf/g3c2e_pass1.cu b/gpu4pyscf/lib/gvhf/g3c2e_pass1.cu
@@ -155,17 +155,11 @@ static void GINTint3c2e_pass1_j_kernel0010(GINTEnvVars envs, JKMatrix jk, BasisP
     }
     int ij, kl;
     int prim_ij0, prim_ij1, prim_kl0, prim_kl1;
-    int nbas = c_bpcache.nbas;
-    double* __restrict__ bas_x = c_bpcache.bas_coords;
-    double* __restrict__ bas_y = bas_x + nbas;
-    double* __restrict__ bas_z = bas_y + nbas;
 
     double gout0 = 0;
     double gout1 = 0;
     double gout2 = 0;
-    double xk = bas_x[ksh];
-    double yk = bas_y[ksh];
-    double zk = bas_z[ksh];
+
     prim_ij0 = prim_ij;
     prim_ij1 = prim_ij + nprim_ij;
     prim_kl0 = prim_kl;
@@ -204,13 +198,11 @@ static void GINTint3c2e_pass1_j_kernel0010(GINTEnvVars envs, JKMatrix jk, BasisP
             root0 = fmt1 / (fmt0 - fmt1);
         }
         double u2 = a0 * root0;
-        double tmp4 = .5 / (u2 * aijkl + a1);
-        double b00 = u2 * tmp4;
-        double tmp1 = 2 * b00;
+        double tmp1 = u2 / (u2 * aijkl + a1);
         double tmp3 = tmp1 * aij;
-        double c0px = xkl - xk + tmp3 * xijxkl;
-        double c0py = ykl - yk + tmp3 * yijykl;
-        double c0pz = zkl - zk + tmp3 * zijzkl;
+        double c0px = tmp3 * xijxkl;
+        double c0py = tmp3 * yijykl;
+        double c0pz = tmp3 * zijzkl;
         double g_0 = 1;
         double g_1 = c0px;
         double g_2 = 1;

diff --git a/gpu4pyscf/lib/gvhf/g3c2e_pass2.cu b/gpu4pyscf/lib/gvhf/g3c2e_pass2.cu
@@ -152,19 +152,14 @@ static void GINTint3c2e_pass2_j_kernel0010(GINTEnvVars envs, JKMatrix jk, BasisP
     double* __restrict__ z12 = c_bpcache.z12;
     int ij, kl;
     int prim_ij0, prim_ij1, prim_kl0, prim_kl1;
-    int nbas = c_bpcache.nbas;
-    double* __restrict__ bas_x = c_bpcache.bas_coords;
-    double* __restrict__ bas_y = bas_x + nbas;
-    double* __restrict__ bas_z = bas_y + nbas;
+
     if (ish == jsh){
         norm *= .5;
     }
     double gout0 = 0;
     double gout1 = 0;
     double gout2 = 0;
-    double xk = bas_x[ksh];
-    double yk = bas_y[ksh];
-    double zk = bas_z[ksh];
+
     prim_ij0 = prim_ij;
     prim_ij1 = prim_ij + nprim_ij;
     prim_kl0 = prim_kl;
@@ -203,13 +198,11 @@ static void GINTint3c2e_pass2_j_kernel0010(GINTEnvVars envs, JKMatrix jk, BasisP
             root0 = fmt1 / (fmt0 - fmt1);
         }
         double u2 = a0 * root0;
-        double tmp4 = .5 / (u2 * aijkl + a1);
-        double b00 = u2 * tmp4;
-        double tmp1 = 2 * b00;
+        double tmp1 = u2 / (u2 * aijkl + a1);
         double tmp3 = tmp1 * aij;
-        double c0px = xkl - xk + tmp3 * xijxkl;
-        double c0py = ykl - yk + tmp3 * yijykl;
-        double c0pz = zkl - zk + tmp3 * zijzkl;
+        double c0px = tmp3 * xijxkl;
+        double c0py = tmp3 * yijykl;
+        double c0pz = tmp3 * zijzkl;
         double g_0 = 1;
         double g_1 = c0px;
         double g_2 = 1;

diff --git a/setup.py b/setup.py
@@ -140,7 +140,7 @@ def initialize_with_default_plat_name(self):
     install_requires=[
         'pyscf~=2.7.0',
         'pyscf-dispersion',
-        f'cupy-cuda{CUDA_VERSION}',
+        f'cupy-cuda{CUDA_VERSION}>=13.0', # Due to expm in cupyx.scipy.linalg and cutensor 2.0
         'geometric',
         f'gpu4pyscf-libxc-cuda{CUDA_VERSION}>=0.5',
     ]