66#include " source_io/para_json.h"
77#include " source_io/print_info.h"
88#include " source_md/run_md.h"
9+ #include " source_base/module_device/device.h"
10+ #include " source_base/module_device/memory_op.h"
11+ #include " source_base/kernels/math_kernel_op.h"
12+ #include " source_hsolver/kernels/dngvd_op.h"
13+
14+ #include < ATen/kernels/blas.h>
15+ #include < ATen/kernels/lapack.h>
16+
17+ #ifdef __DSP
18+ #include " source_base/kernels/dsp/dsp_connector.h"
19+ #endif
920
1021/* *
1122 * @brief This is the driver function which defines the workflow of ABACUS
@@ -47,6 +58,8 @@ void Driver::driver_run()
4758 unitcell::check_atomic_stru (ucell, PARAM .inp .min_dist_coef );
4859
4960 // ! 2: initialize the ESolver (depends on a set-up ucell after `setup_cell`)
61+ this ->init_hardware ();
62+
5063 ModuleESolver::ESolver* p_esolver = ModuleESolver::init_esolver (PARAM .inp , ucell);
5164
5265 // ! 3: initialize Esolver and fill json-structure
@@ -93,9 +106,46 @@ void Driver::driver_run()
93106 p_esolver->after_all_runners (ucell);
94107
95108 ModuleESolver::clean_esolver (p_esolver);
109+ this ->finalize_hardware ();
96110
97111 // ! 6: output the json file
98112 Json::create_Json (&ucell, PARAM );
99113
100114 return ;
101115}
116+
117+ void Driver::init_hardware ()
118+ {
119+ #if ((defined __CUDA) || (defined __ROCM))
120+ if (PARAM .inp .device == " gpu" )
121+ {
122+ ModuleBase::createGpuBlasHandle ();
123+ hsolver::createGpuSolverHandle ();
124+ container::kernels::createGpuBlasHandle ();
125+ container::kernels::createGpuSolverHandle ();
126+ }
127+ #endif
128+
129+ #ifdef __DSP
130+ std::cout << " ** Initializing DSP Hardware..." << std::endl;
131+ mtfunc::dspInitHandle (GlobalV::MY_RANK );
132+ #endif
133+ }
134+
135+ void Driver::finalize_hardware ()
136+ {
137+ #if defined(__CUDA) || defined(__ROCM)
138+ if (PARAM .inp .device == " gpu" )
139+ {
140+ ModuleBase::destoryBLAShandle ();
141+ hsolver::destroyGpuSolverHandle ();
142+ container::kernels::destroyGpuBlasHandle ();
143+ container::kernels::destroyGpuSolverHandle ();
144+ }
145+ #endif
146+
147+ #ifdef __DSP
148+ std::cout << " ** Closing DSP Hardware..." << std::endl;
149+ mtfunc::dspDestoryHandle (GlobalV::MY_RANK );
150+ #endif
151+ }
0 commit comments