cuda-cccl 0.3.0__cp310-cp310-manylinux_2_26_x86_64.whl → 0.3.1__cp310-cp310-manylinux_2_26_x86_64.whl

cuda/cccl/cooperative/__init__.py

@@ -1,3 +1,9 @@
 # Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
 #
-# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+# SPDX-License  -Identifier: Apache-2.0 WITH LLVM-exception
+
+from . import experimental
+
+__all__ = [
+    "experimental",
+]

cuda/cccl/cooperative/experimental/__init__.py

@@ -1,8 +1,24 @@
-# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+# Copyright (c) 2025, NVIDIA CORPORATION.
 #
-# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# alias for backwards compatibility
+
+from warnings import warn
 
-from cuda.cccl.cooperative.experimental import block, warp
-from cuda.cccl.cooperative.experimental._types import StatefulFunction
+from cuda.coop import *  # noqa: F403
 
-__all__ = ["block", "warp", "StatefulFunction"]
+warn(
+    "The module cuda.cccl.cooperative.experimental is deprecated. Use cuda.coop instead.",
+    FutureWarning,
+)

cuda/cccl/headers/include/cub/agent/agent_adjacent_difference.cuh

@@ -64,9 +64,7 @@ struct AgentAdjacentDifferencePolicy
   static constexpr cub::BlockStoreAlgorithm STORE_ALGORITHM = _STORE_ALGORITHM;
 };
 
-namespace detail
-{
-namespace adjacent_difference
+namespace detail::adjacent_difference
 {
 
 template <typename Policy,
@@ -256,7 +254,6 @@ struct AgentDifferenceInit
   }
 };
 
-} // namespace adjacent_difference
-} // namespace detail
+} // namespace detail::adjacent_difference
 
 CUB_NAMESPACE_END

cuda/cccl/headers/include/cub/agent/agent_batch_memcpy.cuh

@@ -62,9 +62,7 @@
 
 CUB_NAMESPACE_BEGIN
 
-namespace detail
-{
-namespace batch_memcpy
+namespace detail::batch_memcpy
 {
 template <bool PTR_IS_FOUR_BYTE_ALIGNED>
 _CCCL_FORCEINLINE _CCCL_DEVICE void
@@ -1179,7 +1177,6 @@ private:
   // buffers
   BLevBlockOffsetTileState blev_block_scan_state;
 };
-} // namespace batch_memcpy
-} // namespace detail
+} // namespace detail::batch_memcpy
 
 CUB_NAMESPACE_END

cuda/cccl/headers/include/cub/agent/agent_for.cuh

@@ -42,9 +42,7 @@
 
 CUB_NAMESPACE_BEGIN
 
-namespace detail
-{
-namespace for_each
+namespace detail::for_each
 {
 
 template <int BlockThreads, int ItemsPerThread>
@@ -78,7 +76,6 @@ struct agent_block_striped_t
   }
 };
 
-} // namespace for_each
-} // namespace detail
+} // namespace detail::for_each
 
 CUB_NAMESPACE_END

cuda/cccl/headers/include/cub/agent/agent_merge.cuh

@@ -53,14 +53,8 @@ struct agent_t
   using policy = Policy;
 
   // key and value type are taken from the first input sequence (consistent with old Thrust behavior)
-  using key_type  = it_value_t<KeysIt1>;
-  using item_type = it_value_t<ItemsIt1>;
-
-  using keys_load_it1  = try_make_cache_modified_iterator_t<Policy::LOAD_MODIFIER, KeysIt1>;
-  using keys_load_it2  = try_make_cache_modified_iterator_t<Policy::LOAD_MODIFIER, KeysIt2>;
-  using items_load_it1 = try_make_cache_modified_iterator_t<Policy::LOAD_MODIFIER, ItemsIt1>;
-  using items_load_it2 = try_make_cache_modified_iterator_t<Policy::LOAD_MODIFIER, ItemsIt2>;
-
+  using key_type          = it_value_t<KeysIt1>;
+  using item_type         = it_value_t<ItemsIt1>;
   using block_store_keys  = typename BlockStoreType<Policy, KeysOutputIt, key_type>::type;
   using block_store_items = typename BlockStoreType<Policy, ItemsOutputIt, item_type>::type;
 
