Merge pull request #734 from rdinnager/get_device_capability

Get device capability
mlverse · Nov 20, 2021 · 5e011e7 · 5e011e7
2 parents b528b29 + 650e5af
commit 5e011e7
Show file tree

Hide file tree

Showing 12 changed files with 94 additions and 3 deletions.
diff --git a/.Rbuildignore b/.Rbuildignore
@@ -34,7 +34,7 @@ tools/torchgen/.Rbuildignore
 # ^vignettes/using-autograd\.Rmd
 
 # uncomment below for CRAN submission
-^inst/deps/.*
+#^inst/deps/.*
 ^doc$
 ^Meta$
 ^vignettes/rsconnect*

diff --git a/NAMESPACE b/NAMESPACE
@@ -122,6 +122,7 @@ export(backends_openmp_is_available)
 export(contrib_sort_vertices)
 export(cuda_current_device)
 export(cuda_device_count)
+export(cuda_get_device_capability)
 export(cuda_is_available)
 export(dataloader)
 export(dataloader_make_iter)

diff --git a/R/RcppExports.R b/R/RcppExports.R
@@ -165,6 +165,10 @@ cpp_cuda_current_device <- function() {
     .Call('_torch_cpp_cuda_current_device', PACKAGE = 'torchpkg')
 }
 
+cpp_cuda_get_device_capability <- function(device) {
+    .Call('_torch_cpp_cuda_get_device_capability', PACKAGE = 'torchpkg', device)
+}
+
 cpp_device_type_to_string <- function(device) {
     .Call('_torch_cpp_device_type_to_string', PACKAGE = 'torchpkg', device)
 }

diff --git a/R/cuda.R b/R/cuda.R
@@ -17,4 +17,18 @@ cuda_current_device <- function() {
 #' @export
 cuda_device_count <- function() {
   cpp_cuda_device_count()
+}
+
+#' Returns the major and minor CUDA capability of `device`
+#'
+#' @param device Integer value of the CUDA device to return capabilities of.
+#' 
+#' @export
+cuda_get_device_capability <- function(device = cuda_current_device()) {
+  if(device < 0 | device >= cuda_device_count()) {
+    stop(paste("device must be an integer between 0 and the number of devices minus 1"))
+  }
+  res <- as.integer(cpp_cuda_get_device_capability(device))
+  names(res) <- c("Major", "Minor")
+  res
 }
diff --git a/lantern/CMakeLists.txt b/lantern/CMakeLists.txt
@@ -129,7 +129,8 @@ if(DEFINED ENV{CUDA})
       src/Contrib/SortVertices/sort_vert_kernel.cu
       src/Contrib/SortVertices/sort_vert.cpp
   )
-
+
+  set_source_files_properties(src/Cuda.cpp PROPERTIES COMPILE_DEFINITIONS __NVCC__)
   cuda_add_library(lantern SHARED ${LANTERN_SRC})
 else()
   set(LANTERN_SRC 

diff --git a/lantern/include/lantern/lantern.h b/lantern/include/lantern/lantern.h
@@ -2163,6 +2163,15 @@ HOST_API void* lantern_optional_tensor_value (void* x)
   return ret;
 }
 
+LANTERN_API void* (LANTERN_PTR _lantern_cuda_get_device_capability) (int64_t device);
+HOST_API void* lantern_cuda_get_device_capability (int64_t device)
+{
+  LANTERN_CHECK_LOADED
+  void* ret = _lantern_cuda_get_device_capability(device);
+  LANTERN_HOST_HANDLER;
+  return ret;
+}
+
   /* Autogen Headers -- Start */
   LANTERN_API void* (LANTERN_PTR _lantern__cast_byte_tensor_bool)(void* self, void* non_blocking);
   HOST_API void* lantern__cast_byte_tensor_bool(void* self, void* non_blocking) { LANTERN_CHECK_LOADED void* ret = _lantern__cast_byte_tensor_bool(self, non_blocking); LANTERN_HOST_HANDLER return ret; }
@@ -7749,6 +7758,7 @@ LOAD_SYMBOL(_lantern_OptionalTensorList_size);
 LOAD_SYMBOL(_lantern_OptionalTensorList_at);
 LOAD_SYMBOL(_lantern_OptionalTensorList_at_is_null);
 LOAD_SYMBOL(_lantern_optional_tensor_value);
+LOAD_SYMBOL(_lantern_cuda_get_device_capability);
   /* Autogen Symbols -- Start */
   LOAD_SYMBOL(_lantern__cast_byte_tensor_bool)
   LOAD_SYMBOL(_lantern__cast_char_tensor_bool)

diff --git a/lantern/src/Cuda.cpp b/lantern/src/Cuda.cpp
@@ -3,7 +3,9 @@
 #define LANTERN_BUILD
 
 #include "lantern/lantern.h"
-
+#ifdef __NVCC__
+#include <ATen/cuda/CUDAContext.h>
+#endif
 #include <torch/torch.h>
 
 #include "utils.hpp"
@@ -35,3 +37,16 @@ void _lantern_cuda_show_config()
     std::cout << at::detail::getCUDAHooks().showConfig() << std::endl;
     LANTERN_FUNCTION_END_VOID
 }
