saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-25 19:48:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Speed up Reduce operators for consecutive reduced axes (#7206)

Browse source 
* Improves Reduction for three specific configurations
* Support ReduceMean
* add ReduceMax, ReduceMin
* refactoring
This commit is contained in:
Xavier Dupré 2021-05-05 09:14:00 +02:00 committed by GitHub
parent 053bada30f
commit ade6ed51eb
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 2292 additions and 264 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

699
onnxruntime/core/providers/cpu/reduction/reduction_ops.cc
View file

@ -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<int64_t>& axes_,
std::vector<int64_t>& axes,
TensorShape& new_input_shape,
std::vector<int64_t>& 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<int64_t>(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<int64_t>(ndim - axes.size()) &&
axes.back() == static_cast<int64_t>(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<FastReduceKind>(static_cast<uint8_t>(a) | static_cast<uint8_t>(b));
}
bool operator==(FastReduceKind a, FastReduceKind b) {
return static_cast<uint8_t>(a) == static_cast<uint8_t>(b);
}
bool operator!=(FastReduceKind a, FastReduceKind b) {
return static_cast<uint8_t>(a) != static_cast<uint8_t>(b);
}
bool IsFastReduceKindAvailable(FastReduceKind scenario, FastReduceKind available) {
return (static_cast<uint8_t>(scenario) & static_cast<uint8_t>(available)) > 0;
}
bool ResultsNoTransposePrepareForReduce::equal(const std::vector<int64_t>& local_input_shape,
const std::vector<int64_t>& 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<int64_t>::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<int64_t>::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<int64_t>& 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<int64_t>& 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<int64_t>& 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<int64_t>&, Tensor&, concurrency::ThreadPool*) {
ValidateMustBeOverloaded();
}
void ReduceAggregatorBase::FastReduceRK(const Tensor&, const std::vector<int64_t>&, Tensor&, concurrency::ThreadPool*) {
ValidateMustBeOverloaded();
}
void ReduceAggregatorBase::FastReduceKRK(const Tensor&, const std::vector<int64_t>&, Tensor&, concurrency::ThreadPool*) {
ValidateMustBeOverloaded();
}
TensorOpCost ParallelReduceFastCost(int64_t n_row, int64_t n_col, int64_t element_size) {
return TensorOpCost{static_cast<double>(n_col * n_row * element_size),
static_cast<double>(n_row * element_size),
static_cast<double>(n_col * n_row * element_size * 2)};
}
void NoTransposePrepareForReduce(const TensorShape& new_input_shape,
const std::vector<int64_t>& reduced_axes,
ResultsNoTransposePrepareForReduce& results) {
@ -355,17 +380,21 @@ void NoTransposePrepareForReduce(const TensorShape& new_input_shape,
}
}
template <typename T, typename AGG>
void NoTransposeReduce(Tensor* output, const TensorShape& new_input_shape, const Tensor& input,
const std::vector<int64_t>& 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 <typename AGG>
void NoTransposeReduce1Loop(Tensor* output, const TensorShape& new_input_shape, const Tensor& input,
const std::vector<int64_t>& reduced_axes, concurrency::ThreadPool* tp,
ResultsNoTransposePrepareForReduce& last_results) {
auto output_shape = output->Shape();
const T* from_data = input.template Data<T>();
const typename AGG::input_type* from_data = input.template Data<typename AGG::input_type>();
typename AGG::value_type* to_data = output->template MutableData<typename AGG::value_type>();
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<int64_t>& input_shape, const std::vector<int64_t>& axes, std::vector<int64_t>& dropped_axes) {
template <typename AGG>
void NoTransposeReduce2Loops(Tensor* output, const TensorShape& new_input_shape, const Tensor& input,
const std::vector<int64_t>& 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::input_type>();
typename AGG::value_type* to_data = output->template MutableData<typename AGG::value_type>();
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<int64_t>& input_shape,
const std::vector<int64_t>& axes,
std::vector<int64_t>& dropped_axes) {
auto dropped_dims = input_shape;
for (auto i : axes) {
dropped_dims[i] = -1;
@ -455,68 +511,273 @@ void DropDimensions(const std::vector<int64_t>& input_shape, const std::vector<i
}
}
template <typename T, typename AGG>
void CommonReduce(OpKernelContext* ctx,
const std::vector<int64_t> axes_, int64_t keepdims_,
ResultsNoTransposePrepareForReduce& last_results,
bool noop_with_empty_axes) {
std::vector<int64_t> axes;
const Tensor* input = ctx->Input<Tensor>(0);
auto reduced_dims = input->Shape().GetDims();
std::vector<int64_t> output_shape;
bool empty_reduce;
TensorShape new_input_shape;
FastReduceKind OptimizeShapeForFastReduce(const std::vector<int64_t>& input_shape,
const std::vector<int64_t>& reduced_axes,
std::vector<int64_t>& fast_shape,
std::vector<int64_t>& fast_output_shape,
std::vector<int64_t>& 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<Tensor>(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<size_t>(axes_tensor->Shape()[0]);
const auto* data = axes_tensor->template Data<int64_t>();
std::vector<int64_t> input_axes(data, data + nDims);
if (input_axes.empty() && noop_with_empty_axes) {
auto* output = ctx->Output(0, input->Shape());
memcpy(output->template MutableData<typename AGG::value_type>(), input->template Data<T>(), input->SizeInBytes());
return;
std::set<int64_t> 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<int64_t>(input_shape.size())));
}
}
fast_output_shape.clear();
fast_output_shape.reserve(input_shape.size());
bool empty_reduce = false;
std::vector<bool> 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<int64_t>());
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 <typename T, typename TVAL>
bool CommonFastReduceCopy(OpKernelContext* ctx, std::vector<int64_t>& input_axes, bool noop_with_empty_axes) {
if (ctx->InputCount() == 2) {
// second input holds the axes.
const Tensor* axes_tensor = ctx->Input<Tensor>(1);
ValidateCommonFastReduce(axes_tensor);
auto nDims = static_cast<size_t>(axes_tensor->Shape()[0]);
const auto* data = axes_tensor->template Data<int64_t>();
input_axes.insert(input_axes.begin(), data, data + nDims);
if (input_axes.empty() && noop_with_empty_axes) {
const Tensor* input = ctx->Input<Tensor>(0);
auto* output = ctx->Output(0, input->Shape());
memcpy(reinterpret_cast<void*>(output->template MutableData<float>()),
reinterpret_cast<const void*>(input->template Data<float>()),
input->SizeInBytes());
return true;
}
}
return false;
}
typedef void fast_reduce_fct(const Tensor& input, const std::vector<int64_t>& fast_shape,
Tensor& output, concurrency::ThreadPool* tp);
bool CommonFastReduceSwitch(OpKernelContext* ctx,
const std::vector<int64_t>& axes_,
int64_t keepdims_,
bool noop_with_empty_axes,
FastReduceKind& fast_kind,
std::vector<int64_t>& fast_shape,
std::vector<int64_t>& output_shape,
std::vector<int64_t>& fast_axes,
FastReduceKind which_fast_reduce,
fast_reduce_fct* case_kr,
fast_reduce_fct* case_rk,
fast_reduce_fct* case_krk) {
std::vector<int64_t> axes;
const Tensor* input = ctx->Input<Tensor>(0);
auto reduced_dims = input->Shape().GetDims();
std::vector<int64_t> 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 <typename AGG>
bool CommonFastReduce(OpKernelContext* ctx,
const std::vector<int64_t>& axes_,
int64_t keepdims_,
bool noop_with_empty_axes,
FastReduceKind& fast_kind,
std::vector<int64_t>& fast_shape,
std::vector<int64_t>& output_shape,
std::vector<int64_t>& 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 <typename AGG>
void CommonReduce1Loop(OpKernelContext* ctx,
const std::vector<int64_t>& axes_, int64_t keepdims_,
bool noop_with_empty_axes) {
FastReduceKind fast_kind;
std::vector<int64_t> fast_shape;
std::vector<int64_t> output_shape;
std::vector<int64_t> fast_axes;
if (CommonFastReduce<AGG>(ctx, axes_, keepdims_, noop_with_empty_axes,
fast_kind, fast_shape, output_shape, fast_axes)) {
return;
}
const Tensor* input = ctx->Input<Tensor>(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<T>();
const typename AGG::input_type* from_data = input->template Data<typename AGG::input_type>();
typename AGG::value_type* to_data = output->template MutableData<typename AGG::value_type>();
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<int64_t> dropped_axes;
DropDimensions(output_shape, axes, dropped_axes);
output = ctx->Output(0, dropped_axes);
ResultsNoTransposePrepareForReduce last_results;
NoTransposeReduce1Loop<AGG>(output, fast_shape, *input, fast_axes, ctx->GetOperatorThreadPool(), last_results);
}
template <typename AGG>
void CommonReduce2Loops(OpKernelContext* ctx,
const std::vector<int64_t>& axes_, int64_t keepdims_,
bool noop_with_empty_axes) {
FastReduceKind fast_kind;
std::vector<int64_t> fast_shape, output_shape, fast_axes;
if (CommonFastReduce<AGG>(ctx, axes_, keepdims_, noop_with_empty_axes,
fast_kind, fast_shape, output_shape, fast_axes)) {
return;
}
NoTransposeReduce<T, AGG>(output, new_input_shape, *input, axes, ctx->GetOperatorThreadPool(), last_results);
const Tensor* input = ctx->Input<Tensor>(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::input_type>();
typename AGG::value_type* to_data = output->template MutableData<typename AGG::value_type>();
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<AGG>(output, fast_shape, *input, fast_axes, ctx->GetOperatorThreadPool(), last_results);
}
template <typename T>
@ -524,64 +785,55 @@ Status ReduceL1<T>::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<T, ReduceAggregatorL1<T>>(ctx, axes_, keepdims_, last_results);
CommonReduce1Loop<ReduceAggregatorL1<T>>(ctx, axes_, keepdims_);
return Status::OK();
}
template <typename T>
Status ReduceL2<T>::Compute(OpKernelContext* ctx) const {
ResultsNoTransposePrepareForReduce last_results;
CommonReduce<T, ReduceAggregatorL2<T>>(ctx, axes_, keepdims_, last_results);
CommonReduce1Loop<ReduceAggregatorL2<T>>(ctx, axes_, keepdims_);
return Status::OK();
}
template <typename T>
Status ReduceLogSum<T>::Compute(OpKernelContext* ctx) const {
ResultsNoTransposePrepareForReduce last_results;
CommonReduce<T, ReduceAggregatorLogSum<T>>(ctx, axes_, keepdims_, last_results);
CommonReduce1Loop<ReduceAggregatorLogSum<T>>(ctx, axes_, keepdims_);
return Status::OK();
}
template <typename T>
Status ReduceLogSumExp<T>::Compute(OpKernelContext* ctx) const {
ResultsNoTransposePrepareForReduce last_results;
CommonReduce<T, ReduceAggregatorLogSumExp<T>>(ctx, axes_, keepdims_, last_results);
CommonReduce2Loops<ReduceAggregatorLogSumExp<T>>(ctx, axes_, keepdims_);
return Status::OK();
}
template <typename T>
Status ReduceMax<T>::Compute(OpKernelContext* ctx) const {
ResultsNoTransposePrepareForReduce last_results;
CommonReduce<T, ReduceAggregatorMax<T>>(ctx, axes_, keepdims_, last_results);
CommonReduce1Loop<ReduceAggregatorMax<T>>(ctx, axes_, keepdims_);
return Status::OK();
}
template <typename T>
Status ReduceMean<T>::Compute(OpKernelContext* ctx) const {
ResultsNoTransposePrepareForReduce last_results;
CommonReduce<T, ReduceAggregatorMean<T>>(ctx, axes_, keepdims_, last_results);
CommonReduce1Loop<ReduceAggregatorMean<T>>(ctx, axes_, keepdims_);
return Status::OK();
}
template <typename T>
Status ReduceMin<T>::Compute(OpKernelContext* ctx) const {
ResultsNoTransposePrepareForReduce last_results;
CommonReduce<T, ReduceAggregatorMin<T>>(ctx, axes_, keepdims_, last_results);
CommonReduce1Loop<ReduceAggregatorMin<T>>(ctx, axes_, keepdims_);
return Status::OK();
}
template <typename T>
Status ReduceProd<T>::Compute(OpKernelContext* ctx) const {
ResultsNoTransposePrepareForReduce last_results;
CommonReduce<T, ReduceAggregatorProd<T>>(ctx, axes_, keepdims_, last_results);
CommonReduce1Loop<ReduceAggregatorProd<T>>(ctx, axes_, keepdims_);
return Status::OK();
}
template <typename T>
Status ReduceSum<T>::Compute(OpKernelContext* ctx) const {
ResultsNoTransposePrepareForReduce last_results;
CommonReduce<T, ReduceAggregatorSum<T>>(ctx, axes_, keepdims_, last_results, noop_with_empty_axes_);
CommonReduce1Loop<ReduceAggregatorSum<T>>(ctx, axes_, keepdims_, noop_with_empty_axes_);
return Status::OK();
}
@ -590,67 +842,79 @@ Tensor ReduceSum<T>::Impl(const Tensor& input, const std::vector<int64_t>& reduc
AllocatorPtr allocator, concurrency::ThreadPool* tp, bool keep_dims,
const TensorShape* input_shape_override) {
std::vector<int64_t> axes;
auto reduced_dims = input.Shape().GetDims();
std::vector<int64_t> 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<int64_t> 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<int64_t>(), 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<int64_t>(), allocator);
if (fast_kind == FastReduceKind::kEmpty) {
if (new_input_shape.Size() == 1) {
const T* from_data = input.template Data<T>();
T* to_data = output.template MutableData<T>();
*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<T, ReduceAggregatorSum<T>>(&output, new_input_shape, input, axes, tp, last_results);
return output;
} else {
ResultsNoTransposePrepareForReduce last_results;
std::vector<int64_t> dropped_axes;
DropDimensions(output_shape, axes, dropped_axes);
Tensor output(input.DataType(), dropped_axes, allocator);
NoTransposeReduce<T, ReduceAggregatorSum<T>>(&output, new_input_shape, input, axes, tp, last_results);
return output;
if (IsFastReduceKindAvailable(fast_kind, ReduceAggregatorSum<T>::WhichFastReduce())) {
switch (fast_kind) {
case FastReduceKind::kKR: {
ValidateFastReduceKR(fast_shape, output);
ReduceAggregatorSum<T>::FastReduceKR(input, fast_shape, output, tp);
return output;
}
case FastReduceKind::kRK: {
ValidateFastReduceRK(fast_shape, output);
ReduceAggregatorSum<T>::FastReduceRK(input, fast_shape, output, tp);
return output;
}
case FastReduceKind::kKRK: {
ValidateFastReduceKRK(fast_shape, output);
ReduceAggregatorSum<T>::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<ReduceAggregatorSum<T>>(&output, fast_shape, input, fast_axes, tp, last_results);
return output;
}
template <typename T>
Status ReduceSumSquare<T>::Compute(OpKernelContext* ctx) const {
ResultsNoTransposePrepareForReduce last_results;
CommonReduce<T, ReduceAggregatorSumSquare<T>>(ctx, axes_, keepdims_, last_results);
CommonReduce1Loop<ReduceAggregatorSumSquare<T>>(ctx, axes_, keepdims_);
return Status::OK();
}
template <typename T>
Status ArgMax<T>::Compute(OpKernelContext* ctx) const {
ResultsNoTransposePrepareForReduce last_results;
if (select_last_index_) {
CommonReduce<T, ReduceAggregatorArgMaxLastIndex<T>>(ctx, axes_, keepdims_, last_results);
CommonReduce1Loop<ReduceAggregatorArgMaxLastIndex<T>>(ctx, axes_, keepdims_);
} else {
CommonReduce<T, ReduceAggregatorArgMax<T>>(ctx, axes_, keepdims_, last_results);
CommonReduce1Loop<ReduceAggregatorArgMax<T>>(ctx, axes_, keepdims_);
}
return Status::OK();
}
template <typename T>
Status ArgMin<T>::Compute(OpKernelContext* ctx) const {
ResultsNoTransposePrepareForReduce last_results;
if (select_last_index_) {
CommonReduce<T, ReduceAggregatorArgMinLastIndex<T>>(ctx, axes_, keepdims_, last_results);
CommonReduce1Loop<ReduceAggregatorArgMinLastIndex<T>>(ctx, axes_, keepdims_);
} else {
CommonReduce<T, ReduceAggregatorArgMin<T>>(ctx, axes_, keepdims_, last_results);
CommonReduce1Loop<ReduceAggregatorArgMin<T>>(ctx, axes_, keepdims_);
}
return Status::OK();
}
@ -664,4 +928,17 @@ template class ReduceSum<int32_t>;
template class ReduceSum<double>;
template class ReduceSum<int64_t>;
template void CommonReduce1Loop<ReduceAggregatorSum<float>>(OpKernelContext* ctx,
const std::vector<int64_t>& axes_, int64_t keepdims_,
bool noop_with_empty_axes);
template void CommonReduce1Loop<ReduceAggregatorSum<int32_t>>(OpKernelContext* ctx,
const std::vector<int64_t>& axes_, int64_t keepdims_,
bool noop_with_empty_axes);
template void CommonReduce1Loop<ReduceAggregatorSum<double>>(OpKernelContext* ctx,
const std::vector<int64_t>& axes_, int64_t keepdims_,
bool noop_with_empty_axes);
template void CommonReduce1Loop<ReduceAggregatorSum<int64_t>>(OpKernelContext* ctx,
const std::vector<int64_t>& axes_, int64_t keepdims_,
bool noop_with_empty_axes);
} // namespace onnxruntime

356
onnxruntime/core/providers/cpu/reduction/reduction_ops.h
View file

@ -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 <cmath>
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<int64_t>& input_shape,
const std::vector<int64_t>& reduced_axes,
std::vector<int64_t>& fast_shape,
std::vector<int64_t>& fast_output_shape,
std::vector<int64_t>& fast_axes,
bool keep_dims, bool noop_with_empty_axes = false);
class ResultsNoTransposePrepareForReduce {
public:
std::vector<int64_t> input_shape;
@ -32,42 +77,35 @@ class ResultsNoTransposePrepareForReduce {
last_loop_inc = 0;
}
bool equal(const std::vector<int64_t>& local_input_shape, const std::vector<int64_t>& 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<int64_t>::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<int64_t>::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<int64_t>& local_input_shape, const std::vector<int64_t>& local_reduced_axes);
void ValidateNotEmpty();
};
template <typename T>
inline T reduce_sqrt(T value) { return std::sqrt(value); }
template <>
inline int64_t reduce_sqrt<int64_t>(int64_t value) { return static_cast<int64_t>(std::sqrt(static_cast<double>(value))); }
inline int64_t reduce_sqrt<int64_t>(int64_t value) {
return static_cast<int64_t>(std::sqrt(static_cast<double>(value)));
}
template <>
inline int32_t reduce_sqrt<int32_t>(int32_t value) { return static_cast<int32_t>(std::sqrt(static_cast<double>(value))); }
inline int32_t reduce_sqrt<int32_t>(int32_t value) {
return static_cast<int32_t>(std::sqrt(static_cast<double>(value)));
}
template <typename T>
inline T reduce_log(T value) { return static_cast<T>(std::log(value)); }
template <>
inline int64_t reduce_log<int64_t>(int64_t value) { return static_cast<int64_t>(std::log(static_cast<double>(value))); }
inline int64_t reduce_log<int64_t>(int64_t value) {
return static_cast<int64_t>(std::log(static_cast<double>(value)));
}
template <>
inline int32_t reduce_log<int32_t>(int32_t value) { return static_cast<int32_t>(std::log(static_cast<double>(value))); }
inline int32_t reduce_log<int32_t>(int32_t value) {
return static_cast<int32_t>(std::log(static_cast<double>(value)));
}
template <typename T>
inline T reduce_exp(T value) { return static_cast<T>(std::exp(value)); }
@ -102,9 +140,19 @@ inline bool reduce_isnan<int32_t>(int32_t) { return false; }
template <>
inline bool reduce_isnan<int64_t>(int64_t) { return false; }
template <typename T, typename TVAL = T>
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<int64_t>&, Tensor&, concurrency::ThreadPool*);
static void FastReduceRK(const Tensor&, const std::vector<int64_t>&, Tensor&, concurrency::ThreadPool*);
static void FastReduceKRK(const Tensor&, const std::vector<int64_t>&, Tensor&, concurrency::ThreadPool*);
};
template <typename T, typename TVAL = T>
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 <typename T, typename TVAL = T>
@ -132,6 +178,60 @@ class ReduceAggregatorSum : public ReduceAggregator<T, TVAL> {
inline TVAL aggall(const T* from_data) {
return Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, 1>>(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<int64_t>& fast_shape,
Tensor& output, concurrency::ThreadPool* tp) {
const T* data = input.Data<T>();
T* out = output.MutableData<T>();
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<T>(data + d * stridei, stridei).sum();
}
});
}
static void FastReduceRK(const Tensor& input, const std::vector<int64_t>& fast_shape,
Tensor& output, concurrency::ThreadPool* tp) {
int64_t N = fast_shape[1];
const T* data = input.Data<T>();
T* out = output.MutableData<T>();
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<T>(out + begin, end - begin) += ConstEigenVectorArrayMap<T>(
data + row * N + begin, end - begin);
}
});
}
static void FastReduceKRK(const Tensor& input, const std::vector<int64_t>& fast_shape,
Tensor& output, concurrency::ThreadPool* tp) {
int64_t N = fast_shape[2];
const T* data = input.Data<T>();
int64_t stridei = fast_shape[1] * fast_shape[2];
int64_t strideo = fast_shape[2];
T* out = output.MutableData<T>();
std::vector<T> 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<T>(1, N, fast_shape[1], one.data(), data + stridei * d, out + strideo * d, nullptr);
}
});
}
};
template <typename T, typename TVAL = T>
@ -152,6 +252,48 @@ class ReduceAggregatorMean : public ReduceAggregatorSum<T, TVAL> {
return Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, 1>>(from_data, this->N_).mean();
}
inline T get_value() { return this->accumulator_ / static_cast<T>(this->N_); }
// Fast reduction
// WhichFastReduce() already defined in ReduceAggregatorSum
static void FastReduceKR(const Tensor& input, const std::vector<int64_t>& fast_shape,
Tensor& output, concurrency::ThreadPool* tp) {
ReduceAggregatorSum<T, TVAL>::FastReduceKR(input, fast_shape, output, tp);
// TODO: use MLAS or BLAS
T* out = output.MutableData<T>();
T* end = out + fast_shape[0];
for (; out != end; ++out) {
*out /= static_cast<T>(fast_shape[1]);
}
}
static void FastReduceRK(const Tensor& input, const std::vector<int64_t>& fast_shape,
Tensor& output, concurrency::ThreadPool* tp) {
ReduceAggregatorSum<T, TVAL>::FastReduceRK(input, fast_shape, output, tp);
// TODO: use MLAS or BLAS
T* out = output.MutableData<T>();
T* end = out + fast_shape[1];
for (; out != end; ++out) {
*out /= static_cast<T>(fast_shape[0]);
}
}
static void FastReduceKRK(const Tensor& input, const std::vector<int64_t>& fast_shape,
Tensor& output, concurrency::ThreadPool* tp) {
ReduceAggregatorSum<T, TVAL>::FastReduceKRK(input, fast_shape, output, tp);
int64_t strideo = fast_shape[2];
T* out = output.MutableData<T>();
T* begin;
T* end;
T div = static_cast<T>(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 <typename T, typename TVAL = T>
@ -162,6 +304,66 @@ class ReduceAggregatorMax : public ReduceAggregator<T, TVAL> {
return Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, 1>>(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<int64_t>& fast_shape,
Tensor& output, concurrency::ThreadPool* tp) {
const T* data = input.Data<T>();
T* out = output.MutableData<T>();
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<T>(out + first, last - first) = ConstEigenMatrixMap<T>(
data + first * stridei, stridei, last - first)
.colwise()
.maxCoeff();
});
}
static void FastReduceRK(const Tensor& input, const std::vector<int64_t>& 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>();
T* out = output.MutableData<T>();
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<int64_t>& fast_shape,
Tensor& output, concurrency::ThreadPool* tp) {
const T* data = input.Data<T>();
T* out = output.MutableData<T>();
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<T>(out + j * strideo, strideo) =
ConstEigenMatrixMap<T>(
data + j * stridei, fast_shape[2], fast_shape[1])
.rowwise()
.maxCoeff();
}
});
}
};
template <typename T, typename TVAL = int64_t>
@ -261,6 +463,66 @@ class ReduceAggregatorMin : public ReduceAggregator<T, TVAL> {
return Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, 1>>(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<int64_t>& fast_shape,
Tensor& output, concurrency::ThreadPool* tp) {
const T* data = input.Data<T>();
T* out = output.MutableData<T>();
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<T>(out + first, last - first) = ConstEigenMatrixMap<T>(
data + first * stridei, stridei, last - first)
.colwise()
.minCoeff();
});
}
static void FastReduceRK(const Tensor& input, const std::vector<int64_t>& 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>();
T* out = output.MutableData<T>();
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<int64_t>& fast_shape,
Tensor& output, concurrency::ThreadPool* tp) {
const T* data = input.Data<T>();
T* out = output.MutableData<T>();
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<T>(out + j * strideo, strideo) =
ConstEigenMatrixMap<T>(
data + j * stridei, fast_shape[2], fast_shape[1])
.rowwise()
.minCoeff();
}
});
}
};
template <typename T, typename TVAL = T>
@ -326,31 +588,33 @@ class ReduceAggregatorLogSumExp : public ReduceAggregator<T, TVAL> {
}
inline void update(const T& v) { this->accumulator_ += reduce_exp(v - max_); }
inline TVAL get_value() { return reduce_log<T>(this->accumulator_) + max_; }
static inline bool two_loops() { return true; }
};
bool SetupForReduce(const Tensor* input_tensor_ptr,
const std::vector<int64_t>& axes_,
std::vector<int64_t>& axes,
TensorShape& new_input_shape,
std::vector<int64_t>& output_shape,
bool& empty_reduce,
const TensorShape* input_shape_override);
void NoTransposePrepareForReduce(const TensorShape& new_input_shape,
const std::vector<int64_t>& reduced_axes,
ResultsNoTransposePrepareForReduce& results);
template <typename T, typename AGG>
void NoTransposeReduce(Tensor* output, const TensorShape& new_input_shape, const Tensor& input,
const std::vector<int64_t>& reduced_axes, concurrency::ThreadPool* tp,
ResultsNoTransposePrepareForReduce& last_results);
template <typename AGG>
void NoTransposeReduce1Loop(Tensor* output, const TensorShape& new_input_shape, const Tensor& input,
const std::vector<int64_t>& reduced_axes, concurrency::ThreadPool* tp,
ResultsNoTransposePrepareForReduce& last_results);
template <typename T, typename AGG>
void CommonReduce(OpKernelContext* ctx,
const std::vector<int64_t> axes_, int64_t keepdims_,
ResultsNoTransposePrepareForReduce& last_results,
bool noop_with_empty_axes = false);
// Specific case for ReduceLogSumExp.
template <typename AGG>
void NoTransposeReduce2Loops(Tensor* output, const TensorShape& new_input_shape, const Tensor& input,
const std::vector<int64_t>& reduced_axes, concurrency::ThreadPool* tp,
ResultsNoTransposePrepareForReduce& last_results);
template <typename AGG>
void CommonReduce1Loop(OpKernelContext* ctx,
const std::vector<int64_t>& axes_, int64_t keepdims_,
bool noop_with_empty_axes = false);
// Specific case for ReduceLogSumExp.
template <typename AGG>
void CommonReduce2Loops(OpKernelContext* ctx,
const std::vector<int64_t>& axes_, int64_t keepdims_,
bool noop_with_empty_axes = false);
template <bool allow_multi_axes>
class ReduceKernelBase {

3
onnxruntime/core/util/math_cpu.cc
View file

@ -859,6 +859,9 @@ DEFINE_BROADCAST_BINARY_FUNCTION(Div, /)
EigenVectorMap<T>(y, N) = ConstEigenMatrixMap<T>(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) \

1
onnxruntime/test/providers/cpu/math/einsum_test.cc
View file

@ -96,6 +96,7 @@ TEST(Einsum, ExplicitEinsumAsReduceOp_2D_input_1) {
test.AddOutput<float>("y", {2}, {4.f, 6.f});
test.Run();
}
TEST(Einsum, ExplicitEinsumAsBatchedReduceOp_3D_input_0) {
OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
test.AddAttribute<std::string>("equation", "...ji->...j");

1494
onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc
View file

File diff suppressed because it is too large Load diff

3
orttraining/orttraining/training_ops/cpu/reduction/reduction_ops.cc
View file

@ -29,8 +29,7 @@ REGISTER_REDUCESUMTRAINING_KERNEL_TYPED(int64_t)
template <typename T>
Status ReduceSumTraining<T>::Compute(OpKernelContext* ctx) const {
ResultsNoTransposePrepareForReduce last_results;
CommonReduce<T, ReduceAggregatorSum<T>>(ctx, axes_, keepdims_, last_results, noop_with_empty_axes_);
CommonReduce1Loop<ReduceAggregatorSum<T>>(ctx, axes_, keepdims_, noop_with_empty_axes_);
return Status::OK();
}