@@ -84,11 +78,11 @@ struct agent_t
 
   // Per thread data
   temp_storages& storage;
-  keys_load_it1 keys1_in;
-  items_load_it1 items1_in;
+  KeysIt1 keys1_in;
+  ItemsIt1 items1_in;
   Offset keys1_count;
-  keys_load_it2 keys2_in;
-  items_load_it2 items2_in;
+  KeysIt2 keys2_in;
+  ItemsIt2 items2_in;
   Offset keys2_count;
   KeysOutputIt keys_out;
   ItemsOutputIt items_out;
@@ -128,10 +122,14 @@ struct agent_t
     }
 
     key_type keys_loc[items_per_thread];
-    merge_sort::gmem_to_reg<threads_per_block, IsFullTile>(
-      keys_loc, keys1_in + keys1_beg, keys2_in + keys2_beg, keys1_count_tile, keys2_count_tile);
-    merge_sort::reg_to_shared<threads_per_block>(&storage.keys_shared[0], keys_loc);
-    __syncthreads();
+    {
+      auto keys1_in_cm = try_make_cache_modified_iterator<Policy::LOAD_MODIFIER>(keys1_in);
+      auto keys2_in_cm = try_make_cache_modified_iterator<Policy::LOAD_MODIFIER>(keys2_in);
+      merge_sort::gmem_to_reg<threads_per_block, IsFullTile>(
+        keys_loc, keys1_in_cm + keys1_beg, keys2_in_cm + keys2_beg, keys1_count_tile, keys2_count_tile);
+      merge_sort::reg_to_shared<threads_per_block>(&storage.keys_shared[0], keys_loc);
+      __syncthreads();
+    }
 
     // now find the merge path for each of thread.
     // we can use int type here, because the number of items in shared memory is limited
@@ -186,11 +184,15 @@ struct agent_t
     if constexpr (have_items)
     {
       item_type items_loc[items_per_thread];
-      merge_sort::gmem_to_reg<threads_per_block, IsFullTile>(
-        items_loc, items1_in + keys1_beg, items2_in + keys2_beg, keys1_count_tile, keys2_count_tile);
-      __syncthreads(); // block_store_keys above uses SMEM, so make sure all threads are done before we write to it
-      merge_sort::reg_to_shared<threads_per_block>(&storage.items_shared[0], items_loc);
-      __syncthreads();
+      {
+        auto items1_in_cm = try_make_cache_modified_iterator<Policy::LOAD_MODIFIER>(items1_in);
+        auto items2_in_cm = try_make_cache_modified_iterator<Policy::LOAD_MODIFIER>(items2_in);
+        merge_sort::gmem_to_reg<threads_per_block, IsFullTile>(
+          items_loc, items1_in_cm + keys1_beg, items2_in_cm + keys2_beg, keys1_count_tile, keys2_count_tile);
+        __syncthreads(); // block_store_keys above uses SMEM, so make sure all threads are done before we write to it
+        merge_sort::reg_to_shared<threads_per_block>(&storage.items_shared[0], items_loc);
+        __syncthreads();
+      }
 
       // gather items from shared mem
       _CCCL_PRAGMA_UNROLL_FULL()

cuda/cccl/headers/include/cub/agent/agent_merge_sort.cuh

@@ -66,9 +66,28 @@ struct AgentMergeSortPolicy
   static constexpr cub::BlockStoreAlgorithm STORE_ALGORITHM = _STORE_ALGORITHM;
 };
 
