optimize for

Author: Qianruipku

source/module_hamilt_pw/hamilt_pwdft/kernels/cuda/ekinetic_op.cu

@@ -21,22 +21,23 @@ __global__ void ekinetic_pw(
 {
   const int block_idx = blockIdx.x;
   const int thread_idx = threadIdx.x;
+  const int start_idx = block_idx * max_npw;
   if(is_first_node)
   {
       for (int ii = thread_idx; ii < npw; ii += blockDim.x)
       {
-          hpsi[block_idx * max_npw + ii] = gk2[ii] * tpiba2 * psi[block_idx * max_npw + ii];
+          hpsi[start_idx + ii] = gk2[ii] * tpiba2 * psi[start_idx + ii];
       }
       for (int ii = npw + thread_idx; ii < max_npw; ii += blockDim.x)
       {
-          hpsi[block_idx * max_npw + ii] = 0.0;
+          hpsi[start_idx + ii] = 0.0;
       }
   }
   else
   {
       for (int ii = thread_idx; ii < npw; ii += blockDim.x)
       {
-          hpsi[block_idx * max_npw + ii] += gk2[ii] * tpiba2 * psi[block_idx * max_npw + ii];
+          hpsi[start_idx + ii] += gk2[ii] * tpiba2 * psi[start_idx + ii];
       }
   }
 }

source/module_hamilt_pw/hamilt_pwdft/kernels/rocm/ekinetic_op.hip.cu

@@ -21,22 +21,23 @@ __global__ void ekinetic_pw(
 {
   const int block_idx = blockIdx.x;
   const int thread_idx = threadIdx.x;
+  const int start_idx = block_idx * max_npw;
   if(is_first_node)
   {
       for (int ii = thread_idx; ii < npw; ii += blockDim.x)
       {
-          hpsi[block_idx * max_npw + ii] = gk2[ii] * tpiba2 * psi[block_idx * max_npw + ii];
+          hpsi[start_idx + ii] = gk2[ii] * tpiba2 * psi[start_idx + ii];
       }
       for (int ii = npw + thread_idx; ii < max_npw; ii += blockDim.x)
       {
-          hpsi[block_idx * max_npw + ii] = 0.0;
+          hpsi[start_idx + ii] = 0.0;
       }
   }
   else
   {
       for (int ii = thread_idx; ii < npw; ii += blockDim.x)
       {
-          hpsi[block_idx * max_npw + ii] += gk2[ii] * tpiba2 * psi[block_idx * max_npw + ii];
+          hpsi[start_idx + ii] += gk2[ii] * tpiba2 * psi[start_idx + ii];
       }
   }
 

source/module_hamilt_pw/hamilt_stodft/hamilt_sdft_pw.cpp

@@ -59,9 +59,10 @@ void HamiltSdftPW<T, Device>::hPsi_norm(const T* psi_in, T* hpsi_norm, const int
     {
         for (int ig = 0; ig < npwk; ++ig)
         {
-            hpsi_norm[ib * npwk_max + ig]
-                = (hpsi_norm[ib * npwk_max + ig] - Ebar * psi_in[ib * npwk_max + ig]) / DeltaE;
+            hpsi_norm[ig] = (hpsi_norm[ig] - Ebar * psi_in[ig]) / DeltaE;
         }
+        hpsi_norm += npwk_max;
+        psi_in += npwk_max;
     }
     ModuleBase::timer::tick("HamiltSdftPW", "hPsi_norm");
 }