cuda-cccl 0.3.1__cp310-cp310-manylinux_2_24_aarch64.whl → 0.3.2__cp310-cp310-manylinux_2_24_aarch64.whl

cuda/cccl/headers/include/thrust/system/cuda/detail/find.h

@@ -39,11 +39,13 @@
 #if _CCCL_HAS_CUDA_COMPILER()
 #  include <thrust/system/cuda/config.h>
 
-#  include <thrust/distance.h>
-#  include <thrust/iterator/counting_iterator.h>
-#  include <thrust/iterator/transform_iterator.h>
 #  include <thrust/system/cuda/detail/execution_policy.h>
 
+#  include <cuda/__iterator/counting_iterator.h>
+#  include <cuda/__iterator/transform_iterator.h>
+#  include <cuda/__iterator/zip_iterator.h>
+#  include <cuda/std/__iterator/distance.h>
+
 THRUST_NAMESPACE_BEGIN
 namespace cuda_cub
 {
@@ -62,7 +64,6 @@ InputIt _CCCL_HOST_DEVICE find(execution_policy<Derived>& policy, InputIt first,
 }; // namespace cuda_cub
 THRUST_NAMESPACE_END
 
-#  include <thrust/iterator/zip_iterator.h>
 #  include <thrust/system/cuda/detail/reduce.h>
 
 THRUST_NAMESPACE_BEGIN
@@ -92,109 +93,13 @@ struct functor
     }
   }
 };
