ExaTrkX/src/GraphStoreHook.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 "Acts/Plugins/ExaTrkX/GraphStoreHook.hpp"
0010
0011 Acts::GraphStoreHook::GraphStoreHook() {
0012   m_storedGraph = std::make_unique<Graph>();
0013 }
0014
0015 void Acts::GraphStoreHook::operator()(const PipelineTensors &tensors,
0016                                       const ExecutionContext &execCtx) const {
0017   auto edgeIndexTensor =
0018       tensors.edgeIndex.clone({Device::Cpu(), execCtx.stream});
0019
0020   // We need to transpose the edges here for the right memory layout
0021   m_storedGraph->first.reserve(edgeIndexTensor.size());
0022   for (auto i = 0ul; i < edgeIndexTensor.shape().at(1); ++i) {
0023     m_storedGraph->first.push_back(*(edgeIndexTensor.data() + i));
0024     m_storedGraph->first.push_back(
0025         *(edgeIndexTensor.data() + edgeIndexTensor.shape().at(1) + i));
0026   }
0027
0028   if (!tensors.edgeScores.has_value()) {
0029     return;
0030   }
0031
0032   auto scoreTensor = tensors.edgeScores->clone({Device::Cpu(), execCtx.stream});
0033
0034   m_storedGraph->second =
0035       std::vector(scoreTensor.data(), scoreTensor.data() + scoreTensor.size());
0036 }