set stream everytime when we get a cuDNN handle (pytorch#31541)

Summary: cudnn version of pytorch#31537 pytorch#31532 is a quick fix and this is a bigger change. This would deprecate pytorch#31532, but we could also merge pytorch#31532 first for a quick fix and then work on this later. Pull Request resolved: pytorch#31541 Differential Revision: D19206753 Pulled By: ngimel fbshipit-source-id: 3352f923d13a9baf0971f64f8b7ce03e9a8b42b1

Author: zasdfgbnm

aten/src/ATen/cudnn/Descriptors.h

@@ -191,7 +191,6 @@ struct TORCH_CUDA_API DropoutDescriptor
  AT_ASSERT(options.device().type() == kCUDA);
  AT_ASSERT(options.dtype() == kByte);
  state = at::empty({static_cast<int64_t>(state_size)}, options);
- setCuDNNStreamToCurrent();
  AT_CUDNN_CHECK(cudnnSetDropoutDescriptor(mut_desc(), handle, dropout, state.data_ptr(), state_size, seed));
  }
 
@@ -202,7 +201,6 @@ struct TORCH_CUDA_API DropoutDescriptor
  void *state_ptr = state.data_ptr();
  size_t state_size = state.size(0);
  // NB: The seed doesn't actually matter, so we give a dummy value
- setCuDNNStreamToCurrent();
  AT_CUDNN_CHECK(cudnnRestoreDropoutDescriptor(mut_desc(), handle, dropout, state_ptr, state_size, 0 /* seed */));
  }
 

aten/src/ATen/cudnn/Handle.cpp

@@ -1,5 +1,6 @@
 #include <ATen/cudnn/Handle.h>
 #include <ATen/cuda/detail/DeviceThreadHandles.h>
+#include <c10/cuda/CUDAStream.h>
 
 namespace at { namespace native {
 namespace {
@@ -40,7 +41,9 @@ cudnnHandle_t getCudnnHandle()
  if (!myPoolWindow)
  myPoolWindow.reset(pool.newPoolWindow());
 
- return myPoolWindow->reserve(device);
+ auto handle = myPoolWindow->reserve(device);
+ AT_CUDNN_CHECK(cudnnSetStream(handle, c10::cuda::getCurrentCUDAStream()));
+ return handle;
 }
 
 }} // namespace at::native

aten/src/ATen/cudnn/Utils.h

