fix: add npol==1 branch to GPU/DCU DFT+U onsite_ps_op kernels

The CUDA and ROCm kernels for the DFT+U onsite_ps_op only had npol==2
(nspin=4) logic. When npol==1 (nspin=1/2), the kernel would access
out-of-bounds memory (ps[psind+1], becp[becpind+tnp], and non-existent
vu matrix blocks). This mirrors the CPU fix in commit 65b738a.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: dyzheng

source/source_pw/module_pwdft/kernels/cuda/onsite_op.cu

@@ -48,26 +48,49 @@ __global__ void onsite_op(const int npm,
     int m1 = ip_m[ip];
     if (m1 >= 0)
     {
-        const int nbands = npm / npol;
-        for (int ib = threadIdx.x; ib < nbands; ib += blockDim.x)
+        if (npol == 2)
         {
-            int ib2 = ib * npol;
-            int iat = ip_iat[ip];
-            const thrust::complex<FPTYPE>* vu_iat = vu + vu_begin_iat[iat];
-            int orb_l = orb_l_iat[iat];
-            int tlp1 = 2 * orb_l + 1;
-            int tlp1_2 = tlp1 * tlp1;
-            int ip2_begin = ip - m1;
-            int ip2_end = ip - m1 + tlp1;
-            const int psind = ip * npm + ib2;
-            for (int ip2 = ip2_begin; ip2 < ip2_end; ip2++)
+            const int nbands = npm / npol;
+            for (int ib = threadIdx.x; ib < nbands; ib += blockDim.x)
             {
-                const int becpind = ib2 * tnp + ip2;
-                int m2 = ip_m[ip2];
-                const int index_mm = m1 * tlp1 + m2;
-                ps[psind] += vu_iat[index_mm] * becp[becpind] + vu_iat[index_mm + tlp1_2 * 2] * becp[becpind + tnp];
-                ps[psind + 1] += vu_iat[index_mm + tlp1_2 * 1] * becp[becpind]
-                                 + vu_iat[index_mm + tlp1_2 * 3] * becp[becpind + tnp];
+                int ib2 = ib * npol;
+                int iat = ip_iat[ip];
+                const thrust::complex<FPTYPE>* vu_iat = vu + vu_begin_iat[iat];
+                int orb_l = orb_l_iat[iat];
+                int tlp1 = 2 * orb_l + 1;
+                int tlp1_2 = tlp1 * tlp1;
+                int ip2_begin = ip - m1;
+                int ip2_end = ip - m1 + tlp1;
+                const int psind = ip * npm + ib2;
+                for (int ip2 = ip2_begin; ip2 < ip2_end; ip2++)
+                {
+                    const int becpind = ib2 * tnp + ip2;
+                    int m2 = ip_m[ip2];
+                    const int index_mm = m1 * tlp1 + m2;
+                    ps[psind] += vu_iat[index_mm] * becp[becpind] + vu_iat[index_mm + tlp1_2 * 2] * becp[becpind + tnp];
+                    ps[psind + 1] += vu_iat[index_mm + tlp1_2 * 1] * becp[becpind]
+                                     + vu_iat[index_mm + tlp1_2 * 3] * becp[becpind + tnp];
+                }
+            }
+        }
+        else // npol == 1, nspin=1 or nspin=2
+        {
+            for (int ib = threadIdx.x; ib < npm; ib += blockDim.x)
+            {
+                int iat = ip_iat[ip];
+                const thrust::complex<FPTYPE>* vu_iat = vu + vu_begin_iat[iat];
+                int orb_l = orb_l_iat[iat];
+                int tlp1 = 2 * orb_l + 1;
+                int ip2_begin = ip - m1;
+                int ip2_end = ip - m1 + tlp1;
+                const int psind = ip * npm + ib;
+                for (int ip2 = ip2_begin; ip2 < ip2_end; ip2++)
+                {
+                    const int becpind = ib * tnp + ip2;
+                    int m2 = ip_m[ip2];
+                    const int index_mm = m1 * tlp1 + m2;
+                    ps[psind] += vu_iat[index_mm] * becp[becpind];
+                }
             }
         }
     }

source/source_pw/module_pwdft/kernels/rocm/onsite_op.hip.cu

@@ -48,26 +48,49 @@ __global__ void onsite_op(const int npm,
     int m1 = ip_m[ip];
     if (m1 >= 0)
     {
-        const int nbands = npm / npol;
-        for (int ib = threadIdx.x; ib < nbands; ib += blockDim.x)
+        if (npol == 2)
         {
-            int ib2 = ib * npol;
-            int iat = ip_iat[ip];
-            const thrust::complex<FPTYPE>* vu_iat = vu + vu_begin_iat[iat];
-            int orb_l = orb_l_iat[iat];
-            int tlp1 = 2 * orb_l + 1;
-            int tlp1_2 = tlp1 * tlp1;
-            int ip2_begin = ip - m1;
-            int ip2_end = ip - m1 + tlp1;
-            const int psind = ip * npm + ib2;
-            for (int ip2 = ip2_begin; ip2 < ip2_end; ip2++)
+            const int nbands = npm / npol;
+            for (int ib = threadIdx.x; ib < nbands; ib += blockDim.x)
             {
-                const int becpind = ib2 * tnp + ip2;
-                int m2 = ip_m[ip2];
-                const int index_mm = m1 * tlp1 + m2;
-                ps[psind] += vu_iat[index_mm] * becp[becpind] + vu_iat[index_mm + tlp1_2 * 2] * becp[becpind + tnp];
-                ps[psind + 1] += vu_iat[index_mm + tlp1_2 * 1] * becp[becpind]
-                                 + vu_iat[index_mm + tlp1_2 * 3] * becp[becpind + tnp];
+                int ib2 = ib * npol;
+                int iat = ip_iat[ip];
+                const thrust::complex<FPTYPE>* vu_iat = vu + vu_begin_iat[iat];
+                int orb_l = orb_l_iat[iat];
+                int tlp1 = 2 * orb_l + 1;
+                int tlp1_2 = tlp1 * tlp1;
+                int ip2_begin = ip - m1;
+                int ip2_end = ip - m1 + tlp1;
+                const int psind = ip * npm + ib2;
+                for (int ip2 = ip2_begin; ip2 < ip2_end; ip2++)
+                {
+                    const int becpind = ib2 * tnp + ip2;
+                    int m2 = ip_m[ip2];
+                    const int index_mm = m1 * tlp1 + m2;
+                    ps[psind] += vu_iat[index_mm] * becp[becpind] + vu_iat[index_mm + tlp1_2 * 2] * becp[becpind + tnp];
+                    ps[psind + 1] += vu_iat[index_mm + tlp1_2 * 1] * becp[becpind]
+                                     + vu_iat[index_mm + tlp1_2 * 3] * becp[becpind + tnp];
+                }
+            }
+        }
+        else // npol == 1, nspin=1 or nspin=2
+        {
+            for (int ib = threadIdx.x; ib < npm; ib += blockDim.x)
+            {
+                int iat = ip_iat[ip];
+                const thrust::complex<FPTYPE>* vu_iat = vu + vu_begin_iat[iat];
+                int orb_l = orb_l_iat[iat];
+                int tlp1 = 2 * orb_l + 1;
+                int ip2_begin = ip - m1;
+                int ip2_end = ip - m1 + tlp1;
+                const int psind = ip * npm + ib;
+                for (int ip2 = ip2_begin; ip2 < ip2_end; ip2++)
+                {
+                    const int becpind = ib * tnp + ip2;
+                    int m2 = ip_m[ip2];
+                    const int index_mm = m1 * tlp1 + m2;
+                    ps[psind] += vu_iat[index_mm] * becp[becpind];
+                }
             }
         }
     }