Skip to content

Port thrust::transform_iterator to cuda #4718

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.

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
merged 6 commits into from
May 21, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
456 changes: 456 additions & 0 deletions libcudacxx/include/cuda/__iterator/transform_iterator.h

Large diffs are not rendered by default.

1 change: 1 addition & 0 deletions libcudacxx/include/cuda/iterator
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
#endif // no system header

#include <cuda/__iterator/discard_iterator.h>
#include <cuda/__iterator/transform_iterator.h>
#include <cuda/std/iterator>

#endif // _CUDA_ITERATOR
Original file line number Diff line number Diff line change
@@ -0,0 +1,91 @@
//===----------------------------------------------------------------------===//
//
// 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) 2025 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//

// transform_iterator::operator{++,--,+=,-=}

#include <cuda/iterator>
#include <cuda/std/cassert>
#include <cuda/std/utility>

#include "test_iterators.h"
#include "test_macros.h"
#include "types.h"

template <class Iter>
_CCCL_CONCEPT can_decrement = _CCCL_REQUIRES_EXPR((Iter), Iter i)((--i));
template <class Iter>
_CCCL_CONCEPT can_post_decrement = _CCCL_REQUIRES_EXPR((Iter), Iter i)((i--));

template <class Iter>
_CCCL_CONCEPT can_plus_equal = _CCCL_REQUIRES_EXPR((Iter), Iter i)((i += 1));
template <class Iter>
_CCCL_CONCEPT can_minus_equal = _CCCL_REQUIRES_EXPR((Iter), Iter i)((i -= 1));

template <class Iter>
__host__ __device__ constexpr void test()
{
int buffer[8] = {0, 1, 2, 3, 4, 5, 6, 7};

cuda::transform_iterator iter{Iter{buffer}, PlusOne{}};
assert((++iter).base() == Iter{buffer + 1});

if constexpr (cuda::std::forward_iterator<Iter>)
{
assert((iter++).base() == Iter{buffer + 1});
}
else
{
iter++;
static_assert(cuda::std::is_same_v<decltype(iter++), void>);
}
assert(iter.base() == Iter{buffer + 2});

if constexpr (cuda::std::bidirectional_iterator<Iter>)
{
assert((--iter).base() == Iter{buffer + 1});
assert((iter--).base() == Iter{buffer + 1});
assert(iter.base() == Iter{buffer});
}
else
{
static_assert(!can_decrement<Iter>);
static_assert(!can_post_decrement<Iter>);
}

if constexpr (cuda::std::random_access_iterator<Iter>)
{
assert((iter += 4).base() == Iter{buffer + 4});
assert((iter -= 3).base() == Iter{buffer + 1});
}
else
{
static_assert(!can_plus_equal<Iter>);
static_assert(!can_minus_equal<Iter>);
}
}

__host__ __device__ constexpr bool test()
{
test<cpp17_input_iterator<int*>>();
test<forward_iterator<int*>>();
test<bidirectional_iterator<int*>>();
test<random_access_iterator<int*>>();
test<int*>();

return true;
}

