|
| 1 | +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. |
| 2 | +// SPDX-License-Identifier: BSD-3 |
| 3 | + |
| 4 | +#pragma once |
| 5 | + |
| 6 | +#include <cuda/std/detail/__config> |
| 7 | + |
| 8 | +#if _CCCL_HAS_CTK() && !_CCCL_COMPILER(NVRTC) |
| 9 | +# include <cuda/__algorithm/copy.h> |
| 10 | +# include <cuda/buffer> |
| 11 | +# include <cuda/devices> |
| 12 | +# include <cuda/std/limits> |
| 13 | +# include <cuda/std/utility> |
| 14 | +# include <cuda/stream> |
| 15 | +#endif // _CCCL_HAS_CTK() && !_CCCL_COMPILER(NVRTC) |
| 16 | + |
| 17 | +#include <cstddef> |
| 18 | + |
| 19 | +#include <c2h/checked_memory_resource.cuh> |
| 20 | +#include <c2h/generator_common.h> |
| 21 | + |
| 22 | +namespace c2h |
| 23 | +{ |
| 24 | +#if _CCCL_HAS_CTK() && !_CCCL_COMPILER(NVRTC) |
| 25 | +namespace detail |
| 26 | +{ |
| 27 | +template <typename T> |
| 28 | +[[nodiscard]] cuda::host_buffer<T> device_buffer_to_host_buffer( |
| 29 | + cuda::stream_ref stream, cuda::device_ref device, const cuda::device_buffer<T>& d_items, std::size_t num_items) |
| 30 | +{ |
| 31 | + // Scope `device` for host allocation checks and possible default-stream copies. |
| 32 | + // Non-default stream/device agreement is part of the public helper contract. |
| 33 | + const ::c2h::detail::scoped_current_device device_scope{device.get()}; |
| 34 | + |
| 35 | + auto h_items = ::c2h::make_host_buffer<T>(stream, device, num_items, cuda::no_init); |
| 36 | + cuda::copy_bytes(stream, d_items.first(num_items), h_items); |
| 37 | + stream.sync(); |
| 38 | + |
| 39 | + return h_items; |
| 40 | +} |
| 41 | + |
| 42 | +template <typename T> |
| 43 | +void gen_into_device_buffer(cuda::stream_ref stream, seed_t seed, cuda::device_buffer<T>& d_items, T min, T max) |
| 44 | +{ |
| 45 | + ::c2h::detail::gen_values_between(stream, seed, d_items.first(d_items.size()), min, max); |
| 46 | +} |
| 47 | + |
| 48 | +template <template <typename> class... Ps> |
| 49 | +void gen_into_device_buffer( |
| 50 | + cuda::stream_ref stream, |
| 51 | + seed_t seed, |
| 52 | + cuda::device_buffer< custom_type_t<Ps...>>& d_items, |
| 53 | + custom_type_t<Ps...> min, |
| 54 | + custom_type_t<Ps...> max) |
| 55 | +{ |
| 56 | + ::c2h::detail::gen_custom_type_state( |
| 57 | + stream, seed, reinterpret_cast< char*>(d_items.data()), min, max, d_items.size(), sizeof( custom_type_t<Ps...>)); |
| 58 | +} |
| 59 | +} // namespace detail |
| 60 | + |
| 61 | +// `size` is the number of generated items shared by both buffers. The owning |
| 62 | +// buffers may contain additional capacity that is not part of the generated sequence. |
| 63 | +template <typename T> |
| 64 | +struct sized_device_host_buffers |
| 65 | +{ |
| 66 | + cuda::device_buffer<T> d_items; |
| 67 | + cuda::host_buffer<T> h_items; |
| 68 | + std::size_t size; |
| 69 | +}; |
| 70 | + |
| 71 | +/** |
| 72 | + * @brief Generates random data with the existing c2h device generator and returns it in device memory. |
| 73 | + * |
| 74 | + * @pre If `stream` is non-default, it must have been created for `device`. |
| 75 | + */ |
| 76 | +template <typename T> |
| 77 | +[[nodiscard]] cuda::device_buffer<T> gen_device_buffer( |
| 78 | + cuda::stream_ref stream, |
| 79 | + cuda::device_ref device, |
| 80 | + seed_t seed, |
| 81 | + std::size_t num_items, |
| 82 | + T min = ::cuda::std::numeric_limits<T>::lowest(), |
| 83 | + T max = ::cuda::std::numeric_limits<T>::max()) |
| 84 | +{ |
| 85 | + // Scope `device` for generator storage backed by current-device allocation. |
| 86 | + const ::c2h::detail::scoped_current_device device_scope{device.get()}; |
| 87 | + |
| 88 | + auto d_items = ::c2h::make_device_buffer<T>(stream, device, num_items, cuda::no_init); |
| 89 | + ::c2h::detail::gen_into_device_buffer(stream, seed, d_items, min, max); |
| 90 | + |
| 91 | + return d_items; |
| 92 | +} |
| 93 | + |
| 94 | +/** |
| 95 | + * @brief Generates random data with the existing c2h device generator and returns device and host buffers. |
| 96 | + * |
| 97 | + * @pre If `stream` is non-default, it must have been created for `device`. |
| 98 | + */ |
| 99 | +template <typename T> |
| 100 | +[[nodiscard]] sized_device_host_buffers<T> gen_buffers( |
| 101 | + cuda::stream_ref stream, |
| 102 | + cuda::device_ref device, |
| 103 | + seed_t seed, |
| 104 | + std::size_t num_items, |
| 105 | + T min = ::cuda::std::numeric_limits<T>::lowest(), |
| 106 | + T max = ::cuda::std::numeric_limits<T>::max()) |
| 107 | +{ |
| 108 | + auto d_items = ::c2h::gen_device_buffer<T>(stream, device, seed, num_items, min, max); |
| 109 | + |
| 110 | + const auto items_count = d_items.size(); |
| 111 | + auto h_items = ::c2h::detail::device_buffer_to_host_buffer(stream, device, d_items, items_count); |
| 112 | + |
| 113 | + return {::cuda::std::move(d_items), ::cuda::std::move(h_items), items_count}; |
| 114 | +} |
| 115 | + |
| 116 | +/** |
| 117 | + * @brief Generates random data with the existing c2h device generator and returns it in host pageable memory. |
| 118 | + * |
| 119 | + * @pre If `stream` is non-default, it must have been created for `device`. |
| 120 | + */ |
| 121 | +template <typename T> |
| 122 | +[[nodiscard]] cuda::host_buffer<T> gen_host_buffer( |
| 123 | + cuda::stream_ref stream, |
| 124 | + cuda::device_ref device, |
| 125 | + seed_t seed, |
| 126 | + std::size_t num_items, |
| 127 | + T min = ::cuda::std::numeric_limits<T>::lowest(), |
| 128 | + T max = ::cuda::std::numeric_limits<T>::max()) |
| 129 | +{ |
| 130 | + auto buffers = ::c2h::gen_buffers<T>(stream, device, seed, num_items, min, max); |
| 131 | + return ::cuda::std::move(buffers.h_items); |
| 132 | +} |
| 133 | +#endif // _CCCL_HAS_CTK() && !_CCCL_COMPILER(NVRTC) |
| 134 | +} // namespace c2h |
0 commit comments