Function get_workspace_bundle

dnn/src/arm_common/lstm/lstm_utils.h:205–243

template <typename CellOpr>
WorkspaceBundle get_workspace_bundle(
        const TensorLayout& input, const TensorLayout& output,
        const TensorLayout& flatten_weights, size_t hidden_size, size_t dir_size,
        size_t states_size) {
    size_t batch_size = input.shape[1];
    size_t input_size = input.shape[2];
    size_t gate_hidden_size = flatten_weights.shape[0];

    // cell workspace
    TensorLayout weight_ih{{gate_hidden_size, input_size}, flatten_weights.dtype};
    TensorLayout weight_hh{
            {gate_hidden_size, dir_size * hidden_size}, flatten_weights.dtype};
    TensorLayout bias{{1, gate_hidden_size}, flatten_weights.dtype};
    TensorLayout hx{{batch_size, dir_size * hidden_size}, input.dtype};

    auto cell_opr = inplace_cpu_handle()->create_operator<CellOpr>();

    TensorLayout h_new, c_new, gates;
    cell_opr->deduce_layout(
            input, weight_ih, bias, hx, weight_hh, bias, hx, h_new, c_new, gates);

    SmallVector<size_t> workspaces;
    //! the cell opr compute workspace
    size_t cell_opr_workspace = cell_opr->get_workspace_in_bytes(
            input, weight_ih, bias, hx, weight_hh, bias, hx, h_new, c_new, gates);
    workspaces.push_back(gates.span().dist_byte());
    workspaces.push_back(cell_opr_workspace);
    //! double tmp output memory
    size_t tmp_output_workspace = output.span().dist_byte();
    workspaces.push_back(tmp_output_workspace);
    workspaces.push_back(tmp_output_workspace);

    //! tmp states memory
    size_t tmp_state_workspace = hx.span().dist_byte();
    for (size_t i = 0; i < states_size; i++) {
        workspaces.push_back(tmp_state_workspace);
    }
    return {nullptr, workspaces};
}

}  // namespace arm_common
}  // namespace megdnn