Method init_param_mapping

src/model/operation/rnn.cc:437–496

}

void CudnnRNNHandle::init_param_mapping(cudnnTensorDescriptor_t *xDesc) {
  int linLayerIDRange = 2;
  if (mode == 0 || mode == 1) {
    // vanilla relu/tanh
    linLayerIDRange = 2;
  } else if (mode == 2) {
    // lstm
    linLayerIDRange = 8;
  } else if (mode == 3) {
    // gru
    linLayerIDRange = 6;
  }
  int pseudoLayerRange = (bidirectional ? 2 : 1) * num_layers;

  // dummy weights for getting the offset
  Tensor weights(
      Shape{
          weights_size,
      },
      dev);
  weights.SetValue(0.0f);
  const void *W_ptr = weights.block()->data();

  void *param_ptr = nullptr;
  int dims[] = {1, 1, 1};
  cudnnDataType_t data_type;
  cudnnTensorFormat_t tensor_format;
  int n_dims;
  cudnnFilterDescriptor_t paramDesc;
  CUDNN_CHECK(cudnnCreateFilterDescriptor(&paramDesc));

  vector<bool> paramTypes{false, true};
  for (int linLayerID = 0; linLayerID < linLayerIDRange; linLayerID++) {
    for (int pseudoLayer = 0; pseudoLayer < pseudoLayerRange; pseudoLayer++) {
      for (const bool &is_bias : paramTypes) {
        // get param ptr
        if (is_bias) {
          CUDNN_CHECK(cudnnGetRNNLinLayerBiasParams(
              ctx->cudnn_handle, rnnDesc, pseudoLayer, xDesc[0], wDesc, W_ptr,
              linLayerID, paramDesc, &param_ptr));
        } else {
          CUDNN_CHECK(cudnnGetRNNLinLayerMatrixParams(
              ctx->cudnn_handle, rnnDesc, pseudoLayer, xDesc[0], wDesc, W_ptr,
              linLayerID, paramDesc, &param_ptr));
        }

        // get param dims
        CUDNN_CHECK(cudnnGetFilterNdDescriptor(paramDesc, 3, &data_type,
                                               &tensor_format, &n_dims, dims));

        // get diff - offset
        size_t offset = (float *)param_ptr - (float *)W_ptr;

        // save in map
        weights_mapping[std::make_tuple(linLayerID, pseudoLayer, is_bias)] =
            std::make_tuple(offset, dims[0] * dims[1] * dims[2]);
      }
    }