aten/src/ATen/native/cudnn/Conv.cpp
// ---------------------------------------------------------------------
//
// Convolution backward / Transposed convolution forward
//
// ---------------------------------------------------------------------
Convolution backward / Transposed convolution forward
Convolution backward / Transposed convolution forward
void raw_cudnn_convolution_backward_input_out(
    const at::Tensor& grad_input,
    const at::Tensor& grad_output,
    const at::Tensor& weight,
    IntList padding, IntList stride, IntList dilation, int64_t groups,
    bool benchmark, bool deterministic) {
  auto dataType = getCudnnDataType(grad_output);
  ConvolutionArgs args{ grad_input, grad_output, weight };
  args.handle = getCudnnHandle();
  setConvolutionParams(&args.params, grad_input, weight, padding, stride, dilation, groups, deterministic);
  args.idesc.set(grad_input);
  args.wdesc.set(weight);
  args.odesc.set(grad_output);
  args.cdesc.set(dataType, grad_output.dim() - 2, args.params.padding, args.params.stride, args.params.dilation, args.params.groups);
  cudnnConvolutionBwdDataAlgo_t bwdDataAlg;
  Workspace workspace = chooseAlgorithm(args, benchmark, &bwdDataAlg);
  Constant one(dataType, 1);
  Constant zero(dataType, 0);
  AT_CUDNN_CHECK(cudnnConvolutionBackwardData(
      args.handle,
      &one, args.wdesc.desc(), weight.data_ptr(),
      args.odesc.desc(), grad_output.data_ptr(),
      args.cdesc.desc(), bwdDataAlg, workspace.data, workspace.size,
      &zero, args.idesc.desc(), grad_input.data_ptr()));
}
void raw_cudnn_convolution_backward_input_out ( const at :: Tensor & grad_input , const at :: Tensor & grad_output , const at :: Tensor & weight , IntList padding , IntList stride , IntList dilation , int64_t groups , bool benchmark , bool deterministic ) { auto dataType = getCudnnDataType ( grad_output ) ; ConvolutionArgs args { grad_input , grad_output , weight } ; args . handle = getCudnnHandle ( ) ; setConvolutionParams ( & args . params , grad_input , weight , padding , stride , dilation , groups , deterministic ) ; args . idesc . set ( grad_input ) ; args . wdesc . set ( weight ) ; args . odesc . set ( grad_output ) ; args . cdesc . set ( dataType , grad_output . dim ( ) - 2 , args . params . padding , args . params . stride , args . params . dilation , args . params . groups ) ; cudnnConvolutionBwdDataAlgo_t bwdDataAlg ; Workspace workspace = chooseAlgorithm ( args , benchmark , & bwdDataAlg ) ; Constant one ( dataType , 1 ) ; Constant zero ( dataType , 0 ) ; AT_CUDNN_CHECK ( cudnnConvolutionBackwardData ( args . handle , & one , args . wdesc . desc ( ) , weight . data_ptr ( ) , args . odesc . desc ( ) , grad_output . data_ptr ( ) , args . cdesc . desc ( ) , bwdDataAlg , workspace . data , workspace . size , & zero , args . idesc . desc ( ) , grad_input . data_ptr ( ) ) ) ; }
Convolution backward / Transposed convolution forward
Convolution backward / Transposed convolution forward