Fix: Fix the errors in building abacus with libtorch-gpu (#6554)

* Fix: Fix the errors in building abacus with libtorch-gpu

* Fix: Mark submodule directory as safe for CI workflow

* Fix: Update CI workflow to manually handle submodules and set safe directories

* Fix: Update CI workflow to handle submodule ownership and initialization

* Fix: Add conditional compilation for CUDA and ROCM in MLKEDF descriptor functions

* Fix: Refactor Psi constructor to separate declaration and implementation for better readability and maintainability

* Fix: Improve code readability by adding colons to constructor comments and formatting memory operations

Author: sunliang98

.github/workflows/test.yml

@@ -17,11 +17,17 @@ jobs:
       volumes:
         - /tmp/ccache:/github/home/.ccache
     steps:
-      - name: Checkout
+      - name: Checkout repository
         uses: actions/checkout@v5
         with:
-          submodules: recursive
           fetch-depth: 0
+          # We will handle submodules manually after fixing ownership
+          submodules: 'false'
+
+      - name: Take ownership of the workspace and update submodules
+        run: |
+          sudo chown -R $(whoami) .
+          git submodule update --init --recursive
 
       - name: Install CI tools
         run: |

source/source_io/write_mlkedf_descriptors.cpp

@@ -67,6 +67,53 @@ void Write_MLKEDF_Descriptors::generateTrainData_KS(
     delete ptempRho;
 }
 
+void Write_MLKEDF_Descriptors::generateTrainData_KS(
+    const std::string& out_dir,
+    psi::Psi<std::complex<float>> *psi,
+    elecstate::ElecState *pelec,
+    ModulePW::PW_Basis_K *pw_psi,
+    ModulePW::PW_Basis *pw_rho,
+    UnitCell& ucell,
+    const double* veff
+)
+{
+    psi::Psi<std::complex<double>, base_device::DEVICE_CPU> psi_double(*psi);
+
+    this->generateTrainData_KS(out_dir, &psi_double, pelec, pw_psi, pw_rho, ucell, veff);
+}
+
+#if ((defined __CUDA) || (defined __ROCM))
+void Write_MLKEDF_Descriptors::generateTrainData_KS(
+    const std::string& out_dir,
+    psi::Psi<std::complex<double>, base_device::DEVICE_GPU>* psi,
+    elecstate::ElecState *pelec,
+    ModulePW::PW_Basis_K *pw_psi,
+    ModulePW::PW_Basis *pw_rho,
+    UnitCell& ucell,
+    const double* veff
+)
+{
+    psi::Psi<std::complex<double>, base_device::DEVICE_CPU> psi_cpu(*psi);
+
+    this->generateTrainData_KS(out_dir, &psi_cpu, pelec, pw_psi, pw_rho, ucell, veff);
+}
+
+void Write_MLKEDF_Descriptors::generateTrainData_KS(
+    const std::string& dir,
+    psi::Psi<std::complex<float>, base_device::DEVICE_GPU>* psi,
+    elecstate::ElecState *pelec,
+    ModulePW::PW_Basis_K *pw_psi,
+    ModulePW::PW_Basis *pw_rho,
+    UnitCell& ucell,
+    const double *veff
+)
+{
+    psi::Psi<std::complex<double>, base_device::DEVICE_CPU> psi_cpu_double(*psi);
+
+    this->generateTrainData_KS(dir, &psi_cpu_double, pelec, pw_psi, pw_rho, ucell, veff);
+}
+#endif
+
 void Write_MLKEDF_Descriptors::generate_descriptor(
     const std::string& out_dir,
     const double * const *prho, 

source/source_io/write_mlkedf_descriptors.h

@@ -40,7 +40,29 @@ class Write_MLKEDF_Descriptors
         ModulePW::PW_Basis *pw_rho,
         UnitCell& ucell,
         const double *veff
-    ){} // a mock function
+    );
+
+#if ((defined __CUDA) || (defined __ROCM))
+    void generateTrainData_KS(
+        const std::string& dir,
+        psi::Psi<std::complex<double>, base_device::DEVICE_GPU>* psi,
+        elecstate::ElecState *pelec,
+        ModulePW::PW_Basis_K *pw_psi,
+        ModulePW::PW_Basis *pw_rho,
+        UnitCell& ucell,
+        const double *veff
+    );
+    void generateTrainData_KS(
+        const std::string& dir,
+        psi::Psi<std::complex<float>, base_device::DEVICE_GPU>* psi,
+        elecstate::ElecState *pelec,
+        ModulePW::PW_Basis_K *pw_psi,
+        ModulePW::PW_Basis *pw_rho,
+        UnitCell& ucell,
+        const double *veff
+    );
+#endif
+
     void generate_descriptor(
         const std::string& out_dir,
         const double * const *prho, 

source/source_psi/psi.cpp

@@ -171,7 +171,6 @@ Psi<T, Device>::Psi(const Psi& psi_in)
     this->psi_current = this->psi + psi_in.get_psi_bias();
 }
 
-
 // Constructor 2-2:
 template <typename T, typename Device>
 template <typename T_in, typename Device_in>
@@ -545,6 +544,8 @@ template Psi<double, base_device::DEVICE_CPU>::Psi(const Psi<double, base_device
 template Psi<double, base_device::DEVICE_GPU>::Psi(const Psi<double, base_device::DEVICE_CPU>&);
 template Psi<std::complex<double>, base_device::DEVICE_CPU>::Psi(
     const Psi<std::complex<double>, base_device::DEVICE_GPU>&);
+template Psi<std::complex<double>, base_device::DEVICE_CPU>::Psi(
+    const Psi<std::complex<float>, base_device::DEVICE_GPU>&);
 template Psi<std::complex<double>, base_device::DEVICE_GPU>::Psi(
     const Psi<std::complex<double>, base_device::DEVICE_CPU>&);
 template Psi<std::complex<float>, base_device::DEVICE_GPU>::Psi(