Milestone: test file using segmented scan compiles

Author: oleksandr-pavlyk

cub/cub/agent/agent_segmented_scan.cuh

@@ -230,7 +230,7 @@ struct AgentSegmentedScan
     AccumT exclusive_prefix{};
     OffsetT n_chunks = cuda::ceil_div(inp_idx_end - inp_idx_begin, TILE_ITEMS);
 
-    InputT thread_values[ITEMS_PER_THREAD] = {};
+    AccumT thread_values[ITEMS_PER_THREAD] = {};
 
     for (OffsetT chunk_id = 0; chunk_id < n_chunks; ++chunk_id)
     {
@@ -245,7 +245,7 @@ struct AgentSegmentedScan
       // execute BlockScan
       AccumT block_aggregate;
       BlockScanT block_scan_algo(temp_storage.scan);
-      if (IS_INCLUSIVE)
+      if constexpr (IS_INCLUSIVE)
       {
         block_scan_algo.InclusiveScan(thread_values, thread_values, scan_op, block_aggregate);
       }
@@ -264,9 +264,9 @@ struct AgentSegmentedScan
       {
         constexpr auto loop_size = static_cast<int>(ITEMS_PER_THREAD);
         cuda::static_for<loop_size>([&](int i) {
-          thread_values[i] += exclusive_prefix;
+          thread_values[i] = thread_values[i] + exclusive_prefix;
         });
-        exclusive_prefix += block_aggregate;
+        exclusive_prefix = exclusive_prefix + block_aggregate;
       }
 
       // write out scan values using BlockStore

cub/cub/device/device_segmented_scan.cuh

@@ -0,0 +1,72 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+
+#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/device/dispatch/dispatch_segmented_scan.cuh>
+
+#include <cuda/std/cstdint>
+
+CUB_NAMESPACE_BEGIN
+
+struct DeviceSegmentedScan
+{
+  template <typename InputIteratorT,
+            typename OutputIteratorT,
+            typename BeginOffsetIteratorInputT,
+            typename EndOffsetIteratorInputT,
+            typename BeginOffsetIteratorOutputT,
+            typename ScanOpT>
+  CUB_RUNTIME_FUNCTION static cudaError_t InclusiveSegmentedScan(
+    void* d_temp_storage,
+    size_t& temp_storage_bytes,
+    InputIteratorT d_in,
+    OutputIteratorT d_out,
+    ::cuda::std::int64_t num_segments,
+    BeginOffsetIteratorInputT d_in_begin_offsets,
+    EndOffsetIteratorInputT d_in_end_offsets,
+    BeginOffsetIteratorOutputT d_out_begin_offsets,
+    ScanOpT scan_op,
+    cudaStream_t stream = 0)
+  {
+    // defined in cub/config.cuh
+    _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSegmentedScan::InclusiveSegmentedScan");
+
+    using OffsetT =
+      detail::common_iterator_value_t<BeginOffsetIteratorInputT, EndOffsetIteratorInputT, BeginOffsetIteratorOutputT>;
+    using integral_offset_check = ::cuda::std::is_integral<OffsetT>;
+
+    static_assert(integral_offset_check::value, "Offset iterator value type should be integral.");
+
+    return cub::DispatchSegmentedScan<
+      InputIteratorT,
+      OutputIteratorT,
+      BeginOffsetIteratorInputT,
+      EndOffsetIteratorInputT,
+      BeginOffsetIteratorOutputT,
+      ScanOpT,
+      NullType>::Dispatch(d_temp_storage,
+                          temp_storage_bytes,
+                          d_in,
+                          d_out,
+                          num_segments,
+                          d_in_begin_offsets,
+                          d_in_end_offsets,
+                          d_out_begin_offsets,
+                          scan_op,
+                          NullType(),
+                          stream);
+  }
+};
+
+CUB_NAMESPACE_END