Feature: Porting abacus to DSP hardware (mtblas part) (#5301)

Critsium-xy · mohanchen · web-flow · commit f03925010d11 · 2024-10-24T01:40:08.000+08:00
* Link mtblas library * Add mtblas gemm kernel usage * Finish memory_op on dsp * Update CMakeLists * Add compilation script * Fix warnings * Fix install script * Initialize DSP hardware * Replace gemm in math_kernel * Fix CMakeLists Bug * Fix bugs #1 * Fix bug 2 * Fix link to shared library error * Stop use gemm_mt globally * Modify op usage * Fix bug * Fix template usage * Fix compilation * Replace all dav_subspace gemm kernels --------- Co-authored-by: Mohan Chen <mohanchen@pku.edu.cn>
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -39,6 +39,7 @@ option(ENABLE_RAPIDJSON "Enable rapid-json usage." OFF)
 option(ENABLE_CNPY "Enable cnpy usage." OFF)
 option(ENABLE_PEXSI "Enable support for PEXSI." OFF)
 option(ENABLE_CUSOLVERMP "Enable cusolvermp." OFF)
+option(USE_DSP "Enable DSP usage." OFF)
 
 # enable json support
 if(ENABLE_RAPIDJSON)
@@ -119,6 +120,12 @@ elseif(ENABLE_LCAO AND NOT ENABLE_MPI)
   set(ABACUS_BIN_NAME abacus_serial)
 endif()
 
+if (USE_DSP)
+  set(USE_ELPA OFF)
+  set(ENABLE_LCAO OFF)
+  set(ABACUS_BIN_NAME abacus_dsp)
+endif()
+
 list(APPEND CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/cmake)
 
 if(ENABLE_COVERAGE)
@@ -240,6 +247,11 @@ if(ENABLE_MPI)
   list(APPEND math_libs MPI::MPI_CXX)
 endif()
 
+if (USE_DSP)
+  target_link_libraries(${ABACUS_BIN_NAME} ${DIR_MTBLAS_LIBRARY})
+  add_compile_definitions(__DSP)
+endif()
+
 find_package(Threads REQUIRED)
 target_link_libraries(${ABACUS_BIN_NAME} Threads::Threads)
 
diff --git a/install_dsp.sh b/install_dsp.sh
@@ -0,0 +1,10 @@
+CXX=mpicxx \
+    cmake -B build \
+        -DUSE_DSP=ON \
+        -DENABLE_LCAO=OFF \
+        -DFFTW3_DIR=/vol8/appsoftware/fftw/ \
+        -DFFTW3_LIBRARY=/vol8/appsoftware/fftw/lib/libfftw3.so \
+        -DFFTW3_OMP_LIBRARY=/vol8/appsoftware/fftw/lib/libfftw3_omp.so \
+        -DFFTW3_FLOAT_LIBRARY=/vol8/appsoftware/fftw/lib/libfftw3f.so \
+        -DLAPACK_DIR=/vol8/appsoftware/openblas/0.3.21/lib \
+        -DDIR_MTBLAS_LIBRARY=/vol8/home/dptech_zyz1/develop/packages/libmtblas_abacus.so
diff --git a/source/module_base/blas_connector.cpp b/source/module_base/blas_connector.cpp
@@ -1,5 +1,9 @@
 #include "blas_connector.h"
 
+#ifdef __DSP
+#include "module_base/kernels/dsp/dsp_connector.h"
+#endif
+
 void BlasConnector::axpy( const int n, const float alpha, const float *X, const int incX, float *Y, const int incY, base_device::AbacusDevice_t device_type)
 {
 	if (device_type == base_device::AbacusDevice_t::CpuDevice) {
@@ -64,13 +68,15 @@ float BlasConnector::dot( const int n, const float *X, const int incX, const flo
 {
 	if (device_type == base_device::AbacusDevice_t::CpuDevice) {
 		return sdot_(&n, X, &incX, Y, &incY);
+	return sdot_(&n, X, &incX, Y, &incY);
 }
 }
 
 double BlasConnector::dot( const int n, const double *X, const int incX, const double *Y, const int incY, base_device::AbacusDevice_t device_type)
 {
 	if (device_type == base_device::AbacusDevice_t::CpuDevice) {
 		return ddot_(&n, X, &incX, Y, &incY);
+	return ddot_(&n, X, &incX, Y, &incY);
 }
 }
 
@@ -83,7 +89,14 @@ void BlasConnector::gemm(const char transa, const char transb, const int m, cons
 		sgemm_(&transb, &transa, &n, &m, &k,
 		&alpha, b, &ldb, a, &lda,
 		&beta, c, &ldc);
-}
+	}
+	#ifdef __DSP
+	else if (device_type == base_device::AbacusDevice_t::DspDevice){
+		sgemm_mt_(&transb, &transa, &n, &m, &k,
+		&alpha, b, &ldb, a, &lda,
+		&beta, c, &ldc);
+	}
+	#endif
 }
 
 void BlasConnector::gemm(const char transa, const char transb, const int m, const int n, const int k,
@@ -94,7 +107,14 @@ void BlasConnector::gemm(const char transa, const char transb, const int m, cons
 		dgemm_(&transb, &transa, &n, &m, &k,
 		&alpha, b, &ldb, a, &lda,
 		&beta, c, &ldc);
-}
+	}
+	#ifdef __DSP
+	else if (device_type == base_device::AbacusDevice_t::DspDevice){
+		dgemm_mt_(&transb, &transa, &n, &m, &k,
+		&alpha, b, &ldb, a, &lda,
+		&beta, c, &ldc);
+	}
+	#endif
 }
 
 void BlasConnector::gemm(const char transa, const char transb, const int m, const int n, const int k,
@@ -105,7 +125,14 @@ void BlasConnector::gemm(const char transa, const char transb, const int m, cons
     	cgemm_(&transb, &transa, &n, &m, &k,
         &alpha, b, &ldb, a, &lda,
         &beta, c, &ldc);
-}
+	}
+	#ifdef __DSP
+	else if (device_type == base_device::AbacusDevice_t::DspDevice) {
+    	cgemm_mt_(&transb, &transa, &n, &m, &k,
+        &alpha, b, &ldb, a, &lda,
+        &beta, c, &ldc);
+	}
+	#endif
 }
 
 void BlasConnector::gemm(const char transa, const char transb, const int m, const int n, const int k,
@@ -116,7 +143,14 @@ void BlasConnector::gemm(const char transa, const char transb, const int m, cons
 		zgemm_(&transb, &transa, &n, &m, &k,
 		&alpha, b, &ldb, a, &lda,
 		&beta, c, &ldc);
-}
+	}
+	#ifdef __DSP
+	else if (device_type == base_device::AbacusDevice_t::DspDevice) {
+    	zgemm_mt_(&transb, &transa, &n, &m, &k,
+        &alpha, b, &ldb, a, &lda,
+        &beta, c, &ldc);
+	}
+	#endif
 }
 
 void BlasConnector::gemv(const char trans, const int m, const int n,
@@ -152,6 +186,7 @@ float BlasConnector::nrm2( const int n, const float *X, const int incX, base_dev
 {
 	if (device_type == base_device::AbacusDevice_t::CpuDevice) {
 		return snrm2_( &n, X, &incX );
+	return snrm2_( &n, X, &incX );
 }
 }
 
@@ -160,6 +195,7 @@ double BlasConnector::nrm2( const int n, const double *X, const int incX, base_d
 {
 	if (device_type == base_device::AbacusDevice_t::CpuDevice) {
 		return dnrm2_( &n, X, &incX );
+	return dnrm2_( &n, X, &incX );
 }
 }
 
@@ -168,6 +204,7 @@ double BlasConnector::nrm2( const int n, const std::complex<double> *X, const in
 {
 	if (device_type == base_device::AbacusDevice_t::CpuDevice) {
 		return dznrm2_( &n, X, &incX );
+	return dznrm2_( &n, X, &incX );
 }
 }
 
diff --git a/source/module_base/kernels/dsp/dsp_connector.h b/source/module_base/kernels/dsp/dsp_connector.h
@@ -0,0 +1,66 @@
+#ifndef DSP_CONNECTOR_H
+#define DSP_CONNECTOR_H
+#ifdef __DSP
+
+// Base dsp functions
+void dspInitHandle(int id);
+void dspDestoryHandle();
+void *malloc_ht(size_t bytes);
+void free_ht(void* ptr);
+
+
+// mtblas functions
+
+void sgemm_mt_(const char *transa, const char *transb,
+	const int *m, const int *n, const int *k,
+	const float *alpha, const float *a, const int *lda,
+	const float *b, const int *ldb, const float *beta,
+	float *c, const int *ldc);
+
+void dgemm_mt_(const char *transa, const char *transb,
+	const int *m, const int *n, const int *k,
+	const double *alpha,const double *a, const int *lda,
+	const double *b, const int *ldb, const double *beta,
+	double *c, const int *ldc);
+
+void zgemm_mt_(const char *transa, const char *transb,
+	const int *m, const int *n, const int *k,
+	const std::complex<double> *alpha, const std::complex<double> *a, const int *lda,
+	const std::complex<double> *b, const int *ldb, const std::complex<double> *beta,
+	std::complex<double> *c, const int *ldc);
+
+void cgemm_mt_(const char *transa, const char *transb,
+	const int *m, const int *n, const int *k,
+	const std::complex<float> *alpha, const std::complex<float> *a, const int *lda,
+	const std::complex<float> *b, const int *ldb, const std::complex<float> *beta,
+	std::complex<float> *c, const int *ldc);
+
+
+void sgemm_mth_(const char *transa, const char *transb,
+	const int *m, const int *n, const int *k,
+	const float *alpha, const float *a, const int *lda,
+	const float *b, const int *ldb, const float *beta,
+	float *c, const int *ldc);
+
+void dgemm_mth_(const char *transa, const char *transb,
+	const int *m, const int *n, const int *k,
+	const double *alpha,const double *a, const int *lda,
+	const double *b, const int *ldb, const double *beta,
+	double *c, const int *ldc);
+
+void zgemm_mth_(const char *transa, const char *transb,
+	const int *m, const int *n, const int *k,
+	const std::complex<double> *alpha, const std::complex<double> *a, const int *lda,
+	const std::complex<double> *b, const int *ldb, const std::complex<double> *beta,
+	std::complex<double> *c, const int *ldc);
+
+void cgemm_mth_(const char *transa, const char *transb,
+	const int *m, const int *n, const int *k,
+	const std::complex<float> *alpha, const std::complex<float> *a, const int *lda,
+	const std::complex<float> *b, const int *ldb, const std::complex<float> *beta,
+	std::complex<float> *c, const int *ldc);
+
+//#define zgemm_ zgemm_mt
+
+#endif
+#endif
diff --git a/source/module_base/module_device/memory_op.cpp b/source/module_base/module_device/memory_op.cpp
@@ -2,6 +2,9 @@
 
 #include "module_base/memory.h"
 #include "module_base/tool_threading.h"
+#ifdef __DSP
+#include "module_base/kernels/dsp/dsp_connector.h"
+#endif
 
 #include <complex>
 #include <cstring>
@@ -18,9 +21,17 @@ struct resize_memory_op<FPTYPE, base_device::DEVICE_CPU>
     {
         if (arr != nullptr)
         {
+#ifdef __DSP
+            free_ht(arr);
+#else
             free(arr);
+#endif
         }
+#ifdef __DSP
+        arr = (FPTYPE*)malloc_ht(sizeof(FPTYPE) * size);
+#else
         arr = (FPTYPE*)malloc(sizeof(FPTYPE) * size);
+#endif
         std::string record_string;
         if (record_in != nullptr)
         {
@@ -92,7 +103,11 @@ struct delete_memory_op<FPTYPE, base_device::DEVICE_CPU>
 {
     void operator()(const base_device::DEVICE_CPU* dev, FPTYPE* arr)
     {
+#ifdef __DSP
+        free_ht(arr);
+#else
         free(arr);
+#endif
     }
 };
 
diff --git a/source/module_base/module_device/types.h b/source/module_base/module_device/types.h
@@ -12,6 +12,7 @@ enum AbacusDevice_t
     UnKnown,
     CpuDevice,
     GpuDevice,
+    DspDevice
 };
 
 } // namespace base_device
diff --git a/source/module_esolver/esolver_ks_pw.cpp b/source/module_esolver/esolver_ks_pw.cpp
@@ -49,6 +49,10 @@
 #include <ATen/kernels/blas.h>
 #include <ATen/kernels/lapack.h>
 
+#ifdef __DSP
+#include "module_base/kernels/dsp/dsp_connector.h"
+#endif
+
 namespace ModuleESolver
 {
 
@@ -67,6 +71,10 @@ ESolver_KS_PW<T, Device>::ESolver_KS_PW()
         container::kernels::createGpuSolverHandle();
     }
 #endif
+#ifdef __DSP
+    std::cout << " ** Initializing DSP Hardware..." << std::endl;
+    dspInitHandle(GlobalV::MY_RANK % 4);
+#endif
 }
 
 template <typename T, typename Device>
@@ -92,7 +100,10 @@ ESolver_KS_PW<T, Device>::~ESolver_KS_PW()
 #endif
         delete reinterpret_cast<psi::Psi<T, Device>*>(this->kspw_psi);
     }
-    
+#ifdef __DSP
+    std::cout << " ** Closing DSP Hardware..." << std::endl;
+    dspDestoryHandle();
+#endif
     if (PARAM.inp.precision == "single")
     {
         delete reinterpret_cast<psi::Psi<std::complex<double>, Device>*>(this->__kspw_psi);
diff --git a/source/module_hsolver/diago_dav_subspace.cpp b/source/module_hsolver/diago_dav_subspace.cpp
@@ -181,7 +181,12 @@ int Diago_DavSubspace<T, Device>::diag_once(const HPsiFunc& hpsi_func,
             // updata eigenvectors of Hamiltonian
             setmem_complex_op()(this->ctx, psi_in, 0, n_band * psi_in_dmax);
 
-            gemm_op<T, Device>()(this->ctx,
+#ifdef __DSP
+    gemm_op_mt<T, Device>()
+#else
+    gemm_op<T, Device>()
+#endif
+                                (this->ctx,
                                  'N',
                                  'N',
                                  this->dim,
@@ -262,7 +267,12 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
         }
     }
 
-    gemm_op<T, Device>()(this->ctx,
+#ifdef __DSP
+    gemm_op_mt<T, Device>()
+#else
+    gemm_op<T, Device>()
+#endif
+                        (this->ctx,
                          'N',
                          'N',
                          this->dim,
@@ -302,7 +312,12 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
         delmem_real_op()(this->ctx, e_temp_hd);
     }
 
-    gemm_op<T, Device>()(this->ctx,
+#ifdef __DSP
+    gemm_op_mt<T, Device>()
+#else
+    gemm_op<T, Device>()
+#endif                  
+                        (this->ctx,
                          'N',
                          'N',
                          this->dim,
@@ -386,7 +401,12 @@ void Diago_DavSubspace<T, Device>::cal_elem(const int& dim,
 {
     ModuleBase::timer::tick("Diago_DavSubspace", "cal_elem");
 
-    gemm_op<T, Device>()(this->ctx,
+#ifdef __DSP
+    gemm_op_mt<T, Device>()
+#else
+    gemm_op<T, Device>()
+#endif 
+                        (this->ctx,
                          'C',
                          'N',
                          nbase + notconv,
@@ -401,7 +421,12 @@ void Diago_DavSubspace<T, Device>::cal_elem(const int& dim,
                          &hcc[nbase * this->nbase_x],
                          this->nbase_x);
 
-    gemm_op<T, Device>()(this->ctx,
+#ifdef __DSP
+    gemm_op_mt<T, Device>()
+#else
+    gemm_op<T, Device>()
+#endif
+                        (this->ctx,
                          'C',
                          'N',
                          nbase + notconv,
@@ -603,7 +628,12 @@ void Diago_DavSubspace<T, Device>::refresh(const int& dim,
 {
     ModuleBase::timer::tick("Diago_DavSubspace", "refresh");
 
-    gemm_op<T, Device>()(this->ctx,
+#ifdef __DSP
+    gemm_op_mt<T, Device>()
+#else
+    gemm_op<T, Device>()
+#endif
+                        (this->ctx,
                          'N',
                          'N',
                          this->dim,
diff --git a/source/module_hsolver/kernels/math_kernel_op.cpp b/source/module_hsolver/kernels/math_kernel_op.cpp
diff --git a/source/module_hsolver/kernels/math_kernel_op.h b/source/module_hsolver/kernels/math_kernel_op.h

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,9 @@`
`1`	`1`	`#include "blas_connector.h"`
`2`	`2`
	`3`	`+#ifdef __DSP`
	`4`	`+#include "module_base/kernels/dsp/dsp_connector.h"`
	`5`	`+#endif`
	`6`	`+`
`3`	`7`	`void BlasConnector::axpy( const int n, const float alpha, const float X, const int incX, float Y, const int incY, base_device::AbacusDevice_t device_type)`
`4`	`8`	`{`
`5`	`9`	`if (device_type == base_device::AbacusDevice_t::CpuDevice) {`
`@@ -64,13 +68,15 @@ float BlasConnector::dot( const int n, const float *X, const int incX, const flo`
`64`	`68`	`{`
`65`	`69`	`if (device_type == base_device::AbacusDevice_t::CpuDevice) {`
`66`	`70`	`return sdot_(&n, X, &incX, Y, &incY);`
	`71`	`+ return sdot_(&n, X, &incX, Y, &incY);`
`67`	`72`	`}`
`68`	`73`	`}`
`69`	`74`
`70`	`75`	`double BlasConnector::dot( const int n, const double X, const int incX, const double Y, const int incY, base_device::AbacusDevice_t device_type)`
`71`	`76`	`{`
`72`	`77`	`if (device_type == base_device::AbacusDevice_t::CpuDevice) {`
`73`	`78`	`return ddot_(&n, X, &incX, Y, &incY);`
	`79`	`+ return ddot_(&n, X, &incX, Y, &incY);`
`74`	`80`	`}`
`75`	`81`	`}`
`76`	`82`
`@@ -83,7 +89,14 @@ void BlasConnector::gemm(const char transa, const char transb, const int m, cons`
`83`	`89`	`sgemm_(&transb, &transa, &n, &m, &k,`
`84`	`90`	`&alpha, b, &ldb, a, &lda,`
`85`	`91`	`&beta, c, &ldc);`
`86`		`-}`
	`92`	`+ }`
	`93`	`+ #ifdef __DSP`
	`94`	`+ else if (device_type == base_device::AbacusDevice_t::DspDevice){`
	`95`	`+ sgemm_mt_(&transb, &transa, &n, &m, &k,`
	`96`	`+ &alpha, b, &ldb, a, &lda,`
	`97`	`+ &beta, c, &ldc);`
	`98`	`+ }`
	`99`	`+ #endif`
`87`	`100`	`}`
`88`	`101`
`89`	`102`	`void BlasConnector::gemm(const char transa, const char transb, const int m, const int n, const int k,`
`@@ -94,7 +107,14 @@ void BlasConnector::gemm(const char transa, const char transb, const int m, cons`
`94`	`107`	`dgemm_(&transb, &transa, &n, &m, &k,`
`95`	`108`	`&alpha, b, &ldb, a, &lda,`
`96`	`109`	`&beta, c, &ldc);`
`97`		`-}`
	`110`	`+ }`
	`111`	`+ #ifdef __DSP`
	`112`	`+ else if (device_type == base_device::AbacusDevice_t::DspDevice){`
	`113`	`+ dgemm_mt_(&transb, &transa, &n, &m, &k,`
	`114`	`+ &alpha, b, &ldb, a, &lda,`
	`115`	`+ &beta, c, &ldc);`
	`116`	`+ }`
	`117`	`+ #endif`
`98`	`118`	`}`
`99`	`119`
`100`	`120`	`void BlasConnector::gemm(const char transa, const char transb, const int m, const int n, const int k,`
`@@ -105,7 +125,14 @@ void BlasConnector::gemm(const char transa, const char transb, const int m, cons`
`105`	`125`	`cgemm_(&transb, &transa, &n, &m, &k,`
`106`	`126`	`&alpha, b, &ldb, a, &lda,`
`107`	`127`	`&beta, c, &ldc);`
`108`		`-}`
	`128`	`+ }`
	`129`	`+ #ifdef __DSP`
	`130`	`+ else if (device_type == base_device::AbacusDevice_t::DspDevice) {`
	`131`	`+ cgemm_mt_(&transb, &transa, &n, &m, &k,`
	`132`	`+ &alpha, b, &ldb, a, &lda,`
	`133`	`+ &beta, c, &ldc);`
	`134`	`+ }`
	`135`	`+ #endif`
`109`	`136`	`}`
`110`	`137`
`111`	`138`	`void BlasConnector::gemm(const char transa, const char transb, const int m, const int n, const int k,`
`@@ -116,7 +143,14 @@ void BlasConnector::gemm(const char transa, const char transb, const int m, cons`
`116`	`143`	`zgemm_(&transb, &transa, &n, &m, &k,`
`117`	`144`	`&alpha, b, &ldb, a, &lda,`
`118`	`145`	`&beta, c, &ldc);`
`119`		`-}`
	`146`	`+ }`
	`147`	`+ #ifdef __DSP`
	`148`	`+ else if (device_type == base_device::AbacusDevice_t::DspDevice) {`
	`149`	`+ zgemm_mt_(&transb, &transa, &n, &m, &k,`
	`150`	`+ &alpha, b, &ldb, a, &lda,`
	`151`	`+ &beta, c, &ldc);`
	`152`	`+ }`
	`153`	`+ #endif`
`120`	`154`	`}`
`121`	`155`
`122`	`156`	`void BlasConnector::gemv(const char trans, const int m, const int n,`
`@@ -152,6 +186,7 @@ float BlasConnector::nrm2( const int n, const float *X, const int incX, base_dev`
`152`	`186`	`{`
`153`	`187`	`if (device_type == base_device::AbacusDevice_t::CpuDevice) {`
`154`	`188`	`return snrm2_( &n, X, &incX );`
	`189`	`+ return snrm2_( &n, X, &incX );`
`155`	`190`	`}`
`156`	`191`	`}`
`157`	`192`
`@@ -160,6 +195,7 @@ double BlasConnector::nrm2( const int n, const double *X, const int incX, base_d`
`160`	`195`	`{`
`161`	`196`	`if (device_type == base_device::AbacusDevice_t::CpuDevice) {`
`162`	`197`	`return dnrm2_( &n, X, &incX );`
	`198`	`+ return dnrm2_( &n, X, &incX );`
`163`	`199`	`}`
`164`	`200`	`}`
`165`	`201`
`@@ -168,6 +204,7 @@ double BlasConnector::nrm2( const int n, const std::complex<double> *X, const in`
`168`	`204`	`{`
`169`	`205`	`if (device_type == base_device::AbacusDevice_t::CpuDevice) {`
`170`	`206`	`return dznrm2_( &n, X, &incX );`
	`207`	`+ return dznrm2_( &n, X, &incX );`
`171`	`208`	`}`
`172`	`209`	`}`
`173`	`210`
Original file line number	Diff line number	Diff line change
`@@ -2,6 +2,9 @@`
`2`	`2`
`3`	`3`	`#include "module_base/memory.h"`
`4`	`4`	`#include "module_base/tool_threading.h"`
	`5`	`+#ifdef __DSP`
	`6`	`+#include "module_base/kernels/dsp/dsp_connector.h"`
	`7`	`+#endif`
`5`	`8`
`6`	`9`	`#include <complex>`
`7`	`10`	`#include <cstring>`
`@@ -18,9 +21,17 @@ struct resize_memory_op<FPTYPE, base_device::DEVICE_CPU>`
`18`	`21`	`{`
`19`	`22`	`if (arr != nullptr)`
`20`	`23`	`{`
	`24`	`+#ifdef __DSP`
	`25`	`+ free_ht(arr);`
	`26`	`+#else`
`21`	`27`	`free(arr);`
	`28`	`+#endif`
`22`	`29`	`}`
	`30`	`+#ifdef __DSP`
	`31`	`+ arr = (FPTYPE)malloc_ht(sizeof(FPTYPE) size);`
	`32`	`+#else`
`23`	`33`	`arr = (FPTYPE)malloc(sizeof(FPTYPE) size);`
	`34`	`+#endif`
`24`	`35`	`std::string record_string;`
`25`	`36`	`if (record_in != nullptr)`
`26`	`37`	`{`
`@@ -92,7 +103,11 @@ struct delete_memory_op<FPTYPE, base_device::DEVICE_CPU>`
`92`	`103`	`{`
`93`	`104`	`void operator()(const base_device::DEVICE_CPU* dev, FPTYPE* arr)`
`94`	`105`	`{`
	`106`	`+#ifdef __DSP`
	`107`	`+ free_ht(arr);`
	`108`	`+#else`
`95`	`109`	`free(arr);`
	`110`	`+#endif`
`96`	`111`	`}`
`97`	`112`	`};`
`98`	`113`
Original file line number	Diff line number	Diff line change
`@@ -49,6 +49,10 @@`
`49`	`49`	`#include <ATen/kernels/blas.h>`
`50`	`50`	`#include <ATen/kernels/lapack.h>`
`51`	`51`
	`52`	`+#ifdef __DSP`
	`53`	`+#include "module_base/kernels/dsp/dsp_connector.h"`
	`54`	`+#endif`
	`55`	`+`
`52`	`56`	`namespace ModuleESolver`
`53`	`57`	`{`
`54`	`58`
`@@ -67,6 +71,10 @@ ESolver_KS_PW<T, Device>::ESolver_KS_PW()`
`67`	`71`	`container::kernels::createGpuSolverHandle();`
`68`	`72`	`}`
`69`	`73`	`#endif`
	`74`	`+#ifdef __DSP`
	`75`	`+ std::cout << " ** Initializing DSP Hardware..." << std::endl;`
	`76`	`+ dspInitHandle(GlobalV::MY_RANK % 4);`
	`77`	`+#endif`
`70`	`78`	`}`
`71`	`79`
`72`	`80`	`template <typename T, typename Device>`
`@@ -92,7 +100,10 @@ ESolver_KS_PW<T, Device>::~ESolver_KS_PW()`
`92`	`100`	`#endif`
`93`	`101`	`delete reinterpret_cast<psi::Psi<T, Device>*>(this->kspw_psi);`
`94`	`102`	`}`
`95`		`-`
	`103`	`+#ifdef __DSP`
	`104`	`+ std::cout << " ** Closing DSP Hardware..." << std::endl;`
	`105`	`+ dspDestoryHandle();`
	`106`	`+#endif`
`96`	`107`	`if (PARAM.inp.precision == "single")`
`97`	`108`	`{`
`98`	`109`	`delete reinterpret_cast<psi::Psi<std::complex<double>, Device>*>(this->__kspw_psi);`