Support nested structs in cuda.compute (#6353)

* Enable nesting of gpu_structs

* Update existing usage of gpu_struct to

* Enable nesting of ZipIterators

* Add tests for nested structs and ZipIterators

* Reorganize struct examples and add new examples for nested structs

* Address review comments

---------

Co-authored-by: Ashwin Srinath <shwina@users.noreply.github.com>

Author: shwina

python/cuda_cccl/cuda/compute/_cccl_interop.py

@@ -313,7 +313,7 @@ def to_cccl_op(op: Callable | OpKind, sig: Signature | None) -> Op:
 
 
 def get_value_type(d_in: IteratorBase | DeviceArrayLike):
-    from .struct import gpu_struct_from_numpy_dtype
+    from .struct import gpu_struct
 
     if isinstance(d_in, IteratorBase):
         return d_in.value_type
@@ -323,7 +323,7 @@ def get_value_type(d_in: IteratorBase | DeviceArrayLike):
         # types directly, as those are passed by pointer to device
         # functions. Instead, we create an anonymous struct type
         # which has the appropriate pass-by-value semantics.
-        return as_numba_type(gpu_struct_from_numpy_dtype("anonymous", dtype))
+        return as_numba_type(gpu_struct(dtype))
     return numba.from_dtype(dtype)
 
 

python/cuda_cccl/cuda/compute/algorithms/_scan.py

@@ -71,14 +71,18 @@ def __init__(
                 init_value_type_info = self.d_in_cccl.value_type
 
             case _bindings.InitKind.FUTURE_VALUE_INIT:
-                self.init_value_cccl = cccl.to_cccl_input_iter(init_value)
+                self.init_value_cccl = cccl.to_cccl_input_iter(
+                    cast(DeviceArrayLike, init_value)
+                )
                 value_type = numba.from_dtype(
                     protocols.get_dtype(cast(DeviceArrayLike, init_value))
                 )
                 init_value_type_info = self.init_value_cccl.value_type
 
             case _bindings.InitKind.VALUE_INIT:
-                self.init_value_cccl = cccl.to_cccl_value(init_value)
+                self.init_value_cccl = cccl.to_cccl_value(
+                    cast(np.ndarray | GpuStruct, init_value)
+                )
                 value_type = (
                     numba.from_dtype(init_value.dtype)
                     if isinstance(init_value, np.ndarray)
@@ -141,7 +145,9 @@ def __call__(
 
             case _bindings.InitKind.VALUE_INIT:
                 self.init_value_cccl = cast(_bindings.Value, self.init_value_cccl)
-                self.init_value_cccl.state = to_cccl_value_state(init_value)
+                self.init_value_cccl.state = to_cccl_value_state(
+                    cast(np.ndarray | GpuStruct, init_value)
+                )
 
         stream_handle = validate_and_get_stream(stream)
 

python/cuda_cccl/cuda/compute/algorithms/_transform.py

@@ -273,7 +273,7 @@ def unary_transform(
         When working with custom struct types, you need to provide type annotations
         to help with type inference. See the binary transform struct example for reference:
 
-        .. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/transform/binary_transform_struct.py
+        .. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/struct/struct_transform.py
            :language: python
            :start-after: # example-begin
 
@@ -312,7 +312,7 @@ def binary_transform(
         When working with custom struct types, you need to provide type annotations
         to help with type inference. See the following example:
 
-        .. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/transform/binary_transform_struct.py
+        .. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/struct/struct_transform.py
            :language: python
            :start-after: # example-begin
 

python/cuda_cccl/cuda/compute/iterators/_factories.py

@@ -232,6 +232,12 @@ def ZipIterator(*iterators):
             :language: python
             :start-after: # example-begin
 
+        ZipIterator can also be used with nested gpu_struct types:
+
+        .. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/struct/nested_struct_zip_iterator.py
+            :language: python
+            :start-after: # example-begin
+
     Args:
         *iterators: Variable number of iterators to zip (at least 1)
 

python/cuda_cccl/cuda/compute/iterators/_iterators.py

@@ -8,7 +8,7 @@
 import numpy as np
 from llvmlite import ir
 from numba import cuda, types
-from numba.core.extending import intrinsic, overload
+from numba.core.extending import as_numba_type, intrinsic, overload
 from numba.core.typing.ctypes_utils import to_ctypes
 from numba.cuda.dispatcher import CUDADispatcher
 
@@ -680,6 +680,8 @@ def make_permutation_iterator(values, indices):
     Returns:
         PermutationIterator: Iterator that yields permuted values
     """
+    from ..struct import make_struct_type
+
     # Convert arrays to iterators if needed
     if hasattr(values, "__cuda_array_interface__"):
         values = pointer(values, numba.from_dtype(get_dtype(values)))
@@ -715,22 +717,20 @@ def make_permutation_iterator(values, indices):
     # The cvalue and state for PermutationIterator are
     # structs composed of the cvalues and states of the
     # value and index iterators.
-    from ..struct import gpu_struct_from_numba_types
-
     class PermutationCValueStruct(ctypes.Structure):
         _fields_ = [
             ("value_state", values.cvalue.__class__),
             ("index_state", indices.cvalue.__class__),
         ]
 
-    PermutationState = gpu_struct_from_numba_types(
+    PermutationState = make_struct_type(
         "PermutationState",
-        ("value_state", "index_state"),
-        (values_state_type, indices.state_type),
+        field_names=("value_state", "index_state"),
+        field_types=(values_state_type, indices.state_type),
     )
 
     cvalue = PermutationCValueStruct(values.cvalue, indices.cvalue)
-    state_type = PermutationState._numba_type
+    state_type = as_numba_type(PermutationState)
     value_type = value_dtype
 
     # Define intrinsics for accessing struct fields

python/cuda_cccl/cuda/compute/iterators/_zip_iterator.py

@@ -4,14 +4,12 @@
 
 import ctypes
 
-import numba
 from llvmlite import ir  # noqa: F401
 from numba import cuda, types  # noqa: F401
 from numba.core.datamodel.registry import default_manager  # noqa: F401
-from numba.core.extending import intrinsic  # noqa: F401
+from numba.core.extending import as_numba_type, intrinsic  # noqa: F401
 
-from .._utils.protocols import get_dtype
-from ..struct import gpu_struct_from_numba_types
+from ..struct import make_struct_type
 from ._iterators import (
     IteratorBase,
     IteratorKind,
@@ -36,22 +34,18 @@ def _get_zip_iterator_metadata(iterators):
     # this iterator's state is a struct composed of the states of the input iterators:
     state_field_names = tuple(f"state_{i}" for i in range(n_iterators))
     state_field_types = tuple(it.state_type for it in iterators)
-    ZipState = gpu_struct_from_numba_types(
-        "ZipState", state_field_names, state_field_types
-    )
+    ZipState = make_struct_type("ZipState", state_field_names, state_field_types)
 
     # this iterator's value is a struct composed of the values of the input iterators:
     value_field_names = tuple(f"value_{i}" for i in range(n_iterators))
     value_field_types = tuple(it.value_type for it in iterators)
-    ZipValue = gpu_struct_from_numba_types(
-        "ZipValue", value_field_names, value_field_types
-    )
+    ZipValue = make_struct_type("ZipValue", value_field_names, value_field_types)
 
     cvalue = ZipCValueStruct(
         **{f"iter_{i}": it.cvalue for i, it in enumerate(iterators)}
     )
-    state_type = ZipState._numba_type
-    value_type = ZipValue._numba_type
+    state_type = as_numba_type(ZipState)
+    value_type = as_numba_type(ZipValue)
     return cvalue, state_type, value_type, ZipValue
 
 
@@ -64,6 +58,15 @@ def _get_advance_and_dereference_functions(iterators):
 
     n_iterators = len(iterators)
 
+    # Create a local namespace for this zip iterator to avoid polluting globals
+    # and prevent name collisions when nesting zip iterators
+    local_ns = {
+        "intrinsic": intrinsic,
+        "types": types,
+        "ir": ir,
+        "default_manager": default_manager,
+    }
+
     # Within the advance and dereference methods of this iterator, we need a way
     # to get pointers to the fields of the state struct (advance and dereference),
     # and the value struct (dereference). This needs `n` custom intrinsics, one
@@ -80,30 +83,29 @@ def codegen(context, builder, sig, args):
         # Use GEP to get pointer to field at index {field_idx}
         field_ptr = builder.gep(
             struct_ptr,
-            [ir.Constant(ir.IntType(32), 0), ir.Constant(ir.IntType(32), {field_idx})],
+            [ir.Constant(ir.IntType(32), 0), ir.Constant(
+                ir.IntType(32), {field_idx})],
         )
         return field_ptr
 
     struct_model = default_manager.lookup(struct_ptr_type.dtype)
     field_type = struct_model._members[{field_idx}]
     return types.CPointer(field_type)(struct_ptr_type), codegen
 """
-        # Execute the code to create the intrinsic function in global namespace
-        exec(intrinsic_code, globals())
+        # Execute the code to create the intrinsic function in local namespace
+        exec(intrinsic_code, local_ns)
 
     # Now we can define the advance and dereference methods of this iterator,
     # which also need to be defined dynamically because they will use the
     # intrinsics defined above.
     for i, it in enumerate(iterators):
-        globals()[f"advance_{i}"] = cuda.jit(it.advance, device=True)
-        globals()[f"input_dereference_{i}"] = cuda.jit(
-            it.input_dereference, device=True
-        )
+        local_ns[f"advance_{i}"] = cuda.jit(it.advance, device=True)
+        local_ns[f"input_dereference_{i}"] = cuda.jit(it.input_dereference, device=True)
         # Also compile output_dereference if available
         try:
             output_deref = it.output_dereference
             if output_deref is not None:
-                globals()[f"output_dereference_{i}"] = cuda.jit(
+                local_ns[f"output_dereference_{i}"] = cuda.jit(
                     output_deref, device=True
                 )
         except AttributeError:
@@ -152,12 +154,12 @@ def input_dereference(state, result):
 {chr(10).join(input_dereference_lines)}
 """  # chr(10) is '\n'
 
-    # Execute the method codes:
-    exec(advance_method_code, globals())
-    exec(input_dereference_method_code, globals())
+    # Execute the method codes in local namespace:
+    exec(advance_method_code, local_ns)
+    exec(input_dereference_method_code, local_ns)
 
-    advance_func = globals()["input_advance"]
-    input_dereference_func = globals()["input_dereference"]
+    advance_func = local_ns["input_advance"]
+    input_dereference_func = local_ns["input_dereference"]
 
     # Generate output_dereference if all iterators support it
     output_dereference_func = None
@@ -167,8 +169,8 @@ def output_dereference(state, x):
     # Write to each iterator using dynamically created field pointer functions
 {chr(10).join(output_dereference_lines)}
 """
-        exec(output_dereference_method_code, globals())
-        output_dereference_func = globals()["output_dereference"]
+        exec(output_dereference_method_code, local_ns)
+        output_dereference_func = local_ns["output_dereference"]
 
     return advance_func, input_dereference_func, output_dereference_func
 
@@ -183,14 +185,16 @@ def make_zip_iterator(*iterators):
     Returns:
         ZipIterator: Iterator that combines all input iterators
     """
+    from .._cccl_interop import get_value_type
+
     if len(iterators) < 1:
         raise ValueError("At least 1 iterator is required")
 
     # Convert arrays to iterators if needed
     processed_iterators = []
     for it in iterators:
         if hasattr(it, "__cuda_array_interface__"):
-            it = pointer(it, numba.from_dtype(get_dtype(it)))
+            it = pointer(it, get_value_type(it))
         processed_iterators.append(it)
 
     # Validate all are iterators

python/cuda_cccl/cuda/compute/numba_utils.py

@@ -2,6 +2,7 @@
 
 import numba
 from numba import cuda
+from numba.core.extending import as_numba_type
 from numpy.typing import DTypeLike
 
 from .typing import GpuStruct
@@ -16,8 +17,8 @@ def to_numba_type(tp: GpuStruct | DTypeLike) -> numba.types.Type:
     """
     Convert a GpuStruct or DtypeLike to a numba type.
     """
-    if hasattr(tp, "_numba_type"):
-        return tp._numba_type  # type: ignore[union-attr]
+    if value := as_numba_type.lookup.get(tp):
+        return value
     return numba.from_dtype(tp)