Fix: error on GPU

dyzheng · Sep 4, 2024 · 87d4614 · 87d4614
1 parent a8b6a9c
commit 87d4614
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 14 deletions.
diff --git a/source/module_hamilt_pw/hamilt_pwdft/kernels/cuda/onsite_op.cu b/source/module_hamilt_pw/hamilt_pwdft/kernels/cuda/onsite_op.cu
@@ -21,12 +21,15 @@ __global__ void onsite_op(const int& npm,
 {
     const int ib = blockIdx.x / tnp;
     const int ip = blockIdx.x % tnp;
-    int ib2 = ib * npol;
-    int iat = ip_iat[ip];
-    const int psind = ip * npm + ib2;
-    const int becpind = ib2 * tnp + ip;
-    ps[psind] += lambda_coeff[iat * 4] * becp[becpind] + lambda_coeff[iat * 4 + 2] * becp[becpind + tnp];
-    ps[psind + 1] += lambda_coeff[iat * 4 + 1] * becp[becpind] + lambda_coeff[iat * 4 + 3] * becp[becpind + tnp];
+    if(ib < npm/npol)
+    {
+        int ib2 = ib * npol;
+        int iat = ip_iat[ip];
+        const int psind = ip * npm + ib2;
+        const int becpind = ib2 * tnp + ip;
+        ps[psind] += lambda_coeff[iat * 4] * becp[becpind] + lambda_coeff[iat * 4 + 2] * becp[becpind + tnp];
+        ps[psind + 1] += lambda_coeff[iat * 4 + 1] * becp[becpind] + lambda_coeff[iat * 4 + 3] * becp[becpind + tnp];
+    }
 }
 
 template <typename FPTYPE>
@@ -44,7 +47,7 @@ __global__ void onsite_op(const int& npm,
     const int ib = blockIdx.x / tnp;
     const int ip = blockIdx.x % tnp;
     int m1 = ip_m[ip];
-    if (m1 >= 0)
+    if (m1 >= 0 && ib < npm/npol)
     {
         int ib2 = ib * npol;
         int iat = ip_iat[ip];

diff --git a/source/module_hamilt_pw/hamilt_pwdft/kernels/rocm/onsite_op.hip.cu b/source/module_hamilt_pw/hamilt_pwdft/kernels/rocm/onsite_op.hip.cu
@@ -21,12 +21,15 @@ __global__ void onsite_op(const int& npm,
 {
     const int ib = blockIdx.x / tnp;
     const int ip = blockIdx.x % tnp;
-    int ib2 = ib * npol;
-    int iat = ip_iat[ip];
-    const int psind = ip * npm + ib2;
-    const int becpind = ib2 * tnp + ip;
-    ps[psind] += lambda_coeff[iat * 4] * becp[becpind] + lambda_coeff[iat * 4 + 2] * becp[becpind + tnp];
-    ps[psind + 1] += lambda_coeff[iat * 4 + 1] * becp[becpind] + lambda_coeff[iat * 4 + 3] * becp[becpind + tnp];
+    if(ib < npm/npol)
+    {
+        int ib2 = ib * npol;
+        int iat = ip_iat[ip];
+        const int psind = ip * npm + ib2;
+        const int becpind = ib2 * tnp + ip;
+        ps[psind] += lambda_coeff[iat * 4] * becp[becpind] + lambda_coeff[iat * 4 + 2] * becp[becpind + tnp];
+        ps[psind + 1] += lambda_coeff[iat * 4 + 1] * becp[becpind] + lambda_coeff[iat * 4 + 3] * becp[becpind + tnp];
+    }
 }
 
 template <typename FPTYPE>
@@ -44,7 +47,7 @@ __global__ void onsite_op(const int& npm,
     const int ib = blockIdx.x / tnp;
     const int ip = blockIdx.x % tnp;
     int m1 = ip_m[ip];
-    if (m1 >= 0)
+    if (m1 >= 0 && ib < npm/npol)
     {
         int ib2 = ib * npol;
         int iat = ip_iat[ip];