Function BuildNearestNeighbor

picasso/point/pi_modules/source/nnquery.cpp:131–171

std::vector<torch::Tensor> BuildNearestNeighbor(
    torch::Tensor database,     // database points: concat_Np * 3
    torch::Tensor query,        // query points: concat_Mp * 3
    torch::Tensor nv_database,  // batch: each element is the vertex number of a database sample
    torch::Tensor nv_query)     // batch: each element is the vertex number of a query sample
{
    CHECK_INPUT(database,2);
    CHECK_INPUT(query,2);
    CHECK_INPUT(nv_database,1);
    CHECK_INPUT(nv_query,1);

    // get the dims required by computations
    int Np = database.size(0);    // number of database points
    int Mp = query.size(0);       // number of query points
    int B  = nv_database.size(0); // batch size

    TORCH_CHECK(database.dim()==2 && database.size(1)==3,
                "Shape of database points requires to be (Np, 3)");
    TORCH_CHECK(query.dim()==2 && query.size(1)==3,
                "Shape of query points requires to be (Mp, 3)");
    TORCH_CHECK(nv_database.dim()==1 && nv_query.dim()==1 && nv_database.size(0)==nv_query.size(0),
                "Shape of nv_database and nv_query should be identical");

    const int nn_out = 3; // find 3 nearest neighbors

    // get pointers the input tensors
    const float* database_ptr   = database.data_ptr<float>();
    const float* query_ptr      = query.data_ptr<float>();
    const int*   nvDatabase_ptr = nv_database.data_ptr<int32_t>();
    const int*   nvQuery_ptr    = nv_query.data_ptr<int32_t>();

    // create output tensors
    auto nn_index = torch::zeros({Mp,nn_out}, database.options().dtype(torch::kInt32));
    auto nn_dist  = torch::zeros({Mp,nn_out}, database.options());
    int* nnIndex_ptr  = nn_index.data_ptr<int32_t>();
    float* nnDist_ptr = nn_dist.data_ptr<float>();

    buildNearestNeighborLauncher(B, Np, Mp, database_ptr, query_ptr,
                                 nvDatabase_ptr, nvQuery_ptr, nnIndex_ptr, nnDist_ptr);
    return {nn_index, nn_dist};
}


PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)