Method get_workspace_bundle

dnn/src/cuda/svd/opr_impl.cpp:43–57

namespace cuda {

WorkspaceBundle SVDForwardImpl::get_workspace_bundle(
        size_t block_cnt, size_t m, size_t n, size_t dtype_size, void* raw_ptr) {
    const size_t max_mn = std::max(m, n);
    const size_t min_mn = std::min(m, n);
    SmallVector<size_t> ws_sizes = {
            block_cnt * m * n * dtype_size,  // copy of src
            get_cusolver_buffer_size(max_mn, min_mn) * dtype_size,
            sizeof(int)  // devInfo
    };
    if (m > n) {
        ws_sizes.push_back(block_cnt * max_mn * max_mn * dtype_size);
        ws_sizes.push_back(block_cnt * min_mn * min_mn * dtype_size);
    }
    return {raw_ptr, std::move(ws_sizes), handle()->alignment_requirement()};
}

size_t SVDForwardImpl::get_cusolver_buffer_size(size_t m, size_t n) {
    int lwork;