refactor(dftu_pw): merge CPU/GPU branches into template, drop unused iter param

The CPU and GPU branches of cal_occ_pw (formerly L32-110 and L113-188)
were near-verbatim duplicates, differing only in the device template
parameter of OnsiteProjector and Psi. Merge them into a single template
member function accumulate_occ_one_k<Device>, with explicit
instantiation controlling CPU always compiled and GPU only under
__CUDA/__ROCM.

Also drop the iter parameter from cal_occ_pw signature since it was
never used in the function body; update the call site in
setup_dftu_pw.cpp accordingly.

Files changed:
- source/source_pw/module_pwdft/dftu_base.h
    - drop const int iter param from cal_occ_pw declaration
    - add template member declaration accumulate_occ_one_k<Device>
- source/source_pw/module_pwdft/dftu_pw.cpp
    - drop iter param from cal_occ_pw implementation
    - replace ~80 lines of CPU/GPU duplicate code with two template calls
    - add accumulate_occ_one_k template definition and explicit
      instantiations at end of file
- source/source_pw/module_pwdft/setup_dftu_pw.cpp
    - drop iter argument at call site

Effects:
- ~80 lines of duplicate code removed from dftu_pw.cpp
- CPU/GPU inner loops (iat/ib/m1/m2 + nspin==4 Pauli block assembly)
  now maintained in a single template implementation
- psi::Psi visibility confirmed transitively via onsite_proj.h

Pending verification (build delegated to user):
- make module_pwdft incremental build
- DFT+U ctest (e.g. MODULE_DFTU label)
- GPU build path (__CUDA/__ROCM) explicit instantiation

Author: abacus_fixer

source/source_pw/module_pwdft/dftu_base.h

@@ -61,10 +61,9 @@ class Plus_U_Base
     void set_energy(const double &e) { energy_u = e; }
     void set_double_energy() { energy_u *= 2.0; }
 
-    /// interface for PW base
+    /// interface for PW basis
     /// calculate the local occupation number matrix for PW based wave functions