@@ -8,11 +8,6 @@
 
 namespace at { namespace native {
 
-inline void setCuDNNStreamToCurrent() {
- // TODO: Should getCurrentStream be a method on Context?
- AT_CUDNN_CHECK(cudnnSetStream(getCudnnHandle(), at::cuda::getCurrentCUDAStream()));
-}
-
 // cuDNN has a buggy check for tensor being contiguous (that is, it does
 // not ignore stride for dimension that is equal to 0). This function
 // makes tensors which have zero stride contiguous, by setting the

aten/src/ATen/native/cudnn/AffineGridGenerator.cpp

@@ -52,8 +52,6 @@ Tensor cudnn_affine_grid_generator_forward(
  const Tensor& theta_t,
  int64_t N, int64_t C, int64_t H, int64_t W)
 {
- setCuDNNStreamToCurrent();
-
  TensorArg theta{ theta_t.contiguous(), "theta", 1 };
  CheckedFrom c = "cudnn_affine_grid_generator_forward";
  checkContiguous(c, theta);
@@ -75,8 +73,6 @@ Tensor cudnn_affine_grid_generator_backward(
  const Tensor& grad_grid_t,
  int64_t N, int64_t C, int64_t H, int64_t W)
 {
- setCuDNNStreamToCurrent();
-
  TensorArg grad_grid{ grad_grid_t.contiguous(), "grad_grid", 1 };
  CheckedFrom c = "cudnn_affine_grid_generator_backward";
  checkContiguous(c, grad_grid);

aten/src/ATen/native/cudnn/BatchNorm.cpp

@@ -60,7 +60,6 @@ std::tuple<Tensor, Tensor, Tensor, Tensor> cudnn_batch_norm(
  running_mean{ running_mean_t, "running_mean", 4 },
  running_var{ running_var_t, "running_var", 5 };
  CheckedFrom c = "cudnn_batch_norm";
- setCuDNNStreamToCurrent();
 
  checkAllDefined(c, {input, weight, bias});
  if (!training) {
@@ -233,7 +232,6 @@ std::tuple<Tensor, Tensor, Tensor> cudnn_batch_norm_backward(
  save_var{ save_var_t, "save_var", 5 },
  reserve{ reserveSpace, "reserve_space", 6 };
  CheckedFrom c = "cudnn_batch_norm_backward";
- setCuDNNStreamToCurrent();
 
  checkAllDefined(c, {input, grad_output, weight, save_mean, save_var});
  checkAllSameGPU(c, {input, grad_output, weight, save_mean, save_var});

aten/src/ATen/native/cudnn/Conv.cpp

@@ -766,7 +766,6 @@ void cudnn_convolution_add_bias_(CheckedFrom c, const TensorArg& output, const T
 // responsibility:
 // - Things that happen in at::Tensor
 // - TensorArg allocation
-// - setCuDNNStreamToCurrent
 // - Things that happen in TensorArg
 // - Check arguments (type, GPU, shape)
 //
@@ -918,7 +917,6 @@ Tensor cudnn_convolution(
  TensorArg input { input_t, "input", 1 },
  weight { weight_t, "weight", 2 },
  bias { bias_t, "bias", 3 };
- setCuDNNStreamToCurrent();
  CheckedFrom c = "cudnn_convolution";
  auto output_t = cudnn_convolution_forward(
  c, input, weight, padding, stride, dilation, groups, benchmark, deterministic);
@@ -937,7 +935,6 @@ Tensor cudnn_convolution_transpose_backward_input(
 {
  TensorArg grad_output { grad_output_t, "grad_output", 1 },
  weight { weight_t, "weight", 2 };
- setCuDNNStreamToCurrent();
  return cudnn_convolution_forward(
  "cudnn_convolution_transpose_backward_input",
  grad_output, weight, padding, stride, dilation, groups, benchmark, deterministic);
@@ -1062,7 +1059,6 @@ Tensor cudnn_convolution_backward_input(
 {
  TensorArg grad_output{ grad_output_t, "grad_output", 1 },
  weight{ weight_t, "weight", 2 };
- setCuDNNStreamToCurrent();
  return cudnn_convolution_backward_input(
  "cudnn_convolution_backward_input",
  input_size, grad_output, weight,
@@ -1192,7 +1188,6 @@ Tensor cudnn_convolution_backward_weight(
 {
  TensorArg grad_output{ grad_output_t, "grad_output", 1 },
  input{ input_t, "input", 2 };
- setCuDNNStreamToCurrent();
  return cudnn_convolution_backward_weight(
  "cudnn_convolution_backward_weight",
  weight_size, grad_output, input,
@@ -1208,7 +1203,6 @@ Tensor cudnn_convolution_transpose_backward_weight(
 {
  TensorArg grad_output{ grad_output_t, "grad_output", 1 },
  input{ input_t, "input", 2 };
- setCuDNNStreamToCurrent();
  return cudnn_convolution_backward_weight(
  "cudnn_convolution_backward_weight",
  weight_size, input, grad_output,
@@ -1225,7 +1219,6 @@ Tensor cudnn_convolution_backward_bias(
  const Tensor& grad_output_t)
 {
  TensorArg grad_output{ grad_output_t, "grad_output", 1 };
- setCuDNNStreamToCurrent();
 
  auto grad_bias_t = at::empty(
  { grad_output->size(output_channels_dim) }, grad_output->options());

aten/src/ATen/native/cudnn/GridSampler.cpp

@@ -69,7 +69,6 @@ Tensor cudnn_grid_sampler_forward(
  TensorArg input{ contiguousIfZeroInStrides(input_t), "input", 1 },
  grid{ grid_t.contiguous(), "grid", 2 };
  CheckedFrom c = "cudnn_grid_sampler_forward";
- setCuDNNStreamToCurrent();
  checkAllSameGPU(c, {input, grid});
  checkAllSameType(c, {input, grid});
  checkGridSize(c, grid, input);
@@ -108,7 +107,6 @@ std::tuple<Tensor, Tensor> cudnn_grid_sampler_backward(
  grid{ grid_t.contiguous(), "grid", 2 },
  grad_output{ contiguousIfZeroInStrides(grad_output_t), "grad_output", 3 };
  CheckedFrom c = "cudnn_grid_sampler_backward";
- setCuDNNStreamToCurrent();
  checkAllSameGPU(c, {input, grad_output, grid});
  checkGridSize(c, grid, input);
  checkDim(c, input, 4);

aten/src/ATen/native/cudnn/LossCTC.cpp

@@ -86,7 +86,6 @@ std::tuple<Tensor, Tensor> _cudnn_ctc_loss(const Tensor& log_probs_t, const Tens
  std::vector<int> input_lengths(input_lengths_.begin(), input_lengths_.end());
  std::vector<int> target_lengths(target_lengths_.begin(), target_lengths_.end());
 
- setCuDNNStreamToCurrent();
  TORCH_CHECK(BLANK == 0, "blank must be label 0 for cudnn_ctc_loss");
  // checked in dispatch:
  // assert other conditions for cudnnCTCLoss: all label lengths <= 256

aten/src/ATen/native/cudnn/RNN.cpp

@@ -778,7 +778,6 @@ std::tuple<Tensor, Tensor, Tensor, Tensor, Tensor> _cudnn_rnn(
  &reserve_size
  ));
  reserve = at::empty(reserve_size, input.options().dtype(kByte));
- setCuDNNStreamToCurrent();
  AT_CUDNN_CHECK(cudnnRNNForwardTraining(
  handle,
  descs.rnn_desc.desc(),
@@ -795,7 +794,6 @@ std::tuple<Tensor, Tensor, Tensor, Tensor, Tensor> _cudnn_rnn(
  ));
  } else { // inference
  reserve = at::empty({0}, input.options().dtype(kByte));
- setCuDNNStreamToCurrent();
  AT_CUDNN_CHECK(cudnnRNNForwardInference(
  handle,
  descs.rnn_desc.desc(),
@@ -914,7 +912,6 @@ std::tuple<Tensor, Tensor, Tensor> _cudnn_rnn_backward_input(
  ));
  // TODO: put this in the correct device???
  Tensor workspace = at::empty(workspace_size, input.options().dtype(kByte));
- setCuDNNStreamToCurrent();
  AT_CUDNN_CHECK(cudnnRNNBackwardData(
  handle,
  descs.rnn_desc.desc(),
@@ -1018,7 +1015,6 @@ std::vector<Tensor> _cudnn_rnn_backward_weight(
  &workspace_size
  ));
  Tensor workspace = at::empty(workspace_size, input.options().dtype(kByte));
- setCuDNNStreamToCurrent();
  AT_CUDNN_CHECK(cudnnRNNBackwardWeights(
  handle,
  descs.rnn_desc.desc(),