Integrate decoupled lookahead warpspeed scan

c/parallel/src/scan.cu

@@ -104,12 +104,27 @@ std::string get_output_iterator_name()
   return iterator_t;
 }
 
-std::string
-get_init_kernel_name(cccl_iterator_t input_it, cccl_iterator_t /*output_it*/, cccl_op_t op, cccl_type_info init)
+std::string get_init_kernel_name(cccl_iterator_t input_it, cccl_iterator_t output_it, cccl_op_t op, cccl_type_info init)
 {
+  std::string chained_policy_t;
+  check(cccl_type_name_from_nvrtc<device_scan_policy>(&chained_policy_t));
+
   const cccl_type_info accum_t  = scan::get_accumulator_type(op, input_it, init);
   const std::string accum_cpp_t = cccl_type_enum_to_name(accum_t.type);
-  return std::format("cub::detail::scan::DeviceScanInitKernel<cub::ScanTileState<{0}>>", accum_cpp_t);
+  const std::string input_iterator_t =
+    (input_it.type == cccl_iterator_kind_t::CCCL_POINTER //
+       ? cccl_type_enum_to_name(input_it.value_type.type, true) //
+       : scan::get_input_iterator_name());
+  const std::string output_iterator_t =
+    output_it.type == cccl_iterator_kind_t::CCCL_POINTER //
+      ? cccl_type_enum_to_name(output_it.value_type.type, true) //
+      : scan::get_output_iterator_name();
+  return std::format(
+    "cub::detail::scan::DeviceScanInitKernel<{0}, {1}, {2}, cub::ScanTileState<{3}>, {3}>",
+    chained_policy_t,
+    input_iterator_t,
+    output_iterator_t,
+    accum_cpp_t);
 }
 
 std::string get_scan_kernel_name(
@@ -208,6 +223,33 @@ struct scan_kernel_source
   {
     return {build.description_bytes_per_tile, build.payload_bytes_per_tile};
   }
+
+  std::size_t look_ahead_tile_state_size() const
+  {
+    return look_ahead_tile_state_alignment();
+  }
+
+  std::size_t look_ahead_tile_state_alignment() const
+  {
+    constexpr int state_size = alignof(cub::detail::warpspeed::scan_state);
+    return ::cuda::next_power_of_two(
+      ::cuda::round_up(state_size, build.accumulator_type.alignment) + build.accumulator_type.size);
+  }
+
+  static auto make_tile_state_kernel_arg(scan_tile_state ts)
+  {
+    cub::detail::scan::tile_state_kernel_arg_t<scan_tile_state, char> arg;
+    ::cuda::std::__construct_at(&arg.lookback, ::cuda::std::move(ts));
+    return arg;
+  }
+
+  static auto look_ahead_make_tile_state_kernel_arg(void* ts)
+  {
+    // we can ignore passing a wrong AccumT, since we only store a pointer, and the kernel will have the right type
+    cub::detail::scan::tile_state_kernel_arg_t<scan_tile_state, char> arg;
+    ::cuda::std::__construct_at(&arg.lookahead, static_cast<cub::detail::warpspeed::tile_state_t<char>*>(ts));
+    return arg;
+  }
 };
 } // namespace scan
 

cub/cub/detail/launcher/cuda_runtime.cuh

@@ -97,6 +97,29 @@ struct TripleChevronFactory
 
     return cudaDeviceGetAttribute(&max_shared_memory, cudaDevAttrMaxSharedMemoryPerBlock, device);
   }
+
+  template <typename Kernel>
+  _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t
+  max_dynamic_smem_size_for(int& max_dynamic_smem_size, [[maybe_unused]] Kernel kernel_ptr)
+  {
+    NV_IF_ELSE_TARGET(NV_IS_HOST, //
+                      ({ return MaxPotentialDynamicSmemBytes(max_dynamic_smem_size, kernel_ptr); }),
+                      ({
+                        ::cudaFuncAttributes func_attrs{};
+                        if (const auto error = CubDebug(::cudaFuncGetAttributes(&func_attrs, kernel_ptr)))
+                        {
+                          return error;
+                        }
+                        max_dynamic_smem_size = func_attrs.maxDynamicSharedSizeBytes;
+                        return cudaSuccess;
+                      }))
+  }
+
+  template <typename Kernel>
+  _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t set_max_dynamic_smem_size_for(Kernel kernel_ptr, int smem_size)
+  {
+    return CubDebug(::cudaFuncSetAttribute(kernel_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+  }
 };
 } // namespace detail
 

cub/cub/detail/warpspeed/allocators/smem_allocator.cuh