int main(int, char**)
{
test();
static_assert(test(), "");

return 0;
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
//===----------------------------------------------------------------------===//
//
// 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) 2025 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//

// transform_iterator::base

#include <cuda/iterator>
#include <cuda/std/cassert>
#include <cuda/std/type_traits>
#include <cuda/std/utility>

#include "test_iterators.h"
#include "test_macros.h"
#include "types.h"

template <class Iter>
__host__ __device__ constexpr void test()
{
int buffer[8] = {0, 1, 2, 3, 4, 5, 6, 7};

cuda::transform_iterator iter{Iter{buffer}, PlusOne{}};
using transform_iter = decltype(iter);
static_assert(cuda::std::is_same_v<decltype(static_cast<transform_iter&>(iter).base()), Iter const&>);
static_assert(cuda::std::is_same_v<decltype(static_cast<transform_iter&&>(iter).base()), Iter>);
static_assert(cuda::std::is_same_v<decltype(static_cast<const transform_iter&>(iter).base()), Iter const&>);
static_assert(cuda::std::is_same_v<decltype(static_cast<const transform_iter&&>(iter).base()), Iter const&>);
static_assert(noexcept(iter.base()));
static_assert(
noexcept(static_cast<transform_iter&&>(iter).base()) == cuda::std::is_nothrow_move_constructible_v<Iter>);
assert(base(iter.base()) == buffer);
assert(base(cuda::std::move(iter).base()) == buffer);
}

__host__ __device__ constexpr bool test()
{
test<cpp17_input_iterator<int*>>();
test<random_access_iterator<int*>>();
test<int*>();

return true;
}

int main(int, char**)
{
test();
static_assert(test(), "");

return 0;
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,82 @@
//===----------------------------------------------------------------------===//
//
// 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) 2025 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//

// transform_iterator::operator{<,>,<=,>=,==,!=,<=>}

#include <cuda/iterator>
#if _LIBCUDACXX_HAS_SPACESHIP_OPERATOR()
# include <cuda/std/compare>
#endif // _LIBCUDACXX_HAS_SPACESHIP_OPERATOR()

#include "test_iterators.h"
#include "test_macros.h"
#include "types.h"

template <class Iter>
__host__ __device__ constexpr void test()
{
int buffer[8] = {0, 1, 2, 3, 4, 5, 6, 7};

cuda::transform_iterator iter1{Iter{buffer}, PlusOne{}};
cuda::transform_iterator iter2{Iter{buffer + 4}, PlusOne{}};

assert(!(iter1 < iter1));
assert(iter1 < iter2);
assert(!(iter2 < iter1));
assert(iter1 <= iter1);
assert(iter1 <= iter2);
assert(!(iter2 <= iter1));
assert(!(iter1 > iter1));
assert(!(iter1 > iter2));
assert(iter2 > iter1);
assert(iter1 >= iter1);
assert(!(iter1 >= iter2));
assert(iter2 >= iter1);
assert(iter1 == iter1);
assert(!(iter1 == iter2));
assert(iter2 == iter2);
assert(!(iter1 != iter1));
assert(iter1 != iter2);
assert(!(iter2 != iter2));

#if TEST_HAS_SPACESHIP()
// Test a new-school iterator with operator<=>; the transform iterator should also have operator<=>.
if constexpr (cuda::std::is_same_v<Iter, three_way_contiguous_iterator<int*>>)
{
static_assert(cuda::std::three_way_comparable<Iter>);
static_assert(cuda::std::three_way_comparable<decltype(iter1)>);

assert((iter1 <=> iter2) == cuda::std::strong_ordering::less);
assert((iter1 <=> iter1) == cuda::std::strong_ordering::equal);
assert((iter2 <=> iter1) == cuda::std::strong_ordering::greater);
}
#endif // TEST_HAS_SPACESHIP()
}

__host__ __device__ constexpr bool test()
{
test<random_access_iterator<int*>>();
test<contiguous_iterator<int*>>();
test<int*>();

#if TEST_HAS_SPACESHIP()
test<three_way_contiguous_iterator<int*>>();
#endif // TEST_HAS_SPACESHIP()

return true;
}

int main(int, char**)
{
test();
static_assert(test(), "");

return 0;
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,126 @@
//===----------------------------------------------------------------------===//
//
// 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) 2025 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//

// transform_iterator::transform_iterator();

#include <cuda/iterator>
#include <cuda/std/cassert>
#include <cuda/std/concepts>

#include "test_iterators.h"
#include "test_macros.h"
#include "types.h"

struct NoDefaultInitIter
{
int* ptr_;
using iterator_category = cuda::std::random_access_iterator_tag;
using value_type = int;
using difference_type = cuda::std::ptrdiff_t;
using pointer = int*;
using reference = int&;
using self = NoDefaultInitIter;

__host__ __device__ constexpr NoDefaultInitIter(int* ptr)
: ptr_(ptr)
{}

__host__ __device__ constexpr reference operator*() const;
__host__ __device__ constexpr pointer operator->() const;
#if TEST_HAS_SPACESHIP()
__host__ __device__ constexpr auto operator<=>(const self&) const = default;
#else // ^^^ TEST_HAS_SPACESHIP() ^^^ / vvv !TEST_HAS_SPACESHIP() vvv
__host__ __device__ constexpr bool operator<(const self&) const;
__host__ __device__ constexpr bool operator<=(const self&) const;
__host__ __device__ constexpr bool operator>(const self&) const;
__host__ __device__ constexpr bool operator>=(const self&) const;
#endif // !TEST_HAS_SPACESHIP()

__host__ __device__ constexpr friend bool operator==(const self& lhs, const self& rhs)
{
return lhs.ptr_ == rhs.ptr_;
}
#if TEST_STD_VER <= 2017
__host__ __device__ constexpr friend bool operator!=(const self& lhs, const self& rhs)
{
return lhs.ptr_ != rhs.ptr_;
}
#endif // TEST_STD_VER <= 2017

__host__ __device__ constexpr self& operator++();
__host__ __device__ constexpr self operator++(int);

__host__ __device__ constexpr self& operator--();
__host__ __device__ constexpr self operator--(int);

__host__ __device__ constexpr self& operator+=(difference_type n);
__host__ __device__ constexpr self operator+(difference_type n) const;
__host__ __device__ constexpr friend self operator+(difference_type n, self x);

__host__ __device__ constexpr self& operator-=(difference_type n);
__host__ __device__ constexpr self operator-(difference_type n) const;
__host__ __device__ constexpr difference_type operator-(const self&) const;

__host__ __device__ constexpr reference operator[](difference_type n) const;
};

struct NoDefaultInitFunc
{
int val_;

__host__ __device__ constexpr NoDefaultInitFunc(int val)
: val_(val)
{}

__host__ __device__ constexpr int operator()(int x) const
{
return x * val_;
}
};

template <class Iter, class Fn>
__host__ __device__ constexpr void test(Fn fun)
{
int buffer[8] = {0, 1, 2, 3, 4, 5, 6, 7};

{ // default initialization
constexpr bool can_default_init = cuda::std::default_initializable<Iter> && cuda::std::default_initializable<Fn>;
static_assert(cuda::std::default_initializable<cuda::transform_iterator<Iter, Fn>> == can_default_init);
if constexpr (can_default_init)
{
[[maybe_unused]] cuda::transform_iterator<Iter, Fn> iter{};
}
}

{ // construction from iter and functor
cuda::transform_iterator iter{Iter{buffer}, fun};
assert(iter.base() == Iter{buffer});
}
}

__host__ __device__ constexpr bool test()
{
test<NoDefaultInitIter>(PlusOne{});
test<random_access_iterator<int*>>(PlusOne{});

NoDefaultInitFunc func{42};
test<NoDefaultInitIter>(func);
test<random_access_iterator<int*>>(func);

return true;
}

int main(int, char**)
{
test();
static_assert(test(), "");

return 0;
}
Loading
Loading