Skip to content

Refactor iterator usage in thrust/cuda find() #6019

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bernhardmgruber merged 5 commits into from
Oct 15, 2025
Merged

Refactor iterator usage in thrust/cuda find() #6019

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
4be9244
Refactor iterator usage in thrust/cuda find()
bernhardmgruber Sep 25, 2025
ee14c8e
Fixes to satisfy concepts
bernhardmgruber Sep 29, 2025
7a557eb
fix
bernhardmgruber Sep 29, 2025
26c790b
Fix warning
bernhardmgruber Sep 29, 2025
e41a347
Fix warning
bernhardmgruber Sep 30, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion cub/test/catch2_test_device_scan_by_key_large_offsets.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ struct expected_sum_op
sum_within_partial_segment = (index_within_segment * (index_within_segment + 1)) / 2;
}
return index_within_segment == 0
? (IsExclusive ? init_value : init_value + full_segments)
? (IsExclusive ? init_value : static_cast<ItemT>(init_value + full_segments))
: static_cast<ItemT>(sum_within_partial_segment + full_segments) + init_value;
}
};
Expand Down
2 changes: 1 addition & 1 deletion thrust/testing/cuda/mismatch.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,7 @@ public:
++my_count;
}

_CCCL_HOST_DEVICE bool operator==(const Wrapper& input) const
_CCCL_HOST_DEVICE bool operator==(const Wrapper&) const
{
return true;
}
Expand Down
128 changes: 13 additions & 115 deletions thrust/thrust/system/cuda/detail/find.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,11 +39,13 @@
#if _CCCL_HAS_CUDA_COMPILER()
# include <thrust/system/cuda/config.h>

# include <thrust/distance.h>
# include <thrust/iterator/counting_iterator.h>
# include <thrust/iterator/transform_iterator.h>
# include <thrust/system/cuda/detail/execution_policy.h>

# include <cuda/__iterator/counting_iterator.h>
# include <cuda/__iterator/transform_iterator.h>
# include <cuda/__iterator/zip_iterator.h>
# include <cuda/std/__iterator/distance.h>

THRUST_NAMESPACE_BEGIN
namespace cuda_cub
{
Expand All @@ -62,7 +64,6 @@ InputIt _CCCL_HOST_DEVICE find(execution_policy<Derived>& policy, InputIt first,
}; // namespace cuda_cub
THRUST_NAMESPACE_END

# include <thrust/iterator/zip_iterator.h>
# include <thrust/system/cuda/detail/reduce.h>

THRUST_NAMESPACE_BEGIN
Expand Down Expand Up @@ -92,109 +93,13 @@ struct functor
}
}
};

template <class ValueType, class InputIt, class UnaryOp>
struct transform_input_iterator_t
{
using self_t = transform_input_iterator_t;
using difference_type = thrust::detail::it_difference_t<InputIt>;
using value_type = ValueType;
using pointer = void;
using reference = value_type;
using iterator_category = ::cuda::std::random_access_iterator_tag;

InputIt input;
mutable UnaryOp op;

_CCCL_HOST_DEVICE _CCCL_FORCEINLINE transform_input_iterator_t(InputIt input, UnaryOp op)
: input(input)
, op(op)
{}

transform_input_iterator_t(const self_t&) = default;

// UnaryOp might not be copy assignable, such as when it is a lambda. Define
// an explicit copy assignment operator that doesn't try to assign it.
_CCCL_HOST_DEVICE self_t& operator=(const self_t& o)
{
input = o.input;
return *this;
}

_CCCL_HOST_DEVICE _CCCL_FORCEINLINE self_t operator++(int)
{
self_t retval = *this;
++input;
return retval;
}

_CCCL_HOST_DEVICE _CCCL_FORCEINLINE self_t operator++()
{
++input;
return *this;
}

_CCCL_HOST_DEVICE _CCCL_FORCEINLINE reference operator*() const
{
thrust::detail::it_value_t<InputIt> x = *input;
return op(x);
}

_CCCL_HOST_DEVICE _CCCL_FORCEINLINE reference operator*()
{
thrust::detail::it_value_t<InputIt> x = *input;
return op(x);
}

_CCCL_HOST_DEVICE _CCCL_FORCEINLINE self_t operator+(difference_type n) const
{
return self_t(input + n, op);
}

_CCCL_HOST_DEVICE _CCCL_FORCEINLINE self_t& operator+=(difference_type n)
{
input += n;
return *this;
}

_CCCL_HOST_DEVICE _CCCL_FORCEINLINE self_t operator-(difference_type n) const
{
return self_t(input - n, op);
}

_CCCL_HOST_DEVICE _CCCL_FORCEINLINE self_t& operator-=(difference_type n)
{
input -= n;
return *this;
}

_CCCL_HOST_DEVICE _CCCL_FORCEINLINE difference_type operator-(self_t other) const
{
return input - other.input;
}

_CCCL_HOST_DEVICE _CCCL_FORCEINLINE reference operator[](difference_type n) const
{
return op(input[n]);
}

_CCCL_HOST_DEVICE _CCCL_FORCEINLINE bool operator==(const self_t& rhs) const
{
return (input == rhs.input);
}

_CCCL_HOST_DEVICE _CCCL_FORCEINLINE bool operator!=(const self_t& rhs) const
{
return (input != rhs.input);
}
};
} // namespace __find_if

