optimize hPsi

Author: Qianruipku

source/module_hamilt_general/operator.cpp

@@ -60,7 +60,7 @@ typename Operator<T, Device>::hpsi_info Operator<T, Device>::hPsi(hpsi_info& inp
         this->hpsi = new psi::Psi<T, Device>(hpsi_pointer, *psi_input, 1, nbands / psi_input->npol);
     }
 
-    auto call_act = [&, this](const Operator* op) -> void {
+    auto call_act = [&, this](const Operator* op, const bool& is_first_node) -> void {
         // a "psi" with the bands of needed range
         psi::Psi<T, Device> psi_wrapper(const_cast<T*>(tmpsi_in), 1, nbands, psi_input->get_nbasis());
         switch (op->get_act_type())
@@ -69,17 +69,17 @@ typename Operator<T, Device>::hpsi_info Operator<T, Device>::hPsi(hpsi_info& inp
             op->act(psi_wrapper, *this->hpsi, nbands);
             break;
         default:
-            op->act(nbands, psi_input->get_nbasis(), psi_input->npol, tmpsi_in, this->hpsi->get_pointer(), psi_input->get_ngk(op->ik));
+            op->act(nbands, psi_input->get_nbasis(), psi_input->npol, tmpsi_in, this->hpsi->get_pointer(), psi_input->get_ngk(op->ik), is_first_node);
             break;
         }
         };
 
     ModuleBase::timer::tick("Operator", "hPsi");
-    call_act(this);
+    call_act(this, true); // first node
     Operator* node((Operator*)this->next_op);
     while (node != nullptr)
     {
-        call_act(node);
+        call_act(node, false); // other nodes
         node = (Operator*)(node->next_op);
     }
     ModuleBase::timer::tick("Operator", "hPsi");
@@ -162,7 +162,7 @@ T* Operator<T, Device>::get_hpsi(const hpsi_info& info) const
     size_t total_hpsi_size = nbands_range * this->hpsi->get_nbasis();
     // ModuleBase::GlobalFunc::ZEROS(hpsi_pointer, total_hpsi_size);
     // denghui replaced at 20221104
-    set_memory_op()(this->ctx, hpsi_pointer, 0, total_hpsi_size);
+    // set_memory_op()(this->ctx, hpsi_pointer, 0, total_hpsi_size);
     return hpsi_pointer;
 }
 

source/module_hamilt_pw/hamilt_pwdft/kernels/ekinetic_op.cpp

@@ -17,28 +17,32 @@ struct ekinetic_pw_op<FPTYPE, base_device::DEVICE_CPU>
     {
         if (is_first_node)
         {
-#ifdef _OPENMP
-#pragma omp parallel for collapse(2) schedule(static, 4096 / sizeof(FPTYPE))
-#endif
             for (int ib = 0; ib < nband; ++ib)
             {
+#ifdef _OPENMP
+#pragma omp parallel for
+#endif
                 for (int ig = 0; ig < npw; ++ig)
                 {
-                    tmhpsi[ib * max_npw + ig] = gk2_ik[ig] * tpiba2 * tmpsi_in[ib * max_npw + ig];
+                    tmhpsi[ig] = gk2_ik[ig] * tpiba2 * tmpsi_in[ig];
                 }
+                tmpsi_in += max_npw;
+                tmhpsi += max_npw;
             }
         }
         else
         {
-#ifdef _OPENMP
-#pragma omp parallel for collapse(2) schedule(static, 4096 / sizeof(FPTYPE))
-#endif
             for (int ib = 0; ib < nband; ++ib)
             {
+#ifdef _OPENMP
+#pragma omp parallel for
+#endif
                 for (int ig = 0; ig < npw; ++ig)
                 {
-                    tmhpsi[ib * max_npw + ig] += gk2_ik[ig] * tpiba2 * tmpsi_in[ib * max_npw + ig];
+                    tmhpsi[ig] += gk2_ik[ig] * tpiba2 * tmpsi_in[ig];
                 }
+                tmpsi_in += max_npw;
+                tmhpsi += max_npw;
             }
         }
     }

source/module_hamilt_pw/hamilt_pwdft/operator_pw/meta_pw.cpp

@@ -52,6 +52,10 @@ void Meta<OperatorPW<T, Device>>::act(
     }
 
     ModuleBase::timer::tick("Operator", "MetaPW");
+    if(is_first_node)
+    {
+        setmem_complex_op()(this->ctx, tmhpsi, 0, nbasis*nbands);
+    }
 
     const int current_spin = this->isk[this->ik];
     int max_npw = nbasis / npol;

source/module_hamilt_pw/hamilt_pwdft/operator_pw/meta_pw.h

@@ -84,6 +84,7 @@ class Meta<OperatorPW<T, Device>> : public OperatorPW<T, Device>
     using vector_mul_vector_op = hsolver::vector_mul_vector_op<T, Device>;
     using resmem_complex_op = base_device::memory::resize_memory_op<T, Device>;
     using delmem_complex_op = base_device::memory::delete_memory_op<T, Device>;
+    using setmem_complex_op = base_device::memory::set_memory_op<T, Device>;
 };
 
 } // namespace hamilt

source/module_hamilt_pw/hamilt_pwdft/operator_pw/nonlocal_pw.cpp

@@ -218,6 +218,10 @@ void Nonlocal<OperatorPW<T, Device>>::act(
     const bool is_first_node)const
 {
     ModuleBase::timer::tick("Operator", "NonlocalPW");
+    if(is_first_node)
+    {
+        setmem_complex_op()(this->ctx, tmhpsi, 0, nbasis*nbands);
+    }
     if(!PARAM.inp.use_paw)
     {
         this->npw = ngk_ik;

source/module_hamilt_pw/hamilt_pwdft/operator_pw/veff_pw.cpp

@@ -45,6 +45,10 @@ void Veff<OperatorPW<T, Device>>::act(
     const bool is_first_node)const
 {
     ModuleBase::timer::tick("Operator", "VeffPW");
+    if(is_first_node)
+    {
+        setmem_complex_op()(this->ctx, tmhpsi, 0, nbasis*nbands);
+    }
 
     int max_npw = nbasis / npol;
     const int current_spin = this->isk[this->ik];

source/module_hamilt_pw/hamilt_pwdft/operator_pw/veff_pw.h

@@ -73,6 +73,7 @@ class Veff<OperatorPW<T, Device>> : public OperatorPW<T, Device>
 
     using resmem_complex_op = base_device::memory::resize_memory_op<T, Device>;
     using delmem_complex_op = base_device::memory::delete_memory_op<T, Device>;
+    using setmem_complex_op = base_device::memory::set_memory_op<T, Device>;
 };
 
 } // namespace hamilt