Fix iterator sigs

shwina · shwina · commit 9d31c800655d · 2025-11-14T20:45:48.000Z
diff --git a/c/parallel/src/jit_templates/mappings/iterator.h b/c/parallel/src/jit_templates/mappings/iterator.h
@@ -19,12 +19,12 @@
 template <typename ValueTp>
 struct cccl_iterator_t_mapping
 {
-  bool is_pointer                             = false;
-  int size                                    = 1;
-  int alignment                               = 1;
-  void (*advance)(void*, cuda::std::uint64_t) = nullptr;
-  void (*dereference)(const void*, ValueTp*)  = nullptr;
-  void (*assign)(const void*, ValueTp);
+  bool is_pointer                            = false;
+  int size                                   = 1;
+  int alignment                              = 1;
+  void (*advance)(void*, const void*)        = nullptr;
+  void (*dereference)(const void*, ValueTp*) = nullptr;
+  void (*assign)(const void*, const void*);
 
   using ValueT = ValueTp;
 };
@@ -68,22 +68,19 @@ struct parameter_mapping<cccl_iterator_t>
     {
       return std::format(
         R"output(
-extern "C" __device__ void {0}(void *, {1});
-extern "C" __device__ void {2}(const void *, {3});
+extern "C" __device__ void {0}(void *, const void*);
+extern "C" __device__ void {1}(const void *, const void*);
 )output",
         arg.advance.name,
-        cccl_type_enum_to_name(cccl_type_enum::CCCL_UINT64),
-        arg.dereference.name,
-        cccl_type_enum_to_name(arg.value_type.type));
+        arg.dereference.name);
     }
 
     return std::format(
       R"input(
-extern "C" __device__ void {0}(void *, {1});
-extern "C" __device__ void {2}(const void *, {3}*);
+extern "C" __device__ void {0}(void *, const void*);
+extern "C" __device__ void {1}(const void *, {2}*);
 )input",
       arg.advance.name,
-      cccl_type_enum_to_name(cccl_type_enum::CCCL_UINT64),
       arg.dereference.name,
       cccl_type_enum_to_name(arg.value_type.type));
   }
diff --git a/c/parallel/src/jit_templates/templates/input_iterator.h b/c/parallel/src/jit_templates/templates/input_iterator.h
@@ -47,7 +47,7 @@ struct input_iterator_t
 
   __device__ input_iterator_t& operator+=(difference_type diff)
   {
-    Iterator.advance(&state, diff);
+    Iterator.advance(&state, &diff);
     return *this;
   }
 
