Allow specification for GPU device index

src/ctorch.cpp

@@ -29,30 +29,41 @@ constexpr auto get_dtype(torch_data_t dtype)
   }
 }
 
-constexpr auto get_device(torch_device_t device)
+const auto get_device(torch_device_t device_type, int device_index)
 {
-  switch (device) {
+  switch (device_type) {
   case torch_kCPU:
-    return torch::kCPU;
+    if (device_index != -1) {
+      std::cerr << "[ERROR]: device index unsupported for CPU-only runs"
+                << std::endl;
+    }
+    return torch::Device(torch::kCPU);
   case torch_kCUDA:
-    return torch::kCUDA;
+    if (device_index >= 0 && device_index < torch::cuda::device_count()) {
+      return torch::Device(torch::kCUDA, device_index);
+    } else {
+      std::cerr << "[ERROR]: invalid device index " << device_index
+                << " for device count " << torch::cuda::device_count()
+                << ", using zero instead" << std::endl;
+      return torch::Device(torch::kCUDA);
+    }
   default:
     std::cerr << "[ERROR]: unknown device type, setting to torch_kCPU"
               << std::endl;
-    return torch::kCPU;
+    return torch::Device(torch::kCPU);
   }
 }
 
 torch_tensor_t torch_zeros(int ndim, const int64_t* shape, torch_data_t dtype,
-                           torch_device_t device)
+                           torch_device_t device_type, int device_index = -1)
 {
   torch::Tensor* tensor = nullptr;
   try {
     // This doesn't throw if shape and dimensions are incompatible
     c10::IntArrayRef vshape(shape, ndim);
     tensor = new torch::Tensor;
     *tensor = torch::zeros(
-        vshape, torch::dtype(get_dtype(dtype))).to(get_device(device));
+        vshape, torch::dtype(get_dtype(dtype))).to(get_device(device_type, device_index));
   } catch (const torch::Error& e) {
     std::cerr << "[ERROR]: " << e.msg() << std::endl;
     delete tensor;
@@ -66,15 +77,15 @@ torch_tensor_t torch_zeros(int ndim, const int64_t* shape, torch_data_t dtype,
 }
 
 torch_tensor_t torch_ones(int ndim, const int64_t* shape, torch_data_t dtype,
-                          torch_device_t device)
+                          torch_device_t device_type, int device_index = -1)
 {
   torch::Tensor* tensor = nullptr;
   try {
     // This doesn't throw if shape and dimensions are incompatible
     c10::IntArrayRef vshape(shape, ndim);
     tensor = new torch::Tensor;
     *tensor = torch::ones(
-        vshape, torch::dtype(get_dtype(dtype))).to(get_device(device));
+        vshape, torch::dtype(get_dtype(dtype))).to(get_device(device_type, device_index));
   } catch (const torch::Error& e) {
     std::cerr << "[ERROR]: " << e.msg() << std::endl;
     delete tensor;
@@ -88,15 +99,15 @@ torch_tensor_t torch_ones(int ndim, const int64_t* shape, torch_data_t dtype,
 }
 
 torch_tensor_t torch_empty(int ndim, const int64_t* shape, torch_data_t dtype,
-                           torch_device_t device)
+                           torch_device_t device_type, int device_index = -1)
 {
   torch::Tensor* tensor = nullptr;
   try {
     // This doesn't throw if shape and dimensions are incompatible
     c10::IntArrayRef vshape(shape, ndim);
     tensor = new torch::Tensor;
     *tensor = torch::empty(
-        vshape, torch::dtype(get_dtype(dtype))).to(get_device(device));
+        vshape, torch::dtype(get_dtype(dtype))).to(get_device(device_type, device_index));
   } catch (const torch::Error& e) {
     std::cerr << "[ERROR]: " << e.msg() << std::endl;
     delete tensor;
@@ -113,7 +124,7 @@ torch_tensor_t torch_empty(int ndim, const int64_t* shape, torch_data_t dtype,
 // data
 torch_tensor_t torch_from_blob(void* data, int ndim, const int64_t* shape,
                                const int64_t* strides, torch_data_t dtype,
-                               torch_device_t device)
+                               torch_device_t device_type, int device_index = -1)
 {
   torch::Tensor* tensor = nullptr;
 
@@ -124,7 +135,7 @@ torch_tensor_t torch_from_blob(void* data, int ndim, const int64_t* shape,
     tensor = new torch::Tensor;
     *tensor = torch::from_blob(
         data, vshape, vstrides,
-        torch::dtype(get_dtype(dtype))).to(get_device(device));
+        torch::dtype(get_dtype(dtype))).to(get_device(device_type, device_index));
 
   } catch (const torch::Error& e) {
     std::cerr << "[ERROR]: " << e.msg() << std::endl;
@@ -144,18 +155,26 @@ void torch_tensor_print(const torch_tensor_t tensor)
   std::cout << *t << std::endl;
 }
 
+int torch_tensor_get_device_index(const torch_tensor_t tensor)
+{
+  auto t = reinterpret_cast<torch::Tensor*>(tensor);
+  return t->device().index();
+}
+
 void torch_tensor_delete(torch_tensor_t tensor)
 {
   auto t = reinterpret_cast<torch::Tensor*>(tensor);
   delete t;
 }
 
-torch_jit_script_module_t torch_jit_load(const char* filename)
+torch_jit_script_module_t torch_jit_load(const char* filename,
+                                         const torch_device_t device_type = torch_kCPU,
+                                         const int device_index = -1)
 {
   torch::jit::script::Module* module = nullptr;
   try {
     module = new torch::jit::script::Module;
-    *module = torch::jit::load(filename);
+    *module = torch::jit::load(filename, get_device(device_type, device_index));
   } catch (const torch::Error& e) {
     std::cerr << "[ERROR]: " << e.msg() << std::endl;
     delete module;

src/ctorch.h

@@ -37,30 +37,36 @@ typedef enum { torch_kCPU, torch_kCUDA } torch_device_t;
  * @param number of dimensions of the Tensor
  * @param shape of the Tensor
  * @param data type of the elements of the Tensor
- * @param device used (cpu, CUDA, etc.)
+ * @param device type used (cpu, CUDA, etc.)
+ * @param device index for the CUDA case
  */
 EXPORT_C torch_tensor_t torch_zeros(int ndim, const int64_t* shape,
-                                    torch_data_t dtype, torch_device_t device);
+                                    torch_data_t dtype, torch_device_t device_type,
+                                    int device_index);
 
 /**
  * Function to generate a Torch Tensor of ones
  * @param number of dimensions of the Tensor
  * @param shape of the Tensor
  * @param data type of the elements of the Tensor
- * @param device used (cpu, CUDA, etc.)
+ * @param device type used (cpu, CUDA, etc.)
+ * @param device index for the CUDA case
  */
 EXPORT_C torch_tensor_t torch_ones(int ndim, const int64_t* shape,
-                                   torch_data_t dtype, torch_device_t device);
+                                   torch_data_t dtype, torch_device_t device_type,
+                                   int device_index);
 
 /**
  * Function to generate an empty Torch Tensor
  * @param number of dimensions of the Tensor
  * @param shape of the Tensor
  * @param data type of the elements of the Tensor
- * @param device used (cpu, CUDA, etc.)
+ * @param device type used (cpu, CUDA, etc.)
+ * @param device index for the CUDA case
  */
 EXPORT_C torch_tensor_t torch_empty(int ndim, const int64_t* shape,
-                                    torch_data_t dtype, torch_device_t device);
+                                    torch_data_t dtype, torch_device_t device_type,
+                                    int device_index);
 
 /**
  * Function to create a Torch Tensor from memory location given extra information
@@ -69,21 +75,30 @@ EXPORT_C torch_tensor_t torch_empty(int ndim, const int64_t* shape,
  * @param shape of the Tensor
  * @param strides to take through data
  * @param data type of the elements of the Tensor
- * @param device used (cpu, CUDA, etc.)
+ * @param device type used (cpu, CUDA, etc.)
+ * @param device index for the CUDA case
  * @return Torch Tensor interpretation of the data pointed at
  */
 EXPORT_C torch_tensor_t torch_from_blob(void* data, int ndim,
                                         const int64_t* shape,
                                         const int64_t* strides,
                                         torch_data_t dtype,
-                                        torch_device_t device);
+                                        torch_device_t device_type,
+                                        int device_index);
 
 /**
  * Function to print out a Torch Tensor
  * @param Torch Tensor to print
  */
 EXPORT_C void torch_tensor_print(const torch_tensor_t tensor);
 
+/**
+ * Function to determine the device index of a Torch Tensor
+ * @param Torch Tensor to determine the device index of
+ * @return device index of the Torch Tensor
+ */
+EXPORT_C int torch_tensor_get_device_index(const torch_tensor_t tensor);
+
 /**
  * Function to delete a Torch Tensor to clean up
  * @param Torch Tensor to delete
@@ -97,9 +112,13 @@ EXPORT_C void torch_tensor_delete(torch_tensor_t tensor);
 /**
  * Function to load in a Torch model from a TorchScript file and store in a Torch Module
  * @param filename where TorchScript description of model is stored
+ * @param device type used (cpu, CUDA, etc.)
+ * @param device index for the CUDA case
  * @return Torch Module loaded in from file
  */
-EXPORT_C torch_jit_script_module_t torch_jit_load(const char* filename);
+EXPORT_C torch_jit_script_module_t torch_jit_load(const char* filename,
+                                                  const torch_device_t device_type,
+                                                  const int device_index);
 
 /**
  * Function to run the `forward` method of a Torch Module