Host incrementable iterator approach 2 (#4697)

* Checkpoint for implementing approach 2 in gh-4148

* Checkpoint 2

* Printing added for debugging

* Almost working version

* Remove superfluous import, increased num_segments to exceed INT_MAX

* Define segment size variable and use it in the test

* Restore INT_MAX increment in streaming iterations in segmented_reduce

* Remove stdio from indirect_arg, remove printf

* Test for large_num_segments is parametrized by three modulus params

The test cooks input sequence in the form of F(i+1) - F(i) for
0 <= i < num_segments * segment_size.

Fixed-size segmented reduction is expected to to be

[ F(segment_size * (k + 1)) - F(segment_size * k) for
     k in range(num_segments) ]

All operations (addition (+), inversion (-)) are assumed to take
place in a ring.

In this test the ring is a cycling ring of integers modulo 7.
F(i) = g(i) * g(i+1) % 7, where g: i -> (i % 5 + i % 3).

Note that `my_add` computes addition modulo 7 too.

Choice of moduli 3, 5, 7 is arbitrary and may be changed. The test
is expected to pass so long as they are changed consistently throughout
the test.

* Add test_large_num_segments3

This example computes segmented reduce of input constant iterator
of int8(1) values. The offsets iterators are crafted to that segment
lengths are

1, 2, ..., (p - 2), (p - 1), 1, 2, ...,  (p - 2), (p - 1), 1, ...

The expected sums would equal these lengths.

The example is run for (2**15 + 2**3) * 2**16 number of segments
which is INT_MAX and about 2**19 elements.

1. Input iterator changed to have type int16
2. Initial value changed to have type int16
3. Changed range of change of segments from [1, p] to
   [m0, m0 + p) and set m0, p as they are set in native test.

   Specifically, m0, p = 265, 163

This means segment sizes vary from 265 elements to 428 elements,
which is an improvement over previous situation where it was varying
from 1 to 113 elements (and most threads in a 256-string block were
unused).

* Change indirect_iterator_t::operator+= to throw if host_advance_fn is null

This ensures meaningful prescriptive cerr output.

* Enable large num_segments test in test_segmented_reduce.cpp

* Add docstrings to test_large_num_segments and test_large_num_segments3

* Fix type in the test name in test_bindings

* Remove make_advance_iterator. Implement IteratorBase.__add__ to copy state and advance cvalue. We rely on cvalue being a ctype for which __add__ is defined

* Add host_advance_ property to iterators

For IteratorBase it is defined as None
For concrete iterators, it is defined as input_iterators, except
for transform iterator, where host_advance does not involve numba-cuda
compiled object.

* Add stubs for Iterators.host_advance_fn getter/setter

* Add function to set host_advance_fn member of bindings Iterator member

* segmented_reduce calls to set host_advance_fn on output and start/end offset iterators

This is done during build step irrespective of the input size, but it should perhaps
be done as needed the first time num_segments > INT_MAX is encountered. TBD

* Remove band-aid setting of host_advance_fn.

The host_advance function is now generated and set during segmented_reduce
build step.

* Use concept to narrow types of supported increment in indirect_iterator_t::operator+=

* Change validation of large num_segments tests in Python

Validation is now performed by constructing an iterator yielding
expected result. The allocating fixed size uint8 buffer and running
binary transform to compare actual to expected and write result
of comparison to the validation buffer. We then use
`cp.all(validation.view(cp.bool_))` to validate this chunk, and
move to the next.

* Mark large num_segments tests as large

* Run large tests, but do it sequentially (-n 0)

Author: oleksandr-pavlyk

c/parallel/include/cccl/c/types.h

@@ -21,6 +21,7 @@
 #endif // !_WIN32
 
 #include <stddef.h>
+#include <stdint.h>
 
 #include <cccl/c/extern_c.h>
 
@@ -77,6 +78,14 @@ typedef struct cccl_value_t
   void* state;
 } cccl_value_t;
 
+typedef union
+{
+  int64_t signed_offset;
+  uint64_t unsigned_offset;
+} cccl_increment_t;
+
+typedef void (*cccl_host_op_fn_ptr_t)(void*, cccl_increment_t);
+
 typedef struct cccl_iterator_t
 {
   size_t size;
@@ -86,6 +95,7 @@ typedef struct cccl_iterator_t
   cccl_op_t dereference;
   cccl_type_info value_type;
   void* state;
+  cccl_host_op_fn_ptr_t host_advance;
 } cccl_iterator_t;
 
 typedef enum cccl_sort_order_t

c/parallel/src/segmented_reduce.cu

