Gnn/src/EdgeLayerConnector.cu

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/EdgeLayerConnector.hpp"
0010 #include "ActsPlugins/Gnn/detail/CudaUtils.hpp"
0011
0012 #include <MMG/CUDA_edge_layer_connector>
0013 #include <thrust/copy.h>
0014 #include <thrust/execution_policy.h>
0015
0016 using namespace Acts;
0017
0018 namespace ActsPlugins {
0019
0020 std::vector<std::vector<int>> EdgeLayerConnector::operator()(
0021     PipelineTensors tensors, std::vector<int>& spacepointIDs,
0022     const ExecutionContext& execContext) {
0023   ACTS_DEBUG("Get event data");
0024
0025   const auto numEdges = static_cast<int>(tensors.edgeIndex.shape().at(1));
0026   const auto numSpacepoints = static_cast<int>(spacepointIDs.size());
0027
0028   auto stream = execContext.stream.value();
0029
0030   // Convert std::int64_t edge indices to int using Tensor for memory management
0031   auto srcInt64Ptr = tensors.edgeIndex.data();
0032   auto tgtInt64Ptr = tensors.edgeIndex.data() + numEdges;
0033
0034   auto edgeSrc =
0035       Tensor<int>::Create({1, static_cast<std::size_t>(numEdges)}, execContext);
0036   auto edgeTgt =
0037       Tensor<int>::Create({1, static_cast<std::size_t>(numEdges)}, execContext);
0038
0039   thrust::copy(thrust::cuda::par.on(stream), srcInt64Ptr,
0040                srcInt64Ptr + numEdges, edgeSrc.data());
0041   thrust::copy(thrust::cuda::par.on(stream), tgtInt64Ptr,
0042                tgtInt64Ptr + numEdges, edgeTgt.data());
0043
0044   // Copy spacepoint IDs to GPU
0045   auto spacepointIDsTensor =
0046       Tensor<int>::Create({1, spacepointIDs.size()}, execContext);
0047   ACTS_CUDA_CHECK(cudaMemcpyAsync(
0048       spacepointIDsTensor.data(), spacepointIDs.data(),
0049       spacepointIDs.size() * sizeof(int), cudaMemcpyHostToDevice, stream));
0050   ACTS_CUDA_CHECK(cudaStreamSynchronize(stream));
0051
0052   ACTS_DEBUG("Setup graph...");
0053   CUDA_graph<float> graph(spacepointIDsTensor.data(), numSpacepoints,
0054                           edgeSrc.data(), edgeTgt.data(),
0055                           tensors.edgeScores->data(), numEdges);
0056   ACTS_CUDA_CHECK(cudaDeviceSynchronize());
0057   ACTS_CUDA_CHECK(cudaGetLastError());
0058
0059   ACTS_DEBUG("Setup EdgeLayerConnector...");
0060   CUDA_edge_layer_connector<float> connector(
0061       &graph, m_cfg.weightsCut, static_cast<int>(m_cfg.minHits),
0062       static_cast<int>(m_cfg.blockSize),
0063       static_cast<int>(m_cfg.maxHitsPerTrack));
0064   ACTS_CUDA_CHECK(cudaDeviceSynchronize());
0065   ACTS_CUDA_CHECK(cudaGetLastError());
0066
0067   ACTS_DEBUG("Build tracks...");
0068   connector.build_tracks();
0069   ACTS_CUDA_CHECK(cudaDeviceSynchronize());
0070   ACTS_CUDA_CHECK(cudaGetLastError());
0071
0072   const int maxHitsPerTrack = connector.cuda_tracks()->max_hits_per_track();
0073   const int tracksSize = connector.cuda_tracks()->size();
0074   const int nbTracks = connector.cuda_tracks()->nb_tracks();
0075
0076   ACTS_DEBUG("maxHitsPerTrack: " << maxHitsPerTrack << ", tracksSize: "
0077                                  << tracksSize << ", nbTracks: " << nbTracks);
0078
0079   std::vector<int> nbHits(nbTracks);
0080   ACTS_CUDA_CHECK(
0081       cudaMemcpyAsync(nbHits.data(), connector.cuda_tracks()->nb_hits(),
0082                       nbTracks * sizeof(int), cudaMemcpyDeviceToHost, stream));
0083
0084   std::vector<int> flatHits(tracksSize);
0085   ACTS_CUDA_CHECK(cudaMemcpyAsync(
0086       flatHits.data(), connector.cuda_tracks()->hits(),
0087       tracksSize * sizeof(int), cudaMemcpyDeviceToHost, stream));
0088
0089   ACTS_CUDA_CHECK(cudaStreamSynchronize(stream));
0090   ACTS_CUDA_CHECK(cudaGetLastError());
0091
0092   std::vector<std::vector<int>> trackCandidates;
0093   trackCandidates.reserve(nbTracks);
0094
0095   std::size_t minTrackSize = std::numeric_limits<std::size_t>::max();
0096   std::size_t maxTrackSize = 0;
0097   std::size_t avgTrackSize = 0;
0098
0099   for (int trackIdx = 0; trackIdx < nbTracks; ++trackIdx) {
0100     const auto* trackBegin = flatHits.data() + trackIdx * maxHitsPerTrack;
0101     const auto* trackEnd = trackBegin + nbHits[trackIdx];
0102     trackCandidates.emplace_back(trackBegin, trackEnd);
0103
0104     // Debug print the first 10 tracks
0105     if (trackIdx < 10) {
0106       ACTS_DEBUG("Track " << trackIdx << ": " << [&]() {
0107         std::ostringstream oss;
0108         for (const auto& hit : trackCandidates.back()) {
0109           oss << hit << " ";
0110         }
0111         return oss.str();
0112       }());
0113     }
0114
0115     const std::size_t trackSize = trackCandidates.back().size();
0116     minTrackSize = std::min(minTrackSize, trackSize);
0117     maxTrackSize = std::max(maxTrackSize, trackSize);
0118     avgTrackSize += trackSize;
0119   }
0120
0121   avgTrackSize /= nbTracks;
0122   ACTS_DEBUG("Min/Avg/Max track size: " << minTrackSize << "/" << avgTrackSize
0123                                         << "/" << maxTrackSize);
0124
0125   return trackCandidates;
0126 }
0127
0128 }  // namespace ActsPlugins