Plugins/ExaTrkX/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/ExaTrkX/Tensor.hpp>
0012
0013 #ifdef ACTS_EXATRKX_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(execContextCpu);
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 = selectedScores.clone(execContextCpu);
0058   auto selectedEdgesHost = selectedEdges.clone(execContextCpu);
0059
0060   BOOST_CHECK(selectedScoresHost.size() == 2);
0061
0062   BOOST_CHECK(selectedEdgesHost.size() == edgeIndexExpected.size());
0063   BOOST_CHECK_EQUAL_COLLECTIONS(
0064       selectedEdgesHost.data(),
0065       selectedEdgesHost.data() + selectedEdgesHost.size(),
0066       edgeIndexExpected.begin(), edgeIndexExpected.end());
0067 }
0068
0069 void testConstructionAndMove(Acts::ExecutionContext execContext) {
0070   auto tensor = Acts::Tensor<float>::Create({10, 1}, execContext);
0071
0072   BOOST_CHECK(tensor.shape()[1] == 1);
0073   BOOST_CHECK(tensor.shape()[0] == 10);
0074
0075   auto tensor2 = std::move(tensor);
0076   BOOST_CHECK(tensor2.shape()[1] == 1);
0077   BOOST_CHECK(tensor2.shape()[0] == 10);
0078   BOOST_CHECK(tensor2.data() != nullptr);
0079   BOOST_CHECK(tensor.data() == nullptr);
0080 }
0081
0082 BOOST_AUTO_TEST_CASE(tensor_create_move_cpu) {
0083   testConstructionAndMove(execContextCpu);
0084 }
0085
0086 BOOST_AUTO_TEST_CASE(test_clone_cpu) {
0087   std::vector<float> data = {1.f, 2.f, 3.f, 4.f, 5.f, 6.f};
0088   auto tensor = createCpuTensor(data, {3, 2});
0089   auto tensorClone = tensor.clone(execContextCpu);
0090
0091   BOOST_CHECK(tensorClone.shape()[0] == 3);
0092   BOOST_CHECK(tensorClone.shape()[1] == 2);
0093   BOOST_CHECK(tensorClone.data() != nullptr);
0094   BOOST_CHECK(tensorClone.data() != tensor.data());
0095   BOOST_CHECK(tensorClone.size() == tensor.size());
0096   BOOST_CHECK(tensorClone.nbytes() == tensor.nbytes());
0097
0098   BOOST_CHECK_EQUAL_COLLECTIONS(tensorClone.data(),
0099                                 tensorClone.data() + tensorClone.size(),
0100                                 data.begin(), data.end());
0101 }
0102
0103 BOOST_AUTO_TEST_CASE(tensor_sigmoid_cpu) {
0104   testSigmoid({-2.f, -1.f, 0.f, 1.f, 2.f}, execContextCpu);
0105 }
0106
0107 const std::vector<float> scores = {0.1f, 0.4f, 0.6f, 0.9f};
0108 const std::vector<std::int64_t> edgeIndex = {0, 1, 2, 3, 4, 5, 6, 7};
0109 const std::vector<std::int64_t> edgeIndexExpected = {2, 3, 6, 7};
0110
0111 BOOST_AUTO_TEST_CASE(tensor_edge_selection_cpu) {
0112   testEdgeSelection(scores, edgeIndex, edgeIndexExpected, execContextCpu);
0113 }
0114
0115 #ifdef ACTS_EXATRKX_WITH_CUDA
0116
0117 const Acts::ExecutionContext execContextCuda{Acts::Device::Cuda(0),
0118                                              cudaStreamLegacy};
0119
0120 BOOST_AUTO_TEST_CASE(tensor_create_move_cuda) {
0121   testConstructionAndMove(execContextCuda);
0122 }
0123
0124 BOOST_AUTO_TEST_CASE(tensor_clone_roundtrip) {
0125   std::vector<float> data = {1.f, 2.f, 3.f, 4.f, 5.f, 6.f};
0126   auto tensorOrigHost = createCpuTensor(data, {3, 2});
0127
0128   auto tensorClone = tensorOrigHost.clone(execContextCuda);
0129   auto tensorCloneCuda = tensorClone.clone(execContextCuda);
0130   auto tensorCloneHost = tensorCloneCuda.clone(execContextCpu);
0131
0132   BOOST_CHECK(tensorCloneHost.shape()[0] == 3);
0133   BOOST_CHECK(tensorCloneHost.shape()[1] == 2);
0134   BOOST_CHECK(tensorCloneHost.data() != nullptr);
0135   BOOST_CHECK(tensorCloneHost.data() != tensorCloneCuda.data());
0136   BOOST_CHECK(tensorCloneHost.size() == tensorCloneCuda.size());
0137   BOOST_CHECK(tensorCloneHost.nbytes() == tensorCloneCuda.nbytes());
0138   BOOST_CHECK_EQUAL_COLLECTIONS(tensorCloneHost.data(),
0139                                 tensorCloneHost.data() + tensorCloneHost.size(),
0140                                 data.begin(), data.end());
0141 }
0142
0143 BOOST_AUTO_TEST_CASE(tensor_sigmoid_cuda) {
0144   testSigmoid({-2.f, -1.f, 0.f, 1.f, 2.f}, execContextCuda);
0145 }
0146
0147 BOOST_AUTO_TEST_CASE(tensor_edge_selection_cuda) {
0148   testEdgeSelection(scores, edgeIndex, edgeIndexExpected, execContextCuda);
0149 }
0150
0151 #endif