-
-template <class ValueType, class InputIt, class UnaryOp>
-struct transform_input_iterator_t
-{
-  using self_t            = transform_input_iterator_t;
-  using difference_type   = thrust::detail::it_difference_t<InputIt>;
-  using value_type        = ValueType;
-  using pointer           = void;
-  using reference         = value_type;
-  using iterator_category = ::cuda::std::random_access_iterator_tag;
-
-  InputIt input;
-  mutable UnaryOp op;
-
-  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE transform_input_iterator_t(InputIt input, UnaryOp op)
-      : input(input)
-      , op(op)
-  {}
-
-  transform_input_iterator_t(const self_t&) = default;
-
-  // UnaryOp might not be copy assignable, such as when it is a lambda.  Define
-  // an explicit copy assignment operator that doesn't try to assign it.
-  _CCCL_HOST_DEVICE self_t& operator=(const self_t& o)
-  {
-    input = o.input;
-    return *this;
-  }
-
-  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE self_t operator++(int)
-  {
-    self_t retval = *this;
-    ++input;
-    return retval;
-  }
-
-  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE self_t operator++()
-  {
-    ++input;
-    return *this;
-  }
-
-  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE reference operator*() const
-  {
-    thrust::detail::it_value_t<InputIt> x = *input;
-    return op(x);
-  }
-
-  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE reference operator*()
-  {
-    thrust::detail::it_value_t<InputIt> x = *input;
-    return op(x);
-  }
-
-  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE self_t operator+(difference_type n) const
-  {
-    return self_t(input + n, op);
-  }
-
-  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE self_t& operator+=(difference_type n)
-  {
-    input += n;
-    return *this;
-  }
-
-  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE self_t operator-(difference_type n) const
-  {
-    return self_t(input - n, op);
-  }
-
-  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE self_t& operator-=(difference_type n)
-  {
-    input -= n;
-    return *this;
-  }
-
-  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE difference_type operator-(self_t other) const
-  {
-    return input - other.input;
-  }
-
-  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE reference operator[](difference_type n) const
-  {
-    return op(input[n]);
-  }
-
-  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE bool operator==(const self_t& rhs) const
-  {
-    return (input == rhs.input);
-  }
-
-  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE bool operator!=(const self_t& rhs) const
-  {
-    return (input != rhs.input);
-  }
-};
 } // namespace __find_if
 
 template <class Derived, class InputIt, class Size, class Predicate>
 InputIt _CCCL_HOST_DEVICE
 find_if_n(execution_policy<Derived>& policy, InputIt first, Size num_items, Predicate predicate)
 {
-  using result_type = typename thrust::tuple<bool, Size>;
+  using result_type = ::cuda::std::tuple<bool, Size>;
 
   // empty sequence
   if (num_items == 0)
@@ -212,27 +117,20 @@ find_if_n(execution_policy<Derived>& policy, InputIt first, Size num_items, Pred
   const Size interval_threshold = 1 << 20;
   const Size interval_size      = (::cuda::std::min) (interval_threshold, num_items);
 
-  // FIXME(bgruber): we should also be able to use transform_iterator here, but it makes nvc++ hang. See:
-  // https://github.com/NVIDIA/cccl/issues/3594. The problem does not occur with nvcc, so we could not add a test :/
-  using XfrmIterator = __find_if::transform_input_iterator_t<bool, InputIt, Predicate>;
-  // using XfrmIterator  = transform_iterator<Predicate, InputIt>;
-  using IteratorTuple = thrust::tuple<XfrmIterator, counting_iterator<Size>>;
-  using ZipIterator   = thrust::zip_iterator<IteratorTuple>;
-
-  IteratorTuple iter_tuple = thrust::make_tuple(XfrmIterator(first, predicate), counting_iterator<Size>(0));
-
-  ZipIterator begin = thrust::make_zip_iterator(iter_tuple);
-  ZipIterator end   = begin + num_items;
+  const auto begin = ::cuda::make_zip_iterator(
+    ::cuda::make_transform_iterator(try_unwrap_contiguous_iterator(first), predicate),
+    ::cuda::counting_iterator<Size>(0));
+  const auto end = begin + num_items;
 
-  for (ZipIterator interval_begin = begin; interval_begin < end; interval_begin += interval_size)
+  for (auto interval_begin = begin; interval_begin < end; interval_begin += interval_size)
   {
-    ZipIterator interval_end = interval_begin + interval_size;
+    auto interval_end = interval_begin + interval_size;
     if (end < interval_end)
     {
       interval_end = end;
     } // end if
 
-    result_type result = reduce(
+    const result_type result = reduce(
       policy, interval_begin, interval_end, result_type(false, interval_end - begin), __find_if::functor<result_type>());
 
     // see if we found something

cuda/cccl/headers/include/thrust/system/cuda/detail/mismatch.h

@@ -73,12 +73,14 @@ struct transform_pair_of_input_iterators_t
   using value_type        = ValueType;
   using pointer           = void;
   using reference         = value_type;
-  using iterator_category = std::random_access_iterator_tag;
+  using iterator_category = ::cuda::std::random_access_iterator_tag;
 
   InputIt1 input1;
   InputIt2 input2;
   mutable BinaryOp op;
 
+  transform_pair_of_input_iterators_t() = default;
+
   _CCCL_HOST_DEVICE _CCCL_FORCEINLINE
   transform_pair_of_input_iterators_t(InputIt1 input1_, InputIt2 input2_, BinaryOp op_)
       : input1(input1_)
@@ -107,7 +109,7 @@ struct transform_pair_of_input_iterators_t
   }
 
   /// Prefix increment
-  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE self_t operator++()
+  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE self_t& operator++()
   {
     ++input1;
     ++input2;
@@ -177,6 +179,10 @@ struct transform_pair_of_input_iterators_t
     return (input1 != rhs.input1) || (input2 != rhs.input2);
   }
 
+  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE bool operator<(const self_t& rhs) const
+  {
+    return input1 < rhs.input1;
+  }
 }; // struct transform_pair_of_input_iterators_t
 } // namespace detail
 

cuda/cccl/headers/include/thrust/type_traits/is_contiguous_iterator.h

@@ -79,7 +79,7 @@ namespace detail
 template <typename Iterator>
 inline constexpr bool is_libcxx_wrap_iter_v = false;
 
-#if defined(_LIBCPP_VERSION)
+#if _CCCL_HOST_STD_LIB(LIBCXX)
 template <typename Iterator>
 inline constexpr bool is_libcxx_wrap_iter_v<
 #  if _LIBCPP_VERSION < 14000
@@ -88,23 +88,23 @@ inline constexpr bool is_libcxx_wrap_iter_v<
   std::__wrap_iter<Iterator>
 #  endif
   > = true;
-#endif
+#endif // _CCCL_HOST_STD_LIB(LIBCXX)
 
 template <typename Iterator>
 inline constexpr bool is_libstdcxx_normal_iterator_v = false;
 
-#if defined(__GLIBCXX__)
+#if _CCCL_HOST_STD_LIB(LIBSTDCXX)
 template <typename Iterator, typename Container>
 inline constexpr bool is_libstdcxx_normal_iterator_v<::__gnu_cxx::__normal_iterator<Iterator, Container>> = true;
-#endif
+#endif // _CCCL_HOST_STD_LIB(LIBSTDCXX)
 
-#if _CCCL_COMPILER(MSVC)
+#if _CCCL_HOST_STD_LIB(STL)
 template <typename Iterator>
 inline constexpr bool is_msvc_contiguous_iterator_v = ::cuda::std::is_pointer_v<::std::_Unwrapped_t<Iterator>>;
-#else
+#else // ^^^ _CCCL_HOST_STD_LIB(STL) ^^^ / vvv !_CCCL_HOST_STD_LIB(STL) vvv
 template <typename Iterator>
 inline constexpr bool is_msvc_contiguous_iterator_v = false;
-#endif
+#endif // ^^^ !_CCCL_HOST_STD_LIB(STL) ^^^
 
 template <typename Iterator>
 inline constexpr bool is_contiguous_iterator_impl_v =

cuda/compute/__init__.py

@@ -32,6 +32,7 @@ from .iterators import (
     CacheModifiedInputIterator,
     ConstantIterator,
     CountingIterator,
+    PermutationIterator,
     ReverseIterator,
     TransformIterator,
     TransformOutputIterator,
@@ -63,6 +64,7 @@ __all__ = [
     "make_unique_by_key",
     "merge_sort",
     "OpKind",
+    "PermutationIterator",
     "radix_sort",
     "reduce_into",
     "ReverseIterator",

cuda/compute/_bindings.pyi

@@ -57,6 +57,12 @@ class SortOrder(IntEnum):
     ASCENDING = ...
     DESCENDING = ...
 
+class InitKind(IntEnum):
+    _value_: int
+    NO_INIT = ...
+    FUTURE_VALUE_INIT = ...
+    VALUE_INIT = ...
+
 class Op:
     def __init__(
         self,
@@ -133,6 +139,8 @@ class Iterator:
     def state(self, value) -> None: ...
     @property
     def type(self) -> IteratorKind: ...
+    @property
+    def value_type(self) -> TypeInfo: ...
     def as_bytes(self) -> bytes: ...
     def is_kind_pointer(self) -> bool: ...
     def is_kind_iterator(self) -> bool: ...
@@ -197,8 +205,9 @@ class DeviceScanBuildResult:
         d_in: Iterator,
         d_out: Iterator,
         binary_op: Op,
-        h_init: Value,
+        init_type: TypeInfo,
         force_inclusive: bool,
+        init_kind: InitKind,
         info: CommonData,
     ): ...
     def compute_inclusive(
@@ -223,6 +232,39 @@ class DeviceScanBuildResult:
         h_init: Value,
         stream,
     ) -> int: ...
+    def compute_inclusive_future_value(
+        self,
+        temp_storage_ptr: int | None,
+        temp_storage_nbytes: int,
+        d_in: Iterator,
+        d_out: Iterator,
+        num_items: int,
+        binary_op: Op,
+        h_init: Iterator,
+        stream,
+    ) -> int: ...
+    def compute_exclusive_future_value(
+        self,
+        temp_storage_ptr: int | None,
+        temp_storage_nbytes: int,
+        d_in: Iterator,
+        d_out: Iterator,
+        num_items: int,
+        binary_op: Op,
+        h_init: Iterator,
+        stream,
+    ) -> int: ...
+    def compute_inclusive_no_init(
+        self,
+        temp_storage_ptr: int | None,
+        temp_storage_nbytes: int,
+        d_in: Iterator,
+        d_out: Iterator,
+        num_items: int,
+        binary_op: Op,
+        h_init: None,
+        stream,
+    ) -> int: ...
 
 # ---------------------
 # DeviceSegmentedReduce

cuda/compute/_bindings_impl.pyx

@@ -120,6 +120,10 @@ cdef extern from "cccl/c/types.h":
         ASCENDING "CCCL_ASCENDING"
         DESCENDING "CCCL_DESCENDING"
 
+    cpdef enum cccl_init_kind_t:
+        VALUE_INIT "CCCL_VALUE_INIT"
+        FUTURE_VALUE_INIT "CCCL_FUTURE_VALUE_INIT"
+        NO_INIT "CCCL_NO_INIT"
 
 cdef void arg_type_check(
     str arg_name,
@@ -136,6 +140,7 @@ OpKind = cccl_op_kind_t
 TypeEnum = cccl_type_enum
 IteratorKind = cccl_iterator_kind_t
 SortOrder = cccl_sort_order_t
+InitKind = cccl_init_kind_t
 
 cdef void _validate_alignment(int alignment) except *:
     """
@@ -724,6 +729,11 @@ cdef class Iterator:
         else:
             return IteratorKind.ITERATOR
 
+    @property
+    def value_type(self):
+        cdef cccl_type_info type_info = self.iter_data.value_type
+        return TypeInfo(type_info.size, type_info.alignment, type_info.type)
+
     def is_kind_pointer(self):
         cdef cccl_iterator_kind_t it_kind = self.iter_data.type
         return (it_kind == cccl_iterator_kind_t.POINTER)
@@ -947,8 +957,9 @@ cdef extern from "cccl/c/scan.h":
         cccl_iterator_t,
         cccl_iterator_t,
         cccl_op_t,
-        cccl_value_t,
+        cccl_type_info,
         _Bool,
+        cccl_init_kind_t,
         int, int, const char*, const char*, const char*, const char*
     ) nogil
 
@@ -976,6 +987,41 @@ cdef extern from "cccl/c/scan.h":
         CUstream
     ) nogil
 
+    cdef CUresult cccl_device_exclusive_scan_future_value(
+        cccl_device_scan_build_result_t,
+        void *,
+        size_t *,
+        cccl_iterator_t,
+        cccl_iterator_t,
+        uint64_t,
+        cccl_op_t,
+        cccl_iterator_t,
+        CUstream
+    ) nogil
+
+    cdef CUresult cccl_device_inclusive_scan_future_value(
+        cccl_device_scan_build_result_t,
+        void *,
+        size_t *,
+        cccl_iterator_t,
+        cccl_iterator_t,
+        uint64_t,
+        cccl_op_t,
+        cccl_iterator_t,
+        CUstream
+    ) nogil
+
+    cdef CUresult cccl_device_inclusive_scan_no_init(
+        cccl_device_scan_build_result_t,
+        void *,
+        size_t *,
+        cccl_iterator_t,
+        cccl_iterator_t,
+        uint64_t,
+        cccl_op_t,
+        CUstream
+    ) nogil
+
     cdef CUresult cccl_device_scan_cleanup(
         cccl_device_scan_build_result_t*
     ) nogil
@@ -989,8 +1035,9 @@ cdef class DeviceScanBuildResult:
         Iterator d_in,
         Iterator d_out,
         Op op,
-        Value h_init,
+        TypeInfo init_type,
         bint force_inclusive,
+        cccl_init_kind_t init_kind,
         CommonData common_data
     ):
         cdef CUresult status = -1
@@ -1008,8 +1055,9 @@ cdef class DeviceScanBuildResult:
                 d_in.iter_data,
                 d_out.iter_data,
                 op.op_data,
-                h_init.value_data,
+                init_type.type_info,
                 force_inclusive,
+                init_kind,
                 cc_major,
                 cc_minor,
                 cub_path,
@@ -1035,7 +1083,7 @@ cdef class DeviceScanBuildResult:
         Iterator d_out,
         size_t num_items,
         Op op,
-        Value h_init,
+        Value init_value,
         stream
     ):
         cdef CUresult status = -1
@@ -1052,7 +1100,7 @@ cdef class DeviceScanBuildResult:
                 d_out.iter_data,
                 <uint64_t>num_items,
                 op.op_data,
-                h_init.value_data,
+                init_value.value_data,
                 c_stream
             )
         if status != 0:
@@ -1069,7 +1117,7 @@ cdef class DeviceScanBuildResult:
         Iterator d_out,
         size_t num_items,
         Op op,
-        Value h_init,
+        Value init_value,
         stream
     ):
         cdef CUresult status = -1
@@ -1086,7 +1134,7 @@ cdef class DeviceScanBuildResult:
                 d_out.iter_data,
                 <uint64_t>num_items,
                 op.op_data,
-                h_init.value_data,
+                init_value.value_data,
                 c_stream
             )
         if status != 0:
@@ -1095,6 +1143,107 @@ cdef class DeviceScanBuildResult:
             )
         return storage_sz
 
+    cpdef int compute_inclusive_future_value(
+        DeviceScanBuildResult self,
+        temp_storage_ptr,
+        temp_storage_bytes,
+        Iterator d_in,
+        Iterator d_out,
+        size_t num_items,
+        Op op,
+        Iterator init_value,
+        stream
+    ):
+        cdef CUresult status = -1
+        cdef void *storage_ptr = (<void *><uintptr_t>temp_storage_ptr) if temp_storage_ptr else NULL
+        cdef size_t storage_sz = <size_t>temp_storage_bytes
+        cdef CUstream c_stream = <CUstream><uintptr_t>(stream) if stream else NULL
+
+        with nogil:
+            status = cccl_device_inclusive_scan_future_value(
+                self.build_data,
+                storage_ptr,
+                &storage_sz,
+                d_in.iter_data,
+                d_out.iter_data,
+                <uint64_t>num_items,
+                op.op_data,
+                init_value.iter_data,
+                c_stream
+            )
+        if status != 0:
+            raise RuntimeError(
+                f"Failed executing inclusive scan, error code: {status}"
+            )
+        return storage_sz
+
+    cpdef int compute_exclusive_future_value(
+        DeviceScanBuildResult self,
+        temp_storage_ptr,
+        temp_storage_bytes,
+        Iterator d_in,
+        Iterator d_out,
+        size_t num_items,
+        Op op,
+        Iterator init_value,
+        stream
+    ):
+        cdef CUresult status = -1
+        cdef void *storage_ptr = (<void *><uintptr_t>temp_storage_ptr) if temp_storage_ptr else NULL
+        cdef size_t storage_sz = <size_t>temp_storage_bytes
+        cdef CUstream c_stream = <CUstream><uintptr_t>(stream) if stream else NULL
+
+        with nogil:
+            status = cccl_device_exclusive_scan_future_value(
+                self.build_data,
+                storage_ptr,
+                &storage_sz,
+                d_in.iter_data,
+                d_out.iter_data,
+                <uint64_t>num_items,
+                op.op_data,
+                init_value.iter_data,
+                c_stream
+            )
+        if status != 0:
+            raise RuntimeError(
+                f"Failed executing exclusive scan, error code: {status}"
+            )
+        return storage_sz
+
+    cpdef int compute_inclusive_no_init(
+        DeviceScanBuildResult self,
+        temp_storage_ptr,
+        temp_storage_bytes,
+        Iterator d_in,
+        Iterator d_out,
+        size_t num_items,
+        Op op,
+        object init_value,
+        stream
+    ):
+        cdef CUresult status = -1
+        cdef void *storage_ptr = (<void *><uintptr_t>temp_storage_ptr) if temp_storage_ptr else NULL
+        cdef size_t storage_sz = <size_t>temp_storage_bytes
+        cdef CUstream c_stream = <CUstream><uintptr_t>(stream) if stream else NULL
+
+        with nogil:
+            status = cccl_device_inclusive_scan_no_init(
+                self.build_data,
+                storage_ptr,
+                &storage_sz,
+                d_in.iter_data,
+                d_out.iter_data,
+                <uint64_t>num_items,
+                op.op_data,
+                c_stream
+            )
+        if status != 0:
+            raise RuntimeError(
+                f"Failed executing inclusive scan, error code: {status}"
+            )
+        return storage_sz
+
     def _get_cubin(self):
         return PyBytes_FromStringAndSize(
             <const char*>self.build_data.cubin,