Get rid of duplicate definitions of DPErrcheck

deepmodeling · Jul 20, 2021 · c0f57f6 · c0f57f6
1 parent eea2ab1
commit c0f57f6
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 36 deletions.
diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc
@@ -1,46 +1,11 @@
 #include "DeepPot.h"
 #include "AtomMap.h"
 #include <stdexcept>	
+#include "device.h"
 
 using namespace tensorflow;
 using namespace deepmd;
 
-#if  GOOGLE_CUDA
-#include "cuda_runtime.h"
-
-#define DPErrcheck(res) { DPAssert((res), __FILE__, __LINE__); }
-inline void DPAssert(cudaError_t code, const char *file, int line, bool abort=true)
-{
-    if (code != cudaSuccess)
-    {
-        fprintf(stderr,"cuda assert: %s %s %d\n", cudaGetErrorString(code), file, line);
-        if (abort) exit(code);
-    }
-}
-
-void DPGetDeviceCount(int &gpu_num) { cudaGetDeviceCount(&gpu_num) ;}
-
-cudaError_t DPSetDevice(int rank) { return  cudaSetDevice(rank); }
-#endif
-
-#if  TENSORFLOW_USE_ROCM
-#include<hip/hip_runtime.h>
-
-#define DPErrcheck(res) { DPAssert((res), __FILE__, __LINE__); }
-inline void DPAssert(hipError_t code, const char *file, int line, bool abort=true)
-{
-    if (code != hipSuccess)
-    {
-        fprintf(stderr,"hip assert: %s %s %d\n", hipGetErrorString(code), file, line);
-        if (abort) exit(code);
-    }
-}
-
-void DPGetDeviceCount(int &gpu_num) { hipGetDeviceCount(&gpu_num) ;}
-hipError_t DPSetDevice(int rank) { return  hipSetDevice(rank); }
-
-#endif //TENSORFLOW_USE_ROCM
-
 static 
 std::vector<int> cum_sum (const std::vector<int32> & n_sel) {
     std::vector<int> sec;

diff --git a/source/lib/include/gpu_cuda.h b/source/lib/include/gpu_cuda.h
@@ -67,6 +67,11 @@ static __inline__ __device__ double atomicAdd(
 #endif
 
 namespace deepmd {
+
+inline void DPGetDeviceCount(int &gpu_num) { cudaGetDeviceCount(&gpu_num) ;}
+
+inline cudaError_t DPSetDevice(int rank) { return  cudaSetDevice(rank); }
+
 template <typename FPTYPE>
 void memcpy_host_to_device(
     FPTYPE * device, 

diff --git a/source/lib/include/gpu_rocm.h b/source/lib/include/gpu_rocm.h
@@ -24,7 +24,12 @@ inline void nborAssert(hipError_t code, const char *file, int line, bool abort=t
     }
 }
 
+
 namespace deepmd {
+inline void DPGetDeviceCount(int &gpu_num) { hipGetDeviceCount(&gpu_num) ;}
+
+inline hipError_t DPSetDevice(int rank) { return  hipSetDevice(rank); }
+
 template <typename FPTYPE>
 void memcpy_host_to_device(
     FPTYPE * device,