aten/src/ATen/native/cudnn/RNN.cpp
// NB: This MUST BE CALLED AFTER _cudnn_rnn_backward_input.
// We'll give a user friendly combined function...
This MUST BE CALLED AFTER _cudnn_rnn_backward_input.
We'll give a user friendly combined function
This MUST BE CALLED AFTER _cudnn_rnn_backward_input . We ' ll give a user friendly combined function
std::vector<Tensor> _cudnn_rnn_backward_weight(
    const Tensor& input_r, TensorList weight_arr, int64_t weight_stride0,
    const Tensor& weight_buf, const Tensor& hx, const Tensor& cx,
    const Tensor& output_r,
    int64_t fn_mode, int64_t fn_hidden_size,
    int64_t fn_num_layers, bool batch_first, double fn_dropout,
    bool fn_train, bool fn_bidirectional, IntList fn_batch_sizes,
    const Tensor& fn_dropout_state, const Tensor& fn_reserve
    ) {
  MatrixRef<Tensor> weight{ weight_arr, static_cast<size_t>(weight_stride0) };
  auto input = input_r;
  auto output = output_r;
  RNNParams fn;
  fn.rnn.set(fn_mode, fn_hidden_size, fn_num_layers, fn_bidirectional, getCudnnDataType(input));
  fn.dropout.set(fn_train, fn_dropout, fn_dropout_state);
  fn.tensors.set(input.sizes(), fn_batch_sizes, batch_first);
  auto handle = getCudnnHandle();
  if (fn.rnn.mode != CUDNN_LSTM) {
    if (cx.defined()) {
      throw std::runtime_error("rnn: illegal defined cx for non-LSTM RNN");
    }
  }
  auto is_input_packed = fn_batch_sizes.size() != 0;
  if (batch_first && !is_input_packed) {
    input = input.transpose(0, 1);
    output = output.transpose(0, 1);
  }
  auto input_size = _input_size(fn.tensors);
  auto hidden_size = _hidden_size(fn.rnn, fn.tensors);
  if (!fn_train) {
    throw std::runtime_error("cudnn RNN backward can only be called in training mode");
  }
  if (!input.sizes().equals(input_size)) {
    std::ostringstream oss;
    oss << "Expected input size " << IntList{input_size} << ", got " << input.sizes();
    throw std::runtime_error(oss.str());
  }
  if (hx.defined() && !hx.sizes().equals(hidden_size)) {
    std::ostringstream oss;
    oss << "Expected hidden size " << IntList{hidden_size} << ", got " << hx.sizes();
    throw std::runtime_error(oss.str());
  }
  if (!hx.is_contiguous()) {
    throw std::runtime_error("rnn: hx is not contiguous");
  }
  if (cx.defined() && !cx.is_contiguous()) {
    throw std::runtime_error("rnn: cx is not contiguous");
  }
  auto x = input.contiguous();
  const auto& y = output;
  auto dw = weight_buf.type().tensor(weight_buf.sizes()).zero_();
  RNNDescriptors descs(fn, handle, x, y, hx, cx);
  FilterDescriptor w_desc;
  w_desc.set(weight_buf, 3);
  size_t workspace_size;
  auto x_descs_arr = descs.get_x_descs();
  auto y_descs_arr = descs.get_y_descs();
  AT_CUDNN_CHECK(cudnnGetRNNWorkspaceSize(
        handle,
        descs.rnn_desc.desc(),
        fn.tensors.seq_length,
        x_descs_arr.data(),
        &workspace_size
        ));
  Tensor workspace = input.type().toScalarType(kByte).tensor(workspace_size);
  AT_CUDNN_CHECK(cudnnRNNBackwardWeights(
        handle,
        descs.rnn_desc.desc(),
        fn.tensors.seq_length,
        x_descs_arr.data(), x.data_ptr(),
        descs.hx_desc.desc(), hx.data_ptr(),
        y_descs_arr.data(), y.data_ptr(),
        workspace.data_ptr(), workspace.size(0),
        w_desc.desc(), dw.data_ptr(),
        fn_reserve.data_ptr(), fn_reserve.size(0)
        ));
  std::vector<Tensor> grad_weight_arr;
  grad_weight_arr.reserve( weight.numel() );
  for (const auto& w : weight_arr) {
    grad_weight_arr.emplace_back(w.type().tensor(w.sizes()).zero_());
  }
  std::vector<Tensor> grad_params_arr;
  size_t grad_params_stride0;
  std::tie(grad_params_arr, grad_params_stride0) = get_parameters(handle, fn.rnn, descs.rnn_desc, descs.x_descs[0], w_desc, dw);
  _copyParams(MatrixRef<Tensor>{grad_params_arr, grad_params_stride0},
              MatrixRef<Tensor>{grad_weight_arr, static_cast<size_t>(weight_stride0)});
  return grad_weight_arr; 
}
std :: vector < Tensor > _cudnn_rnn_backward_weight ( const Tensor & input_r , TensorList weight_arr , int64_t weight_stride0 , const Tensor & weight_buf , const Tensor & hx , const Tensor & cx , const Tensor & output_r , int64_t fn_mode , int64_t fn_hidden_size , int64_t fn_num_layers , bool batch_first , double fn_dropout , bool fn_train , bool fn_bidirectional , IntList fn_batch_sizes , const Tensor & fn_dropout_state , const Tensor & fn_reserve ) { MatrixRef < Tensor > weight { weight_arr , static_cast < size_t > ( weight_stride0 ) } ; auto input = input_r ; auto output = output_r ; RNNParams fn ; fn . rnn . set ( fn_mode , fn_hidden_size , fn_num_layers , fn_bidirectional , getCudnnDataType ( input ) ) ; fn . dropout . set ( fn_train , fn_dropout , fn_dropout_state ) ; fn . tensors . set ( input . sizes ( ) , fn_batch_sizes , batch_first ) ; auto handle = getCudnnHandle ( ) ; if ( fn . rnn . mode != CUDNN_LSTM ) { if ( cx . defined ( ) ) { throw std :: runtime_error ( " " ) ; } } auto is_input_packed = fn_batch_sizes . size ( ) != 0 ; if ( batch_first && ! is_input_packed ) { input = input . transpose ( 0 , 1 ) ; output = output . transpose ( 0 , 1 ) ; } auto input_size = _input_size ( fn . tensors ) ; auto hidden_size = _hidden_size ( fn . rnn , fn . tensors ) ; if ( ! fn_train ) { throw std :: runtime_error ( " " ) ; } if ( ! input . sizes ( ) . equals ( input_size ) ) { std :: ostringstream oss ; oss << " " << IntList { input_size } << " " << input . sizes ( ) ; throw std :: runtime_error ( oss . str ( ) ) ; } if ( hx . defined ( ) && ! hx . sizes ( ) . equals ( hidden_size ) ) { std :: ostringstream oss ; oss << " " << IntList { hidden_size } << " " << hx . sizes ( ) ; throw std :: runtime_error ( oss . str ( ) ) ; } if ( ! hx . is_contiguous ( ) ) { throw std :: runtime_error ( " " ) ; } if ( cx . defined ( ) && ! cx . is_contiguous ( ) ) { throw std :: runtime_error ( " " ) ; } auto x = input . contiguous ( ) ; const auto & y = output ; auto dw = weight_buf . type ( ) . tensor ( weight_buf . sizes ( ) ) . zero_ ( ) ; RNNDescriptors descs ( fn , handle , x , y , hx , cx ) ; FilterDescriptor w_desc ; w_desc . set ( weight_buf , 3 ) ; size_t workspace_size ; auto x_descs_arr = descs . get_x_descs ( ) ; auto y_descs_arr = descs . get_y_descs ( ) ; AT_CUDNN_CHECK ( cudnnGetRNNWorkspaceSize ( handle , descs . rnn_desc . desc ( ) , fn . tensors . seq_length , x_descs_arr . data ( ) , & workspace_size ) ) ; Tensor workspace = input . type ( ) . toScalarType ( kByte ) . tensor ( workspace_size ) ; AT_CUDNN_CHECK ( cudnnRNNBackwardWeights ( handle , descs . rnn_desc . desc ( ) , fn . tensors . seq_length , x_descs_arr . data ( ) , x . data_ptr ( ) , descs . hx_desc . desc ( ) , hx . data_ptr ( ) , y_descs_arr . data ( ) , y . data_ptr ( ) , workspace . data_ptr ( ) , workspace . size ( 0 ) , w_desc . desc ( ) , dw . data_ptr ( ) , fn_reserve . data_ptr ( ) , fn_reserve . size ( 0 ) ) ) ; std :: vector < Tensor > grad_weight_arr ; grad_weight_arr . reserve ( weight . numel ( ) ) ; for ( const auto & w : weight_arr ) { grad_weight_arr . emplace_back ( w . type ( ) . tensor ( w . sizes ( ) ) . zero_ ( ) ) ; } std :: vector < Tensor > grad_params_arr ; size_t grad_params_stride0 ; std :: tie ( grad_params_arr , grad_params_stride0 ) = get_parameters ( handle , fn . rnn , descs . rnn_desc , descs . x_descs [ 0 ] , w_desc , dw ) ; _copyParams ( MatrixRef < Tensor > { grad_params_arr , grad_params_stride0 } , MatrixRef < Tensor > { grad_weight_arr , static_cast < size_t > ( weight_stride0 ) } ) ; return grad_weight_arr ; }
NB: This MUST BE CALLED AFTER _cudnn_rnn_backward_input.
NB : This MUST BE CALLED AFTER _cudnn_rnn_backward_input .
// TODO: I think tensor geometry sufficient for weight_buf/weight // TODO: the above were the only checks in rnn.py, but it doesn't seem // like these checks are enough // stride is known from call site (and also inconvenient to return)