wehub-resource-sync
809 lines
32 KiB
C++
809 lines
32 KiB
C++
/* Copyright 2022 The TensorFlow Authors. 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 TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_REDUCE_H_
|
|
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_REDUCE_H_
|
|
|
|
#include <stdint.h>
|
|
|
|
#include <algorithm>
|
|
#include <limits>
|
|
#include <vector>
|
|
|
|
#include "ruy/profiler/instrumentation.h" // from @ruy
|
|
#include "tensorflow/lite/kernels/cpu_backend_threadpool.h"
|
|
#include "tensorflow/lite/kernels/internal/optimized/optimized_ops_utils.h"
|
|
#include "tensorflow/lite/kernels/internal/optimized/reduce_utils.h"
|
|
#include "tensorflow/lite/kernels/internal/reduce_common.h"
|
|
#include "tensorflow/lite/kernels/internal/reference/reduce.h"
|
|
#include "tensorflow/lite/kernels/internal/runtime_shape.h"
|
|
#include "tensorflow/lite/kernels/internal/types.h"
|
|
#include "tensorflow/lite/kernels/kernel_util.h"
|
|
|
|
namespace tflite {
|
|
namespace optimized_ops {
|
|
|
|
inline void MeanImpl(const tflite::MeanParams& op_params,
|
|
const RuntimeShape& input_shape, const uint8_t* input_data,
|
|
int32 multiplier, int32 shift, int32 bias,
|
|
const RuntimeShape& output_shape, uint8_t* output_data,
|
|
int start_depth, int end_depth) {
|
|
ruy::profiler::ScopeLabel label("Mean4D/Uint8/MeanImpl");
|
|
|
|
// Current implementation only supports dimension equals 4 and simultaneous
|
|
// reduction over width and height.
|
|
const int output_batch = output_shape.Dims(0);
|
|
const int output_height = output_shape.Dims(2);
|
|
const int output_width = output_shape.Dims(2);
|
|
const int input_height = input_shape.Dims(1);
|
|
const int input_width = input_shape.Dims(2);
|
|
|
|
TFLITE_CHECK_EQ(op_params.axis_count, 2);
|
|
TFLITE_CHECK((op_params.axis[0] == 1 && op_params.axis[1] == 2) ||
|
|
(op_params.axis[0] == 2 && op_params.axis[1] == 1));
|
|
TFLITE_CHECK_EQ(output_height, 1);
|
|
TFLITE_CHECK_EQ(output_width, 1);
|
|
|
|
constexpr int32_t kMinValue = std::numeric_limits<uint8_t>::min();
|
|
constexpr int32_t kMaxValue = std::numeric_limits<uint8_t>::max();
|
|
|
|
#ifdef USE_NEON
|
|
const int32x4_t bias_dup = vdupq_n_s32(bias);
|
|
const int32x4_t min_dup = vdupq_n_s32(kMinValue);
|
|
const int32x4_t max_dup = vdupq_n_s32(kMaxValue);
|
|
#endif // USE_NEON
|
|
|
|
for (int out_b = 0; out_b < output_batch; ++out_b) {
|
|
int out_d = start_depth;
|
|
#ifdef USE_NEON
|
|
|
|
for (; out_d <= end_depth - 16; out_d += 16) {
|
|
int32x4x4_t temp_sum;
|
|
temp_sum.val[0] = vdupq_n_s32(0);
|
|
temp_sum.val[1] = vdupq_n_s32(0);
|
|
temp_sum.val[2] = vdupq_n_s32(0);
|
|
temp_sum.val[3] = vdupq_n_s32(0);
|
|
for (int in_h = 0; in_h < input_height; ++in_h) {
|
|
for (int in_w = 0; in_w < input_width; ++in_w) {
|
|
const uint8_t* input_data_ptr =
|
|
input_data + Offset(input_shape, out_b, in_h, in_w, out_d);
|
|
uint8x16_t input_data_val = vld1q_u8(input_data_ptr);
|
|
|
|
int16x8_t input_data_low_shift =
|
|
vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(input_data_val)));
|
|
int16x8_t input_data_high_shift =
|
|
vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(input_data_val)));
|
|
|
|
int32x4_t input_low_low =
|
|
vmovl_s16(vget_low_s16(input_data_low_shift));
|
|
int32x4_t input_high_low =
|
|
vmovl_s16(vget_high_s16(input_data_low_shift));
|
|
int32x4_t input_low_high =
|
|
vmovl_s16(vget_low_s16(input_data_high_shift));
|
|
int32x4_t input_high_high =
|
|
vmovl_s16(vget_high_s16(input_data_high_shift));
|
|
|
|
temp_sum.val[0] = vaddq_s32(temp_sum.val[0], input_low_low);
|
|
temp_sum.val[1] = vaddq_s32(temp_sum.val[1], input_high_low);
|
|
temp_sum.val[2] = vaddq_s32(temp_sum.val[2], input_low_high);
|
|
temp_sum.val[3] = vaddq_s32(temp_sum.val[3], input_high_high);
|
|
}
|
|
}
|
|
|
|
temp_sum =
|
|
MultiplyByQuantizedMultiplier4Rows(temp_sum, multiplier, shift);
|
|
|
|
temp_sum.val[0] = vaddq_s32(temp_sum.val[0], bias_dup);
|
|
temp_sum.val[1] = vaddq_s32(temp_sum.val[1], bias_dup);
|
|
temp_sum.val[2] = vaddq_s32(temp_sum.val[2], bias_dup);
|
|
temp_sum.val[3] = vaddq_s32(temp_sum.val[3], bias_dup);
|
|
|
|
temp_sum.val[0] = vminq_s32(vmaxq_s32(temp_sum.val[0], min_dup), max_dup);
|
|
temp_sum.val[1] = vminq_s32(vmaxq_s32(temp_sum.val[1], min_dup), max_dup);
|
|
temp_sum.val[2] = vminq_s32(vmaxq_s32(temp_sum.val[2], min_dup), max_dup);
|
|
temp_sum.val[3] = vminq_s32(vmaxq_s32(temp_sum.val[3], min_dup), max_dup);
|
|
|
|
uint16x4_t narrowed_low_low =
|
|
vmovn_u32(vreinterpretq_u32_s32(temp_sum.val[0]));
|
|
uint16x4_t narrowed_high_low =
|
|
vmovn_u32(vreinterpretq_u32_s32(temp_sum.val[1]));
|
|
uint16x4_t narrowed_low_high =
|
|
vmovn_u32(vreinterpretq_u32_s32(temp_sum.val[2]));
|
|
uint16x4_t narrowed_high_high =
|
|
vmovn_u32(vreinterpretq_u32_s32(temp_sum.val[3]));
|
|
|
|
uint16x8_t combined_low =
|
|
vcombine_u16(narrowed_low_low, narrowed_high_low);
|
|
uint16x8_t combined_high =
|
|
vcombine_u16(narrowed_low_high, narrowed_high_high);
|
|
|
|
uint8x8_t narrowed_low = vmovn_u16(combined_low);
|
|
uint8x8_t narrowed_high = vmovn_u16(combined_high);
|
|
|
|
uint8x16_t combined_output = vcombine_u8(narrowed_low, narrowed_high);
|
|
|
|
uint8_t* output_data_ptr =
|
|
output_data + Offset(output_shape, out_b, 0, 0, out_d);
|
|
vst1q_u8(output_data_ptr, combined_output);
|
|
}
|
|
#endif // USE_NEON
|
|
|
|
for (; out_d < end_depth; ++out_d) {
|
|
int acc = 0;
|
|
for (int in_h = 0; in_h < input_height; ++in_h) {
|
|
for (int in_w = 0; in_w < input_width; ++in_w) {
|
|
acc += input_data[Offset(input_shape, out_b, in_h, in_w, out_d)];
|
|
}
|
|
}
|
|
|
|
acc = MultiplyByQuantizedMultiplier(acc, multiplier, shift);
|
|
acc += bias;
|
|
acc = std::min(std::max(acc, kMinValue), kMaxValue);
|
|
output_data[Offset(output_shape, out_b, 0, 0, out_d)] =
|
|
static_cast<uint8_t>(acc);
|
|
}
|
|
}
|
|
}
|
|
|
|
struct MeanWorkerTask : cpu_backend_threadpool::Task {
|
|
MeanWorkerTask(const tflite::MeanParams& op_params,
|
|
const RuntimeShape& input_shape, const uint8_t* input_data,
|
|
int32 multiplier, int32 shift, int32 bias,
|
|
const RuntimeShape& output_shape, uint8_t* output_data,
|
|
int start_height, int end_height)
|
|
: op_params(op_params),
|
|
input_shape(input_shape),
|
|
input_data(input_data),
|
|
multiplier(multiplier),
|
|
shift(shift),
|
|
bias(bias),
|
|
output_shape(output_shape),
|
|
output_data(output_data),
|
|
start_height(start_height),
|
|
end_height(end_height) {}
|
|
|
|
void Run() override {
|
|
MeanImpl(op_params, input_shape, input_data, multiplier, shift, bias,
|
|
output_shape, output_data, start_height, end_height);
|
|
}
|
|
|
|
private:
|
|
const tflite::MeanParams& op_params;
|
|
const RuntimeShape& input_shape;
|
|
const uint8_t* input_data;
|
|
int32 multiplier;
|
|
int32 shift;
|
|
int32 bias;
|
|
const RuntimeShape& output_shape;
|
|
uint8_t* output_data;
|
|
int start_height;
|
|
int end_height;
|
|
};
|
|
|
|
inline void Mean(const tflite::MeanParams& op_params,
|
|
const RuntimeShape& unextended_input_shape,
|
|
const uint8_t* input_data, int32 input_zero_point,
|
|
float input_scale, const RuntimeShape& unextended_output_shape,
|
|
uint8_t* output_data, int32 output_zero_point,
|
|
float output_scale, CpuBackendContext* cpu_backend_context) {
|
|
ruy::profiler::ScopeLabel label("Mean4D/Uint8");
|
|
// Current implementation only supports dimension equals 4 and simultaneous
|
|
// reduction over width and height.
|
|
TFLITE_CHECK_EQ(unextended_input_shape.DimensionsCount(), 4);
|
|
TFLITE_CHECK_LE(unextended_output_shape.DimensionsCount(), 4);
|
|
const RuntimeShape input_shape =
|
|
RuntimeShape::ExtendedShape(4, unextended_input_shape);
|
|
const RuntimeShape output_shape =
|
|
RuntimeShape::ExtendedShape(4, unextended_output_shape);
|
|
const int output_height = output_shape.Dims(1);
|
|
const int output_width = output_shape.Dims(2);
|
|
const int output_depth = output_shape.Dims(3);
|
|
|
|
TFLITE_CHECK_EQ(op_params.axis_count, 2);
|
|
TFLITE_CHECK((op_params.axis[0] == 1 && op_params.axis[1] == 2) ||
|
|
(op_params.axis[0] == 2 && op_params.axis[1] == 1));
|
|
TFLITE_CHECK_EQ(output_height, 1);
|
|
TFLITE_CHECK_EQ(output_width, 1);
|
|
|
|
const int input_height = input_shape.Dims(1);
|
|
const int input_width = input_shape.Dims(2);
|
|
const float num_elements_in_axis = input_width * input_height;
|
|
|
|
float temp = input_zero_point * input_scale / output_scale;
|
|
temp = temp > 0 ? temp + 0.5f : temp - 0.5f;
|
|
int32_t bias = output_zero_point - static_cast<int32_t>(temp);
|
|
float real_scale = input_scale / (num_elements_in_axis * output_scale);
|
|
|
|
int32 multiplier, shift;
|
|
QuantizeMultiplier(real_scale, &multiplier, &shift);
|
|
|
|
constexpr int kMinDepthPerThread = 8;
|
|
int thread_count = output_depth / kMinDepthPerThread;
|
|
thread_count = thread_count > 0 ? thread_count : 1;
|
|
const int capped_thread_count =
|
|
std::min(thread_count, cpu_backend_context->max_num_threads());
|
|
|
|
if (capped_thread_count == 1) {
|
|
MeanImpl(op_params, input_shape, input_data, multiplier, shift, bias,
|
|
output_shape, output_data, 0, output_depth);
|
|
} else {
|
|
// Instead parallel for batch, we loop for the output_depth since batch
|
|
// is typical 1.
|
|
std::vector<MeanWorkerTask> tasks;
|
|
// TODO(b/131746020) don't create new heap allocations every time.
|
|
// At least we make it a single heap allocation by using reserve().
|
|
tasks.reserve(capped_thread_count);
|
|
int depth_start = 0;
|
|
for (int i = 0; i < capped_thread_count; ++i) {
|
|
// Try to distribute the tasks as even as possible.
|
|
int depth_end = depth_start +
|
|
(output_depth - depth_start) / (capped_thread_count - i);
|
|
tasks.emplace_back(op_params, input_shape, input_data, multiplier, shift,
|
|
bias, output_shape, output_data, depth_start,
|
|
depth_end);
|
|
depth_start = depth_end;
|
|
}
|
|
cpu_backend_threadpool::Execute(tasks.size(), tasks.data(),
|
|
cpu_backend_context);
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
struct SumOp {
|
|
inline T operator()(const T& a) const { return a; }
|
|
inline T operator()(const T& a, const T& b) const { return a + b; }
|
|
static constexpr T kNeutralElement = T(0);
|
|
};
|
|
|
|
template <typename T, typename U>
|
|
struct CastSumOp {
|
|
inline U operator()(const T& a) const { return static_cast<U>(a); }
|
|
inline U operator()(const U& a, const T& b) const {
|
|
return a + static_cast<U>(b);
|
|
}
|
|
static constexpr U kNeutralElement = U(0);
|
|
};
|
|
|
|
template <typename T>
|
|
struct ProdOp {
|
|
inline T operator()(const T& a) const { return a; }
|
|
inline T operator()(const T& a, const T& b) const { return a * b; }
|
|
static constexpr T kNeutralElement = T(1);
|
|
};
|
|
|
|
template <typename T>
|
|
struct MaxOp {
|
|
inline T operator()(const T& a) const { return a; }
|
|
inline T operator()(const T& a, const T& b) const { return (a > b) ? a : b; }
|
|
static constexpr T kNeutralElement = std::numeric_limits<T>::lowest();
|
|
};
|
|
|
|
template <typename T>
|
|
struct MinOp {
|
|
inline T operator()(const T& a) const { return a; }
|
|
inline T operator()(const T& a, const T& b) const { return (a < b) ? a : b; }
|
|
static constexpr T kNeutralElement = std::numeric_limits<T>::max();
|
|
};
|
|
|
|
struct AndOp {
|
|
inline bool operator()(bool a) const { return a; }
|
|
inline bool operator()(bool a, bool b) const { return a && b; }
|
|
static constexpr bool kNeutralElement = true;
|
|
};
|
|
|
|
struct OrOp {
|
|
inline bool operator()(bool a) const { return a; }
|
|
inline bool operator()(bool a, bool b) const { return a || b; }
|
|
static constexpr bool kNeutralElement = false;
|
|
};
|
|
|
|
// When the number of axis is zero, the reduction is simply a copy.
|
|
template <typename T>
|
|
void ReduceIsCopy(const T* input_data, const int* input_dims,
|
|
const int input_num_dims, T* output_data) {
|
|
int num_elems = NumElements(input_dims, input_num_dims);
|
|
memcpy(output_data, input_data, num_elems * sizeof(T));
|
|
}
|
|
|
|
// Reduces the input over either odd or even dimensions using Op.
|
|
// One recursive call for each dimension is made.
|
|
// 'depth' is the depth of recursion.
|
|
// 'parity' indicates whether odd or even dimensions are being reduced.
|
|
// ReducerFirst is applied to the first element to be written to each output
|
|
// position.
|
|
// ReducerNext is applied to each subsequent element to be written to each
|
|
// output position.
|
|
template <typename T, typename U, typename ReducerFirst, typename ReducerNext>
|
|
inline std::pair<const T*, U*> ReduceImpl(const T* input_data,
|
|
const int* input_dims, U* output_data,
|
|
int depth, int parity, bool next,
|
|
const ReducerFirst& reducer_first,
|
|
const ReducerNext& reducer_next) {
|
|
// The output pointer is incremented conditionally depending on whether the
|
|
// odd or even dimension is being reduced.
|
|
// The input pointer is always incremented as each input is read once.
|
|
if (depth > 0) {
|
|
U* future_output = output_data;
|
|
bool update_output = (depth % 2) == parity;
|
|
for (int i = 0; i < input_dims[0]; ++i) {
|
|
if (i > 0 && !update_output) {
|
|
next = true;
|
|
}
|
|
std::tie(input_data, future_output) =
|
|
ReduceImpl(input_data, &input_dims[1], output_data, depth - 1, parity,
|
|
next, reducer_first, reducer_next);
|
|
if (update_output) {
|
|
output_data = future_output;
|
|
}
|
|
}
|
|
output_data = future_output;
|
|
} else {
|
|
// Reduce the final dimension.
|
|
if (parity) {
|
|
// Reduce the even dimension. The entire dimension is reduced into one
|
|
// value.
|
|
U res = next ? reducer_next(*output_data, *input_data++)
|
|
: reducer_first(*input_data++);
|
|
for (int i = 1; i < input_dims[0]; ++i) {
|
|
res = reducer_next(res, *input_data++);
|
|
}
|
|
*output_data++ = res;
|
|
} else {
|
|
// Reduce the odd dimension. Each input is accumulated into a separate
|
|
// output.
|
|
if (!next) {
|
|
for (int i = 0; i < input_dims[0]; ++i) {
|
|
U res = reducer_first(*input_data++);
|
|
*output_data++ = res;
|
|
}
|
|
} else {
|
|
for (int i = 0; i < input_dims[0]; ++i) {
|
|
U res = *output_data;
|
|
res = reducer_next(res, *input_data++);
|
|
*output_data++ = res;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return {input_data, output_data};
|
|
}
|
|
|
|
// A generic reduce method that can be used for reduce_sum, reduce_mean, etc.
|
|
// This method iterates through input data and reduce elements along the
|
|
// dimensions given in axis. ReducerFirst is used the first time each output
|
|
// element is written and ReducerNext is used for all subsequent writes.
|
|
template <typename In, typename Out, typename ReducerFirst,
|
|
typename ReducerNext>
|
|
inline bool Reduce(const In* input_data, const int* input_dims,
|
|
const int input_num_dims, const int* axis,
|
|
const int num_axis, Out* output_data,
|
|
const ReducerFirst& reducer_first,
|
|
const ReducerNext& reducer_next) {
|
|
const int parity = (axis[num_axis - 1] == input_num_dims - 1) ? 1 : 0;
|
|
ReduceImpl(input_data, input_dims, output_data, input_num_dims - 1, parity,
|
|
/*next=*/false, reducer_first, reducer_next);
|
|
return true;
|
|
}
|
|
|
|
// Computes the mean or sum of elements across dimensions given in axis.
|
|
// It does so in two stages, first calculates the sum of elements along the axis
|
|
// then divides it by the number of element in axis for quantized values.
|
|
template <typename T, typename U>
|
|
bool QuantizedMeanOrSum(const T* input_data, int32_t input_zero_point,
|
|
float input_scale, const int* input_dims,
|
|
const int input_num_dims, T* output_data,
|
|
int32_t output_zero_point, float output_scale,
|
|
const int* output_dims, const int output_num_dims,
|
|
const int* axis, const int num_axis_dimensions,
|
|
bool keep_dims, int* normalized_dims,
|
|
int* resolved_axis, U* temp_sum, bool compute_sum) {
|
|
const int32_t kMinValue = std::numeric_limits<T>::min();
|
|
const int32_t kMaxValue = std::numeric_limits<T>::max();
|
|
ruy::profiler::ScopeLabel label(compute_sum ? "QuantizedSum"
|
|
: "QuantizedMean");
|
|
// Reset output data.
|
|
size_t num_outputs = 1;
|
|
for (int idx = 0; idx < output_num_dims; ++idx) {
|
|
size_t current = static_cast<size_t>(output_dims[idx]);
|
|
// Overflow prevention.
|
|
if (num_outputs > std::numeric_limits<size_t>::max() / current) {
|
|
return false;
|
|
}
|
|
num_outputs *= current;
|
|
}
|
|
|
|
// Return early when input shape has zero dim. This is done after initializing
|
|
// data for output tensor because there are cases that the input tensor is
|
|
// empty but output tensor is not. In that case, output tensor should be
|
|
// filled with init_value.
|
|
for (int i = 0; i < input_num_dims; ++i) {
|
|
if (input_dims[i] == 0) return true;
|
|
}
|
|
|
|
// Resolve axis.
|
|
int num_resolved_axis = 0;
|
|
int normalized_num_dims = 0;
|
|
if (!reduce_utils::ResolveAxis(input_num_dims, axis, num_axis_dimensions,
|
|
resolved_axis, num_resolved_axis, input_dims,
|
|
normalized_dims, normalized_num_dims)) {
|
|
return false;
|
|
}
|
|
|
|
if (num_resolved_axis == 0) {
|
|
int count = NumElements(input_dims, input_num_dims);
|
|
for (int i = 0; i < count; ++i) {
|
|
temp_sum[i] = U(input_data[i]);
|
|
}
|
|
} else {
|
|
if (!Reduce<T, U, CastSumOp<T, U>, CastSumOp<T, U>>(
|
|
input_data, normalized_dims, normalized_num_dims, resolved_axis,
|
|
num_resolved_axis, temp_sum, CastSumOp<T, U>(),
|
|
CastSumOp<T, U>())) {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
// Calculate mean by dividing output_data by num of aggregated element.
|
|
size_t num_elements_in_axis = 1;
|
|
for (int idx = 0; idx < num_resolved_axis; ++idx) {
|
|
size_t current = static_cast<size_t>(normalized_dims[resolved_axis[idx]]);
|
|
// Overflow prevention.
|
|
if (current > (std::numeric_limits<size_t>::max() / num_elements_in_axis)) {
|
|
return false;
|
|
}
|
|
num_elements_in_axis *= current;
|
|
}
|
|
|
|
if (num_elements_in_axis > 0) {
|
|
const float scale = input_scale / output_scale;
|
|
if (compute_sum) {
|
|
const float bias = -input_zero_point * scale * num_elements_in_axis;
|
|
for (size_t idx = 0; idx < num_outputs; ++idx) {
|
|
U value = static_cast<U>(TfLiteRound(temp_sum[idx] * scale + bias)) +
|
|
output_zero_point;
|
|
value = std::min(std::max(value, kMinValue), kMaxValue);
|
|
output_data[idx] = static_cast<T>(value);
|
|
}
|
|
} else {
|
|
const float bias = -input_zero_point * scale;
|
|
for (size_t idx = 0; idx < num_outputs; ++idx) {
|
|
float float_mean = static_cast<float>(temp_sum[idx]) /
|
|
static_cast<float>(num_elements_in_axis);
|
|
float result = TfLiteMin(
|
|
TfLiteRound(float_mean * scale + bias) + output_zero_point,
|
|
static_cast<float>(std::numeric_limits<T>::max()));
|
|
result = TfLiteMax(result,
|
|
static_cast<float>(std::numeric_limits<T>::min()));
|
|
output_data[idx] = static_cast<T>(result);
|
|
}
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
using ops::builtin::reduce::ReduceType;
|
|
|
|
template <typename T>
|
|
inline bool ReduceDispatcher(const T* input_data, const int* input_dims,
|
|
const int input_num_dims, const int* output_dims,
|
|
int output_num_dims, T* output_data,
|
|
const int* axis, const int64_t num_axis_dimensions,
|
|
ReduceType reduce_type) {
|
|
T init_value;
|
|
switch (reduce_type) {
|
|
case ReduceType::kProd:
|
|
init_value = ProdOp<T>::kNeutralElement;
|
|
break;
|
|
case ReduceType::kSum:
|
|
init_value = SumOp<T>::kNeutralElement;
|
|
break;
|
|
case ReduceType::kMin:
|
|
init_value = MinOp<T>::kNeutralElement;
|
|
break;
|
|
case ReduceType::kMax:
|
|
init_value = MaxOp<T>::kNeutralElement;
|
|
break;
|
|
default:
|
|
return false;
|
|
}
|
|
// Return early when input shape has zero dim. This is done after initializing
|
|
// data for output tensor because there are cases that the input tensor is
|
|
// empty but output tensor is not. In that case, output tensor should be
|
|
// filled with Op::kNeutralElement.
|
|
for (int i = 0; i < input_num_dims; ++i) {
|
|
if (input_dims[i] == 0) {
|
|
return reference_ops::InitTensorDataForReduce(
|
|
output_dims, output_num_dims, init_value, output_data);
|
|
}
|
|
}
|
|
|
|
switch (reduce_type) {
|
|
case ReduceType::kProd:
|
|
return Reduce<T, T, ProdOp<T>, ProdOp<T>>(
|
|
input_data, input_dims, input_num_dims, axis, num_axis_dimensions,
|
|
output_data, ProdOp<T>(), ProdOp<T>());
|
|
case ReduceType::kSum:
|
|
return Reduce<T, T, SumOp<T>, SumOp<T>>(
|
|
input_data, input_dims, input_num_dims, axis, num_axis_dimensions,
|
|
output_data, SumOp<T>(), SumOp<T>());
|
|
case ReduceType::kMin:
|
|
return Reduce<T, T, MinOp<T>, MinOp<T>>(
|
|
input_data, input_dims, input_num_dims, axis, num_axis_dimensions,
|
|
output_data, MinOp<T>(), MinOp<T>());
|
|
case ReduceType::kMax:
|
|
return Reduce<T, T, MaxOp<T>, MaxOp<T>>(
|
|
input_data, input_dims, input_num_dims, axis, num_axis_dimensions,
|
|
output_data, MaxOp<T>(), MaxOp<T>());
|
|
default:
|
|
return false;
|
|
}
|
|
}
|
|
|
|
template <>
|
|
inline bool ReduceDispatcher<bool>(const bool* input_data,
|
|
const int* input_dims,
|
|
const int input_num_dims,
|
|
const int* output_dims, int output_num_dims,
|
|
bool* output_data, const int* axis,
|
|
const int64_t num_axis_dimensions,
|
|
ReduceType reduce_type) {
|
|
bool init_value;
|
|
switch (reduce_type) {
|
|
case ReduceType::kAny:
|
|
init_value = OrOp::kNeutralElement;
|
|
break;
|
|
case ReduceType::kAll:
|
|
init_value = AndOp::kNeutralElement;
|
|
break;
|
|
default:
|
|
return false;
|
|
}
|
|
// Return early when input shape has zero dim. This is done after initializing
|
|
// data for output tensor because there are cases that the input tensor is
|
|
// empty but output tensor is not. In that case, output tensor should be
|
|
// filled with Op::kNeutralElement.
|
|
for (int i = 0; i < input_num_dims; ++i) {
|
|
if (input_dims[i] == 0) {
|
|
return reference_ops::InitTensorDataForReduce(
|
|
output_dims, output_num_dims, init_value, output_data);
|
|
}
|
|
}
|
|
switch (reduce_type) {
|
|
case ReduceType::kAll:
|
|
return Reduce<bool, bool, AndOp, AndOp>(
|
|
input_data, input_dims, input_num_dims, axis, num_axis_dimensions,
|
|
output_data, AndOp(), AndOp());
|
|
case ReduceType::kAny:
|
|
return Reduce<bool, bool, OrOp, OrOp>(
|
|
input_data, input_dims, input_num_dims, axis, num_axis_dimensions,
|
|
output_data, OrOp(), OrOp());
|
|
default:
|
|
return false;
|
|
}
|
|
}
|
|
|
|
// Calculate the reduced product by rescaling each multiplication step to
|
|
// avoid an overflow.
|
|
template <typename T>
|
|
struct ReducerFirst {
|
|
explicit ReducerFirst(int input_zero_point_arg)
|
|
: input_zero_point(input_zero_point_arg) {}
|
|
int32_t operator()(T in) const { return in - input_zero_point; }
|
|
int input_zero_point;
|
|
};
|
|
|
|
template <typename T>
|
|
struct ReducerNext {
|
|
ReducerNext(int32_t input_zero_point_arg, int32_t scaling_multiplier_arg,
|
|
int32_t scaling_shift_arg)
|
|
: input_zero_point(input_zero_point_arg),
|
|
scaling_multiplier(scaling_multiplier_arg),
|
|
scaling_shift(scaling_shift_arg) {}
|
|
int32_t operator()(int32_t current, T in) const {
|
|
const int64_t result =
|
|
static_cast<int64_t>(current) * (in - input_zero_point);
|
|
return MultiplyByQuantizedMultiplier(result, scaling_multiplier,
|
|
scaling_shift);
|
|
}
|
|
int32_t input_zero_point, scaling_multiplier, scaling_shift;
|
|
};
|
|
|
|
template <typename T>
|
|
inline bool QuantizedReduceProd(
|
|
const T* input_data, int32_t input_zero_point,
|
|
const RuntimeShape& input_shape, T* output_data, int32_t output_zero_point,
|
|
const RuntimeShape& output_shape, const int* axis,
|
|
const int64_t num_axis_dimensions, int* resolved_axis, int* normalized_dims,
|
|
int32_t* temp_prod, int32_t scaling_multiplier, int scaling_shift) {
|
|
const int32_t kMinValue = std::numeric_limits<T>::min();
|
|
const int32_t kMaxValue = std::numeric_limits<T>::max();
|
|
|
|
// Resolve axis.
|
|
int num_resolved_axis = 0;
|
|
int normalized_num_dims = 0;
|
|
if (!reduce_utils::ResolveAxis(input_shape.DimensionsCount(), axis,
|
|
num_axis_dimensions, resolved_axis,
|
|
num_resolved_axis, input_shape.DimsData(),
|
|
normalized_dims, normalized_num_dims)) {
|
|
return false;
|
|
}
|
|
|
|
if (!Reduce<T, int32_t, ReducerFirst<T>, ReducerNext<T>>(
|
|
input_data, normalized_dims, normalized_num_dims, resolved_axis,
|
|
num_resolved_axis, temp_prod, ReducerFirst<T>(input_zero_point),
|
|
ReducerNext<T>(input_zero_point, scaling_multiplier,
|
|
scaling_shift))) {
|
|
return false;
|
|
}
|
|
|
|
for (int i = 0; i < output_shape.FlatSize(); i++) {
|
|
int32_t result =
|
|
MultiplyByQuantizedMultiplier(static_cast<int64_t>(temp_prod[i]),
|
|
scaling_multiplier, scaling_shift) +
|
|
output_zero_point;
|
|
result = std::min(std::max(result, kMinValue), kMaxValue);
|
|
output_data[i] = static_cast<T>(result);
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
template <typename T>
|
|
inline void Mean(const tflite::MeanParams& op_params,
|
|
const RuntimeShape& input_shape, const T* input_data,
|
|
const RuntimeShape& output_shape, T* output_data) {
|
|
return reference_ops::Mean(op_params, input_shape, input_data, output_shape,
|
|
output_data);
|
|
}
|
|
|
|
// Computes the mean of elements across dimensions given in axis.
|
|
// It does so in two stages, first calculates the sum of elements along the axis
|
|
// then divides it by the number of element in axis.
|
|
template <typename T, typename U>
|
|
inline bool MeanGeneral(const T* input_data, const int* input_dims,
|
|
const int input_num_dims, T* output_data,
|
|
const int* output_dims, const int output_num_dims,
|
|
const int* axis, const int num_axis_dimensions,
|
|
bool keep_dims, int* normalized_dims,
|
|
int* resolved_axis, U* temp_sum) {
|
|
ruy::profiler::ScopeLabel label("Mean");
|
|
// Resolve axis.
|
|
int num_resolved_axis = 0;
|
|
int normalized_num_dims = 0;
|
|
if (!reduce_utils::ResolveAxis(input_num_dims, axis, num_axis_dimensions,
|
|
resolved_axis, num_resolved_axis, input_dims,
|
|
normalized_dims, normalized_num_dims)) {
|
|
return false;
|
|
}
|
|
if (num_resolved_axis == 0) {
|
|
optimized_ops::ReduceIsCopy(input_data, input_dims, input_num_dims,
|
|
output_data);
|
|
return true;
|
|
}
|
|
// Reset output data.
|
|
size_t num_outputs = 1;
|
|
for (int idx = 0; idx < output_num_dims; ++idx) {
|
|
size_t current = static_cast<size_t>(output_dims[idx]);
|
|
// Overflow prevention.
|
|
if (num_outputs > std::numeric_limits<size_t>::max() / current) {
|
|
return false;
|
|
}
|
|
num_outputs *= current;
|
|
}
|
|
|
|
if (!Reduce<T, U, CastSumOp<T, U>, CastSumOp<T, U>>(
|
|
input_data, normalized_dims, normalized_num_dims, resolved_axis,
|
|
num_resolved_axis, temp_sum, CastSumOp<T, U>(), CastSumOp<T, U>())) {
|
|
return false;
|
|
}
|
|
|
|
// Calculate mean by dividing output_data by num of aggregated element.
|
|
size_t num_elements_in_axis = 1;
|
|
for (int idx = 0; idx < num_resolved_axis; ++idx) {
|
|
size_t current = static_cast<size_t>(normalized_dims[resolved_axis[idx]]);
|
|
// Overflow prevention.
|
|
if (current > (std::numeric_limits<size_t>::max() / num_elements_in_axis)) {
|
|
return false;
|
|
}
|
|
num_elements_in_axis *= current;
|
|
}
|
|
|
|
if (num_elements_in_axis > 0) {
|
|
for (size_t idx = 0; idx < num_outputs; ++idx) {
|
|
output_data[idx] =
|
|
static_cast<T>(temp_sum[idx] / static_cast<U>(num_elements_in_axis));
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
template <typename T, typename U>
|
|
inline bool Mean(const T* input_data, const int* input_dims,
|
|
const int input_num_dims, T* output_data,
|
|
const int* output_dims, const int output_num_dims,
|
|
const int* axis, const int num_axis_dimensions, bool keep_dims,
|
|
int* normalized_dims, int* resolved_axis, U* temp_sum) {
|
|
return MeanGeneral(input_data, input_dims, input_num_dims, output_data,
|
|
output_dims, output_num_dims, axis, num_axis_dimensions,
|
|
false, normalized_dims, resolved_axis, temp_sum);
|
|
}
|
|
|
|
// Use Eigen when Mean is calculated over the last dimension only of a float
|
|
// tensor.
|
|
template <>
|
|
inline bool Mean<float, float>(const float* input_data, const int* input_dims,
|
|
const int input_num_dims, float* output_data,
|
|
const int* output_dims,
|
|
const int output_num_dims, const int* axis,
|
|
const int num_axis_dimensions, bool keep_dims,
|
|
int* normalized_dims, int* resolved_axis,
|
|
float* temp_sum) {
|
|
// Handle reduce_mean for the last dimensions.
|
|
int num_resolved_axis = 0;
|
|
int normalized_num_dims = 0;
|
|
if (!reduce_utils::ResolveAxis(input_num_dims, axis, num_axis_dimensions,
|
|
resolved_axis, num_resolved_axis, input_dims,
|
|
normalized_dims, normalized_num_dims)) {
|
|
return false;
|
|
}
|
|
if (normalized_num_dims > 1 && num_resolved_axis == 1 &&
|
|
resolved_axis[0] == (normalized_num_dims - 1)) {
|
|
ruy::profiler::ScopeLabel label("MeanLastDim/Float");
|
|
int output_size = normalized_dims[0];
|
|
const int last_input_dim = normalized_dims[1];
|
|
|
|
// TODO(b/152563685): Consider use eigen to cover more general cases.
|
|
const MatrixMap<const float> in_mat(input_data, last_input_dim,
|
|
output_size);
|
|
VectorMap<float> out(output_data, output_size, 1);
|
|
out = (in_mat.array().colwise().sum()) / static_cast<float>(last_input_dim);
|
|
return true;
|
|
}
|
|
|
|
return MeanGeneral(input_data, input_dims, input_num_dims, output_data,
|
|
output_dims, output_num_dims, axis, num_axis_dimensions,
|
|
false, normalized_dims, resolved_axis, temp_sum);
|
|
}
|
|
|
|
// Computes the generic value (i.e., sum/max/min/prod) of elements across
|
|
// dimensions given in axis. It needs to pass in init_value and reducer.
|
|
template <typename T>
|
|
inline bool ReduceGeneric(const T* input_data, const int* input_dims,
|
|
const int input_num_dims, T* output_data,
|
|
const int* output_dims, const int output_num_dims,
|
|
const int* axis, const int64_t num_axis_dimensions,
|
|
int* resolved_axis, int* normalized_dims,
|
|
ReduceType reduce_type) {
|
|
int num_resolved_axis = 0;
|
|
int normalized_num_dims = 0;
|
|
if (!reduce_utils::ResolveAxis(input_num_dims, axis, num_axis_dimensions,
|
|
resolved_axis, num_resolved_axis, input_dims,
|
|
normalized_dims, normalized_num_dims)) {
|
|
return false;
|
|
}
|
|
if (num_resolved_axis == 0) {
|
|
optimized_ops::ReduceIsCopy(input_data, input_dims, input_num_dims,
|
|
output_data);
|
|
return true;
|
|
}
|
|
return ReduceDispatcher(input_data, normalized_dims, normalized_num_dims,
|
|
output_dims, output_num_dims, output_data,
|
|
resolved_axis, num_resolved_axis, reduce_type);
|
|
}
|
|
|
|
} // namespace optimized_ops
|
|
} // namespace tflite
|
|
|
|
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_REDUCE_H_
|