Feature: MPI available on ABACUS DSP version (#5351)

Critsium-xy · web-flow · commit 72d9d1dd02bb · 2024-10-26T11:49:13.000+08:00
* Fix parallel function

* Fix parallel usage

* Temporarily remove memory_op porting
diff --git a/source/module_base/blas_connector.cpp b/source/module_base/blas_connector.cpp
@@ -2,6 +2,7 @@
 
 #ifdef __DSP
 #include "module_base/kernels/dsp/dsp_connector.h"
+#include "module_base/global_variable.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)
@@ -94,7 +95,7 @@ void BlasConnector::gemm(const char transa, const char transb, const int m, cons
 	else if (device_type == base_device::AbacusDevice_t::DspDevice){
 		sgemm_mt_(&transb, &transa, &n, &m, &k,
 		&alpha, b, &ldb, a, &lda,
-		&beta, c, &ldc);
+		&beta, c, &ldc, GlobalV::MY_RANK);
 	}
 	#endif
 }
@@ -112,7 +113,7 @@ void BlasConnector::gemm(const char transa, const char transb, const int m, cons
 	else if (device_type == base_device::AbacusDevice_t::DspDevice){
 		dgemm_mt_(&transb, &transa, &n, &m, &k,
 		&alpha, b, &ldb, a, &lda,
-		&beta, c, &ldc);
+		&beta, c, &ldc, GlobalV::MY_RANK);
 	}
 	#endif
 }
@@ -130,7 +131,7 @@ void BlasConnector::gemm(const char transa, const char transb, const int m, cons
 	else if (device_type == base_device::AbacusDevice_t::DspDevice) {
     	cgemm_mt_(&transb, &transa, &n, &m, &k,
         &alpha, b, &ldb, a, &lda,
-        &beta, c, &ldc);
+        &beta, c, &ldc, GlobalV::MY_RANK);
 	}
 	#endif
 }
@@ -148,7 +149,7 @@ void BlasConnector::gemm(const char transa, const char transb, const int m, cons
 	else if (device_type == base_device::AbacusDevice_t::DspDevice) {
     	zgemm_mt_(&transb, &transa, &n, &m, &k,
         &alpha, b, &ldb, a, &lda,
-        &beta, c, &ldc);
+        &beta, c, &ldc, GlobalV::MY_RANK);
 	}
 	#endif
 }
diff --git a/source/module_base/kernels/dsp/dsp_connector.h b/source/module_base/kernels/dsp/dsp_connector.h
@@ -4,8 +4,8 @@
 
 // Base dsp functions
 void dspInitHandle(int id);
-void dspDestoryHandle();
-void *malloc_ht(size_t bytes);
+void dspDestoryHandle(int id);
+void *malloc_ht(size_t bytes, int cluster_id);
 void free_ht(void* ptr);
 
 
@@ -15,50 +15,50 @@ 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);
+	float *c, const int *ldc, int cluster_id);
 
 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);
+	double *c, const int *ldc, int cluster_id);
 
 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);
+	std::complex<double> *c, const int *ldc, int cluster_id);
 
 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);
+	std::complex<float> *c, const int *ldc, int cluster_id);
 
 
 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);
+	float *c, const int *ldc, int cluster_id);
 
 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);
+	double *c, const int *ldc, int cluster_id);
 
 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);
+	std::complex<double> *c, const int *ldc, int cluster_id);
 
 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);
+	std::complex<float> *c, const int *ldc, int cluster_id);
 
 //#define zgemm_ zgemm_mt
 
diff --git a/source/module_base/module_device/memory_op.cpp b/source/module_base/module_device/memory_op.cpp
@@ -4,6 +4,7 @@
 #include "module_base/tool_threading.h"
 #ifdef __DSP
 #include "module_base/kernels/dsp/dsp_connector.h"
+#include "module_base/global_variable.h"
 #endif
 
 #include <complex>
@@ -21,17 +22,9 @@ 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)
         {
@@ -103,11 +96,7 @@ 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_esolver/esolver_ks_pw.cpp b/source/module_esolver/esolver_ks_pw.cpp
@@ -73,7 +73,7 @@ ESolver_KS_PW<T, Device>::ESolver_KS_PW()
 #endif
 #ifdef __DSP
     std::cout << " ** Initializing DSP Hardware..." << std::endl;
-    dspInitHandle(GlobalV::MY_RANK % 4);
+    dspInitHandle(GlobalV::MY_RANK);
 #endif
 }
 
@@ -102,7 +102,7 @@ ESolver_KS_PW<T, Device>::~ESolver_KS_PW()
     }
 #ifdef __DSP
     std::cout << " ** Closing DSP Hardware..." << std::endl;
