| 990 | }; |
| 991 | |
| 992 | inline void Mean(const tflite::MeanParams& op_params, |
| 993 | const RuntimeShape& unextended_input_shape, |
| 994 | const uint8_t* input_data, int32 input_zero_point, |
| 995 | float input_scale, const RuntimeShape& unextended_output_shape, |
| 996 | uint8_t* output_data, int32 output_zero_point, |
| 997 | float output_scale, CpuBackendContext* cpu_backend_context) { |
| 998 | gemmlowp::ScopedProfilingLabel label("Mean4D/Uint8"); |
| 999 | // Current implementation only supports dimension equals 4 and simultaneous |
| 1000 | // reduction over width and height. |
| 1001 | TFLITE_CHECK_EQ(unextended_input_shape.DimensionsCount(), 4); |
| 1002 | TFLITE_CHECK_LE(unextended_output_shape.DimensionsCount(), 4); |
| 1003 | const RuntimeShape input_shape = |
| 1004 | RuntimeShape::ExtendedShape(4, unextended_input_shape); |
| 1005 | const RuntimeShape output_shape = |
| 1006 | RuntimeShape::ExtendedShape(4, unextended_output_shape); |
| 1007 | const int output_height = output_shape.Dims(1); |
| 1008 | const int output_width = output_shape.Dims(2); |
| 1009 | const int output_depth = output_shape.Dims(3); |
| 1010 | |
| 1011 | TFLITE_DCHECK_EQ(op_params.axis_count, 2); |
| 1012 | TFLITE_DCHECK((op_params.axis[0] == 1 && op_params.axis[1] == 2) || |
| 1013 | (op_params.axis[0] == 2 && op_params.axis[1] == 1)); |
| 1014 | TFLITE_DCHECK_EQ(output_height, 1); |
| 1015 | TFLITE_DCHECK_EQ(output_width, 1); |
| 1016 | |
| 1017 | constexpr int kMinDepthPerThread = 8; |
| 1018 | int thread_count = output_depth / kMinDepthPerThread; |
| 1019 | thread_count = thread_count > 0 ? thread_count : 1; |
| 1020 | const int capped_thread_count = |
| 1021 | std::min(thread_count, cpu_backend_context->max_num_threads()); |
| 1022 | |
| 1023 | if (capped_thread_count == 1) { |
| 1024 | MeanImpl(op_params, input_shape, input_data, input_zero_point, input_scale, |
| 1025 | output_shape, output_data, output_zero_point, output_scale, 0, |
| 1026 | output_depth); |
| 1027 | } else { |
| 1028 | // Instead parrallel for batch, we loop for the output_depth since batch |
| 1029 | // is typical 1. |
| 1030 | std::vector<MeanWorkerTask> tasks; |
| 1031 | // TODO(b/131746020) don't create new heap allocations every time. |
| 1032 | // At least we make it a single heap allocation by using reserve(). |
| 1033 | tasks.reserve(capped_thread_count); |
| 1034 | int depth_start = 0; |
| 1035 | for (int i = 0; i < capped_thread_count; ++i) { |
| 1036 | // Try to distribute the tasks as even as possible. |
| 1037 | int depth_end = depth_start + |
| 1038 | (output_depth - depth_start) / (capped_thread_count - i); |
| 1039 | tasks.emplace_back(op_params, input_shape, input_data, input_zero_point, |
| 1040 | input_scale, output_shape, output_data, |
| 1041 | output_zero_point, output_scale, depth_start, |
| 1042 | depth_end); |
| 1043 | depth_start = depth_end; |
| 1044 | } |
| 1045 | cpu_backend_threadpool::Execute(tasks.size(), tasks.data(), |
| 1046 | cpu_backend_context); |
| 1047 | } |
| 1048 | } |
| 1049 |
no test coverage detected