template <class Derived, class InputIt, class Size, class Predicate>
InputIt _CCCL_HOST_DEVICE
find_if_n(execution_policy<Derived>& policy, InputIt first, Size num_items, Predicate predicate)
{
using result_type = typename thrust::tuple<bool, Size>;
using result_type = ::cuda::std::tuple<bool, Size>;

// empty sequence
if (num_items == 0)
Expand All @@ -212,27 +117,20 @@ find_if_n(execution_policy<Derived>& policy, InputIt first, Size num_items, Pred
const Size interval_threshold = 1 << 20;
const Size interval_size = (::cuda::std::min) (interval_threshold, num_items);

// FIXME(bgruber): we should also be able to use transform_iterator here, but it makes nvc++ hang. See:
// https://github.com/NVIDIA/cccl/issues/3594. The problem does not occur with nvcc, so we could not add a test :/
Comment on lines -215 to -216

@bernhardmgruber bernhardmgruber Sep 25, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nvc++ is fine with this code

using XfrmIterator = __find_if::transform_input_iterator_t<bool, InputIt, Predicate>;
// using XfrmIterator = transform_iterator<Predicate, InputIt>;
using IteratorTuple = thrust::tuple<XfrmIterator, counting_iterator<Size>>;
using ZipIterator = thrust::zip_iterator<IteratorTuple>;

IteratorTuple iter_tuple = thrust::make_tuple(XfrmIterator(first, predicate), counting_iterator<Size>(0));

ZipIterator begin = thrust::make_zip_iterator(iter_tuple);
ZipIterator end = begin + num_items;
const auto begin = ::cuda::make_zip_iterator(
::cuda::make_transform_iterator(try_unwrap_contiguous_iterator(first), predicate),
::cuda::counting_iterator<Size>(0));
const auto end = begin + num_items;

for (ZipIterator interval_begin = begin; interval_begin < end; interval_begin += interval_size)
for (auto interval_begin = begin; interval_begin < end; interval_begin += interval_size)
{
ZipIterator interval_end = interval_begin + interval_size;
auto interval_end = interval_begin + interval_size;
if (end < interval_end)
{
interval_end = end;
} // end if

result_type result = reduce(
const result_type result = reduce(
policy, interval_begin, interval_end, result_type(false, interval_end - begin), __find_if::functor<result_type>());

@bernhardmgruber bernhardmgruber Sep 25, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I get an error here at interval_end - begin because zip_iterator::operator- is disabled because _Constraints::__all_sized_sentinel fails because the transform iterator above is not default constructible (the predicate isn't default constructible). The latter is fine, so I wonder why the constraint on zip_iterator::operator- is needed.

@miscco

@miscco miscco Sep 25, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is from ::std::ranges::zip_view

It kind of also boils down to the iterator concepts, because we want a forward iterator to be regular which requires default constructibility

@bernhardmgruber bernhardmgruber Sep 25, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's a problem now, because if we replace any thrust zip iterator by a cuda one, it can be a source breaking change for a user.

@bernhardmgruber bernhardmgruber Sep 25, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Discussed with @miscco offline: he will relax the constraint on operator-

@miscco miscco Sep 26, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Put this in #6036 because we had to fix zip_iterator for NVHPC anyway


// see if we found something
Expand Down
10 changes: 8 additions & 2 deletions thrust/thrust/system/cuda/detail/mismatch.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,12 +73,14 @@ struct transform_pair_of_input_iterators_t
using value_type = ValueType;
using pointer = void;
using reference = value_type;
using iterator_category = std::random_access_iterator_tag;
using iterator_category = ::cuda::std::random_access_iterator_tag;

InputIt1 input1;
InputIt2 input2;
mutable BinaryOp op;

transform_pair_of_input_iterators_t() = default;

_CCCL_HOST_DEVICE _CCCL_FORCEINLINE
transform_pair_of_input_iterators_t(InputIt1 input1_, InputIt2 input2_, BinaryOp op_)
: input1(input1_)
Expand Down Expand Up @@ -107,7 +109,7 @@ struct transform_pair_of_input_iterators_t
}

/// Prefix increment
_CCCL_HOST_DEVICE _CCCL_FORCEINLINE self_t operator++()
_CCCL_HOST_DEVICE _CCCL_FORCEINLINE self_t& operator++()
{
++input1;
++input2;
Expand Down Expand Up @@ -177,6 +179,10 @@ struct transform_pair_of_input_iterators_t
return (input1 != rhs.input1) || (input2 != rhs.input2);
}

_CCCL_HOST_DEVICE _CCCL_FORCEINLINE bool operator<(const self_t& rhs) const
{
return input1 < rhs.input1;
}
}; // struct transform_pair_of_input_iterators_t
} // namespace detail

Expand Down