From a09195376d90d551c42f2748a4d60d44eec4eefd Mon Sep 17 00:00:00 2001 From: Raimondas Galvelis Date: Tue, 2 Aug 2022 15:40:45 +0200 Subject: [PATCH 1/3] Factor out the CUDA accessors --- src/pytorch/common/accessor.cuh | 14 ++++++++++++++ src/pytorch/neighbors/getNeighborPairsCUDA.cu | 12 ++---------- 2 files changed, 16 insertions(+), 10 deletions(-) create mode 100644 src/pytorch/common/accessor.cuh diff --git a/src/pytorch/common/accessor.cuh b/src/pytorch/common/accessor.cuh new file mode 100644 index 0000000..81957b3 --- /dev/null +++ b/src/pytorch/common/accessor.cuh @@ -0,0 +1,14 @@ +#ifndef NNPOPS_ACCESSOR_H +#define NNPOPS_ACCESSOR_H + +#include + +template + using Accessor = torch::PackedTensorAccessor32; + +template +inline Accessor get_accessor(const torch::Tensor& tensor) { + return tensor.packed_accessor32(); +}; + +#endif \ No newline at end of file diff --git a/src/pytorch/neighbors/getNeighborPairsCUDA.cu b/src/pytorch/neighbors/getNeighborPairsCUDA.cu index b87a221..bdc9aab 100644 --- a/src/pytorch/neighbors/getNeighborPairsCUDA.cu +++ b/src/pytorch/neighbors/getNeighborPairsCUDA.cu @@ -4,6 +4,8 @@ #include #include +#include "common/accessor.cuh" + using c10::cuda::CUDAStreamGuard; using c10::cuda::getCurrentCUDAStream; using std::make_tuple; @@ -14,21 +16,11 @@ using torch::autograd::tensor_list; using torch::empty; using torch::full; using torch::kInt32; -using torch::PackedTensorAccessor32; -using torch::RestrictPtrTraits; using torch::Scalar; using torch::Tensor; using torch::TensorOptions; using torch::zeros; -template - using Accessor = PackedTensorAccessor32; - -template -inline Accessor get_accessor(const Tensor& tensor) { - return tensor.packed_accessor32(); -}; - template __device__ __forceinline__ scalar_t sqrt_(scalar_t x) {}; template<> __device__ __forceinline__ float sqrt_(float x) { return ::sqrtf(x); }; template<> __device__ __forceinline__ double sqrt_(double x) { return ::sqrt(x); }; From e0ce04a31eb46eb1850c6b2780ad8983c23e84b0 Mon Sep 17 00:00:00 2001 From: Raimondas Galvelis Date: Tue, 2 Aug 2022 15:48:31 +0200 Subject: [PATCH 2/3] Add an include path --- CMakeLists.txt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 05c44da..894d792 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -31,7 +31,8 @@ set(SRC_FILES src/ani/CpuANISymmetryFunctions.cpp set(LIBRARY ${NAME}PyTorch) add_library(${LIBRARY} SHARED ${SRC_FILES}) set_property(TARGET ${LIBRARY} PROPERTY CXX_STANDARD 14) -target_include_directories(${LIBRARY} PRIVATE ${PYTHON_INCLUDE_DIRS} src/ani src/schnet) +target_include_directories(${LIBRARY} PRIVATE ${PYTHON_INCLUDE_DIRS} + src/ani src/pytorch src/schnet) target_link_libraries(${LIBRARY} ${TORCH_LIBRARIES} ${PYTHON_LIBRARIES}) if(ENABLE_CUDA) set_property(TARGET ${LIBRARY} PROPERTY CUDA_STANDARD 14) From 322a9287cf1826006423fbe0824646d24adf9b98 Mon Sep 17 00:00:00 2001 From: Raimondas Galvelis Date: Tue, 2 Aug 2022 15:50:37 +0200 Subject: [PATCH 3/3] Factor out atomicAdd --- src/pytorch/common/atomicAdd.cuh | 24 +++++++++++++++++++ src/pytorch/neighbors/getNeighborPairsCUDA.cu | 16 +------------ 2 files changed, 25 insertions(+), 15 deletions(-) create mode 100644 src/pytorch/common/atomicAdd.cuh diff --git a/src/pytorch/common/atomicAdd.cuh b/src/pytorch/common/atomicAdd.cuh new file mode 100644 index 0000000..80b79ce --- /dev/null +++ b/src/pytorch/common/atomicAdd.cuh @@ -0,0 +1,24 @@ +#ifndef NNPOPS_ATOMICADD_H +#define NNPOPS_ATOMICADD_H + +/* +Implement atomicAdd with double precision numbers for pre-Pascal GPUs. +Taken from https://stackoverflow.com/questions/37566987/cuda-atomicadd-for-doubles-definition-error +NOTE: remove when the support of CUDA 11 is dropped. +*/ + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 600 +__device__ double atomicAdd(double* address, double val) +{ + unsigned long long int* address_as_ull = (unsigned long long int*)address; + unsigned long long int old = *address_as_ull, assumed; + do { + assumed = old; + old = atomicCAS(address_as_ull, assumed, + __double_as_longlong(val + __longlong_as_double(assumed))); + } while (assumed != old); + return __longlong_as_double(old); +} +#endif + +#endif \ No newline at end of file diff --git a/src/pytorch/neighbors/getNeighborPairsCUDA.cu b/src/pytorch/neighbors/getNeighborPairsCUDA.cu index bdc9aab..33e0058 100644 --- a/src/pytorch/neighbors/getNeighborPairsCUDA.cu +++ b/src/pytorch/neighbors/getNeighborPairsCUDA.cu @@ -5,6 +5,7 @@ #include #include "common/accessor.cuh" +#include "common/atomicAdd.cuh" using c10::cuda::CUDAStreamGuard; using c10::cuda::getCurrentCUDAStream; @@ -25,21 +26,6 @@ template __device__ __forceinline__ scalar_t sqrt_(scalar_t template<> __device__ __forceinline__ float sqrt_(float x) { return ::sqrtf(x); }; template<> __device__ __forceinline__ double sqrt_(double x) { return ::sqrt(x); }; -// Support pre-Pascal GPUs. Remove when the support of CUDA 11 is dropped. -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 600 -__device__ double atomicAdd(double* address, double val) -{ - unsigned long long int* address_as_ull = (unsigned long long int*)address; - unsigned long long int old = *address_as_ull, assumed; - do { - assumed = old; - old = atomicCAS(address_as_ull, assumed, - __double_as_longlong(val + __longlong_as_double(assumed))); - } while (assumed != old); - return __longlong_as_double(old); -} -#endif - template __global__ void forward_kernel( const int32_t num_all_pairs, const Accessor positions,