@@ -0,0 +1,83 @@
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
+#pragma once
+
+#include <cub/config.cuh>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#include <cub/detail/warpspeed/optimize_smem_ptr.cuh>
+
+#include <cuda/std/__type_traits/is_constant_evaluated.h>
+#include <cuda/std/cstdint>
+
+#include <nv/target>
+
+CUB_NAMESPACE_BEGIN
+
+namespace detail::warpspeed
+{
+struct SmemAllocator
+{
+  ::cuda::std::uint32_t mPtrSmem32 = 0;
+  int mAllocatedSize               = 0;
+
+  _CCCL_API constexpr SmemAllocator() noexcept
+  {
+    // we only need the real pointer at runtime in device code
+    if (!::cuda::std::is_constant_evaluated())
+    {
+      NV_IF_TARGET(NV_IS_DEVICE, mPtrSmem32 = dynamic_smem_base();)
+    }
+  }
+
+  [[nodiscard]] _CCCL_DEVICE_API static ::cuda::std::uint32_t dynamic_smem_base() noexcept
+  {
+    extern __shared__ char warpSpeedDynamicSmemBase[];
+    return __cvta_generic_to_shared(warpSpeedDynamicSmemBase);
+  }
+
+  // SmemAllocator is a non-copyable, non-movable type. It must be passed by
+  // (mutable) reference to be useful.
+  SmemAllocator(const SmemAllocator&)            = delete; // Delete copy constructor
+  SmemAllocator(SmemAllocator&&)                 = delete; // Delete move constructor
+  SmemAllocator& operator=(const SmemAllocator&) = delete; // Delete copy assignment
+  SmemAllocator& operator=(SmemAllocator&&)      = delete; // Delete move assignment
+
+  [[nodiscard]] _CCCL_API constexpr void* alloc(::cuda::std::uint32_t size, ::cuda::std::uint32_t align = 0)
+  {
+    // Align mPtrSmem32 to requested alignment (round-up)
+    ::cuda::std::uint32_t ptrAllocation32 = (mPtrSmem32 + (align - 1)) & ~(align - 1);
+
+    // Move base pointer and update allocated size
+    mAllocatedSize += size + ptrAllocation32 - mPtrSmem32;
+    mPtrSmem32 = ptrAllocation32 + size;
+
+    // we only need the pointer at runtime in device code
+    if (!::cuda::std::is_constant_evaluated())
+    {
+      NV_IF_TARGET(
+        NV_IS_DEVICE,
+        (
+          // Convert allocated smem address to generic pointer
+          void* mPtrAllocation = __cvta_shared_to_generic(ptrAllocation32);
+          // Ensure alignment calculation does not move down into rest of kernel code.
+          return optimizeSmemPtr(mPtrAllocation);))
+    }
+    return nullptr;
+  }
+
+  [[nodiscard]] _CCCL_API constexpr ::cuda::std::uint32_t sizeBytes() const
+  {
+    return mAllocatedSize;
+  }
+};
+} // namespace detail::warpspeed
+
+CUB_NAMESPACE_END

cub/cub/detail/warpspeed/constant_assert.cuh

@@ -0,0 +1,52 @@
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
+#pragma once
+
+#include <cub/config.cuh>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#include <cuda/std/cassert>
+
+/*
+ * _WS_CONSTANT_ASSERT: an assertion that is intended to be verified at compile time.
+ *
+ * A _WS_CONSTANT_ASSERT asserts something that the compiler (optimizer) can verify
+ * at compile time. Therefore, it does not result in an actual call to assert in
+ * the compiled binary. This allows checking various properties that cannot be
+ * verified using static_assert.
+ *
+ * To ensure that all _WS_CONSTANT_ASSERTs are in fact eliminated, compile with
+ * -D_WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME. With this macro defined, any
+ * _WS_CONSTANT_ASSERT failure will output illegal PTX containing the error message.
+ * As a result, compilation will fail.
+ *
+ * Compiling with -D_WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME has the additional
+ * advantage that violating any of the assertions can be detected at compile
+ * time and before even running the code.
+ *
+ */
+
+#if defined(_WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME) && defined(__CUDA_ARCH__)
+// When _WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME is defined and compiling for device, output illegal PTX.
+// This causes the compilation to fail.
+#  define _WS_CONSTANT_ASSERT(expr, msg)                                                       \
+    do                                                                                         \
+    {                                                                                          \
+      if (!(expr))                                                                             \
+      {                                                                                        \
+        asm volatile(".pragma \"\n" __FILE__ "(" _CCCL_TO_STRING(                              \
+          __LINE__) "): %0"                                                                    \
+                    ": error: constant assertion failed with '" msg "'\n\";" ::"C"(__func__)); \
+      }                                                                                        \
+    } while (0)
+#else
+// Host or !_WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME
+#  define _WS_CONSTANT_ASSERT(expr, msg) (assert((expr) && (msg)))
+#endif