tests/cuda/test_thrust.cu

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