Skip to content

Enable fixed-size reductions for expressions #1245

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
cliffburdick merged 1 commit into from
Sep 4, 2026
Merged

Enable fixed-size reductions for expressions #1245

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions docs_input/api/manipulation/selecting/reduce.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,10 @@ For CUDA tensor views whose final two dimensions are a transpose of contiguous
storage, a reduction of the innermost dimension uses coalesced tiled loads.
Other permutations and strided layouts retain the general CUB iterator path.

Fixed-size reductions over the innermost dimensions of a random-access
expression are evaluated without an intermediate tensor when the CUDA backend
supports fixed-size segmented reductions.

.. versionadded:: 0.6.0

.. doxygenfunction:: reduce(const InType &in, ReduceOp op, bool init = true)
Expand Down
4 changes: 4 additions & 0 deletions docs_input/api/math/sumprod/sum.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,10 @@ transposing the final two dimensions uses a tiled fast path. This keeps reads
coalesced without materializing the transpose. Other permutations and strided
views continue to use the general reduction path.

Fixed-size reductions over the innermost dimensions of a random-access
expression are evaluated without an intermediate tensor when the CUDA backend
supports fixed-size segmented reductions.

.. versionadded:: 0.6.0

.. doxygenfunction:: sum(const InType &in, const int (&dims)[D])
Expand Down
39 changes: 38 additions & 1 deletion include/matx/core/reduce_utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,14 @@