+
+void* _lantern_cuda_get_device_capability(int64_t device)
+{
+    LANTERN_FUNCTION_START
+    #ifdef __NVCC__
+        cudaDeviceProp * devprop = at::cuda::getDeviceProperties(device);
+        std::vector<int64_t> cap = {devprop->major, devprop->minor};
+        return (void*) new LanternObject<std::vector<int64_t>>(cap);
+    #else
+        throw std::runtime_error("`cuda_get_device` is only supported on CUDA runtimes.");
+    #endif
+    LANTERN_FUNCTION_END
+}
diff --git a/man/cuda_get_device_capability.Rd b/man/cuda_get_device_capability.Rd
diff --git a/src/RcppExports.cpp b/src/RcppExports.cpp
@@ -478,6 +478,17 @@ BEGIN_RCPP
     return rcpp_result_gen;
 END_RCPP
 }
+// cpp_cuda_get_device_capability
+XPtrTorchvector_int64_t cpp_cuda_get_device_capability(int64_t device);
+RcppExport SEXP _torch_cpp_cuda_get_device_capability(SEXP deviceSEXP) {
+BEGIN_RCPP
+    Rcpp::RObject rcpp_result_gen;
+    Rcpp::RNGScope rcpp_rngScope_gen;
+    Rcpp::traits::input_parameter< int64_t >::type device(deviceSEXP);
+    rcpp_result_gen = Rcpp::wrap(cpp_cuda_get_device_capability(device));
+    return rcpp_result_gen;
+END_RCPP
+}
 // cpp_device_type_to_string
 std::string cpp_device_type_to_string(Rcpp::XPtr<XPtrTorchDevice> device);
 RcppExport SEXP _torch_cpp_device_type_to_string(SEXP deviceSEXP) {
@@ -35172,6 +35183,7 @@ static const R_CallMethodDef CallEntries[] = {
     {"_torch_cpp_cuda_is_available", (DL_FUNC) &_torch_cpp_cuda_is_available, 0},
     {"_torch_cpp_cuda_device_count", (DL_FUNC) &_torch_cpp_cuda_device_count, 0},
     {"_torch_cpp_cuda_current_device", (DL_FUNC) &_torch_cpp_cuda_current_device, 0},
+    {"_torch_cpp_cuda_get_device_capability", (DL_FUNC) &_torch_cpp_cuda_get_device_capability, 1},
     {"_torch_cpp_device_type_to_string", (DL_FUNC) &_torch_cpp_device_type_to_string, 1},
     {"_torch_cpp_device_index_to_int", (DL_FUNC) &_torch_cpp_device_index_to_int, 1},
     {"_torch_cpp_torch_device", (DL_FUNC) &_torch_cpp_torch_device, 2},

diff --git a/src/cuda.cpp b/src/cuda.cpp
@@ -15,3 +15,8 @@ int cpp_cuda_device_count () {
 int64_t cpp_cuda_current_device() {
   return lantern_cuda_current_device();
 }
+
+// [[Rcpp::export]]
+XPtrTorchvector_int64_t cpp_cuda_get_device_capability(int64_t device) {
+  return lantern_cuda_get_device_capability(device);
+}
diff --git a/src/lantern/lantern.h b/src/lantern/lantern.h
@@ -2163,6 +2163,15 @@ HOST_API void* lantern_optional_tensor_value (void* x)
   return ret;
 }
 
+LANTERN_API void* (LANTERN_PTR _lantern_cuda_get_device_capability) (int64_t device);
+HOST_API void* lantern_cuda_get_device_capability (int64_t device)
+{
+  LANTERN_CHECK_LOADED
+  void* ret = _lantern_cuda_get_device_capability(device);
+  LANTERN_HOST_HANDLER;
+  return ret;
+}
+
   /* Autogen Headers -- Start */
   LANTERN_API void* (LANTERN_PTR _lantern__cast_byte_tensor_bool)(void* self, void* non_blocking);
   HOST_API void* lantern__cast_byte_tensor_bool(void* self, void* non_blocking) { LANTERN_CHECK_LOADED void* ret = _lantern__cast_byte_tensor_bool(self, non_blocking); LANTERN_HOST_HANDLER return ret; }
@@ -7749,6 +7758,7 @@ LOAD_SYMBOL(_lantern_OptionalTensorList_size);
 LOAD_SYMBOL(_lantern_OptionalTensorList_at);
 LOAD_SYMBOL(_lantern_OptionalTensorList_at_is_null);
 LOAD_SYMBOL(_lantern_optional_tensor_value);
+LOAD_SYMBOL(_lantern_cuda_get_device_capability);
   /* Autogen Symbols -- Start */
   LOAD_SYMBOL(_lantern__cast_byte_tensor_bool)
   LOAD_SYMBOL(_lantern__cast_char_tensor_bool)

diff --git a/tests/testthat/test-cuda.R b/tests/testthat/test-cuda.R
@@ -6,6 +6,11 @@ test_that("cuda", {
   expect_true(cuda_device_count() > 0)
   expect_true(cuda_current_device() >= 0)
   expect_true(cuda_is_available())
+
+  capability <- cuda_get_device_capability(cuda_current_device())
+  expect_type(capability, "integer")
+  expect_length(capability, 2)
+  expect_error(cuda_get_device_capability(cuda_device_count() + 1), "device must be an integer between 0 and")
 })
 
 test_that("cuda tensors", {