Method make_bundle

dnn/src/cuda/non_zero/opr_impl.cpp:16–27

using namespace megdnn::cuda::non_zero;

WorkspaceBundle NonZeroImpl::make_bundle(const TensorLayout& data) {
    size_t nr_item = data.total_nr_elems();
    cuda_check(cudaSetDevice(concrete_handle(handle())->device_id()));
    auto gen_idx_wk_size = cuda::cond_take::gen_idx_get_workspace_size(nr_item);
    SmallVector<size_t> sizes_in_bytes;
    sizes_in_bytes.push_back((nr_item + 1) * sizeof(megdnn::cuda::cond_take::IdxType));
    sizes_in_bytes.push_back(gen_idx_wk_size);
    // the two ele is the shape of arr and the reverse multiply arr of the shape
    sizes_in_bytes.push_back(sizeof(TensorLayout::shape));
    sizes_in_bytes.push_back(sizeof(TensorLayout::shape));
    return {nullptr, sizes_in_bytes, handle()->alignment_requirement()};
}

size_t NonZeroImpl::get_workspace_in_bytes(const TensorLayout& data) {
    return make_bundle(data).total_size_in_bytes();