-    void cal_occ_pw(const int iter,
-                    const void* psi_in,
+    void cal_occ_pw(const void* psi_in,
                     const ModuleBase::matrix& wg_in,
                     const UnitCell& cell,
                     Charge_Mixing* p_chgmix,
@@ -178,6 +177,13 @@ class Plus_U_Base
     void mix_occ_mat(const UnitCell& ucell, const double& mixing_beta);
     void set_occ_mat(const UnitCell& ucell);
 
+    /// accumulate occ_mat from psi for all k-points (per-device template)
+    template <typename Device>
+    void accumulate_occ_one_k(const void* psi_in,
+                              const ModuleBase::matrix& wg_in,
+                              const UnitCell& cell,
+                              const int* isk);
+
     std::vector<std::complex<double>> eff_pot_pw;
     std::vector<int> eff_pot_pw_index;
     std::vector<double> uom_array;

source/source_pw/module_pwdft/dftu_pw.cpp

@@ -18,8 +18,7 @@
 /// nspin=4 (npol=2): spinor calculation;
 ///   occ_mat has a single matrix of size (2*tlp1) x (2*tlp1) per atom
 ///   storing all 4 Pauli blocks contiguously.
-void Plus_U_Base::cal_occ_pw(const int iter,
-            const void* psi_in,
+void Plus_U_Base::cal_occ_pw(const void* psi_in,
             const ModuleBase::matrix& wg_in,
             const UnitCell& cell,
             Charge_Mixing* p_chgmix,
@@ -31,160 +30,12 @@ void Plus_U_Base::cal_occ_pw(const int iter,
 
     if(this->device == "cpu")
     {
-        auto* onsite_p = projectors::OnsiteProjector<double, base_device::DEVICE_CPU>::get_instance();
-        const psi::Psi<std::complex<double>>* psi_p = (const psi::Psi<std::complex<double>>*)psi_in;
-        const int nbands = psi_p->get_nbands();
-        const int npol = psi_p->get_npol();
-        for(int ik = 0; ik < psi_p->get_nk(); ik++)
-        {
-            int is = (this->nspin == 2) ? isk[ik] : 0;
-            psi_p->fix_k(ik);
-            onsite_p->tabulate_atomic(ik);
-
-            onsite_p->overlap_proj_psi(nbands*npol, psi_p->get_pointer());
-            const std::complex<double>* becp = onsite_p->get_h_becp();
-            int nkb = onsite_p->get_size_becp() / nbands / npol;
-
-            int begin_ih = 0;
-            for(int iat = 0; iat < cell.nat; iat++)
-            {
-                const int it = cell.iat2it[iat];
-                const int nh = onsite_p->get_nh(iat);
-                const int target_l = get_orbital_corr(it);
-                if(!has_correlated_orbital(it))
-                {
-                    begin_ih += nh;
-                    continue;
-                }
-                const int m_begin = target_l * target_l;
-                const int tlp1 = 2 * target_l + 1;
-                const int tlp1_2 = tlp1 * tlp1;
-                if(this->nspin == 4)
-                {
-                    for(int ib = 0;ib<nbands;ib++)
-                    {
-                        const double weight = wg_in(ik, ib);
-                        int ind_m1m2 = 0;
-                        for(int m1 = 0; m1 < tlp1; m1++)
-                        {
-                            const int index_m1 = ib*npol*nkb + begin_ih + m_begin + m1;
-                            for(int m2 = 0; m2 < tlp1; m2++)
-                            {
-                                const int index_m2 = ib*npol*nkb + begin_ih + m_begin + m2;
-                                std::complex<double> occ[4];
-                                occ[0] = weight * conj(becp[index_m1]) * becp[index_m2];
-                                occ[1] = weight * conj(becp[index_m1]) * becp[index_m2 + nkb];
-                                occ[2] = weight * conj(becp[index_m1 + nkb]) * becp[index_m2];
-                                occ[3] = weight * conj(becp[index_m1 + nkb]) * becp[index_m2 + nkb];
-                                this->occ_mat[iat][target_l][0][0].c[ind_m1m2] += (occ[0] + occ[3]).real();
-                                this->occ_mat[iat][target_l][0][0].c[ind_m1m2 + tlp1_2] += (occ[1] + occ[2]).real();
-                                this->occ_mat[iat][target_l][0][0].c[ind_m1m2 + 2 * tlp1_2] += (occ[1] - occ[2]).imag();
-                                this->occ_mat[iat][target_l][0][0].c[ind_m1m2 + 3 * tlp1_2] += (occ[0] - occ[3]).real();
-                                ind_m1m2++;
-                            }
-                        }
-                    }// ib
-                }
-                else // nspin=1 or nspin=2
-                {
-                    for(int ib = 0;ib<nbands;ib++)
-                    {
-                        const double weight = wg_in(ik, ib);
-                        int ind_m1m2 = 0;
-                        for(int m1 = 0; m1 < tlp1; m1++)
-                        {
-                            const int index_m1 = ib*nkb + begin_ih + m_begin + m1;
-                            for(int m2 = 0; m2 < tlp1; m2++)
-                            {
-                                const int index_m2 = ib*nkb + begin_ih + m_begin + m2;
-                                this->occ_mat[iat][target_l][0][is].c[ind_m1m2] += weight * (conj(becp[index_m1]) * becp[index_m2]).real();
-                                ind_m1m2++;
-                            }
-                        }
-                    }// ib
-                }
-                begin_ih += nh;
-            }// iat
-
-        }// ik
+        this->accumulate_occ_one_k<base_device::DEVICE_CPU>(psi_in, wg_in, cell, isk);
     }
 #if defined(__CUDA) || defined(__ROCM)
     else
     {
-        auto* onsite_p = projectors::OnsiteProjector<double, base_device::DEVICE_GPU>::get_instance();
-        const psi::Psi<std::complex<double>, base_device::DEVICE_GPU>* psi_p = (const psi::Psi<std::complex<double>, base_device::DEVICE_GPU>*)psi_in;
-        const int nbands = psi_p->get_nbands();
-        const int npol = psi_p->get_npol();
-        for(int ik = 0; ik < psi_p->get_nk(); ik++)
-        {
-            int is = (this->nspin == 2) ? isk[ik] : 0;
-            psi_p->fix_k(ik);
-            onsite_p->tabulate_atomic(ik);
-
-            onsite_p->overlap_proj_psi(nbands*npol, psi_p->get_pointer());
-            const std::complex<double>* becp = onsite_p->get_h_becp();
-            int nkb = onsite_p->get_size_becp() / nbands / npol;
-            int begin_ih = 0;
-            for(int iat = 0; iat < cell.nat; iat++)
-            {
-                const int it = cell.iat2it[iat];
-                const int nh = onsite_p->get_nh(iat);
-                const int target_l = get_orbital_corr(it);
-                if(!has_correlated_orbital(it))
-                {
-                    begin_ih += nh;
-                    continue;
-                }
-                const int m_begin = target_l * target_l;
-                const int tlp1 = 2 * target_l + 1;
-                const int tlp1_2 = tlp1 * tlp1;
-                if(this->nspin == 4)
-                {
-                    for(int ib = 0;ib<nbands;ib++)
-                    {
-                        const double weight = wg_in(ik, ib);
-                        int ind_m1m2 = 0;
-                        for(int m1 = 0; m1 < tlp1; m1++)
-                        {
-                            const int index_m1 = ib*npol*nkb + begin_ih + m_begin + m1;
-                            for(int m2 = 0; m2 < tlp1; m2++)
-                            {
-                                const int index_m2 = ib*npol*nkb + begin_ih + m_begin + m2;
-                                std::complex<double> occ[4];
-                                occ[0] = weight * conj(becp[index_m1]) * becp[index_m2];
-                                occ[1] = weight * conj(becp[index_m1]) * becp[index_m2 + nkb];
-                                occ[2] = weight * conj(becp[index_m1 + nkb]) * becp[index_m2];
-                                occ[3] = weight * conj(becp[index_m1 + nkb]) * becp[index_m2 + nkb];
-                                this->occ_mat[iat][target_l][0][0].c[ind_m1m2] += (occ[0] + occ[3]).real();
-                                this->occ_mat[iat][target_l][0][0].c[ind_m1m2 + tlp1_2] += (occ[1] + occ[2]).real();
-                                this->occ_mat[iat][target_l][0][0].c[ind_m1m2 + 2 * tlp1_2] += (occ[1] - occ[2]).imag();
-                                this->occ_mat[iat][target_l][0][0].c[ind_m1m2 + 3 * tlp1_2] += (occ[0] - occ[3]).real();
-                                ind_m1m2++;
-                            }
-                        }
-                    }// ib
-                }
-                else // nspin=1 or nspin=2
-                {
-                    for(int ib = 0;ib<nbands;ib++)
-                    {
-                        const double weight = wg_in(ik, ib);
-                        int ind_m1m2 = 0;
-                        for(int m1 = 0; m1 < tlp1; m1++)
-                        {
-                            const int index_m1 = ib*nkb + begin_ih + m_begin + m1;
-                            for(int m2 = 0; m2 < tlp1; m2++)
-                            {
-                                const int index_m2 = ib*nkb + begin_ih + m_begin + m2;
-                                this->occ_mat[iat][target_l][0][is].c[ind_m1m2] += weight * (conj(becp[index_m1]) * becp[index_m2]).real();
-                                ind_m1m2++;
-                            }
-                        }
-                    }// ib
-                }
-                begin_ih += nh;
-            }// iat
-        }// ik
+        this->accumulate_occ_one_k<base_device::DEVICE_GPU>(psi_in, wg_in, cell, isk);
     }
 #endif
 
@@ -347,3 +198,95 @@ void Plus_U_Base::cal_occ_pw(const int iter,
 
     ModuleBase::timer::end("Plus_U_Base", "cal_occ_pw");
 }
+
+template <typename Device>
+void Plus_U_Base::accumulate_occ_one_k(const void* psi_in,
+                                       const ModuleBase::matrix& wg_in,
+                                       const UnitCell& cell,
+                                       const int* isk)
+{
+    auto* onsite_p = projectors::OnsiteProjector<double, Device>::get_instance();
+    const psi::Psi<std::complex<double>, Device>* psi_p =
+        (const psi::Psi<std::complex<double>, Device>*)psi_in;
+    const int nbands = psi_p->get_nbands();
+    const int npol = psi_p->get_npol();
+    for(int ik = 0; ik < psi_p->get_nk(); ik++)
+    {
+        int is = (this->nspin == 2) ? isk[ik] : 0;
+        psi_p->fix_k(ik);
+        onsite_p->tabulate_atomic(ik);
+
+        onsite_p->overlap_proj_psi(nbands*npol, psi_p->get_pointer());
+        const std::complex<double>* becp = onsite_p->get_h_becp();
+        int nkb = onsite_p->get_size_becp() / nbands / npol;
+
+        int begin_ih = 0;
+        for(int iat = 0; iat < cell.nat; iat++)
+        {
+            const int it = cell.iat2it[iat];
+            const int nh = onsite_p->get_nh(iat);
+            const int target_l = get_orbital_corr(it);
+            if(!has_correlated_orbital(it))
+            {
+                begin_ih += nh;
+                continue;
+            }
+            const int m_begin = target_l * target_l;
+            const int tlp1 = 2 * target_l + 1;
+            const int tlp1_2 = tlp1 * tlp1;
+            if(this->nspin == 4)
+            {
+                for(int ib = 0; ib < nbands; ib++)
+                {
+                    const double weight = wg_in(ik, ib);
+                    int ind_m1m2 = 0;
+                    for(int m1 = 0; m1 < tlp1; m1++)
+                    {
+                        const int index_m1 = ib*npol*nkb + begin_ih + m_begin + m1;
+                        for(int m2 = 0; m2 < tlp1; m2++)
+                        {
+                            const int index_m2 = ib*npol*nkb + begin_ih + m_begin + m2;
+                            std::complex<double> occ[4];
+                            occ[0] = weight * conj(becp[index_m1]) * becp[index_m2];
+                            occ[1] = weight * conj(becp[index_m1]) * becp[index_m2 + nkb];
+                            occ[2] = weight * conj(becp[index_m1 + nkb]) * becp[index_m2];
+                            occ[3] = weight * conj(becp[index_m1 + nkb]) * becp[index_m2 + nkb];
+                            this->occ_mat[iat][target_l][0][0].c[ind_m1m2] += (occ[0] + occ[3]).real();
+                            this->occ_mat[iat][target_l][0][0].c[ind_m1m2 + tlp1_2] += (occ[1] + occ[2]).real();
+                            this->occ_mat[iat][target_l][0][0].c[ind_m1m2 + 2 * tlp1_2] += (occ[1] - occ[2]).imag();
+                            this->occ_mat[iat][target_l][0][0].c[ind_m1m2 + 3 * tlp1_2] += (occ[0] - occ[3]).real();
+                            ind_m1m2++;
+                        }
+                    }
+                }
+            }
+            else // nspin=1 or nspin=2
+            {
+                for(int ib = 0; ib < nbands; ib++)
+                {
+                    const double weight = wg_in(ik, ib);
+                    int ind_m1m2 = 0;
+                    for(int m1 = 0; m1 < tlp1; m1++)
+                    {
+                        const int index_m1 = ib*nkb + begin_ih + m_begin + m1;
+                        for(int m2 = 0; m2 < tlp1; m2++)
+                        {
+                            const int index_m2 = ib*nkb + begin_ih + m_begin + m2;
+                            this->occ_mat[iat][target_l][0][is].c[ind_m1m2] += weight * (conj(becp[index_m1]) * becp[index_m2]).real();
+                            ind_m1m2++;
+                        }
+                    }
+                }
+            }
+            begin_ih += nh;
+        }
+    }
+}
+
+// explicit instantiations
+template void Plus_U_Base::accumulate_occ_one_k<base_device::DEVICE_CPU>(
+    const void*, const ModuleBase::matrix&, const UnitCell&, const int*);
+#if defined(__CUDA) || defined(__ROCM)
+template void Plus_U_Base::accumulate_occ_one_k<base_device::DEVICE_GPU>(
+    const void*, const ModuleBase::matrix&, const UnitCell&, const int*);
+#endif

source/source_pw/module_pwdft/setup_dftu_pw.cpp

@@ -27,7 +27,7 @@ void iter_init_dftu_pw(const int iter,
 
     if (dftu.get_occ_mat_ctrl() != 2)
     {
-        dftu.cal_occ_pw(iter, psi, wg, ucell, p_chgmix, isk);
+        dftu.cal_occ_pw(psi, wg, ucell, p_chgmix, isk);
     }
     dftu_io::output(dftu, ucell, PARAM.inp.out_chg[0], PARAM.globalv.global_out_dir, PARAM.inp.nspin, PARAM.globalv.npol);
 }