+#if defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON)
 namespace detail
 {
-namespace merge_sort
+// Only define this when needed.
+// Because of overload woes, this depends on C++20 concepts. util_device.h checks that concepts are available when
+// either runtime policies or PTX JSON information are enabled, so if they are, this is always valid. The generic
+// version is always defined, and that's the only one needed for regular CUB operations.
+//
+// TODO: enable this unconditionally once concepts are always available
+CUB_DETAIL_POLICY_WRAPPER_DEFINE(
+  MergeSortAgentPolicy,
+  (GenericAgentPolicy),
+  (BLOCK_THREADS, BlockThreads, int),
+  (ITEMS_PER_THREAD, ItemsPerThread, int),
+  (ITEMS_PER_TILE, ItemsPerTile, int),
+  (LOAD_ALGORITHM, LoadAlgorithm, cub::BlockLoadAlgorithm),
+  (LOAD_MODIFIER, LoadModifier, cub::CacheLoadModifier),
+  (STORE_ALGORITHM, StoreAlgorithm, cub::BlockStoreAlgorithm))
+} // namespace detail
+#endif // defined(CUB_DEFINE_RUNTIME_POLICIES
+
+namespace detail::merge_sort
 {
 
 template <typename Policy,
@@ -724,7 +743,6 @@ struct AgentMerge
   }
 };
 
-} // namespace merge_sort
-} // namespace detail
+} // namespace detail::merge_sort
 
 CUB_NAMESPACE_END

cuda/cccl/headers/include/cub/agent/agent_radix_sort_downsweep.cuh

