From f68a62f8a78a41d74effb298e564e765dd9b725e Mon Sep 17 00:00:00 2001 From: cliffburdick Date: Thu, 11 Sep 2025 19:33:50 -0700 Subject: [PATCH 1/2] Use cub's new fixed size segmented reduce --- cmake/versions.json | 4 ++-- include/matx/transforms/cub.h | 19 ++++++++++++++----- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/cmake/versions.json b/cmake/versions.json index 4c1c5eb8b..0cc4eb57e 100644 --- a/cmake/versions.json +++ b/cmake/versions.json @@ -1,10 +1,10 @@ { "packages": { "CCCL": { - "version": "3.0.0", + "version": "3.2.0", "git_shallow": false, "git_url": "https://github.com/NVIDIA/cccl.git", - "git_tag": "e944297" + "git_tag": "4071a73" }, "nvbench" : { "version" : "0.0", diff --git a/include/matx/transforms/cub.h b/include/matx/transforms/cub.h index ace85c0de..f606c778e 100644 --- a/include/matx/transforms/cub.h +++ b/include/matx/transforms/cub.h @@ -710,11 +710,20 @@ inline void ExecSort(OutputTensor &a_out, // type of reduction where there's not a single output, since any type of reduction can be generalized // to a segmented type if constexpr (OutputTensor::Rank() > 0) { - auto ft = [&](auto &&in, auto &&out, auto &&begin, auto &&end) { - return cub::DeviceSegmentedReduce::Sum(d_temp, temp_storage_bytes, in, out, static_cast(TotalSize(out_base)), begin, end, stream); - }; - [[maybe_unused]] auto rv = ReduceInput(ft, out_base, in_base); - MATX_ASSERT_STR_EXP(rv, cudaSuccess, matxCudaError, "Error in cub::DeviceSegmentedReduce::Sum"); + [[maybe_unused]] cudaError_t err; + if (is_tensor_view_v && a.IsContiguous() && a_out.IsContiguous()) { + const int seg_size = static_cast(TotalSize(a) / TotalSize(out_base)); + err = cub::DeviceSegmentedReduce::Sum(d_temp, temp_storage_bytes, in_base.Data(), out_base.Data(), static_cast(TotalSize(out_base)), seg_size, stream); + } + else { + auto ft = [&](auto &&in, auto &&out, auto &&begin, auto &&end) { + return cub::DeviceSegmentedReduce::Sum(d_temp, temp_storage_bytes, in, out, static_cast(TotalSize(out_base)), begin, end, stream); + }; + + err = ReduceInput(ft, out_base, in_base); + } + + MATX_ASSERT_STR_EXP(err, cudaSuccess, matxCudaError, "Error in cub::DeviceSegmentedReduce::Sum"); } else { auto ft = [&](auto &&in, auto &&out, [[maybe_unused]] auto &&unused1, [[maybe_unused]] auto &&unused2) { From a06a7c3a5afd535143e2b30e870da30d48deb744 Mon Sep 17 00:00:00 2001 From: cliffburdick Date: Mon, 29 Sep 2025 15:19:25 -0700 Subject: [PATCH 2/2] Added fixed-size segmented reduce --- cmake/versions.json | 4 ++-- include/matx/transforms/cub.h | 27 +++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/cmake/versions.json b/cmake/versions.json index 0cc4eb57e..4c1c5eb8b 100644 --- a/cmake/versions.json +++ b/cmake/versions.json @@ -1,10 +1,10 @@ { "packages": { "CCCL": { - "version": "3.2.0", + "version": "3.0.0", "git_shallow": false, "git_url": "https://github.com/NVIDIA/cccl.git", - "git_tag": "4071a73" + "git_tag": "e944297" }, "nvbench" : { "version" : "0.0", diff --git a/include/matx/transforms/cub.h b/include/matx/transforms/cub.h index f606c778e..dca74a043 100644 --- a/include/matx/transforms/cub.h +++ b/include/matx/transforms/cub.h @@ -659,12 +659,30 @@ inline void ExecSort(OutputTensor &a_out, // type of reduction where there's not a single output, since any type of reduction can be generalized // to a segmented type if constexpr (OutputTensor::Rank() > 0) { +#if CUB_MAJOR_VERSION >= 3 && CUB_MINOR_VERSION >= 2 + [[maybe_unused]] cudaError_t err; + if (is_tensor_view_v && a.IsContiguous() && a_out.IsContiguous()) { + const int seg_size = static_cast(TotalSize(a) / TotalSize(out_base)); + err = cub::DeviceSegmentedReduce::Reduce(d_temp, temp_storage_bytes, in_base.Data(), out_base.Data(), static_cast(TotalSize(out_base)), seg_size, cparams_.reduce_op, + cparams_.init, stream); + } + else { + auto ft = [&](auto &&in, auto &&out, auto &&begin, auto &&end) { + return cub::DeviceSegmentedReduce::Reduce(d_temp, temp_storage_bytes, in, out, static_cast(TotalSize(out_base)), begin, end, cparams_.reduce_op, + cparams_.init, stream); + }; + err = ReduceInput(ft, out_base, in_base); + } + + MATX_ASSERT_STR_EXP(err, cudaSuccess, matxCudaError, "Error in cub::DeviceSegmentedReduce::Reduce"); +#else auto ft = [&](auto &&in, auto &&out, auto &&begin, auto &&end) { return cub::DeviceSegmentedReduce::Reduce(d_temp, temp_storage_bytes, in, out, static_cast(TotalSize(out_base)), begin, end, cparams_.reduce_op, cparams_.init, stream); }; [[maybe_unused]] auto rv = ReduceInput(ft, out_base, in_base); MATX_ASSERT_STR_EXP(rv, cudaSuccess, matxCudaError, "Error in cub::DeviceSegmentedReduce::Reduce"); +#endif } else { auto ft = [&](auto &&in, auto &&out, [[maybe_unused]] auto &&unused1, [[maybe_unused]] auto &&unused2) { @@ -710,6 +728,8 @@ inline void ExecSort(OutputTensor &a_out, // type of reduction where there's not a single output, since any type of reduction can be generalized // to a segmented type if constexpr (OutputTensor::Rank() > 0) { + // Check if fixed-size reductions are supported +#if CUB_MAJOR_VERSION >= 3 && CUB_MINOR_VERSION >= 2 [[maybe_unused]] cudaError_t err; if (is_tensor_view_v && a.IsContiguous() && a_out.IsContiguous()) { const int seg_size = static_cast(TotalSize(a) / TotalSize(out_base)); @@ -724,6 +744,13 @@ inline void ExecSort(OutputTensor &a_out, } MATX_ASSERT_STR_EXP(err, cudaSuccess, matxCudaError, "Error in cub::DeviceSegmentedReduce::Sum"); +#else + auto ft = [&](auto &&in, auto &&out, auto &&begin, auto &&end) { + return cub::DeviceSegmentedReduce::Sum(d_temp, temp_storage_bytes, in, out, static_cast(TotalSize(out_base)), begin, end, stream); + }; + [[maybe_unused]] auto rv = ReduceInput(ft, out_base, in_base); + MATX_ASSERT_STR_EXP(rv, cudaSuccess, matxCudaError, "Error in cub::DeviceSegmentedReduce::Sum"); +#endif } else { auto ft = [&](auto &&in, auto &&out, [[maybe_unused]] auto &&unused1, [[maybe_unused]] auto &&unused2) {