@@ -276,6 +276,7 @@ CUresult cccl_device_segmented_reduce_build(
 
     const std::string input_iterator_src =
       make_kernel_input_iterator(offset_t, "input_iterator_t", input_it_value_t, input_it);
+
     const std::string output_iterator_src =
       make_kernel_output_iterator(offset_t, "output_iterator_t", accum_cpp, output_it);
 
@@ -407,9 +408,9 @@ CUresult cccl_device_segmented_reduce(
 
     auto exec_status = cub::DispatchSegmentedReduce<
       indirect_arg_t, // InputIteratorT
-      indirect_arg_t, // OutputIteratorT
-      indirect_arg_t, // BeginSegmentIteratorT
-      indirect_arg_t, // EndSegmentIteratorT
+      indirect_iterator_t, // OutputIteratorT
+      indirect_iterator_t, // BeginSegmentIteratorT
+      indirect_iterator_t, // EndSegmentIteratorT
       OffsetT, // OffsetT
       indirect_arg_t, // ReductionOpT
       indirect_arg_t, // InitT
@@ -421,10 +422,10 @@ CUresult cccl_device_segmented_reduce(
         d_temp_storage,
         *temp_storage_bytes,
         d_in,
-        d_out,
+        indirect_iterator_t{d_out},
         num_segments,
-        start_offset,
-        end_offset,
+        indirect_iterator_t{start_offset},
+        indirect_iterator_t{end_offset},
         op,
         init,
         stream,

c/parallel/src/util/indirect_arg.h

@@ -10,6 +10,10 @@
 
 #pragma once
 
+#include <cstdint>
+#include <stdexcept>
+#include <type_traits>
+
 #include <cccl/c/types.h>
 
 struct indirect_arg_t
@@ -33,3 +37,69 @@ struct indirect_arg_t
     return ptr;
   }
 };
+
+template <typename U>
+concept Increment64 = std::is_integral_v<U> && sizeof(U) == sizeof(int64_t);
+
+struct indirect_iterator_t
+{
+  void* ptr;
+  size_t value_size;
+  cccl_host_op_fn_ptr_t host_advance_fn_p;
+
+  indirect_iterator_t(cccl_iterator_t& it)
+      : ptr{nullptr}
+      , value_size{0}
+      , host_advance_fn_p{nullptr}
+  {
+    if (it.type == cccl_iterator_kind_t::CCCL_POINTER)
+    {
+      value_size = it.value_type.size;
+      ptr        = &it.state;
+    }
+    else
+    {
+      ptr               = it.state;
+      host_advance_fn_p = it.host_advance;
+    }
+  }
+
+  void* operator&() const
+  {
+    return ptr;
+  }
+
+  template <Increment64 U>
+  void operator+=(U offset)
+  {
+    if (value_size)
+    {
+      // CCCL_POINTER case
+      // ptr is a pointer to pointer we need to increment
+      // read the iterator pointer value
+      char*& p = *static_cast<char**>(ptr);
+      // increment the value
+      p += (offset * value_size);
+    }
+    else
+    {
+      if (host_advance_fn_p)
+      {
+        if constexpr (std::is_signed_v<U>)
+        {
+          cccl_increment_t incr{.signed_offset = offset};
+          (*host_advance_fn_p)(ptr, incr);
+        }
+        else
+        {
+          cccl_increment_t incr{.unsigned_offset = offset};
+          (*host_advance_fn_p)(ptr, incr);
+        }
+      }
+      else
+      {
+        throw std::runtime_error("Attempt to increment iterator from host, but host advance function is not defined");
+      }
+    }
+  }
+};

c/parallel/test/test_segmented_reduce.cpp

@@ -462,6 +462,22 @@ struct host_check_functor_state
   DataT* m_ptr;
 };
 
+template <typename StateT>
+void host_advance_transform_it_state(void* state, cccl_increment_t offset)
+{
+  auto st      = reinterpret_cast<StateT*>(state);
+  using IndexT = decltype(st->base_it_state.value);
+
+  if constexpr (std::is_signed_v<IndexT>)
+  {
+    st->base_it_state.value += offset.signed_offset;
+  }
+  else
+  {
+    st->base_it_state.value += offset.unsigned_offset;
+  }
+}
+
 namespace validate
 {
 
@@ -535,14 +551,12 @@ C2H_TEST("SegmentedReduce works with large num_segments", "[segmented_reduce]")
 
   static constexpr IndexT n_segments_base          = (IndexT(1) << 15) + (IndexT(1) << 3);
   static constexpr IndexT n_segments_under_int_max = n_segments_base << 10;
+  static_assert(n_segments_under_int_max < INT_MAX);
 
-#ifdef SUPPORTS_HOST_INCREMENT
   static constexpr IndexT n_segments_over_int_max = n_segments_base << 16;
+  static_assert(n_segments_over_int_max > INT_MAX);
 
   const IndexT n_segments = GENERATE(n_segments_under_int_max, n_segments_over_int_max);
-#else
-  const IndexT n_segments = GENERATE(n_segments_under_int_max, n_segments_under_int_max * 8);
-#endif
 
   // first define constant iterator:
   //   iterators.ConstantIterator(np.int8(1))
@@ -660,6 +674,10 @@ extern "C" __device__ void {0}(const void *x1_p, const void *x2_p, void *out_p)
   auto cccl_start_offsets_it = static_cast<cccl_iterator_t>(start_offsets_it);
   auto cccl_end_offsets_it   = static_cast<cccl_iterator_t>(end_offsets_it);
 
+  // set host_advance functions
+  cccl_start_offsets_it.host_advance = &host_advance_transform_it_state<HostTransformStateT>;
+  cccl_end_offsets_it.host_advance   = &host_advance_transform_it_state<HostTransformStateT>;
+
   value_t<DataT> h_init{DataT{0}};
 
   auto& build_cache    = get_cache<SegmentedReduce_LargeNumSegments_Fixture_Tag>();

ci/test_cuda_parallel_python.sh

@@ -21,3 +21,4 @@ python -m pip install "${CUDA_PARALLEL_WHEEL_PATH}[test]"
 # Run tests
 cd "/home/coder/cccl/python/cuda_parallel/tests/"
 python -m pytest -n 6 -v -m "not large"
+python -m pytest -n 0 -v -m "large"

python/cuda_parallel/cuda/parallel/experimental/_bindings.pyi

@@ -132,6 +132,7 @@ class Iterator:
         dereference_fn: Op,
         value_type: TypeInfo,
         state=None,
+        host_advance_fn=None,
     ):
         pass
 
@@ -148,6 +149,10 @@ class Iterator:
     def as_bytes(self) -> bytes: ...
     def is_kind_pointer(self) -> bool: ...
     def is_kind_iterator(self) -> bool: ...
+    @property
+    def host_advance_fn(self): ...
+    @host_advance_fn.setter
+    def host_advance_fn(self, value) -> None: ...
 
 class CommonData:
     def __init__(