Method backward_workspace_size_in_bytes

dnn/src/naive/rnn/rnn.cpp:104–139

}

size_t CellWeightsWrapperBase::backward_workspace_size_in_bytes(
        Handle* handle, size_t batch_size, size_t hidden_size, size_t input_size,
        size_t num_chunks, DType dtype) {
    size_t gate_hidden_size = hidden_size * num_chunks;
    TensorLayout tmp = {{batch_size, gate_hidden_size}, dtype};
    TensorLayout bias_expanded = {{1, gate_hidden_size}, dtype};
    TensorLayout wih = {{gate_hidden_size, input_size}, dtype};
    TensorLayout whh = {{gate_hidden_size, hidden_size}, dtype};
    TensorLayout x = {{batch_size, input_size}, dtype};
    TensorLayout hx = {{batch_size, hidden_size}, dtype};
    size_t workspace_size = 0;
    auto matrixmul_opr = handle->create_operator<MatrixMulForward>();
    matrixmul_opr->param().transposeA = false;
    matrixmul_opr->param().transposeB = false;
    // dx
    workspace_size = std::max(
            workspace_size, matrixmul_opr->get_workspace_in_bytes(tmp, wih, x));
    // dhx
    workspace_size = std::max(
            workspace_size, matrixmul_opr->get_workspace_in_bytes(tmp, whh, hx));
    // dwi
    matrixmul_opr->param().transposeA = true;
    workspace_size = std::max(
            workspace_size, matrixmul_opr->get_workspace_in_bytes(tmp, x, wih));
    // dwh
    workspace_size = std::max(
            workspace_size, matrixmul_opr->get_workspace_in_bytes(tmp, hx, whh));
    // dbias
    auto sum_opr = handle->create_operator<ReduceForward>();
    sum_opr->param().mode = ReduceForward::Mode::SUM;
    sum_opr->param().axis = 0;
    workspace_size = std::max(
            workspace_size, sum_opr->get_workspace_in_bytes(tmp, bias_expanded));
    workspace_size += tmp.span().dist_byte();
    return workspace_size;
}

RNNCellWeightWrapper::RNNCellWeightWrapper(
        void* weight_ptr, size_t hidden_size, size_t input_size, bool has_bias,