RT-TDDFT GPU Acceleration: RT-TD now fully support GPU computation

source/module_esolver/esolver_ks_lcao_tddft.cpp

@@ -24,12 +24,12 @@
 #include "module_io/print_info.h"
 
 //-----HSolver ElecState Hamilt--------
+#include "module_elecstate/cal_ux.h"
 #include "module_elecstate/elecstate_lcao.h"
 #include "module_hamilt_lcao/hamilt_lcaodft/hamilt_lcao.h"
 #include "module_hsolver/hsolver_lcao.h"
 #include "module_parameter/parameter.h"
 #include "module_psi/psi.h"
-#include "module_elecstate/cal_ux.h"
 
 //-----force& stress-------------------
 #include "module_hamilt_lcao/hamilt_lcaodft/FORCE_STRESS.h"
@@ -290,7 +290,12 @@ void ESolver_KS_LCAO_TDDFT::after_scf(UnitCell& ucell, const int istep)
         {
             std::stringstream ss_dipole;
             ss_dipole << PARAM.globalv.global_out_dir << "SPIN" << is + 1 << "_DIPOLE";
-            ModuleIO::write_dipole(ucell,pelec->charge->rho_save[is], pelec->charge->rhopw, is, istep, ss_dipole.str());
+            ModuleIO::write_dipole(ucell,
+                                   pelec->charge->rho_save[is],
+                                   pelec->charge->rhopw,
+                                   is,
+                                   istep,
+                                   ss_dipole.str());
         }
     }
     if (TD_Velocity::out_current == true)

source/module_hamilt_lcao/module_tddft/bandenergy.cpp

@@ -1,12 +1,12 @@
 #include "bandenergy.h"
 
-#include <complex>
-#include <iostream>
-
 #include "evolve_elec.h"
 #include "module_base/lapack_connector.h"
 #include "module_base/scalapack_connector.h"
 
+#include <complex>
+#include <iostream>
+
 namespace module_tddft
 {
 #ifdef __MPI
@@ -133,14 +133,144 @@ void compute_ekb(const Parallel_Orbitals* pv,
                 }
             }
         } // loop ipcol
-    }     // loop iprow
+    } // loop iprow
     info = MPI_Allreduce(Eii, ekb, nband, MPI_DOUBLE, MPI_SUM, pv->comm());
 
     delete[] tmp1;
     delete[] Eij;
     delete[] Eii;
 }
 
