tests/hip/test_thrust.hip

0001 /** TRACCC library, part of the ACTS project (R&D line)
0002  *
0003  * (c) 2024-2026 CERN for the benefit of the ACTS project
0004  *
0005  * Mozilla Public License Version 2.0
0006  */
0007
0008 // VecMem include(s).
0009 #include <vecmem/containers/data/vector_buffer.hpp>
0010 #include <vecmem/containers/device_vector.hpp>
0011 #include <vecmem/containers/vector.hpp>
0012 #include <vecmem/memory/hip/device_memory_resource.hpp>
0013 #include <vecmem/memory/host_memory_resource.hpp>
0014 #include <vecmem/utils/copy.hpp>
0015 #include <vecmem/utils/hip/copy.hpp>
0016
0017 // Thrust include(s).
0018 #include <thrust/execution_policy.h>
0019 #include <thrust/fill.h>
0020 #include <thrust/scan.h>
0021 #include <thrust/sort.h>
0022
0023 // GTest include(s).
0024 #include <gtest/gtest.h>
0025
0026 // This defines the local frame test suite
0027
0028 namespace {
0029 vecmem::hip::copy copy;
0030 vecmem::host_memory_resource host_resource;
0031 vecmem::hip::device_memory_resource device_resource;
0032
0033 }  // namespace
0034
0035 TEST(thrust, sort) {
0036   vecmem::vector<unsigned int> host_vector{{3, 2, 1, 8, 4}, &host_resource};
0037
0038   auto host_buffer = vecmem::get_data(host_vector);
0039   auto device_buffer = copy.to(vecmem::get_data(host_vector), device_resource,
0040                                vecmem::copy::type::host_to_device);
0041
0042   vecmem::device_vector<unsigned int> device_vector(device_buffer);
0043
0044   thrust::sort(thrust::device, device_vector.begin(), device_vector.end());
0045
0046   copy(device_buffer, host_buffer, vecmem::copy::type::device_to_host)->wait();
0047
0048   ASSERT_EQ(host_vector[0], 1);
0049   ASSERT_EQ(host_vector[1], 2);
0050   ASSERT_EQ(host_vector[2], 3);
0051   ASSERT_EQ(host_vector[3], 4);
0052   ASSERT_EQ(host_vector[4], 8);
0053 }
0054
0055 TEST(thrust, scan) {
0056   vecmem::vector<unsigned int> host_vector{{3, 2, 1, 8, 4}, &host_resource};
0057
0058   auto host_buffer = vecmem::get_data(host_vector);
0059   auto device_buffer = copy.to(vecmem::get_data(host_vector), device_resource,
0060                                vecmem::copy::type::host_to_device);
0061
0062   vecmem::device_vector<unsigned int> device_vector(device_buffer);
0063
0064   thrust::inclusive_scan(thrust::device, device_vector.begin(),
0065                          device_vector.end(), device_vector.begin());
0066
0067   copy(device_buffer, host_buffer, vecmem::copy::type::device_to_host)->wait();
0068
0069   ASSERT_EQ(host_vector[0], 3);
0070   ASSERT_EQ(host_vector[1], 5);
0071   ASSERT_EQ(host_vector[2], 6);
0072   ASSERT_EQ(host_vector[3], 14);
0073   ASSERT_EQ(host_vector[4], 18);
0074 }
0075
0076 TEST(thrust, fill) {
0077   vecmem::vector<unsigned int> host_vector{{1, 1, 1, 1, 1, 1, 1},
0078                                            &host_resource};
0079
0080   auto host_buffer = vecmem::get_data(host_vector);
0081   auto device_buffer = copy.to(vecmem::get_data(host_vector), device_resource,
0082                                vecmem::copy::type::host_to_device);
0083
0084   vecmem::device_vector<unsigned int> device_vector(device_buffer);
0085
0086   thrust::fill(thrust::device, device_vector.begin(), device_vector.end(), 112);
0087
0088   copy(device_buffer, host_buffer, vecmem::copy::type::device_to_host)->wait();
0089
0090   ASSERT_EQ(host_vector[0], 112);
0091   ASSERT_EQ(host_vector[1], 112);
0092   ASSERT_EQ(host_vector[2], 112);
0093   ASSERT_EQ(host_vector[3], 112);
0094   ASSERT_EQ(host_vector[4], 112);
0095   ASSERT_EQ(host_vector[5], 112);
0096   ASSERT_EQ(host_vector[6], 112);
0097 }