diff --git a/docs_input/api/manipulation/selecting/reduce.rst b/docs_input/api/manipulation/selecting/reduce.rst index 9e7a25e70..95c3e62ca 100644 --- a/docs_input/api/manipulation/selecting/reduce.rst +++ b/docs_input/api/manipulation/selecting/reduce.rst @@ -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) diff --git a/docs_input/api/math/sumprod/sum.rst b/docs_input/api/math/sumprod/sum.rst index 39883f29a..d77117024 100644 --- a/docs_input/api/math/sumprod/sum.rst +++ b/docs_input/api/math/sumprod/sum.rst @@ -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]) diff --git a/include/matx/core/reduce_utils.h b/include/matx/core/reduce_utils.h index 7d96c5b31..9a68c9acc 100644 --- a/include/matx/core/reduce_utils.h +++ b/include/matx/core/reduce_utils.h @@ -41,6 +41,14 @@ namespace matx { + template + __MATX_HOST__ __MATX_INLINE__ bool IsFixedSizeReductionCompatible(const InputOp &in) { + if constexpr (requires { in.IsIdentityPermutation(); }) { + return in.IsIdentityPermutation(); + } + return true; + } + template __MATX_HOST__ __MATX_INLINE__ auto ReduceOutput(Func &&func, OutputOp &&out, InputOp &&in, BeginIter &&bi, EndIter &&ei) { @@ -66,6 +74,26 @@ namespace matx { return func(in, iter, bi, ei); } + template + __MATX_HOST__ __MATX_INLINE__ auto ReduceOutputFixed(Func &&func, OutputOp &&out, InputOp &&in) { + if constexpr (is_tensor_view_v) { + if (out.IsContiguous()) { + if constexpr(ConvertType) { + return func(in, + reinterpret_cast::value_type> *>(out.Data())); + } + else { + return func(in, + reinterpret_cast::value_type *>(out.Data())); + } + } + } + + detail::base_type_t out_base = out; + auto iter = RandomOperatorOutputIterator{out_base}; + return func(in, iter); + } + template __MATX_HOST__ __MATX_INLINE__ auto ReduceInput(Func &&func, OutputOp &&out, InputOp &&in) { typename detail::base_type_t in_base = in; @@ -96,9 +124,18 @@ namespace matx { return ReduceOutput(std::forward(func), std::forward(out), iter, BeginOffset{iter}, EndOffset{iter}); } + template + __MATX_HOST__ __MATX_INLINE__ auto ReduceInputFixed(Func &&func, OutputOp &&out, InputOp &&in) { + typename detail::base_type_t in_base = in; + auto collapsed = matx::lcollapse::Rank()>(rcollapse::Rank() - + remove_cvref_t::Rank()>(in_base)); + const auto &iter = matx::RandomOperatorIterator{collapsed}; + return ReduceOutputFixed(std::forward(func), std::forward(out), iter); + } + template __MATX_HOST__ __MATX_INLINE__ auto ReduceInputNoConvert(Func &&func, OutputOp &&out, InputOp &&in) { return ReduceInput(std::forward(func), std::forward(out), std::forward(in)); } } -#endif \ No newline at end of file +#endif diff --git a/include/matx/transforms/cub.h b/include/matx/transforms/cub.h index 389eef07e..f8e09d846 100644 --- a/include/matx/transforms/cub.h +++ b/include/matx/transforms/cub.h @@ -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) { @@ -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(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(TotalSize(a) / TotalSize(out_base)); + return cub::DeviceSegmentedReduce::Reduce(d_temp, temp_storage_bytes, in, out, + static_cast(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(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) { @@ -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; @@ -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(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(TotalSize(a) / TotalSize(out_base)); + return cub::DeviceSegmentedReduce::Sum(d_temp, temp_storage_bytes, in, out, + static_cast(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(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) { @@ -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) { @@ -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(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(TotalSize(a) / TotalSize(out_base)); + return cub::DeviceSegmentedReduce::Min(d_temp, temp_storage_bytes, in, out, + static_cast(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(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) { @@ -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) { @@ -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(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(TotalSize(a) / TotalSize(out_base)); + return cub::DeviceSegmentedReduce::Max(d_temp, temp_storage_bytes, in, out, + static_cast(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(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) { diff --git a/test/00_operators/ReductionTests.cu b/test/00_operators/ReductionTests.cu index c03409d75..4bb5e0d69 100644 --- a/test/00_operators/ReductionTests.cu +++ b/test/00_operators/ReductionTests.cu @@ -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::type; + + ExecType exec{}; + auto input = make_tensor({4, 3, 5}); + auto inner_output = make_tensor({4, 3}); + auto outer_output = make_tensor({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(batch + 1), + static_cast(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((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((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({4, 3, 5}); + auto minimum = make_tensor({4, 3}); + auto maximum = make_tensor({4, 3}); + auto product = make_tensor({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(batch + row + col + 1); + } + } + } + + auto expression = input + static_cast(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(batch + row + 2); + const auto last = static_cast(batch + row + 6); + TestType expected_product = static_cast(1); + for (index_t col = 0; col < input.Size(2); ++col) { + expected_product *= static_cast(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();