+void compute_ekb_tensor(const Parallel_Orbitals* pv,
+                        const int nband,
+                        const int nlocal,
+                        const container::Tensor& Htmp,
+                        const container::Tensor& psi_k,
+                        container::Tensor& ekb)
+{
+    // Create Tensor objects for temporary data
+    container::Tensor tmp1(container::DataType::DT_COMPLEX_DOUBLE,
+                           container::DeviceType::CpuDevice,
+                           container::TensorShape({pv->nloc_wfc}));
+    tmp1.zero();
+
+    container::Tensor Eij(container::DataType::DT_COMPLEX_DOUBLE,
+                          container::DeviceType::CpuDevice,
+                          container::TensorShape({pv->nloc}));
+    Eij.zero();
+
+    // Perform matrix multiplication: tmp1 = Htmp * psi_k
+    ScalapackConnector::gemm('N',
+                             'N',
+                             nlocal,
+                             nband,
+                             nlocal,
+                             1.0,
+                             Htmp.data<std::complex<double>>(),
+                             1,
+                             1,
+                             pv->desc,
+                             psi_k.data<std::complex<double>>(),
+                             1,
+                             1,
+                             pv->desc_wfc,
+                             0.0,
+                             tmp1.data<std::complex<double>>(),
+                             1,
+                             1,
+                             pv->desc_wfc);
+
+    // Perform matrix multiplication: Eij = psi_k^dagger * tmp1
+    ScalapackConnector::gemm('C',
+                             'N',
+                             nband,
+                             nband,
+                             nlocal,
+                             1.0,
+                             psi_k.data<std::complex<double>>(),
+                             1,
+                             1,
+                             pv->desc_wfc,
+                             tmp1.data<std::complex<double>>(),
+                             1,
+                             1,
+                             pv->desc_wfc,
+                             0.0,
+                             Eij.data<std::complex<double>>(),
+                             1,
+                             1,
+                             pv->desc_Eij);
+
+    if (Evolve_elec::td_print_eij >= 0.0)
+    {
+        GlobalV::ofs_running
+            << "------------------------------------------------------------------------------------------------"
+            << std::endl;
+        GlobalV::ofs_running << " Eij:" << std::endl;
+        for (int i = 0; i < pv->nrow_bands; i++)
+        {
+            for (int j = 0; j < pv->ncol_bands; j++)
+            {
+                double aa, bb;
+                aa = Eij.data<std::complex<double>>()[i * pv->ncol + j].real();
+                bb = Eij.data<std::complex<double>>()[i * pv->ncol + j].imag();
+                if (std::abs(aa) < Evolve_elec::td_print_eij)
+                    aa = 0.0;
+                if (std::abs(bb) < Evolve_elec::td_print_eij)
+                    bb = 0.0;
+                if (aa > 0.0 || bb > 0.0)
+                {
+                    GlobalV::ofs_running << i << " " << j << " " << aa << "+" << bb << "i " << std::endl;
+                }
+            }
+        }
+        GlobalV::ofs_running << std::endl;
+        GlobalV::ofs_running
+            << "------------------------------------------------------------------------------------------------"
+            << std::endl;
+    }
+
+    int info;
+    int naroc[2];
+
+    // Create a Tensor for Eii
+    container::Tensor Eii(container::DataType::DT_DOUBLE,
+                          container::DeviceType::CpuDevice,
+                          container::TensorShape({nband}));
+    Eii.zero();
+
+    for (int iprow = 0; iprow < pv->dim0; ++iprow)
+    {
+        for (int ipcol = 0; ipcol < pv->dim1; ++ipcol)
+        {
+            if (iprow == pv->coord[0] && ipcol == pv->coord[1])
+            {
+                naroc[0] = pv->nrow;
+                naroc[1] = pv->ncol;
+                for (int j = 0; j < naroc[1]; ++j)
+                {
+                    int igcol = globalIndex(j, pv->nb, pv->dim1, ipcol);
+                    if (igcol >= nband)
+                        continue;
+                    for (int i = 0; i < naroc[0]; ++i)
+                    {
+                        int igrow = globalIndex(i, pv->nb, pv->dim0, iprow);
+                        if (igrow >= nband)
+                            continue;
+                        if (igcol == igrow)
+                        {
+                            Eii.data<double>()[igcol] = Eij.data<std::complex<double>>()[j * naroc[0] + i].real();
+                        }
+                    }
+                }
+            }
+        } // loop ipcol
+    } // loop iprow
+
+    // Perform MPI reduction to compute ekb
+    info = MPI_Allreduce(Eii.data<double>(), ekb.data<double>(), nband, MPI_DOUBLE, MPI_SUM, pv->comm());
+}
+
 #endif
 
 } // namespace module_tddft

source/module_hamilt_lcao/module_tddft/bandenergy.h

@@ -6,6 +6,7 @@
 #ifndef BANDENERGY_H
 #define BANDENERGY_H
 
+#include "module_base/module_container/ATen/core/tensor.h" // container::Tensor
 #include "module_basis/module_ao/parallel_orbitals.h"
 
 #include <complex>
@@ -29,6 +30,13 @@ void compute_ekb(const Parallel_Orbitals* pv,
                  const std::complex<double>* Htmp,
                  const std::complex<double>* psi_k,
                  double* ekb);
+
+void compute_ekb_tensor(const Parallel_Orbitals* pv,
+                        const int nband,
+                        const int nlocal,
+                        const container::Tensor& Htmp,
+                        const container::Tensor& psi_k,
+                        container::Tensor& ekb);
 #endif
 } // namespace module_tddft
 #endif

source/module_hamilt_lcao/module_tddft/evolve_elec.cpp

