Plugins/Gnn/TensorTests.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 <boost/test/unit_test.hpp>
0010
0011 #include "ActsPlugins/Gnn/Tensor.hpp"
0012
0013 using namespace Acts;
0014 using namespace ActsPlugins;
0015
0016 #ifdef ACTS_GNN_WITH_CUDA
0017 #include <cuda_runtime_api.h>
0018 #endif
0019
0020 const ExecutionContext execContextCpu{Device::Cpu(), {}};
0021
0022 template <typename T>
0023 Tensor<T> createCpuTensor(const std::vector<T>& data,
0024                           std::array<std::size_t, 2> shape) {
0025   auto tensor = Tensor<T>::Create(shape, execContextCpu);
0026   std::copy(data.begin(), data.end(), tensor.data());
0027   return tensor;
0028 }
0029
0030 void testSigmoid(std::vector<float> input, ExecutionContext execContext) {
0031   auto tensor = createCpuTensor(input, {input.size(), 1ul});
0032
0033   auto tensorTarget = tensor.clone(execContext);
0034   sigmoid(tensorTarget, execContext.stream);
0035   auto result = tensorTarget.clone({Device::Cpu(), execContext.stream});
0036
0037   std::vector<float> expected(input.size());
0038   std::transform(input.begin(), input.end(), expected.begin(),
0039                  [](float x) { return 1.f / (1.f + std::exp(-x)); });
0040
0041   BOOST_CHECK(result.size() == expected.size());
0042   for (std::size_t i = 0; i < result.size(); ++i) {
0043     BOOST_CHECK_CLOSE(result.data()[i], expected[i], 1e-4);
0044   }
0045 }
0046
0047 void testEdgeSelection(const std::vector<float>& scores,
0048                        const std::vector<std::int64_t>& edgeIndex,
0049                        const std::vector<std::int64_t>& edgeIndexExpected,
0050                        ExecutionContext execContext) {
0051   auto scoreTensor = createCpuTensor<float>(scores, {scores.size(), 1ul});
0052   auto edgeTensor = createCpuTensor(edgeIndex, {2, edgeIndex.size() / 2});
0053
0054   auto scoreTensorTarget = scoreTensor.clone(execContext);
0055   auto edgeTensorTarget = edgeTensor.clone(execContext);
0056
0057   auto [selectedScores, selectedEdges] = applyScoreCut(
0058       scoreTensorTarget, edgeTensorTarget, 0.5f, execContext.stream);
0059
0060   auto selectedScoresHost =
0061       selectedScores.clone({Device::Cpu(), execContext.stream});
0062   auto selectedEdgesHost =
0063       selectedEdges.clone({Device::Cpu(), execContext.stream});
0064
0065   BOOST_CHECK(selectedScoresHost.size() == 2);
0066
0067   BOOST_CHECK(selectedEdgesHost.size() == edgeIndexExpected.size());
0068   BOOST_CHECK_EQUAL_COLLECTIONS(
0069       selectedEdgesHost.data(),
0070       selectedEdgesHost.data() + selectedEdgesHost.size(),
0071       edgeIndexExpected.begin(), edgeIndexExpected.end());
0072 }
0073
0074 void testConstructionAndMove(ExecutionContext execContext) {
0075   auto tensor = Tensor<float>::Create({10, 1}, execContext);
0076
0077   BOOST_CHECK(tensor.shape()[1] == 1);
0078   BOOST_CHECK(tensor.shape()[0] == 10);
0079
0080   auto tensor2 = std::move(tensor);
0081   BOOST_CHECK(tensor2.shape()[1] == 1);
0082   BOOST_CHECK(tensor2.shape()[0] == 10);
0083   BOOST_CHECK(tensor2.data() != nullptr);
0084   BOOST_CHECK(tensor.data() == nullptr);
0085 }
0086
0087 void testEdgeLimit(ExecutionContext execContext) {
0088   // Create a sample edge tensor with 10 edges on CPU
0089   auto edgeTensor = Tensor<std::int64_t>::Create({2, 10}, execContextCpu);
0090   for (std::size_t i = 0; i < 10; ++i) {
0091     edgeTensor.data()[i] = i;
0092     edgeTensor.data()[i + 10] = 2 * i;
0093   }
0094
0095   // Create a sample edge feature tensor with 3 };
0096   auto edgeFeatureTensor = Tensor<float>::Create({10, 3}, execContextCpu);
0097   for (std::size_t i = 0; i < 10; ++i) {
0098     edgeFeatureTensor.data()[i * 3] = static_cast<float>(i);  // Feature 1
0099     edgeFeatureTensor.data()[i * 3 + 1] =
0100         static_cast<float>(i + 1);  // Feature 2
0101     edgeFeatureTensor.data()[i * 3 + 2] =
0102         static_cast<float>(i + 2);  // Feature 3
0103   }
0104
0105   // Clone to execContext
0106   auto edgeTensorTarget = edgeTensor.clone(execContext);
0107   std::optional<Tensor<float>> edgeFeatureTensorTarget =
0108       edgeFeatureTensor.clone(execContext);
0109
0110   // Apply edge limit
0111   auto [limitedEdges, limitedEdgeFeatures] = applyEdgeLimit(
0112       edgeTensorTarget, edgeFeatureTensorTarget, 5, execContext.stream);
0113
0114   // Clone results back to CPU
0115   auto limitedEdgesHost =
0116       limitedEdges.clone({Device::Cpu(), execContext.stream});
0117   auto limitedEdgeFeaturesHost =
0118       limitedEdgeFeatures->clone({Device::Cpu(), execContext.stream});
0119
0120   // Check the size of the limited edges
0121   BOOST_CHECK(limitedEdgesHost.shape()[1] == 5);
0122   BOOST_CHECK(limitedEdgesHost.shape()[0] == 2);
0123   BOOST_CHECK(limitedEdgeFeaturesHost.shape()[0] == 5);
0124   BOOST_CHECK(limitedEdgeFeaturesHost.shape()[1] == 3);
0125
0126   // Check the data of the limited edges
0127   for (std::size_t i = 0; i < 5; ++i) {
0128     BOOST_CHECK(limitedEdgesHost.data()[i] == edgeTensor.data()[i]);
0129     BOOST_CHECK(limitedEdgesHost.data()[i + 5] == edgeTensor.data()[i + 10]);
0130
0131     BOOST_CHECK(limitedEdgeFeaturesHost.data()[i * 3] ==
0132                 edgeFeatureTensor.data()[i * 3]);
0133     BOOST_CHECK(limitedEdgeFeaturesHost.data()[i * 3 + 1] ==
0134                 edgeFeatureTensor.data()[i * 3 + 1]);
0135     BOOST_CHECK(limitedEdgeFeaturesHost.data()[i * 3 + 2] ==
0136                 edgeFeatureTensor.data()[i * 3 + 2]);
0137   }
0138 }
0139
0140 namespace ActsTests {
0141
0142 BOOST_AUTO_TEST_SUITE(GnnSuite)
0143
0144 BOOST_AUTO_TEST_CASE(tensor_create_move_cpu) {
0145   testConstructionAndMove(execContextCpu);
0146 }
0147
0148 BOOST_AUTO_TEST_CASE(test_clone_cpu) {
0149   std::vector<float> data = {1.f, 2.f, 3.f, 4.f, 5.f, 6.f};
0150   auto tensor = createCpuTensor(data, {3, 2});
0151   auto tensorClone = tensor.clone(execContextCpu);
0152
0153   BOOST_CHECK(tensorClone.shape()[0] == 3);
0154   BOOST_CHECK(tensorClone.shape()[1] == 2);
0155   BOOST_CHECK(tensorClone.data() != nullptr);
0156   BOOST_CHECK(tensorClone.data() != tensor.data());
0157   BOOST_CHECK(tensorClone.size() == tensor.size());
0158   BOOST_CHECK(tensorClone.nbytes() == tensor.nbytes());
0159
0160   BOOST_CHECK_EQUAL_COLLECTIONS(tensorClone.data(),
0161                                 tensorClone.data() + tensorClone.size(),
0162                                 data.begin(), data.end());
0163 }
0164
0165 BOOST_AUTO_TEST_CASE(tensor_sigmoid_cpu) {
0166   testSigmoid({-2.f, -1.f, 0.f, 1.f, 2.f}, execContextCpu);
0167 }
0168
0169 const std::vector<float> scores = {0.1f, 0.4f, 0.6f, 0.9f};
0170 const std::vector<std::int64_t> edgeIndex = {0, 1, 2, 3, 4, 5, 6, 7};
0171 const std::vector<std::int64_t> edgeIndexExpected = {2, 3, 6, 7};
0172
0173 BOOST_AUTO_TEST_CASE(tensor_edge_selection_cpu) {
0174   testEdgeSelection(scores, edgeIndex, edgeIndexExpected, execContextCpu);
0175 }
0176
0177 BOOST_AUTO_TEST_CASE(tensor_edge_limit_cpu) {
0178   testEdgeLimit(execContextCpu);
0179 }
0180
0181 #ifdef ACTS_GNN_WITH_CUDA
0182
0183 const ExecutionContext execContextCuda{Device::Cuda(0), cudaStreamLegacy};
0184
0185 BOOST_AUTO_TEST_CASE(tensor_create_move_cuda) {
0186   testConstructionAndMove(execContextCuda);
0187 }
0188
0189 BOOST_AUTO_TEST_CASE(tensor_clone_roundtrip) {
0190   std::vector<float> data = {1.f, 2.f, 3.f, 4.f, 5.f, 6.f};
0191   auto tensorOrigHost = createCpuTensor(data, {3, 2});
0192   BOOST_CHECK(tensorOrigHost.device().isCpu());
0193
0194   auto tensorClone = tensorOrigHost.clone(execContextCuda);
0195   BOOST_CHECK(tensorClone.device().isCuda());
0196
0197   auto tensorCloneCuda = tensorClone.clone(execContextCuda);
0198   BOOST_CHECK(tensorCloneCuda.device().isCuda());
0199
0200   auto tensorCloneHost =
0201       tensorCloneCuda.clone({Device::Cpu(), cudaStreamLegacy});
0202   BOOST_CHECK(tensorCloneHost.device().isCpu());
0203
0204   BOOST_CHECK(tensorCloneHost.shape()[0] == 3);
0205   BOOST_CHECK(tensorCloneHost.shape()[1] == 2);
0206   BOOST_CHECK(tensorCloneHost.data() != nullptr);
0207   BOOST_CHECK(tensorCloneHost.data() != tensorCloneCuda.data());
0208   BOOST_CHECK(tensorCloneHost.size() == tensorCloneCuda.size());
0209   BOOST_CHECK(tensorCloneHost.nbytes() == tensorCloneCuda.nbytes());
0210   BOOST_CHECK_EQUAL_COLLECTIONS(tensorCloneHost.data(),
0211                                 tensorCloneHost.data() + tensorCloneHost.size(),
0212                                 data.begin(), data.end());
0213 }
0214
0215 BOOST_AUTO_TEST_CASE(tensor_sigmoid_cuda) {
0216   testSigmoid({-2.f, -1.f, 0.f, 1.f, 2.f}, execContextCuda);
0217 }
0218
0219 BOOST_AUTO_TEST_CASE(tensor_edge_selection_cuda) {
0220   testEdgeSelection(scores, edgeIndex, edgeIndexExpected, execContextCuda);
0221 }
0222
0223 BOOST_AUTO_TEST_CASE(tensor_edge_limit_cuda) {
0224   testEdgeLimit(execContextCuda);
0225 }
0226
0227 #endif
0228
0229 BOOST_AUTO_TEST_SUITE_END()
0230
0231 }  // namespace ActsTests