Add ATen overload to AutoGPU. (pytorch#2234)

* Add ATen overload to AutoGPU. Signed-off-by: Edward Z. Yang <ezyang@fb.com> * Use new AutoGPU overload. Signed-off-by: Edward Z. Yang <ezyang@fb.com>

Author: ezyang

torch/csrc/autograd/functions/accumulate_grad.cpp

@@ -19,7 +19,7 @@ auto AccumulateGrad::acc_inplace(std::shared_ptr<Variable>& grad,
  std::shared_ptr<Variable>& new_grad) -> void {
  auto& grad_data = grad->data;
  auto& new_grad_data = new_grad->data;
- AutoGPU guard(grad_data.type().isCuda() ? grad_data.get_device() : -1);
+ AutoGPU guard(grad_data);
 
  if (grad_data.type().isSparse() && !new_grad_data.type().isSparse()) {
  grad->data = new_grad_data + grad_data;

torch/csrc/autograd/functions/basic_ops.cpp

@@ -25,7 +25,7 @@ auto Add::apply(const variable_list& inputs) -> variable_list {
  check_input_variables("Add", inputs, 2);
  auto& input1 = inputs[0]->data;
  auto& input2 = inputs[1]->data;
- AutoGPU guard(input1.type().isCuda() ? input1.get_device() : -1);
+ AutoGPU guard(input1);
 
  at::Tensor output;
  if (input1.type().isSparse()) {

torch/csrc/autograd/functions/batch_normalization.cpp

@@ -39,7 +39,7 @@ auto BatchNormForward::apply(const variable_list& inputs) -> variable_list {
  auto& input = inputs[0];
  auto& weight = inputs[1];
  auto& bias = inputs[2];
- AutoGPU guard(input->data.type().isCuda() ? input->data.get_device() : -1);
+ AutoGPU guard(input->data);
 
  auto num_features = input->data.sizes()[1];
  check_dims_match_num_input_features("running_mean", num_features, running_mean.numel());
@@ -117,7 +117,7 @@ auto BatchNormBackward::apply(const variable_list& grad_outputs) -> variable_lis
  auto weight = weight_var ? weight_var->data : at::Tensor();
  auto bias = bias_var ? bias_var->data : at::Tensor();
 
- AutoGPU guard(input.type().isCuda() ? input.get_device() : -1);
+ AutoGPU guard(input);
 
  bool use_cudnn = false;
 #ifdef WITH_CUDNN

torch/csrc/autograd/functions/convolution.cpp

@@ -148,7 +148,7 @@ auto ConvForward::apply(const variable_list& inputs) -> variable_list {
  check_input_variables("ConvNd", inputs, 3, 2);
  if (is_padding_neg()) throw std::runtime_error("negative padding is not supported");
  if (is_output_padding_neg()) throw std::runtime_error("negative output_padding is not supported");
- AutoGPU guard(inputs[0]->data.type().isCuda() ? inputs[0]->data.get_device() : -1);
+ AutoGPU guard(inputs[0]->data);
  auto input = inputs[0]->data.contiguous();
  auto weight = inputs[1]->data;
  auto bias = inputs[2] ? inputs[2]->data : at::Tensor();
@@ -249,7 +249,7 @@ auto ConvBackward::apply(const variable_list& grad_outputs) -> variable_list {
  auto weight = weight_var->data;
  auto bias = bias_var ? bias_var->data : at::Tensor();
 
- AutoGPU guard(input.type().isCuda() ? input.get_device() : -1);
+ AutoGPU guard(input);
 
  input = input.contiguous();
  auto grad_output = grad_outputs[0]->data.contiguous();

torch/csrc/autograd/functions/tensor.cpp

@@ -14,7 +14,7 @@ auto Identity::apply(const variable_list& inputs) -> variable_list {
 auto Clone::apply(const variable_list& inputs) -> variable_list {
  check_input_variables("Clone", inputs, 1);
  auto& input = inputs[0]->data;
- AutoGPU guard(input.type().isCuda() ? input.get_device() : -1);
+ AutoGPU guard(input);
 
  at::Tensor output = input.clone();
 
@@ -26,7 +26,7 @@ auto Clone::apply(const variable_list& inputs) -> variable_list {
 auto Contiguous::apply(const variable_list& inputs) -> variable_list {
  check_input_variables("Contiguous", inputs, 1);
  auto& input = inputs[0]->data;
- AutoGPU guard(input.type().isCuda() ? input.get_device() : -1);
+ AutoGPU guard(input);
 
  at::Tensor output = input.contiguous();
 
@@ -39,7 +39,7 @@ auto Transpose::apply(const variable_list& inputs) -> variable_list {
  check_input_variables("Transpose", inputs, 1);
 
  auto& input = inputs[0]->data;
- AutoGPU guard(input.type().isCuda() ? input.get_device() : -1);
+ AutoGPU guard(input);
 
  at::Tensor output = input.transpose(dim1, dim2);
 
@@ -52,7 +52,7 @@ auto View::apply(const variable_list& inputs) -> variable_list {
  check_input_variables("View", inputs, 1);
 
  auto& input = inputs[0]->data;
- AutoGPU guard(input.type().isCuda() ? input.get_device() : -1);
+ AutoGPU guard(input);
 
  at::Tensor output = input.view(size);
 
@@ -65,7 +65,7 @@ auto Expand::apply(const variable_list& inputs) -> variable_list {
  check_input_variables("Expand", inputs, 1);
 
  auto& input = inputs[0]->data;
- AutoGPU guard(input.type().isCuda() ? input.get_device() : -1);
+ AutoGPU guard(input);
 
  at::Tensor output = input.expand(size);
 
@@ -78,7 +78,7 @@ auto Narrow::apply(const variable_list& inputs) -> variable_list {
  check_input_variables("Narrow", inputs, 1);
 
  auto& input = inputs[0]->data;
- AutoGPU guard(input.type().isCuda() ? input.get_device() : -1);
+ AutoGPU guard(input);
 
  at::Tensor output = input.narrow(dim, start, size);
 
@@ -94,7 +94,7 @@ auto Cat::apply(const variable_list& inputs) -> variable_list {
  }
 
  auto& input = inputs[0]->data;
- AutoGPU guard(input.type().isCuda() ? input.get_device() : -1);
+ AutoGPU guard(input);
 
  std::vector<at::Tensor> tensors(num_inputs);
  for (int i = 0; i < num_inputs; ++i) {

torch/csrc/utils/auto_gpu.h

@@ -5,6 +5,8 @@
 #include <string>
 #include <stdexcept>
 
+#include <ATen/ATen.h>
+
 #ifdef WITH_CUDA
 #include <cuda.h>
 #include <cuda_runtime.h>
@@ -15,6 +17,10 @@ struct AutoGPU {
  setDevice(device);
  }
 
+ explicit AutoGPU(const at::Tensor& t) {
+ setDevice(t.type().isCuda() ? t.get_device() : -1);
+ }
+
  ~AutoGPU() {
 #ifdef WITH_CUDA
  if (original_device != -1) {