Skip to content

Commit 1beb059

Browse files
misccomiscco
authored
Drop thrust::counting_iterator in favor of cuda::counting_iterator (#2958)
Authors: - Michael Schellenberger Costa (https://github.com/miscco) - Anupam (https://github.com/aamijar) Approvers: - Anupam (https://github.com/aamijar) - Divye Gala (https://github.com/divyegala) URL: #2958
1 parent f1db23c commit 1beb059

File tree

12 files changed

+64
-67
lines changed

12 files changed

+64
-67
lines changed

‎cpp/include/raft/linalg/detail/transpose.cuh‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* SPDX-FileCopyrightText: Copyright (c) 2022-2024, NVIDIA CORPORATION.
2+
* SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

@@ -14,8 +14,8 @@
1414

1515
#include <rmm/exec_policy.hpp>
1616

17+
#include <cuda/iterator>
1718
#include <thrust/for_each.h>
18-
#include <thrust/iterator/counting_iterator.h>
1919

2020
#include <cmath>
2121

@@ -167,7 +167,7 @@ void transpose(math_t* inout, int n, cudaStream_t stream)
167167
auto m = n;
168168
auto size = n * n;
169169
auto d_inout = inout;
170-
auto counting = thrust::make_counting_iterator<int>(0);
170+
auto counting = cuda::make_counting_iterator<int>(0);
171171

172172
thrust::for_each(rmm::exec_policy(stream), counting, counting + size, [=] __device__(int idx) {
173173
int s_row = idx % m;

‎cpp/include/raft/matrix/detail/gather_inplace.cuh‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
2+
* SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55
#pragma once
@@ -9,8 +9,8 @@
99
#include <raft/linalg/map.cuh>
1010
#include <raft/util/fast_int_div.cuh>
1111

12+
#include <cuda/iterator>
1213
#include <thrust/for_each.h>
13-
#include <thrust/iterator/counting_iterator.h>
1414

1515
namespace raft {
1616
namespace matrix {
@@ -81,7 +81,7 @@ void gatherInplaceImpl(raft::resources const& handle,
8181
inout[row * ld + batch_offset + col] = scratch_space[idx];
8282
return;
8383
};
84-
auto counting = thrust::make_counting_iterator<IndexT>(0);
84+
auto counting = cuda::make_counting_iterator<IndexT>(0);
8585
thrust::for_each(exec_policy, counting, counting + map_length * cols_per_batch, copy_op);
8686
}
8787
}

‎cpp/include/raft/matrix/detail/matrix.cuh‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* SPDX-FileCopyrightText: Copyright (c) 2021-2024, NVIDIA CORPORATION.
2+
* SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

@@ -14,9 +14,9 @@
1414

1515
#include <rmm/exec_policy.hpp>
1616

17+
#include <cuda/iterator>
1718
#include <cuda_runtime.h>
1819
#include <thrust/for_each.h>
19-
#include <thrust/iterator/counting_iterator.h>
2020

2121
#include <cusolverDn.h>
2222

@@ -46,7 +46,7 @@ void copyRows(const m_t* in,
4646
}
4747

4848
idx_t size = n_rows_indices * n_cols;
49-
auto counting = thrust::make_counting_iterator<idx_t>(0);
49+
auto counting = cuda::make_counting_iterator<idx_t>(0);
5050

5151
thrust::for_each(rmm::exec_policy(stream), counting, counting + size, [=] __device__(idx_t idx) {
5252
idx_t row = idx % n_rows_indices;
@@ -65,7 +65,7 @@ void truncZeroOrigin(
6565
idx_t size = out_n_rows * out_n_cols;
6666
auto d_q = in;
6767
auto d_q_trunc = out;
68-
auto counting = thrust::make_counting_iterator<idx_t>(0);
68+
auto counting = cuda::make_counting_iterator<idx_t>(0);
6969

7070
thrust::for_each(rmm::exec_policy(stream), counting, counting + size, [=] __device__(idx_t idx) {
7171
idx_t row = idx % m;
@@ -82,7 +82,7 @@ void colReverse(m_t* inout, idx_t n_rows, idx_t n_cols, cudaStream_t stream)
8282
idx_t size = n_rows * n_cols;
8383
auto d_q = inout;
8484
auto d_q_reversed = inout;
85-
auto counting = thrust::make_counting_iterator<idx_t>(0);
85+
auto counting = cuda::make_counting_iterator<idx_t>(0);
8686

8787
thrust::for_each(
8888
rmm::exec_policy(stream), counting, counting + (size / 2), [=] __device__(idx_t idx) {
@@ -103,7 +103,7 @@ void rowReverse(m_t* inout, idx_t n_rows, idx_t n_cols, cudaStream_t stream)
103103
idx_t size = n_rows * n_cols;
104104
auto d_q = inout;
105105
auto d_q_reversed = inout;
106-
auto counting = thrust::make_counting_iterator<idx_t>(0);
106+
auto counting = cuda::make_counting_iterator<idx_t>(0);
107107

108108
thrust::for_each(
109109
rmm::exec_policy(stream), counting, counting + (size / 2), [=] __device__(idx_t idx) {

‎cpp/include/raft/matrix/detail/print.hpp‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* SPDX-FileCopyrightText: Copyright (c) 2021-2024, NVIDIA CORPORATION.
2+
* SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

@@ -14,8 +14,6 @@
1414
#include <rmm/exec_policy.hpp>
1515

1616
#include <cuda_runtime.h>
17-
#include <thrust/for_each.h>
18-
#include <thrust/iterator/counting_iterator.h>
1917

2018
#include <cusolverDn.h>
2119

‎cpp/include/raft/matrix/detail/scatter_inplace.cuh‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
2+
* SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55
#pragma once
@@ -10,8 +10,8 @@
1010
#include <raft/util/cuda_dev_essentials.cuh>
1111
#include <raft/util/fast_int_div.cuh>
1212

13+
#include <cuda/iterator>
1314
#include <thrust/for_each.h>
14-
#include <thrust/iterator/counting_iterator.h>
1515

1616
#include <cstdint>
1717

@@ -100,7 +100,7 @@ void scatterInplaceImpl(
100100
inout[map_val * n + batch_offset + col] = scratch_space[idx];
101101
return;
102102
};
103-
auto counting = thrust::make_counting_iterator<IndexT>(0);
103+
auto counting = cuda::make_counting_iterator<IndexT>(0);
104104
thrust::for_each(exec_policy, counting, counting + m * cols_per_batch, scatter_op);
105105
}
106106
}

‎cpp/tests/core/mdarray.cu‎

Lines changed: 22 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,10 @@
2323
#include <rmm/device_vector.hpp>
2424
#include <rmm/exec_policy.hpp>
2525

26+
#include <cuda/iterator>
2627
#include <thrust/device_vector.h>
2728
#include <thrust/execution_policy.h>
2829
#include <thrust/for_each.h>
29-
#include <thrust/iterator/counting_iterator.h>
3030
#include <thrust/sequence.h>
3131

3232
#include <gtest/gtest.h>
@@ -51,7 +51,7 @@ void test_mdspan()
5151
thrust::device_vector<int32_t> status(1, 0);
5252
auto p_status = status.data().get();
5353
thrust::for_each_n(
54-
rmm::exec_policy(stream), thrust::make_counting_iterator(0ul), 4, [=] __device__(size_t i) {
54+
rmm::exec_policy(stream), cuda::make_counting_iterator(0ul), 4, [=] __device__(size_t i) {
5555
auto v = span(0, i);
5656
if (v != i) { raft::myAtomicAdd(p_status, 1); }
5757
auto k = cuda::std::submdspan(span, 0, cuda::std::full_extent);
@@ -101,7 +101,7 @@ void test_mdarray_basic()
101101
thrust::device_vector<int32_t> status(1, 0);
102102
auto p_status = status.data().get();
103103
thrust::for_each_n(rmm::exec_policy(s),
104-
thrust::make_counting_iterator(0ul),
104+
cuda::make_counting_iterator(0ul),
105105
1,
106106
[d_view, p_status] __device__(auto i) {
107107
if (d_view(0, 3) != 1) { myAtomicAdd(p_status, 1); }
@@ -115,7 +115,7 @@ void test_mdarray_basic()
115115
ASSERT_EQ(arr(0, 3), 1);
116116
auto const_d_view = arr.view();
117117
thrust::for_each_n(rmm::exec_policy(s),
118-
thrust::make_counting_iterator(0ul),
118+
cuda::make_counting_iterator(0ul),
119119
1,
120120
[const_d_view, p_status] __device__(auto i) {
121121
if (const_d_view(0, 3) != 1) { myAtomicAdd(p_status, 1); }
@@ -162,7 +162,7 @@ void test_mdarray_basic()
162162
ASSERT_EQ(array(0, 3), 1);
163163
auto h_view = array.view();
164164
static_assert(decltype(h_view)::accessor_type::is_host_type::value);
165-
thrust::for_each_n(thrust::host, thrust::make_counting_iterator(0ul), 1, [h_view](auto i) {
165+
thrust::for_each_n(thrust::host, cuda::make_counting_iterator(0ul), 1, [h_view](auto i) {
166166
ASSERT_EQ(h_view(0, 3), 1);
167167
});
168168
}
@@ -359,7 +359,7 @@ void test_factory_methods()
359359
thrust::device_vector<int32_t> status(1, 0);
360360
auto p_status = status.data().get();
361361
thrust::for_each_n(rmm::exec_policy(resource::get_cuda_stream(handle)),
362-
thrust::make_counting_iterator(0),
362+
cuda::make_counting_iterator(0),
363363
1,
364364
[=] __device__(auto i) {
365365
if (view(i) != 17.0) { myAtomicAdd(p_status, 1); }
@@ -518,7 +518,7 @@ void test_mdarray_padding()
518518
thrust::device_vector<int32_t> status(1, 0);
519519
auto p_status = status.data().get();
520520
thrust::for_each_n(rmm::exec_policy(),
521-
thrust::make_counting_iterator(0ul),
521+
cuda::make_counting_iterator(0ul),
522522
1,
523523
[d_view, p_status] __device__(size_t i) {
524524
if (d_view(0, 3) != 1) { myAtomicAdd(p_status, 1); }
@@ -532,7 +532,7 @@ void test_mdarray_padding()
532532
ASSERT_EQ(arr(0, 3), 1);
533533
auto const_d_view = arr.view();
534534
thrust::for_each_n(rmm::exec_policy(s),
535-
thrust::make_counting_iterator(0ul),
535+
cuda::make_counting_iterator(0ul),
536536
1,
537537
[const_d_view, p_status] __device__(size_t i) {
538538
if (const_d_view(0, 3) != 1) { myAtomicAdd(p_status, 1); }
@@ -542,7 +542,7 @@ void test_mdarray_padding()
542542
// initialize with sequence
543543
thrust::for_each_n(
544544
rmm::exec_policy(s),
545-
thrust::make_counting_iterator(0ul),
545+
cuda::make_counting_iterator(0ul),
546546
rows * cols,
547547
[d_view, rows, cols] __device__(size_t i) { d_view(i / cols, i % cols) = i; });
548548

@@ -552,7 +552,7 @@ void test_mdarray_padding()
552552
using padded_mdspan_type = device_mdspan<float, extents_type, padded_layout_row_major>;
553553
auto padded_span = padded_mdspan_type(data_padded, layout);
554554
thrust::for_each_n(rmm::exec_policy(s),
555-
thrust::make_counting_iterator(0ul),
555+
cuda::make_counting_iterator(0ul),
556556
rows * cols,
557557
[padded_span, rows, cols, p_status] __device__(size_t i) {
558558
if (padded_span(i / cols, i % cols) != i) myAtomicAdd(p_status, 1);
@@ -608,7 +608,7 @@ TEST(MDArray, Padding) { test_mdarray_padding(); }
608608
static_assert(std::is_same_v<typename decltype(d_view)::layout_type, layout_padded_general>);
609609
thrust::for_each_n(
610610
rmm::exec_policy(s),
611-
thrust::make_counting_iterator(0ul),
611+
cuda::make_counting_iterator(0ul),
612612
rows * cols,
613613
[d_view, rows, cols] __device__(size_t i) { d_view(i / cols, i % cols) = i; });
614614
}
@@ -618,7 +618,7 @@ TEST(MDArray, Padding) { test_mdarray_padding(); }
618618
auto data_padded = padded_device_array.data();
619619
auto padded_span = padded_mdspan_type(data_padded, layout);
620620
thrust::for_each_n(rmm::exec_policy(s),
621-
thrust::make_counting_iterator(0ul),
621+
cuda::make_counting_iterator(0ul),
622622
rows * cols,
623623
[padded_span, rows, cols, p_status] __device__(size_t i) {
624624
if (padded_span(i / cols, i % cols) != i) myAtomicAdd(p_status, 1);
@@ -631,7 +631,7 @@ TEST(MDArray, Padding) { test_mdarray_padding(); }
631631
auto padded_span = padded_device_array.view();
632632
auto subspan_full = stdex::submdspan(padded_span, stdex::full_extent, stdex::full_extent);
633633
thrust::for_each_n(rmm::exec_policy(s),
634-
thrust::make_counting_iterator(0ul),
634+
cuda::make_counting_iterator(0ul),
635635
cols * rows,
636636
[subspan_full, padded_span, rows, cols, p_status] __device__(size_t i) {
637637
if (subspan_full(i / cols, i % cols) != padded_span(i / cols, i % cols))
@@ -649,7 +649,7 @@ TEST(MDArray, Padding) { test_mdarray_padding(); }
649649
auto padded_span = padded_device_array.view();
650650
auto row3 = stdex::submdspan(padded_span, 3, stdex::full_extent);
651651
thrust::for_each_n(rmm::exec_policy(s),
652-
thrust::make_counting_iterator(0ul),
652+
cuda::make_counting_iterator(0ul),
653653
cols,
654654
[row3, padded_span, p_status] __device__(size_t i) {
655655
if (row3(i) != padded_span(3, i)) myAtomicAdd(p_status, 1);
@@ -665,7 +665,7 @@ TEST(MDArray, Padding) { test_mdarray_padding(); }
665665
auto padded_span = padded_device_array.view();
666666
auto col1 = stdex::submdspan(padded_span, stdex::full_extent, 1);
667667
thrust::for_each_n(rmm::exec_policy(s),
668-
thrust::make_counting_iterator(0ul),
668+
cuda::make_counting_iterator(0ul),
669669
rows,
670670
[col1, padded_span, p_status] __device__(size_t i) {
671671
if (col1(i) != padded_span(i, 1)) myAtomicAdd(p_status, 1);
@@ -682,7 +682,7 @@ TEST(MDArray, Padding) { test_mdarray_padding(); }
682682
auto subspan =
683683
stdex::submdspan(padded_span, std::make_tuple(1ul, 4ul), std::make_tuple(2ul, 5ul));
684684
thrust::for_each_n(rmm::exec_policy(s),
685-
thrust::make_counting_iterator(0ul),
685+
cuda::make_counting_iterator(0ul),
686686
(rows - 1) * (cols - 2),
687687
[subspan, rows, cols, padded_span, p_status] __device__(size_t i) {
688688
size_t idx = i / (cols - 2);
@@ -703,7 +703,7 @@ TEST(MDArray, Padding) { test_mdarray_padding(); }
703703
auto subspan =
704704
stdex::submdspan(padded_span, std::make_tuple(1ul, 4ul), std::make_tuple(2ul, 5ul));
705705
thrust::for_each_n(rmm::exec_policy(s),
706-
thrust::make_counting_iterator(0ul),
706+
cuda::make_counting_iterator(0ul),
707707
(rows - 1) * (cols - 2),
708708
[subspan, rows, cols, padded_span, p_status] __device__(size_t i) {
709709
size_t idx = i / (cols - 2);
@@ -756,7 +756,7 @@ void test_mdspan_padding_by_type()
756756
auto padded_span = padded_device_array.view();
757757
thrust::for_each_n(
758758
rmm::exec_policy(s),
759-
thrust::make_counting_iterator(0ul),
759+
cuda::make_counting_iterator(0ul),
760760
rows * cols,
761761
[rows, cols, padded_span, alignment_elements, p_status] __device__(size_t i) {
762762
size_t idx = i / cols;
@@ -783,7 +783,7 @@ void test_mdspan_padding_by_type()
783783
auto padded_span = padded_device_array.view();
784784
thrust::for_each_n(
785785
rmm::exec_policy(s),
786-
thrust::make_counting_iterator(0ul),
786+
cuda::make_counting_iterator(0ul),
787787
rows * cols,
788788
[rows, cols, padded_span, alignment_elements, p_status] __device__(size_t i) {
789789
size_t idx = i / cols;
@@ -843,7 +843,7 @@ void test_mdspan_aligned_matrix()
843843
thrust::device_vector<int32_t> status(1, 0);
844844
auto p_status = status.data().get();
845845
thrust::for_each_n(rmm::exec_policy(s),
846-
thrust::make_counting_iterator(0ul),
846+
cuda::make_counting_iterator(0ul),
847847
rows * cols,
848848
[rows, cols, my_aligned_device_span, p_status] __device__(size_t i) {
849849
size_t idx = i / cols;
@@ -922,7 +922,7 @@ void test_mdarray_unravel()
922922
auto m = make_device_matrix<float, size_t>(handle, 7, 6);
923923
auto m_v = m.view();
924924
thrust::for_each_n(resource::get_thrust_policy(handle),
925-
thrust::make_counting_iterator(0ul),
925+
cuda::make_counting_iterator(0ul),
926926
m_v.size(),
927927
[=] HD(size_t i) {
928928
auto coord =
@@ -932,7 +932,7 @@ void test_mdarray_unravel()
932932
thrust::device_vector<int32_t> status(1, 0);
933933
auto p_status = status.data().get();
934934
thrust::for_each_n(resource::get_thrust_policy(handle),
935-
thrust::make_counting_iterator(0ul),
935+
cuda::make_counting_iterator(0ul),
936936
m_v.size(),
937937
[=] __device__(size_t i) {
938938
auto coord =

0 commit comments

Comments
 (0)