refactor

NVIDIA · Mar 5, 2025 · 8e16c10 · 8e16c10
1 parent 1f1292a
commit 8e16c10
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 32 deletions.
diff --git a/thrust/testing/counting_iterator.cu b/thrust/testing/counting_iterator.cu
@@ -321,19 +321,4 @@ void TestCountingIteratorPointer()
 }
 DECLARE_UNITTEST(TestCountingIteratorPointer);
 
-void TestCountingIteratorStridedPointer()
-{
-  ::cuda::std::array<std::pair<int, double>, 4> arr{{{1, 2}, {3, 4}, {5, 6}, {7, 8}}};
-
-  // iterate over all second elements
-  auto iter = thrust::make_counting_iterator(
-    &arr[0].second, ::cuda::std::integral_constant<int, sizeof(std::pair<int, double>)>{});
-
-  cuda::std::fill(iter, iter + 4, 1337);
-
-  const ::cuda::std::array<std::pair<int, double>, 4> reference{{{1, 1337}, {3, 1337}, {5, 1337}, {7, 1337}}};
-  ASSERT_EQUAL(arr, reference);
-}
-DECLARE_UNITTEST(TestCountingIteratorStridedPointer);
-
 _CCCL_DIAG_POP
diff --git a/thrust/thrust/iterator/counting_iterator.h b/thrust/thrust/iterator/counting_iterator.h
@@ -188,17 +188,6 @@ struct counting_iterator_equal<Difference,
 
 struct empty
 {};
-
-template <typename T>
-struct value_holder
-{
-  T value;
-
-  _CCCL_HOST_DEVICE auto operator()() const
-  {
-    return value;
-  }
-};
 } // namespace detail
 
 //! \addtogroup iterators
@@ -336,7 +325,7 @@ class _CCCL_DECLSPEC_EMPTY_BASES counting_iterator
   auto step() const
   {
     static_assert(!::cuda::std::is_same_v<Step, detail::empty>);
-    return static_cast<const Step&>(*this)();
+    return static_cast<const Step&>(*this).value;
   }
 
   _CCCL_EXEC_CHECK_DISABLE
@@ -422,25 +411,46 @@ inline _CCCL_HOST_DEVICE counting_iterator<Incrementable> make_counting_iterator
   return counting_iterator<Incrementable>(x);
 }
 
+namespace detail
+{
+// Holds a runtime step
+template <typename T>
+struct runtime_step_holder
+{
+  T value;
+};
+
+// Holds a compile-time step
+// (we cannot use ::cuda::std::integral_constant, because it has a conversion operator to T that causes an ambiguity
+// with operator+(counting_iterator, counting_iterator::difference_type) in any expression `counting_iterator +
+// integral`.
+template <typename T, T Value>
+struct compile_time_step_holder
+{
+  static constexpr T value = Value;
+};
+} // namespace detail
+
+//! Constructs a counting_iterator with a runtime stride
 template <typename Incrementable, typename Stride>
-inline _CCCL_HOST_DEVICE auto make_counting_iterator(Incrementable x, Stride stride)
+_CCCL_HOST_DEVICE auto make_counting_iterator(Incrementable x, Stride stride)
 {
   return counting_iterator<Incrementable,
                            use_default,
                            random_access_traversal_tag,
                            use_default,
-                           detail::value_holder<Stride>>(x, {stride});
+                           detail::runtime_step_holder<Stride>>(x, {stride});
 }
 
+//! Constructs a counting_iterator with a compile-time stride
 template <typename Incrementable, typename Stride, Stride Value>
-inline _CCCL_HOST_DEVICE auto
-make_counting_iterator(Incrementable x, ::cuda::std::integral_constant<Stride, Value> stride)
+_CCCL_HOST_DEVICE auto make_counting_iterator(Incrementable x, ::cuda::std::integral_constant<Stride, Value>)
 {
   return counting_iterator<Incrementable,
                            use_default,
                            random_access_traversal_tag,
                            use_default,
-                           ::cuda::std::integral_constant<Stride, Value>>(x, stride);
+                           detail::compile_time_step_holder<Stride, Value>>(x, {});
 }
 
 //! \} // end fancyiterators