|
| 1 | +#include <thrust/detail/config.h> |
| 2 | + |
| 3 | +#if THRUST_CPP_DIALECT >= 2011 |
| 4 | + |
| 5 | +# include <unittest/unittest.h> |
| 6 | + |
| 7 | +# include <thrust/allocate_unique.h> |
| 8 | +# include <thrust/memory/detail/device_system_resource.h> |
| 9 | +# include <thrust/mr/allocator.h> |
| 10 | +# include <thrust/type_traits/is_contiguous_iterator.h> |
| 11 | + |
| 12 | +# include <numeric> |
| 13 | +# include <vector> |
| 14 | + |
| 15 | +namespace |
| 16 | +{ |
| 17 | + |
| 18 | +template <typename T> |
| 19 | +using allocator = |
| 20 | + thrust::mr::stateless_resource_allocator<T, thrust::universal_memory_resource>; |
| 21 | + |
| 22 | +// The managed_memory_pointer class should be identified as a |
| 23 | +// contiguous_iterator |
| 24 | +THRUST_STATIC_ASSERT( |
| 25 | + thrust::is_contiguous_iterator<allocator<int>::pointer>::value); |
| 26 | + |
| 27 | +template <typename T> |
| 28 | +struct some_object { |
| 29 | + some_object(T data) |
| 30 | + : m_data(data) |
| 31 | + {} |
| 32 | + |
| 33 | + void setter(T data) { m_data = data; } |
| 34 | + T getter() const { return m_data; } |
| 35 | + |
| 36 | +private: |
| 37 | + T m_data; |
| 38 | +}; |
| 39 | + |
| 40 | +} // namespace |
| 41 | + |
| 42 | +template <typename T> |
| 43 | +void TestAllocateUnique() |
| 44 | +{ |
| 45 | + // Simple test to ensure that pointers created with universal_memory_resource |
| 46 | + // can be dereferenced and used with STL code. This is necessary as some |
| 47 | + // STL implementations break when using fancy references that overload |
| 48 | + // `operator&`, so universal_memory_resource uses a special pointer type that |
| 49 | + // returns regular C++ references that can be safely used host-side. |
| 50 | + |
| 51 | + // These operations fail to compile with fancy references: |
| 52 | + auto pRaw = thrust::allocate_unique<T>(allocator<T>{}, 42); |
| 53 | + auto pObj = |
| 54 | + thrust::allocate_unique<some_object<T> >(allocator<some_object<T> >{}, 42); |
| 55 | + |
| 56 | + static_assert( |
| 57 | + std::is_same<decltype(pRaw.get()), |
| 58 | + thrust::system::cuda::detail::managed_memory_pointer<T> >::value, |
| 59 | + "Unexpected pointer returned from unique_ptr::get."); |
| 60 | + static_assert( |
| 61 | + std::is_same<decltype(pObj.get()), |
| 62 | + thrust::system::cuda::detail::managed_memory_pointer< |
| 63 | + some_object<T> > >::value, |
| 64 | + "Unexpected pointer returned from unique_ptr::get."); |
| 65 | + |
| 66 | + ASSERT_EQUAL(*pRaw, T(42)); |
| 67 | + ASSERT_EQUAL(*pRaw.get(), T(42)); |
| 68 | + ASSERT_EQUAL(pObj->getter(), T(42)); |
| 69 | + ASSERT_EQUAL((*pObj).getter(), T(42)); |
| 70 | + ASSERT_EQUAL(pObj.get()->getter(), T(42)); |
| 71 | + ASSERT_EQUAL((*pObj.get()).getter(), T(42)); |
| 72 | +} |
| 73 | +DECLARE_GENERIC_UNITTEST(TestAllocateUnique); |
| 74 | + |
| 75 | +template <typename T> |
| 76 | +void TestIterationRaw() |
| 77 | +{ |
| 78 | + auto array = thrust::allocate_unique_n<T>(allocator<T>{}, 6, 42); |
| 79 | + |
| 80 | + static_assert( |
| 81 | + std::is_same<decltype(array.get()), |
| 82 | + thrust::system::cuda::detail::managed_memory_pointer<T> >::value, |
| 83 | + "Unexpected pointer returned from unique_ptr::get."); |
| 84 | + |
| 85 | + for (auto iter = array.get(), end = array.get() + 6; iter < end; ++iter) |
| 86 | + { |
| 87 | + ASSERT_EQUAL(*iter, T(42)); |
| 88 | + ASSERT_EQUAL(*iter.get(), T(42)); |
| 89 | + } |
| 90 | +} |
| 91 | +DECLARE_GENERIC_UNITTEST(TestIterationRaw); |
| 92 | + |
| 93 | +template <typename T> |
| 94 | +void TestIterationObj() |
| 95 | +{ |
| 96 | + auto array = |
| 97 | + thrust::allocate_unique_n<some_object<T> >(allocator<some_object<T> >{}, |
| 98 | + 6, |
| 99 | + 42); |
| 100 | + |
| 101 | + static_assert( |
| 102 | + std::is_same<decltype(array.get()), |
| 103 | + thrust::system::cuda::detail::managed_memory_pointer< |
| 104 | + some_object<T> > >::value, |
| 105 | + "Unexpected pointer returned from unique_ptr::get."); |
| 106 | + |
| 107 | + for (auto iter = array.get(), end = array.get() + 6; iter < end; ++iter) |
| 108 | + { |
| 109 | + ASSERT_EQUAL(iter->getter(), T(42)); |
| 110 | + ASSERT_EQUAL((*iter).getter(), T(42)); |
| 111 | + ASSERT_EQUAL(iter.get()->getter(), T(42)); |
| 112 | + ASSERT_EQUAL((*iter.get()).getter(), T(42)); |
| 113 | + } |
| 114 | +} |
| 115 | +DECLARE_GENERIC_UNITTEST(TestIterationObj); |
| 116 | + |
| 117 | +template <typename T> |
| 118 | +void TestStdVector() |
| 119 | +{ |
| 120 | + // Verify that a std::vector using the universal allocator will work with |
| 121 | + // STL algorithms. |
| 122 | + std::vector<T, allocator<T> > v0; |
| 123 | + |
| 124 | + static_assert( |
| 125 | + std::is_same<typename std::decay<decltype(v0)>::type::pointer, |
| 126 | + thrust::system::cuda::detail::managed_memory_pointer< |
| 127 | + T > >::value, |
| 128 | + "Unexpected pointer returned from unique_ptr::get."); |
| 129 | + |
| 130 | + v0.resize(6); |
| 131 | + std::iota(v0.begin(), v0.end(), 0); |
| 132 | + ASSERT_EQUAL(v0[0], T(0)); |
| 133 | + ASSERT_EQUAL(v0[1], T(1)); |
| 134 | + ASSERT_EQUAL(v0[2], T(2)); |
| 135 | + ASSERT_EQUAL(v0[3], T(3)); |
| 136 | + ASSERT_EQUAL(v0[4], T(4)); |
| 137 | + ASSERT_EQUAL(v0[5], T(5)); |
| 138 | +} |
| 139 | +DECLARE_GENERIC_UNITTEST(TestStdVector); |
| 140 | + |
| 141 | +#endif // C++11 |
0 commit comments