Fix GPU, but so ugly...

Author: Flying-dragon-boxing

.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: |

docs/advanced/input_files/input-main.md

@@ -1487,15 +1487,20 @@ These variables are used to control the geometry relaxation.
 
 ### relax_method
 
-- **Type**: String
+- **Type**: Vector of string
 - **Description**: The methods to do geometry optimization.
+  the first element:
   - cg: using the conjugate gradient (CG) algorithm. Note that there are two implementations of the conjugate gradient (CG) method, see [relax_new](#relax_new).
-  - bfgs: using the Broyden–Fletcher–Goldfarb–Shanno (BFGS) algorithm.
-  - bfgs_trad: using the traditional Broyden–Fletcher–Goldfarb–Shanno (BFGS) algorithm.
+  - bfgs : using the Broyden–Fletcher–Goldfarb–Shanno (BFGS) algorithm.
+  - lbfgs: using the Limited-memory Broyden–Fletcher–Goldfarb–Shanno (LBFGS) algorithm.
   - cg_bfgs: using the CG method for the initial steps, and switching to BFGS method when the force convergence is smaller than [relax_cg_thr](#relax_cg_thr).
   - sd: using the steepest descent (SD) algorithm.
   - fire: the Fast Inertial Relaxation Engine method (FIRE), a kind of molecular-dynamics-based relaxation algorithm, is implemented in the molecular dynamics (MD) module. The algorithm can be used by setting [calculation](#calculation) to `md` and [md_type](#md_type) to `fire`. Also ionic velocities should be set in this case. See [fire](../md.md#fire) for more details.
-- **Default**: cg
+
+  the second element:
+  when the first element is bfgs, if the second parameter is 1, it indicates the use of the new BFGS algorithm; if the second parameter is not 1, it indicates the use of the old BFGS algorithm.
+- **Default**: cg 1
+- **Note**:In the 3.10-LTS version, the type of this parameter is std::string. It can be set to "cg","bfgs","cg_bfgs","bfgs_trad","lbfgs","sd","fire".
 
 ### relax_new
 

source/source_io/module_parameter/input_parameter.h

@@ -6,6 +6,7 @@
 #include <string>
 #include <vector>
 
+
 // It stores all input parameters both defined in INPUT file and not defined in
 // INPUT file
 struct Input_para
@@ -150,7 +151,7 @@ struct Input_para
     //  int		bessel_nao_lmax;		///< lmax used in descriptor
 
     // ==============   #Parameters (4.Relaxation) ===========================
-    std::string relax_method = "cg"; ///< methods to move_ion: sd, bfgs, cg...
+    std::vector<std::string> relax_method = {"cg","1"}; ///< methods to move_ion: sd, bfgs, cg...
     bool relax_new = true;
     bool relax = false; ///< allow relaxation along the specific direction
     double relax_scale_force = 0.5;

source/source_io/read_input_item_relax.cpp

@@ -5,28 +5,53 @@
 
 namespace ModuleIO
 {
+
+
 void ReadInput::item_relax()
 {
     {
         Input_Item item("relax_method");
         item.annotation = "cg; bfgs; sd; cg; cg_bfgs;";
-        read_sync_string(input.relax_method);
-        item.check_value = [](const Input_Item& item, const Parameter& para) {
-            const std::vector<std::string> relax_methods = {"cg", "bfgs", "sd", "cg_bfgs","bfgs_trad","lbfgs"};
-            if (std::find(relax_methods.begin(),relax_methods.end(), para.input.relax_method)==relax_methods.end())
-            {
-                const std::string warningstr = nofound_str(relax_methods, "relax_method");
-                ModuleBase::WARNING_QUIT("ReadInput", warningstr);
-            }
+        item.read_value = [](const Input_Item& item, Parameter& para) {
+        if(item.get_size()==1)
+        {
+            para.input.relax_method[0] = item.str_values[0];
+            para.input.relax_method[1] = "1"; 
+        }
+        else if(item.get_size()>=2)
+        {
+            para.input.relax_method[0] = item.str_values[0];
+            para.input.relax_method[1] = item.str_values[1];
+        }
+        };
+        item.check_value = [](const Input_Item& item, const Parameter& para) {          
+        const std::vector<std::string> relax_methods = {"cg", "sd", "cg_bfgs","lbfgs","bfgs"};
+        if (std::find(relax_methods.begin(), relax_methods.end(), para.input.relax_method[0]) == relax_methods.end()) {
+            const std::string warningstr = nofound_str(relax_methods, "relax_method");
+            ModuleBase::WARNING_QUIT("ReadInput", warningstr);
+        }
         };
         this->add_item(item);
+        
+        // Input_Item item("relax_method");
+        // item.annotation = "cg; bfgs; sd; cg; cg_bfgs;";
+        // read_sync_string(input.relax_method);
+        // item.check_value = [](const Input_Item& item, const Parameter& para) {
+        //     const std::vector<std::string> relax_methods = {"cg", "bfgs_old", "sd", "cg_bfgs","bfgs","lbfgs"};
+        //     if (std::find(relax_methods.begin(),relax_methods.end(), para.input.relax_method)==relax_methods.end())
+        //     {
+        //         const std::string warningstr = nofound_str(relax_methods, "relax_method");
+        //         ModuleBase::WARNING_QUIT("ReadInput", warningstr);
+        //     }
+        // };
+        // this->add_item(item);
     }
     {
         Input_Item item("relax_new");
         item.annotation = "whether to use the new relaxation method";
         read_sync_bool(input.relax_new);
         item.reset_value = [](const Input_Item& item, Parameter& para) {
-            if (para.input.relax_new && para.input.relax_method != "cg")
+            if (para.input.relax_new && para.input.relax_method[0] != "cg")
             {
                 para.input.relax_new = false;
             }

source/source_io/test/read_input_ptest.cpp

@@ -114,7 +114,7 @@ TEST_F(InputParaTest, ParaRead)
     EXPECT_EQ(param.inp.fixed_axes, "None");
     EXPECT_FALSE(param.inp.fixed_ibrav);
     EXPECT_FALSE(param.inp.fixed_atoms);
-    EXPECT_EQ(param.inp.relax_method, "cg");
+    EXPECT_EQ(param.inp.relax_method[0], "cg");
     EXPECT_DOUBLE_EQ(param.inp.relax_cg_thr, 0.5);
     EXPECT_EQ(param.inp.out_level, "ie");
     EXPECT_TRUE(param.globalv.out_md_control);

source/source_io/test_serial/read_input_item_test.cpp

@@ -766,7 +766,7 @@ TEST_F(InputTest, Item_test)
     }
     { // relax_method
         auto it = find_label("relax_method", readinput.input_lists);
-        param.input.relax_method = "none";
+        param.input.relax_method[0] = "none";
         testing::internal::CaptureStdout();
         EXPECT_EXIT(it->second.check_value(it->second, param), ::testing::ExitedWithCode(1), "");
         output = testing::internal::GetCapturedStdout();
@@ -775,12 +775,12 @@ TEST_F(InputTest, Item_test)
     { //relax_new
         auto it = find_label("relax_new", readinput.input_lists);
         param.input.relax_new = true;
-        param.input.relax_method = "cg";
+        param.input.relax_method[0] = "cg";
         it->second.reset_value(it->second, param);
         EXPECT_EQ(param.input.relax_new, true);
 
         param.input.relax_new = true;
-        param.input.relax_method = "none";
+        param.input.relax_method[0] = "none";
         it->second.reset_value(it->second, param);
         EXPECT_EQ(param.input.relax_new, false);
     }

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(

source/source_relax/ions_move_basic.cpp

@@ -23,7 +23,7 @@ double Ions_Move_Basic::relax_bfgs_init = -1.0; // default is 0.5
 double Ions_Move_Basic::best_xxx = 1.0;
 
 int Ions_Move_Basic::out_stru = 0;
-std::string Ions_Move_Basic::relax_method = "bfgs";
+std::vector<std::string> Ions_Move_Basic::relax_method = {"bfgs","2"};
 
 void Ions_Move_Basic::setup_gradient(const UnitCell &ucell, const ModuleBase::matrix &force, double *pos, double *grad)
 {