Skip to content

Commit d9f3077

Browse files
bernhardmgruberelstehle
bernhardmgruber
and
elstehle
authored
Add thrust::offset_iterator (#4073)
* Add thrust::offset_iterator * Support custom offset types in offset_iterator * Move CUB-using test to CUDA * Remove mutation * Add example loading offset via transform_iterator and extend doc * Fx1# Please enter the commit message for your changes. Lines starting * MSVC workaround * Update after discssion with elstehle * Add select example with offset_iterator * Apply suggestions from code review Co-authored-by: Elias Stehle <3958403+elstehle@users.noreply.github.com> --------- Co-authored-by: Elias Stehle <3958403+elstehle@users.noreply.github.com>
1 parent 2c16c6b commit d9f3077

File tree

4 files changed

+561
-0
lines changed

4 files changed

+561
-0
lines changed

‎cub/test/catch2_test_device_select_if.cu‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,9 @@
3232
#include <cub/device/dispatch/dispatch_select_if.cuh>
3333

3434
#include <thrust/distance.h>
35+
#include <thrust/functional.h>
3536
#include <thrust/iterator/counting_iterator.h>
37+
#include <thrust/iterator/offset_iterator.h>
3638
#include <thrust/iterator/transform_iterator.h>
3739
#include <thrust/logical.h>
3840
#include <thrust/partition.h>
@@ -414,3 +416,29 @@ catch (std::bad_alloc&)
414416
{
415417
// Exceeding memory is not a failure.
416418
}
419+
420+
C2H_TEST("DeviceSelect::If works with iterators", "[device][select_if]")
421+
{
422+
using type = int;
423+
424+
const int num_items = 10'000;
425+
c2h::device_vector<type> in(num_items);
426+
thrust::sequence(in.begin(), in.end());
427+
c2h::device_vector<type> out(num_items);
428+
using thrust::placeholders::_1;
429+
430+
// select twice, appending the second selection to the first one without bringing the first selection's count to the
431+
// host
432+
c2h::device_vector<int> num_selected_out(2);
433+
select_if(in.begin(), out.begin(), num_selected_out.begin(), num_items, _1 < 1000); // [0;999]
434+
auto output_end = thrust::offset_iterator{out.begin(), num_selected_out.begin()};
435+
select_if(in.begin(), output_end, num_selected_out.begin() + 1, num_items, _1 >= 9000); // [9000;9999]
436+
437+
c2h::device_vector<type> expected(2000);
438+
thrust::sequence(expected.begin(), expected.begin() + 1000);
439+
thrust::sequence(expected.begin() + 1000, expected.end(), 9000);
440+
441+
out.resize(2000);
442+
REQUIRE(num_selected_out == c2h::device_vector<int>{1000, 1000});
443+
REQUIRE(out == expected);
444+
}

‎thrust/testing/cuda/offset_iterator.cu‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
#include <thrust/distance.h>
2+
#include <thrust/functional.h>
3+
#include <thrust/iterator/offset_iterator.h>
4+
#include <thrust/iterator/transform_iterator.h>
5+
6+
#include <cuda/std/iterator>
7+
8+
#include <unittest/unittest.h>
9+
10+
struct device_only_iterator
11+
{
12+
using iterator_category = cuda::std::random_access_iterator_tag;
13+
using difference_type = cuda::std::ptrdiff_t;
14+
using value_type = int;
15+
using pointer = int*;
16+
using reference = int&;
17+
18+
_CCCL_HOST_DEVICE device_only_iterator(pointer ptr)
19+
: m_ptr(ptr)
20+
{}
21+
22+
_CCCL_DEVICE reference operator*() const
23+
{
24+
return *m_ptr;
25+
}
26+
27+
_CCCL_DEVICE device_only_iterator& operator++()
28+
{
29+
++m_ptr;
30+
return *this;
31+
}
32+
33+
_CCCL_DEVICE device_only_iterator operator++(int)
34+
{
35+
device_only_iterator tmp = *this;
36+
++*this;
37+
return tmp;
38+
}
39+
40+
_CCCL_DEVICE device_only_iterator& operator--()
41+
{
42+
--m_ptr;
43+
return *this;
44+
}
45+
46+
_CCCL_DEVICE device_only_iterator operator--(int)
47+
{
48+
device_only_iterator tmp = *this;
49+
--*this;
50+
return tmp;
51+
}
52+
53+
_CCCL_DEVICE device_only_iterator& operator+=(difference_type n)
54+
{
55+
m_ptr += n;
56+
return *this;
57+
}
58+
59+
_CCCL_DEVICE friend bool operator-(const device_only_iterator& a, const device_only_iterator& b)
60+
{
61+
return a.m_ptr - b.m_ptr;
62+
}
63+
64+
_CCCL_DEVICE friend bool operator==(const device_only_iterator& a, const device_only_iterator& b)
65+
{
66+
return a.m_ptr == b.m_ptr;
67+
}
68+
69+
_CCCL_DEVICE friend bool operator!=(const device_only_iterator& a, const device_only_iterator& b)
70+
{
71+
return a.m_ptr != b.m_ptr;
72+
}
73+
74+
private:
75+
pointer m_ptr;
76+
};
77+
78+
_CCCL_HOST_DEVICE void TestOffsetIteratorBoth(thrust::offset_iterator<device_only_iterator> iter)
79+
{
80+
assert(iter.offset() == 0);
81+
++iter;
82+
assert(iter.offset() == 1);
83+
iter++;
84+
assert(iter.offset() == 2);
85+
--iter;
86+
assert(iter.offset() == 1);
87+
iter--;
88+
assert(iter.offset() == 0);
89+
iter += 100;
90+
assert(iter.offset() == 100);
91+
}
92+
93+
__global__ void TestOffsetIteratorDevice(thrust::offset_iterator<device_only_iterator> iter)
94+
{
95+
TestOffsetIteratorBoth(iter);
96+
97+
// access
98+
assert(*iter == 1);
99+
100+
auto iter2 = iter;
101+
iter2 += 3;
102+
assert(*iter2 == 1);
103+
104+
// difference
105+
assert(iter2 - iter == 3);
106+
107+
// comparison
108+
assert(!(iter2 == iter));
109+
assert(iter2 != iter);
110+
}
111+
112+
void TestOffsetIteratorWithDeviceOnlyIterator()
113+
{
114+
thrust::device_vector<int> v{1, 2, 3, 4, 5};
115+
device_only_iterator base(thrust::raw_pointer_cast(v.data()));
116+
thrust::offset_iterator iter(base);
117+
TestOffsetIteratorBoth(iter);
118+
TestOffsetIteratorDevice<<<1, 1>>>(iter);
119+
}
120+
DECLARE_UNITTEST(TestOffsetIteratorWithDeviceOnlyIterator);

‎thrust/testing/offset_iterator.cu‎

Lines changed: 221 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,221 @@
1+
2+
#include <thrust/distance.h>
3+
#include <thrust/iterator/offset_iterator.h>
4+
5+
#include <cuda/std/iterator>
6+
7+
#include <unittest/unittest.h>
8+
9+
// ensure that we properly support thrust::counting_iterator from cuda::std
10+
void TestOffsetIteratorTraits()
11+
{
12+
using base_it = thrust::host_vector<int>::iterator;
13+
using it = thrust::offset_iterator<base_it>;
14+
using traits = cuda::std::iterator_traits<it>;
15+
using vec_traits = cuda::std::iterator_traits<base_it>;
16+
17+
static_assert(cuda::std::is_same_v<traits::difference_type, vec_traits::difference_type>);
18+
static_assert(cuda::std::is_same_v<traits::value_type, vec_traits::value_type>);
19+
static_assert(cuda::std::is_same_v<traits::pointer, vec_traits::pointer>);
20+
static_assert(cuda::std::is_same_v<traits::reference, vec_traits::reference>);
21+
static_assert(cuda::std::is_same_v<traits::iterator_category, vec_traits::iterator_category>);
22+
23+
static_assert(cuda::std::is_same_v<thrust::iterator_traversal_t<it>, thrust::random_access_traversal_tag>);
24+
25+
static_assert(cuda::std::__is_cpp17_random_access_iterator<it>::value);
26+
27+
static_assert(cuda::std::output_iterator<it, int>);
28+
static_assert(cuda::std::input_iterator<it>);
29+
static_assert(cuda::std::forward_iterator<it>);
30+
static_assert(cuda::std::bidirectional_iterator<it>);
31+
static_assert(cuda::std::random_access_iterator<it>);
32+
static_assert(!cuda::std::contiguous_iterator<it>);
33+
}
34+
DECLARE_UNITTEST(TestOffsetIteratorTraits);
35+
36+
template <typename Vector>
37+
void TestOffsetConstructor()
38+
{
39+
thrust::offset_iterator<int*> iter0;
40+
ASSERT_EQUAL(iter0.base(), static_cast<int*>(nullptr));
41+
ASSERT_EQUAL(iter0.offset(), 0);
42+
43+
Vector v{42, 43};
44+
thrust::offset_iterator iter1(v.begin());
45+
ASSERT_EQUAL_QUIET(iter1.base(), v.begin());
46+
ASSERT_EQUAL(iter1.offset(), 0);
47+
ASSERT_EQUAL(*iter1, 42);
48+
49+
thrust::offset_iterator iter2(v.begin(), 1);
50+
ASSERT_EQUAL_QUIET(iter2.base(), v.begin());
51+
ASSERT_EQUAL(iter2.offset(), 1);
52+
ASSERT_EQUAL(*iter2, 43);
53+
54+
ptrdiff_t offset = 1;
55+
thrust::offset_iterator iter3(v.begin(), &offset);
56+
ASSERT_EQUAL_QUIET(iter3.base(), v.begin());
57+
ASSERT_EQUAL(iter3.offset(), &offset);
58+
ASSERT_EQUAL(*iter3.offset(), 1);
59+
ASSERT_EQUAL(*iter3, 43);
60+
}
61+
DECLARE_VECTOR_UNITTEST(TestOffsetConstructor);
62+
63+
template <typename Vector>
64+
void TestOffsetIteratorCopyConstructorAndAssignment()
65+
{
66+
Vector v{42, 43};
67+
68+
// value offset
69+
{
70+
thrust::offset_iterator iter0(v.begin());
71+
#if _CCCL_COMPILER(MSVC) // MSVC cannot deduce the template arguments from the copy ctor
72+
decltype(iter0) iter1(iter0);
73+
#else // _CCCL_COMPILER(MSVC)
74+
thrust::offset_iterator iter1(iter0);
75+
#endif // _CCCL_COMPILER(MSVC)
76+
ASSERT_EQUAL(iter0 == iter1, true);
77+
ASSERT_EQUAL(*iter0 == *iter1, true);
78+
79+
thrust::offset_iterator iter2(v.begin() + 1);
80+
ASSERT_EQUAL(iter0 != iter2, true);
81+
ASSERT_EQUAL(*iter0 != *iter2, true);
82+
83+
iter2 = iter0;
84+
ASSERT_EQUAL(iter0 == iter2, true);
85+
ASSERT_EQUAL(*iter0 == *iter2, true);
86+
}
87+
88+
// indirect offset
89+
{
90+
const typename Vector::iterator::difference_type offset = 0;
91+
thrust::offset_iterator iter0(v.begin(), &offset);
92+
93+
#if _CCCL_COMPILER(MSVC) // MSVC cannot deduce the template arguments from the copy ctor
94+
decltype(iter0) iter1(iter0);
95+
#else // _CCCL_COMPILER(MSVC)
96+
thrust::offset_iterator iter1(iter0);
97+
#endif // _CCCL_COMPILER(MSVC)
98+
ASSERT_EQUAL(iter0 == iter1, true);
99+
ASSERT_EQUAL(*iter0 == *iter1, true);
100+
101+
thrust::offset_iterator iter2(v.begin() + 1, &offset);
102+
ASSERT_EQUAL(iter0 != iter2, true);
103+
ASSERT_EQUAL(*iter0 != *iter2, true);
104+
105+
iter2 = iter0;
106+
ASSERT_EQUAL(iter0 == iter2, true);
107+
ASSERT_EQUAL(*iter0 == *iter2, true);
108+
}
109+
}
110+
DECLARE_VECTOR_UNITTEST(TestOffsetIteratorCopyConstructorAndAssignment);
111+
112+
template <typename Vector>
113+
void TestOffsetIteratorIncrement()
114+
{
115+
auto test = [](auto iter) {
116+
ASSERT_EQUAL(*iter, 0);
117+
iter++;
118+
ASSERT_EQUAL(*iter, 1);
119+
iter++;
120+
iter++;
121+
ASSERT_EQUAL(*iter, 3);
122+
iter += 5;
123+
ASSERT_EQUAL(*iter, 8);
124+
iter -= 10;
125+
ASSERT_EQUAL(*iter, -2);
126+
};
127+
128+
const Vector v{-2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8};
129+
test(thrust::offset_iterator(v.begin() + 1, 1));
130+
const typename Vector::iterator::difference_type offset = 1;
131+
test(thrust::offset_iterator(v.begin() + 1, &offset));
132+
}
133+
DECLARE_VECTOR_UNITTEST(TestOffsetIteratorIncrement);
134+
135+
template <typename Vector>
136+
void TestOffsetIteratorMutation()
137+
{
138+
{
139+
Vector v{-2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8};
140+
thrust::offset_iterator it(v.begin() + 1, 1);
141+
*it = 42;
142+
++it;
143+
*it = 43;
144+
++it.offset();
145+
*it = 44;
146+
ASSERT_EQUAL(v, (Vector{-2, -1, 42, 43, 44, 3, 4, 5, 6, 7, 8}));
147+
}
148+
{
149+
Vector v{-2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8};
150+
typename Vector::iterator::difference_type offset = 1;
151+
thrust::offset_iterator it(v.begin() + 1, &offset);
152+
*it = 42;
153+
++it;
154+
*it = 43;
155+
offset = 2;
156+
*it = 44;
157+
ASSERT_EQUAL(v, (Vector{-2, -1, 42, 43, 44, 3, 4, 5, 6, 7, 8}));
158+
}
159+
}
160+
DECLARE_VECTOR_UNITTEST(TestOffsetIteratorMutation);
161+
162+
template <typename Vector>
163+
void TestOffsetIteratorComparisonAndDistance()
164+
{
165+
auto test = [](auto iter1, auto iter2) {
166+
ASSERT_EQUAL(iter1 == iter2, true);
167+
ASSERT_EQUAL(iter1 - iter2, 0);
168+
ASSERT_EQUAL(thrust::distance(iter1, iter2), 0);
169+
170+
iter1++;
171+
ASSERT_EQUAL(iter1 == iter2, false);
172+
ASSERT_EQUAL(iter1 - iter2, 1);
173+
ASSERT_EQUAL(thrust::distance(iter1, iter2), -1);
174+
175+
iter2++;
176+
ASSERT_EQUAL(iter1 == iter2, true);
177+
ASSERT_EQUAL(iter1 - iter2, 0);
178+
ASSERT_EQUAL(thrust::distance(iter1, iter2), 0);
179+
180+
iter1 += 100;
181+
iter2 += 100;
182+
ASSERT_EQUAL(iter1 == iter2, true);
183+
ASSERT_EQUAL(iter1 - iter2, 0);
184+
ASSERT_EQUAL(thrust::distance(iter1, iter2), 0);
185+
186+
iter1 -= 5;
187+
ASSERT_EQUAL(iter1 == iter2, false);
188+
ASSERT_EQUAL(iter1 - iter2, -5);
189+
ASSERT_EQUAL(thrust::distance(iter1, iter2), 5);
190+
};
191+
192+
Vector v(101);
193+
test(thrust::offset_iterator(v.begin()), thrust::offset_iterator(v.begin()));
194+
const typename Vector::iterator::difference_type offset = 0;
195+
test(thrust::offset_iterator(v.begin(), &offset), thrust::offset_iterator(v.begin(), &offset));
196+
}
197+
DECLARE_VECTOR_UNITTEST(TestOffsetIteratorComparisonAndDistance);
198+
199+
template <typename Vector>
200+
void TestOffsetIteratorLateValue()
201+
{
202+
typename Vector::difference_type offset;
203+
Vector v{0, 1, 2, 3, 4, 5, 6, 7, 8};
204+
thrust::offset_iterator iter(v.begin(), &offset);
205+
offset = 2; // we provide the offset value **after** constructing the iterator
206+
ASSERT_EQUAL(*iter, 2);
207+
}
208+
DECLARE_VECTOR_UNITTEST(TestOffsetIteratorLateValue);
209+
210+
template <typename Vector>
211+
void TestOffsetIteratorIndirectValueFancyIterator()
212+
{
213+
using thrust::placeholders::_1;
214+
215+
Vector v{0, 1, 2, 3, 4, 5, 6, 7, 8};
216+
thrust::device_vector<typename Vector::difference_type> offsets{2};
217+
auto it = thrust::make_transform_iterator(offsets.begin(), _1 * 3);
218+
thrust::offset_iterator iter(v.begin(), it);
219+
ASSERT_EQUAL(*iter, 6);
220+
}
221+
DECLARE_VECTOR_UNITTEST(TestOffsetIteratorIndirectValueFancyIterator);

0 commit comments

Comments
 (0)