| /****************************************************************************** | |
| * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. | |
| * | |
| * Redistribution and use in source and binary forms, with or without | |
| * modification, are permitted provided that the following conditions are met: | |
| * * Redistributions of source code must retain the above copyright | |
| * notice, this list of conditions and the following disclaimer. | |
| * * Redistributions in binary form must reproduce the above copyright | |
| * notice, this list of conditions and the following disclaimer in the | |
| * documentation and/or other materials provided with the distribution. | |
| * * Neither the name of the NVIDIA CORPORATION nor the | |
| * names of its contributors may be used to endorse or promote products | |
| * derived from this software without specific prior written permission. | |
| * | |
| * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND | |
| * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED | |
| * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE | |
| * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY | |
| * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES | |
| * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; | |
| * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND | |
| * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT | |
| * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS | |
| * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. | |
| * | |
| ******************************************************************************/ | |
| /** | |
| * \file | |
| * cub::DeviceCopy provides device-wide, parallel operations for copying data. | |
| */ | |
| #pragma once | |
| #include <cub/config.cuh> | |
| #include <cub/device/dispatch/dispatch_batch_memcpy.cuh> | |
| #include <thrust/system/cuda/detail/core/triple_chevron_launch.h> | |
| #include <cstdint> | |
| CUB_NAMESPACE_BEGIN | |
| /** | |
| * @brief cub::DeviceCopy provides device-wide, parallel operations for copying data. | |
| * \ingroup SingleModule | |
| */ | |
| struct DeviceCopy | |
| { | |
| /** | |
| * @brief Copies data from a batch of given source ranges to their corresponding destination | |
| * ranges. | |
| * @note If any input range aliases any output range the behavior is undefined. If | |
| * any output range aliases another output range the behavior is undefined. Input | |
| * ranges can alias one another. | |
| * | |
| * @par Snippet | |
| * The code snippet below illustrates usage of DeviceCopy::Batched to perform a DeviceRunLength | |
| * Decode operation. | |
| * @par | |
| * @code | |
| * struct GetIteratorToRange | |
| * { | |
| * __host__ __device__ __forceinline__ auto operator()(uint32_t index) | |
| * { | |
| * return thrust::make_constant_iterator(d_data_in[index]); | |
| * } | |
| * int32_t *d_data_in; | |
| * }; | |
| * | |
| * struct GetPtrToRange | |
| * { | |
| * __host__ __device__ __forceinline__ auto operator()(uint32_t index) | |
| * { | |
| * return d_data_out + d_offsets[index]; | |
| * } | |
| * int32_t *d_data_out; | |
| * uint32_t *d_offsets; | |
| * }; | |
| * | |
| * struct GetRunLength | |
| * { | |
| * __host__ __device__ __forceinline__ uint32_t operator()(uint32_t index) | |
| * { | |
| * return d_offsets[index + 1] - d_offsets[index]; | |
| * } | |
| * uint32_t *d_offsets; | |
| * }; | |
| * | |
| * uint32_t num_ranges = 5; | |
| * int32_t *d_data_in; // e.g., [4, 2, 7, 3, 1] | |
| * int32_t *d_data_out; // e.g., [0, ... ] | |
| * uint32_t *d_offsets; // e.g., [0, 2, 5, 6, 9, 14] | |
| * | |
| * // Returns a constant iterator to the element of the i-th run | |
| * thrust::counting_iterator<uint32_t> iota(0); | |
| * auto iterators_in = thrust::make_transform_iterator(iota, GetIteratorToRange{d_data_in}); | |
| * | |
| * // Returns the run length of the i-th run | |
| * auto sizes = thrust::make_transform_iterator(iota, GetRunLength{d_offsets}); | |
| * | |
| * // Returns pointers to the output range for each run | |
| * auto ptrs_out = thrust::make_transform_iterator(iota, GetPtrToRange{d_data_out, d_offsets}); | |
| * | |
| * // Determine temporary device storage requirements | |
| * void *d_temp_storage = nullptr; | |
| * size_t temp_storage_bytes = 0; | |
| * cub::DeviceCopy::Batched(d_temp_storage, temp_storage_bytes, iterators_in, ptrs_out, sizes, | |
| * num_ranges); | |
| * | |
| * // Allocate temporary storage | |
| * cudaMalloc(&d_temp_storage, temp_storage_bytes); | |
| * | |
| * // Run batched copy algorithm (used to perform runlength decoding) | |
| * cub::DeviceCopy::Batched(d_temp_storage, temp_storage_bytes, iterators_in, ptrs_out, sizes, | |
| * num_ranges); | |
| * | |
| * // d_data_out <-- [4, 4, 2, 2, 2, 7, 3, 3, 3, 1, 1, 1, 1, 1] | |
| * @endcode | |
| * @tparam InputIt <b>[inferred]</b> Device-accessible random-access input iterator type | |
| * providing the iterators to the source ranges | |
| * @tparam OutputIt <b>[inferred]</b> Device-accessible random-access input iterator type | |
| * providing the iterators to the destination ranges | |
| * @tparam SizeIteratorT <b>[inferred]</b> Device-accessible random-access input iterator | |
| * type providing the number of items to be copied for each pair of ranges | |
| * @param d_temp_storage [in] Device-accessible allocation of temporary storage. When NULL, the | |
| * required allocation size is written to \p temp_storage_bytes and no work is done. | |
| * @param temp_storage_bytes [in,out] Reference to size in bytes of \p d_temp_storage allocation | |
| * @param input_it [in] Device-accessible iterator providing the iterators to the source | |
| * ranges | |
| * @param output_it [in] Device-accessible iterator providing the iterators to the | |
| * destination ranges | |
| * @param sizes [in] Device-accessible iterator providing the number of elements to be copied | |
| * for each pair of ranges | |
| * @param num_ranges [in] The total number of range pairs | |
| * @param stream [in] <b>[optional]</b> CUDA stream to launch kernels within. Default is | |
| * stream<sub>0</sub>. | |
| */ | |
| template <typename InputIt, typename OutputIt, typename SizeIteratorT> | |
| CUB_RUNTIME_FUNCTION static cudaError_t Batched(void *d_temp_storage, | |
| size_t &temp_storage_bytes, | |
| InputIt input_it, | |
| OutputIt output_it, | |
| SizeIteratorT sizes, | |
| uint32_t num_ranges, | |
| cudaStream_t stream = 0) | |
| { | |
| // Integer type large enough to hold any offset in [0, num_ranges) | |
| using RangeOffsetT = uint32_t; | |
| // Integer type large enough to hold any offset in [0, num_thread_blocks_launched), where a safe | |
| // uppper bound on num_thread_blocks_launched can be assumed to be given by | |
| // IDIV_CEIL(num_ranges, 64) | |
| using BlockOffsetT = uint32_t; | |
| return detail::DispatchBatchMemcpy<InputIt, | |
| OutputIt, | |
| SizeIteratorT, | |
| RangeOffsetT, | |
| BlockOffsetT, | |
| detail::DeviceBatchMemcpyPolicy, | |
| false>::Dispatch(d_temp_storage, | |
| temp_storage_bytes, | |
| input_it, | |
| output_it, | |
| sizes, | |
| num_ranges, | |
| stream); | |
| } | |
| }; | |
| CUB_NAMESPACE_END | |