namespace matx {

template <typename InputOp>
__MATX_HOST__ __MATX_INLINE__ bool IsFixedSizeReductionCompatible(const InputOp &in) {
if constexpr (requires { in.IsIdentityPermutation(); }) {
return in.IsIdentityPermutation();
}
return true;
}


template <bool ConvertType, typename Func, typename OutputOp, typename InputOp, typename BeginIter, typename EndIter>
__MATX_HOST__ __MATX_INLINE__ auto ReduceOutput(Func &&func, OutputOp &&out, InputOp &&in, BeginIter &&bi, EndIter &&ei) {
Expand All @@ -66,6 +74,26 @@ namespace matx {
return func(in, iter, bi, ei);
}

template <bool ConvertType, typename Func, typename OutputOp, typename InputOp>
__MATX_HOST__ __MATX_INLINE__ auto ReduceOutputFixed(Func &&func, OutputOp &&out, InputOp &&in) {
if constexpr (is_tensor_view_v<OutputOp>) {
if (out.IsContiguous()) {
if constexpr(ConvertType) {
return func(in,
reinterpret_cast<detail::convert_matx_type_t<typename remove_cvref_t<OutputOp>::value_type> *>(out.Data()));
}
else {
return func(in,
reinterpret_cast<typename remove_cvref_t<OutputOp>::value_type *>(out.Data()));
}
}
}

detail::base_type_t<OutputOp> out_base = out;
auto iter = RandomOperatorOutputIterator<decltype(out_base), ConvertType>{out_base};
return func(in, iter);
}

template <typename Func, typename OutputOp, typename InputOp, bool ConvertType = true>
__MATX_HOST__ __MATX_INLINE__ auto ReduceInput(Func &&func, OutputOp &&out, InputOp &&in) {
typename detail::base_type_t<InputOp> in_base = in;
Expand Down Expand Up @@ -96,9 +124,18 @@ namespace matx {
return ReduceOutput<ConvertType>(std::forward<Func>(func), std::forward<OutputOp>(out), iter, BeginOffset{iter}, EndOffset{iter});
}

template <typename Func, typename OutputOp, typename InputOp, bool ConvertType = true>
__MATX_HOST__ __MATX_INLINE__ auto ReduceInputFixed(Func &&func, OutputOp &&out, InputOp &&in) {
typename detail::base_type_t<InputOp> in_base = in;
auto collapsed = matx::lcollapse<remove_cvref_t<decltype(out)>::Rank()>(rcollapse<remove_cvref_t<decltype(in)>::Rank() -
remove_cvref_t<decltype(out)>::Rank()>(in_base));
const auto &iter = matx::RandomOperatorIterator<decltype(collapsed), ConvertType>{collapsed};
return ReduceOutputFixed<ConvertType>(std::forward<Func>(func), std::forward<OutputOp>(out), iter);
}

template <typename Func, typename OutputOp, typename InputOp>
__MATX_HOST__ __MATX_INLINE__ auto ReduceInputNoConvert(Func &&func, OutputOp &&out, InputOp &&in) {
return ReduceInput<Func, OutputOp, InputOp, false>(std::forward<Func>(func), std::forward<OutputOp>(out), std::forward<InputOp>(in));
}
}
#endif
#endif
91 changes: 74 additions & 17 deletions include/matx/transforms/cub.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -849,6 +849,9 @@ 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 (TotalSize(out_base) == 0) {
return;
}
#if CUB_MAJOR_VERSION > 3 || (CUB_MAJOR_VERSION == 3 && CUB_MINOR_VERSION >= 2)
[[maybe_unused]] cudaError_t err;
if constexpr(is_tensor_view_v<InputOperator>) {
Expand All @@ -861,11 +864,23 @@ inline void ExecSort(OutputTensor &a_out,
}
}

auto ft = [&](auto &&in, auto &&out, auto &&begin, auto &&end) {
return cub::DeviceSegmentedReduce::Reduce(d_temp, temp_storage_bytes, in, out, static_cast<int>(TotalSize(out_base)), begin, end, Params().reduce_op,
Params().init, stream);
};
err = ReduceInput(ft, out_base, in_base);
if (IsFixedSizeReductionCompatible(a)) {
auto fixed_size_ft = [&](auto &&in, auto &&out) {
const int seg_size = static_cast<int>(TotalSize(a) / TotalSize(out_base));
return cub::DeviceSegmentedReduce::Reduce(d_temp, temp_storage_bytes, in, out,
static_cast<cuda::std::int64_t>(TotalSize(out_base)), seg_size, Params().reduce_op,
Params().init, stream);
};
err = ReduceInputFixed(fixed_size_ft, out_base, in_base);
}
else {
auto ft = [&](auto &&in, auto &&out, auto &&begin, auto &&end) {
return cub::DeviceSegmentedReduce::Reduce(d_temp, temp_storage_bytes, in, out,
static_cast<int>(TotalSize(out_base)), begin, end, Params().reduce_op,
Params().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) {
Expand Down Expand Up @@ -920,6 +935,9 @@ 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 (TotalSize(out_base) == 0) {
return;
}
// Check if fixed-size reductions are supported
#if CUB_MAJOR_VERSION > 3 || (CUB_MAJOR_VERSION == 3 && CUB_MINOR_VERSION >= 2)
[[maybe_unused]] cudaError_t err;
Expand All @@ -932,10 +950,21 @@ inline void ExecSort(OutputTensor &a_out,
}
}

auto ft = [&](auto &&in, auto &&out, auto &&begin, auto &&end) {
return cub::DeviceSegmentedReduce::Sum(d_temp, temp_storage_bytes, in, out, static_cast<int>(TotalSize(out_base)), begin, end, stream);
};
err = ReduceInput(ft, out_base, in_base);
if (IsFixedSizeReductionCompatible(a)) {
auto fixed_size_ft = [&](auto &&in, auto &&out) {
const int seg_size = static_cast<int>(TotalSize(a) / TotalSize(out_base));
return cub::DeviceSegmentedReduce::Sum(d_temp, temp_storage_bytes, in, out,
static_cast<cuda::std::int64_t>(TotalSize(out_base)), seg_size, stream);
};
err = ReduceInputFixed(fixed_size_ft, out_base, in_base);
}
else {
auto ft = [&](auto &&in, auto &&out, auto &&begin, auto &&end) {
return cub::DeviceSegmentedReduce::Sum(d_temp, temp_storage_bytes, in, out,
static_cast<int>(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, auto &&begin, auto &&end) {
Expand Down Expand Up @@ -985,6 +1014,9 @@ 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 (TotalSize(out_base) == 0) {
return;
}
#if CUB_MAJOR_VERSION > 3 || (CUB_MAJOR_VERSION == 3 && CUB_MINOR_VERSION >= 2)
[[maybe_unused]] cudaError_t err;
if constexpr (is_tensor_view_v<InputOperator>) {
Expand All @@ -996,10 +1028,21 @@ inline void ExecSort(OutputTensor &a_out,
}
}

auto ft = [&](auto &&in, auto &&out, auto &&begin, auto &&end) {
return cub::DeviceSegmentedReduce::Min(d_temp, temp_storage_bytes, in, out, static_cast<int>(TotalSize(out_base)), begin, end, stream);
};
err = ReduceInput(ft, out_base, in_base);
if (IsFixedSizeReductionCompatible(a)) {
auto fixed_size_ft = [&](auto &&in, auto &&out) {
const int seg_size = static_cast<int>(TotalSize(a) / TotalSize(out_base));
return cub::DeviceSegmentedReduce::Min(d_temp, temp_storage_bytes, in, out,
static_cast<cuda::std::int64_t>(TotalSize(out_base)), seg_size, stream);
};
err = ReduceInputFixed(fixed_size_ft, out_base, in_base);
}
else {
auto ft = [&](auto &&in, auto &&out, auto &&begin, auto &&end) {
return cub::DeviceSegmentedReduce::Min(d_temp, temp_storage_bytes, in, out,
static_cast<int>(TotalSize(out_base)), begin, end, stream);
};
err = ReduceInput(ft, out_base, in_base);
}
MATX_ASSERT_STR_EXP(err, cudaSuccess, matxCudaError, "Error in cub::DeviceSegmentedReduce::Min");
#else
auto ft = [&](auto &&in, auto &&out, auto &&begin, auto &&end) {
Expand Down Expand Up @@ -1049,6 +1092,9 @@ 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 (TotalSize(out_base) == 0) {
return;
}
#if CUB_MAJOR_VERSION > 3 || (CUB_MAJOR_VERSION == 3 && CUB_MINOR_VERSION >= 2)
[[maybe_unused]] cudaError_t err;
if constexpr (is_tensor_view_v<InputOperator>) {
Expand All @@ -1060,10 +1106,21 @@ inline void ExecSort(OutputTensor &a_out,
}
}

auto ft = [&](auto &&in, auto &&out, auto &&begin, auto &&end) {
return cub::DeviceSegmentedReduce::Max(d_temp, temp_storage_bytes, in, out, static_cast<int>(TotalSize(out_base)), begin, end, stream);
};
err = ReduceInput(ft, out_base, in_base);
if (IsFixedSizeReductionCompatible(a)) {
auto fixed_size_ft = [&](auto &&in, auto &&out) {
const int seg_size = static_cast<int>(TotalSize(a) / TotalSize(out_base));
return cub::DeviceSegmentedReduce::Max(d_temp, temp_storage_bytes, in, out,
static_cast<cuda::std::int64_t>(TotalSize(out_base)), seg_size, stream);
};
err = ReduceInputFixed(fixed_size_ft, out_base, in_base);
}
else {
auto ft = [&](auto &&in, auto &&out, auto &&begin, auto &&end) {
return cub::DeviceSegmentedReduce::Max(d_temp, temp_storage_bytes, in, out,
static_cast<int>(TotalSize(out_base)), begin, end, stream);
};
err = ReduceInput(ft, out_base, in_base);
}
MATX_ASSERT_STR_EXP(err, cudaSuccess, matxCudaError, "Error in cub::DeviceSegmentedReduce::Max");
#else
auto ft = [&](auto &&in, auto &&out, auto &&begin, auto &&end) {
Expand Down
93 changes: 93 additions & 0 deletions test/00_operators/ReductionTests.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1198,6 +1198,99 @@ TYPED_TEST(ReductionTestsFloatNonComplexNonHalfAllExecs, SegmentedSum)
MATX_EXIT_HANDLER();
}

TYPED_TEST(ReductionTestsComplexNonHalfTypes, FixedSizeExpressionSum)
{
MATX_ENTER_HANDLER();
using TestType = cuda::std::tuple_element_t<0, TypeParam>;
using ExecType = cuda::std::tuple_element_t<1, TypeParam>;
using ValueType = typename inner_op_type_t<TestType>::type;

ExecType exec{};
auto input = make_tensor<TestType>({4, 3, 5});
auto inner_output = make_tensor<ValueType>({4, 3});
auto outer_output = make_tensor<ValueType>({3, 5});

for (index_t batch = 0; batch < input.Size(0); ++batch) {
for (index_t row = 0; row < input.Size(1); ++row) {
for (index_t col = 0; col < input.Size(2); ++col) {
input(batch, row, col) = TestType(static_cast<ValueType>(batch + 1),
static_cast<ValueType>(row + col + 1));
}
}
}
(inner_output = sum(abs2(input), {2})).run(exec);
(outer_output = sum(abs2(input), {0})).run(exec);
exec.sync();

for (index_t batch = 0; batch < inner_output.Size(0); ++batch) {
for (index_t row = 0; row < inner_output.Size(1); ++row) {
ValueType expected{};
for (index_t col = 0; col < input.Size(2); ++col) {
expected += static_cast<ValueType>((batch + 1) * (batch + 1) +
(row + col + 1) * (row + col + 1));
}
ASSERT_TRUE(MatXUtils::MatXTypeCompare(inner_output(batch, row), expected));
}
}

for (index_t row = 0; row < outer_output.Size(0); ++row) {
for (index_t col = 0; col < outer_output.Size(1); ++col) {
ValueType expected{};
for (index_t batch = 0; batch < input.Size(0); ++batch) {
expected += static_cast<ValueType>((batch + 1) * (batch + 1) +
(row + col + 1) * (row + col + 1));
}
ASSERT_TRUE(MatXUtils::MatXTypeCompare(outer_output(row, col), expected));
}
}

MATX_EXIT_HANDLER();
}

TYPED_TEST(ReductionTestsFloatNonComplexNonHalf, FixedSizeExpressionReductions)
{
MATX_ENTER_HANDLER();
using TestType = cuda::std::tuple_element_t<0, TypeParam>;
using ExecType = cuda::std::tuple_element_t<1, TypeParam>;

ExecType exec{};
auto input = make_tensor<TestType>({4, 3, 5});
auto minimum = make_tensor<TestType>({4, 3});
auto maximum = make_tensor<TestType>({4, 3});
auto product = make_tensor<TestType>({4, 3});

for (index_t batch = 0; batch < input.Size(0); ++batch) {
for (index_t row = 0; row < input.Size(1); ++row) {
for (index_t col = 0; col < input.Size(2); ++col) {
input(batch, row, col) = static_cast<TestType>(batch + row + col + 1);
}
}
}

auto expression = input + static_cast<TestType>(1);
(minimum = min(expression, {2})).run(exec);
(maximum = max(expression, {2})).run(exec);
(product = prod(expression, {2})).run(exec);
exec.sync();

for (index_t batch = 0; batch < input.Size(0); ++batch) {
for (index_t row = 0; row < input.Size(1); ++row) {
const auto first = static_cast<TestType>(batch + row + 2);
const auto last = static_cast<TestType>(batch + row + 6);
TestType expected_product = static_cast<TestType>(1);
for (index_t col = 0; col < input.Size(2); ++col) {
expected_product *= static_cast<TestType>(batch + row + col + 2);
}

ASSERT_TRUE(MatXUtils::MatXTypeCompare(minimum(batch, row), first));
ASSERT_TRUE(MatXUtils::MatXTypeCompare(maximum(batch, row), last));
ASSERT_TRUE(MatXUtils::MatXTypeCompare(product(batch, row), expected_product));
}
}

MATX_EXIT_HANDLER();
}

TYPED_TEST(ReductionTestsFloatNonComplexNonHalfAllExecs, SegmentedMin)
{
MATX_ENTER_HANDLER();
Expand Down