Function gather_points_v2

pointersect/pr/cuda/pr_cuda.cpp:257–271

Gather the points belonging to each grid cell. Args: grid_idxs: (b, n), grid linear index of each point total_cells: (b,) total grid cells valid_mask: (b, n), whether to record the point cell_counts: (b, max_num_cells), number of points in a cell Returns: membank: (b, n), should use the pidx of cell_idx is from cell_start_idx[b][cell_idx] to cell_counts[b][cell_idx+1] (excluded). cell_start_idx:

//         (b, n_cells+1), 
// 
std::tuple<torch::Tensor, torch::Tensor> gather_points_v2(
        const torch::Tensor & grid_idxs,  // (b, n), long
        const torch::Tensor & total_cells, // (b,) int64_t
        const torch::Tensor & valid_mask,  // (b, n), bool
        const torch::Tensor & cell_counts  // (b, n_cells), int32
) { 

    CHECK_INPUT(grid_idxs);
    CHECK_INPUT(total_cells);
    CHECK_INPUT(valid_mask);
    CHECK_INPUT(cell_counts);


    return gather_points_cuda_v2(grid_idxs, total_cells, valid_mask, cell_counts);
}