Check that all tensors are on the same GPU in cuDNN bindings

Author: apaszke

torch/csrc/cudnn/BatchNorm.cpp

@@ -62,6 +62,8 @@ void cudnn_batch_norm_forward(
     THVoidTensor* save_mean, THVoidTensor* save_var, bool training,
     double exponential_average_factor, double epsilon)
 {
+  assertSameGPU(dataType, input, output, weight, bias, running_mean, running_var,
+      save_mean, save_var);
   cudnnBatchNormMode_t mode;
   if (input->nDimension == 2) {
     mode = CUDNN_BATCHNORM_PER_ACTIVATION;
@@ -120,6 +122,8 @@ void cudnn_batch_norm_backward(
     THVoidTensor* save_mean, THVoidTensor* save_var, bool training,
     double epsilon)
 {
+  assertSameGPU(dataType, input, grad_output, grad_input, grad_weight, grad_bias, weight,
+      running_mean, running_var, save_mean, save_var);
   cudnnBatchNormMode_t mode;
   if (input->nDimension == 2) {
     mode = CUDNN_BATCHNORM_PER_ACTIVATION;

torch/csrc/cudnn/Conv.cpp

@@ -285,6 +285,7 @@ void cudnn_convolution_forward(
     THVoidTensor* input, THVoidTensor* weight, THVoidTensor* output,
     Convolution* info, bool benchmark)
 {
+  assertSameGPU(dataType, input, weight, output);
   int groups = info->groups;
 
   cudnnConvolutionFwdAlgo_t fwdAlg;
@@ -309,6 +310,7 @@ void cudnn_convolution_add_bias(
     THVoidTensor* bias, THVoidTensor* output,
     Convolution* info)
 {
+  assertSameGPU(dataType, bias, output);
   CHECK_ARG(output->nDimension <= 5);
   TensorDescriptor& bdesc = info->bdesc;
 
@@ -329,6 +331,7 @@ void cudnn_convolution_backward_data(
     THVoidTensor* gradOutput, THVoidTensor* gradInput, THVoidTensor* weight,
     Convolution* info, bool benchmark)
 {
+  assertSameGPU(dataType, gradOutput, gradInput, weight);
   int groups = info->params.groups;
 
   cudnnConvolutionBwdDataAlgo_t bwdDataAlg;
@@ -353,6 +356,7 @@ void cudnn_convolution_backward_filter(
     THVoidTensor* gradOutput, THVoidTensor* input, THVoidTensor* gradWeight,
     Convolution* info, bool benchmark)
 {
+  assertSameGPU(dataType, gradOutput, input, gradWeight);
   int groups = info->params.groups;
 
   cudnnConvolutionBwdFilterAlgo_t bwdFilterAlg;
@@ -380,6 +384,7 @@ void cudnn_convolution_backward_bias(
     THCState* state, cudnnHandle_t handle, cudnnDataType_t dataType,
     THVoidTensor* gradOutput, THVoidTensor* gradBias, Convolution* info)
 {
+  assertSameGPU(dataType, gradOutput, gradBias);
   Constant one(dataType, 1);
   Constant zero(dataType, 0);
   void* gradOutput_ptr = tensorPointer(dataType, gradOutput, 0, 1, 0);

torch/csrc/cudnn/Exceptions.h

@@ -1,17 +1,42 @@
 #ifndef THP_CUDNN_EXCEPTIONS_INC
 #define THP_CUDNN_EXCEPTIONS_INC
 
+#include <THC/THC.h>
 #include <cudnn.h>
 #include <string>
 #include <stdexcept>
 #include <sstream>
 
+#include "Types.h"
 
 #define CHECK_ARG(cond) _CHECK_ARG(cond, #cond, __FILE__, __LINE__)
 
+extern THCState* state;
 
 namespace torch { namespace cudnn {
 
+template<typename ...T>
+void assertSameGPU(cudnnDataType_t dataType, T* ... tensors) {
+  static_assert(std::is_same<THVoidTensor, typename std::common_type<T...>::type>::value,
+      "all arguments to assertSameGPU have to be THVoidTensor*");
+  int is_same;
+  if (dataType == CUDNN_DATA_FLOAT) {
+    is_same = THCudaTensor_checkGPU(state, sizeof...(T),
+        reinterpret_cast<THCudaTensor*>(tensors)...);
+  } else if (dataType == CUDNN_DATA_HALF) {
+    is_same = THCudaHalfTensor_checkGPU(state, sizeof...(T),
+        reinterpret_cast<THCudaHalfTensor*>(tensors)...);
+  } else if (dataType == CUDNN_DATA_DOUBLE) {
+    is_same = THCudaDoubleTensor_checkGPU(state, sizeof...(T),
+        reinterpret_cast<THCudaDoubleTensor*>(tensors)...);
+  } else {
+    throw std::runtime_error("unknown cuDNN data type");
+  }
+  if (!is_same) {
+    throw std::runtime_error("tensors are on different GPUs");
+  }
+}
+
 class cudnn_exception : public std::runtime_error {
 public:
   cudnnStatus_t status;