@@ -10,8 +10,8 @@
 
 namespace module_tddft
 {
-Evolve_elec::Evolve_elec(){};
-Evolve_elec::~Evolve_elec(){};
+Evolve_elec::Evolve_elec() {};
+Evolve_elec::~Evolve_elec() {};
 
 double Evolve_elec::td_force_dt;
 bool Evolve_elec::td_vext;
@@ -73,6 +73,67 @@ void Evolve_elec::solve_psi(const int& istep,
                        &(ekb(ik, 0)),
                        htype,
                        propagator);
+
+            const bool use_tensor = false;
+            if (use_tensor)
+            {
+                std::cout << "Print ekb: " << std::endl;
+                ekb.print(std::cout);
+                std::cout << "nband = " << nband << std::endl;
+                std::cout << "psi->get_nbands() = " << psi->get_nbands() << std::endl;
+                std::cout << "nlocal = " << nlocal << std::endl;
+                std::cout << "psi->get_nbasis() = " << psi->get_nbasis() << std::endl;
+                std::cout << "ekb.nr = " << ekb.nr << std::endl;
+                std::cout << "ekb.nc = " << ekb.nc << std::endl;
+
+                // Create TensorMap for psi_k, psi_k_laststep, H_laststep, S_laststep, ekb
+                container::TensorMap psi_k_tensor(psi[0].get_pointer(),
+                                                  container::DataType::DT_COMPLEX_DOUBLE,
+                                                  container::DeviceType::CpuDevice,
+                                                  container::TensorShape({psi->get_nbands(), psi->get_nbasis()}));
+                container::TensorMap psi_k_laststep_tensor(
+                    psi_laststep[0].get_pointer(),
+                    container::DataType::DT_COMPLEX_DOUBLE,
+                    container::DeviceType::CpuDevice,
+                    container::TensorShape({psi->get_nbands(), psi->get_nbasis()}));
+                container::TensorMap H_laststep_tensor(Hk_laststep[ik],
+                                                       container::DataType::DT_COMPLEX_DOUBLE,
+                                                       container::DeviceType::CpuDevice,
+                                                       container::TensorShape({para_orb.nloc}));
+                container::TensorMap S_laststep_tensor(Sk_laststep[ik],
+                                                       container::DataType::DT_COMPLEX_DOUBLE,
+                                                       container::DeviceType::CpuDevice,
+                                                       container::TensorShape({para_orb.nloc}));
+                container::TensorMap ekb_tensor(&(ekb(ik, 0)),
+                                                container::DataType::DT_DOUBLE,
+                                                container::DeviceType::CpuDevice,
+                                                container::TensorShape({nband}));
+
+                evolve_psi_tensor(nband,
+                                  nlocal,
+                                  &(para_orb),
+                                  phm,
+                                  psi_k_tensor,
+                                  psi_k_laststep_tensor,
+                                  H_laststep_tensor,
+                                  S_laststep_tensor,
+                                  ekb_tensor,
+                                  htype,
+                                  propagator);
+                // evolve_psi_tensor(nband,
+                //                   nlocal,
+                //                   &(para_orb),
+                //                   phm,
+                //                   psi[0].get_pointer(),
+                //                   psi_laststep[0].get_pointer(),
+                //                   Hk_laststep[ik],
+                //                   Sk_laststep[ik],
+                //                   &(ekb(ik, 0)),
+                //                   htype,
+                //                   propagator);
+                std::cout << "Print ekb tensor: " << std::endl;
+                ekb.print(std::cout);
+            }
         }
         else
         {

source/module_hamilt_lcao/module_tddft/evolve_elec.h

@@ -3,6 +3,8 @@
 
 #include "module_base/global_function.h"
 #include "module_base/global_variable.h"
+#include "module_base/module_container/ATen/core/tensor.h"     // container::Tensor
+#include "module_base/module_container/ATen/core/tensor_map.h" // TensorMap
 #include "module_esolver/esolver_ks_lcao.h"
 #include "module_esolver/esolver_ks_lcao_tddft.h"
 #include "module_hamilt_lcao/hamilt_lcaodft/hamilt_lcao.h"