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