@@ -146,9 +146,7 @@ CUB_DETAIL_POLICY_WRAPPER_DEFINE(
  * Thread block abstractions
  ******************************************************************************/
 
-namespace detail
-{
-namespace radix_sort
+namespace detail::radix_sort
 {
 
 /**
@@ -783,7 +781,6 @@ struct AgentRadixSortDownsweep
   }
 };
 
-} // namespace radix_sort
-} // namespace detail
+} // namespace detail::radix_sort
 
 CUB_NAMESPACE_END

cuda/cccl/headers/include/cub/agent/agent_radix_sort_histogram.cuh

@@ -85,9 +85,7 @@ struct AgentRadixSortExclusiveSumPolicy
   };
 };
 
-namespace detail
-{
-namespace radix_sort
+namespace detail::radix_sort
 {
 
 template <typename AgentRadixSortHistogramPolicy,
@@ -283,7 +281,6 @@ struct AgentRadixSortHistogram
   }
 };
 
-} // namespace radix_sort
-} // namespace detail
+} // namespace detail::radix_sort
 
 CUB_NAMESPACE_END

cuda/cccl/headers/include/cub/agent/agent_radix_sort_onesweep.cuh

@@ -100,9 +100,7 @@ struct AgentRadixSortOnesweepPolicy : ScalingType
   static constexpr RadixSortStoreAlgorithm STORE_ALGORITHM = _STORE_ALGORITHM;
 };
 
-namespace detail
-{
-namespace radix_sort
+namespace detail::radix_sort
 {
 
 template <typename AgentRadixSortOnesweepPolicy,
@@ -700,7 +698,6 @@ struct AgentRadixSortOnesweep
   }
 };
 
-} // namespace radix_sort
-} // namespace detail
+} // namespace detail::radix_sort
 
 CUB_NAMESPACE_END

cuda/cccl/headers/include/cub/agent/agent_radix_sort_upsweep.cuh

@@ -103,9 +103,7 @@ struct AgentRadixSortUpsweepPolicy : ScalingType
  * Thread block abstractions
  ******************************************************************************/
 
-namespace detail
-{
-namespace radix_sort
+namespace detail::radix_sort
 {
 
 /**
@@ -552,7 +550,6 @@ struct AgentRadixSortUpsweep
   }
 };
 
-} // namespace radix_sort
-} // namespace detail
+} // namespace detail::radix_sort
 
 CUB_NAMESPACE_END

cuda/cccl/headers/include/cub/agent/agent_rle.cuh

@@ -134,9 +134,7 @@ struct AgentRlePolicy
  * Thread block abstractions
  ******************************************************************************/
 
-namespace detail
-{
-namespace rle
+namespace detail::rle
 {
 
 /**
@@ -1121,7 +1119,6 @@ struct AgentRle
   }
 };
 
-} // namespace rle
-} // namespace detail
+} // namespace detail::rle
 
 CUB_NAMESPACE_END

cuda/cccl/headers/include/cub/agent/agent_scan.cuh

@@ -51,6 +51,10 @@
 #include <cub/iterator/cache_modified_input_iterator.cuh>
 #include <cub/util_device.cuh>
 
+#if defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON)
+#  include <cub/agent/agent_unique_by_key.cuh> // for UniqueByKeyAgentPolicy
+#endif
+
 #include <cuda/std/__type_traits/conditional.h>
 #include <cuda/std/__type_traits/is_pointer.h>
 #include <cuda/std/__type_traits/is_same.h>
@@ -123,7 +127,7 @@ namespace detail
 // TODO: enable this unconditionally once concepts are always available
 CUB_DETAIL_POLICY_WRAPPER_DEFINE(
   ScanAgentPolicy,
-  (GenericAgentPolicy),
+  (UniqueByKeyAgentPolicy),
   (BLOCK_THREADS, BlockThreads, int),
   (ITEMS_PER_THREAD, ItemsPerThread, int),
   (LOAD_ALGORITHM, LoadAlgorithm, cub::BlockLoadAlgorithm),

cuda/cccl/headers/include/cub/agent/agent_scan_by_key.cuh

@@ -96,9 +96,7 @@ struct AgentScanByKeyPolicy
  * Thread block abstractions
  ******************************************************************************/
 
-namespace detail
-{
-namespace scan_by_key
+namespace detail::scan_by_key
 {
 
 /**
@@ -471,7 +469,6 @@ struct AgentScanByKey
   }
 };
 
-} // namespace scan_by_key
-} // namespace detail
+} // namespace detail::scan_by_key
 
 CUB_NAMESPACE_END

cuda/cccl/headers/include/cub/agent/agent_segmented_radix_sort.cuh

@@ -45,9 +45,7 @@
 
 CUB_NAMESPACE_BEGIN
 
-namespace detail
-{
-namespace radix_sort
+namespace detail::radix_sort
 {
 
 /**
@@ -286,7 +284,6 @@ struct AgentSegmentedRadixSort
   }
 };
 
-} // namespace radix_sort
-} // namespace detail
+} // namespace detail::radix_sort
 
 CUB_NAMESPACE_END

cuda/cccl/headers/include/cub/agent/agent_select_if.cuh

@@ -126,9 +126,7 @@ struct AgentSelectIfPolicy
  * Thread block abstractions
  ******************************************************************************/
 
-namespace detail
-{
-namespace select
+namespace detail::select
 {
 
 template <typename EqualityOpT>
@@ -1114,7 +1112,6 @@ struct AgentSelectIf
   }
 };
 
-} // namespace select
-} // namespace detail
+} // namespace detail::select
 
 CUB_NAMESPACE_END

cuda/cccl/headers/include/cub/agent/agent_sub_warp_merge_sort.cuh

@@ -84,9 +84,7 @@ CUB_DETAIL_POLICY_WRAPPER_DEFINE(
 } // namespace detail
 #endif // defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON)
 
-namespace detail
-{
-namespace sub_warp_merge_sort
+namespace detail::sub_warp_merge_sort
 {
 
 /**
@@ -343,7 +341,6 @@ private:
   }
 };
 
-} // namespace sub_warp_merge_sort
-} // namespace detail
+} // namespace detail::sub_warp_merge_sort
 
 CUB_NAMESPACE_END

cuda/cccl/headers/include/cub/agent/agent_three_way_partition.cuh

@@ -91,9 +91,7 @@ CUB_DETAIL_POLICY_WRAPPER_DEFINE(
 } // namespace detail
 #endif // defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON)
 
-namespace detail
-{
-namespace three_way_partition
+namespace detail::three_way_partition
 {
 
 template <class OffsetT>
@@ -603,7 +601,6 @@ struct AgentThreeWayPartition
   }
 };
 
-} // namespace three_way_partition
-} // namespace detail
+} // namespace detail::three_way_partition
 
 CUB_NAMESPACE_END

cuda/cccl/headers/include/cub/agent/agent_unique_by_key.cuh

@@ -85,13 +85,31 @@ struct AgentUniqueByKeyPolicy
   };
 };
 
+#if defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON)
+namespace detail
+{
+// Only define this when needed.
+// Because of overload woes, this depends on C++20 concepts. util_device.h checks that concepts are available when
+// either runtime policies or PTX JSON information are enabled, so if they are, this is always valid. The generic
+// version is always defined, and that's the only one needed for regular CUB operations.
+//
+// TODO: enable this unconditionally once concepts are always available
+CUB_DETAIL_POLICY_WRAPPER_DEFINE(
+  UniqueByKeyAgentPolicy,
+  (GenericAgentPolicy),
+  (BLOCK_THREADS, BlockThreads, int),
+  (ITEMS_PER_THREAD, ItemsPerThread, int),
+  (LOAD_ALGORITHM, LoadAlgorithm, cub::BlockLoadAlgorithm),
+  (LOAD_MODIFIER, LoadModifier, cub::CacheLoadModifier),
+  (SCAN_ALGORITHM, ScanAlgorithm, cub::BlockScanAlgorithm))
+} // namespace detail
+#endif // defined(CUB_DEFINE_RUNTIME_POLICIES) || defined(CUB_ENABLE_POLICY_PTX_JSON)
+
 /******************************************************************************
  * Thread block abstractions
  ******************************************************************************/
 
-namespace detail
-{
-namespace unique_by_key
+namespace detail::unique_by_key
 {
 
 /**
@@ -608,7 +626,6 @@ struct AgentUniqueByKey
   }
 };
 
-} // namespace unique_by_key
-} // namespace detail
+} // namespace detail::unique_by_key
 
 CUB_NAMESPACE_END

cuda/cccl/headers/include/cub/block/block_radix_rank.cuh

@@ -50,6 +50,7 @@
 
 #include <cuda/__ptx/instructions/get_sreg.h>
 #include <cuda/std/__algorithm/max.h>
+#include <cuda/std/__bit/integral.h>
 #include <cuda/std/__functional/operations.h>
 #include <cuda/std/__type_traits/conditional.h>
 #include <cuda/std/__type_traits/is_same.h>
@@ -1072,7 +1073,7 @@ struct BlockRadixRankMatchEarlyCounts
         atomicOr(p_match_mask, lane_mask);
         __syncwarp(WARP_MASK);
         int bin_mask    = *p_match_mask;
-        int leader      = (WARP_THREADS - 1) - __clz(bin_mask);
+        int leader      = ::cuda::std::__bit_log2(static_cast<unsigned>(bin_mask));
         int warp_offset = 0;
         int popc        = __popc(bin_mask & ::cuda::ptx::get_sreg_lanemask_le());
         if (lane == leader)
@@ -1102,7 +1103,7 @@ struct BlockRadixRankMatchEarlyCounts
         ::cuda::std::uint32_t bin = Digit(keys[u]);
         int bin_mask =
           detail::warp_in_block_matcher_t<RADIX_BITS, PARTIAL_WARP_THREADS, BLOCK_WARPS - 1>::match_any(bin, warp);
-        int leader      = (WARP_THREADS - 1) - __clz(bin_mask);
+        int leader      = ::cuda::std::__bit_log2(static_cast<unsigned>(bin_mask));
         int warp_offset = 0;
         int popc        = __popc(bin_mask & ::cuda::ptx::get_sreg_lanemask_le());
         if (lane == leader)

cuda/cccl/headers/include/cub/block/block_radix_sort.cuh

@@ -174,10 +174,12 @@ CUB_NAMESPACE_BEGIN
 //!
 //!    .. code-block:: python
 //!
-//!        import cuda.cccl.cooperative.experimental as cudax
+//!        from cuda import coop
+//!        from pynvjitlink import patch
+//!        patch.patch_numba_linker(lto=True)
 //!
 //!        # Specialize radix sort for a 1D block of 128 threads owning 4 integer items each
-//!        block_radix_sort = cudax.block.radix_sort_keys(numba.int32, 128, 4)
+//!        block_radix_sort = coop.block.radix_sort_keys(numba.int32, 128, 4)
 //!        temp_storage_bytes = block_radix_sort.temp_storage_bytes
 //!
 //!        @cuda.jit(link=block_radix_sort.files)

cuda/cccl/headers/include/cub/detail/device_memory_resource.cuh

@@ -14,6 +14,7 @@
 #endif // no system header
 
 #include <cuda/__stream/stream_ref.h>
+#include <cuda/std/__cuda/api_wrapper.h>
 #include <cuda/std/cstdint>
 
 CUB_NAMESPACE_BEGIN