Gnn/src/buildEdgesFRNN.cpp

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008
0009 #include "ActsPlugins/Gnn/detail/buildEdges.hpp"
0010
0011 #include <cstdint>
0012 #include <stdexcept>
0013 #include <utility>
0014
0015 #include <cuda.h>
0016 #include <cuda_runtime_api.h>
0017 #include <grid/counting_sort.h>
0018 #include <grid/find_nbrs.h>
0019 #include <grid/grid.h>
0020 #include <grid/insert_points.h>
0021 #include <grid/prefix_sum.h>
0022 #include <torch/torch.h>
0023
0024 using namespace torch::indexing;
0025
0026 torch::Tensor ActsPlugins::detail::buildEdgesFRNN(torch::Tensor &embedFeatures,
0027                                                   float rVal, int kVal,
0028                                                   bool flipDirections) {
0029   const auto device = embedFeatures.device();
0030
0031   const std::int64_t numSpacePoints = embedFeatures.size(0);
0032   const int dim = embedFeatures.size(1);
0033
0034   const int grid_params_size = 8;
0035   const int grid_delta_idx = 3;
0036   const int grid_total_idx = 7;
0037   const int grid_max_res = 128;
0038   const int grid_dim = 3;
0039
0040   if (dim < 3) {
0041     throw std::runtime_error("DIM < 3 is not supported for now.\n");
0042   }
0043
0044   const float radius_cell_ratio = 2.0;
0045   const int batch_size = 1;
0046   int G = -1;
0047
0048   // Set up grid properties
0049   torch::Tensor grid_min;
0050   torch::Tensor grid_max;
0051   torch::Tensor grid_size;
0052
0053   torch::Tensor embedTensor = embedFeatures.reshape({1, numSpacePoints, dim});
0054   torch::Tensor gridParamsCuda =
0055       torch::zeros({batch_size, grid_params_size}, device).to(torch::kFloat32);
0056   torch::Tensor r_tensor = torch::full({batch_size}, rVal, device);
0057   torch::Tensor lengths = torch::full({batch_size}, numSpacePoints, device);
0058
0059   // build the grid
0060   for (int i = 0; i < batch_size; i++) {
0061     torch::Tensor allPoints =
0062         embedTensor.index({i, Slice(None, lengths.index({i}).item().to<long>()),
0063                            Slice(None, grid_dim)});
0064     grid_min = std::get<0>(allPoints.min(0));
0065     grid_max = std::get<0>(allPoints.max(0));
0066     gridParamsCuda.index_put_({i, Slice(None, grid_delta_idx)}, grid_min);
0067
0068     grid_size = grid_max - grid_min;
0069
0070     float cell_size =
0071         r_tensor.index({i}).item().to<float>() / radius_cell_ratio;
0072
0073     if (cell_size < (grid_size.min().item().to<float>() / grid_max_res)) {
0074       cell_size = grid_size.min().item().to<float>() / grid_max_res;
0075     }
0076
0077     gridParamsCuda.index_put_({i, grid_delta_idx}, 1 / cell_size);
0078
0079     gridParamsCuda.index_put_({i, Slice(1 + grid_delta_idx, grid_total_idx)},
0080                               floor(grid_size / cell_size) + 1);
0081
0082     gridParamsCuda.index_put_(
0083         {i, grid_total_idx},
0084         gridParamsCuda.index({i, Slice(1 + grid_delta_idx, grid_total_idx)})
0085             .prod());
0086
0087     if (G < gridParamsCuda.index({i, grid_total_idx}).item().to<int>()) {
0088       G = gridParamsCuda.index({i, grid_total_idx}).item().to<int>();
0089     }
0090   }
0091
0092   torch::Tensor pc_grid_cnt =
0093       torch::zeros({batch_size, G}, device).to(torch::kInt32);
0094   torch::Tensor pc_grid_cell =
0095       torch::full({batch_size, numSpacePoints}, -1, device).to(torch::kInt32);
0096   torch::Tensor pc_grid_idx =
0097       torch::full({batch_size, numSpacePoints}, -1, device).to(torch::kInt32);
0098
0099   // put space points into the grid
0100   InsertPointsCUDA(embedTensor, lengths.to(torch::kInt64), gridParamsCuda,
0101                    pc_grid_cnt, pc_grid_cell, pc_grid_idx, G);
0102
0103   torch::Tensor pc_grid_off =
0104       torch::full({batch_size, G}, 0, device).to(torch::kInt32);
0105   torch::Tensor grid_params = gridParamsCuda.to(torch::kCPU);
0106
0107   // for loop seems not to be necessary anymore
0108   pc_grid_off = PrefixSumCUDA(pc_grid_cnt, grid_params);
0109
0110   torch::Tensor sorted_points =
0111       torch::zeros({batch_size, numSpacePoints, dim}, device)
0112           .to(torch::kFloat32);
0113   torch::Tensor sorted_points_idxs =
0114       torch::full({batch_size, numSpacePoints}, -1, device).to(torch::kInt32);
0115
0116   CountingSortCUDA(embedTensor, lengths.to(torch::kInt64), pc_grid_cell,
0117                    pc_grid_idx, pc_grid_off, sorted_points, sorted_points_idxs);
0118
0119   auto [indices, distances] = FindNbrsCUDA(
0120       sorted_points, sorted_points, lengths.to(torch::kInt64),
0121       lengths.to(torch::kInt64), pc_grid_off.to(torch::kInt32),
0122       sorted_points_idxs, sorted_points_idxs,
0123       gridParamsCuda.to(torch::kFloat32), kVal, r_tensor, r_tensor * r_tensor);
0124   torch::Tensor positiveIndices = indices >= 0;
0125
0126   torch::Tensor repeatRange = torch::arange(positiveIndices.size(1), device)
0127                                   .repeat({1, positiveIndices.size(2), 1})
0128                                   .transpose(1, 2);
0129
0130   torch::Tensor stackedEdges = torch::stack(
0131       {repeatRange.index({positiveIndices}), indices.index({positiveIndices})});
0132
0133   return postprocessEdgeTensor(std::move(stackedEdges), true, true,
0134                                flipDirections);
0135 }