Some api function require cuda11.2+

nliber · Jul 27, 2023 · 3f19a97 · 3f19a97
1 parent 03039fe
commit 3f19a97
Showing 1 changed file with 17 additions and 14 deletions.
diff --git a/core/src/Cuda/Kokkos_Cuda_Instance.hpp b/core/src/Cuda/Kokkos_Cuda_Instance.hpp
@@ -273,13 +273,6 @@ class CudaInternal {
     return cudaFree(devPtr);
   }
 
-  template <bool setCudaDevice = true>
-  cudaError_t cuda_free_async_wrapper(void* devPtr,
-                                      cudaStream_t hStream) const {
-    if (setCudaDevice) set_cuda_device();
-    return cudaFreeAsync(devPtr, hStream);
-  }
-
   template <bool setCudaDevice = true>
   cudaError_t cuda_free_host_wrapper(void* ptr) const {
     if (setCudaDevice) set_cuda_device();
@@ -383,13 +376,6 @@ class CudaInternal {
     return cudaMalloc(devPtr, size);
   }
 
-  template <bool setCudaDevice = true>
-  cudaError_t cuda_malloc_async_wrapper(void** devPtr, size_t size,
-                                        cudaStream_t hStream) const {
-    if (setCudaDevice) set_cuda_device();
-    return cudaMallocAsync(devPtr, size, hStream);
-  }
-
   template <bool setCudaDevice = true>
   cudaError_t cuda_malloc_host_wrapper(void** ptr, size_t size) const {
     if (setCudaDevice) set_cuda_device();
@@ -481,6 +467,23 @@ class CudaInternal {
     return cudaStreamSynchronize(stream);
   }
 
+  // The following are only available for cuda 11.2 and greater
+#if (defined(KOKKOS_ENABLE_IMPL_CUDA_MALLOC_ASYNC) && CUDART_VERSION >= 11020)
+  template <bool setCudaDevice = true>
+  cudaError_t cuda_malloc_async_wrapper(void** devPtr, size_t size,
+                                        cudaStream_t hStream) const {
+    if (setCudaDevice) set_cuda_device();
+    return cudaMallocAsync(devPtr, size, hStream);
+  }
+
+  template <bool setCudaDevice = true>
+  cudaError_t cuda_free_async_wrapper(void* devPtr,
+                                      cudaStream_t hStream) const {
+    if (setCudaDevice) set_cuda_device();
+    return cudaFreeAsync(devPtr, hStream);
+  }
+#endif
+
   // C++ API routines
   template <typename T, bool setCudaDevice = true>
   cudaError_t cuda_func_get_attributes_wrapper(cudaFuncAttributes* attr,