-    dspDestoryHandle();
+    dspDestoryHandle(GlobalV::MY_RANK);
 #endif
     if (PARAM.inp.precision == "single")
     {

Original file line number	Diff line number	Diff line change
`@@ -2,6 +2,7 @@`
`2`	`2`
`3`	`3`	`#ifdef __DSP`
`4`	`4`	`#include "module_base/kernels/dsp/dsp_connector.h"`
	`5`	`+#include "module_base/global_variable.h"`
`5`	`6`	`#endif`
`6`	`7`
`7`	`8`	`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)`
`@@ -94,7 +95,7 @@ void BlasConnector::gemm(const char transa, const char transb, const int m, cons`
`94`	`95`	`else if (device_type == base_device::AbacusDevice_t::DspDevice){`
`95`	`96`	`sgemm_mt_(&transb, &transa, &n, &m, &k,`
`96`	`97`	`&alpha, b, &ldb, a, &lda,`
`97`		`- &beta, c, &ldc);`
	`98`	`+ &beta, c, &ldc, GlobalV::MY_RANK);`
`98`	`99`	`}`
`99`	`100`	`#endif`
`100`	`101`	`}`
`@@ -112,7 +113,7 @@ void BlasConnector::gemm(const char transa, const char transb, const int m, cons`
`112`	`113`	`else if (device_type == base_device::AbacusDevice_t::DspDevice){`
`113`	`114`	`dgemm_mt_(&transb, &transa, &n, &m, &k,`
`114`	`115`	`&alpha, b, &ldb, a, &lda,`
`115`		`- &beta, c, &ldc);`
	`116`	`+ &beta, c, &ldc, GlobalV::MY_RANK);`
`116`	`117`	`}`
`117`	`118`	`#endif`
`118`	`119`	`}`
`@@ -130,7 +131,7 @@ void BlasConnector::gemm(const char transa, const char transb, const int m, cons`
`130`	`131`	`else if (device_type == base_device::AbacusDevice_t::DspDevice) {`
`131`	`132`	`cgemm_mt_(&transb, &transa, &n, &m, &k,`
`132`	`133`	`&alpha, b, &ldb, a, &lda,`
`133`		`- &beta, c, &ldc);`
	`134`	`+ &beta, c, &ldc, GlobalV::MY_RANK);`
`134`	`135`	`}`
`135`	`136`	`#endif`
`136`	`137`	`}`
`@@ -148,7 +149,7 @@ void BlasConnector::gemm(const char transa, const char transb, const int m, cons`
`148`	`149`	`else if (device_type == base_device::AbacusDevice_t::DspDevice) {`
`149`	`150`	`zgemm_mt_(&transb, &transa, &n, &m, &k,`
`150`	`151`	`&alpha, b, &ldb, a, &lda,`
`151`		`- &beta, c, &ldc);`
	`152`	`+ &beta, c, &ldc, GlobalV::MY_RANK);`
`152`	`153`	`}`
`153`	`154`	`#endif`
`154`	`155`	`}`
Original file line number	Diff line number	Diff line change
`@@ -4,6 +4,7 @@`
`4`	`4`	`#include "module_base/tool_threading.h"`
`5`	`5`	`#ifdef __DSP`
`6`	`6`	`#include "module_base/kernels/dsp/dsp_connector.h"`
	`7`	`+#include "module_base/global_variable.h"`
`7`	`8`	`#endif`
`8`	`9`
`9`	`10`	`#include <complex>`
`@@ -21,17 +22,9 @@ struct resize_memory_op<FPTYPE, base_device::DEVICE_CPU>`
`21`	`22`	`{`
`22`	`23`	`if (arr != nullptr)`
`23`	`24`	`{`
`24`		`-#ifdef __DSP`
`25`		`- free_ht(arr);`
`26`		`-#else`
`27`	`25`	`free(arr);`
`28`		`-#endif`
`29`	`26`	`}`
`30`		`-#ifdef __DSP`
`31`		`- arr = (FPTYPE)malloc_ht(sizeof(FPTYPE) size);`
`32`		`-#else`
`33`	`27`	`arr = (FPTYPE)malloc(sizeof(FPTYPE) size);`
`34`		`-#endif`
`35`	`28`	`std::string record_string;`
`36`	`29`	`if (record_in != nullptr)`
`37`	`30`	`{`
`@@ -103,11 +96,7 @@ struct delete_memory_op<FPTYPE, base_device::DEVICE_CPU>`
`103`	`96`	`{`
`104`	`97`	`void operator()(const base_device::DEVICE_CPU* dev, FPTYPE* arr)`
`105`	`98`	`{`
`106`		`-#ifdef __DSP`
`107`		`- free_ht(arr);`
`108`		`-#else`
`109`	`99`	`free(arr);`
`110`		`-#endif`
`111`	`100`	`}`
`112`	`101`	`};`
`113`	`102`
Original file line number	Diff line number	Diff line change
`@@ -73,7 +73,7 @@ ESolver_KS_PW<T, Device>::ESolver_KS_PW()`
`73`	`73`	`#endif`
`74`	`74`	`#ifdef __DSP`
`75`	`75`	`std::cout << " ** Initializing DSP Hardware..." << std::endl;`
`76`		`- dspInitHandle(GlobalV::MY_RANK % 4);`
	`76`	`+ dspInitHandle(GlobalV::MY_RANK);`
`77`	`77`	`#endif`
`78`	`78`	`}`
`79`	`79`
`@@ -102,7 +102,7 @@ ESolver_KS_PW<T, Device>::~ESolver_KS_PW()`
`102`	`102`	`}`
`103`	`103`	`#ifdef __DSP`
`104`	`104`	`std::cout << " ** Closing DSP Hardware..." << std::endl;`
`105`		`- dspDestoryHandle();`
	`105`	`+ dspDestoryHandle(GlobalV::MY_RANK);`
`106`	`106`	`#endif`
`107`	`107`	`if (PARAM.inp.precision == "single")`
`108`	`108`	`{`