diff --git a/cub/benchmarks/bench/load/block_load.cu b/cub/benchmarks/bench/load/block_load.cu new file mode 100644 index 000000000000..66a58ffe1a89 --- /dev/null +++ b/cub/benchmarks/bench/load/block_load.cu @@ -0,0 +1,136 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include + +#include + +#include +#include + +#include + +struct increment +{ + template + _CCCL_DEVICE T operator()(T value) const + { + return static_cast(value + T{1}); + } +}; + +struct twice +{ + template + _CCCL_DEVICE T operator()(T value) const + { + return static_cast(value * T{2}); + } +}; + +// Load BaseT storage through a nested transform and keep a per-thread sum. The sum is written once per block so the +// transformed loads cannot be discarded. +template +__global__ void +block_load_transform(const BaseT* input, ::cuda::std::int64_t num_tiles, ::cuda::std::uint64_t* block_sums) +{ + constexpr int tile_size = BlockSize * ItemsPerThread; + using block_load_t = cub::BlockLoad; + using storage_t = typename block_load_t::TempStorage; + + __shared__ storage_t storage; + ::cuda::std::uint64_t sum = 0; + for (::cuda::std::int64_t tile = blockIdx.x; tile < num_tiles; tile += gridDim.x) + { + const BaseT* tile_input = input + tile * tile_size; + auto inner = ::cuda::transform_iterator{tile_input, increment{}}; + auto in = ::cuda::transform_iterator{inner, twice{}}; + + BaseT data[ItemsPerThread]; + block_load_t{storage}.Load(in, data); + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < ItemsPerThread; ++item) + { + sum += static_cast<::cuda::std::uint64_t>(data[item]); + } + } + + if (threadIdx.x == 0) + { + block_sums[blockIdx.x] = sum; + } +} + +template +void run_block_load(nvbench::state& state) +{ + constexpr int tile_size = BlockSize * ItemsPerThread; + const auto requested = state.get_int64("Elements{io}"); + const auto num_tiles = requested / tile_size; + if (num_tiles == 0) + { + state.skip("Skipping: element count is smaller than one tile."); + return; + } + const auto num_items = num_tiles * tile_size; + + thrust::device_vector input(static_cast(num_items)); + const auto kernel = block_load_transform; + const int num_sms = state.get_device().value().get_number_of_sms(); // NOLINT(bugprone-unchecked-optional-access) + int max_blocks_per_sm = 0; + NVBENCH_CUDA_CALL_NOEXCEPT(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&max_blocks_per_sm, kernel, BlockSize, 0)); + if (max_blocks_per_sm == 0) + { + state.skip("Skipping: no resident blocks for this load on this device."); + return; + } + const int grid_size = max_blocks_per_sm * num_sms; + thrust::device_vector<::cuda::std::uint64_t> block_sums(static_cast(grid_size)); + + state.add_element_count(num_items); + state.add_global_memory_reads(num_items); + state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) { + kernel<<>>( + thrust::raw_pointer_cast(input.data()), num_tiles, thrust::raw_pointer_cast(block_sums.data())); + }); +} + +template +void block_load(nvbench::state& state, nvbench::type_list) +{ + const bool vectorize = state.get_string("Algorithm") == "vectorize"; + const auto items_per_thread = state.get_int64("ItemsPerThread"); + const auto block_size = state.get_int64("BlockSize"); + + if (block_size == 128 && items_per_thread == 4) + { + return vectorize ? run_block_load(state) + : run_block_load(state); + } + if (block_size == 128 && items_per_thread == 8) + { + return vectorize ? run_block_load(state) + : run_block_load(state); + } + if (block_size == 256 && items_per_thread == 4) + { + return vectorize ? run_block_load(state) + : run_block_load(state); + } + if (block_size == 256 && items_per_thread == 8) + { + return vectorize ? run_block_load(state) + : run_block_load(state); + } + state.skip("Skipping: unsupported block size or items per thread."); +} + +NVBENCH_BENCH_TYPES(block_load, NVBENCH_TYPE_AXES(integral_types)) + .set_name("transform") + .set_type_axes_names({"BaseT{ct}"}) + .add_string_axis("Algorithm", {"vectorize", "direct"}) + .add_int64_axis("BlockSize", {128, 256}) + .add_int64_axis("ItemsPerThread", {4, 8}) + .add_int64_power_of_two_axis("Elements{io}", nvbench::range(20, 28, 4)); diff --git a/cub/benchmarks/bench/reduce/contiguous_adaptor.cu b/cub/benchmarks/bench/reduce/contiguous_adaptor.cu new file mode 100644 index 000000000000..a22a22a84496 --- /dev/null +++ b/cub/benchmarks/bench/reduce/contiguous_adaptor.cu @@ -0,0 +1,65 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#include + +#include + +#include + +// Same-type transform over a device pointer. The stored values are contiguous; the iterator itself is not. +template +struct plus_one_t +{ + __host__ __device__ T operator()(T value) const + { + return static_cast(value + T{1}); + } +}; + +template +void sum(nvbench::state& state, nvbench::type_list) +try +{ + const auto elements = state.get_int64("Elements{io}"); + const auto source = state.get_string("Source"); + const bool as_transform = source == "transform"; + + thrust::device_vector in = generate(static_cast(elements)); + thrust::device_vector out(1); + + auto* d_in = thrust::raw_pointer_cast(in.data()); + auto* d_out = thrust::raw_pointer_cast(out.data()); + + state.add_element_count(elements); + state.add_global_memory_reads(elements, "Size"); + state.add_global_memory_writes(1); + + caching_allocator_t alloc; + state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) { + auto env = cub_bench_env(alloc, launch); + if (as_transform) + { + auto it = cuda::transform_iterator(d_in, plus_one_t{}); + _CCCL_TRY_RUNTIME_API( + cub::DeviceReduce::Sum, "Reduce failed", it, d_out, static_cast(elements), env); + } + else + { + _CCCL_TRY_RUNTIME_API( + cub::DeviceReduce::Sum, "Reduce failed", d_in, d_out, static_cast(elements), env); + } + }); +} +catch (const std::bad_alloc&) +{ + state.skip("Skipping: out of memory."); +} + +using value_types = nvbench::type_list; + +NVBENCH_BENCH_TYPES(sum, NVBENCH_TYPE_AXES(value_types)) + .set_name("contiguous_adaptor") + .set_type_axes_names({"T{ct}"}) + .add_string_axis("Source", {"pointer", "transform"}) + .add_int64_power_of_two_axis("Elements{io}", nvbench::range(24, 28, 4)); diff --git a/cub/benchmarks/bench/topk/transform_iterator.cu b/cub/benchmarks/bench/topk/transform_iterator.cu new file mode 100644 index 000000000000..b303ea39d9cf --- /dev/null +++ b/cub/benchmarks/bench/topk/transform_iterator.cu @@ -0,0 +1,96 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include + +// DeviceTopK is the device algorithm whose shipped tuning selects BLOCK_LOAD_VECTORIZE. The input is a +// transform_iterator over const char*, so the vectorized load copies the underlying bytes and the functor runs when +// the loaded items are read back. +struct widen +{ + [[nodiscard]] _CCCL_HOST_DEVICE uint32_t operator()(char value) const + { + return static_cast(static_cast(value)); + } +}; + +// Same tile shape as the shipped uint32 policy, with the load forced to BLOCK_LOAD_DIRECT. +struct direct_load_policy +{ + [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(::cuda::compute_capability cc) const + { + auto policy = cub::detail::topk::policy_selector_from_types{}(cc); + return cub::detail::topk::topk_policy{ + policy.threads_per_block, + policy.items_per_thread, + cub::BLOCK_LOAD_DIRECT, + policy.scan_algorithm, + policy.bits_per_pass}; + } +}; + +template +void max_keys(nvbench::state& state, int64_t elements, int selected_elements, Env env) +{ + thrust::device_vector input(static_cast(elements), thrust::no_init); + thrust::device_vector output(static_cast(selected_elements), thrust::no_init); + + const auto in = cuda::transform_iterator{thrust::raw_pointer_cast(input.data()), widen{}}; + auto* out = thrust::raw_pointer_cast(output.data()); + + state.add_element_count(elements, "NumElements"); + state.add_element_count(selected_elements, "NumSelectedElements"); + // Traffic is the underlying char storage, not the uint32 key the functor returns. + state.add_global_memory_reads(elements, "InputKeys"); + state.add_global_memory_writes(selected_elements, "OutputKeys"); + + size_t temp_size{}; + cub::DeviceTopK::MaxKeys(nullptr, temp_size, in, out, elements, selected_elements, env); + thrust::device_vector temp(temp_size, thrust::no_init); + + auto* temp_storage = thrust::raw_pointer_cast(temp.data()); + state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) { + auto env_with_stream = cuda::std::execution::env{cuda::stream_ref{launch.get_stream().get_stream()}, env}; + cub::DeviceTopK::MaxKeys(temp_storage, temp_size, in, out, elements, selected_elements, env_with_stream); + }); +} + +void topk_transform(nvbench::state& state) +{ + const auto elements = state.get_int64("Elements{io}"); + constexpr int selected = 1024; + const bool vectorize = state.get_string("Load") == "vectorize"; + + if (selected >= elements) + { + state.skip("Selected count must be smaller than the input."); + return; + } + + auto requirements = + cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted); + if (vectorize) + { + max_keys(state, elements, selected, cuda::std::execution::env{requirements}); + } + else + { + max_keys( + state, elements, selected, cuda::std::execution::env{requirements, cuda::execution::tune(direct_load_policy{})}); + } +} + +NVBENCH_BENCH(topk_transform) + .set_name("transform") + .add_string_axis("Load", {"vectorize", "direct"}) + .add_int64_power_of_two_axis("Elements{io}", nvbench::range(20, 28, 4)); diff --git a/cub/benchmarks/bench/transform/contiguous_adaptor.cu b/cub/benchmarks/bench/transform/contiguous_adaptor.cu new file mode 100644 index 000000000000..8ddb5a269947 --- /dev/null +++ b/cub/benchmarks/bench/transform/contiguous_adaptor.cu @@ -0,0 +1,77 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#include + +#include +#include + +#include + +// Same-type transform over a device pointer. The stored values are contiguous; the iterator itself is not. +template +struct plus_one_t +{ + __host__ __device__ T operator()(T value) const + { + return static_cast(value + T{1}); + } +}; + +template +void copy(nvbench::state& state, nvbench::type_list) +try +{ + const auto elements = state.get_int64("Elements{io}"); + const auto source = state.get_string("Source"); + const bool as_transform = source == "transform"; + + thrust::device_vector in = generate(static_cast(elements)); + thrust::device_vector out = generate(static_cast(elements)); + + auto* d_in = thrust::raw_pointer_cast(in.data()); + auto* d_out = thrust::raw_pointer_cast(out.data()); + + state.add_element_count(elements); + state.add_global_memory_reads(elements, "Size"); + state.add_global_memory_writes(elements); + + state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) { + auto env = cuda::std::execution::env{cuda::stream_ref{launch.get_stream().get_stream()}}; + if (as_transform) + { + auto it = cuda::transform_iterator(d_in, plus_one_t{}); + _CCCL_TRY_RUNTIME_API( + cub::DeviceTransform::Transform, + "Transform failed", + cuda::std::make_tuple(it), + d_out, + static_cast(elements), + cuda::std::identity{}, + env); + } + else + { + _CCCL_TRY_RUNTIME_API( + cub::DeviceTransform::Transform, + "Transform failed", + cuda::std::make_tuple(d_in), + d_out, + static_cast(elements), + plus_one_t{}, + env); + } + }); +} +catch (const std::bad_alloc&) +{ + state.skip("Skipping: out of memory."); +} + +using value_types = nvbench::type_list; + +NVBENCH_BENCH_TYPES(copy, NVBENCH_TYPE_AXES(value_types)) + .set_name("contiguous_adaptor") + .set_type_axes_names({"T{ct}"}) + .add_string_axis("Source", {"pointer", "transform"}) + .add_int64_power_of_two_axis("Elements{io}", nvbench::range(24, 28, 4)); diff --git a/cub/cub/agent/agent_find.cuh b/cub/cub/agent/agent_find.cuh index 36311311767d..cbecacca1d76 100644 --- a/cub/cub/agent/agent_find.cuh +++ b/cub/cub/agent/agent_find.cuh @@ -4,17 +4,15 @@ #pragma once #include -#include +#include #include -#include #include #include #include #include -#include -#include +#include #include #if !_CCCL_HAS_NV_ATOMIC_BUILTINS() @@ -33,18 +31,12 @@ template struct agent_t { - // The input value type - using InputT = typename ::cuda::std::iterator_traits::value_type; - - // Vector type of InputT for data movement - using VectorT = typename CubVector::Type; - static constexpr int tile_size = ThreadsPerBlock * ItemsPerThread; - // Can vectorize according to the policy if the input iterator is a native pointer to a primitive type + // Only contiguous inputs use the arrangement of vectors, because item-wise loads of other inputs in that arrangement + // are slower than coalesced ones static constexpr bool attempt_vectorization = - (VecSize > 1) && (ItemsPerThread % VecSize == 0) && (THRUST_NS_QUALIFIER::is_contiguous_iterator_v) - && ::cuda::is_trivially_copyable_v; + (VecSize > 1) && (ItemsPerThread % VecSize == 0) && THRUST_NS_QUALIFIER::is_contiguous_iterator_v; static constexpr CacheLoadModifier load_modifier = LoadModifier; @@ -64,47 +56,19 @@ struct agent_t OffsetT* found_pos_ptr; OffsetT num_items; - template - _CCCL_DEVICE _CCCL_FORCEINLINE bool is_aligned_and_full_tile(OffsetT tile_offset) - { - if constexpr (CanVectorize) - { - static_assert(::cuda::std::is_pointer_v); - - // Retrieve the value type from the iterator to determine the vector type - using InputT = typename ::cuda::std::iterator_traits::value_type; - using VectorT = typename CubVector::Type; - - const bool full_tile = (tile_offset + tile_size) <= num_items; - - // Check alignment at the actual load position (d_in + tile_offset) - return full_tile && ::cuda::is_aligned(d_in + tile_offset, sizeof(VectorT)); - } - else - { - return false; - } - } - _CCCL_DEVICE _CCCL_FORCEINLINE bool ConsumeTile(OffsetT tile_offset, ::cuda::std::true_type /*CAN_VECTORIZE*/) { - using InputT = typename ::cuda::std::iterator_traits::value_type; - using VectorT = typename CubVector::Type; - - // vectorized loads begin - auto load_ptr = reinterpret_cast(d_in + tile_offset + (threadIdx.x * VecSize)); - const CacheModifiedInputIterator d_vec_in(load_ptr); - - alignas(VectorT) unsigned char input_bytes[ItemsPerThread * sizeof(InputT)]; - auto* vec_items = reinterpret_cast(input_bytes); - - constexpr int number_of_vectors = ItemsPerThread / VecSize; - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < number_of_vectors; ++i) + // Only full tiles are loaded with vectors. Checking this here instead of in Process() avoids a duplicated + // item-wise loop when vectorization is not attempted + if (tile_offset + tile_size > num_items) { - vec_items[i] = d_vec_in[ThreadsPerBlock * i]; + return ConsumeTile(tile_offset, ::cuda::std::false_type{}); } + bool predicates[ItemsPerThread]; + LoadDirectStripedVectorized( + static_cast(threadIdx.x), ::cuda::transform_iterator{d_in + tile_offset, predicate}, predicates); + for (int i = 0; i < ItemsPerThread; ++i) { OffsetT nth_vector_of_thread = i / VecSize; @@ -113,8 +77,7 @@ struct agent_t OffsetT index = tile_offset + vector_of_tile * VecSize + element_in_vector; - auto* input_items = reinterpret_cast(input_bytes); - if (index < num_items && predicate(input_items[i])) + if (index < num_items && predicates[i]) { atomicMin(&temp_storage.block_result, index); // every thread goes over multiple elements per thread for every tile. If a thread finds a local minimum it @@ -179,10 +142,7 @@ struct agent_t return; } - const bool found_thread = - is_aligned_and_full_tile(tile_offset) - ? ConsumeTile(tile_offset, ::cuda::std::bool_constant{}) - : ConsumeTile(tile_offset, ::cuda::std::false_type{}); + const bool found_thread = ConsumeTile(tile_offset, ::cuda::std::bool_constant{}); const bool found_block = __syncthreads_or(found_thread); if (found_block) diff --git a/cub/cub/agent/agent_histogram.cuh b/cub/cub/agent/agent_histogram.cuh index d4e0f83b2785..549c5b6c3d82 100644 --- a/cub/cub/agent/agent_histogram.cuh +++ b/cub/cub/agent/agent_histogram.cuh @@ -20,7 +20,6 @@ #include #include -#include #include #include #include @@ -144,18 +143,22 @@ using AgentHistogramPolicy namespace detail::histogram { -// Return a native pixel pointer (specialized for CacheModifiedInputIterator types) -template -_CCCL_DEVICE _CCCL_FORCEINLINE auto NativePointer(CacheModifiedInputIterator itr) +// The vectorized counterparts of the direct and striped arrangements check the alignment of every tile and load +// misaligned tiles item by item. Multiple channels are loaded in a blocked arrangement instead of a striped one, +// because all samples of a pixel must belong to the same thread +[[nodiscard]] _CCCL_HOST_DEVICE_API _CCCL_CONSTEVAL BlockLoadAlgorithm +vectorized_load_algorithm(BlockLoadAlgorithm algorithm, int num_channels) { - return itr.ptr; -} - -// Return a native pixel pointer (specialized for other types) -template -_CCCL_DEVICE _CCCL_FORCEINLINE auto NativePointer(IteratorT itr) -{ - return nullptr; + switch (algorithm) + { + case BLOCK_LOAD_DIRECT: + return BLOCK_LOAD_VECTORIZE; + case BLOCK_LOAD_STRIPED: + case BLOCK_LOAD_STRIPED_VECTORIZED: + return num_channels == 1 ? BLOCK_LOAD_STRIPED_VECTORIZED : BLOCK_LOAD_VECTORIZE; + default: + return algorithm; + } } //! @brief AgentHistogram implements a stateful abstraction of CUDA thread blocks for participating @@ -201,11 +204,9 @@ template struct AgentHistogram { - static constexpr int vec_size = AgentHistogramPolicyT::VEC_SIZE; static constexpr int threads_per_block = AgentHistogramPolicyT::BLOCK_THREADS; static constexpr int pixels_per_thread = AgentHistogramPolicyT::PIXELS_PER_THREAD; static constexpr int samples_per_thread = pixels_per_thread * NumChannels; - static constexpr int vecs_per_thread = samples_per_thread / vec_size; static constexpr int tile_pixels = pixels_per_thread * threads_per_block; static constexpr int tile_samples = samples_per_thread * threads_per_block; static constexpr bool is_rle_compress = AgentHistogramPolicyT::IS_RLE_COMPRESS; @@ -213,10 +214,10 @@ struct AgentHistogram static constexpr CacheLoadModifier load_modifier = AgentHistogramPolicyT::LOAD_MODIFIER; static constexpr auto mem_preference = (PrivatizedSmemBins > 0) ? BlockHistogramMemoryPreference{AgentHistogramPolicyT::MEM_PREFERENCE} : GMEM; + static constexpr BlockLoadAlgorithm load_algorithm = + vectorized_load_algorithm(AgentHistogramPolicyT::LOAD_ALGORITHM, NumChannels); using SampleT = it_value_t; - using PixelT = typename CubVector::Type; - using VecT = typename CubVector::Type; /// Input iterator wrapper type (for applying cache modifier) // Wrap the native input pointer with CacheModifiedInputIterator or directly use the supplied input iterator type @@ -225,33 +226,20 @@ struct AgentHistogram ::cuda::std::_If<::cuda::std::is_pointer_v, CacheModifiedInputIterator, SampleIteratorT>; - using WrappedPixelIteratorT = CacheModifiedInputIterator; - using WrappedVecsIteratorT = CacheModifiedInputIterator; - using BlockLoadSampleT = - BlockLoad; - using BlockLoadPixelT = - BlockLoad; - using BlockLoadVecT = BlockLoad; + using BlockLoadSampleT = BlockLoad; struct _TempStorage { // Smem needed for block-privatized smem histogram (with 1 word of padding) CounterT histograms[NumActiveChannels][PrivatizedSmemBins + 1]; int tile_idx; - - union - { - typename BlockLoadSampleT::TempStorage sample_load; - typename BlockLoadPixelT::TempStorage pixel_load; - typename BlockLoadVecT::TempStorage vec_load; - }; + typename BlockLoadSampleT::TempStorage sample_load; }; using TempStorage = Uninitialized<_TempStorage>; _TempStorage& temp_storage; WrappedSampleIteratorT d_wrapped_samples; // with cache modifier applied, if possible - SampleT* d_native_samples; // possibly nullptr if unavailable const int* num_output_bins; // one for each channel const int* num_privatized_bins; // one for each channel CounterT* d_privatized_histograms[NumActiveChannels]; // one for each channel @@ -382,64 +370,20 @@ struct AgentHistogram } } - // Load full, aligned tile using pixel iterator - _CCCL_DEVICE _CCCL_FORCEINLINE void - LoadFullAlignedTile(OffsetT block_offset, SampleT (&samples)[pixels_per_thread][NumChannels]) - { - // The vectorized load can only load whole tiles if the samples per thread are a multiple of the vector size - if constexpr (NumActiveChannels == 1 && samples_per_thread % vec_size == 0) - { - using AliasedVecs = VecT[vecs_per_thread]; - const WrappedVecsIteratorT d_wrapped_vecs(reinterpret_cast(d_native_samples + block_offset)); - // Load using a wrapped vec iterator - BlockLoadVecT{temp_storage.vec_load}.Load(d_wrapped_vecs, reinterpret_cast(samples)); - } - else - { - using AliasedPixels = PixelT[pixels_per_thread]; - const WrappedPixelIteratorT d_wrapped_pixels(reinterpret_cast(d_native_samples + block_offset)); - // Load using a wrapped pixel iterator - BlockLoadPixelT{temp_storage.pixel_load}.Load(d_wrapped_pixels, reinterpret_cast(samples)); - } - } - - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void LoadTile(OffsetT block_offset, int valid_samples, SampleT (&samples)[pixels_per_thread][NumChannels]) { + using AliasedSamples = SampleT[samples_per_thread]; if constexpr (IsFullTile) { - if constexpr (IsAligned) - { - LoadFullAlignedTile(block_offset, samples); - } - else - { - // Load using sample iterator - using AliasedSamples = SampleT[samples_per_thread]; - BlockLoadSampleT{temp_storage.sample_load}.Load( - d_wrapped_samples + block_offset, reinterpret_cast(samples)); - } + BlockLoadSampleT{temp_storage.sample_load}.Load( + d_wrapped_samples + block_offset, reinterpret_cast(samples)); } else { - if constexpr (IsAligned) - { - // Load partially-full, aligned tile using the pixel iterator - using AliasedPixels = PixelT[pixels_per_thread]; - const WrappedPixelIteratorT d_wrapped_pixels((PixelT*) (d_native_samples + block_offset)); - const int valid_pixels = valid_samples / NumChannels; - - // Load using a wrapped pixel iterator - BlockLoadPixelT{temp_storage.pixel_load}.Load( - d_wrapped_pixels, reinterpret_cast(samples), valid_pixels); - } - else - { - using AliasedSamples = SampleT[samples_per_thread]; - BlockLoadSampleT{temp_storage.sample_load}.Load( - d_wrapped_samples + block_offset, reinterpret_cast(samples), valid_samples); - } + BlockLoadSampleT{temp_storage.sample_load}.Load( + d_wrapped_samples + block_offset, reinterpret_cast(samples), valid_samples); } } @@ -462,19 +406,17 @@ struct AgentHistogram //! @brief Consume a tile of data samples //! - //! @tparam IsAligned - //! Whether the tile offset is aligned (vec-aligned for single-channel, pixel-aligned for multi-channel) - //! //! @tparam IsFullTile //! Whether the tile is full - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumeTile(OffsetT block_offset, int valid_samples) { SampleT samples[pixels_per_thread][NumChannels]; bool is_valid[pixels_per_thread]; - LoadTile(block_offset, valid_samples, samples); - MarkValid(is_valid, valid_samples); + LoadTile(block_offset, valid_samples, samples); + // Partial tiles of the striped vectorized algorithm are loaded in the striped arrangement of single samples + MarkValid(is_valid, valid_samples); if (prefer_smem) { @@ -499,7 +441,6 @@ struct AgentHistogram //! //! @param tiles_per_row //! Number of image tiles per row - template _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumeTiles( OffsetT num_row_pixels, OffsetT num_rows, @@ -524,12 +465,12 @@ struct AgentHistogram { // Consume a partially-full tile at the end of the row OffsetT num_remaining = (num_row_pixels * NumChannels) - col_offset; - ConsumeTile(tile_offset, num_remaining); + ConsumeTile(tile_offset, num_remaining); } else { // Consume full tile - ConsumeTile(tile_offset, tile_samples); + ConsumeTile(tile_offset, tile_samples); } __syncthreads(); @@ -550,7 +491,7 @@ struct AgentHistogram // already synchronizes between tiles. _CCCL_DEVICE _CCCL_FORCEINLINE void SyncBeforeNextTile() { - if constexpr (AgentHistogramPolicyT::LOAD_ALGORITHM == BLOCK_LOAD_TRANSPOSE) + if constexpr (load_algorithm == BLOCK_LOAD_TRANSPOSE) { __syncthreads(); } @@ -566,7 +507,6 @@ struct AgentHistogram //! //! @param row_stride_samples //! The number of samples between starts of consecutive rows in the region of interest - template _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumeTiles( OffsetT num_row_pixels, OffsetT num_rows, OffsetT row_stride_samples, int, GridQueue, ::cuda::std::false_type) { @@ -583,13 +523,13 @@ struct AgentHistogram if (num_remaining < tile_samples) { // Consume partial tile - ConsumeTile(tile_offset, num_remaining); + ConsumeTile(tile_offset, num_remaining); SyncBeforeNextTile(); break; } // Consume full tile - ConsumeTile(tile_offset, tile_samples); + ConsumeTile(tile_offset, tile_samples); SyncBeforeNextTile(); tile_offset += gridDim.x * tile_samples; } @@ -636,7 +576,6 @@ struct AgentHistogram const PrivatizedDecodeOpT* privatized_decode_op) : temp_storage(temp_storage.Alias()) , d_wrapped_samples(d_samples) - , d_native_samples(NativePointer(d_wrapped_samples)) , num_output_bins(num_output_bins) , num_privatized_bins(num_privatized_bins) , d_output_histograms(d_output_histograms) @@ -677,36 +616,9 @@ struct AgentHistogram _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumeTiles( OffsetT num_row_pixels, OffsetT num_rows, OffsetT row_stride_samples, int tiles_per_row, GridQueue tile_queue) { - // Check whether all row starting offsets are vec-aligned (in single-channel) or pixel-aligned (in multi-channel) - constexpr int vec_mask = alignof(VecT) - 1; - constexpr int pixel_mask = alignof(PixelT) - 1; - const size_t row_bytes = sizeof(SampleT) * row_stride_samples; - - const bool vec_aligned_rows = - (NumChannels == 1) && (samples_per_thread % vec_size == 0) && // Single channel - ((size_t(d_native_samples) & vec_mask) == 0) && // ptr is quad-aligned - ((num_rows == 1) || ((row_bytes & vec_mask) == 0)); // number of row-samples is a multiple of the alignment of the - // quad - - const bool pixel_aligned_rows = - (NumChannels > 1) && // Multi channel - ((size_t(d_native_samples) & pixel_mask) == 0) && // ptr is pixel-aligned - ((row_bytes & pixel_mask) == 0); // number of row-samples is a multiple of the alignment of the pixel - _CCCL_PDL_GRID_DEPENDENCY_SYNC(); - - // Whether rows are aligned and can be vectorized - if ((d_native_samples != nullptr) && (vec_aligned_rows || pixel_aligned_rows)) - { - ConsumeTiles( - num_row_pixels, num_rows, row_stride_samples, tiles_per_row, tile_queue, bool_constant_v); - } - else - { - ConsumeTiles( - num_row_pixels, num_rows, row_stride_samples, tiles_per_row, tile_queue, bool_constant_v); - } - + ConsumeTiles( + num_row_pixels, num_rows, row_stride_samples, tiles_per_row, tile_queue, bool_constant_v); _CCCL_PDL_TRIGGER_NEXT_LAUNCH(); // omitting makes no difference in cub.bench.histogram.even.base } diff --git a/cub/cub/agent/agent_reduce.cuh b/cub/cub/agent/agent_reduce.cuh index 3cb7e80454c5..316a58356a1e 100644 --- a/cub/cub/agent/agent_reduce.cuh +++ b/cub/cub/agent/agent_reduce.cuh @@ -29,12 +29,15 @@ #include #include +#include +#include #include #include #include #include #include #include +#include #include #include @@ -202,8 +205,9 @@ struct AgentReduceImpl /// The input value type using InputT = it_value_t; - /// Vector type of InputT for data movement - using VectorT = typename CubVector::Type; + /// The value type of the storage behind the input iterator. Differs from InputT for an iterator adaptor over + /// contiguous storage, like cuda::transform_iterator over a pointer + using StorageT = ::cuda::__unwrapped_iter_value_t; /// Input iterator wrapper type (for applying cache modifier) // Wrap the native input pointer with CacheModifiedInputIterator @@ -218,17 +222,23 @@ struct AgentReduceImpl static constexpr int TILE_ITEMS = NumThreads * ITEMS_PER_THREAD; static constexpr int vec_size = ::cuda::std::min(ITEMS_PER_THREAD, AgentReducePolicy::VECTOR_LOAD_LENGTH); - // Can vectorize according to the policy if the input iterator is a native - // pointer to a primitive type + /// Vector type of StorageT for data movement. Cache modified loads are at most 16 bytes wide, so they gain nothing + /// from a stricter alignment + using VectorT = ::cuda::__as_bytes_vectorized_t; + + // Can vectorize according to the policy if the input iterator is a native pointer to a primitive type, or an iterator + // adaptor over one // TODO(bgruber): we should not check for `is_pointer_v` but `contiguous_iterator` and unwrap it static constexpr bool ATTEMPT_VECTORIZATION = - (vec_size > 1) && (ITEMS_PER_THREAD % vec_size == 0) - && (::cuda::std::is_pointer_v) - // TODO(bgruber): remove the check for is_primitive in CCCL 4.0 - &&(is_primitive::value || ::cuda::is_trivially_copyable_v) + (vec_size > 1) && (ITEMS_PER_THREAD % vec_size == 0) && !::cuda::std::is_volatile_v + && (::cuda::std::is_pointer_v || ::cuda::__iterator_can_rebase) + // The loaded items are stored in a local array of InputT + &&::cuda::std::is_default_constructible_v + // TODO(bgruber): remove the check for is_primitive in CCCL 4.0 + && (is_primitive::value || ::cuda::is_trivially_copyable_v) // vectorizing large types leads to regressions again, see https://github.com/NVIDIA/cccl/issues/9761 // TODO(bgruber): this should be decided by tuning - &&sizeof(InputT) + &&sizeof(StorageT) <= 8; static constexpr CacheLoadModifier LOAD_MODIFIER = AgentReducePolicy::LOAD_MODIFIER; @@ -263,7 +273,7 @@ struct AgentReduceImpl { if constexpr (AttemptVectorization) { - return ::cuda::std::is_sufficiently_aligned(d_in); + return ::cuda::std::is_sufficiently_aligned(::cuda::__unwrap_iterator_adaptor(d_in)); } else { @@ -304,45 +314,29 @@ struct AgentReduceImpl template _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumeFullTile(AccumT& thread_aggregate, OffsetT block_offset) { + // load items in striped fashion and reduce items within each thread stripe + AccumT items[ITEMS_PER_THREAD]; if constexpr (CanVectorize) { - // Fabricate a vectorized input iterator - InputT* d_in_unqualified = const_cast(d_in) + block_offset + (lane_id * vec_size); - const CacheModifiedInputIterator d_vec_in( - reinterpret_cast(d_in_unqualified)); - - // Load items as vector items - alignas(VectorT) InputT input_items[ITEMS_PER_THREAD]; - VectorT* vec_items = reinterpret_cast(input_items); - - // Alias items as an array of VectorT and load it in striped fashion - static constexpr int words = ITEMS_PER_THREAD / vec_size; - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < words; ++i) - { - vec_items[i] = d_vec_in[NumThreads * i]; - } + InputT input_items[ITEMS_PER_THREAD]; + // lane_id is the warp lane for a warp agent and the block thread id otherwise. + InternalLoadDirectStripedVectorized( + static_cast(lane_id), d_in + block_offset, input_items); // Convert from input type to output type - AccumT items[ITEMS_PER_THREAD]; _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < ITEMS_PER_THREAD; ++i) { items[i] = transform_op(input_items[i]); } - - // Reduce items within each thread stripe - thread_aggregate = - IsFirstTile ? cub::ThreadReduce(items, reduction_op) : cub::ThreadReduce(items, reduction_op, thread_aggregate); } else { - // Scalar path: load items in striped fashion and reduce items within each thread stripe - AccumT items[ITEMS_PER_THREAD]; load_transform_direct_striped(lane_id, d_wrapped_in + block_offset, items, transform_op); - thread_aggregate = - IsFirstTile ? cub::ThreadReduce(items, reduction_op) : cub::ThreadReduce(items, reduction_op, thread_aggregate); } + // Reduce items within each thread stripe + thread_aggregate = + IsFirstTile ? cub::ThreadReduce(items, reduction_op) : cub::ThreadReduce(items, reduction_op, thread_aggregate); } /** diff --git a/cub/cub/block/block_load.cuh b/cub/cub/block/block_load.cuh index e42653a89ac5..adbe2a498cb7 100644 --- a/cub/cub/block/block_load.cuh +++ b/cub/cub/block/block_load.cuh @@ -28,14 +28,17 @@ #include #include +#include +#include +#include #include #include #include #include +#include #include -#include -#include -#include + +#include CUB_NAMESPACE_BEGIN @@ -182,59 +185,55 @@ _CCCL_DEVICE _CCCL_FORCEINLINE void LoadDirectBlocked( #ifndef _CCCL_DOXYGEN_INVOKED // Do not document -//! @brief Internal implementation for load vectorization +//! @brief Internal implementation of LoadDirectBlockedVectorized for input aligned to the vector type //! //! @param[in] linear_tid //! A suitable 1D thread-identifier for the calling thread (e.g., `(threadIdx.y * blockDim.x) + linear_tid` for 2D //! thread blocks) //! -//! @param[in] block_src_ptr -//! Input pointer for loading from +//! @param[in] block_src_it +//! Input iterator for loading from. Either a contiguous iterator or an iterator adaptor that can be rebased, whose +//! storage is aligned to ``cuda::__as_bytes_vectorized_t`` of ``ItemsPerThread`` storage items with vectors of at +//! most ``MaxVectorAccess`` bytes //! //! @param[out] dst_items //! Destination to load data into -template +template _CCCL_DEVICE _CCCL_FORCEINLINE void -InternalLoadDirectBlockedVectorized(int linear_tid, const T* block_src_ptr, T (&dst_items)[ItemsPerThread]) +InternalLoadDirectBlockedVectorized(int linear_tid, RandomAccessIterator block_src_it, T (&dst_items)[ItemsPerThread]) { - // Find biggest memory access word that T is a whole multiple of - using device_word_t = typename UnitWord::DeviceWord; - _CCCL_DIAG_PUSH -# if _CCCL_COMPILER(CLANG, >=, 10) - _CCCL_DIAG_SUPPRESS_CLANG("-Wsizeof-array-div") -# endif // _CCCL_COMPILER(CLANG, >=, 10) - // NOLINTNEXTLINE(bugprone-sizeof-expression) - constexpr int total_words = static_cast(sizeof(dst_items) / sizeof(device_word_t)); - _CCCL_DIAG_POP - constexpr int vector_size = (total_words % 4 == 0) ? 4 : (total_words % 2 == 0) ? 2 : 1; - constexpr int vectors_per_thread = total_words / vector_size; - - // Load into an array of vectors in thread-blocked order - using vector_t = typename CubVector::Type; - - // Add the alignment check to ensure the vectorized loading can proceed. - if (::cuda::std::is_sufficiently_aligned(block_src_ptr)) - { - vector_t vec_items[vectors_per_thread]; - // Load into an array of vectors in thread-blocked order - const vector_t* vec_ptr = reinterpret_cast(block_src_ptr) + linear_tid * vectors_per_thread; + // The storage of all items of a thread is loaded with the widest vectors that evenly divide it + using storage_t = ::cuda::__unwrapped_iter_value_t; + using vector_t = ::cuda::__as_bytes_vectorized_t; + const storage_t* block_src_ptr = + ::cuda::std::__runtime_assume_aligned(::cuda::__unwrap_iterator_adaptor(block_src_it), alignof(vector_t)); - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < vectors_per_thread; i++) - { - vec_items[i] = ThreadLoad(vec_ptr + i); - } + // Load in thread-blocked order + vector_t vec_items = ThreadLoad(reinterpret_cast(block_src_ptr) + linear_tid); + if constexpr (THRUST_NS_QUALIFIER::is_contiguous_iterator_v) + { // Copy to destination _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < ItemsPerThread; i++) { - dst_items[i] = *(reinterpret_cast(vec_items) + i); + dst_items[i] = vec_items.template __alias(i); } } else { - LoadDirectBlocked(linear_tid, block_src_ptr, dst_items); + // Copy to destination through the iterator rebased onto the loaded storage + ::cuda::__rebase_contiguous_iterator_adaptor( + block_src_it, ::cuda::std::addressof(vec_items.template __alias(0))); + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + dst_items[i] = block_src_it[i]; + } } } @@ -248,36 +247,111 @@ InternalLoadDirectBlockedVectorized(int linear_tid, const T* block_src_ptr, T (& //! //! @blocked //! -//! The input offset (``block_ptr + block_offset``) must be quad-item aligned -//! //! The following conditions will prevent vectorization and loading will fall back to cub::BLOCK_LOAD_DIRECT: //! -//! - ``ItemsPerThread`` is odd -//! - The data type ``T`` is not a built-in primitive or CUDA vector type -//! (e.g., ``short``, ``int2``, ``double``, ``float2``, etc.) +//! - The ``RandomAccessIterator`` is neither contiguous nor an iterator adaptor over contiguous storage, like +//! ``cuda::transform_iterator`` over a pointer +//! - ``block_src_it`` is not aligned to the largest power of two that divides ``sizeof(T) * VecSize`` +//! (``sizeof(T) * ItemsPerThread`` if ``VecSize`` is ``0``), up to 16 bytes +//! - The data type ``T`` is not trivially copyable according to ``cuda::is_trivially_copyable_v`` +//! - The data type ``T`` is aligned to more than 32 bytes (16 bytes before CUDA 13) +//! +//! With CUDA 13, sm_100 and newer load ``block_src_it`` with vectors of up to 32 bytes if it is aligned to 32 bytes and +//! ``Modifier`` is ``LOAD_DEFAULT``. +//! +//! The storage of ``block_src_it`` is loaded with vectors and the items are read from it through ``block_src_it``, so +//! the conditions on ``T`` and ``block_src_it`` apply to its storage. //! //! @endrst //! +//! @tparam Modifier +//! **[optional]** Cache load modifier for the vectorized loads (default: ``LOAD_DEFAULT``) +//! +//! @tparam VecSize +//! **[optional]** The number of consecutive items in each vector. Must evenly divide ``ItemsPerThread``. Each thread +//! loads its items with ``ItemsPerThread / VecSize`` vectors, and ``0`` loads them with a single vector +//! (default: ``0``) +//! //! @tparam T //! **[inferred]** The data type to load. //! //! @tparam ItemsPerThread //! **[inferred]** The number of consecutive items partitioned onto each thread. //! +//! @tparam RandomAccessIterator +//! **[inferred]** The random-access iterator type for input iterator. +//! //! @param[in] linear_tid //! A suitable 1D thread-identifier for the calling thread (e.g., `(threadIdx.y * blockDim.x) + //! linear_tid` for 2D thread blocks) //! -//! @param[in] block_src_ptr -//! The thread block's base pointer for loading from +//! @param[in] block_src_it +//! The thread block's base iterator for loading from //! //! @param[out] dst_items //! destination to load data into -template +template _CCCL_DEVICE _CCCL_FORCEINLINE void -LoadDirectBlockedVectorized(int linear_tid, T* block_src_ptr, T (&dst_items)[ItemsPerThread]) +LoadDirectBlockedVectorized(int linear_tid, RandomAccessIterator block_src_it, T (&dst_items)[ItemsPerThread]) { - InternalLoadDirectBlockedVectorized(linear_tid, block_src_ptr, dst_items); + constexpr int vec_size = VecSize == 0 ? ItemsPerThread : VecSize; + static_assert(ItemsPerThread % vec_size == 0, "VecSize must evenly divide ItemsPerThread"); + using storage_t = ::cuda::__unwrapped_iter_value_t; + + constexpr bool can_vectorize = + ::cuda::is_trivially_copyable_v && !::cuda::std::is_volatile_v + && (THRUST_NS_QUALIFIER::is_contiguous_iterator_v + || ::cuda::__iterator_can_rebase); + + if constexpr (can_vectorize) + { + // Vector i of a thread is the blocked vector of thread linear_tid * vectors_per_thread + i + constexpr int vectors_per_thread = ItemsPerThread / vec_size; + const auto block_src_ptr = ::cuda::__unwrap_iterator_adaptor(block_src_it); + + // Only sm_100 and newer load vectors wider than 16 bytes with a single access, and only without a cache modifier + using vector16_t = ::cuda::__as_bytes_vectorized_t; + +#if _CCCL_CTK_AT_LEAST(13, 0) + using vector32_t = ::cuda::__as_bytes_vectorized_t; + if constexpr (Modifier == LOAD_DEFAULT && alignof(vector16_t) < alignof(vector32_t)) + { + bool is_vector32_aligned = false; + NV_IF_TARGET(NV_PROVIDES_SM_100, + (is_vector32_aligned = ::cuda::std::is_sufficiently_aligned(block_src_ptr);)) + if (is_vector32_aligned) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < vectors_per_thread; i++) + { + InternalLoadDirectBlockedVectorized( + linear_tid * vectors_per_thread + i, + block_src_it, + reinterpret_cast(dst_items[i * vec_size])); + } + return; + } + } +#endif // _CCCL_CTK_AT_LEAST(13, 0) + + // Add the alignment check on the storage to ensure the vectorized loading can proceed. + if (::cuda::std::is_sufficiently_aligned(block_src_ptr)) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < vectors_per_thread; i++) + { + InternalLoadDirectBlockedVectorized( + linear_tid * vectors_per_thread + i, block_src_it, reinterpret_cast(dst_items[i * vec_size])); + } + return; + } + } + + LoadDirectBlocked(linear_tid, block_src_it, dst_items); } //! @} @@ -326,6 +400,158 @@ LoadDirectStriped(int linear_tid, RandomAccessIterator block_src_it, T (&dst_ite } } +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + +//! @brief Internal implementation of LoadDirectStripedVectorized for input aligned to the vector type +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread (e.g., `(threadIdx.y * blockDim.x) + linear_tid` for 2D +//! thread blocks) +//! +//! @param[in] block_src_it +//! Input iterator for loading from. Either a contiguous iterator or an iterator adaptor that can be rebased, whose +//! storage is aligned to ``cuda::__as_bytes_vectorized_t`` of ``VecSize`` storage items with vectors of at most +//! ``MaxVectorAccess`` bytes +//! +//! @param[out] dst_items +//! Destination to load data into +template +_CCCL_DEVICE _CCCL_FORCEINLINE void +InternalLoadDirectStripedVectorized(int linear_tid, RandomAccessIterator block_src_it, T (&dst_items)[ItemsPerThread]) +{ + static_assert(ItemsPerThread % VecSize == 0, "VecSize must evenly divide ItemsPerThread"); + + // Vector i of a thread is the blocked vector of thread linear_tid + i * ThreadsPerBlock + // The thread offset is computed in unsigned arithmetic like offsets from threadIdx.x, because a signed multiplication + // generates additional address arithmetic in the calling kernels + const auto thread_src_it = block_src_it + static_cast(linear_tid) * VecSize; + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread / VecSize; i++) + { + InternalLoadDirectBlockedVectorized( + i * ThreadsPerBlock, thread_src_it, reinterpret_cast(dst_items[i * VecSize])); + } +} + +#endif // _CCCL_DOXYGEN_INVOKED + +//! @rst +//! Load a linear segment of items into a striped arrangement of vectors across the thread block. +//! +//! .. versionadded:: 3.6.0 +//! +//! Assumes a :ref:`striped arrangement ` of vectors of ``VecSize`` consecutive items across +//! the thread block, where *thread*\ :sub:`i` owns the vectors (*i*), (*i* + *block-threads*), ..., +//! (*i* + (*block-threads* * (*items-per-thread* / ``VecSize`` - 1))). With ``VecSize == 1``, this is the arrangement +//! of ``LoadDirectStriped``. +//! +//! The following conditions will prevent vectorization and the same arrangement is loaded item by item: +//! +//! - The ``RandomAccessIterator`` is neither contiguous nor an iterator adaptor over contiguous storage, like +//! ``cuda::transform_iterator`` over a pointer +//! - ``block_src_it`` is not aligned to the largest power of two that divides ``sizeof(T) * VecSize``, up to 16 bytes +//! - The data type ``T`` is not trivially copyable according to ``cuda::is_trivially_copyable_v`` +//! - The data type ``T`` is aligned to more than 32 bytes (16 bytes before CUDA 13) +//! +//! With CUDA 13, sm_100 and newer load ``block_src_it`` with vectors of up to 32 bytes if it is aligned to 32 bytes and +//! ``Modifier`` is ``LOAD_DEFAULT``. +//! +//! The storage of ``block_src_it`` is loaded with vectors and the items are read from it through ``block_src_it``, so +//! the conditions on ``T`` and ``block_src_it`` apply to its storage. +//! +//! @endrst +//! +//! @tparam ThreadsPerBlock +//! The thread block size in threads +//! +//! @tparam VecSize +//! The number of consecutive items in each vector. Must evenly divide ``ItemsPerThread``. +//! +//! @tparam Modifier +//! **[optional]** Cache load modifier for the vectorized loads (default: ``LOAD_DEFAULT``) +//! +//! @tparam RandomAccessIterator +//! **[inferred]** The random-access iterator type for input iterator. +//! +//! @tparam T +//! **[inferred]** The data type to load. +//! +//! @tparam ItemsPerThread +//! **[inferred]** The number of items partitioned onto each thread. +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread (e.g., `(threadIdx.y * blockDim.x) + linear_tid` for 2D +//! thread blocks) +//! +//! @param[in] block_src_it +//! The thread block's base iterator for loading from +//! +//! @param[out] dst_items +//! Destination to load data into +template +_CCCL_DEVICE _CCCL_FORCEINLINE void +LoadDirectStripedVectorized(int linear_tid, RandomAccessIterator block_src_it, T (&dst_items)[ItemsPerThread]) +{ + static_assert(ItemsPerThread % VecSize == 0, "VecSize must evenly divide ItemsPerThread"); + using storage_t = ::cuda::__unwrapped_iter_value_t; + + constexpr bool can_vectorize = + ::cuda::is_trivially_copyable_v && !::cuda::std::is_volatile_v + && (THRUST_NS_QUALIFIER::is_contiguous_iterator_v + || ::cuda::__iterator_can_rebase); + + if constexpr (can_vectorize) + { + const auto block_src_ptr = ::cuda::__unwrap_iterator_adaptor(block_src_it); + + // Only sm_100 and newer load vectors wider than 16 bytes with a single access, and only without a cache modifier + using vector16_t = ::cuda::__as_bytes_vectorized_t; + +#if _CCCL_CTK_AT_LEAST(13, 0) + using vector32_t = ::cuda::__as_bytes_vectorized_t; + if constexpr (Modifier == LOAD_DEFAULT && alignof(vector16_t) < alignof(vector32_t)) + { + bool is_vector32_aligned = false; + NV_IF_TARGET(NV_PROVIDES_SM_100, + (is_vector32_aligned = ::cuda::std::is_sufficiently_aligned(block_src_ptr);)) + if (is_vector32_aligned) + { + InternalLoadDirectStripedVectorized( + linear_tid, block_src_it, dst_items); + return; + } + } +#endif // _CCCL_CTK_AT_LEAST(13, 0) + + // Add the alignment check on the storage to ensure the vectorized loading can proceed. + if (::cuda::std::is_sufficiently_aligned(block_src_ptr)) + { + InternalLoadDirectStripedVectorized( + linear_tid, block_src_it, dst_items); + return; + } + } + + // Load the same arrangement item by item + const auto thread_src_it = block_src_it + static_cast(linear_tid) * VecSize; + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread / VecSize; i++) + { + LoadDirectBlocked(i * ThreadsPerBlock, thread_src_it, reinterpret_cast(dst_items[i * VecSize])); + } +} + namespace detail { template @@ -663,12 +889,11 @@ enum BlockLoadAlgorithm // NOLINT(cppcoreguidelines-use-enum-class) //! is lower). //! - The following conditions will prevent vectorization and loading will fall back to cub::BLOCK_LOAD_DIRECT: //! - //! - ``ItemsPerThread`` is odd - //! - The ``RandomAccessIterator`` is not a contiguous iterator (``cub::CacheModifiedInputIterator`` is an exception - //! and is still vectorized) - //! - The block input offset is not aligned to the vector access width - //! - The data type ``T`` is not a built-in primitive or CUDA vector type - //! (e.g., ``short``, ``int2``, ``double``, ``float2``, etc.) + //! - The ``RandomAccessIterator`` is neither contiguous nor an iterator adaptor over contiguous storage, like + //! ``cuda::transform_iterator`` over a pointer + //! - The block input offset is not aligned to the vector, up to 16 bytes + //! - The data type ``T`` is not trivially copyable according to ``cuda::is_trivially_copyable_v`` + //! - The data type ``T`` is aligned to more than 32 bytes (16 bytes before CUDA 13) //! //! @endrst BLOCK_LOAD_VECTORIZE, @@ -737,6 +962,34 @@ enum BlockLoadAlgorithm // NOLINT(cppcoreguidelines-use-enum-class) //! //! @endrst BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED, + + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! A :ref:`striped arrangement ` of vectors is read from memory using CUDA's built-in + //! vectorized loads. Each vector holds consecutive items, and thread ``t`` owns the vectors at positions ``t``, + //! ``t + ThreadsPerBlock``, and so on. The vectors are as wide as the largest power of two that divides + //! ``sizeof(T) * ItemsPerThread``, up to 32 bytes (16 bytes before CUDA 13). If that does not hold a whole number of + //! items, each vector is a single item, which is the arrangement of cub::BLOCK_LOAD_STRIPED. + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! - Consecutive threads load consecutive vectors, so the accesses stay coalesced regardless of the number of items + //! per thread. + //! - The following conditions will prevent vectorization and the same arrangement is loaded item by item: + //! + //! - The ``RandomAccessIterator`` is neither contiguous nor an iterator adaptor over contiguous storage, like + //! ``cuda::transform_iterator`` over a pointer + //! - The block input offset is not aligned to the vector, up to 16 bytes + //! - The data type ``T`` is not trivially copyable according to ``cuda::is_trivially_copyable_v`` + //! - The data type ``T`` is aligned to more than 32 bytes (16 bytes before CUDA 13) + //! + //! - Guarded loads use the arrangement of cub::BLOCK_LOAD_STRIPED. + //! + //! @endrst + BLOCK_LOAD_STRIPED_VECTORIZED, }; #if _CCCL_HOSTED() && !defined(_CCCL_DOXYGEN_INVOKED) @@ -758,6 +1011,8 @@ namespace detail return "BLOCK_LOAD_WARP_TRANSPOSE"; case BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED: return "BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED"; + case BLOCK_LOAD_STRIPED_VECTORIZED: + return "BLOCK_LOAD_STRIPED_VECTORIZED"; } return ""; } @@ -805,6 +1060,9 @@ CUB_NAMESPACE_BEGIN //! #. :cpp:enumerator:`cub::BLOCK_LOAD_VECTORIZE`: //! A :ref:`blocked arrangement ` of data is read directly from memory //! using CUDA's built-in vectorized loads as a coalescing optimization. +//! #. :cpp:enumerator:`cub::BLOCK_LOAD_STRIPED_VECTORIZED`: +//! A :ref:`striped arrangement ` of vectors is read directly from memory +//! using CUDA's built-in vectorized loads as a coalescing optimization. //! #. :cpp:enumerator:`cub::BLOCK_LOAD_TRANSPOSE`: //! A :ref:`striped arrangement ` of data is read directly from memory and is then //! locally transposed into a :ref:`blocked arrangement `. @@ -899,8 +1157,8 @@ class BlockLoad _CCCL_HOST_DEVICE_API static constexpr auto temp_storage_helper() { - if constexpr (Algorithm == BLOCK_LOAD_DIRECT || Algorithm == BLOCK_LOAD_STRIPED - || Algorithm == BLOCK_LOAD_VECTORIZE) + if constexpr (Algorithm == BLOCK_LOAD_DIRECT || Algorithm == BLOCK_LOAD_STRIPED || Algorithm == BLOCK_LOAD_VECTORIZE + || Algorithm == BLOCK_LOAD_STRIPED_VECTORIZED) { return NullType{}; } @@ -1015,22 +1273,34 @@ public: } else if constexpr (Algorithm == BLOCK_LOAD_VECTORIZE) { - if constexpr (detail::is_CacheModifiedInputIterator - && ::cuda::std::is_same_v, T>) + if constexpr (detail::is_CacheModifiedInputIterator) { - InternalLoadDirectBlockedVectorized(linear_tid, block_src_it.ptr, dst_items); + LoadDirectBlockedVectorized(linear_tid, block_src_it.ptr, dst_items); } - else if constexpr (THRUST_NS_QUALIFIER::is_contiguous_iterator_v - && ::cuda::std::is_same_v, T> - // handles contiguous iterators wrapping volatile pointers, like span::iterator - && !::cuda::std::is_volatile_v< - ::cuda::std::remove_reference_t>>) + else if constexpr (THRUST_NS_QUALIFIER::is_contiguous_iterator_v) { - InternalLoadDirectBlockedVectorized(linear_tid, ::cuda::std::to_address(block_src_it), dst_items); + LoadDirectBlockedVectorized(linear_tid, ::cuda::std::to_address(block_src_it), dst_items); } else { - LoadDirectBlocked(linear_tid, block_src_it, dst_items); + LoadDirectBlockedVectorized(linear_tid, block_src_it, dst_items); + } + } + else if constexpr (Algorithm == BLOCK_LOAD_STRIPED_VECTORIZED) + { + // Vectors of the widest vector access that evenly divides the items of a thread, or single items if that does not + // hold a whole number of T + using vector_word_t = typename ::cuda::__as_bytes_vectorized_t::value_type; + constexpr int vec_size = + sizeof(vector_word_t) % sizeof(T) == 0 ? static_cast(sizeof(vector_word_t) / sizeof(T)) : 1; + if constexpr (detail::is_CacheModifiedInputIterator) + { + LoadDirectStripedVectorized( + linear_tid, block_src_it.ptr, dst_items); + } + else + { + LoadDirectStripedVectorized(linear_tid, block_src_it, dst_items); } } else if constexpr (Algorithm == BLOCK_LOAD_TRANSPOSE) @@ -1102,7 +1372,7 @@ public: { LoadDirectBlocked(linear_tid, block_src_it, dst_items, block_items_end); } - else if constexpr (Algorithm == BLOCK_LOAD_STRIPED) + else if constexpr (Algorithm == BLOCK_LOAD_STRIPED || Algorithm == BLOCK_LOAD_STRIPED_VECTORIZED) { LoadDirectStriped(linear_tid, block_src_it, dst_items, block_items_end); } @@ -1178,7 +1448,7 @@ public: { LoadDirectBlocked(linear_tid, block_src_it, dst_items, block_items_end, oob_default); } - else if constexpr (Algorithm == BLOCK_LOAD_STRIPED) + else if constexpr (Algorithm == BLOCK_LOAD_STRIPED || Algorithm == BLOCK_LOAD_STRIPED_VECTORIZED) { LoadDirectStriped(linear_tid, block_src_it, dst_items, block_items_end, oob_default); } diff --git a/cub/cub/block/block_load_to_shared.cuh b/cub/cub/block/block_load_to_shared.cuh index e1464df7edc2..d96d582517cb 100644 --- a/cub/cub/block/block_load_to_shared.cuh +++ b/cub/cub/block/block_load_to_shared.cuh @@ -25,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -81,10 +82,7 @@ private: static constexpr int threads_per_block = BlockDimX * BlockDimY * BlockDimZ; // Helper for fallback to gmem->reg->smem - struct alignas(detail::bulk_copy_min_align) vec_load_t - { - char c_array[detail::bulk_copy_min_align]; - }; + using vec_load_t = ::cuda::__as_bytes_vectorized_t; struct _TempStorage { diff --git a/cub/cub/device/device_for.cuh b/cub/cub/device/device_for.cuh index 24a7e5b9d260..d9304a6130bd 100644 --- a/cub/cub/device/device_for.cuh +++ b/cub/cub/device/device_for.cuh @@ -32,6 +32,7 @@ #include #include #include +#include #include #include #include @@ -103,8 +104,9 @@ struct DeviceFor // TODO Can be extracted into tuning constexpr static int vec_size = 4; - // Type of the vector that is used to load the input data - using vector_t = typename CubVector::Type; + // Type of the vector that is used to load the input data, with vectors of at most 16 bytes, so that input aligned + // to 16 bytes is vectorized + using vector_t = ::cuda::__as_bytes_vectorized_t; _CCCL_DEVICE _CCCL_FORCEINLINE void operator()(OffsetT i) { @@ -124,7 +126,7 @@ struct DeviceFor _CCCL_PRAGMA_UNROLL_FULL() for (int j = 0; j < vec_size; j++) { - (void) op(*(reinterpret_cast(&vec) + j)); + (void) op(vec.template __alias(j)); } } } diff --git a/cub/cub/device/dispatch/dispatch_transform.cuh b/cub/cub/device/dispatch/dispatch_transform.cuh index 81210bfd8f6e..882d364faedc 100644 --- a/cub/cub/device/dispatch/dispatch_transform.cuh +++ b/cub/cub/device/dispatch/dispatch_transform.cuh @@ -153,12 +153,13 @@ struct TransformKernelSource - CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static auto is_pointer_aligned(T it, [[maybe_unused]] int alignment) + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static auto is_pointer_aligned(It it, [[maybe_unused]] int alignment) { - if constexpr (THRUST_NS_QUALIFIER::is_contiguous_iterator_v) + if constexpr (::cuda::__is_contiguous_iterator_or_adaptor_v + || THRUST_NS_QUALIFIER::is_contiguous_iterator_v) { - return ::cuda::is_aligned(THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(it), alignment); + return ::cuda::is_aligned(::cuda::__unwrap_iterator_adaptor(it), alignment); } else { diff --git a/cub/cub/device/dispatch/kernels/kernel_transform.cuh b/cub/cub/device/dispatch/kernels/kernel_transform.cuh index e59f4e8e6829..b638ef6916f0 100644 --- a/cub/cub/device/dispatch/kernels/kernel_transform.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_transform.cuh @@ -13,6 +13,7 @@ # pragma system_header #endif // no system header +#include #include #include #include @@ -49,6 +50,7 @@ #include #include #include +#include #include #include #include @@ -284,35 +286,14 @@ _CCCL_DEVICE void transform_kernel_vectorized( auto provide_array = [&](auto... inputs) { // load inputs [[maybe_unused]] auto load_tile = [](auto in, auto& input) { - using it_t = decltype(in); - using value_t = it_value_t; - if constexpr (THRUST_NS_QUALIFIER::is_contiguous_iterator_v) - { - // TODO(bgruber): we could add a max_load_store_size to the policy to avoid huge load types and huge alignment - // requirements - using load_t = decltype(load_store_type()); - auto in_vec = reinterpret_cast(in) + threadIdx.x; - auto input_vec = reinterpret_cast(input.data()); - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < load_store_count; ++i) - { - input_vec[i] = in_vec[i * ThreadsPerBlock]; - } - } - else - { - in += threadIdx.x * VecSize; // NOLINT(bugprone-misplaced-widening-cast) - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < load_store_count; ++i) - { - _CCCL_PRAGMA_UNROLL_FULL() - for (int j = 0; j < VecSize; ++j) - { - // NOLINTNEXTLINE(bugprone-misplaced-widening-cast) - input[i * VecSize + j] = in[i * VecSize * ThreadsPerBlock + j]; - } - } + using value_t = it_value_t; + if constexpr (::cuda::__is_contiguous_iterator_or_adaptor_v) + { // CanVectorize only selects this path if contiguous inputs are aligned to their vectors + const auto* base_ptr = ::cuda::__unwrap_iterator_adaptor(in); + _CCCL_ASSUME(reinterpret_cast<::cuda::std::uintptr_t>(base_ptr) % (sizeof(value_t) * VecSize) == 0); } + auto& items = *reinterpret_cast(input.data()); + LoadDirectStripedVectorized(static_cast(threadIdx.x), in, items); }; _CCCL_PDL_GRID_DEPENDENCY_SYNC(); (load_tile(ins, inputs), ...); diff --git a/cub/cub/warp/warp_load.cuh b/cub/cub/warp/warp_load.cuh index 3f44400ce1ec..b8fcff1810cb 100644 --- a/cub/cub/warp/warp_load.cuh +++ b/cub/cub/warp/warp_load.cuh @@ -335,19 +335,19 @@ private: _CCCL_DEVICE _CCCL_FORCEINLINE void Load(InputT* block_ptr, InputT (&items)[ItemsPerThread]) { - InternalLoadDirectBlockedVectorized(linear_tid, block_ptr, items); + LoadDirectBlockedVectorized(linear_tid, block_ptr, items); } _CCCL_DEVICE _CCCL_FORCEINLINE void Load(const InputT* block_ptr, InputT (&items)[ItemsPerThread]) { - InternalLoadDirectBlockedVectorized(linear_tid, block_ptr, items); + LoadDirectBlockedVectorized(linear_tid, block_ptr, items); } template _CCCL_DEVICE _CCCL_FORCEINLINE void Load(CacheModifiedInputIterator block_itr, InputT (&items)[ItemsPerThread]) { - InternalLoadDirectBlockedVectorized(linear_tid, block_itr.ptr, items); + LoadDirectBlockedVectorized(linear_tid, block_itr.ptr, items); } template diff --git a/cub/test/catch2_test_block_load.cu b/cub/test/catch2_test_block_load.cu index b5ff586dbded..46ae7903bd55 100644 --- a/cub/test/catch2_test_block_load.cu +++ b/cub/test/catch2_test_block_load.cu @@ -6,18 +6,100 @@ #include #include +#include +#include + +#include +#include #include +#include #include "cub_test_macros.h" -template -static __device__ int get_output_idx(int item) +struct widen_and_increment { - if (LoadAlgorithm == cub::BlockLoadAlgorithm::BLOCK_LOAD_STRIPED) + _CCCL_HOST_DEVICE ::cuda::std::uint32_t operator()(char value) const + { + return static_cast<::cuda::std::uint32_t>(value) + 1; + } +}; + +struct double_value +{ + _CCCL_HOST_DEVICE ::cuda::std::uint32_t operator()(::cuda::std::uint32_t value) const + { + return value * 2; + } +}; + +struct plus_one +{ + _CCCL_HOST_DEVICE int operator()(int value) const + { + return value + 1; + } +}; + +struct base_item_t +{ + int value; +}; + +struct derived_item_t : base_item_t +{ + int extra; +}; + +// Records whether it was assigned from an item in global memory, rather than from a copy of its bytes +struct source_tracking_item_t +{ + int value; + bool assigned_from_global; + + source_tracking_item_t() = default; + + _CCCL_HOST_DEVICE explicit source_tracking_item_t(int value) + : value{value} + , assigned_from_global{false} + {} + + source_tracking_item_t(const source_tracking_item_t&) = default; + + _CCCL_HOST_DEVICE source_tracking_item_t& operator=(const source_tracking_item_t& other) + { + value = other.value; + NV_IF_ELSE_TARGET( + NV_IS_DEVICE, + (assigned_from_global = ::cuda::device::is_address_from(&other, ::cuda::device::address_space::global);), + (assigned_from_global = false;)); + return *this; + } +}; +static_assert(!::cuda::is_trivially_copyable_v); + +template +static __device__ int get_output_idx(int item, bool full_tile) +{ + if constexpr (LoadAlgorithm == cub::BlockLoadAlgorithm::BLOCK_LOAD_STRIPED) { return static_cast(threadIdx.x) + ThreadsInBlock * item; } - return static_cast(threadIdx.x) * ItemsPerThread + item; + else if constexpr (LoadAlgorithm == cub::BlockLoadAlgorithm::BLOCK_LOAD_STRIPED_VECTORIZED) + { + // Full tiles are striped vectors of the widest vector access that evenly divides the items of a thread, or of + // single items if that does not hold a whole number of T. Guarded loads are striped. + using vector_word_t = typename cuda::__as_bytes_vectorized_t::value_type; + constexpr int vec_size = + sizeof(vector_word_t) % sizeof(T) == 0 ? static_cast(sizeof(vector_word_t) / sizeof(T)) : 1; + const int items_per_vector = full_tile ? vec_size : 1; + const int vector = item / items_per_vector; + const int element = item % items_per_vector; + return (static_cast(threadIdx.x) + ThreadsInBlock * vector) * items_per_vector + element; + } + else + { + return static_cast(threadIdx.x) * ItemsPerThread + item; + } } template (i); + const int idx = get_output_idx(i, full_tile); if (idx < num_items) { output[idx] = data[i]; @@ -61,9 +144,11 @@ template __global__ void kernel(cuda::std::false_type, InputIteratorT input, OutputIteratorT output, int num_items) { + using input_t = cub::detail::it_value_t; for (int i = 0; i < ItemsPerThread; i++) { - const int idx = get_output_idx(i); + const int idx = + get_output_idx(i, true); if (idx < num_items) { output[idx] = input[idx]; @@ -102,6 +187,7 @@ using load_algorithm = cub::BlockLoadAlgorithm::BLOCK_LOAD_DIRECT, cub::BlockLoadAlgorithm::BLOCK_LOAD_STRIPED, cub::BlockLoadAlgorithm::BLOCK_LOAD_VECTORIZE, + cub::BlockLoadAlgorithm::BLOCK_LOAD_STRIPED_VECTORIZED, cub::BlockLoadAlgorithm::BLOCK_LOAD_TRANSPOSE, cub::BlockLoadAlgorithm::BLOCK_LOAD_WARP_TRANSPOSE, cub::BlockLoadAlgorithm::BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED>; @@ -111,8 +197,14 @@ using odd_load_algorithm = cub::BlockLoadAlgorithm::BLOCK_LOAD_DIRECT, cub::BlockLoadAlgorithm::BLOCK_LOAD_STRIPED, cub::BlockLoadAlgorithm::BLOCK_LOAD_VECTORIZE, + cub::BlockLoadAlgorithm::BLOCK_LOAD_STRIPED_VECTORIZED, cub::BlockLoadAlgorithm::BLOCK_LOAD_TRANSPOSE>; +using vectorized_load_algorithm = + c2h::enum_type_list; + template struct params_t { @@ -207,20 +299,111 @@ CUB_TEST("Block load works with caching iterators", "[load][block]", CUB_SMALL, test_block_load(d_input, in); } +CUB_TEST("Block load works with fancy iterators", "[load][block]", CUB_SMALL, items_per_thread, load_algorithm) +{ + using type = int; + constexpr int items_per_thread = c2h::get<0, TestType>::value; + constexpr int threads_in_block = 64; + constexpr int tile_size = items_per_thread * threads_in_block; + static constexpr cub::BlockLoadAlgorithm load_algorithm = c2h::get<1, TestType>::value; + + c2h::device_vector d_input(GENERATE_COPY(take(10, random(0, tile_size)))); + c2h::gen(c2h::modulo_t{d_input.size()}, d_input); + test_block_load(d_input, cuda::counting_iterator{0}); + test_block_load(d_input, thrust::counting_iterator{0}); +} + +template +void test_widening_adaptor(int offset) +{ + constexpr int tile_size = ItemsPerThread * ThreadsInBlock; + + c2h::host_vector h_storage(static_cast(tile_size + offset)); + c2h::host_vector<::cuda::std::uint32_t> h_expected(tile_size); + for (int item = 0; item < tile_size + offset; ++item) + { + h_storage[item] = static_cast(item % 100); + } + for (int item = 0; item < tile_size; ++item) + { + h_expected[item] = (static_cast<::cuda::std::uint32_t>(h_storage[item + offset]) + 1) * 2; + } + + c2h::device_vector d_storage = h_storage; + c2h::device_vector<::cuda::std::uint32_t> d_expected = h_expected; + const char* base = thrust::raw_pointer_cast(d_storage.data()) + offset; + auto inner = ::cuda::transform_iterator{base, widen_and_increment{}}; + auto input = ::cuda::transform_iterator{inner, double_value{}}; + static_assert(::cuda::__iterator_can_rebase); + + test_block_load(d_expected, input); +} + +template +void test_same_size_adaptor(int offset) +{ + constexpr int tile_size = ItemsPerThread * ThreadsInBlock; + + c2h::host_vector h_storage(static_cast(tile_size + offset)); + for (int item = 0; item < tile_size + offset; ++item) + { + h_storage[item] = item * 17 + 3; + } + c2h::device_vector d_storage = h_storage; + c2h::host_vector h_expected(tile_size); + for (int item = 0; item < tile_size; ++item) + { + h_expected[item] = h_storage[item + offset] + 1; + } + + c2h::device_vector d_expected = h_expected; + auto input = + ::cuda::transform_iterator{thrust::raw_pointer_cast(d_storage.data()) + offset, plus_one{}}; + static_assert(::cuda::__iterator_can_rebase); + test_block_load(d_expected, input); + + // A thrust adaptor over the vector iterator is rebased through the device pointer of the vector + auto thrust_input = thrust::make_transform_iterator(d_storage.cbegin() + offset, plus_one{}); + static_assert(::cuda::__iterator_can_rebase); + test_block_load(d_expected, thrust_input); +} + #if IPT == 1 +// Offset 0 takes the vectorized load. Offset 1 is misaligned for the wider vector types and uses the scalar fallback. +// Offset 4 aligns the storage of 8 ints to 16 but not 32 bytes, so their vectors fall back to 16 bytes. +CUB_TEST("Vectorized block load works with contiguous iterator adaptors", + "[load][block]", + CUB_SMALL, + vectorized_load_algorithm) +{ + static constexpr cub::BlockLoadAlgorithm load_algorithm = c2h::get<0, TestType>::value; + + const int offset = GENERATE(0, 1, 4); + test_widening_adaptor<1, 32, load_algorithm>(offset); + test_widening_adaptor<4, 64, load_algorithm>(offset); + test_widening_adaptor<8, 128, load_algorithm>(offset); + test_same_size_adaptor<4, 64, load_algorithm>(offset); + test_same_size_adaptor<8, 64, load_algorithm>(offset); +} +#endif + +#if IPT == 1 +// With 8 items per thread, offset 4 aligns the input to 16 but not 32 bytes, so the vectors fall back to 16 bytes CUB_TEST("Vectorized block load with const and non-const datatype and different alignment cases", "[load][block]", CUB_SMALL, - c2h::type_list) + c2h::type_list, + vectorized_load_algorithm, + c2h::enum_type_list) { using type = int; using input_ptr_type = c2h::get<0, TestType>; const int offset_for_elements = GENERATE_COPY(0, 1, 2, 3, 4); - constexpr int items_per_thread = 4; + constexpr int items_per_thread = c2h::get<2, TestType>::value; constexpr int threads_in_block = 64; constexpr int tile_size = items_per_thread * threads_in_block; - static constexpr cub::BlockLoadAlgorithm load_algorithm = cub::BlockLoadAlgorithm::BLOCK_LOAD_VECTORIZE; + static constexpr cub::BlockLoadAlgorithm load_algorithm = c2h::get<1, TestType>::value; c2h::device_vector d_input_ref(tile_size); c2h::gen(C2H_SEED(10), d_input_ref); @@ -232,18 +415,18 @@ CUB_TEST("Vectorized block load with const and non-const datatype and different d_input_ref, thrust::raw_pointer_cast(d_input.data()) + offset_for_elements); } -template +template __global__ void vectorized_full_tile_load_kernel(InputIteratorT input, T* output) { T data[ItemsPerThread]; - cub::BlockLoad().Load(input, data); + cub::BlockLoad().Load(input, data); for (int i = 0; i < ItemsPerThread; i++) { - output[threadIdx.x * ItemsPerThread + i] = data[i]; + output[get_output_idx(i, true)] = data[i]; } } -template +template void test_vectorized_full_tile_load(const c2h::device_vector& d_input, InputIteratorT input) { constexpr int items_per_thread = 4; @@ -251,7 +434,7 @@ void test_vectorized_full_tile_load(const c2h::device_vector& d_input, I REQUIRE(d_input.size() == items_per_thread * threads_in_block); c2h::device_vector d_output(d_input.size(), thrust::no_init); - vectorized_full_tile_load_kernel + vectorized_full_tile_load_kernel <<<1, threads_in_block>>>(input, thrust::raw_pointer_cast(d_output.data())); REQUIRE(cudaSuccess == cudaPeekAtLastError()); REQUIRE(cudaSuccess == cudaDeviceSynchronize()); @@ -259,16 +442,201 @@ void test_vectorized_full_tile_load(const c2h::device_vector& d_input, I REQUIRE(expected == d_output); } -CUB_TEST_CASE("Vectorized block load converts from iterators of a different value type", "[load][block]", CUB_SMALL) +CUB_TEST("Vectorized block load converts from iterators of a different value type", + "[load][block]", + CUB_SMALL, + vectorized_load_algorithm) { + static constexpr cub::BlockLoadAlgorithm load_algorithm = c2h::get<0, TestType>::value; + c2h::device_vector d_input(256, thrust::no_init); c2h::gen(C2H_SEED(1), d_input); int* ptr = thrust::raw_pointer_cast(d_input.data()); - test_vectorized_full_tile_load(d_input, ptr); - test_vectorized_full_tile_load(d_input, cuda::std::span(ptr, d_input.size()).begin()); - test_vectorized_full_tile_load( + test_vectorized_full_tile_load(d_input, ptr); + test_vectorized_full_tile_load(d_input, cuda::std::span(ptr, d_input.size()).begin()); + test_vectorized_full_tile_load( d_input, cub::CacheModifiedInputIterator(ptr)); - test_vectorized_full_tile_load(d_input, cuda::std::span(ptr, d_input.size()).begin()); + test_vectorized_full_tile_load( + d_input, cuda::std::span(ptr, d_input.size()).begin()); + test_vectorized_full_tile_load(d_input, static_cast(ptr)); +} + +template +__global__ void load_direct_blocked_vectorized_kernel(InputIteratorT input, T* output) +{ + T data[ItemsPerThread]; + cub::LoadDirectBlockedVectorized( + static_cast(threadIdx.x), input, data); + for (int i = 0; i < ItemsPerThread; i++) + { + output[threadIdx.x * ItemsPerThread + i] = data[i]; + } +} + +template +void test_load_direct_blocked_vectorized(const c2h::device_vector& d_expected, InputIteratorT input) +{ + REQUIRE(d_expected.size() == ItemsPerThread * ThreadsInBlock); + + c2h::device_vector d_output(d_expected.size(), thrust::no_init); + load_direct_blocked_vectorized_kernel + <<<1, ThreadsInBlock>>>(input, thrust::raw_pointer_cast(d_output.data())); + REQUIRE(cudaSuccess == cudaPeekAtLastError()); + REQUIRE(cudaSuccess == cudaDeviceSynchronize()); + REQUIRE(d_expected == d_output); +} + +// Offset 0 takes the vectorized loads. Offset 1 misaligns the vectors of multiple items, which use the scalar fallback. +// An offset of 16 bytes misaligns the vectors wider than 16 bytes, which fall back to 16-byte vectors. +CUB_TEST("LoadDirectBlockedVectorized loads the items of a thread with vectors of VecSize items", + "[load][block]", + CUB_SMALL, + types) +{ + using type = c2h::get<0, TestType>; + constexpr int items_per_thread = 16; + constexpr int threads_in_block = 64; + constexpr int tile_size = items_per_thread * threads_in_block; + + const int offset = GENERATE(0, 1, static_cast(16 / sizeof(type))); + c2h::device_vector d_expected(tile_size); + c2h::gen(C2H_SEED(1), d_expected); + c2h::device_vector d_input(tile_size + offset); + thrust::copy_n(d_expected.begin(), tile_size, d_input.begin() + offset); + const type* input = thrust::raw_pointer_cast(d_input.data()) + offset; + + test_load_direct_blocked_vectorized<0, items_per_thread, threads_in_block>(d_expected, input); + test_load_direct_blocked_vectorized<1, items_per_thread, threads_in_block>(d_expected, input); + test_load_direct_blocked_vectorized<4, items_per_thread, threads_in_block>(d_expected, input); + test_load_direct_blocked_vectorized<16, items_per_thread, threads_in_block>(d_expected, input); +} + +// Every vector of a thread reads its items through a copy of the adaptor that is rebased onto that vector +CUB_TEST_CASE("LoadDirectBlockedVectorized loads contiguous iterator adaptors with vectors of VecSize items", + "[load][block]", + CUB_SMALL) +{ + constexpr int items_per_thread = 16; + constexpr int threads_in_block = 64; + constexpr int tile_size = items_per_thread * threads_in_block; + + const int offset = GENERATE(0, 1, 4); + c2h::host_vector h_storage(static_cast(tile_size + offset)); + for (int item = 0; item < tile_size + offset; ++item) + { + h_storage[item] = static_cast(item % 100); + } + c2h::host_vector h_expected(tile_size); + for (int item = 0; item < tile_size; ++item) + { + h_expected[item] = double_value{}(widen_and_increment{}(h_storage[item + offset])); + } + + const c2h::device_vector d_storage = h_storage; + const c2h::device_vector d_expected = h_expected; + const char* base = thrust::raw_pointer_cast(d_storage.data()) + offset; + auto inner = cuda::transform_iterator{base, widen_and_increment{}}; + auto input = cuda::transform_iterator{inner, double_value{}}; + static_assert(cuda::__iterator_can_rebase); + + test_load_direct_blocked_vectorized<1, items_per_thread, threads_in_block>(d_expected, input); + test_load_direct_blocked_vectorized<4, items_per_thread, threads_in_block>(d_expected, input); + test_load_direct_blocked_vectorized<16, items_per_thread, threads_in_block>(d_expected, input); +} +#endif + +template +__global__ void base_item_kernel(const derived_item_t* input, int* output) +{ + using block_load_t = cub::BlockLoad; + __shared__ typename block_load_t::TempStorage storage; + base_item_t data[ItemsPerThread]; + block_load_t(storage).Load(input, data); + + for (int i = 0; i < ItemsPerThread; i++) + { + output[get_output_idx(i, true)] = data[i].value; + } +} + +#if IPT == 1 +// Every item is converted from an element of a contiguous input of a different type +CUB_TEST("Vectorized block load converts the elements of a contiguous input to the item type", + "[load][block]", + CUB_SMALL, + vectorized_load_algorithm) +{ + static constexpr cub::BlockLoadAlgorithm load_algorithm = c2h::get<0, TestType>::value; + constexpr int items_per_thread = 4; + constexpr int threads_in_block = 64; + constexpr int tile_size = items_per_thread * threads_in_block; + + c2h::host_vector h_input(tile_size); + for (int item = 0; item < tile_size; ++item) + { + h_input[item].value = item; + h_input[item].extra = -1; + } + c2h::device_vector d_input = h_input; + c2h::device_vector d_expected(tile_size); + c2h::gen(c2h::modulo_t{d_expected.size()}, d_expected); + c2h::device_vector d_output(tile_size); + + base_item_kernel + <<<1, threads_in_block>>>(thrust::raw_pointer_cast(d_input.data()), thrust::raw_pointer_cast(d_output.data())); + REQUIRE(cudaSuccess == cudaPeekAtLastError()); + REQUIRE(cudaSuccess == cudaDeviceSynchronize()); + REQUIRE(d_expected == d_output); +} +#endif + +template +__global__ void source_tracking_item_kernel(const source_tracking_item_t* input, int* values, bool* assigned_from_global) +{ + using block_load_t = cub::BlockLoad; + __shared__ typename block_load_t::TempStorage storage; + source_tracking_item_t data[ItemsPerThread]; + block_load_t(storage).Load(input, data); + + for (int i = 0; i < ItemsPerThread; i++) + { + const int idx = get_output_idx(i, true); + values[idx] = data[i].value; + assigned_from_global[idx] = data[i].assigned_from_global; + } +} + +#if IPT == 1 +// Vectorized loads would assign the items from copies of their bytes in registers +CUB_TEST("Vectorized block load assigns items that are not trivially copyable from their source", + "[load][block]", + CUB_SMALL, + vectorized_load_algorithm) +{ + static constexpr cub::BlockLoadAlgorithm load_algorithm = c2h::get<0, TestType>::value; + constexpr int items_per_thread = 4; + constexpr int threads_in_block = 64; + constexpr int tile_size = items_per_thread * threads_in_block; + + c2h::host_vector h_input(tile_size); + for (int item = 0; item < tile_size; ++item) + { + h_input[item] = source_tracking_item_t{item}; + } + c2h::device_vector d_input = h_input; + c2h::device_vector d_expected(tile_size); + c2h::gen(c2h::modulo_t{d_expected.size()}, d_expected); + c2h::device_vector d_values(tile_size); + c2h::device_vector d_assigned_from_global(tile_size); + + source_tracking_item_kernel<<<1, threads_in_block>>>( + thrust::raw_pointer_cast(d_input.data()), + thrust::raw_pointer_cast(d_values.data()), + thrust::raw_pointer_cast(d_assigned_from_global.data())); + REQUIRE(cudaSuccess == cudaPeekAtLastError()); + REQUIRE(cudaSuccess == cudaDeviceSynchronize()); + REQUIRE(d_expected == d_values); + REQUIRE(c2h::device_vector(tile_size, true) == d_assigned_from_global); } #endif diff --git a/cub/test/catch2_test_device_for_copy.cu b/cub/test/catch2_test_device_for_copy.cu index aab0414d294d..43a800005ea4 100644 --- a/cub/test/catch2_test_device_for_copy.cu +++ b/cub/test/catch2_test_device_for_copy.cu @@ -76,8 +76,14 @@ CUB_TEST("Device for each works", "[for_copy][device]", CUB_SMALL) REQUIRE(num_of_once_marked_items == num_items); } -CUB_TEST("Device for each works with unaligned vectors", "[for_copy][device]", CUB_SMALL) +// An offset of 2 aligns 8-byte items to 16 but not 32 bytes, which suffices for their vectors of at most 16 bytes +CUB_TEST("Device for each works with unaligned vectors", + "[for_copy][device]", + CUB_SMALL, + c2h::type_list) { + using item_t = c2h::get<0, TestType>; + constexpr int max_items = 5000000; constexpr int min_items = 1; @@ -91,11 +97,11 @@ CUB_TEST("Device for each works with unaligned vectors", "[for_copy][device]", C const int offset = GENERATE(1, 2, 3); c2h::device_vector counts(num_items); - c2h::device_vector input(num_items + offset); + c2h::device_vector input(num_items + offset); thrust::sequence(c2h::device_policy, input.begin() + offset, input.end()); - int* d_counts = thrust::raw_pointer_cast(counts.data()); - int* d_input = thrust::raw_pointer_cast(input.data()) + offset; + int* d_counts = thrust::raw_pointer_cast(counts.data()); + item_t* d_input = thrust::raw_pointer_cast(input.data()) + offset; device_for_each_copy(d_input, d_input + num_items, incrementer_t{d_counts}); diff --git a/cub/test/catch2_test_device_histogram.cu b/cub/test/catch2_test_device_histogram.cu index 231af3e7f678..0252e616a257 100644 --- a/cub/test/catch2_test_device_histogram.cu +++ b/cub/test/catch2_test_device_histogram.cu @@ -726,6 +726,38 @@ CUB_TEST("DeviceHistogram::HistogramRange levels/samples aliasing", "[histogram_ } } +CUB_TEST("DeviceHistogram::HistogramEven works with unaligned samples", + "[histogram_even][device]", + CUB_SMALL, + c2h::type_list) +{ + using sample_t = c2h::get<0, TestType>; + constexpr int num_levels = 257; + const int num_samples = GENERATE(1000, 1 << 20); + const int offset = GENERATE(0, 1, 2, 3); + CAPTURE(c2h::type_name(), num_samples, offset); + + c2h::device_vector d_samples(num_samples + offset); + c2h::gen(C2H_SEED(1), d_samples, sample_t{0}, sample_t{num_levels - 2}); + auto d_histogram = c2h::device_vector(num_levels - 1); + + histogram_even( + thrust::raw_pointer_cast(d_samples.data()) + offset, + thrust::raw_pointer_cast(d_histogram.data()), + num_levels, + sample_t{0}, + sample_t{num_levels - 1}, + num_samples); + + const auto h_samples = c2h::host_vector(d_samples); + auto h_histogram = c2h::host_vector(num_levels - 1); + for (int i = offset; i < num_samples + offset; ++i) + { + ++h_histogram[static_cast(h_samples[i])]; + } + CHECK(d_histogram == h_histogram); +} + // Limit this large-memory reproducer to the host launch path. #if TEST_LAUNCH == 0 CUB_TEST("DeviceHistogram::MultiHistogramEven large privatized offsets", "[histogram_even][device]", CUB_LARGE) diff --git a/cub/test/catch2_test_device_reduce_iterators.cu b/cub/test/catch2_test_device_reduce_iterators.cu index c797d49b5ed2..ba4d42082afe 100644 --- a/cub/test/catch2_test_device_reduce_iterators.cu +++ b/cub/test/catch2_test_device_reduce_iterators.cu @@ -5,6 +5,8 @@ #include +#include + #include #include @@ -119,3 +121,62 @@ CUB_TEST("Device reduce compiles with discard output iterator", "[reduce][device // Run test device_reduce(in_it, cuda::discard_iterator(), num_items, reduction_op, init_value_t{}); } + +struct plus_one_t +{ + __host__ __device__ int operator()(int value) const + { + return value + 1; + } +}; + +struct widen_t +{ + __host__ __device__ int operator()(std::uint8_t value) const + { + return static_cast(value) + 1; + } +}; + +template +void check_sum(Iterator in, int num_items, std::int64_t expected) +{ + c2h::device_vector out(1, thrust::no_init); + device_sum(in, thrust::raw_pointer_cast(out.data()), num_items); + REQUIRE(out[0] == expected); +} + +// Full tiles exercise the vectorized load. The extra elements exercise the scalar tail. An offset of one int is +// aligned for the element type and misaligned for a 16-byte vector load. +CUB_TEST("Device sum of a transform iterator over contiguous storage", "[reduce][device]", CUB_SMALL) +{ + const int num_items = GENERATE(4096, 1000003, 1 << 20); + c2h::device_vector in(num_items); + thrust::sequence(c2h::device_policy, in.begin(), in.end()); + auto* d_in = thrust::raw_pointer_cast(in.data()); + + const auto expected = static_cast(num_items) * (num_items + 1) / 2; + check_sum(cuda::transform_iterator(d_in, plus_one_t{}), num_items, expected); + + if (num_items > 1) + { + const int tail = num_items - 1; + const auto tail_expected = expected - 1; + check_sum(cuda::transform_iterator(d_in + 1, plus_one_t{}), tail, tail_expected); + } +} + +CUB_TEST("Device sum of a widening transform iterator over contiguous storage", "[reduce][device]", CUB_SMALL) +{ + constexpr int num_items = 1 << 20; + c2h::device_vector in(num_items, 1); + auto* d_in = thrust::raw_pointer_cast(in.data()); + check_sum(cuda::transform_iterator(d_in, widen_t{}), num_items, static_cast(num_items) * 2); +} + +CUB_TEST("Device sum of a counting iterator stays scalar", "[reduce][device]", CUB_SMALL) +{ + constexpr int num_items = 1 << 20; + const auto expected = static_cast(num_items) * (num_items - 1) / 2; + check_sum(cuda::counting_iterator{0}, num_items, expected); +} diff --git a/cub/test/catch2_test_device_transform_vectorized.cu b/cub/test/catch2_test_device_transform_vectorized.cu index 1d84c0c46d6d..41198bb32464 100644 --- a/cub/test/catch2_test_device_transform_vectorized.cu +++ b/cub/test/catch2_test_device_transform_vectorized.cu @@ -5,7 +5,13 @@ #include +#include + +#include +#include + #include +#include #include "catch2_test_launch_helper.h" #include "cub_test_macros.h" @@ -76,6 +82,74 @@ CUB_TEST("DeviceTransform::Transform vectorized store widening from uint8", REQUIRE(reference_h == result); } +struct plus_one_t +{ + __host__ __device__ int operator()(int value) const + { + return value + 1; + } +}; + +// The input is a transform iterator over a pointer, so the kernel loads the stored elements and the functor runs on +// that copy. Counts include a partial last tile and a pointer that is not 16-byte aligned. +CUB_TEST("DeviceTransform vectorized load of a contiguous transform iterator", "[device][transform]", CUB_SMALL) +{ + const int num_items = GENERATE(4096, 4097, 100000); + c2h::device_vector in(num_items + 1); + thrust::sequence(c2h::device_policy, in.begin(), in.end()); + auto* d_in = thrust::raw_pointer_cast(in.data()); + + c2h::device_vector result(num_items, thrust::no_init); + transform_many(cuda::std::make_tuple(cuda::transform_iterator(d_in, plus_one_t{})), + result.begin(), + num_items, + cuda::std::identity{}); + + c2h::host_vector reference(num_items); + std::iota(reference.begin(), reference.end(), 1); + REQUIRE(reference == result); + + c2h::device_vector shifted(num_items - 1, thrust::no_init); + transform_many(cuda::std::make_tuple(cuda::transform_iterator(d_in + 1, plus_one_t{})), + shifted.begin(), + num_items - 1, + cuda::std::identity{}); + + c2h::host_vector shifted_reference(num_items - 1); + std::iota(shifted_reference.begin(), shifted_reference.end(), 2); + REQUIRE(shifted_reference == shifted); +} + +// Widening changes the stored type and the value type. The kernel has to load the stored bytes, then apply the functor. +CUB_TEST("DeviceTransform contiguous transform iterator widens the stored type", "[device][transform]", CUB_SMALL) +{ + const int num_items = GENERATE(4096, 4097, 100000); + c2h::device_vector in(num_items + 1); + thrust::sequence(c2h::device_policy, in.begin(), in.end()); + auto* d_in = thrust::raw_pointer_cast(in.data()); + + c2h::device_vector result(num_items, thrust::no_init); + transform_many(cuda::std::make_tuple(cuda::transform_iterator(d_in, cast_to{})), + result.begin(), + num_items, + cuda::std::identity{}); + + c2h::host_vector in_h = in; + c2h::host_vector reference(num_items); + std::transform(in_h.begin(), in_h.begin() + num_items, reference.begin(), cast_to{}); + REQUIRE(reference == result); + + c2h::device_vector shifted(num_items - 1, thrust::no_init); + transform_many(cuda::std::make_tuple(cuda::transform_iterator(d_in + 1, cast_to{})), + shifted.begin(), + num_items - 1, + cuda::std::identity{}); + + c2h::host_vector shifted_reference(num_items - 1); + std::transform(in_h.begin() + 1, in_h.begin() + num_items, shifted_reference.begin(), cast_to{}); + REQUIRE(shifted_reference == shifted); +} + #if TEST_LAUNCH == 0 struct ublkcp_store_vec_size_2_selector { diff --git a/cub/test/catch2_test_enum_formatting.cu b/cub/test/catch2_test_enum_formatting.cu index 32efbeb6aef9..f4c0b108fde3 100644 --- a/cub/test/catch2_test_enum_formatting.cu +++ b/cub/test/catch2_test_enum_formatting.cu @@ -84,6 +84,7 @@ void do_test(const Tester& tester) tester(cub::BlockLoadAlgorithm::BLOCK_LOAD_TRANSPOSE, "BLOCK_LOAD_TRANSPOSE"); tester(cub::BlockLoadAlgorithm::BLOCK_LOAD_WARP_TRANSPOSE, "BLOCK_LOAD_WARP_TRANSPOSE"); tester(cub::BlockLoadAlgorithm::BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED, "BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED"); + tester(cub::BlockLoadAlgorithm::BLOCK_LOAD_STRIPED_VECTORIZED, "BLOCK_LOAD_STRIPED_VECTORIZED"); tester(cub::BlockLoadAlgorithm(100), ""); } diff --git a/cudax/include/cuda/experimental/__cuco/detail/hyperloglog/hyperloglog_impl.cuh b/cudax/include/cuda/experimental/__cuco/detail/hyperloglog/hyperloglog_impl.cuh index c5d4712c6658..6e7aae62d640 100644 --- a/cudax/include/cuda/experimental/__cuco/detail/hyperloglog/hyperloglog_impl.cuh +++ b/cudax/include/cuda/experimental/__cuco/detail/hyperloglog/hyperloglog_impl.cuh @@ -28,6 +28,7 @@ #include #include #include +#include #include #include #include @@ -223,10 +224,9 @@ public: }(); if constexpr (__can_vectorize) { - const auto __ptr = ::cuda::std::to_address(__first); - constexpr auto __max_vector_bytes = 32; + const auto __ptr = ::cuda::std::to_address(__first); const auto __alignment = - 1u << ::cuda::std::countr_zero(reinterpret_cast<::cuda::std::uintptr_t>(__ptr) | __max_vector_bytes); + 1u << ::cuda::std::countr_zero(reinterpret_cast<::cuda::std::uintptr_t>(__ptr) | ::cuda::__max_vector_access); const auto __vector_size = __alignment / sizeof(__value_type); switch (__vector_size) diff --git a/cudax/include/cuda/experimental/__cuco/detail/hyperloglog/kernels.cuh b/cudax/include/cuda/experimental/__cuco/detail/hyperloglog/kernels.cuh index a33e2f47455a..2c2894928c4f 100644 --- a/cudax/include/cuda/experimental/__cuco/detail/hyperloglog/kernels.cuh +++ b/cudax/include/cuda/experimental/__cuco/detail/hyperloglog/kernels.cuh @@ -21,8 +21,8 @@ # pragma system_header #endif // no system header +#include #include -#include #include #include @@ -67,9 +67,8 @@ template _CCCL_KERNEL_ATTRIBUTES void __add_shmem_vectorized(const typename _RefType::__value_type* __first, ::cuda::std::int64_t __n, _RefType __ref) { - using __value_type = typename _RefType::__value_type; - // TODO: replace with ::cuda::__vector_type - using __vector_type = ::cuda::std::array<__value_type, _VectorSize>; + using __value_type = typename _RefType::__value_type; + using __vector_type = ::cuda::__as_bytes_vectorized_t<__value_type, _VectorSize>; using __local_ref_type = typename _RefType::template __rebind_scope<::cuda::std::thread_scope_block>; // Base address of dynamic shared memory is guaranteed to be aligned to at least 16 bytes which is @@ -93,7 +92,7 @@ __add_shmem_vectorized(const typename _RefType::__value_type* __first, ::cuda::s ::cuda::std::assume_aligned(__first + __idx * _VectorSize)); for (int i = 0; i < _VectorSize; ++i) { - __local_ref.__add(__vec[i]); + __local_ref.__add(__vec.template __alias<__value_type>(i)); } __idx += __loop_stride; diff --git a/libcudacxx/include/cuda/__iterator/contiguous_iterator_adaptor.h b/libcudacxx/include/cuda/__iterator/contiguous_iterator_adaptor.h new file mode 100644 index 000000000000..243d6e283dc6 --- /dev/null +++ b/libcudacxx/include/cuda/__iterator/contiguous_iterator_adaptor.h @@ -0,0 +1,122 @@ +//===----------------------------------------------------------------------===// +// +// Part of libcu++, the C++ Standard Library for your entire system, +// under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +// +//===----------------------------------------------------------------------===// + +#ifndef _CUDA___ITERATOR_CONTIGUOUS_ITERATOR_ADAPTOR_H +#define _CUDA___ITERATOR_CONTIGUOUS_ITERATOR_ADAPTOR_H + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +_CCCL_BEGIN_NAMESPACE_CUDA + +//! @brief True when `_Iter` is an iterator adaptor that exposes its base iterator through `base()`. +template +_CCCL_CONCEPT __iterator_can_unwrap = _CCCL_REQUIRES_EXPR((_Iter), const _Iter& __iter)((__iter.base())); + +//! @brief Peel iterator adaptors until the base iterator remains. +//! @param[in] __iter The iterator to unwrap. +//! @return The address of the storage of a contiguous iterator. Otherwise the unwrapped base iterator of an iterator +//! adaptor, or `__iter` when it is not an iterator adaptor. +template +[[nodiscard]] _CCCL_API constexpr auto __unwrap_iterator_adaptor(const _Iter& __iter) +{ + if constexpr (::cuda::std::__can_to_address<_Iter>) + { // Support standard contiguous iterators + return ::cuda::std::to_address(__iter); + } + else if constexpr (__iterator_can_unwrap<_Iter>) + { + return ::cuda::__unwrap_iterator_adaptor(__iter.base()); + } + else + { + return __iter; + } +} + +//! @brief The contiguous iterator that remains after peeling adaptors from `_Iter`. +template +using __unwrapped_iterator_t = decltype(::cuda::__unwrap_iterator_adaptor(::cuda::std::declval<_Iter>())); + +//! @brief The value type of the storage behind `_Iter`. That is the element type of the pointer `_Iter` unwraps to, or +//! the value type of `_Iter` if it does not unwrap to a pointer, e.g. because the base of an adaptor is a counter. +//! Only `const` is removed from the element type, because volatile storage must be accessed item by item. +template +using __unwrapped_iter_value_t = + ::cuda::std::conditional_t<::cuda::std::is_pointer_v<__unwrapped_iterator_t<_Iter>>, + ::cuda::std::remove_const_t<::cuda::std::remove_pointer_t<__unwrapped_iterator_t<_Iter>>>, + typename ::cuda::std::iterator_traits<_Iter>::value_type>; + +template > +_CCCL_CONCEPT __iterator_can_rebase = + _CCCL_REQUIRES_EXPR((_Iter, _Tp), _Iter& __iter, _Tp* __ptr)((__iter.__rebase(__ptr))); + +template +_CCCL_API constexpr void __rebase_contiguous_iterator_adaptor(_Iter& __iter, _Tp* __ptr) +{ + if constexpr (__iterator_can_rebase<_Iter, _Tp>) + { + __iter.__rebase(__ptr); + } + else if constexpr (::cuda::std::is_assignable_v<_Iter&, _Tp*>) + { + __iter = __ptr; + } + else if constexpr (::cuda::std::is_constructible_v<_Iter, _Tp*>) + { + __iter = _Iter(__ptr); + } + else + { + static_assert(::cuda::std::__always_false_v<_Iter>, "Invalid pointer to rebase"); + } +} + +template +inline constexpr bool __is_contiguous_iterator_or_adaptor_impl = + ::cuda::std::is_pointer_v<_Iter> || ::cuda::std::__has_contiguous_traversal<_Iter>; + +template +inline constexpr bool __is_contiguous_iterator_or_adaptor_impl<::cuda::transform_iterator<_Fn, _Iter>> = + __is_contiguous_iterator_or_adaptor_impl<::cuda::std::remove_cvref_t<_Iter>>; + +//! @brief True when `_Iter` is contiguous storage, or an adaptor whose stored iterator is. +//! @tparam _Iter The iterator type to query. cv-qualifiers and references are ignored. +template +inline constexpr bool __is_contiguous_iterator_or_adaptor_v = + __is_contiguous_iterator_or_adaptor_impl<::cuda::std::remove_cvref_t<_Iter>>; + +_CCCL_END_NAMESPACE_CUDA + +#include + +#endif // _CUDA___ITERATOR_CONTIGUOUS_ITERATOR_ADAPTOR_H diff --git a/libcudacxx/include/cuda/__iterator/transform_input_output_iterator.h b/libcudacxx/include/cuda/__iterator/transform_input_output_iterator.h index 643c2b50ca22..e4faa4fbacf2 100644 --- a/libcudacxx/include/cuda/__iterator/transform_input_output_iterator.h +++ b/libcudacxx/include/cuda/__iterator/transform_input_output_iterator.h @@ -24,6 +24,7 @@ #if _LIBCUDACXX_HAS_SPACESHIP_OPERATOR() # include #endif // _LIBCUDACXX_HAS_SPACESHIP_OPERATOR() +#include #include #include #include @@ -248,6 +249,16 @@ class transform_input_output_iterator return ::cuda::std::move(__iter()); } + //! @brief Retargets the stored iterator at @p __ptr. The stored functors are left unchanged. The input functor is + //! then invoked on the elements at @p __ptr. + //! @param[in] __ptr The address that replaces the contiguous storage this adaptor reads from and writes to. + _CCCL_TEMPLATE(class _Tp) + _CCCL_REQUIRES(__iterator_can_rebase<_Iter, _Tp> || ::cuda::std::is_pointer_v<_Iter>) + _CCCL_API constexpr void __rebase(_Tp* __ptr) & + { + ::cuda::__rebase_contiguous_iterator_adaptor(__iter(), __ptr); + } + //! @brief Dereferences the @c transform_input_output_iterator. Returns a proxy that transforms values read from the //! stored iterator via the stored input functor and transforms assigned values via the output functor _CCCL_EXEC_CHECK_DISABLE diff --git a/libcudacxx/include/cuda/__iterator/transform_iterator.h b/libcudacxx/include/cuda/__iterator/transform_iterator.h index 640e57aa0dd3..77beb29b90e0 100644 --- a/libcudacxx/include/cuda/__iterator/transform_iterator.h +++ b/libcudacxx/include/cuda/__iterator/transform_iterator.h @@ -25,6 +25,7 @@ #if _LIBCUDACXX_HAS_SPACESHIP_OPERATOR() # include #endif // _LIBCUDACXX_HAS_SPACESHIP_OPERATOR() +#include #include #include #include @@ -40,6 +41,7 @@ #include #include #include +#include #include #include #include @@ -233,6 +235,16 @@ class transform_iterator : public __transform_iterator_category_base<_Fn, _Iter> return ::cuda::std::move(__iter()); } + //! @brief Retargets the stored iterator at @p __ptr. The stored functor is left unchanged. It is then invoked on the + //! elements at @p __ptr. + //! @param[in] __ptr The address that replaces the contiguous storage this adaptor reads from. + _CCCL_TEMPLATE(class _Tp) + _CCCL_REQUIRES(__iterator_can_rebase<_Iter, _Tp> || ::cuda::std::is_pointer_v<_Iter>) + _CCCL_API constexpr void __rebase(_Tp* __ptr) & + { + ::cuda::__rebase_contiguous_iterator_adaptor(__iter(), __ptr); + } + //! @brief Dereferences the stored iterator and applies the stored functor to the result _CCCL_EXEC_CHECK_DISABLE #ifndef _CCCL_DOXYGEN_INVOKED // Doxygen has issues with the constraint diff --git a/libcudacxx/include/cuda/__iterator/transform_output_iterator.h b/libcudacxx/include/cuda/__iterator/transform_output_iterator.h index db9103d21138..3bbe61960b77 100644 --- a/libcudacxx/include/cuda/__iterator/transform_output_iterator.h +++ b/libcudacxx/include/cuda/__iterator/transform_output_iterator.h @@ -24,6 +24,7 @@ #if _LIBCUDACXX_HAS_SPACESHIP_OPERATOR() # include #endif // _LIBCUDACXX_HAS_SPACESHIP_OPERATOR() +#include #include #include #include @@ -214,6 +215,15 @@ class transform_output_iterator return ::cuda::std::move(__iter()); } + //! @brief Retargets the stored iterator at @p __ptr. The stored functor is left unchanged. + //! @param[in] __ptr The address that replaces the contiguous storage this adaptor writes to. + _CCCL_TEMPLATE(class _Tp) + _CCCL_REQUIRES(__iterator_can_rebase<_Iter, _Tp> || ::cuda::std::is_pointer_v<_Iter>) + _CCCL_API constexpr void __rebase(_Tp* __ptr) & + { + ::cuda::__rebase_contiguous_iterator_adaptor(__iter(), __ptr); + } + //! @brief Returns a proxy that transforms the input upon assignment _CCCL_EXEC_CHECK_DISABLE [[nodiscard]] _CCCL_API constexpr auto operator*() const noexcept(::cuda::std::is_nothrow_copy_constructible_v<_Iter>) diff --git a/libcudacxx/include/cuda/__mdspan/__copy/copy_contiguous.h b/libcudacxx/include/cuda/__mdspan/__copy/copy_contiguous.h index 2d0a810d2f55..bc8f475949c2 100644 --- a/libcudacxx/include/cuda/__mdspan/__copy/copy_contiguous.h +++ b/libcudacxx/include/cuda/__mdspan/__copy/copy_contiguous.h @@ -27,6 +27,8 @@ #include #include #include +#include +#include #include #include #include diff --git a/libcudacxx/include/cuda/__mdspan/__copy/dispatch_by_vector.h b/libcudacxx/include/cuda/__mdspan/__copy/dispatch_by_vector.h index a91dc7ab4ad2..61e6ded9f1fc 100644 --- a/libcudacxx/include/cuda/__mdspan/__copy/dispatch_by_vector.h +++ b/libcudacxx/include/cuda/__mdspan/__copy/dispatch_by_vector.h @@ -25,6 +25,7 @@ # include # include +# include # include # include # include diff --git a/libcudacxx/include/cuda/__mdspan/__copy/mdspan_d2d.h b/libcudacxx/include/cuda/__mdspan/__copy/mdspan_d2d.h index b17211d1e7d7..fbecded6a05c 100644 --- a/libcudacxx/include/cuda/__mdspan/__copy/mdspan_d2d.h +++ b/libcudacxx/include/cuda/__mdspan/__copy/mdspan_d2d.h @@ -37,9 +37,9 @@ # include # include # include -# include # include # include +# include # include # include # include diff --git a/libcudacxx/include/cuda/__mdspan/__copy/tensor_copy_utils.h b/libcudacxx/include/cuda/__mdspan/__copy/tensor_copy_utils.h index 1452b257089a..b73ac873f847 100644 --- a/libcudacxx/include/cuda/__mdspan/__copy/tensor_copy_utils.h +++ b/libcudacxx/include/cuda/__mdspan/__copy/tensor_copy_utils.h @@ -25,7 +25,7 @@ # include # include -# include +# include # include # include # include @@ -89,8 +89,8 @@ __max_alignment(const __raw_tensor<_ExtentT, _StrideT, _Tp, _MaxRank>& __tensor) template <::cuda::std::size_t _VectorBytes, typename _Tp> using __reshape_vector_type _CCCL_NODEBUG = ::cuda::std::conditional_t<::cuda::std::is_const_v<_Tp>, - const ::cuda::__vector_access_t<_VectorBytes>, - ::cuda::__vector_access_t<_VectorBytes>>; + const ::cuda::__as_bytes_vectorized_t, + ::cuda::__as_bytes_vectorized_t>; //! @brief Reshape a raw tensor for vectorized access by widening the element type. //! diff --git a/libcudacxx/include/cuda/__mdspan/__copy/vector_access.h b/libcudacxx/include/cuda/__mdspan/__copy/vector_access.h index 1d4efd7e713f..9c5b60c20974 100644 --- a/libcudacxx/include/cuda/__mdspan/__copy/vector_access.h +++ b/libcudacxx/include/cuda/__mdspan/__copy/vector_access.h @@ -32,25 +32,8 @@ _CCCL_BEGIN_NAMESPACE_CUDA -//! @brief Aligned storage type for vectorized memory access of a given byte width. -template <::cuda::std::size_t _VectorBytes> -struct alignas(_VectorBytes) __vector_access -{ - char __data[_VectorBytes]; -}; - -// 32-byte accesses are supported since CTK 13.0 -#if _CCCL_CTK_AT_LEAST(13, 0) -inline constexpr auto __max_vector_access = 32; -#else -inline constexpr auto __max_vector_access = 16; -#endif // _CCCL_CTK_AT_LEAST(13, 0) - #if !_CCCL_COMPILER(NVRTC) -template <::cuda::std::size_t _VectorBytes> -using __vector_access_t _CCCL_NODEBUG = __vector_access<_VectorBytes>; - //! @brief Query the maximum vector access width supported by the current GPU architecture. //! //! @return Maximum vector width in bytes (32 for SM >= 10.0, 16 otherwise) diff --git a/libcudacxx/include/cuda/__memory/as_bytes_vectorized.h b/libcudacxx/include/cuda/__memory/as_bytes_vectorized.h new file mode 100644 index 000000000000..afea0ece1abc --- /dev/null +++ b/libcudacxx/include/cuda/__memory/as_bytes_vectorized.h @@ -0,0 +1,157 @@ +//===----------------------------------------------------------------------===// +// +// Part of libcu++, the C++ Standard Library for your entire system, +// under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +// +//===----------------------------------------------------------------------===// + +#ifndef _CUDA___MEMORY_AS_BYTES_VECTORIZED_H +#define _CUDA___MEMORY_AS_BYTES_VECTORIZED_H + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +_CCCL_BEGIN_NAMESPACE_CUDA + +// 32-byte accesses are supported since CTK 13.0 +#if _CCCL_CTK_AT_LEAST(13, 0) +inline constexpr size_t __max_vector_access = 32ull; +#else +inline constexpr size_t __max_vector_access = 16ull; +#endif // _CCCL_CTK_AT_LEAST(13, 0) + +//! Uninitialized storage of `_NumVectors` vectors of `_VectorSize` words, aligned for the widest vector access that +//! evenly divides it, at most to `_MaxVectorAccess` but at least to `_MinAlignment` +template +struct __bytes_vectorized_t +{ + static_assert(sizeof(_WordType) * _VectorSize <= _MaxVectorAccess, "A vector must fit into a single vector access"); + + using value_type = _WordType[_VectorSize]; + + // Aligned to the widest vector access that evenly divides the storage size + static constexpr size_t __full_alignment = size_t{1} << ::cuda::std::countr_zero(sizeof(value_type) * _NumVectors); + static constexpr size_t __alignment = + ::cuda::std::max(_MinAlignment, ::cuda::std::min(__full_alignment, _MaxVectorAccess)); + + // A single array, so that a pointer to the first word can be iterated over the whole storage + alignas(__alignment) _WordType __storage_[_NumVectors * _VectorSize]; + + static_assert(sizeof(__storage_) % __alignment == 0, "Storage size must be a multiple of its alignment"); + + [[nodiscard]] _CCCL_API constexpr _WordType* data() noexcept + { + return __storage_; + } + + [[nodiscard]] _CCCL_API constexpr const _WordType* data() const noexcept + { + return __storage_; + } + + //! Returns the `__index`-th object of type `_Tp` in the storage + template + [[nodiscard]] _CCCL_API _Tp& __alias(const _Index __index) noexcept + { + static_assert(::cuda::std::__cccl_is_integer_v<_Index>, "The index must be an integer"); + static_assert(sizeof(_Tp) <= sizeof(__storage_), "T does not fit in vectorized storage"); + static_assert(alignof(_Tp) <= __alignment, "T is more aligned than vectorized storage"); + const auto __i = static_cast(__index); + if constexpr (::cuda::std::is_same_v<::cuda::std::remove_cvref_t<_Tp>, _WordType>) + { + return __storage_[__i]; + } + else + { + return ::cuda::std::launder(reinterpret_cast<_Tp*>(__storage_))[__i]; + } + } + + //! Returns the `__index`-th object of type `_Tp` in the storage + template + [[nodiscard]] _CCCL_API const _Tp& __alias(const _Index __index) const noexcept + { + static_assert(::cuda::std::__cccl_is_integer_v<_Index>, "The index must be an integer"); + static_assert(sizeof(_Tp) <= sizeof(__storage_), "T does not fit in vectorized storage"); + static_assert(alignof(_Tp) <= __alignment, "T is more aligned than vectorized storage"); + const auto __i = static_cast(__index); + if constexpr (::cuda::std::is_same_v<::cuda::std::remove_cvref_t<_Tp>, _WordType>) + { + return __storage_[__i]; + } + else + { + return ::cuda::std::launder(reinterpret_cast(__storage_))[__i]; + } + } +}; + +template +[[nodiscard]] _CCCL_API constexpr auto __as_bytes_vectorized_impl() noexcept +{ + if constexpr (alignof(_Tp) <= _MaxVectorAccess) + { + using _Word = + // Arithmetic types use themselves as the word, except for 128bit integers which show regressions on older GPUs. + ::cuda::std::conditional_t<::cuda::std::is_arithmetic_v<_Tp> && sizeof(_Tp) <= 8ull, + _Tp, + ::cuda::std::__make_nbit_uint_t<::cuda::std::min(alignof(_Tp), size_t{8}) * CHAR_BIT>>; + constexpr size_t __num_words = sizeof(_Tp) / sizeof(_Word) * _VecSize; + constexpr size_t __words_per_vector = + ::cuda::std::min(size_t{1} << ::cuda::std::countr_zero(__num_words), _MaxVectorAccess / sizeof(_Word)); + return __bytes_vectorized_t<_Word, + __num_words / __words_per_vector, + __words_per_vector, + alignof(_Tp), + _MaxVectorAccess>{}; + } + else + { // Default case: use 32 bit words as the vectorized storage + return __bytes_vectorized_t<::cuda::std::uint32_t, + ::cuda::ceil_div(sizeof(_Tp), 4) * _VecSize, + 1, + alignof(_Tp), + _MaxVectorAccess>{}; + } +} + +//! Storage for `_VecSize` objects of type `_Tp`. Arithmetic types use themselves as the word. Other types use unsigned +//! integer words of the alignment of `_Tp`, at most 8 bytes. Words are grouped into the widest vectors that evenly +//! divide the storage. The storage is aligned for the widest vector access of its size, at most to +//! `_MaxVectorAccess` but at least to `alignof(_Tp)` +template +using __as_bytes_vectorized_t = decltype(::cuda::__as_bytes_vectorized_impl<_Tp, _VecSize, _MaxVectorAccess>()); + +_CCCL_END_NAMESPACE_CUDA + +#include + +#endif // _CUDA___MEMORY_AS_BYTES_VECTORIZED_H diff --git a/libcudacxx/include/cuda/iterator b/libcudacxx/include/cuda/iterator index a9b0630598d4..15176155f6e6 100644 --- a/libcudacxx/include/cuda/iterator +++ b/libcudacxx/include/cuda/iterator @@ -23,6 +23,7 @@ // IWYU pragma: begin_exports #include +#include #include #include #include diff --git a/libcudacxx/include/cuda/std/__iterator/iterator_traits.h b/libcudacxx/include/cuda/std/__iterator/iterator_traits.h index 344417eb8505..3b80dce313ea 100644 --- a/libcudacxx/include/cuda/std/__iterator/iterator_traits.h +++ b/libcudacxx/include/cuda/std/__iterator/iterator_traits.h @@ -532,6 +532,11 @@ inline constexpr bool __has_random_access_traversal = __has_iterator_category_convertible_to<_Iter, random_access_iterator_tag> || __has_iterator_concept_convertible_to<_Iter, random_access_iterator_tag>; +// in C++17 the standard does not yet know about the contiguous iterator concept, so we need to check for it explicitly +template +_CCCL_CONCEPT __is_cpp17_contiguous_iterator = _CCCL_REQUIRES_EXPR((_Iter)) // + (typename(_ITER_CONCEPT<_Iter>), requires(is_convertible_v<_ITER_CONCEPT<_Iter>, contiguous_iterator_tag>)); + // __has_contiguous_traversal determines if an iterator is known by // libc++ to be contiguous, either because it advertises itself as such // (in C++20) or because it is a pointer type or a known trivial wrapper @@ -542,7 +547,7 @@ inline constexpr bool __has_random_access_traversal = template inline constexpr bool __has_contiguous_traversal = __has_iterator_category_convertible_to<_Iter, contiguous_iterator_tag> - || __has_iterator_concept_convertible_to<_Iter, contiguous_iterator_tag>; + || __has_iterator_concept_convertible_to<_Iter, contiguous_iterator_tag> || __is_cpp17_contiguous_iterator<_Iter>; // Any native pointer which is an iterator is also a contiguous iterator. template diff --git a/libcudacxx/test/libcudacxx/cuda/iterators/contiguous_iterator_adaptor.pass.cpp b/libcudacxx/test/libcudacxx/cuda/iterators/contiguous_iterator_adaptor.pass.cpp new file mode 100644 index 000000000000..41543350c13e --- /dev/null +++ b/libcudacxx/test/libcudacxx/cuda/iterators/contiguous_iterator_adaptor.pass.cpp @@ -0,0 +1,329 @@ +//===----------------------------------------------------------------------===// +// +// Part of libcu++, the C++ Standard Library for your entire system, +// under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +// +//===----------------------------------------------------------------------===// + +// Transform iterators satisfy __is_contiguous_iterator_or_adaptor_v when their stored iterator does, transform output +// and input output iterators never do. __unwrap_iterator_adaptor peels all of these adaptors down to the contiguous +// iterator, and __rebase aims the adaptor at a different contiguous region. + +#include +#include +#include +#include + +#include "test_iterators.h" +#include "test_macros.h" + +struct PlusOne +{ + TEST_FUNC constexpr int operator()(int value) const + { + return value + 1; + } +}; + +struct TimesTwo +{ + TEST_FUNC constexpr int operator()(int value) const + { + return value * 2; + } +}; + +struct Narrow +{ + TEST_FUNC constexpr char operator()(int value) const + { + return static_cast(value); + } +}; + +// Contiguous iterator with no base(). Unwrap reaches its pointer through to_address. +struct addressable_iterator +{ + using iterator_category = cuda::std::contiguous_iterator_tag; + using value_type = int; + using difference_type = cuda::std::ptrdiff_t; + using pointer = int*; + using reference = int&; + + int* ptr_; + + TEST_FUNC constexpr explicit addressable_iterator(int* ptr) + : ptr_(ptr) + {} + + TEST_FUNC constexpr reference operator*() const + { + return *ptr_; + } + + TEST_FUNC constexpr pointer operator->() const + { + return ptr_; + } +}; + +template +TEST_FUNC constexpr void test_trait() +{ + static_assert(cuda::__is_contiguous_iterator_or_adaptor_v == Expected); + static_assert(cuda::__is_contiguous_iterator_or_adaptor_v == Expected); + static_assert(cuda::__is_contiguous_iterator_or_adaptor_v == Expected); + static_assert(cuda::__is_contiguous_iterator_or_adaptor_v == Expected); + static_assert(cuda::__is_contiguous_iterator_or_adaptor_v == Expected); + static_assert(cuda::__is_contiguous_iterator_or_adaptor_v == Expected); + static_assert(cuda::__is_contiguous_iterator_or_adaptor_v == Expected); +} + +TEST_FUNC constexpr bool test() +{ + int buffer[4] = {1, 2, 3, 4}; + int staged[4] = {10, 20, 30, 40}; + + test_trait(); + test_trait(); + test_trait, true>(); + test_trait, false>(); + test_trait, false>(); + + { + int* input = buffer; + auto unwrapped = cuda::__unwrap_iterator_adaptor(input); + static_assert(cuda::std::is_same_v, int*>); + static_assert(cuda::std::is_same_v, int>); + static_assert(cuda::std::is_same_v, int>); + static_assert(cuda::std::is_same_v, volatile int>); + static_assert(cuda::std::is_same_v, volatile int>); + static_assert(cuda::std::is_same_v); + assert(unwrapped == buffer); + + cuda::counting_iterator counting{4}; + auto same = cuda::__unwrap_iterator_adaptor(counting); + static_assert(cuda::std::is_same_v>); + static_assert(cuda::std::is_same_v>, int>); + static_assert(!cuda::__iterator_can_rebase>); + assert(*same == 4); + } + + { + using Iter = cuda::transform_iterator; + test_trait(); + static_assert( + cuda::__is_contiguous_iterator_or_adaptor_v == cuda::__is_contiguous_iterator_or_adaptor_v); + + Iter iter{buffer, PlusOne{}}; + auto advanced = iter + 2; + auto unwrapped = cuda::__unwrap_iterator_adaptor(advanced); + static_assert(cuda::std::is_same_v, int*>); + static_assert(cuda::std::is_same_v, int>); + static_assert(cuda::std::is_same_v); + assert(unwrapped == buffer + 2); + + auto rebased = iter; + rebased.__rebase(staged); + assert(*iter == 2); + assert(*rebased == 11); + assert(rebased.base() == staged); + assert(rebased[3] == 41); + } + + { + using Iter = cuda::transform_iterator; + test_trait(); + const int* input = buffer; + Iter iter{input, PlusOne{}}; + auto unwrapped = cuda::__unwrap_iterator_adaptor(iter); + static_assert(cuda::std::is_same_v); + assert(unwrapped == buffer); + + auto rebased = iter; + rebased.__rebase(staged); + assert(*iter == 2); + assert(*rebased == 11); + assert(rebased.base() == staged); + } + + { + using Leaf = contiguous_iterator; + using Iter = cuda::transform_iterator; + test_trait(); + Iter iter{Leaf{buffer}, PlusOne{}}; + auto unwrapped = cuda::__unwrap_iterator_adaptor(iter); + static_assert(cuda::std::is_same_v); + assert(unwrapped == buffer); + + static_assert(!cuda::__iterator_can_rebase); + } + + { + using Inner = cuda::transform_iterator; + using Outer = cuda::transform_iterator; + test_trait(); + static_assert( + cuda::__is_contiguous_iterator_or_adaptor_v == cuda::__is_contiguous_iterator_or_adaptor_v); + + Outer iter{Inner{buffer, PlusOne{}}, TimesTwo{}}; + auto unwrapped = cuda::__unwrap_iterator_adaptor(iter); + static_assert(cuda::std::is_same_v); + assert(unwrapped == buffer); + assert(*iter == 4); + + auto rebased = iter; + rebased.__rebase(staged); + assert(*iter == 4); + assert(*rebased == 22); + assert(rebased.base().base() == staged); + } + + { + using Iter = addressable_iterator; + test_trait(); + Iter iter{buffer + 2}; + auto unwrapped = cuda::__unwrap_iterator_adaptor(iter); + static_assert(cuda::std::is_same_v, int*>); + static_assert(cuda::std::is_same_v, int>); + static_assert(cuda::std::is_same_v); + assert(unwrapped == buffer + 2); + + using Transformed = cuda::transform_iterator; + test_trait(); + Transformed transformed{iter, PlusOne{}}; + auto transformed_unwrapped = cuda::__unwrap_iterator_adaptor(transformed); + static_assert(cuda::std::is_same_v); + assert(transformed_unwrapped == buffer + 2); + assert(*transformed == 4); + static_assert(!cuda::__iterator_can_rebase); + } + + { + using Bad = cuda::transform_iterator>; + test_trait(); + test_trait, false>(); + test_trait>, false>(); + // A reverse iterator unwraps to a pointer, but does not traverse the storage forward + test_trait>, false>(); + + Bad iter{cuda::counting_iterator{5}, PlusOne{}}; + auto unwrapped = cuda::__unwrap_iterator_adaptor(iter); + static_assert(cuda::std::is_same_v>); + assert(*unwrapped == 5); + + // Without storage behind the adaptor, its own value type is used + using Narrowed = cuda::transform_iterator>; + static_assert(cuda::std::is_same_v, char>); + static_assert(!cuda::__iterator_can_rebase); + } + + { + using Iter = cuda::transform_iterator; + test_trait(); + static_assert(cuda::std::is_same_v, int>); + static_assert(cuda::__iterator_can_rebase); + static_assert(cuda::__iterator_can_rebase>); + static_assert( + cuda::__iterator_can_rebase>>); + } + + { + using Iter = cuda::transform_output_iterator; + test_trait(); + + char output[2] = {}; + char staged_output[2] = {}; + Iter iter{output, Narrow{}}; + auto unwrapped = cuda::__unwrap_iterator_adaptor(iter); + static_assert(cuda::std::is_same_v, char*>); + static_assert(cuda::std::is_same_v, char>); + static_assert(cuda::std::is_same_v); + assert(unwrapped == output); + + auto rebased = iter; + rebased.__rebase(staged_output); + *iter = 3; + rebased[1] = 7; + assert(output[0] == 3); + assert(staged_output[1] == 7); + } + + { + using Iter = cuda::transform_input_output_iterator; + test_trait(); + + Iter iter{buffer, PlusOne{}, PlusOne{}}; + auto unwrapped = cuda::__unwrap_iterator_adaptor(iter); + static_assert(cuda::std::is_same_v, int*>); + static_assert(cuda::std::is_same_v, int>); + static_assert(cuda::std::is_same_v); + assert(unwrapped == buffer); + assert(*iter == 2); + + auto rebased = iter; + rebased.__rebase(staged); + assert(*rebased == 11); + *rebased = 5; + assert(buffer[0] == 1); + assert(staged[0] == 6); + assert(rebased.base() == staged); + } + + { + using Inner = cuda::transform_input_output_iterator; + using Outer = cuda::transform_input_output_iterator; + test_trait(); + + staged[0] = 10; + Outer iter{Inner{buffer, PlusOne{}, PlusOne{}}, TimesTwo{}, PlusOne{}}; + auto unwrapped = cuda::__unwrap_iterator_adaptor(iter); + static_assert(cuda::std::is_same_v); + assert(unwrapped == buffer); + assert(*iter == 4); + + auto rebased = iter; + rebased.__rebase(staged); + assert(*iter == 4); + assert(*rebased == 22); + *rebased = 3; + assert(buffer[0] == 1); + assert(staged[0] == 5); + assert(rebased.base().base() == staged); + } + + { + using Iter = cuda::transform_input_output_iterator; + test_trait(); + static_assert(!cuda::__iterator_can_rebase); + } + + { + using Iter = cuda::transform_input_output_iterator; + test_trait(); + static_assert(cuda::__iterator_can_rebase); + } + + { + using Bad = cuda::transform_input_output_iterator>; + test_trait(); + test_trait, false>(); + + Bad iter{cuda::counting_iterator{5}, PlusOne{}, PlusOne{}}; + auto unwrapped = cuda::__unwrap_iterator_adaptor(iter); + static_assert(cuda::std::is_same_v>); + assert(*unwrapped == 5); + } + + return true; +} + +int main(int, char**) +{ + test(); + static_assert(test()); + return 0; +} diff --git a/libcudacxx/test/libcudacxx/libcxx/iterators/has_contiguous_traversal.compile.pass.cpp b/libcudacxx/test/libcudacxx/libcxx/iterators/has_contiguous_traversal.compile.pass.cpp new file mode 100644 index 000000000000..2db656828922 --- /dev/null +++ b/libcudacxx/test/libcudacxx/libcxx/iterators/has_contiguous_traversal.compile.pass.cpp @@ -0,0 +1,59 @@ +//===----------------------------------------------------------------------===// +// +// Part of libcu++, the C++ Standard Library for your entire system, +// under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +// +//===----------------------------------------------------------------------===// + +// __has_contiguous_traversal +// +// Contiguous iterators have contiguous traversal. That includes the iterators of contiguous standard library +// containers, which only advertise it through iterator_concept since C++20. + +#include + +#include "test_iterators.h" +#include "test_macros.h" + +#if !TEST_COMPILER(NVRTC) +# include +# include +# include +# include +# include +# include +#endif // !TEST_COMPILER(NVRTC) + +static_assert(cuda::std::__has_contiguous_traversal); +static_assert(cuda::std::__has_contiguous_traversal); +static_assert(cuda::std::__has_contiguous_traversal>); +static_assert(!cuda::std::__has_contiguous_traversal>); +static_assert(!cuda::std::__has_contiguous_traversal>); +static_assert(!cuda::std::__has_contiguous_traversal>); + +// Types that are not iterators +static_assert(!cuda::std::__has_contiguous_traversal); +static_assert(!cuda::std::__has_contiguous_traversal); + +#if !TEST_COMPILER(NVRTC) +static_assert(cuda::std::__has_contiguous_traversal::iterator>); +static_assert(cuda::std::__has_contiguous_traversal::const_iterator>); +static_assert(cuda::std::__has_contiguous_traversal::iterator>); +static_assert(cuda::std::__has_contiguous_traversal::const_iterator>); +static_assert(cuda::std::__has_contiguous_traversal); +static_assert(cuda::std::__has_contiguous_traversal); +static_assert(cuda::std::__has_contiguous_traversal); + +static_assert(!cuda::std::__has_contiguous_traversal::reverse_iterator>); +static_assert(!cuda::std::__has_contiguous_traversal::iterator>); +static_assert(!cuda::std::__has_contiguous_traversal::iterator>); +static_assert(!cuda::std::__has_contiguous_traversal::iterator>); +#endif // !TEST_COMPILER(NVRTC) + +int main(int, char**) +{ + return 0; +} diff --git a/libcudacxx/test/libcudacxx/libcxx/memory/as_bytes_vectorized.compile.pass.cpp b/libcudacxx/test/libcudacxx/libcxx/memory/as_bytes_vectorized.compile.pass.cpp new file mode 100644 index 000000000000..fe4672a60a3f --- /dev/null +++ b/libcudacxx/test/libcudacxx/libcxx/memory/as_bytes_vectorized.compile.pass.cpp @@ -0,0 +1,306 @@ +//===----------------------------------------------------------------------===// +// +// Part of libcu++, the C++ Standard Library for your entire system, +// under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +// +//===----------------------------------------------------------------------===// + +#include +#include +#include +#include +#include + +#include "test_macros.h" + +struct three_chars_t +{ + char data[3]; +}; + +struct short_and_char_t +{ + short s; + char c; +}; + +struct alignas(32) aligned_32_t +{ + long long data[4]; +}; + +struct alignas(64) aligned_64_t +{ + char data[64]; +}; + +template +TEST_FUNC constexpr void test_vector_word() +{ + using storage_t = cuda::__bytes_vectorized_t; + static_assert(cuda::std::is_same_v); + static_assert(sizeof(storage_t) == sizeof(Word) * Count); + static_assert(alignof(storage_t) == sizeof(Word) * Count); +} + +// Every word of the storage is reachable through a pointer to its first word +template +TEST_FUNC constexpr bool test_flat_iteration() +{ + Storage storage{}; + const auto* const words = storage.data(); + constexpr auto num_words = sizeof(Storage) / sizeof(*words); + return words[num_words - 1] == words[0]; +} + +template +TEST_FUNC constexpr void test_vec_size() +{ + using storage_t = cuda::__as_bytes_vectorized_t; + static_assert( + cuda::std::is_same_v>); + static_assert(test_flat_iteration()); +} + +template +TEST_FUNC constexpr void test_scalar_words() +{ + static_assert(cuda::std::is_same_v, + cuda::__bytes_vectorized_t>); + static_assert( + cuda::std::is_same_v, + cuda::__bytes_vectorized_t>); + static_assert( + cuda::std::is_same_v, + cuda::__bytes_vectorized_t>); +} + +template +TEST_FUNC constexpr void test_1_byte_word() +{ + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); +#if _CCCL_CTK_AT_LEAST(13, 0) + test_vec_size(); + test_vec_size(); +#else // ^^^ _CCCL_CTK_AT_LEAST(13, 0) ^^^ / vvv _CCCL_CTK_BELOW(13, 0) vvv + test_vec_size(); + test_vec_size(); +#endif // ^^^ _CCCL_CTK_BELOW(13, 0) ^^^ +} + +template +TEST_FUNC constexpr void test_2_byte_word() +{ + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); +#if _CCCL_CTK_AT_LEAST(13, 0) + test_vec_size(); +#else // ^^^ _CCCL_CTK_AT_LEAST(13, 0) ^^^ / vvv _CCCL_CTK_BELOW(13, 0) vvv + test_vec_size(); +#endif // ^^^ _CCCL_CTK_BELOW(13, 0) ^^^ +} + +template +TEST_FUNC constexpr void test_4_byte_word() +{ + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); +#if _CCCL_CTK_AT_LEAST(13, 0) + test_vec_size(); + test_vec_size(); +#else // ^^^ _CCCL_CTK_AT_LEAST(13, 0) ^^^ / vvv _CCCL_CTK_BELOW(13, 0) vvv + test_vec_size(); + test_vec_size(); +#endif // ^^^ _CCCL_CTK_BELOW(13, 0) ^^^ +} + +template +TEST_FUNC constexpr void test_8_byte_word() +{ + test_vec_size(); + test_vec_size(); + test_vec_size(); +#if _CCCL_CTK_AT_LEAST(13, 0) + test_vec_size(); + test_vec_size(); + test_vec_size(); +#else // ^^^ _CCCL_CTK_AT_LEAST(13, 0) ^^^ / vvv _CCCL_CTK_BELOW(13, 0) vvv + test_vec_size(); + test_vec_size(); + test_vec_size(); +#endif // ^^^ _CCCL_CTK_BELOW(13, 0) ^^^ +} + +template +TEST_FUNC constexpr void test_16_byte_word() +{ + test_vec_size(); + test_vec_size(); +#if _CCCL_CTK_AT_LEAST(13, 0) + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); +#else // ^^^ _CCCL_CTK_AT_LEAST(13, 0) ^^^ / vvv _CCCL_CTK_BELOW(13, 0) vvv + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); +#endif // ^^^ _CCCL_CTK_BELOW(13, 0) ^^^ +} + +TEST_FUNC void test() +{ + test_vector_word(); + test_vector_word(); + test_vector_word(); + test_vector_word(); + test_vector_word(); +#if _CCCL_CTK_AT_LEAST(13, 0) + test_vector_word(); + test_vector_word(); +#endif // _CCCL_CTK_AT_LEAST(13, 0) + + test_1_byte_word(); + test_1_byte_word(); + test_1_byte_word(); + test_1_byte_word(); + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); + + test_2_byte_word(); + test_2_byte_word(); + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); +#if _CCCL_CTK_AT_LEAST(13, 0) + test_vec_size(); + test_vec_size(); +#else // ^^^ _CCCL_CTK_AT_LEAST(13, 0) ^^^ / vvv _CCCL_CTK_BELOW(13, 0) vvv + test_vec_size(); + test_vec_size(); +#endif // ^^^ _CCCL_CTK_BELOW(13, 0) ^^^ + + test_4_byte_word(); + test_4_byte_word(); + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); +#if _CCCL_CTK_AT_LEAST(13, 0) + test_vec_size(); + test_vec_size(); +#else // ^^^ _CCCL_CTK_AT_LEAST(13, 0) ^^^ / vvv _CCCL_CTK_BELOW(13, 0) vvv + test_vec_size(); + test_vec_size(); +#endif // ^^^ _CCCL_CTK_BELOW(13, 0) ^^^ + test_4_byte_word(); + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); +#if _CCCL_CTK_AT_LEAST(13, 0) + test_vec_size(); + test_vec_size(); +#else // ^^^ _CCCL_CTK_AT_LEAST(13, 0) ^^^ / vvv _CCCL_CTK_BELOW(13, 0) vvv + test_vec_size(); + test_vec_size(); +#endif // ^^^ _CCCL_CTK_BELOW(13, 0) ^^^ + + test_8_byte_word(); + test_8_byte_word(); + test_vec_size(); + test_vec_size(); + test_vec_size(); +#if _CCCL_CTK_AT_LEAST(13, 0) + test_vec_size(); + test_vec_size(); + test_vec_size(); +#else // ^^^ _CCCL_CTK_AT_LEAST(13, 0) ^^^ / vvv _CCCL_CTK_BELOW(13, 0) vvv + test_vec_size(); + test_vec_size(); + test_vec_size(); +#endif // ^^^ _CCCL_CTK_BELOW(13, 0) ^^^ + test_8_byte_word(); + test_8_byte_word(); + test_8_byte_word, cuda::std::uint64_t>(); + test_vec_size(); + test_vec_size(); + test_vec_size(); +#if _CCCL_CTK_AT_LEAST(13, 0) + test_vec_size(); + test_vec_size(); + test_vec_size(); +#else // ^^^ _CCCL_CTK_AT_LEAST(13, 0) ^^^ / vvv _CCCL_CTK_BELOW(13, 0) vvv + test_vec_size(); + test_vec_size(); + test_vec_size(); +#endif // ^^^ _CCCL_CTK_BELOW(13, 0) ^^^ + +#if _CCCL_HAS_INT128() + test_16_byte_word<__int128_t>(); + test_16_byte_word<__uint128_t>(); +#endif // _CCCL_HAS_INT128() + test_16_byte_word(); + test_16_byte_word(); + test_16_byte_word>(); + +#if _CCCL_CTK_AT_LEAST(13, 0) + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); +#else // ^^^ _CCCL_CTK_AT_LEAST(13, 0) ^^^ / vvv _CCCL_CTK_BELOW(13, 0) vvv + test_scalar_words(); +#endif // ^^^ _CCCL_CTK_BELOW(13, 0) ^^^ + + test_scalar_words(); + + // A lower vector access limit splits the storage into narrower vectors, which are less aligned + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); + test_vec_size(); + static_assert(alignof(cuda::__as_bytes_vectorized_t) == 16); + + // Types aligned to more than the vector access limit keep their alignment and are stored in scalar words + test_scalar_words(); + static_assert(alignof(cuda::__as_bytes_vectorized_t) == 32); +} + +int main(int, char**) +{ + test(); + return 0; +} diff --git a/thrust/testing/contiguous_iterator_adaptor.cu b/thrust/testing/contiguous_iterator_adaptor.cu new file mode 100644 index 000000000000..c1deef4b12d9 --- /dev/null +++ b/thrust/testing/contiguous_iterator_adaptor.cu @@ -0,0 +1,233 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include + +struct ContiguousAdaptorPlusOne +{ + _CCCL_HOST_DEVICE int operator()(int value) const + { + return value + 1; + } +}; + +TEST_CASE("contiguous iterator adaptor thrust storage", "[iterators]") +{ + static_assert(cuda::std::__has_contiguous_traversal>); + static_assert(cuda::__is_contiguous_iterator_or_adaptor_v>); + static_assert(cuda::std::__has_contiguous_traversal>); + static_assert(cuda::__is_contiguous_iterator_or_adaptor_v>); + static_assert(cuda::std::__has_contiguous_traversal>>); + static_assert(cuda::std::__has_contiguous_traversal>); + static_assert(cuda::__is_contiguous_iterator_or_adaptor_v>); + static_assert(cuda::__is_contiguous_iterator_or_adaptor_v::iterator>); + static_assert(cuda::__is_contiguous_iterator_or_adaptor_v::iterator>); + + using Transformed = thrust::transform_iterator>; + using Nested = thrust::transform_iterator; + using Normal = thrust::detail::normal_iterator; + using TransformedNormal = thrust::transform_iterator; + static_assert(cuda::__is_contiguous_iterator_or_adaptor_v); + static_assert(cuda::__is_contiguous_iterator_or_adaptor_v + == cuda::__is_contiguous_iterator_or_adaptor_v>); + static_assert(cuda::__is_contiguous_iterator_or_adaptor_v); + static_assert(cuda::__is_contiguous_iterator_or_adaptor_v); + static_assert(!cuda::__is_contiguous_iterator_or_adaptor_v< + thrust::transform_iterator>>); + // A reverse iterator unwraps to a pointer, but does not traverse the storage forward + static_assert(!cuda::__is_contiguous_iterator_or_adaptor_v< + thrust::transform_iterator>>); + static_assert( + cuda::__is_contiguous_iterator_or_adaptor_v< + thrust::transform_iterator>>); + static_assert( + cuda::__is_contiguous_iterator_or_adaptor_v>); + + using TransformedOutput = thrust::transform_output_iterator; + using TransformedInOut = + thrust::transform_input_output_iterator; + using NestedOutput = thrust::transform_output_iterator; + using NestedInOut = + thrust::transform_input_output_iterator; + // Output adaptors are not contiguous adaptors, even over contiguous storage + static_assert(!cuda::__is_contiguous_iterator_or_adaptor_v); + static_assert(!cuda::__is_contiguous_iterator_or_adaptor_v); + static_assert(!cuda::__is_contiguous_iterator_or_adaptor_v); + static_assert(!cuda::__is_contiguous_iterator_or_adaptor_v); + static_assert(!cuda::__is_contiguous_iterator_or_adaptor_v); + static_assert(!cuda::__is_contiguous_iterator_or_adaptor_v); + static_assert(!cuda::__is_contiguous_iterator_or_adaptor_v< + thrust::transform_output_iterator>>); + static_assert(!cuda::__is_contiguous_iterator_or_adaptor_v< + thrust::transform_input_output_iterator>>); + + // The base of a counting iterator is its counter, so there is no storage to rebase + using Counting = thrust::counting_iterator; + static_assert(cuda::std::is_same_v, int>); + static_assert(!cuda::__iterator_can_rebase); + static_assert(!cuda::__iterator_can_rebase>); + static_assert( + cuda::std::is_same_v())), int*>); + static_assert( + cuda::std::is_same_v())), int*>); + static_assert( + cuda::std::is_same_v())), int*>); + static_assert( + cuda::std::is_same_v())), int*>); + + static_assert( + cuda::std::is_same_v>())), + int*>); + static_assert(cuda::std::is_same_v>())), + int*>); + static_assert(cuda::std::is_same_v())), int*>); + static_assert( + cuda::std::is_same_v>>())), + int*>); + static_assert( + cuda::std::is_same_v())), int*>); + static_assert(cuda::std::is_same_v())), int*>); + static_assert( + cuda::std::is_same_v())), int*>); + + static_assert(cuda::__iterator_can_rebase); + static_assert(cuda::__iterator_can_rebase); + static_assert(cuda::__iterator_can_rebase); + static_assert(cuda::__iterator_can_rebase); + + int input[] = {1, 2, 3, 4}; + int staged[] = {10, 20, 30, 40}; + + { + auto iter = thrust::make_transform_iterator(input, ContiguousAdaptorPlusOne{}); + auto unwrapped = cuda::__unwrap_iterator_adaptor(iter); + REQUIRE(unwrapped == input); + + auto rebased = iter; + rebased.__rebase(staged); + REQUIRE(*iter == 2); + REQUIRE(*rebased == 11); + REQUIRE(rebased.base() == staged); + } + + { + auto inner = thrust::make_transform_iterator(input, ContiguousAdaptorPlusOne{}); + auto outer = thrust::make_transform_iterator(inner, ContiguousAdaptorPlusOne{}); + REQUIRE(cuda::__unwrap_iterator_adaptor(outer) == input); + auto rebased = outer; + rebased.__rebase(staged); + REQUIRE(*outer == 3); + REQUIRE(*rebased == 12); + } + + { + Normal iter{input}; + auto rebased = iter; + rebased.__rebase(staged); + REQUIRE(iter.base() == input); + REQUIRE(rebased.base() == staged); + REQUIRE(*rebased == 10); + + auto fancy = thrust::make_transform_iterator(iter, ContiguousAdaptorPlusOne{}); + REQUIRE(cuda::__unwrap_iterator_adaptor(fancy) == input); + auto fancy_rebased = fancy; + fancy_rebased.__rebase(staged); + REQUIRE(*fancy == 2); + REQUIRE(*fancy_rebased == 11); + } + + { + thrust::pointer stored{input}; + auto stored_rebased = stored; + stored_rebased.__rebase(staged); + REQUIRE(stored.get() == input); + REQUIRE(stored_rebased.get() == staged); + + thrust::device_ptr ptr{input}; + auto ptr_rebased = ptr; + ptr_rebased.__rebase(staged); + REQUIRE(ptr.get() == input); + REQUIRE(ptr_rebased.get() == staged); + + auto fancy = thrust::make_transform_iterator(ptr, ContiguousAdaptorPlusOne{}); + REQUIRE(cuda::__unwrap_iterator_adaptor(fancy) == input); + auto fancy_rebased = fancy; + fancy_rebased.__rebase(staged); + REQUIRE(fancy.base().get() == input); + REQUIRE(fancy_rebased.base().get() == staged); + } + + { + int alternative = 7; + thrust::device_vector values(1); + auto rebased = values.begin(); + rebased.__rebase(&alternative); + REQUIRE(rebased.base().get() == &alternative); + + auto fancy = thrust::make_transform_iterator(values.begin(), ContiguousAdaptorPlusOne{}); + REQUIRE(cuda::__unwrap_iterator_adaptor(fancy) == values.begin().base().get()); + auto fancy_rebased = fancy; + fancy_rebased.__rebase(&alternative); + REQUIRE(fancy_rebased.base().base().get() == &alternative); + } + + { + auto iter = thrust::make_transform_output_iterator(input, ContiguousAdaptorPlusOne{}); + REQUIRE(cuda::__unwrap_iterator_adaptor(iter) == input); + auto rebased = iter; + rebased.__rebase(staged); + *rebased = 5; + REQUIRE(input[0] == 1); + REQUIRE(staged[0] == 6); + REQUIRE(rebased.base() == staged); + + auto nested = thrust::make_transform_output_iterator( + thrust::make_transform_iterator(input, ContiguousAdaptorPlusOne{}), ContiguousAdaptorPlusOne{}); + REQUIRE(cuda::__unwrap_iterator_adaptor(nested) == input); + auto nested_rebased = nested; + nested_rebased.__rebase(staged); + REQUIRE(nested_rebased.base().base() == staged); + } + + { + staged[0] = 10; + auto iter = + thrust::make_transform_input_output_iterator(input, ContiguousAdaptorPlusOne{}, ContiguousAdaptorPlusOne{}); + REQUIRE(cuda::__unwrap_iterator_adaptor(iter) == input); + REQUIRE(*iter == 2); + auto rebased = iter; + rebased.__rebase(staged); + REQUIRE(*rebased == 11); + *rebased = 7; + REQUIRE(input[0] == 1); + REQUIRE(staged[0] == 8); + REQUIRE(rebased.base() == staged); + + staged[0] = 10; + Normal stored{input}; + auto nested = + thrust::make_transform_input_output_iterator(stored, ContiguousAdaptorPlusOne{}, ContiguousAdaptorPlusOne{}); + REQUIRE(cuda::__unwrap_iterator_adaptor(nested) == input); + auto nested_rebased = nested; + nested_rebased.__rebase(staged); + REQUIRE(*nested_rebased == 11); + REQUIRE(nested_rebased.base().base() == staged); + } +} diff --git a/thrust/testing/is_contiguous_iterator.cu b/thrust/testing/is_contiguous_iterator.cu index 1e0a95f8909f..6c22686a2105 100644 --- a/thrust/testing/is_contiguous_iterator.cu +++ b/thrust/testing/is_contiguous_iterator.cu @@ -5,7 +5,9 @@ #include #include +#include #include +#include #include #include @@ -32,6 +34,7 @@ void check_is_contiguous() if constexpr (Expected) { STATIC_REQUIRE(cuda::std::__can_to_address == Expected); + STATIC_REQUIRE(cuda::std::__has_contiguous_traversal>); } } @@ -98,6 +101,7 @@ void check_iterator_unwrapping() { STATIC_REQUIRE(std::is_same_v); STATIC_REQUIRE(cuda::std::__can_to_address); + STATIC_REQUIRE(cuda::std::__has_contiguous_traversal); } else { diff --git a/thrust/testing/trivial_sequence.cu b/thrust/testing/trivial_sequence.cu index cffec28c4451..3b517379a282 100644 --- a/thrust/testing/trivial_sequence.cu +++ b/thrust/testing/trivial_sequence.cu @@ -2,6 +2,8 @@ #include #include +#include + #include template @@ -28,6 +30,7 @@ void test(Iterator first, Iterator last) REQUIRE_FALSE(thrust::is_contiguous_iterator_v); REQUIRE(thrust::is_contiguous_iterator_v); + REQUIRE(cuda::std::__has_contiguous_traversal); } template diff --git a/thrust/testing/type_traits.cu b/thrust/testing/type_traits.cu index 5084f4403b35..52c02f9f8903 100644 --- a/thrust/testing/type_traits.cu +++ b/thrust/testing/type_traits.cu @@ -12,6 +12,7 @@ #include #include +#include #include #include @@ -23,13 +24,19 @@ TEST_CASE("TestIsContiguousIterator", "[type_traits]") using DeviceVector = thrust::device_vector; REQUIRE(thrust::is_contiguous_iterator_v); + REQUIRE(cuda::std::__has_contiguous_traversal); REQUIRE(thrust::is_contiguous_iterator_v>); + REQUIRE(cuda::std::__has_contiguous_traversal>); REQUIRE(thrust::is_contiguous_iterator_v); + REQUIRE(cuda::std::__has_contiguous_traversal); REQUIRE(thrust::is_contiguous_iterator_v); + REQUIRE(cuda::std::__has_contiguous_traversal); REQUIRE(thrust::is_contiguous_iterator_v); + REQUIRE(cuda::std::__has_contiguous_traversal); REQUIRE(thrust::is_contiguous_iterator_v); + REQUIRE(cuda::std::__has_contiguous_traversal); REQUIRE(thrust::is_contiguous_iterator_v>); diff --git a/thrust/thrust/detail/pointer.h b/thrust/thrust/detail/pointer.h index b5ee38779da5..b0d3f56f09e2 100644 --- a/thrust/thrust/detail/pointer.h +++ b/thrust/thrust/detail/pointer.h @@ -158,11 +158,12 @@ class pointer : public detail::pointer_base::t } } - // don't provide access to this part of super_t's interface - using super_t::base; + // base_type stays private. Callers use raw_pointer. using typename super_t::base_type; public: + //! @brief The stored raw pointer. Public so contiguous unwrapping can reach it. + using super_t::base; /*! The type of the raw pointer */ using raw_pointer = typename super_t::base_type; @@ -391,6 +392,14 @@ class pointer : public detail::pointer_base::t !detail::ptr_can_compare_less_than>, bool> operator<=(pointer const& lhs, pointer const& rhs) = delete; + + //! @brief Retargets this pointer at @p __ptr. + //! @param[in] __ptr The address that replaces the stored pointer. + template + _CCCL_HOST_DEVICE void __rebase(_Tp* __ptr) & + { + this->base_reference() = static_cast(__ptr); + } }; /*! \} // memory_management @@ -431,4 +440,8 @@ struct pointer_traits> } }; +template +inline constexpr bool __has_contiguous_traversal> = + true; + _CCCL_END_NAMESPACE_CUDA_STD diff --git a/thrust/thrust/device_ptr.h b/thrust/thrust/device_ptr.h index 63d1be41599f..502bd7a7f82a 100644 --- a/thrust/thrust/device_ptr.h +++ b/thrust/thrust/device_ptr.h @@ -225,3 +225,8 @@ THRUST_NAMESPACE_END #include #include + +_CCCL_BEGIN_NAMESPACE_CUDA_STD +template +inline constexpr bool __has_contiguous_traversal> = true; +_CCCL_END_NAMESPACE_CUDA_STD diff --git a/thrust/thrust/iterator/detail/normal_iterator.h b/thrust/thrust/iterator/detail/normal_iterator.h index f9328aa7a5e2..e834fa8cba35 100644 --- a/thrust/thrust/iterator/detail/normal_iterator.h +++ b/thrust/thrust/iterator/detail/normal_iterator.h @@ -21,6 +21,7 @@ #include #include +#include #include #include @@ -44,6 +45,15 @@ class normal_iterator : public iterator_adaptor, Pointe normal_iterator(const normal_iterator& other, enable_if_convertible_t* = nullptr) : super_t(other.base()) {} + + //! @brief Retargets the stored pointer at @p __ptr. + //! @param[in] __ptr The address that replaces the contiguous storage. + _CCCL_TEMPLATE(class _Tp) + _CCCL_REQUIRES(::cuda::__iterator_can_rebase || ::cuda::std::is_pointer_v) + _CCCL_HOST_DEVICE void __rebase(_Tp* __ptr) & + { + ::cuda::__rebase_contiguous_iterator_adaptor(this->base_reference(), __ptr); + } }; template @@ -87,4 +97,8 @@ struct pointer_traits> } }; +template +inline constexpr bool __has_contiguous_traversal> = + __has_contiguous_traversal<_Pointer>; + _CCCL_END_NAMESPACE_CUDA_STD diff --git a/thrust/thrust/iterator/transform_input_output_iterator.h b/thrust/thrust/iterator/transform_input_output_iterator.h index 808d36099c3a..e2a3ba02fe65 100644 --- a/thrust/thrust/iterator/transform_input_output_iterator.h +++ b/thrust/thrust/iterator/transform_input_output_iterator.h @@ -20,6 +20,7 @@ #include #include +#include #include THRUST_NAMESPACE_BEGIN @@ -185,6 +186,16 @@ class transform_input_output_iterator , output_function(output_function) {} + //! @brief Retargets the stored iterator at @p __ptr. The stored functors are left unchanged. The input functor is + //! then invoked on the elements at @p __ptr. + //! @param[in] __ptr The address that replaces the contiguous storage this adaptor reads from and writes to. + _CCCL_TEMPLATE(class _Tp) + _CCCL_REQUIRES(::cuda::__iterator_can_rebase || ::cuda::std::is_pointer_v) + _CCCL_HOST_DEVICE void __rebase(_Tp* __ptr) & + { + ::cuda::__rebase_contiguous_iterator_adaptor(this->base_reference(), __ptr); + } + //! \cond private: diff --git a/thrust/thrust/iterator/transform_iterator.h b/thrust/thrust/iterator/transform_iterator.h index 685bc84c1ca6..f5a362a25c97 100644 --- a/thrust/thrust/iterator/transform_iterator.h +++ b/thrust/thrust/iterator/transform_iterator.h @@ -34,6 +34,7 @@ #include #include +#include #include #include #include @@ -284,6 +285,16 @@ class transform_iterator return m_f; } + //! @brief Retargets the stored iterator at @p __ptr. The stored functor is left unchanged. It is then invoked on the + //! elements at @p __ptr. + //! @param[in] __ptr The address that replaces the contiguous storage this adaptor reads from. + _CCCL_TEMPLATE(class _Tp) + _CCCL_REQUIRES(::cuda::__iterator_can_rebase || ::cuda::std::is_pointer_v) + _CCCL_HOST_DEVICE void __rebase(_Tp* __ptr) & + { + ::cuda::__rebase_contiguous_iterator_adaptor(this->base_reference(), __ptr); + } + //! \cond private: @@ -342,3 +353,10 @@ make_transform_iterator(Iterator it, AdaptableUnaryFunction fun) //! \} // end iterators THRUST_NAMESPACE_END + +_CCCL_BEGIN_NAMESPACE_CUDA +template +inline constexpr bool + __is_contiguous_iterator_or_adaptor_impl> = + __is_contiguous_iterator_or_adaptor_impl<::cuda::std::remove_cvref_t<_Iter>>; +_CCCL_END_NAMESPACE_CUDA diff --git a/thrust/thrust/iterator/transform_output_iterator.h b/thrust/thrust/iterator/transform_output_iterator.h index c77a24ccdfad..f60202568bcb 100644 --- a/thrust/thrust/iterator/transform_output_iterator.h +++ b/thrust/thrust/iterator/transform_output_iterator.h @@ -20,6 +20,8 @@ #include #include +#include + THRUST_NAMESPACE_BEGIN template @@ -149,6 +151,15 @@ class transform_output_iterator : public detail::transform_output_iterator_base< , fun(fun) {} + //! @brief Retargets the stored iterator at @p __ptr. The stored functor is left unchanged. + //! @param[in] __ptr The address that replaces the contiguous storage this adaptor writes to. + _CCCL_TEMPLATE(class _Tp) + _CCCL_REQUIRES(::cuda::__iterator_can_rebase || ::cuda::std::is_pointer_v) + _CCCL_HOST_DEVICE void __rebase(_Tp* __ptr) & + { + ::cuda::__rebase_contiguous_iterator_adaptor(this->base_reference(), __ptr); + } + //! \cond private: