Plugins/Gnn/GnnRemoveUnusedNodesTests.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/Stages.hpp"
0012
0013 #ifdef ACTS_GNN_WITH_CUDA
0014 #include <cuda_runtime_api.h>
0015 #endif
0016
0017 using namespace ActsPlugins;
0018
0019 namespace ActsTests {
0020
0021 namespace {
0022
0023 const ExecutionContext execContextCpu{Device::Cpu(), {}};
0024 #ifdef ACTS_GNN_WITH_CUDA
0025 const ExecutionContext execContextCuda{Device::Cuda(0), cudaStreamLegacy};
0026 #endif
0027
0028 // Build a PipelineTensors on the given device from flat CPU data.
0029 PipelineTensors makeTensors(std::size_t nNodes, std::size_t nFeatures,
0030                             std::vector<std::int64_t> edges,
0031                             const ExecutionContext &ctx) {
0032   const auto nEdges = edges.size() / 2;
0033
0034   auto nodeFeaturesCpu =
0035       Tensor<float>::Create({nNodes, nFeatures}, execContextCpu);
0036   for (std::size_t i = 0; i < nNodes * nFeatures; ++i) {
0037     nodeFeaturesCpu.data()[i] = static_cast<float>(i);
0038   }
0039
0040   auto edgeIndexCpu = Tensor<std::int64_t>::Create({2, nEdges}, execContextCpu);
0041   std::copy(edges.begin(), edges.end(), edgeIndexCpu.data());
0042
0043   return {std::move(nodeFeaturesCpu).clone(ctx),
0044           std::move(edgeIndexCpu).clone(ctx), std::nullopt, std::nullopt};
0045 }
0046
0047 // Read tensor data back to a host vector regardless of device.
0048 template <typename T>
0049 std::vector<T> toHost(const Tensor<T> &tensor, const ExecutionContext &srcCtx) {
0050   auto host = tensor.clone({Device::Cpu(), srcCtx.stream});
0051   return {host.data(), host.data() + host.size()};
0052 }
0053
0054 void testAllNodesUsed(const ExecutionContext &ctx) {
0055   // 3 nodes, all referenced: edges 0→1, 1→2
0056   auto tensors = makeTensors(3, 4, {0, 1, 1, 2}, ctx);
0057   const auto originalNodeData = toHost(tensors.nodeFeatures, ctx);
0058
0059   std::vector<int> spIds = {10, 11, 12};
0060   auto result = removeUnusedNodes(std::move(tensors), spIds, ctx);
0061
0062   BOOST_CHECK_EQUAL(result.nodeFeatures.shape()[0], 3u);
0063   BOOST_CHECK_EQUAL(result.nodeFeatures.shape()[1], 4u);
0064   BOOST_CHECK_EQUAL(result.edgeIndex.shape()[1], 2u);
0065
0066   BOOST_REQUIRE_EQUAL(spIds.size(), 3u);
0067   BOOST_CHECK_EQUAL(spIds[0], 10);
0068   BOOST_CHECK_EQUAL(spIds[1], 11);
0069   BOOST_CHECK_EQUAL(spIds[2], 12);
0070
0071   const auto nodeData = toHost(result.nodeFeatures, ctx);
0072   BOOST_CHECK_EQUAL_COLLECTIONS(nodeData.begin(), nodeData.end(),
0073                                 originalNodeData.begin(),
0074                                 originalNodeData.end());
0075
0076   const auto edgeData = toHost(result.edgeIndex, ctx);
0077   BOOST_CHECK_EQUAL(edgeData[0], 0);
0078   BOOST_CHECK_EQUAL(edgeData[1], 1);
0079   BOOST_CHECK_EQUAL(edgeData[2], 1);
0080   BOOST_CHECK_EQUAL(edgeData[3], 2);
0081 }
0082
0083 void testUnusedNodesRemoved(const ExecutionContext &ctx) {
0084   // 5 nodes, nodes 1 and 3 unreferenced: edges 0→2, 2→4
0085   // Used: {0,2,4} → compacted to {0,1,2}
0086   auto tensors = makeTensors(5, 2, {0, 2, 2, 4}, ctx);
0087   std::vector<int> spIds = {100, 101, 102, 103, 104};
0088
0089   auto result = removeUnusedNodes(std::move(tensors), spIds, ctx);
0090
0091   BOOST_CHECK_EQUAL(result.nodeFeatures.shape()[0], 3u);
0092   BOOST_CHECK_EQUAL(result.nodeFeatures.shape()[1], 2u);
0093   BOOST_CHECK_EQUAL(result.edgeIndex.shape()[1], 2u);
0094
0095   BOOST_REQUIRE_EQUAL(spIds.size(), 3u);
0096   BOOST_CHECK_EQUAL(spIds[0], 100);
0097   BOOST_CHECK_EQUAL(spIds[1], 102);
0098   BOOST_CHECK_EQUAL(spIds[2], 104);
0099
0100   const auto edgeData = toHost(result.edgeIndex, ctx);
0101   BOOST_CHECK_EQUAL(edgeData[0], 0);  // source of edge 0
0102   BOOST_CHECK_EQUAL(edgeData[1], 1);  // source of edge 1
0103   BOOST_CHECK_EQUAL(edgeData[2], 1);  // dest of edge 0
0104   BOOST_CHECK_EQUAL(edgeData[3], 2);  // dest of edge 1
0105
0106   // With nFeatures=2, row i has values {2*i, 2*i+1}; kept rows: 0, 2, 4
0107   const auto nodeData = toHost(result.nodeFeatures, ctx);
0108   BOOST_CHECK_EQUAL(nodeData[0], 0.f);  // old row 0, col 0
0109   BOOST_CHECK_EQUAL(nodeData[1], 1.f);  // old row 0, col 1
0110   BOOST_CHECK_EQUAL(nodeData[2], 4.f);  // old row 2, col 0
0111   BOOST_CHECK_EQUAL(nodeData[3], 5.f);  // old row 2, col 1
0112   BOOST_CHECK_EQUAL(nodeData[4], 8.f);  // old row 4, col 0
0113   BOOST_CHECK_EQUAL(nodeData[5], 9.f);  // old row 4, col 1
0114 }
0115
0116 void testEdgeAndScoreTensorsPreserved(const ExecutionContext &ctx) {
0117   // 4 nodes, single edge 1→2; nodes 0 and 3 unused
0118   auto nodeFeaturesCpu = Tensor<float>::Create({4, 1}, execContextCpu);
0119   auto edgeIndexCpu = Tensor<std::int64_t>::Create({2, 1}, execContextCpu);
0120   edgeIndexCpu.data()[0] = 1;
0121   edgeIndexCpu.data()[1] = 2;
0122   auto edgeFeaturesCpu = Tensor<float>::Create({1, 3}, execContextCpu);
0123   edgeFeaturesCpu.data()[0] = 7.f;
0124   edgeFeaturesCpu.data()[1] = 8.f;
0125   edgeFeaturesCpu.data()[2] = 9.f;
0126   auto edgeScoresCpu = Tensor<float>::Create({1, 1}, execContextCpu);
0127   edgeScoresCpu.data()[0] = 0.99f;
0128
0129   PipelineTensors tensors{nodeFeaturesCpu.clone(ctx), edgeIndexCpu.clone(ctx),
0130                           edgeFeaturesCpu.clone(ctx), edgeScoresCpu.clone(ctx)};
0131
0132   std::vector<int> spIds = {10, 11, 12, 13};
0133   auto result = removeUnusedNodes(std::move(tensors), spIds, ctx);
0134
0135   BOOST_CHECK_EQUAL(result.nodeFeatures.shape()[0], 2u);
0136   BOOST_REQUIRE_EQUAL(spIds.size(), 2u);
0137   BOOST_CHECK_EQUAL(spIds[0], 11);
0138   BOOST_CHECK_EQUAL(spIds[1], 12);
0139
0140   const auto edgeData = toHost(result.edgeIndex, ctx);
0141   BOOST_CHECK_EQUAL(edgeData[0], 0);
0142   BOOST_CHECK_EQUAL(edgeData[1], 1);
0143
0144   BOOST_REQUIRE(result.edgeFeatures.has_value());
0145   const auto featData = toHost(*result.edgeFeatures, ctx);
0146   BOOST_CHECK_EQUAL(featData[0], 7.f);
0147   BOOST_CHECK_EQUAL(featData[1], 8.f);
0148   BOOST_CHECK_EQUAL(featData[2], 9.f);
0149
0150   BOOST_REQUIRE(result.edgeScores.has_value());
0151   const auto scoreData = toHost(*result.edgeScores, ctx);
0152   BOOST_CHECK_EQUAL(scoreData[0], 0.99f);
0153 }
0154
0155 }  // namespace
0156
0157 BOOST_AUTO_TEST_SUITE(GnnRemoveUnusedNodesSuite)
0158
0159 BOOST_AUTO_TEST_CASE(test_all_nodes_used_cpu) {
0160   testAllNodesUsed(execContextCpu);
0161 }
0162
0163 BOOST_AUTO_TEST_CASE(test_unused_nodes_removed_cpu) {
0164   testUnusedNodesRemoved(execContextCpu);
0165 }
0166
0167 BOOST_AUTO_TEST_CASE(test_zero_edges_throws) {
0168   auto nodeFeatures = Tensor<float>::Create({4, 2}, execContextCpu);
0169   auto edgeIndex = Tensor<std::int64_t>::Create({2, 0}, execContextCpu);
0170   PipelineTensors tensors{std::move(nodeFeatures), std::move(edgeIndex),
0171                           std::nullopt, std::nullopt};
0172   std::vector<int> spIds = {0, 1, 2, 3};
0173   BOOST_CHECK_THROW(
0174       removeUnusedNodes(std::move(tensors), spIds, execContextCpu),
0175       NoEdgesError);
0176 }
0177
0178 BOOST_AUTO_TEST_CASE(test_edge_and_score_tensors_preserved_cpu) {
0179   testEdgeAndScoreTensorsPreserved(execContextCpu);
0180 }
0181
0182 #ifdef ACTS_GNN_WITH_CUDA
0183
0184 BOOST_AUTO_TEST_CASE(test_all_nodes_used_cuda) {
0185   testAllNodesUsed(execContextCuda);
0186 }
0187
0188 BOOST_AUTO_TEST_CASE(test_unused_nodes_removed_cuda) {
0189   testUnusedNodesRemoved(execContextCuda);
0190 }
0191
0192 BOOST_AUTO_TEST_CASE(test_edge_and_score_tensors_preserved_cuda) {
0193   testEdgeAndScoreTensorsPreserved(execContextCuda);
0194 }
0195
0196 #endif  // ACTS_GNN_WITH_CUDA
0197
0198 BOOST_AUTO_TEST_SUITE_END()
0199
0200 }  // namespace ActsTests