diff --git a/c/parallel/src/jit_templates/templates/output_iterator.h b/c/parallel/src/jit_templates/templates/output_iterator.h
@@ -35,7 +35,7 @@ struct output_iterator_proxy_t
 {
   __device__ output_iterator_proxy_t& operator=(AssignT x)
   {
-    AssignF(&state, cuda::std::move(x));
+    AssignF(&state, &x);
     return *this;
   }
 
@@ -59,7 +59,7 @@ struct output_iterator_t
 
   __device__ output_iterator_t& operator+=(difference_type diff)
   {
-    Iterator.advance(&state, diff);
+    Iterator.advance(&state, &diff);
     return *this;
   }
 
diff --git a/c/parallel/src/kernels/iterators.cpp b/c/parallel/src/kernels/iterators.cpp
@@ -45,7 +45,7 @@ std::string make_kernel_input_iterator(
 {
   const std::string iter_def = std::format(R"XXX(
 extern "C" __device__ void DEREF(const void *self_ptr, VALUE_T* result);
-extern "C" __device__ void ADVANCE(void *self_ptr, DIFF_T offset);
+extern "C" __device__ void ADVANCE(void *self_ptr, const void* offset);
 struct __align__(OP_ALIGNMENT) {0} {{
   using iterator_category = cuda::std::random_access_iterator_tag;
   using value_type = VALUE_T;
@@ -58,7 +58,7 @@ struct __align__(OP_ALIGNMENT) {0} {{
     return result;
   }}
   __device__ inline {0}& operator+=(difference_type diff) {{
-      ADVANCE(data, diff);
+      ADVANCE(data, &diff);
       return *this;
   }}
   __device__ inline value_type operator[](difference_type diff) const {{
@@ -99,14 +99,14 @@ std::string make_kernel_output_iterator(
   std::string_view advance)
 {
   const std::string iter_def = std::format(R"XXX(
-extern "C" __device__ void DEREF(const void *self_ptr, VALUE_T x);
-extern "C" __device__ void ADVANCE(void *self_ptr, DIFF_T offset);
+extern "C" __device__ void DEREF(const void *self_ptr, const void* x);
+extern "C" __device__ void ADVANCE(void *self_ptr, const void* offset);
 struct __align__(OP_ALIGNMENT) {0}_state_t {{
   char data[OP_SIZE];
 }};
 struct {0}_proxy_t {{
   __device__ {0}_proxy_t operator=(VALUE_T x) {{
-    DEREF(&state, x);
+    DEREF(&state, &x);
     return *this;
   }}
   {0}_state_t state;
@@ -119,7 +119,7 @@ struct {0} {{
   using reference         = {0}_proxy_t;
   __device__ {0}_proxy_t operator*() const {{ return {{state}}; }}
   __device__ {0}& operator+=(difference_type diff) {{
-      ADVANCE(&state, diff);
+      ADVANCE(&state, &diff);
       return *this;
   }}
   __device__ {0}_proxy_t operator[](difference_type diff) const {{
diff --git a/c/parallel/test/test_util.h b/c/parallel/test/test_util.h
@@ -931,8 +931,10 @@ inline std::tuple<std::string, std::string, std::string> make_random_access_iter
 {
   std::string state_def_src      = std::format("struct {0} {{ {1}* data; }};\n", iterator_state_name, value_type);
   std::string advance_fn_def_src = std::format(
-    "extern \"C\" __device__ void {0}({1}* state, unsigned long long offset) {{\n"
-    "  state->data += offset;\n"
+    "extern \"C\" __device__ void {0}(void* state, const void* offset) {{\n"
+    "  auto* typed_state = static_cast<{1}*>(state);\n"
+    "  auto offset_val = *static_cast<const unsigned long long*>(offset);\n"
+    "  typed_state->data += offset_val;\n"
     "}}",
     advance_fn_name,
     iterator_state_name);
@@ -941,19 +943,22 @@ inline std::tuple<std::string, std::string, std::string> make_random_access_iter
   if (kind == iterator_kind::INPUT)
   {
     dereference_fn_def_src = std::format(
-      "extern \"C\" __device__ void {0}({1}* state, {2}* result) {{\n"
-      "  *result = (*state->data){3};\n"
+      "extern \"C\" __device__ void {0}(const void* state, {1}* result) {{\n"
+      "  auto* typed_state = static_cast<const {2}*>(state);\n"
+      "  *result = (*typed_state->data){3};\n"
       "}}",
       dereference_fn_name,
-      iterator_state_name,
       value_type,
+      iterator_state_name,
       transform);
   }
   else
   {
     dereference_fn_def_src = std::format(
-      "extern \"C\" __device__ void {0}({1}* state, {2} x) {{\n"
-      "  *state->data = x{3};\n"
+      "extern \"C\" __device__ void {0}(const void* state, const void* x) {{\n"
+      "  auto* typed_state = static_cast<const {1}*>(state);\n"
+      "  auto x_val = *static_cast<const {2}*>(x);\n"
+      "  *typed_state->data = x_val{3};\n"
       "}}",
       dereference_fn_name,
       iterator_state_name,
@@ -1033,17 +1038,16 @@ inline std::tuple<std::string, std::string, std::string> make_constant_iterator_
   std::string_view dereference_fn_name)
 {
   std::string iterator_state_src = std::format("struct {0} {{ {1} value; }};\n", iterator_state_name, value_type);
-  std::string advance_fn_src     = std::format(
-    "extern \"C\" __device__ void {0}({1}* state, unsigned long long offset) {{ }}",
-    advance_fn_name,
-    iterator_state_name);
+  std::string advance_fn_src =
+    std::format("extern \"C\" __device__ void {0}(void* state, const void* offset) {{ }}", advance_fn_name);
   std::string dereference_fn_src = std::format(
-    "extern \"C\" __device__ void {0}({1}* state, {2}* result) {{ \n"
-    "  *result = state->value;\n"
+    "extern \"C\" __device__ void {0}(const void* state, {1}* result) {{ \n"
+    "  auto* typed_state = static_cast<const {2}*>(state);\n"
+    "  *result = typed_state->value;\n"
     "}}",
     dereference_fn_name,
-    iterator_state_name,
-    value_type);
+    value_type,
+    iterator_state_name);
 
   return std::make_tuple(iterator_state_src, advance_fn_src, dereference_fn_src);
 }
@@ -1124,24 +1128,27 @@ struct {0} {{
   const std::string it_state_def_src = std::format(it_state_src_tmpl, state_name, index_ty_name);
 
   static constexpr std::string_view it_def_src_tmpl = R"XXX(
-extern "C" __device__ void {0}({1}* state, {2} offset)
+extern "C" __device__ void {0}(void* state, const void* offset)
 {{
-  state->linear_id += offset;
+  auto* typed_state = static_cast<{1}*>(state);
+  auto offset_val = *static_cast<const {2}*>(offset);
+  typed_state->linear_id += offset_val;
 }}
 )XXX";
 
   const std::string it_advance_fn_def_src =
     std::format(it_def_src_tmpl, /*0*/ advance_fn_name, state_name, index_ty_name);
 
   static constexpr std::string_view it_deref_src_tmpl = R"XXX(
-extern "C" __device__ void {0}({1}* state, {2}* result)
+extern "C" __device__ void {0}(const void* state, {1}* result)
 {{
-  *result = (state->linear_id) * (state->segment_size);
+  auto* typed_state = static_cast<const {2}*>(state);
+  *result = (typed_state->linear_id) * (typed_state->segment_size);
 }}
 )XXX";
 
   const std::string it_deref_fn_def_src =
-    std::format(it_deref_src_tmpl, dereference_fn_name, state_name, index_ty_name);
+    std::format(it_deref_src_tmpl, dereference_fn_name, index_ty_name, state_name);
 
   return std::make_tuple(it_state_def_src, it_advance_fn_def_src, it_deref_fn_def_src);
 }
@@ -1229,8 +1236,9 @@ struct {0} {{
 
   static constexpr std::string_view transform_it_advance_fn_src_tmpl = R"XXX(
 {3}
-extern "C" __device__ void {0}({1} *transform_it_state, unsigned long long offset) {{
-    {2}(&(transform_it_state->base_it_state), offset);
+extern "C" __device__ void {0}(void* transform_it_state, const void* offset) {{
+    auto* typed_state = static_cast<{1}*>(transform_it_state);
+    {2}(&(typed_state->base_it_state), offset);
 }}
 )XXX";
 
@@ -1244,11 +1252,12 @@ extern "C" __device__ void {0}({1} *transform_it_state, unsigned long long offse
   static constexpr std::string_view transform_it_dereference_fn_src_tmpl = R"XXX(
 {5}
 {6}
-extern "C" __device__ void {0}({1} *transform_it_state, {2}* result) {{
+extern "C" __device__ void {0}(const void* transform_it_state, {2}* result) {{
+    auto* typed_state = static_cast<const {1}*>(transform_it_state);
     {7} base_result;
-    {4}(&(transform_it_state->base_it_state), &base_result);
+    {4}(&(typed_state->base_it_state), &base_result);
     *result = {3}(
-        &(transform_it_state->functor_state),
+        &(typed_state->functor_state),
         base_result
     );
 }}
diff --git a/python/cuda_cccl/cuda/compute/_cccl_interop.py b/python/cuda_cccl/cuda/compute/_cccl_interop.py
@@ -277,6 +277,142 @@ def wrapped_{op.__name__}({arg_str}):
     return wrapper_func, void_sig
 
 
+def _create_advance_wrapper(advance_fn, state_ptr_type):
+    """Creates a wrapper function for iterator advance that takes void* arguments.
+
+    The wrapper takes 2 void* arguments:
+    - state pointer
+    - offset pointer (points to uint64 value)
+    """
+    void_sig = types.void(types.voidptr, types.voidptr)
+
+    wrapper_src = textwrap.dedent(f"""
+    @intrinsic
+    def impl(typingctx, state_arg, offset_arg):
+        def codegen(context, builder, impl_sig, args):
+            state_type_llvm = context.get_value_type(state_ptr_type.dtype)
+            offset_type_llvm = context.get_value_type(types.uint64)
+
+            state_ptr = builder.bitcast(args[0], state_type_llvm.as_pointer())
+            offset_ptr = builder.bitcast(args[1], offset_type_llvm.as_pointer())
+            offset_val = builder.load(offset_ptr)
+
+            sig = types.void(state_ptr_type, types.uint64)
+            context.compile_internal(builder, advance_fn, sig, [state_ptr, offset_val])
+
+            return context.get_dummy_value()
+        return void_sig, codegen
+
+    def wrapped_{advance_fn.__name__}(state_arg, offset_arg):
+        return impl(state_arg, offset_arg)
+    """)
+
+    local_dict = {
+        "types": types,
+        "state_ptr_type": state_ptr_type,
+        "advance_fn": advance_fn,
+        "intrinsic": intrinsic,
+        "void_sig": void_sig,
+    }
+    exec(wrapper_src, globals(), local_dict)
+
+    wrapper_func = local_dict[f"wrapped_{advance_fn.__name__}"]
+    wrapper_func.__globals__.update(local_dict)
+
+    return wrapper_func, void_sig
+
+
+def _create_input_dereference_wrapper(deref_fn, state_ptr_type, value_type):
+    """Creates a wrapper function for input iterator dereference that takes void* arguments.
+
+    The wrapper takes 2 void* arguments:
+    - state pointer
+    - result pointer
+    """
+    void_sig = types.void(types.voidptr, types.voidptr)
+
+    wrapper_src = textwrap.dedent(f"""
+    @intrinsic
+    def impl(typingctx, state_arg, result_arg):
+        def codegen(context, builder, impl_sig, args):
+            state_type_llvm = context.get_value_type(state_ptr_type.dtype)
+            value_type_llvm = context.get_value_type(value_type)
+
+            state_ptr = builder.bitcast(args[0], state_type_llvm.as_pointer())
+            result_ptr = builder.bitcast(args[1], value_type_llvm.as_pointer())
+
+            sig = types.void(state_ptr_type, types.CPointer(value_type))
+            context.compile_internal(builder, deref_fn, sig, [state_ptr, result_ptr])
+
+            return context.get_dummy_value()
+        return void_sig, codegen
+
+    def wrapped_{deref_fn.__name__}(state_arg, result_arg):
+        return impl(state_arg, result_arg)
+    """)
+
+    local_dict = {
+        "types": types,
+        "state_ptr_type": state_ptr_type,
+        "value_type": value_type,
+        "deref_fn": deref_fn,
+        "intrinsic": intrinsic,
+        "void_sig": void_sig,
+    }
+    exec(wrapper_src, globals(), local_dict)
+
+    wrapper_func = local_dict[f"wrapped_{deref_fn.__name__}"]
+    wrapper_func.__globals__.update(local_dict)
+
+    return wrapper_func, void_sig
+
+
+def _create_output_dereference_wrapper(deref_fn, state_ptr_type, value_type):
+    """Creates a wrapper function for output iterator dereference that takes void* arguments.
+
+    The wrapper takes 2 void* arguments:
+    - state pointer
+    - value pointer (points to value)
+    """
+    void_sig = types.void(types.voidptr, types.voidptr)
+
+    wrapper_src = textwrap.dedent(f"""
+    @intrinsic
+    def impl(typingctx, state_arg, value_arg):
+        def codegen(context, builder, impl_sig, args):
+            state_type_llvm = context.get_value_type(state_ptr_type.dtype)
+            value_type_llvm = context.get_value_type(value_type)
+
+            state_ptr = builder.bitcast(args[0], state_type_llvm.as_pointer())
+            value_ptr = builder.bitcast(args[1], value_type_llvm.as_pointer())
+            value_val = builder.load(value_ptr)
+
+            sig = types.void(state_ptr_type, value_type)
+            context.compile_internal(builder, deref_fn, sig, [state_ptr, value_val])
+
+            return context.get_dummy_value()
+        return void_sig, codegen
+
+    def wrapped_{deref_fn.__name__}(state_arg, value_arg):
+        return impl(state_arg, value_arg)
+    """)
+
+    local_dict = {
+        "types": types,
+        "state_ptr_type": state_ptr_type,
+        "value_type": value_type,
+        "deref_fn": deref_fn,
+        "intrinsic": intrinsic,
+        "void_sig": void_sig,
+    }
+    exec(wrapper_src, globals(), local_dict)
+
+    wrapper_func = local_dict[f"wrapped_{deref_fn.__name__}"]
+    wrapper_func.__globals__.update(local_dict)
+
+    return wrapper_func, void_sig
+
+
 def to_cccl_op(op: Callable | OpKind, sig: Signature | None) -> Op:
     """Return an `Op` object corresponding to the given callable or well-known operation.
 
diff --git a/python/cuda_cccl/cuda/compute/iterators/_iterators.py b/python/cuda_cccl/cuda/compute/iterators/_iterators.py
diff --git a/python/cuda_cccl/tests/compute/test_zip_iterator.py b/python/cuda_cccl/tests/compute/test_zip_iterator.py

Original file line number	Diff line number	Diff line change
`@@ -47,7 +47,7 @@ struct input_iterator_t`
`47`	`47`
`48`	`48`	`__device__ input_iterator_t& operator+=(difference_type diff)`
`49`	`49`	`{`
`50`		`- Iterator.advance(&state, diff);`
	`50`	`+ Iterator.advance(&state, &diff);`
`51`	`51`	`return *this;`
`52`	`52`	`}`
`53`	`53`
Original file line number	Diff line number	Diff line change
`@@ -35,7 +35,7 @@ struct output_iterator_proxy_t`
`35`	`35`	`{`
`36`	`36`	`__device__ output_iterator_proxy_t& operator=(AssignT x)`
`37`	`37`	`{`
`38`		`- AssignF(&state, cuda::std::move(x));`
	`38`	`+ AssignF(&state, &x);`
`39`	`39`	`return *this;`
`40`	`40`	`}`
`41`	`41`
`@@ -59,7 +59,7 @@ struct output_iterator_t`
`59`	`59`
`60`	`60`	`__device__ output_iterator_t& operator+=(difference_type diff)`
`61`	`61`	`{`
`62`		`- Iterator.advance(&state, diff);`
	`62`	`+ Iterator.advance(&state, &diff);`
`63`	`63`	`return *this;`
`64`	`64`	`}`
`65`	`65`