scatter_points_cpu.cpp

#include <torch/extension.h>
#include <ATen/TensorUtils.h>
// #include "voxelization.h"


namespace {


template <typename T_int>
void determin_max_points_kernel(torch::TensorAccessor<T_int,2> coor,
                                torch::TensorAccessor<T_int,1> point_to_voxelidx,
                                torch::TensorAccessor<T_int,1> num_points_per_voxel,
                                torch::TensorAccessor<T_int,3> coor_to_voxelidx,
                                int& voxel_num,
                                int& max_points,
                                const int num_points
                                ) {

    int voxelidx, num;
    for (int i = 0; i < num_points; ++i) {
        if (coor[i][0] == -1)
            continue;
        voxelidx = coor_to_voxelidx[coor[i][0]][coor[i][1]][coor[i][2]];

        // record voxel
        if (voxelidx == -1) {
            voxelidx = voxel_num;
            voxel_num += 1;
            coor_to_voxelidx[coor[i][0]][coor[i][1]][coor[i][2]] = voxelidx;
        }

        // put points into voxel
        num = num_points_per_voxel[voxelidx];
        point_to_voxelidx[i] = num;
        num_points_per_voxel[voxelidx] += 1;

        // update max points per voxel
        max_points = std::max(max_points, num+1);
    }

    return;
}


template <typename T, typename T_int>
void scatter_point_to_voxel_kernel(
                const torch::TensorAccessor<T,2> points,
                torch::TensorAccessor<T_int,2> coor,
                torch::TensorAccessor<T_int,1> point_to_voxelidx,
                torch::TensorAccessor<T_int,3> coor_to_voxelidx,
                torch::TensorAccessor<T,3> voxels,
                torch::TensorAccessor<T_int,2> voxel_coors,
                const int num_features,
                const int num_points,
                const int NDim
                ){
    for (int i = 0; i < num_points; ++i) {
        int num = point_to_voxelidx[i];
        int voxelidx = coor_to_voxelidx[coor[i][0]][coor[i][1]][coor[i][2]];
        for (int k = 0; k < num_features; ++k) {
            voxels[voxelidx][num][k] = points[i][k];
        }
        for (int k = 0; k < NDim; ++k) {
            voxel_coors[voxelidx][k] = coor[i][k];
        }
    }
}

} // namespace


namespace voxelization {

std::vector<at::Tensor> dynamic_point_to_voxel_cpu(
    const at::Tensor& points,
    const at::Tensor& voxel_mapping,
    const std::vector<float> voxel_size,
    const std::vector<float> coors_range) {
    // current version tooks about 0.02s_0.03s for one frame on cpu
    // check device
    AT_ASSERTM(points.device().is_cpu(), "points must be a CPU tensor");

    const int NDim = voxel_mapping.size(1);
    const int num_points = points.size(0);
    const int num_features = points.size(1);

    std::vector<int> grid_size(NDim);
    for (int i = 0; i < NDim; ++i) {
        grid_size[i] = round((coors_range[NDim + i] - coors_range[i]) / voxel_size[i]);
    }

    at::Tensor num_points_per_voxel = at::zeros({num_points,}, voxel_mapping.options());
    at::Tensor coor_to_voxelidx = -at::ones({grid_size[2], grid_size[1], grid_size[0]}, voxel_mapping.options());
    at::Tensor point_to_voxelidx = -at::ones({num_points,}, voxel_mapping.options());

    int voxel_num = 0;
    int max_points = 0;
    AT_DISPATCH_ALL_TYPES(voxel_mapping.type(), "determin_max_point", [&] {
        determin_max_points_kernel<scalar_t>(
            voxel_mapping.accessor<scalar_t,2>(),
            point_to_voxelidx.accessor<scalar_t,1>(),
            num_points_per_voxel.accessor<scalar_t,1>(),
            coor_to_voxelidx.accessor<scalar_t,3>(),
            voxel_num,
            max_points,
            num_points
        );
    });

    at::Tensor voxels = at::zeros({voxel_num, max_points, num_features}, points.options());
    at::Tensor voxel_coors = at::zeros({voxel_num, NDim}, points.options().dtype(at::kInt));

    AT_DISPATCH_ALL_TYPES(points.type(), "scatter_point_to_voxel", [&] {
        scatter_point_to_voxel_kernel<scalar_t, int>(
            points.accessor<scalar_t,2>(),
            voxel_mapping.accessor<int,2>(),
            point_to_voxelidx.accessor<int,1>(),
            coor_to_voxelidx.accessor<int,3>(),
            voxels.accessor<scalar_t,3>(),
            voxel_coors.accessor<int,2>(),
            num_features,
            num_points,
            NDim
        );
    });

    at::Tensor num_points_per_voxel_out = num_points_per_voxel.slice(/*dim=*/0, /*start=*/0, /*end=*/voxel_num);
    return {voxels, voxel_coors, num_points_per_voxel_out};
}

}