Checkpoint for implementing approach 2 in NVIDIAgh-4148

Author: oleksandr-pavlyk

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

@@ -22,6 +22,7 @@
 
 #include <cccl/c/extern_c.h>
 #include <stddef.h>
+#include <stdint.h>
 
 CCCL_C_EXTERN_C_BEGIN
 
@@ -76,6 +77,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;
@@ -84,6 +93,7 @@ typedef struct cccl_iterator_t
   cccl_op_t advance;
   cccl_op_t dereference;
   cccl_type_info value_type;
+  cccl_host_op_fn_ptr_t host_advance;
   void* state;
 } cccl_iterator_t;
 

c/parallel/src/segmented_reduce.cu

@@ -406,9 +406,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

c/parallel/src/util/indirect_arg.h

@@ -33,3 +33,66 @@ struct indirect_arg_t
     return ptr;
   }
 };
+
+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;
+  }
+
+  void operator+=(int64_t signed_offset)
+  {
+    if (value_size)
+    {
+      // CCCL_POINTER case
+      ptr = reinterpret_cast<void*>(reinterpret_cast<char*>(ptr) + (signed_offset * value_size));
+    }
+    else
+    {
+      if (host_advance_fn_p)
+      {
+        cccl_increment_t incr{.signed_offset = signed_offset};
+        (*host_advance_fn_p)(ptr, incr);
+      }
+    }
+  }
+
+  void operator+=(uint64_t unsigned_offset)
+  {
+    if (value_size)
+    {
+      // CCCL_POINTER case
+      ptr = reinterpret_cast<void*>(reinterpret_cast<char*>(ptr) + (unsigned_offset * value_size));
+    }
+    else
+    {
+      if (host_advance_fn_p)
+      {
+        cccl_increment_t incr{.unsigned_offset = unsigned_offset};
+        (*host_advance_fn_p)(ptr, incr);
+      }
+    }
+  }
+};

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
 
@@ -287,7 +288,7 @@ class DeviceMergeSortBuildResult:
         num_items: int,
         binary_op: Op,
         stream,
-    ) -> tuple[int, int]: ...
+    ) -> int: ...
 
 # -----------------
 # DeviceUniqueByKey
@@ -317,7 +318,7 @@ class DeviceUniqueByKeyBuildResult:
         binary_op: Op,
         num_items: int,
         stream,
-    ) -> tuple[int, int]: ...
+    ) -> int: ...
 
 # -----------------
 # DeviceRadixSort

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

@@ -7,7 +7,7 @@
 # static type checker tools like mypy green-lights cuda.parallel
 
 from libc.string cimport memset, memcpy
-from libc.stdint cimport uint8_t, uint32_t, uint64_t
+from libc.stdint cimport uint8_t, uint32_t, uint64_t, int64_t
 from cpython.bytes cimport PyBytes_FromStringAndSize
 
 from cpython.buffer cimport (
@@ -68,13 +68,20 @@ cdef extern from "cccl/c/types.h":
         cccl_type_info type
         void *state
 
+    cdef union cccl_increment_t:
+        int64_t signed_offset
+        uint64_t unsigned_offset
+
+    ctypedef void (*cccl_host_op_fn_ptr_t)(void *, cccl_increment_t) nogil
+
     cdef struct cccl_iterator_t:
         size_t size
         size_t alignment
         cccl_iterator_kind_t type
         cccl_op_t advance
         cccl_op_t dereference
         cccl_type_info value_type
+        cccl_host_op_fn_ptr_t host_advance
         void *state
 
     ctypedef enum cccl_sort_order_t:
@@ -799,10 +806,16 @@ cdef class IteratorState(StateBase):
         pass
 
 
+
+cdef cccl_host_op_fn_ptr_t unbox_host_advance_fn(object host_fn_obj) except *:
+    return <cccl_host_op_fn_ptr_t>NULL
+
+
 cdef class Iterator:
     cdef Op advance
     cdef Op dereference
     cdef object state_obj
+    cdef object host_advance_obj
     cdef cccl_iterator_t iter_data
 
     def __cinit__(self,
@@ -811,7 +824,8 @@ cdef class Iterator:
         Op advance_fn,
         Op dereference_fn,
         TypeInfo value_type,
-        state = None
+        state=None,
+        host_advance_fn=None
     ):
         cdef cccl_iterator_kind_t it_kind
         _validate_alignment(alignment)
@@ -836,6 +850,12 @@ cdef class Iterator:
                     "Expect for Iterator of kind POINTER, state must have type Pointer or int, "
                     f"got {type(state)}"
                 )
+            if host_advance_fn is not None:
+                raise ValueError(
+                    "host_advance_fn must be set to None for iterators of kind POINTER"
+                )
+            self.iter_data.host_advance = NULL
+            self.host_advance_obj = None
         elif it_kind == cccl_iterator_kind_t.CCCL_ITERATOR:
             if state is None:
                 self.state_obj = None
@@ -850,6 +870,12 @@ cdef class Iterator:
                     "For Iterator of kind ITERATOR, state must have type IteratorState, "
                     f"got type {type(state)}"
                 )
+            if host_advance_fn is not None:
+                self.iter_data.host_advance = unbox_host_advance_fn(host_advance_fn)
+                self.host_advance_obj = host_advance_fn
+            else:
+                self.iter_data.host_advance = NULL
+                self.host_advance_obj = None
         else:  # pragma: no cover
             raise ValueError("Unrecognized iterator kind")
         self.advance = advance_fn