Method get_workspace_in_bytes

dnn/src/naive/rnn/opr_impl.cpp:55–72

}

size_t RNNImpl::get_workspace_in_bytes(
        const TensorLayout& input, const TensorLayout& hx,
        const TensorLayout& flatten_weights, const TensorLayout& output,
        const TensorLayout& /*hy*/, const TensorLayout& /*reserve_space*/) {
    auto _param = param();
    size_t D = _param.bidirectional ? 2 : 1;
    size_t last_dim = std::max(input.shape[2], D * hx.shape[1]);
    TensorLayout last_input = {{input.shape[0], input.shape[1], last_dim}, input.dtype};
    size_t workspace_size = rnn::get_workspace_in_bytes<RNNCellForward>(
            last_input, flatten_weights, param().hidden_size,
            param().bidirectional ? 2 : 1, this->handle());
    if (!param().bias) {  // use fake bias (all 0)
        TensorLayout bias_layout = {{param().hidden_size}, flatten_weights.dtype};
        workspace_size += bias_layout.span().dist_byte();
    }
    workspace_size += output.span().dist_byte();
    return workspace_size;
}

size_t RNNImpl::get_reserve_size_in_bytes(const TensorLayout& input) {
    size_t num_layers = param().num_layers;