From ade6ed51eb81415f9273c77b2eaa7b47ef54c677 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Wed, 5 May 2021 09:14:00 +0200 Subject: [PATCH] Speed up Reduce operators for consecutive reduced axes (#7206) * Improves Reduction for three specific configurations * Support ReduceMean * add ReduceMax, ReduceMin * refactoring --- .../providers/cpu/reduction/reduction_ops.cc | 699 +++++--- .../providers/cpu/reduction/reduction_ops.h | 356 +++- onnxruntime/core/util/math_cpu.cc | 3 + .../test/providers/cpu/math/einsum_test.cc | 1 + .../cpu/reduction/reduction_ops_test.cc | 1494 ++++++++++++++++- .../cpu/reduction/reduction_ops.cc | 3 +- 6 files changed, 2292 insertions(+), 264 deletions(-) diff --git a/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc b/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc index ca89296dc2..d55a13a857 100644 --- a/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc @@ -207,57 +207,82 @@ REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL_DOUBLE_ONLY(ArgMin, 11, 12) REGISTER_UNARY_ELEMENTWISE_KERNEL(ArgMin, 13); REGISTER_UNARY_ELEMENTWISE_KERNEL_DOUBLE_ONLY(ArgMin, 13); -bool SetupForReduce(const Tensor* input_tensor_ptr, - const std::vector& axes_, - std::vector& axes, - TensorShape& new_input_shape, - std::vector& output_shape, - bool& empty_reduce, - const TensorShape* input_shape_override) { - ORT_ENFORCE(input_tensor_ptr != nullptr, "Input to be reduced is null"); - - if (input_shape_override) { - ORT_ENFORCE(input_tensor_ptr->Shape().Size() == input_shape_override->Size(), - "The input shape override's size does not match the input tensor's shape size"); - } - - new_input_shape = input_shape_override ? *input_shape_override : input_tensor_ptr->Shape(); - size_t ndim = new_input_shape.NumDimensions(); - if (ndim == 0) { - empty_reduce = true; - return false; - } - - axes.reserve(axes_.size()); - for (int64_t axis : axes_) { - axes.push_back(HandleNegativeAxis(axis, static_cast(ndim))); - } - - if (axes.empty()) { - // This is the default case for non-arg kind reductions. Reduce on all dimensions. - for (size_t i = 0; i < ndim; i++) { - axes.push_back(i); - } - } - - std::sort(axes.begin(), axes.end()); - - // If all reduced axes are located at the tail of the input shape, then copy could be skipped is required - bool need_copy = true; - if (axes.size() <= ndim && ndim > 0 && - axes.front() == static_cast(ndim - axes.size()) && - axes.back() == static_cast(ndim) - 1) { - need_copy = false; - } - - empty_reduce = false; - output_shape = new_input_shape.GetDims(); - for (auto a : axes) { - output_shape[a] = new_input_shape[a] > 0 ? 1 : 0; - empty_reduce |= output_shape[a] == 0; - } - return need_copy; +FastReduceKind operator|(FastReduceKind a, FastReduceKind b) { + return static_cast(static_cast(a) | static_cast(b)); } + +bool operator==(FastReduceKind a, FastReduceKind b) { + return static_cast(a) == static_cast(b); +} + +bool operator!=(FastReduceKind a, FastReduceKind b) { + return static_cast(a) != static_cast(b); +} + +bool IsFastReduceKindAvailable(FastReduceKind scenario, FastReduceKind available) { + return (static_cast(scenario) & static_cast(available)) > 0; +} + +bool ResultsNoTransposePrepareForReduce::equal(const std::vector& local_input_shape, + const std::vector& local_reduced_axes) { + if (input_shape.size() != local_input_shape.size()) + return false; + if (reduced_axes.size() != local_reduced_axes.size()) + return false; + for (std::vector::const_iterator it1 = input_shape.begin(), it2 = local_input_shape.begin(); + it1 != input_shape.end(); ++it1, ++it2) { + if (*it1 != *it2) + return false; + } + for (std::vector::const_iterator it1 = reduced_axes.begin(), it2 = local_reduced_axes.begin(); + it1 != reduced_axes.end(); ++it1, ++it2) { + if (*it1 != *it2) + return false; + } + return true; +} + +void ResultsNoTransposePrepareForReduce::ValidateNotEmpty() { + ORT_ENFORCE(last_loop_red_size > 0); + ORT_ENFORCE(last_loop_size > 0); + ORT_ENFORCE(projected_index.size() > 0); +} + +static void ValidateMustBeOverloaded() { + ORT_ENFORCE(false, "must be overloaded."); +} + +static void ValidateFastReduceKR(const std::vector& fast_shape, const Tensor& output) { + ORT_ENFORCE(fast_shape.size() == 2, "Only works on matrices with two dimensions."); + ORT_ENFORCE(fast_shape[0] == output.Shape().Size(), "Output size mismatch."); +} + +static void ValidateFastReduceRK(const std::vector& fast_shape, const Tensor& output) { + ORT_ENFORCE(fast_shape.size() == 2, "Only works on matrices with two dimensions."); + ORT_ENFORCE(fast_shape[1] == output.Shape().Size(), "Output size mismatch."); +} + +static void ValidateFastReduceKRK(const std::vector& fast_shape, const Tensor& output) { + ORT_ENFORCE(fast_shape.size() == 3, "Only works on matrices with two dimensions."); + ORT_ENFORCE(fast_shape[0] * fast_shape[2] == output.Shape().Size(), "Output size mismatch."); +} + +void ReduceAggregatorBase::FastReduceKR(const Tensor&, const std::vector&, Tensor&, concurrency::ThreadPool*) { + ValidateMustBeOverloaded(); +} +void ReduceAggregatorBase::FastReduceRK(const Tensor&, const std::vector&, Tensor&, concurrency::ThreadPool*) { + ValidateMustBeOverloaded(); +} +void ReduceAggregatorBase::FastReduceKRK(const Tensor&, const std::vector&, Tensor&, concurrency::ThreadPool*) { + ValidateMustBeOverloaded(); +} + +TensorOpCost ParallelReduceFastCost(int64_t n_row, int64_t n_col, int64_t element_size) { + return TensorOpCost{static_cast(n_col * n_row * element_size), + static_cast(n_row * element_size), + static_cast(n_col * n_row * element_size * 2)}; +} + void NoTransposePrepareForReduce(const TensorShape& new_input_shape, const std::vector& reduced_axes, ResultsNoTransposePrepareForReduce& results) { @@ -355,17 +380,21 @@ void NoTransposePrepareForReduce(const TensorShape& new_input_shape, } } -template -void NoTransposeReduce(Tensor* output, const TensorShape& new_input_shape, const Tensor& input, - const std::vector& reduced_axes, concurrency::ThreadPool* tp, - ResultsNoTransposePrepareForReduce& last_results) { +void ValidateNoTransposeReduce(int64_t count) { + ORT_ENFORCE(count == 1, "Reduction on all axes, output size should be 1."); +} + +template +void NoTransposeReduce1Loop(Tensor* output, const TensorShape& new_input_shape, const Tensor& input, + const std::vector& reduced_axes, concurrency::ThreadPool* tp, + ResultsNoTransposePrepareForReduce& last_results) { auto output_shape = output->Shape(); - const T* from_data = input.template Data(); + const typename AGG::input_type* from_data = input.template Data(); typename AGG::value_type* to_data = output->template MutableData(); int64_t count = output_shape.Size(); if (reduced_axes.size() == 0 || reduced_axes.size() == new_input_shape.NumDimensions()) { - ORT_ENFORCE(count == 1, "Reduction on all axes, output size should be 1."); + ValidateNoTransposeReduce(count); int64_t input_size = new_input_shape.Size(); to_data[0] = AGG(input_size, from_data[0]).aggall(from_data); return; @@ -376,74 +405,101 @@ void NoTransposeReduce(Tensor* output, const TensorShape& new_input_shape, const if (last_results.last_loop_red_size == 0 || last_results.last_loop_size == 0) return; } - ORT_ENFORCE(last_results.last_loop_red_size > 0); - ORT_ENFORCE(last_results.last_loop_size > 0); - ORT_ENFORCE(last_results.projected_index.size() > 0); + last_results.ValidateNotEmpty(); int64_t denominator = last_results.last_loop_red_size * last_results.projected_index.size(); - if (AGG::two_loops()) { - auto fn = [&](std::ptrdiff_t first, std::ptrdiff_t end) { - int64_t loop; - const T* loop_red_ptr; - const T* loop_red_ptr_end; - int64_t current_index = first * last_results.last_loop_size; - for (int64_t main_index = first; main_index < end; ++main_index) { - for (loop = 0; loop < last_results.last_loop_size; ++loop, ++current_index) { - int64_t origin = last_results.unprojected_index[main_index] + loop * last_results.last_loop_inc; - AGG accumulator(denominator, from_data[origin + last_results.projected_index[0]]); - for (auto it = last_results.projected_index.begin(); it != last_results.projected_index.end(); ++it) { - loop_red_ptr = from_data + (origin + *it); - loop_red_ptr_end = loop_red_ptr + last_results.last_loop_red_size * last_results.last_loop_red_inc; - for (; loop_red_ptr != loop_red_ptr_end; loop_red_ptr += last_results.last_loop_red_inc) { - accumulator.update0(*loop_red_ptr); - } + auto fn = [&](std::ptrdiff_t first, std::ptrdiff_t end) { + int64_t loop; + const typename AGG::input_type* loop_red_ptr; + const typename AGG::input_type* loop_red_ptr_end; + int64_t current_index = first * last_results.last_loop_size; + for (int64_t main_index = first; main_index < end; ++main_index) { + for (loop = 0; loop < last_results.last_loop_size; ++loop, ++current_index) { + int64_t origin = last_results.unprojected_index[main_index] + loop * last_results.last_loop_inc; + AGG accumulator(denominator, from_data[origin + last_results.projected_index[0]]); + for (auto it = last_results.projected_index.begin(); it != last_results.projected_index.end(); ++it) { + loop_red_ptr = from_data + (origin + *it); + loop_red_ptr_end = loop_red_ptr + last_results.last_loop_red_size * last_results.last_loop_red_inc; + for (; loop_red_ptr != loop_red_ptr_end; loop_red_ptr += last_results.last_loop_red_inc) { + accumulator.update(*loop_red_ptr); } - for (auto it = last_results.projected_index.begin(); it != last_results.projected_index.end(); ++it) { - loop_red_ptr = from_data + (origin + *it); - loop_red_ptr_end = loop_red_ptr + last_results.last_loop_red_size * last_results.last_loop_red_inc; - for (; loop_red_ptr != loop_red_ptr_end; loop_red_ptr += last_results.last_loop_red_inc) { - accumulator.update(*loop_red_ptr); - } - } - to_data[current_index] = accumulator.get_value(); } + to_data[current_index] = accumulator.get_value(); } - }; + } + }; - auto cost = TensorOpCost{(double)(last_results.projected_index.size() * sizeof(T) * last_results.last_loop_size * last_results.last_loop_red_size), - (double)last_results.last_loop_size * last_results.last_loop_red_size, - (double)last_results.projected_index.size() * last_results.last_loop_size * last_results.last_loop_red_size * 2}; - concurrency::ThreadPool::TryParallelFor(tp, count / last_results.last_loop_size, cost, fn); - } else { - auto fn = [&](std::ptrdiff_t first, std::ptrdiff_t end) { - int64_t loop; - const T* loop_red_ptr; - const T* loop_red_ptr_end; - int64_t current_index = first * last_results.last_loop_size; - for (int64_t main_index = first; main_index < end; ++main_index) { - for (loop = 0; loop < last_results.last_loop_size; ++loop, ++current_index) { - int64_t origin = last_results.unprojected_index[main_index] + loop * last_results.last_loop_inc; - AGG accumulator(denominator, from_data[origin + last_results.projected_index[0]]); - for (auto it = last_results.projected_index.begin(); it != last_results.projected_index.end(); ++it) { - loop_red_ptr = from_data + (origin + *it); - loop_red_ptr_end = loop_red_ptr + last_results.last_loop_red_size * last_results.last_loop_red_inc; - for (; loop_red_ptr != loop_red_ptr_end; loop_red_ptr += last_results.last_loop_red_inc) { - accumulator.update(*loop_red_ptr); - } - } - to_data[current_index] = accumulator.get_value(); - } - } - }; - - auto cost = TensorOpCost{(double)(last_results.projected_index.size() * sizeof(T) * last_results.last_loop_size * last_results.last_loop_red_size), - (double)last_results.last_loop_size * last_results.last_loop_red_size, - (double)last_results.projected_index.size() * last_results.last_loop_size * last_results.last_loop_red_size}; - concurrency::ThreadPool::TryParallelFor(tp, count / last_results.last_loop_size, cost, fn); - } + auto cost = TensorOpCost{(double)(last_results.projected_index.size() * sizeof(typename AGG::input_type) * + last_results.last_loop_size * last_results.last_loop_red_size), + (double)last_results.last_loop_size * last_results.last_loop_red_size, + (double)last_results.projected_index.size() * last_results.last_loop_size * + last_results.last_loop_red_size}; + concurrency::ThreadPool::TryParallelFor(tp, count / last_results.last_loop_size, cost, fn); } -void DropDimensions(const std::vector& input_shape, const std::vector& axes, std::vector& dropped_axes) { +template +void NoTransposeReduce2Loops(Tensor* output, const TensorShape& new_input_shape, const Tensor& input, + const std::vector& reduced_axes, concurrency::ThreadPool* tp, + ResultsNoTransposePrepareForReduce& last_results) { + auto output_shape = output->Shape(); + const typename AGG::input_type* from_data = input.template Data(); + typename AGG::value_type* to_data = output->template MutableData(); + int64_t count = output_shape.Size(); + + if (reduced_axes.size() == 0 || reduced_axes.size() == new_input_shape.NumDimensions()) { + ValidateNoTransposeReduce(count); + int64_t input_size = new_input_shape.Size(); + to_data[0] = AGG(input_size, from_data[0]).aggall(from_data); + return; + } + + if (!last_results.equal(new_input_shape.GetDims(), reduced_axes)) { + NoTransposePrepareForReduce(new_input_shape, reduced_axes, last_results); + if (last_results.last_loop_red_size == 0 || last_results.last_loop_size == 0) + return; + } + last_results.ValidateNotEmpty(); + int64_t denominator = last_results.last_loop_red_size * last_results.projected_index.size(); + + auto fn = [&](std::ptrdiff_t first, std::ptrdiff_t end) { + int64_t loop; + const typename AGG::input_type* loop_red_ptr; + const typename AGG::input_type* loop_red_ptr_end; + int64_t current_index = first * last_results.last_loop_size; + for (int64_t main_index = first; main_index < end; ++main_index) { + for (loop = 0; loop < last_results.last_loop_size; ++loop, ++current_index) { + int64_t origin = last_results.unprojected_index[main_index] + loop * last_results.last_loop_inc; + AGG accumulator(denominator, from_data[origin + last_results.projected_index[0]]); + for (auto it = last_results.projected_index.begin(); it != last_results.projected_index.end(); ++it) { + loop_red_ptr = from_data + (origin + *it); + loop_red_ptr_end = loop_red_ptr + last_results.last_loop_red_size * last_results.last_loop_red_inc; + for (; loop_red_ptr != loop_red_ptr_end; loop_red_ptr += last_results.last_loop_red_inc) { + accumulator.update0(*loop_red_ptr); + } + } + for (auto it = last_results.projected_index.begin(); it != last_results.projected_index.end(); ++it) { + loop_red_ptr = from_data + (origin + *it); + loop_red_ptr_end = loop_red_ptr + last_results.last_loop_red_size * last_results.last_loop_red_inc; + for (; loop_red_ptr != loop_red_ptr_end; loop_red_ptr += last_results.last_loop_red_inc) { + accumulator.update(*loop_red_ptr); + } + } + to_data[current_index] = accumulator.get_value(); + } + } + }; + + auto cost = TensorOpCost{(double)(last_results.projected_index.size() * sizeof(typename AGG::input_type) * + last_results.last_loop_size * last_results.last_loop_red_size), + (double)last_results.last_loop_size * last_results.last_loop_red_size, + (double)last_results.projected_index.size() * last_results.last_loop_size * + last_results.last_loop_red_size * 2}; + concurrency::ThreadPool::TryParallelFor(tp, count / last_results.last_loop_size, cost, fn); +} + +void DropDimensions(const std::vector& input_shape, + const std::vector& axes, + std::vector& dropped_axes) { auto dropped_dims = input_shape; for (auto i : axes) { dropped_dims[i] = -1; @@ -455,68 +511,273 @@ void DropDimensions(const std::vector& input_shape, const std::vector -void CommonReduce(OpKernelContext* ctx, - const std::vector axes_, int64_t keepdims_, - ResultsNoTransposePrepareForReduce& last_results, - bool noop_with_empty_axes) { - std::vector axes; - const Tensor* input = ctx->Input(0); - auto reduced_dims = input->Shape().GetDims(); - std::vector output_shape; - bool empty_reduce; - TensorShape new_input_shape; +FastReduceKind OptimizeShapeForFastReduce(const std::vector& input_shape, + const std::vector& reduced_axes, + std::vector& fast_shape, + std::vector& fast_output_shape, + std::vector& fast_axes, + bool keep_dims, bool noop_with_empty_axes) { + if (input_shape.empty()) { + fast_shape = input_shape; + fast_output_shape = input_shape; + fast_axes = reduced_axes; + return FastReduceKind::kNone; + } - if (ctx->InputCount() == 2) { - // second input holds the axes. - const Tensor* axes_tensor = ctx->Input(1); - ORT_ENFORCE(axes_tensor != nullptr, "Axes input is null"); - ORT_ENFORCE(axes_tensor->Shape().NumDimensions() == 1, - "An axes tensor must be a vector tensor."); - auto nDims = static_cast(axes_tensor->Shape()[0]); - const auto* data = axes_tensor->template Data(); - std::vector input_axes(data, data + nDims); - if (input_axes.empty() && noop_with_empty_axes) { - auto* output = ctx->Output(0, input->Shape()); - memcpy(output->template MutableData(), input->template Data(), input->SizeInBytes()); - return; + std::set axes; + if (reduced_axes.size() == 0 && !noop_with_empty_axes) { + for (int64_t i = 0; i < (int64_t)input_shape.size(); ++i) { + axes.insert(i); } - SetupForReduce(input, input_axes, axes, new_input_shape, output_shape, empty_reduce, nullptr); } else { - SetupForReduce(input, axes_, axes, new_input_shape, output_shape, empty_reduce, nullptr); + for (auto it = reduced_axes.begin(); it != reduced_axes.end(); ++it) { + axes.insert(HandleNegativeAxis(*it, static_cast(input_shape.size()))); + } + } + + fast_output_shape.clear(); + fast_output_shape.reserve(input_shape.size()); + bool empty_reduce = false; + std::vector reduce(input_shape.size()); + for (int64_t i = 0; i < (int64_t)input_shape.size(); ++i) { + reduce[i] = axes.find(i) != axes.end(); + if (reduce[i]) { + empty_reduce |= input_shape[i] == 0; + if (keep_dims) + fast_output_shape.push_back(input_shape[i] > 0 ? 1 : 0); + } else { + fast_output_shape.push_back(input_shape[i]); + } } if (empty_reduce) { - Tensor* output = ctx->Output(0, keepdims_ ? output_shape : std::vector()); + return FastReduceKind::kEmpty; + } + + if (reduced_axes.empty()) { + fast_shape.resize(1); + fast_shape[0] = 1; + for (auto a : input_shape) { + fast_shape[0] *= a; + } + if (noop_with_empty_axes) { + fast_axes.clear(); + fast_output_shape = input_shape; + return FastReduceKind::kK; + } else { + if (keep_dims) { + fast_output_shape.resize(input_shape.size(), 1); + } else { + fast_output_shape.clear(); + } + fast_axes.resize(1); + fast_axes[0] = 0; + return FastReduceKind::kR; + } + } + + fast_shape.clear(); + fast_axes.clear(); + fast_shape.reserve(input_shape.size()); + fast_axes.reserve(reduced_axes.size()); + + fast_shape.push_back(input_shape[0]); + if (reduce[0]) + fast_axes.push_back(0); + for (size_t i = 1; i < input_shape.size(); ++i) { + if (reduce[i] == reduce[i - 1]) { + fast_shape[fast_shape.size() - 1] *= input_shape[i]; + } else { + if (reduce[i]) { + fast_axes.push_back(fast_shape.size()); + } + fast_shape.push_back(input_shape[i]); + } + } + if (fast_shape.size() == 1) { + return reduce[0] ? FastReduceKind::kR : FastReduceKind::kK; + } + if (fast_shape.size() == 2) { + return reduce[0] ? FastReduceKind::kRK : FastReduceKind::kKR; + } + if (fast_shape.size() == 3 && !reduce[0]) { + return FastReduceKind::kKRK; + } + return FastReduceKind::kNone; +} + +void ValidateCommonFastReduce(const Tensor* axes_tensor) { + ORT_ENFORCE(axes_tensor != nullptr, "Axes input is null"); + ORT_ENFORCE(axes_tensor->Shape().NumDimensions() == 1, + "An axes tensor must be a vector tensor."); +} + +//template +bool CommonFastReduceCopy(OpKernelContext* ctx, std::vector& input_axes, bool noop_with_empty_axes) { + if (ctx->InputCount() == 2) { + // second input holds the axes. + const Tensor* axes_tensor = ctx->Input(1); + ValidateCommonFastReduce(axes_tensor); + auto nDims = static_cast(axes_tensor->Shape()[0]); + const auto* data = axes_tensor->template Data(); + input_axes.insert(input_axes.begin(), data, data + nDims); + if (input_axes.empty() && noop_with_empty_axes) { + const Tensor* input = ctx->Input(0); + auto* output = ctx->Output(0, input->Shape()); + memcpy(reinterpret_cast(output->template MutableData()), + reinterpret_cast(input->template Data()), + input->SizeInBytes()); + return true; + } + } + return false; +} + +typedef void fast_reduce_fct(const Tensor& input, const std::vector& fast_shape, + Tensor& output, concurrency::ThreadPool* tp); + +bool CommonFastReduceSwitch(OpKernelContext* ctx, + const std::vector& axes_, + int64_t keepdims_, + bool noop_with_empty_axes, + FastReduceKind& fast_kind, + std::vector& fast_shape, + std::vector& output_shape, + std::vector& fast_axes, + FastReduceKind which_fast_reduce, + fast_reduce_fct* case_kr, + fast_reduce_fct* case_rk, + fast_reduce_fct* case_krk) { + std::vector axes; + const Tensor* input = ctx->Input(0); + auto reduced_dims = input->Shape().GetDims(); + std::vector input_axes; + + if (CommonFastReduceCopy(ctx, input_axes, noop_with_empty_axes)) { + return true; + } + + fast_kind = OptimizeShapeForFastReduce( + reduced_dims, input_axes.empty() ? axes_ : input_axes, + fast_shape, output_shape, fast_axes, keepdims_, noop_with_empty_axes); + if (which_fast_reduce != FastReduceKind::kNone) { + if (IsFastReduceKindAvailable(fast_kind, which_fast_reduce)) { + Tensor* output = ctx->Output(0, output_shape); + switch (fast_kind) { + case FastReduceKind::kKR: { + ValidateFastReduceKR(fast_shape, *output); + case_kr(*input, fast_shape, *output, ctx->GetOperatorThreadPool()); + return true; + } + case FastReduceKind::kRK: { + ValidateFastReduceRK(fast_shape, *output); + case_rk(*input, fast_shape, *output, ctx->GetOperatorThreadPool()); + return true; + } + case FastReduceKind::kKRK: { + ValidateFastReduceKRK(fast_shape, *output); + case_krk(*input, fast_shape, *output, ctx->GetOperatorThreadPool()); + return true; + } + case FastReduceKind::kR: + case FastReduceKind::kK: + case FastReduceKind::kNone: + default: + // Former implementation prevails in this case. + break; + } + } + } + return false; +} + +template +bool CommonFastReduce(OpKernelContext* ctx, + const std::vector& axes_, + int64_t keepdims_, + bool noop_with_empty_axes, + FastReduceKind& fast_kind, + std::vector& fast_shape, + std::vector& output_shape, + std::vector& fast_axes) { + return CommonFastReduceSwitch(ctx, axes_, keepdims_, noop_with_empty_axes, fast_kind, fast_shape, output_shape, fast_axes, + AGG::WhichFastReduce(), &AGG::FastReduceKR, &AGG::FastReduceRK, &AGG::FastReduceKRK); +} + +static void ValidateKeepDims(const TensorShape& shape, int64_t keepdims) { + ORT_ENFORCE(keepdims, + "Can't reduce on dim with value of 0 if 'keepdims' is false. " + "Invalid output shape would be produced. input_shape:", + shape); +} + +static void ValidateKeepDims(const Tensor* input, int64_t keepdims) { + ValidateKeepDims(input->Shape(), keepdims); +} + +template +void CommonReduce1Loop(OpKernelContext* ctx, + const std::vector& axes_, int64_t keepdims_, + bool noop_with_empty_axes) { + FastReduceKind fast_kind; + std::vector fast_shape; + std::vector output_shape; + std::vector fast_axes; + if (CommonFastReduce(ctx, axes_, keepdims_, noop_with_empty_axes, + fast_kind, fast_shape, output_shape, fast_axes)) { + return; + } + + const Tensor* input = ctx->Input(0); + Tensor* output = ctx->Output(0, output_shape); + if (fast_kind == FastReduceKind::kEmpty) { + const TensorShape& new_input_shape = input->Shape(); if (new_input_shape.Size() == 1) { - const T* from_data = input->template Data(); + const typename AGG::input_type* from_data = input->template Data(); typename AGG::value_type* to_data = output->template MutableData(); AGG agg(1, *from_data); - if (agg.two_loops()) { - agg.update0(*from_data); - agg.update(*from_data); - } else { - agg.update(*from_data); - } + agg.update(*from_data); *to_data = agg.get_value(); } else { - ORT_ENFORCE(keepdims_, - "Can't reduce on dim with value of 0 if 'keepdims' is false. " - "Invalid output shape would be produced. input_shape:", - input->Shape()); + ValidateKeepDims(input, keepdims_); } return; } - Tensor* output; - if (keepdims_) { - output = ctx->Output(0, output_shape); - } else { - std::vector dropped_axes; - DropDimensions(output_shape, axes, dropped_axes); - output = ctx->Output(0, dropped_axes); + ResultsNoTransposePrepareForReduce last_results; + NoTransposeReduce1Loop(output, fast_shape, *input, fast_axes, ctx->GetOperatorThreadPool(), last_results); +} + +template +void CommonReduce2Loops(OpKernelContext* ctx, + const std::vector& axes_, int64_t keepdims_, + bool noop_with_empty_axes) { + FastReduceKind fast_kind; + std::vector fast_shape, output_shape, fast_axes; + if (CommonFastReduce(ctx, axes_, keepdims_, noop_with_empty_axes, + fast_kind, fast_shape, output_shape, fast_axes)) { + return; } - NoTransposeReduce(output, new_input_shape, *input, axes, ctx->GetOperatorThreadPool(), last_results); + + const Tensor* input = ctx->Input(0); + Tensor* output = ctx->Output(0, output_shape); + if (fast_kind == FastReduceKind::kEmpty) { + const TensorShape& new_input_shape = input->Shape(); + if (new_input_shape.Size() == 1) { + const typename AGG::input_type* from_data = input->template Data(); + typename AGG::value_type* to_data = output->template MutableData(); + AGG agg(1, *from_data); + agg.update0(*from_data); + agg.update(*from_data); + *to_data = agg.get_value(); + } else { + ValidateKeepDims(input, keepdims_); + } + return; + } + + ResultsNoTransposePrepareForReduce last_results; + NoTransposeReduce2Loops(output, fast_shape, *input, fast_axes, ctx->GetOperatorThreadPool(), last_results); } template @@ -524,64 +785,55 @@ Status ReduceL1::Compute(OpKernelContext* ctx) const { // The following variable does not change if the input tensor and the // axes do not either. It could be either cached in ctx or precomputed // in the constructor if shape and axes are known at this stage. - ResultsNoTransposePrepareForReduce last_results; - CommonReduce>(ctx, axes_, keepdims_, last_results); + CommonReduce1Loop>(ctx, axes_, keepdims_); return Status::OK(); } template Status ReduceL2::Compute(OpKernelContext* ctx) const { - ResultsNoTransposePrepareForReduce last_results; - CommonReduce>(ctx, axes_, keepdims_, last_results); + CommonReduce1Loop>(ctx, axes_, keepdims_); return Status::OK(); } template Status ReduceLogSum::Compute(OpKernelContext* ctx) const { - ResultsNoTransposePrepareForReduce last_results; - CommonReduce>(ctx, axes_, keepdims_, last_results); + CommonReduce1Loop>(ctx, axes_, keepdims_); return Status::OK(); } template Status ReduceLogSumExp::Compute(OpKernelContext* ctx) const { - ResultsNoTransposePrepareForReduce last_results; - CommonReduce>(ctx, axes_, keepdims_, last_results); + CommonReduce2Loops>(ctx, axes_, keepdims_); return Status::OK(); } template Status ReduceMax::Compute(OpKernelContext* ctx) const { - ResultsNoTransposePrepareForReduce last_results; - CommonReduce>(ctx, axes_, keepdims_, last_results); + CommonReduce1Loop>(ctx, axes_, keepdims_); return Status::OK(); } template Status ReduceMean::Compute(OpKernelContext* ctx) const { - ResultsNoTransposePrepareForReduce last_results; - CommonReduce>(ctx, axes_, keepdims_, last_results); + CommonReduce1Loop>(ctx, axes_, keepdims_); return Status::OK(); } template Status ReduceMin::Compute(OpKernelContext* ctx) const { - ResultsNoTransposePrepareForReduce last_results; - CommonReduce>(ctx, axes_, keepdims_, last_results); + CommonReduce1Loop>(ctx, axes_, keepdims_); return Status::OK(); } template Status ReduceProd::Compute(OpKernelContext* ctx) const { - ResultsNoTransposePrepareForReduce last_results; - CommonReduce>(ctx, axes_, keepdims_, last_results); + CommonReduce1Loop>(ctx, axes_, keepdims_); return Status::OK(); } template Status ReduceSum::Compute(OpKernelContext* ctx) const { - ResultsNoTransposePrepareForReduce last_results; - CommonReduce>(ctx, axes_, keepdims_, last_results, noop_with_empty_axes_); + CommonReduce1Loop>(ctx, axes_, keepdims_, noop_with_empty_axes_); return Status::OK(); } @@ -590,67 +842,79 @@ Tensor ReduceSum::Impl(const Tensor& input, const std::vector& reduc AllocatorPtr allocator, concurrency::ThreadPool* tp, bool keep_dims, const TensorShape* input_shape_override) { std::vector axes; - auto reduced_dims = input.Shape().GetDims(); - std::vector output_shape; - TensorShape new_input_shape; - bool empty_reduce; - SetupForReduce(&input, reduce_axes, axes, new_input_shape, output_shape, empty_reduce, input_shape_override); + std::vector output_shape, fast_shape, fast_axes; + TensorShape new_input_shape = input_shape_override == nullptr ? input.Shape() : *input_shape_override; + auto reduced_dims = new_input_shape.GetDims(); - if (empty_reduce) { - Tensor output(input.DataType(), keep_dims ? output_shape : std::vector(), allocator); + FastReduceKind fast_kind = OptimizeShapeForFastReduce( + reduced_dims, reduce_axes, fast_shape, output_shape, fast_axes, keep_dims, false); + + Tensor output(input.DataType(), keep_dims ? output_shape : std::vector(), allocator); + + if (fast_kind == FastReduceKind::kEmpty) { if (new_input_shape.Size() == 1) { const T* from_data = input.template Data(); T* to_data = output.template MutableData(); *to_data = *from_data; } else { - ORT_ENFORCE(keep_dims, - "Can't reduce on dim with value of 0 if 'keepdims' is false. " - "Invalid output shape would be produced. input_shape:", - new_input_shape); + ValidateKeepDims(new_input_shape, keep_dims); } return output; } - if (keep_dims) { - ResultsNoTransposePrepareForReduce last_results; - Tensor output(input.DataType(), output_shape, allocator); - NoTransposeReduce>(&output, new_input_shape, input, axes, tp, last_results); - return output; - } else { - ResultsNoTransposePrepareForReduce last_results; - std::vector dropped_axes; - DropDimensions(output_shape, axes, dropped_axes); - Tensor output(input.DataType(), dropped_axes, allocator); - NoTransposeReduce>(&output, new_input_shape, input, axes, tp, last_results); - return output; + if (IsFastReduceKindAvailable(fast_kind, ReduceAggregatorSum::WhichFastReduce())) { + switch (fast_kind) { + case FastReduceKind::kKR: { + ValidateFastReduceKR(fast_shape, output); + ReduceAggregatorSum::FastReduceKR(input, fast_shape, output, tp); + return output; + } + case FastReduceKind::kRK: { + ValidateFastReduceRK(fast_shape, output); + ReduceAggregatorSum::FastReduceRK(input, fast_shape, output, tp); + return output; + } + case FastReduceKind::kKRK: { + ValidateFastReduceKRK(fast_shape, output); + ReduceAggregatorSum::FastReduceKRK(input, fast_shape, output, tp); + return output; + } + case FastReduceKind::kR: + case FastReduceKind::kK: + case FastReduceKind::kNone: + default: + // Former implementation prevails in this case. + break; + } } + + ResultsNoTransposePrepareForReduce last_results; + NoTransposeReduce1Loop>(&output, fast_shape, input, fast_axes, tp, last_results); + return output; } template Status ReduceSumSquare::Compute(OpKernelContext* ctx) const { - ResultsNoTransposePrepareForReduce last_results; - CommonReduce>(ctx, axes_, keepdims_, last_results); + CommonReduce1Loop>(ctx, axes_, keepdims_); return Status::OK(); } template Status ArgMax::Compute(OpKernelContext* ctx) const { - ResultsNoTransposePrepareForReduce last_results; if (select_last_index_) { - CommonReduce>(ctx, axes_, keepdims_, last_results); + CommonReduce1Loop>(ctx, axes_, keepdims_); } else { - CommonReduce>(ctx, axes_, keepdims_, last_results); + CommonReduce1Loop>(ctx, axes_, keepdims_); } return Status::OK(); } template Status ArgMin::Compute(OpKernelContext* ctx) const { - ResultsNoTransposePrepareForReduce last_results; if (select_last_index_) { - CommonReduce>(ctx, axes_, keepdims_, last_results); + CommonReduce1Loop>(ctx, axes_, keepdims_); } else { - CommonReduce>(ctx, axes_, keepdims_, last_results); + CommonReduce1Loop>(ctx, axes_, keepdims_); } return Status::OK(); } @@ -664,4 +928,17 @@ template class ReduceSum; template class ReduceSum; template class ReduceSum; +template void CommonReduce1Loop>(OpKernelContext* ctx, + const std::vector& axes_, int64_t keepdims_, + bool noop_with_empty_axes); +template void CommonReduce1Loop>(OpKernelContext* ctx, + const std::vector& axes_, int64_t keepdims_, + bool noop_with_empty_axes); +template void CommonReduce1Loop>(OpKernelContext* ctx, + const std::vector& axes_, int64_t keepdims_, + bool noop_with_empty_axes); +template void CommonReduce1Loop>(OpKernelContext* ctx, + const std::vector& axes_, int64_t keepdims_, + bool noop_with_empty_axes); + } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/reduction/reduction_ops.h b/onnxruntime/core/providers/cpu/reduction/reduction_ops.h index 092e6d0abe..e3ddfac4c4 100644 --- a/onnxruntime/core/providers/cpu/reduction/reduction_ops.h +++ b/onnxruntime/core/providers/cpu/reduction/reduction_ops.h @@ -8,12 +8,57 @@ #include "core/common/optional.h" #include "core/framework/op_kernel.h" #include "core/providers/cpu/containers.h" +#include "core/util/math.h" #include "core/util/math_cpuonly.h" #include "core/platform/threadpool.h" +#include "core/common/safeint.h" #include namespace onnxruntime { +enum FastReduceKind { + kNone = 0, // no fast implementation + kK = 1, // kept dim = no reduce + kR = 2, // reduced dim = all reduced + kKR = 4, // kept dim, reduced dim + kRK = 8, // reduced dim, kept dim + kKRK = 16, // kept dim, reduced dim, kept dim + kEmpty = 32 // empty reduce +}; + +FastReduceKind operator|(FastReduceKind a, FastReduceKind b); + +bool operator==(FastReduceKind a, FastReduceKind b); + +bool operator!=(FastReduceKind a, FastReduceKind b); + +bool IsFastReduceKindAvailable(FastReduceKind scenario, FastReduceKind available); + +/* Evaluate the cost of parallelized FastReduce implementations. */ +TensorOpCost ParallelReduceFastCost(int64_t n_row, int64_t n_col, int64_t element_size); + +/** + This only improves reduce function when reduced axes are contiguous: + if len(shape) == 4, any single axis is ok, axes=(0, 1) or (1, 2) or (2, 3) is ok, + axes=(0, 2) is not covered by this change, former implementation prevails. + In that case, the shape can be compressed into three cases: + (K = axis not reduced, R = reduced axis): + + * KR - reduction on the last dimensions + * RK - reduction on the first dimensions + * KRK - reduction on the middle dimensions. + + For these three configuration, the reduction may be optimized + with vectors operations. Method WhichFastReduce() returns which case + case be optimized for which aggregator. +*/ +FastReduceKind OptimizeShapeForFastReduce(const std::vector& input_shape, + const std::vector& reduced_axes, + std::vector& fast_shape, + std::vector& fast_output_shape, + std::vector& fast_axes, + bool keep_dims, bool noop_with_empty_axes = false); + class ResultsNoTransposePrepareForReduce { public: std::vector input_shape; @@ -32,42 +77,35 @@ class ResultsNoTransposePrepareForReduce { last_loop_inc = 0; } - bool equal(const std::vector& local_input_shape, const std::vector& local_reduced_axes) { - if (input_shape.size() != local_input_shape.size()) - return false; - if (reduced_axes.size() != local_reduced_axes.size()) - return false; - for (std::vector::const_iterator it1 = input_shape.begin(), it2 = local_input_shape.begin(); - it1 != input_shape.end(); ++it1, ++it2) { - if (*it1 != *it2) - return false; - } - for (std::vector::const_iterator it1 = reduced_axes.begin(), it2 = local_reduced_axes.begin(); - it1 != reduced_axes.end(); ++it1, ++it2) { - if (*it1 != *it2) - return false; - } - return true; - } + bool equal(const std::vector& local_input_shape, const std::vector& local_reduced_axes); + void ValidateNotEmpty(); }; template inline T reduce_sqrt(T value) { return std::sqrt(value); } template <> -inline int64_t reduce_sqrt(int64_t value) { return static_cast(std::sqrt(static_cast(value))); } +inline int64_t reduce_sqrt(int64_t value) { + return static_cast(std::sqrt(static_cast(value))); +} template <> -inline int32_t reduce_sqrt(int32_t value) { return static_cast(std::sqrt(static_cast(value))); } +inline int32_t reduce_sqrt(int32_t value) { + return static_cast(std::sqrt(static_cast(value))); +} template inline T reduce_log(T value) { return static_cast(std::log(value)); } template <> -inline int64_t reduce_log(int64_t value) { return static_cast(std::log(static_cast(value))); } +inline int64_t reduce_log(int64_t value) { + return static_cast(std::log(static_cast(value))); +} template <> -inline int32_t reduce_log(int32_t value) { return static_cast(std::log(static_cast(value))); } +inline int32_t reduce_log(int32_t value) { + return static_cast(std::log(static_cast(value))); +} template inline T reduce_exp(T value) { return static_cast(std::exp(value)); } @@ -102,9 +140,19 @@ inline bool reduce_isnan(int32_t) { return false; } template <> inline bool reduce_isnan(int64_t) { return false; } -template -class ReduceAggregator { +class ReduceAggregatorBase { public: + // Fast reduction: see OptimizeShapeForFastReduce's comment. + static inline FastReduceKind WhichFastReduce() { return FastReduceKind::kNone; } + static void FastReduceKR(const Tensor&, const std::vector&, Tensor&, concurrency::ThreadPool*); + static void FastReduceRK(const Tensor&, const std::vector&, Tensor&, concurrency::ThreadPool*); + static void FastReduceKRK(const Tensor&, const std::vector&, Tensor&, concurrency::ThreadPool*); +}; + +template +class ReduceAggregator : public ReduceAggregatorBase { + public: + typedef T input_type; typedef TVAL value_type; protected: @@ -116,12 +164,10 @@ class ReduceAggregator { N_ = N; accumulator_ = init; } - inline void update(const T&) { ORT_ENFORCE(false, "must be overloaded."); } - inline void update0(const T&) { ORT_ENFORCE(false, "must be overloaded."); } - inline TVAL aggall(const T*) { ORT_ENFORCE(false, "must be overloaded."); } + inline void update(const T&) {} + inline void update0(const T&) {} + inline TVAL aggall(const T*) {} inline TVAL get_value() { return accumulator_; } - inline void enforce(const ResultsNoTransposePrepareForReduce&) {} - static inline bool two_loops() { return false; } }; template @@ -132,6 +178,60 @@ class ReduceAggregatorSum : public ReduceAggregator { inline TVAL aggall(const T* from_data) { return Eigen::Map>(from_data, this->N_).sum(); } + + // Fast reduction + static inline FastReduceKind WhichFastReduce() { + return FastReduceKind::kKR | FastReduceKind::kRK | FastReduceKind::kKRK; + } + + static void FastReduceKR(const Tensor& input, const std::vector& fast_shape, + Tensor& output, concurrency::ThreadPool* tp) { + const T* data = input.Data(); + T* out = output.MutableData(); + int64_t stridei = fast_shape[1]; + concurrency::ThreadPool::TryParallelFor( + tp, fast_shape[0], ParallelReduceFastCost(1, stridei, sizeof(T)), + [data, stridei, out](ptrdiff_t first, ptrdiff_t last) { + for (ptrdiff_t d = first; d < last; ++d) { + out[d] = ConstEigenVectorArrayMap(data + d * stridei, stridei).sum(); + } + }); + } + + static void FastReduceRK(const Tensor& input, const std::vector& fast_shape, + Tensor& output, concurrency::ThreadPool* tp) { + int64_t N = fast_shape[1]; + const T* data = input.Data(); + T* out = output.MutableData(); + + int64_t n_rows = fast_shape[0]; + memcpy(out, data, N * sizeof(T)); + concurrency::ThreadPool::TryParallelFor( + tp, N, ParallelReduceFastCost(1, n_rows, sizeof(T)), + [data, out, N, n_rows](ptrdiff_t begin, ptrdiff_t end) { + for (int64_t row = 1; row < n_rows; ++row) { + EigenVectorArrayMap(out + begin, end - begin) += ConstEigenVectorArrayMap( + data + row * N + begin, end - begin); + } + }); + } + + static void FastReduceKRK(const Tensor& input, const std::vector& fast_shape, + Tensor& output, concurrency::ThreadPool* tp) { + int64_t N = fast_shape[2]; + const T* data = input.Data(); + int64_t stridei = fast_shape[1] * fast_shape[2]; + int64_t strideo = fast_shape[2]; + T* out = output.MutableData(); + std::vector one(fast_shape[1], 1); + concurrency::ThreadPool::TryParallelFor( + tp, fast_shape[0], ParallelReduceFastCost(fast_shape[1], fast_shape[2], sizeof(T)), + [one, data, fast_shape, stridei, strideo, out, N](ptrdiff_t begin, ptrdiff_t last) { + for (ptrdiff_t d = begin; d < last; ++d) { + math::MatMul(1, N, fast_shape[1], one.data(), data + stridei * d, out + strideo * d, nullptr); + } + }); + } }; template @@ -152,6 +252,48 @@ class ReduceAggregatorMean : public ReduceAggregatorSum { return Eigen::Map>(from_data, this->N_).mean(); } inline T get_value() { return this->accumulator_ / static_cast(this->N_); } + + // Fast reduction + // WhichFastReduce() already defined in ReduceAggregatorSum + + static void FastReduceKR(const Tensor& input, const std::vector& fast_shape, + Tensor& output, concurrency::ThreadPool* tp) { + ReduceAggregatorSum::FastReduceKR(input, fast_shape, output, tp); + // TODO: use MLAS or BLAS + T* out = output.MutableData(); + T* end = out + fast_shape[0]; + for (; out != end; ++out) { + *out /= static_cast(fast_shape[1]); + } + } + + static void FastReduceRK(const Tensor& input, const std::vector& fast_shape, + Tensor& output, concurrency::ThreadPool* tp) { + ReduceAggregatorSum::FastReduceRK(input, fast_shape, output, tp); + // TODO: use MLAS or BLAS + T* out = output.MutableData(); + T* end = out + fast_shape[1]; + for (; out != end; ++out) { + *out /= static_cast(fast_shape[0]); + } + } + + static void FastReduceKRK(const Tensor& input, const std::vector& fast_shape, + Tensor& output, concurrency::ThreadPool* tp) { + ReduceAggregatorSum::FastReduceKRK(input, fast_shape, output, tp); + int64_t strideo = fast_shape[2]; + T* out = output.MutableData(); + T* begin; + T* end; + T div = static_cast(fast_shape[1]); + for (int64_t d = 0; d < fast_shape[0]; ++d) { + begin = out + strideo * d; + end = begin + strideo; + for (; begin != end; ++begin) { + *begin /= div; + } + } + } }; template @@ -162,6 +304,66 @@ class ReduceAggregatorMax : public ReduceAggregator { return Eigen::Map>(from_data, this->N_).maxCoeff(); } inline void update(const T& v) { this->accumulator_ = v > this->accumulator_ ? v : this->accumulator_; } + + // Fast reduction + static inline FastReduceKind WhichFastReduce() { + return FastReduceKind::kKR | FastReduceKind::kRK | FastReduceKind::kKRK; + } + + static void FastReduceKR(const Tensor& input, const std::vector& fast_shape, + Tensor& output, concurrency::ThreadPool* tp) { + const T* data = input.Data(); + T* out = output.MutableData(); + int64_t stridei = fast_shape[1]; + concurrency::ThreadPool::TryParallelFor( + tp, fast_shape[0], ParallelReduceFastCost(1, stridei, sizeof(T)), + [data, stridei, out](std::ptrdiff_t first, std::ptrdiff_t last) { + EigenVectorMap(out + first, last - first) = ConstEigenMatrixMap( + data + first * stridei, stridei, last - first) + .colwise() + .maxCoeff(); + }); + } + + static void FastReduceRK(const Tensor& input, const std::vector& fast_shape, + Tensor& output, concurrency::ThreadPool* tp) { + int64_t n_rows = fast_shape[0]; + int64_t N = fast_shape[1]; + const T* data = input.Data(); + T* out = output.MutableData(); + memcpy(out, data, N * sizeof(T)); + + concurrency::ThreadPool::TryParallelFor( + tp, N, ParallelReduceFastCost(1, n_rows, sizeof(T)), + [data, out, N, n_rows](ptrdiff_t begin, ptrdiff_t end) { + const T* p; + for (int64_t row = 1; row < n_rows; ++row) { + p = data + row * N; + for (int64_t j = begin; j < end; ++j) { + out[j] = out[j] > p[j] ? out[j] : p[j]; + } + } + }); + } + + static void FastReduceKRK(const Tensor& input, const std::vector& fast_shape, + Tensor& output, concurrency::ThreadPool* tp) { + const T* data = input.Data(); + T* out = output.MutableData(); + int64_t stridei = fast_shape[1] * fast_shape[2]; + int64_t strideo = fast_shape[2]; + concurrency::ThreadPool::TryParallelFor( + tp, fast_shape[0], ParallelReduceFastCost(fast_shape[1], fast_shape[2], sizeof(T)), + [data, fast_shape, stridei, strideo, out](ptrdiff_t begin, ptrdiff_t end) { + for (ptrdiff_t j = begin; j < end; ++j) { + EigenVectorMap(out + j * strideo, strideo) = + ConstEigenMatrixMap( + data + j * stridei, fast_shape[2], fast_shape[1]) + .rowwise() + .maxCoeff(); + } + }); + } }; template @@ -261,6 +463,66 @@ class ReduceAggregatorMin : public ReduceAggregator { return Eigen::Map>(from_data, this->N_).minCoeff(); } inline void update(const T& v) { this->accumulator_ = v < this->accumulator_ ? v : this->accumulator_; } + + // Fast reduction + static inline FastReduceKind WhichFastReduce() { + return FastReduceKind::kKR | FastReduceKind::kRK | FastReduceKind::kKRK; + } + + static void FastReduceKR(const Tensor& input, const std::vector& fast_shape, + Tensor& output, concurrency::ThreadPool* tp) { + const T* data = input.Data(); + T* out = output.MutableData(); + int64_t stridei = fast_shape[1]; + concurrency::ThreadPool::TryParallelFor( + tp, fast_shape[0], ParallelReduceFastCost(1, stridei, sizeof(T)), + [data, stridei, out](std::ptrdiff_t first, std::ptrdiff_t last) { + EigenVectorMap(out + first, last - first) = ConstEigenMatrixMap( + data + first * stridei, stridei, last - first) + .colwise() + .minCoeff(); + }); + } + + static void FastReduceRK(const Tensor& input, const std::vector& fast_shape, + Tensor& output, concurrency::ThreadPool* tp) { + int64_t n_rows = fast_shape[0]; + int64_t N = fast_shape[1]; + const T* data = input.Data(); + T* out = output.MutableData(); + memcpy(out, data, N * sizeof(T)); + + concurrency::ThreadPool::TryParallelFor( + tp, N, ParallelReduceFastCost(1, n_rows, sizeof(T)), + [data, out, N, n_rows](ptrdiff_t begin, ptrdiff_t end) { + const T* p; + for (int64_t row = 1; row < n_rows; ++row) { + p = data + row * N; + for (int64_t j = begin; j < end; ++j) { + out[j] = out[j] < p[j] ? out[j] : p[j]; + } + } + }); + } + + static void FastReduceKRK(const Tensor& input, const std::vector& fast_shape, + Tensor& output, concurrency::ThreadPool* tp) { + const T* data = input.Data(); + T* out = output.MutableData(); + int64_t stridei = fast_shape[1] * fast_shape[2]; + int64_t strideo = fast_shape[2]; + concurrency::ThreadPool::TryParallelFor( + tp, fast_shape[0], ParallelReduceFastCost(fast_shape[1], fast_shape[2], sizeof(T)), + [data, fast_shape, stridei, strideo, out](ptrdiff_t begin, ptrdiff_t end) { + for (ptrdiff_t j = begin; j < end; ++j) { + EigenVectorMap(out + j * strideo, strideo) = + ConstEigenMatrixMap( + data + j * stridei, fast_shape[2], fast_shape[1]) + .rowwise() + .minCoeff(); + } + }); + } }; template @@ -326,31 +588,33 @@ class ReduceAggregatorLogSumExp : public ReduceAggregator { } inline void update(const T& v) { this->accumulator_ += reduce_exp(v - max_); } inline TVAL get_value() { return reduce_log(this->accumulator_) + max_; } - static inline bool two_loops() { return true; } }; -bool SetupForReduce(const Tensor* input_tensor_ptr, - const std::vector& axes_, - std::vector& axes, - TensorShape& new_input_shape, - std::vector& output_shape, - bool& empty_reduce, - const TensorShape* input_shape_override); - void NoTransposePrepareForReduce(const TensorShape& new_input_shape, const std::vector& reduced_axes, ResultsNoTransposePrepareForReduce& results); -template -void NoTransposeReduce(Tensor* output, const TensorShape& new_input_shape, const Tensor& input, - const std::vector& reduced_axes, concurrency::ThreadPool* tp, - ResultsNoTransposePrepareForReduce& last_results); +template +void NoTransposeReduce1Loop(Tensor* output, const TensorShape& new_input_shape, const Tensor& input, + const std::vector& reduced_axes, concurrency::ThreadPool* tp, + ResultsNoTransposePrepareForReduce& last_results); -template -void CommonReduce(OpKernelContext* ctx, - const std::vector axes_, int64_t keepdims_, - ResultsNoTransposePrepareForReduce& last_results, - bool noop_with_empty_axes = false); +// Specific case for ReduceLogSumExp. +template +void NoTransposeReduce2Loops(Tensor* output, const TensorShape& new_input_shape, const Tensor& input, + const std::vector& reduced_axes, concurrency::ThreadPool* tp, + ResultsNoTransposePrepareForReduce& last_results); + +template +void CommonReduce1Loop(OpKernelContext* ctx, + const std::vector& axes_, int64_t keepdims_, + bool noop_with_empty_axes = false); + +// Specific case for ReduceLogSumExp. +template +void CommonReduce2Loops(OpKernelContext* ctx, + const std::vector& axes_, int64_t keepdims_, + bool noop_with_empty_axes = false); template class ReduceKernelBase { diff --git a/onnxruntime/core/util/math_cpu.cc b/onnxruntime/core/util/math_cpu.cc index fbed516ea7..ed8e90d664 100644 --- a/onnxruntime/core/util/math_cpu.cc +++ b/onnxruntime/core/util/math_cpu.cc @@ -859,6 +859,9 @@ DEFINE_BROADCAST_BINARY_FUNCTION(Div, /) EigenVectorMap(y, N) = ConstEigenMatrixMap(x, D, N).colwise().sum(); \ } SPECIALIZED_ROWWISESUM(float) +SPECIALIZED_ROWWISESUM(int32_t) +SPECIALIZED_ROWWISESUM(int64_t) +SPECIALIZED_ROWWISESUM(double) #undef SPECIALIZED_ROWWISESUM #define SPECIALIZED_SUM(T) \ diff --git a/onnxruntime/test/providers/cpu/math/einsum_test.cc b/onnxruntime/test/providers/cpu/math/einsum_test.cc index 79446b5273..25688375f2 100644 --- a/onnxruntime/test/providers/cpu/math/einsum_test.cc +++ b/onnxruntime/test/providers/cpu/math/einsum_test.cc @@ -96,6 +96,7 @@ TEST(Einsum, ExplicitEinsumAsReduceOp_2D_input_1) { test.AddOutput("y", {2}, {4.f, 6.f}); test.Run(); } + TEST(Einsum, ExplicitEinsumAsBatchedReduceOp_3D_input_0) { OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); test.AddAttribute("equation", "...ji->...j"); diff --git a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc index 8f3596dea2..14785ed2e5 100644 --- a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc +++ b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc @@ -8,6 +8,7 @@ #include "test/common/tensor_op_test_utils.h" #include "test/providers/provider_test_utils.h" #include "test/providers/cpu/reduction/reduction_test_cases.h" +#include "core/providers/cpu/reduction/reduction_ops.h" namespace onnxruntime { namespace test { @@ -717,7 +718,7 @@ TEST(ReductionOpTest, ReduceMax_int32) { #if defined(OPENVINO_CONFIG_GPU_FP32) || defined(OPENVINO_CONFIG_GPU_FP16) || defined(OPENVINO_CONFIG_MYRIAD) test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); // OpenVINO: Disabled temporarily #else - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: axis must be 0 + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: axis must be 0 #endif } @@ -738,7 +739,7 @@ TEST(ReductionOpTest, ReduceMax_int64) { #if defined(OPENVINO_CONFIG_GPU_FP32) || defined(OPENVINO_CONFIG_GPU_FP16) || defined(OPENVINO_CONFIG_MYRIAD) test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); // OpenVINO: Disabled temporarily #else - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: axis must be 0 + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: axis must be 0 #endif } @@ -759,7 +760,7 @@ TEST(ReductionOpTest, ReduceMax_int8) { #if defined(OPENVINO_CONFIG_MYRIAD) test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); // OpenVINO: Disabled temporarily #else - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: axis must be 0 + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: axis must be 0 #endif } @@ -780,7 +781,7 @@ TEST(ReductionOpTest, ReduceMax_uint8) { #if defined(OPENVINO_CONFIG_MYRIAD) test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); // OpenVINO: Disabled temporarily #else - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: axis must be 0 + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: axis must be 0 #endif } @@ -1377,6 +1378,7 @@ TEST(ReductionOpTest, ReduceSum_half_bert) { // Add more UTs for half as needed #endif + TEST(ReductionOpTest, ReduceSum_apex_reduction) { OpTester test("ReduceSum"); test.AddAttribute("keepdims", (int64_t)0); @@ -2266,8 +2268,44 @@ TEST(ReductionOpTest, ArgMin_int32_neg_axis) { test.Run(); } +TEST(ReductionOpTest, OptimizeShapeForFastReduce_ReduceDimWithZero1) { + FastReduceKind fast_kind; + std::vector fast_shape, fast_output_shape, fast_axes; + std::vector expected_fast_shape, expected_fast_output_shape, expected_fast_axes; + + // R - keep_dims=1 - noop=false + fast_kind = OptimizeShapeForFastReduce( + std::vector{3, 0, 2}, std::vector(), + fast_shape, fast_output_shape, fast_axes, true); + expected_fast_shape = std::vector{}; + expected_fast_axes = std::vector{}; + expected_fast_output_shape = std::vector{1, 0, 1}; + ASSERT_EQ(fast_kind, FastReduceKind::kEmpty); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); +} + +TEST(ReductionOpTest, OptimizeShapeForFastReduce_ReduceDimWithZero1b) { + FastReduceKind fast_kind; + std::vector fast_shape, fast_output_shape, fast_axes; + std::vector expected_fast_shape, expected_fast_output_shape, expected_fast_axes; + + // R - keep_dims=1 - noop=false + fast_kind = OptimizeShapeForFastReduce( + std::vector{3, 0, 2}, std::vector{1}, + fast_shape, fast_output_shape, fast_axes, true); + expected_fast_shape = std::vector{}; + expected_fast_axes = std::vector{}; + expected_fast_output_shape = std::vector{3, 0, 2}; + ASSERT_EQ(fast_kind, FastReduceKind::kEmpty); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); +} + // test that PrepareForReduce handles this case. Called by all reduction ops so any op can be used in the test -TEST(ReductionOpTest, ReduceDimWithZero) { +TEST(ReductionOpTest, ReduceDimWithZero1) { auto run = [](OpTester& tester, const std::string& error_msg = "") { auto expect = error_msg.empty() ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure; @@ -2283,6 +2321,34 @@ TEST(ReductionOpTest, ReduceDimWithZero) { test.AddInput("data", {3, 0, 2}, {}); test.AddOutput("reduced", {1, 0, 1}, {}); run(test); +} + +TEST(ReductionOpTest, OptimizeShapeForFastReduce_ReduceDimWithZero2) { + FastReduceKind fast_kind; + std::vector fast_shape, fast_output_shape, fast_axes; + std::vector expected_fast_shape, expected_fast_output_shape, expected_fast_axes; + + // R - keep_dims=0 - noop=false + fast_kind = OptimizeShapeForFastReduce( + std::vector{3, 0, 2}, std::vector(), + fast_shape, fast_output_shape, fast_axes, false); + expected_fast_shape = std::vector{}; + expected_fast_axes = std::vector{}; + expected_fast_output_shape = std::vector{}; + ASSERT_EQ(fast_kind, FastReduceKind::kEmpty); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); +} + +TEST(ReductionOpTest, ReduceDimWithZero2) { + auto run = [](OpTester& tester, const std::string& error_msg = "") { + auto expect = error_msg.empty() ? OpTester::ExpectResult::kExpectSuccess + : OpTester::ExpectResult::kExpectFailure; + + // exclude OpenVINO and TensorRT as this isn't handled by those EPs + tester.Run(expect, error_msg, {kTensorrtExecutionProvider, kOpenVINOExecutionProvider, kNupharExecutionProvider}); + }; // reduction without keeping dims on all axes. can't reduce on an axis with value of 0 OpTester test2("ReduceSum", 10); @@ -2293,6 +2359,34 @@ TEST(ReductionOpTest, ReduceDimWithZero) { run(test2, "Can't reduce on dim with value of 0 if 'keepdims' is false. " "Invalid output shape would be produced. input_shape:{3,0,2}"); +} + +TEST(ReductionOpTest, OptimizeShapeForFastReduce_ReduceDimWithZero3) { + FastReduceKind fast_kind; + std::vector fast_shape, fast_output_shape, fast_axes; + std::vector expected_fast_shape, expected_fast_output_shape, expected_fast_axes; + + // R - keep_dims=0 - noop=false + fast_kind = OptimizeShapeForFastReduce( + std::vector{3, 0, 2}, std::vector{2}, + fast_shape, fast_output_shape, fast_axes, false); + expected_fast_shape = std::vector{0, 2}; + expected_fast_axes = std::vector{1}; + expected_fast_output_shape = std::vector{3, 0}; + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + ASSERT_EQ(fast_kind, FastReduceKind::kKR); +} + +TEST(ReductionOpTest, ReduceDimWithZero3) { + auto run = [](OpTester& tester, const std::string& error_msg = "") { + auto expect = error_msg.empty() ? OpTester::ExpectResult::kExpectSuccess + : OpTester::ExpectResult::kExpectFailure; + + // exclude OpenVINO and TensorRT as this isn't handled by those EPs + tester.Run(expect, error_msg, {kTensorrtExecutionProvider, kOpenVINOExecutionProvider, kNupharExecutionProvider}); + }; // reduction is possible without keeping dims if we only reduce on non-zero dims OpTester test3("ReduceSum", 10); @@ -2430,5 +2524,1395 @@ TEST(ReductionOpTest, ReduceInfLogSumExp_double) { test.Run(); } +// Specific cases for Reduce. + +TEST(ReductionOpTest, OptimizeShapeForFastReduce_R_K) { + FastReduceKind fast_kind; + std::vector fast_shape, fast_output_shape, fast_axes; + std::vector expected_fast_shape, expected_fast_output_shape, expected_fast_axes; + + // R - keep_dims=1 + fast_kind = OptimizeShapeForFastReduce( + std::vector{10}, std::vector{0}, + fast_shape, fast_output_shape, fast_axes, true); + expected_fast_shape = std::vector{10}; + expected_fast_output_shape = std::vector{1}; + expected_fast_axes = std::vector{0}; + ASSERT_EQ(fast_kind, FastReduceKind::kR); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{10, 11}, std::vector{0, 1}, + fast_shape, fast_output_shape, fast_axes, true); + expected_fast_shape = std::vector{110}; + expected_fast_output_shape = std::vector{1, 1}; + expected_fast_axes = std::vector{0}; + ASSERT_EQ(fast_kind, FastReduceKind::kR); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + // R - keep_dims=0 + fast_kind = OptimizeShapeForFastReduce( + std::vector{10}, std::vector{0}, + fast_shape, fast_output_shape, fast_axes, false); + expected_fast_shape = std::vector{10}; + expected_fast_output_shape = std::vector(); + expected_fast_axes = std::vector{0}; + ASSERT_EQ(fast_kind, FastReduceKind::kR); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{10, 11}, std::vector{0, 1}, + fast_shape, fast_output_shape, fast_axes, false); + expected_fast_shape = std::vector{110}; + expected_fast_output_shape = std::vector(); + expected_fast_axes = std::vector{0}; + ASSERT_EQ(fast_kind, FastReduceKind::kR); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); +} + +TEST(ReductionOpTest, OptimizeShapeForFastReduce_R_empty) { + FastReduceKind fast_kind; + std::vector fast_shape, fast_output_shape, fast_axes; + std::vector expected_fast_shape, expected_fast_output_shape, expected_fast_axes; + + // R - keep_dims=1 - noop=false + fast_kind = OptimizeShapeForFastReduce( + std::vector{10}, std::vector(), + fast_shape, fast_output_shape, fast_axes, true); + expected_fast_axes = std::vector{0}; + ASSERT_EQ(fast_kind, FastReduceKind::kR); + ASSERT_EQ(fast_shape, std::vector{10}); + ASSERT_EQ(fast_output_shape, std::vector{1}); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{10, 11}, std::vector(), + fast_shape, fast_output_shape, fast_axes, true); + expected_fast_shape = std::vector{110}; + expected_fast_output_shape = std::vector{1, 1}; + expected_fast_axes = std::vector{0}; + ASSERT_EQ(fast_kind, FastReduceKind::kR); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + // R - keep_dims=0 - noop=false + fast_kind = OptimizeShapeForFastReduce( + std::vector{10}, std::vector{}, + fast_shape, fast_output_shape, fast_axes, false); + expected_fast_shape = std::vector{10}; + expected_fast_output_shape = std::vector{}; + expected_fast_axes = std::vector{0}; + ASSERT_EQ(fast_kind, FastReduceKind::kR); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{10, 11}, std::vector{}, + fast_shape, fast_output_shape, fast_axes, false); + expected_fast_shape = std::vector{110}; + expected_fast_output_shape = std::vector{}; + expected_fast_axes = std::vector{0}; + ASSERT_EQ(fast_kind, FastReduceKind::kR); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); +} + +TEST(ReductionOpTest, OptimizeShapeForFastReduce_K_empty) { + FastReduceKind fast_kind; + std::vector fast_shape, fast_output_shape, fast_axes; + std::vector expected_fast_shape, expected_fast_output_shape, expected_fast_axes; + + // R - keep_dims=1 - noop=true + fast_kind = OptimizeShapeForFastReduce( + std::vector{10}, std::vector(), + fast_shape, fast_output_shape, fast_axes, true, true); + expected_fast_shape = std::vector{10}; + expected_fast_output_shape = std::vector{10}; + expected_fast_axes = std::vector{}; + ASSERT_EQ(fast_kind, FastReduceKind::kK); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{10, 11}, std::vector(), + fast_shape, fast_output_shape, fast_axes, true, true); + expected_fast_shape = std::vector{110}; + expected_fast_output_shape = std::vector{10, 11}; + expected_fast_axes = std::vector{}; + ASSERT_EQ(fast_kind, FastReduceKind::kK); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + // R - keep_dims=0 - noop=true + fast_kind = OptimizeShapeForFastReduce( + std::vector{10}, std::vector{}, + fast_shape, fast_output_shape, fast_axes, false, true); + expected_fast_shape = std::vector{10}; + expected_fast_output_shape = std::vector{10}; + expected_fast_axes = std::vector{}; + ASSERT_EQ(fast_kind, FastReduceKind::kK); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{10, 11}, std::vector{}, + fast_shape, fast_output_shape, fast_axes, false, true); + expected_fast_shape = std::vector{110}; + expected_fast_output_shape = std::vector{10, 11}; + expected_fast_axes = std::vector{}; + ASSERT_EQ(fast_kind, FastReduceKind::kK); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); +} + +TEST(ReductionOpTest, OptimizeShapeForFastReduce_KR) { + FastReduceKind fast_kind; + std::vector fast_shape, fast_output_shape, fast_axes; + std::vector expected_fast_shape, expected_fast_output_shape, expected_fast_axes; + + // KR - keep_dims=1 + fast_kind = OptimizeShapeForFastReduce( + std::vector{10, 11}, std::vector{1}, + fast_shape, fast_output_shape, fast_axes, true); + expected_fast_shape = std::vector{10, 11}; + expected_fast_output_shape = std::vector{10, 1}; + expected_fast_axes = std::vector{1}; + ASSERT_EQ(fast_kind, FastReduceKind::kKR); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{9, 10, 11}, std::vector{1, 2}, + fast_shape, fast_output_shape, fast_axes, true); + expected_fast_shape = std::vector{9, 110}; + expected_fast_output_shape = std::vector{9, 1, 1}; + ASSERT_EQ(fast_kind, FastReduceKind::kKR); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{9, 10, 11}, std::vector{2}, + fast_shape, fast_output_shape, fast_axes, true); + expected_fast_shape = std::vector{90, 11}; + expected_fast_output_shape = std::vector{9, 10, 1}; + ASSERT_EQ(fast_kind, FastReduceKind::kKR); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + // KR - keep_dims=0 + fast_kind = OptimizeShapeForFastReduce( + std::vector{10, 11}, std::vector{1}, + fast_shape, fast_output_shape, fast_axes, false); + expected_fast_shape = std::vector{10, 11}; + expected_fast_output_shape = std::vector{10}; + ASSERT_EQ(fast_kind, FastReduceKind::kKR); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{9, 10, 11}, std::vector{1, 2}, + fast_shape, fast_output_shape, fast_axes, false); + expected_fast_shape = std::vector{9, 110}; + expected_fast_output_shape = std::vector{9}; + ASSERT_EQ(fast_kind, FastReduceKind::kKR); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{9, 10, 11}, std::vector{2}, + fast_shape, fast_output_shape, fast_axes, false); + expected_fast_shape = std::vector{90, 11}; + expected_fast_output_shape = std::vector{9, 10}; + ASSERT_EQ(fast_kind, FastReduceKind::kKR); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); +} + +TEST(ReductionOpTest, OptimizeShapeForFastReduce_KR_neg) { + FastReduceKind fast_kind; + std::vector fast_shape, fast_output_shape, fast_axes; + std::vector expected_fast_shape, expected_fast_output_shape, expected_fast_axes; + + // KR - keep_dims=1 + fast_kind = OptimizeShapeForFastReduce( + std::vector{10, 11}, std::vector{-1}, + fast_shape, fast_output_shape, fast_axes, true); + expected_fast_shape = std::vector{10, 11}; + expected_fast_output_shape = std::vector{10, 1}; + expected_fast_axes = std::vector{1}; + ASSERT_EQ(fast_kind, FastReduceKind::kKR); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); +} + +TEST(ReductionOpTest, OptimizeShapeForFastReduce_RK) { + FastReduceKind fast_kind; + std::vector fast_shape, fast_output_shape, fast_axes; + std::vector expected_fast_shape, expected_fast_output_shape, expected_fast_axes; + + // RK - keep_dims=1 + fast_kind = OptimizeShapeForFastReduce( + std::vector{10, 11}, std::vector{0}, + fast_shape, fast_output_shape, fast_axes, true); + expected_fast_shape = std::vector{10, 11}; + expected_fast_output_shape = std::vector{1, 11}; + expected_fast_axes = std::vector{0}; + ASSERT_EQ(fast_kind, FastReduceKind::kRK); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{9, 10, 11}, std::vector{0, 1}, + fast_shape, fast_output_shape, fast_axes, true); + expected_fast_shape = std::vector{90, 11}; + expected_fast_output_shape = std::vector{1, 1, 11}; + ASSERT_EQ(fast_kind, FastReduceKind::kRK); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{9, 10, 11}, std::vector{0}, + fast_shape, fast_output_shape, fast_axes, true); + expected_fast_shape = std::vector{9, 110}; + expected_fast_output_shape = std::vector{1, 10, 11}; + ASSERT_EQ(fast_kind, FastReduceKind::kRK); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + // RK - keep_dims=0 + fast_kind = OptimizeShapeForFastReduce( + std::vector{10, 11}, std::vector{0}, + fast_shape, fast_output_shape, fast_axes, false); + expected_fast_shape = std::vector{10, 11}; + expected_fast_output_shape = std::vector{11}; + ASSERT_EQ(fast_kind, FastReduceKind::kRK); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{9, 10, 11}, std::vector{0, 1}, + fast_shape, fast_output_shape, fast_axes, false); + expected_fast_shape = std::vector{90, 11}; + expected_fast_output_shape = std::vector{11}; + ASSERT_EQ(fast_kind, FastReduceKind::kRK); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{9, 10, 11}, std::vector{0}, + fast_shape, fast_output_shape, fast_axes, false); + expected_fast_shape = std::vector{9, 110}; + expected_fast_output_shape = std::vector{10, 11}; + ASSERT_EQ(fast_kind, FastReduceKind::kRK); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); +} + +TEST(ReductionOpTest, OptimizeShapeForFastReduce_KRK) { + FastReduceKind fast_kind; + std::vector fast_shape, fast_output_shape, fast_axes; + std::vector expected_fast_shape, expected_fast_output_shape, expected_fast_axes; + + // KRK - keep_dims=1 + fast_kind = OptimizeShapeForFastReduce( + std::vector{9, 10, 11}, std::vector{1}, + fast_shape, fast_output_shape, fast_axes, true); + expected_fast_shape = std::vector{9, 10, 11}; + expected_fast_output_shape = std::vector{9, 1, 11}; + expected_fast_axes = std::vector{1}; + ASSERT_EQ(fast_kind, FastReduceKind::kKRK); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{7, 9, 10, 11}, std::vector{1, 2}, + fast_shape, fast_output_shape, fast_axes, true); + expected_fast_shape = std::vector{7, 90, 11}; + expected_fast_output_shape = std::vector{7, 1, 1, 11}; + ASSERT_EQ(fast_kind, FastReduceKind::kKRK); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{7, 9, 10, 11}, std::vector{1}, + fast_shape, fast_output_shape, fast_axes, true); + expected_fast_shape = std::vector{7, 9, 110}; + expected_fast_output_shape = std::vector{7, 1, 10, 11}; + ASSERT_EQ(fast_kind, FastReduceKind::kKRK); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{7, 9, 10, 11}, std::vector{2}, + fast_shape, fast_output_shape, fast_axes, true); + expected_fast_shape = std::vector{63, 10, 11}; + expected_fast_output_shape = std::vector{7, 9, 1, 11}; + ASSERT_EQ(fast_kind, FastReduceKind::kKRK); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + // KRK - keep_dims=0 + fast_kind = OptimizeShapeForFastReduce( + std::vector{9, 10, 11}, std::vector{1}, + fast_shape, fast_output_shape, fast_axes, false); + expected_fast_shape = std::vector{9, 10, 11}; + expected_fast_output_shape = std::vector{9, 11}; + ASSERT_EQ(fast_kind, FastReduceKind::kKRK); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{7, 9, 10, 11}, std::vector{1, 2}, + fast_shape, fast_output_shape, fast_axes, false); + expected_fast_shape = std::vector{7, 90, 11}; + expected_fast_output_shape = std::vector{7, 11}; + ASSERT_EQ(fast_kind, FastReduceKind::kKRK); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{7, 9, 10, 11}, std::vector{1}, + fast_shape, fast_output_shape, fast_axes, false); + expected_fast_shape = std::vector{7, 9, 110}; + expected_fast_output_shape = std::vector{7, 10, 11}; + ASSERT_EQ(fast_kind, FastReduceKind::kKRK); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{7, 9, 10, 11}, std::vector{2}, + fast_shape, fast_output_shape, fast_axes, false); + expected_fast_shape = std::vector{63, 10, 11}; + expected_fast_output_shape = std::vector{7, 9, 11}; + ASSERT_EQ(fast_kind, FastReduceKind::kKRK); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); +} + +TEST(ReductionOpTest, OptimizeShapeForFastReduce_NONE) { + FastReduceKind fast_kind; + std::vector fast_shape, fast_output_shape, fast_axes; + std::vector expected_fast_shape, expected_fast_output_shape, expected_fast_axes; + + // RKRK + fast_kind = OptimizeShapeForFastReduce( + std::vector{7, 9, 10, 11}, std::vector{0, 2}, + fast_shape, fast_output_shape, fast_axes, false); + expected_fast_shape = std::vector{7, 9, 10, 11}; + expected_fast_output_shape = std::vector{9, 11}; + expected_fast_axes = std::vector{0, 2}; + ASSERT_EQ(fast_kind, FastReduceKind::kNone); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{7, 9, 10, 11}, std::vector{1, 3}, + fast_shape, fast_output_shape, fast_axes, true); + expected_fast_shape = std::vector{7, 9, 10, 11}; + expected_fast_output_shape = std::vector{7, 1, 10, 1}; + expected_fast_axes = std::vector{1, 3}; + ASSERT_EQ(fast_kind, FastReduceKind::kNone); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + // RRKKRRKK + fast_kind = OptimizeShapeForFastReduce( + std::vector{7, 9, 10, 11, 2, 3, 4, 6}, std::vector{0, 1, 4, 5}, + fast_shape, fast_output_shape, fast_axes, false); + expected_fast_shape = std::vector{63, 110, 6, 24}; + expected_fast_output_shape = std::vector{10, 11, 4, 6}; + expected_fast_axes = std::vector{0, 2}; + ASSERT_EQ(fast_kind, FastReduceKind::kNone); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); + + fast_kind = OptimizeShapeForFastReduce( + std::vector{7, 9, 10, 11, 2, 3, 4, 6}, std::vector{0, 1, 4, 5}, + fast_shape, fast_output_shape, fast_axes, true); + expected_fast_shape = std::vector{63, 110, 6, 24}; + expected_fast_output_shape = std::vector{1, 1, 10, 11, 1, 1, 4, 6}; + expected_fast_axes = std::vector{0, 2}; + ASSERT_EQ(fast_kind, FastReduceKind::kNone); + ASSERT_EQ(fast_shape, expected_fast_shape); + ASSERT_EQ(fast_output_shape, expected_fast_output_shape); + ASSERT_EQ(fast_axes, expected_fast_axes); +} + +TEST(ReductionOpTest, EigenMax) { + std::vector mat{1, 2, 3, 4}; + + auto res1 = ConstEigenMatrixMap(mat.data(), 2, 2).rowwise().maxCoeff(); + std::vector expected{3, 4}; + std::vector out1(res1.begin(), res1.end()); + ASSERT_EQ(out1, expected); + + auto res2 = ConstEigenMatrixMap(mat.data(), 2, 2).colwise().maxCoeff(); + expected = std::vector{2, 4}; + std::vector out2(res2.begin(), res2.end()); + ASSERT_EQ(out2, expected); + + mat = std::vector{1, 2, 3, 4, 5, 6}; + + auto res3 = ConstEigenMatrixMap(mat.data(), 2, 3).rowwise().maxCoeff(); + expected = std::vector{5, 6}; + std::vector out3(res3.begin(), res3.end()); + ASSERT_EQ(out3, expected); + + auto res4 = ConstEigenMatrixMap(mat.data(), 2, 3).colwise().maxCoeff(); + expected = std::vector{2, 4, 6}; + std::vector out4(res4.begin(), res4.end()); + ASSERT_EQ(out4, expected); + + auto res5 = ConstEigenMatrixMap(mat.data(), 3, 2).rowwise().maxCoeff(); + expected = std::vector{4, 5, 6}; + std::vector out5(res5.begin(), res5.end()); + ASSERT_EQ(out5, expected); + + auto res6 = ConstEigenMatrixMap(mat.data(), 2, 3).colwise().maxCoeff(); + expected = std::vector{2, 4, 6}; + std::vector out6(res6.begin(), res6.end()); + ASSERT_EQ(out6, expected); +} + +TEST(ReductionOpTest, EigenSum) { + std::vector mat{1, 10, 100, 1000}; + + auto res1 = ConstEigenMatrixMap(mat.data(), 2, 2).rowwise().sum(); + std::vector expected{101, 1010}; + std::vector out1(res1.begin(), res1.end()); + ASSERT_EQ(out1, expected); + + auto res2 = ConstEigenMatrixMap(mat.data(), 2, 2).colwise().sum(); + expected = std::vector{11, 1100}; + std::vector out2(res2.begin(), res2.end()); + ASSERT_EQ(out2, expected); + + mat = std::vector{1, 10, 100, 1000, 10000, 100000}; + + auto res3 = ConstEigenMatrixMap(mat.data(), 2, 3).rowwise().sum(); + expected = std::vector{10101, 101010}; + std::vector out3(res3.begin(), res3.end()); + ASSERT_EQ(out3, expected); + + auto res4 = ConstEigenMatrixMap(mat.data(), 2, 3).colwise().sum(); + expected = std::vector{11, 1100, 110000}; + std::vector out4(res4.begin(), res4.end()); + ASSERT_EQ(out4, expected); + + auto res5 = ConstEigenMatrixMap(mat.data(), 3, 2).rowwise().sum(); + expected = std::vector{1001, 10010, 100100}; + std::vector out5(res5.begin(), res5.end()); + ASSERT_EQ(out5, expected); + + auto res6 = ConstEigenMatrixMap(mat.data(), 2, 3).colwise().sum(); + expected = std::vector{11, 1100, 110000}; + std::vector out6(res6.begin(), res6.end()); + ASSERT_EQ(out6, expected); +} + +TEST(ReductionOpTest, ReduceMax_KR_parallel) { + OpTester test("ReduceMax"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {4, 3}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {4}, {3.f, 6.f, 9.f, 12.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMax_KR) { + OpTester test("ReduceMax"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3, 4}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {3}, {4.f, 8.f, 12.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMax_KR_keepdims) { + OpTester test("ReduceMax"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 4}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {3, 1}, {4.f, 8.f, 12.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMax_RK) { + OpTester test("ReduceMax"); + test.AddAttribute("axes", std::vector{0}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3, 4}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {4}, {9.f, 10.f, 11.f, 12.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMax_RK_keepdims) { + OpTester test("ReduceMax"); + test.AddAttribute("axes", std::vector{0}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 4}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {1, 4}, {9.f, 10.f, 11.f, 12.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMax_RK_parallel) { + OpTester test("ReduceMax"); + test.AddAttribute("axes", std::vector{0}); + test.AddAttribute("keepdims", (int64_t)0); + std::vector in_data(128); + for (size_t i = 0; i < in_data.size(); ++i) + in_data[i] = (float)i; + test.AddInput("data", {4, 32}, in_data); + std::vector expected(32); + for (size_t i = 0; i < expected.size(); ++i) { + expected[i] = 0; + for (size_t j = 0; j < 4; ++j) { + if (in_data[i + j * expected.size()] > expected[i]) + expected[i] = in_data[i + j * expected.size()]; + } + } + test.AddOutput("reduced", {32}, expected); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMax_KRK) { + OpTester test("ReduceMax"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {3, 2}, {3.f, 4.f, 7.f, 8.f, 11.f, 12.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMax_KRK_keepdims) { + OpTester test("ReduceMax"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {3, 1, 2}, {3.f, 4.f, 7.f, 8.f, 11.f, 12.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMax_RKRK) { + OpTester test("ReduceMax"); + test.AddAttribute("axes", std::vector{0, 2}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3, 2, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f, + + 13.0f, 14.0f, + 15.0f, 16.0f, + + 17.0f, 18.0f, + 19.0f, 20.0f, + + 21.0f, 22.0f, + 23.0f, 24.0f}); + test.AddOutput("reduced", {2, 2}, {19.f, 20.f, 23.f, 24.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMax_RKRK_keepdims) { + OpTester test("ReduceMax"); + test.AddAttribute("axes", std::vector{0, 2}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 2, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f, + + 13.0f, 14.0f, + 15.0f, 16.0f, + + 17.0f, 18.0f, + 19.0f, 20.0f, + + 21.0f, 22.0f, + 23.0f, 24.0f}); + test.AddOutput("reduced", {1, 2, 1, 2}, {19.f, 20.f, 23.f, 24.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMean_KR) { + OpTester test("ReduceMean"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3, 4}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {3}, {2.5f, 6.5f, 10.5f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMean_KR_keepdims) { + OpTester test("ReduceMean"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 4}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {3, 1}, {2.5f, 6.5f, 10.5f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMean_RK) { + OpTester test("ReduceMean"); + test.AddAttribute("axes", std::vector{0}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3, 4}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {4}, {5.f, 6.f, 7.f, 8.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMean_RK_keepdims) { + OpTester test("ReduceMean"); + test.AddAttribute("axes", std::vector{0}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 4}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {1, 4}, {5.f, 6.f, 7.f, 8.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMean_KRK) { + OpTester test("ReduceMean"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {3, 2}, {2.f, 3.f, 6.f, 7.f, 10.f, 11.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMean_KRK_keepdims) { + OpTester test("ReduceMean"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {3, 1, 2}, {2.f, 3.f, 6.f, 7.f, 10.f, 11.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMean_RKRK) { + OpTester test("ReduceMean"); + test.AddAttribute("axes", std::vector{0, 2}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3, 2, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f, + + 13.0f, 14.0f, + 15.0f, 16.0f, + + 17.0f, 18.0f, + 19.0f, 20.0f, + + 21.0f, 22.0f, + 23.0f, 24.0f}); + test.AddOutput("reduced", {2, 2}, {10.f, 11.f, 14.f, 15.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMean_RKRK_keepdims) { + OpTester test("ReduceMean"); + test.AddAttribute("axes", std::vector{0, 2}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 2, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f, + + 13.0f, 14.0f, + 15.0f, 16.0f, + + 17.0f, 18.0f, + 19.0f, 20.0f, + + 21.0f, 22.0f, + 23.0f, 24.0f}); + test.AddOutput("reduced", {1, 2, 1, 2}, {10.f, 11.f, 14.f, 15.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMin_KR) { + OpTester test("ReduceMin"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3, 4}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {3}, {1.f, 5.f, 9.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMin_KR_parallel) { + OpTester test("ReduceMin"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {4, 3}, + {11.0f, 12.0f, + 13.0f, 14.0f, + + 15.0f, 16.0f, + 17.0f, 18.0f, + + 19.0f, 20.0f, + 21.0f, 22.0f}); + test.AddOutput("reduced", {4}, {11.f, 14.f, 17.f, 20.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMin_KR_keepdims) { + OpTester test("ReduceMin"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 4}, + {11.0f, 12.0f, + 13.0f, 14.0f, + + 15.0f, 16.0f, + 17.0f, 18.0f, + + 19.0f, 20.0f, + 21.0f, 22.0f}); + test.AddOutput("reduced", {3, 1}, {11.f, 15.f, 19.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMin_RK) { + OpTester test("ReduceMin"); + test.AddAttribute("axes", std::vector{0}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3, 4}, + {11.0f, 12.0f, + 13.0f, 14.0f, + + 15.0f, 16.0f, + 17.0f, 18.0f, + + 19.0f, 20.0f, + 21.0f, 22.0f}); + test.AddOutput("reduced", {4}, {11.f, 12.f, 13.f, 14.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMin_RK_parallel) { + OpTester test("ReduceMin"); + test.AddAttribute("axes", std::vector{0}); + test.AddAttribute("keepdims", (int64_t)0); + std::vector in_data(128); + for (size_t i = 0; i < in_data.size(); ++i) + in_data[i] = (float)i; + test.AddInput("data", {4, 32}, in_data); + std::vector expected(32); + for (size_t i = 0; i < expected.size(); ++i) { + expected[i] = 1000000000; + for (size_t j = 0; j < 4; ++j) { + if (in_data[i + j * expected.size()] < expected[i]) + expected[i] = in_data[i + j * expected.size()]; + } + } + test.AddOutput("reduced", {32}, expected); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMin_RK_keepdims) { + OpTester test("ReduceMin"); + test.AddAttribute("axes", std::vector{0}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 4}, + {11.0f, 12.0f, + 13.0f, 14.0f, + + 15.0f, 16.0f, + 17.0f, 18.0f, + + 19.0f, 20.0f, + 21.0f, 22.0f}); + test.AddOutput("reduced", {1, 4}, {11.f, 12.f, 13.f, 14.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMin_KRK) { + OpTester test("ReduceMin"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3, 2, 2}, + {11.0f, 12.0f, + 13.0f, 14.0f, + + 15.0f, 16.0f, + 17.0f, 18.0f, + + 19.0f, 20.0f, + 21.0f, 22.0f}); + test.AddOutput("reduced", {3, 2}, {11.f, 12.f, 15.f, 16.f, 19.f, 20.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMin_KRK_keepdims) { + OpTester test("ReduceMin"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 2, 2}, + {11.0f, 12.0f, + 13.0f, 14.0f, + + 15.0f, 16.0f, + 17.0f, 18.0f, + + 19.0f, 20.0f, + 21.0f, 22.0f}); + test.AddOutput("reduced", {3, 1, 2}, {11.f, 12.f, 15.f, 16.f, 19.f, 20.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMin_RKRK) { + OpTester test("ReduceMin"); + test.AddAttribute("axes", std::vector{0, 2}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3, 2, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f, + + 13.0f, 14.0f, + 15.0f, 16.0f, + + 17.0f, 18.0f, + 19.0f, 20.0f, + + 21.0f, 22.0f, + 23.0f, 24.0f}); + test.AddOutput("reduced", {2, 2}, {1.f, 2.f, 5.f, 6.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceMin_RKRK_keepdims) { + OpTester test("ReduceMin"); + test.AddAttribute("axes", std::vector{0, 2}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 2, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f, + + 13.0f, 14.0f, + 15.0f, 16.0f, + + 17.0f, 18.0f, + 19.0f, 20.0f, + + 21.0f, 22.0f, + 23.0f, 24.0f}); + test.AddOutput("reduced", {1, 2, 1, 2}, {1.f, 2.f, 5.f, 6.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceSum_KR) { + OpTester test("ReduceSum"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3, 4}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {3}, {10.0f, 26.0f, 42.0f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceSum_KR_parallel) { + OpTester test("ReduceSum"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {4, 3}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {4}, {6.0f, 15.0f, 24.0f, 33.0f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceSum_KR_keepdims) { + OpTester test("ReduceSum"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 4}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {3, 1}, {10.0f, 26.0f, 42.0f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceSum_KR2) { + OpTester test("ReduceSum"); + test.AddAttribute("axes", std::vector{1, 2}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {3}, {10.0f, 26.0f, 42.0f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceSum_KR2_keepdims) { + OpTester test("ReduceSum"); + test.AddAttribute("axes", std::vector{1, 2}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {3, 1, 1}, {10.0f, 26.0f, 42.0f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceSum_RK) { + OpTester test("ReduceSum"); + test.AddAttribute("axes", std::vector{0}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3, 4}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {4}, {15.f, 18.f, 21.f, 24.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceSum_RK_parallel) { + OpTester test("ReduceSum"); + test.AddAttribute("axes", std::vector{0}); + test.AddAttribute("keepdims", (int64_t)0); + std::vector in_data(128); + for (size_t i = 0; i < in_data.size(); ++i) + in_data[i] = (float)i; + test.AddInput("data", {4, 32}, in_data); + std::vector expected(32); + for (size_t i = 0; i < expected.size(); ++i) { + expected[i] = 0; + for (size_t j = 0; j < 4; ++j) { + expected[i] += in_data[i + j * expected.size()]; + } + } + test.AddOutput("reduced", {32}, expected); + test.Run(); +} + +TEST(ReductionOpTest, ReduceSum_RK_keepdims) { + OpTester test("ReduceSum"); + test.AddAttribute("axes", std::vector{0}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 4}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {1, 4}, {15.f, 18.f, 21.f, 24.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceSum_RK2) { + OpTester test("ReduceSum"); + test.AddAttribute("axes", std::vector{0, 1}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {2}, {36.f, 42.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceSum_RK2_keepdims) { + OpTester test("ReduceSum"); + test.AddAttribute("axes", std::vector{0, 1}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {1, 1, 2}, {36.f, 42.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceSum_KRK) { + OpTester test("ReduceSum"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {3, 2}, {4.f, 6.f, 12.f, 14.f, 20.f, 22.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceSum_KRK_parallel) { + OpTester test("ReduceSum"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {4, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f, + + 13.0f, 14.0f, + 15.0f, 16.0f}); + test.AddOutput("reduced", {4, 2}, {4.f, 6.f, 12.f, 14.f, 20.f, 22.f, 28.f, 30.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceSum_KRK_keepdims) { + OpTester test("ReduceSum"); + test.AddAttribute("axes", std::vector{1}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f}); + test.AddOutput("reduced", {3, 1, 2}, {4.f, 6.f, 12.f, 14.f, 20.f, 22.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceSum_KRK2) { + OpTester test("ReduceSum"); + test.AddAttribute("axes", std::vector{1, 2}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3, 2, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f, + + 13.0f, 14.0f, + 15.0f, 16.0f, + + 17.0f, 18.0f, + 19.0f, 20.0f, + + 21.0f, 22.0f, + 23.0f, 24.0f}); + test.AddOutput("reduced", {3, 2}, {16.f, 20.f, 48.f, 52.f, 80.f, 84.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceSum_KRK2_keepdims) { + OpTester test("ReduceSum"); + test.AddAttribute("axes", std::vector{1, 2}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 2, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f, + + 13.0f, 14.0f, + 15.0f, 16.0f, + + 17.0f, 18.0f, + 19.0f, 20.0f, + + 21.0f, 22.0f, + 23.0f, 24.0f}); + test.AddOutput("reduced", {3, 1, 1, 2}, {16.f, 20.f, 48.f, 52.f, 80.f, 84.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceSum_RKRK) { + OpTester test("ReduceSum"); + test.AddAttribute("axes", std::vector{0, 2}); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3, 2, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f, + + 13.0f, 14.0f, + 15.0f, 16.0f, + + 17.0f, 18.0f, + 19.0f, 20.0f, + + 21.0f, 22.0f, + 23.0f, 24.0f}); + test.AddOutput("reduced", {2, 2}, {60.f, 66.f, 84.f, 90.f}); + test.Run(); +} + +TEST(ReductionOpTest, ReduceSum_RKRK_keepdims) { + OpTester test("ReduceSum"); + test.AddAttribute("axes", std::vector{0, 2}); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 2, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f, + + 5.0f, 6.0f, + 7.0f, 8.0f, + + 9.0f, 10.0f, + 11.0f, 12.0f, + + 13.0f, 14.0f, + 15.0f, 16.0f, + + 17.0f, 18.0f, + 19.0f, 20.0f, + + 21.0f, 22.0f, + 23.0f, 24.0f}); + test.AddOutput("reduced", {1, 2, 1, 2}, {60.f, 66.f, 84.f, 90.f}); + test.Run(); +} + } // namespace test } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cpu/reduction/reduction_ops.cc b/orttraining/orttraining/training_ops/cpu/reduction/reduction_ops.cc index 400f7688c1..d1b6cea1d9 100644 --- a/orttraining/orttraining/training_ops/cpu/reduction/reduction_ops.cc +++ b/orttraining/orttraining/training_ops/cpu/reduction/reduction_ops.cc @@ -29,8 +29,7 @@ REGISTER_REDUCESUMTRAINING_KERNEL_TYPED(int64_t) template Status ReduceSumTraining::Compute(OpKernelContext* ctx) const { - ResultsNoTransposePrepareForReduce last_results; - CommonReduce>(ctx, axes_, keepdims_, last_results, noop_with_empty_axes_); + CommonReduce1Loop>(ctx, axes_, keepdims_, noop_with_empty_axes_); return Status::OK(); }