diff options
Diffstat (limited to 'compiler/mir-interpreter/src/ops/ReduceMean.cpp')
-rw-r--r-- | compiler/mir-interpreter/src/ops/ReduceMean.cpp | 98 |
1 files changed, 98 insertions, 0 deletions
diff --git a/compiler/mir-interpreter/src/ops/ReduceMean.cpp b/compiler/mir-interpreter/src/ops/ReduceMean.cpp new file mode 100644 index 000000000..ebaa3b48f --- /dev/null +++ b/compiler/mir-interpreter/src/ops/ReduceMean.cpp @@ -0,0 +1,98 @@ +/* + * Copyright (c) 2020 Samsung Electronics Co., Ltd. All Rights Reserved + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef _NNC_CORE_BACKEND_INTERPRETER_REDUCE_MEAN_ +#define _NNC_CORE_BACKEND_INTERPRETER_REDUCE_MEAN_ + +#include "ReduceMean.h" +#include "Common.h" + +#include "mir/ops/ReduceMeanOp.h" +#include "mir/Tensor.h" +#include "mir/ShapeRange.h" + +namespace mir_interpreter +{ + +template <typename T> struct ReduceMeanImpl +{ + static void run(const mir::TensorVariant &inputv, const mir::ops::ReduceMeanOp &op, + mir::TensorVariant &output); +}; + +template <typename T> +void ReduceMeanImpl<T>::run(const mir::TensorVariant &inputv, const mir::ops::ReduceMeanOp &op, + mir::TensorVariant &output) +{ + const auto &input_shape = op.getInputShape(0); + const auto &output_shape = op.getOutputShape(0); + const auto &reduction_dims = op.getReductionDims(); + const bool keep_dims = op.getKeepDims(); + + const auto reductor = [](T result, T x) { return result + x; }; + + mir::Tensor<T> input(inputv); + mir::Tensor<T> res_accessor(output); + + erase<T>(output); + + // This mask contains 'true' for dimensions that should be reduced. For example, if we want + // to reduce dimensions 1 and 3 with total number of dimensions of 4, the mask will be + // [false, true, false, true]. + std::vector<bool> reduction_dims_mask(input_shape.rank(), false); + for (const int dim : reduction_dims) + { + reduction_dims_mask[dim] = true; + } + + mir::Index out_index(output_shape.rank()); + for (const mir::Index &in_index : mir::ShapeRange(input_shape)) + { + int out_index_dim = 0; + for (int dim = 0; dim < input_shape.rank(); ++dim) + { + if (keep_dims) + { + out_index.at(out_index_dim++) = reduction_dims_mask[dim] ? 0 : in_index.at(dim); + } + else + { + if (!reduction_dims_mask[dim]) + { + out_index.at(out_index_dim++) = in_index.at(dim); + } + } + } + res_accessor.at(out_index) = reductor(res_accessor.at(out_index), input.at(in_index)); + } + + const std::int32_t reduction_factor = input_shape.numElements() / output_shape.numElements(); + + for (const auto &index : mir::ShapeRange(output_shape)) + { + res_accessor.at(index) /= reduction_factor; + } +} + +void ReduceMean(const mir::TensorVariant &input, const mir::ops::ReduceMeanOp &op, + mir::TensorVariant &output) +{ + dispatch<ReduceMeanImpl>(input.getElementType(), input, op, output); +}; + +} // namespace mir_interpreter + +#endif //_NNC_CORE_BACKEND_INTERPRETER_REDUCE_MEAN_ |