From 1821835c6962e409a554640f5e26d01042fcd065 Mon Sep 17 00:00:00 2001 From: TFLM-bot Date: Sun, 2 Aug 2026 14:29:42 +0000 Subject: [PATCH] Sync from upstream TF. --- .../lite/core/api/flatbuffer_conversions.cc | 58 ++++--- tensorflow/lite/core/c/common.cc | 11 +- tensorflow/lite/core/macros.h | 11 ++ tensorflow/lite/kernels/internal/common.cc | 15 ++ tensorflow/lite/kernels/internal/common.h | 62 +++++++ .../lite/kernels/internal/quantization_util.h | 49 +++--- .../lite/kernels/internal/reference/add.h | 6 +- .../kernels/internal/reference/arg_min_max.h | 29 +++- .../lite/kernels/internal/reference/cumsum.h | 6 +- .../lite/kernels/internal/reference/mul.h | 6 +- .../lite/kernels/internal/reference/reduce.h | 98 ++++++----- .../reference/resize_nearest_neighbor.h | 8 +- .../lite/kernels/internal/reference/select.h | 8 +- tensorflow/lite/kernels/kernel_util.cc | 12 +- tensorflow/lite/kernels/kernel_util.h | 28 ++-- tensorflow/lite/kernels/padding.h | 158 ++++++++++++++---- .../lite/tools/flatbuffer_utils_test.py | 6 +- tensorflow/lite/tools/test_utils.py | 2 +- tensorflow/lite/tools/visualize_test.py | 4 +- 19 files changed, 409 insertions(+), 168 deletions(-) diff --git a/tensorflow/lite/core/api/flatbuffer_conversions.cc b/tensorflow/lite/core/api/flatbuffer_conversions.cc index 6790a92df44..d47767c79b0 100644 --- a/tensorflow/lite/core/api/flatbuffer_conversions.cc +++ b/tensorflow/lite/core/api/flatbuffer_conversions.cc @@ -2330,34 +2330,42 @@ TfLiteStatus ParseStablehloGather(const Operator* op, op->builtin_options_2_as_StablehloGatherOptions(); if (schema_params != nullptr) { - TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( - /*max_size_of_buffer=*/schema_params->offset_dims()->size() * - sizeof(int64_t), - /*flat_vector=*/schema_params->offset_dims(), - /*buffer=*/params->offset_dims, /*error_reporter=*/error_reporter, - /*op_name=*/"stablehlo_gather")); - params->num_offset_dims = schema_params->offset_dims()->size(); - - TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->collapsed_slice_dims()->size() * sizeof(int64_t), - schema_params->collapsed_slice_dims(), params->collapsed_slice_dims, - error_reporter, "stablehlo_gather")); - params->num_collapsed_slice_dims = - schema_params->collapsed_slice_dims()->size(); - - TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->start_index_map()->size() * sizeof(int64_t), - schema_params->start_index_map(), params->start_index_map, - error_reporter, "stablehlo_gather")); - params->num_start_index_map = schema_params->start_index_map()->size(); + if (schema_params->offset_dims()) { + TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( + /*max_size_of_buffer=*/schema_params->offset_dims()->size() * + sizeof(int64_t), + /*flat_vector=*/schema_params->offset_dims(), + /*buffer=*/params->offset_dims, /*error_reporter=*/error_reporter, + /*op_name=*/"stablehlo_gather")); + params->num_offset_dims = schema_params->offset_dims()->size(); + } + + if (schema_params->collapsed_slice_dims()) { + TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( + schema_params->collapsed_slice_dims()->size() * sizeof(int64_t), + schema_params->collapsed_slice_dims(), params->collapsed_slice_dims, + error_reporter, "stablehlo_gather")); + params->num_collapsed_slice_dims = + schema_params->collapsed_slice_dims()->size(); + } + + if (schema_params->start_index_map()) { + TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( + schema_params->start_index_map()->size() * sizeof(int64_t), + schema_params->start_index_map(), params->start_index_map, + error_reporter, "stablehlo_gather")); + params->num_start_index_map = schema_params->start_index_map()->size(); + } params->index_vector_dim = schema_params->index_vector_dim(); - TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->slice_sizes()->size() * sizeof(int64_t), - schema_params->slice_sizes(), params->slice_sizes, error_reporter, - "stablehlo_gather")); - params->num_slice_sizes = schema_params->slice_sizes()->size(); + if (schema_params->slice_sizes()) { + TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( + schema_params->slice_sizes()->size() * sizeof(int64_t), + schema_params->slice_sizes(), params->slice_sizes, error_reporter, + "stablehlo_gather")); + params->num_slice_sizes = schema_params->slice_sizes()->size(); + } params->indices_are_sorted = schema_params->indices_are_sorted(); } else { diff --git a/tensorflow/lite/core/c/common.cc b/tensorflow/lite/core/c/common.cc index 66e3d4eedb2..19996dba56b 100644 --- a/tensorflow/lite/core/c/common.cc +++ b/tensorflow/lite/core/c/common.cc @@ -20,6 +20,7 @@ limitations under the License. #endif // TF_LITE_STATIC_MEMORY #include +#include #include #include #include @@ -460,15 +461,23 @@ TfLiteStatus TfLiteTensorCopy(const TfLiteTensor* src, TfLiteTensor* dst) { TfLiteStatus TfLiteTensorResizeMaybeCopy(size_t num_bytes, TfLiteTensor* tensor, bool preserve_data) { + if (tensor == nullptr) { + return kTfLiteError; + } if (tensor->allocation_type != kTfLiteDynamic && tensor->allocation_type != kTfLitePersistentRo) { return kTfLiteOk; } + // Guard against integer overflow: num_bytes + XNN_EXTRA_BYTES must not wrap. + constexpr size_t kXnnExtraBytes = 16; + if (num_bytes > std::numeric_limits::max() - kXnnExtraBytes) { + return kTfLiteError; + } #ifdef TF_LITE_TENSORFLOW_PROFILER tflite::PauseHeapMonitoring(/*pause=*/true); #endif // This buffer may be consumed by XNNPack. - size_t alloc_bytes = num_bytes + /*XNN_EXTRA_BYTES=*/16; + size_t alloc_bytes = num_bytes + kXnnExtraBytes; // TODO(b/145340303): Tensor data should be aligned. if (!tensor->data.data) { tensor->data.data = (char*)malloc(alloc_bytes); diff --git a/tensorflow/lite/core/macros.h b/tensorflow/lite/core/macros.h index 86de4daefe7..c2ecce91ae4 100644 --- a/tensorflow/lite/core/macros.h +++ b/tensorflow/lite/core/macros.h @@ -65,4 +65,15 @@ limitations under the License. #define TFLITE_HAS_ATTRIBUTE_WEAK 0 #endif +// Disables UBSan's integer-overflow checks for an annotated function. Only use +// this where the overflow merely yields a wrong numeric result and poses no +// safety risk. Expands to nothing on compilers without the attribute. +#if TFLITE_HAS_ATTRIBUTE(no_sanitize) +#define TFLITE_NO_SANITIZE_INTEGER_OVERFLOW \ + __attribute__(( \ + no_sanitize("signed-integer-overflow", "unsigned-integer-overflow"))) +#else +#define TFLITE_NO_SANITIZE_INTEGER_OVERFLOW +#endif + #endif // TENSORFLOW_LITE_CORE_MACROS_H_ diff --git a/tensorflow/lite/kernels/internal/common.cc b/tensorflow/lite/kernels/internal/common.cc index fabb0208b7d..73d91dea8bc 100644 --- a/tensorflow/lite/kernels/internal/common.cc +++ b/tensorflow/lite/kernels/internal/common.cc @@ -15,10 +15,22 @@ limitations under the License. #include "tensorflow/lite/kernels/internal/common.h" +#include "tensorflow/lite/core/macros.h" + namespace tflite { +// Note on TFLITE_NO_SANITIZE_INTEGER_OVERFLOW below: +// +// These MultiplyByQuantizedMultiplier overloads do not intentionally wrap, so +// they deliberately do NOT use WrappingMul/WrappingAdd. Under their documented +// input contract (asserted below) the arithmetic does not overflow; the +// attribute only silences UBSan reports at the contract's boundaries. Using +// wrapping here would be wrong -- it would mask a contract violation (a real +// bug) instead of letting it surface. + // Single-rounding MultiplyByQuantizedMultiplier #if TFLITE_SINGLE_ROUNDING +TFLITE_NO_SANITIZE_INTEGER_OVERFLOW int32_t MultiplyByQuantizedMultiplier(int32_t x, int32_t quantized_multiplier, int shift) { TFLITE_DCHECK(quantized_multiplier >= 0); @@ -34,6 +46,7 @@ int32_t MultiplyByQuantizedMultiplier(int32_t x, int32_t quantized_multiplier, return static_cast(result); } +TFLITE_NO_SANITIZE_INTEGER_OVERFLOW int32_t MultiplyByQuantizedMultiplier(int64_t x, int32_t quantized_multiplier, int shift) { // Inputs: @@ -64,6 +77,7 @@ int32_t MultiplyByQuantizedMultiplier(int64_t x, int32_t quantized_multiplier, } // Double-rounding MultiplyByQuantizedMultiplier #else +TFLITE_NO_SANITIZE_INTEGER_OVERFLOW int32_t MultiplyByQuantizedMultiplier(int32_t x, int32_t quantized_multiplier, int shift) { using gemmlowp::RoundingDivideByPOT; @@ -75,6 +89,7 @@ int32_t MultiplyByQuantizedMultiplier(int32_t x, int32_t quantized_multiplier, right_shift); } +TFLITE_NO_SANITIZE_INTEGER_OVERFLOW int32_t MultiplyByQuantizedMultiplier(int64_t x, int32_t quantized_multiplier, int shift) { // Inputs: diff --git a/tensorflow/lite/kernels/internal/common.h b/tensorflow/lite/kernels/internal/common.h index aa5248470fe..2eeca0ffcf9 100644 --- a/tensorflow/lite/kernels/internal/common.h +++ b/tensorflow/lite/kernels/internal/common.h @@ -40,6 +40,68 @@ namespace tflite { constexpr int kReverseShift = -1; +// Well-defined wrapping integer arithmetic helpers. +// +// Several elementwise/reduction kernels intentionally allow their value +// arithmetic to wrap around for narrow integer types (the result is later +// clamped, or wrapping is the documented behavior). Performing that wrap +// directly on signed types is Undefined Behavior in C++, which the optimizer is +// free to miscompile. These helpers instead perform the arithmetic in the +// corresponding unsigned type -- where wrapping is fully defined -- and convert +// the result back, so there is no UB even in optimized release builds. +// +// For floating-point T these are plain a+b / a*b (no overflow concern); the +// unsigned path is only taken for integral types. +// +// Note on integral promotion: for integer types narrower than int (e.g. +// int16_t), the corresponding unsigned type (uint16_t) is still promoted to the +// signed `int` before the arithmetic is performed. That promotion would +// reintroduce signed overflow UB (e.g. 65535 * 65535 overflows int). To avoid +// it we widen the operands to an unsigned type at least as wide as unsigned int +// so the arithmetic itself stays unsigned, then truncate back down to the +// narrow type -- both steps are well-defined. +template +inline T WrappingAdd(T a, T b) { + if constexpr (std::is_integral_v && !std::is_same_v) { + using U = std::make_unsigned_t; + using P = std::common_type_t; + return static_cast(static_cast(static_cast

(static_cast(a)) + + static_cast

(static_cast(b)))); + } else { + return a + b; + } +} + +template +inline T WrappingMul(T a, T b) { + if constexpr (std::is_integral_v && !std::is_same_v) { + using U = std::make_unsigned_t; + using P = std::common_type_t; + return static_cast(static_cast(static_cast

(static_cast(a)) * + static_cast

(static_cast(b)))); + } else { + return a * b; + } +} + +// Tensor value arithmetic helpers for ops where integer overflow is documented +// numeric behavior and not used for memory allocation, indexing, shape +// computation, or control-flow bounds. These intentionally preserve the direct +// arithmetic expression and only suppress UBSan's integer-overflow diagnostic. +// Do not use these helpers for shape, element count, allocation, or pointer +// offset arithmetic; those paths need checked integer helpers instead. +template +TFLITE_NO_SANITIZE_INTEGER_OVERFLOW inline T +AddTensorValuesWithExpectedOverflow(T a, T b) { + return a + b; +} + +template +TFLITE_NO_SANITIZE_INTEGER_OVERFLOW inline T +MulTensorValuesWithExpectedOverflow(T a, T b) { + return a * b; +} + // Reduces and compresses dimensions so that broadcast handling becomes more // efficient. Returns true if the output shape is broadcastable; it doesn't // contain any degenerate dimension, i.e. shape dimension = 0. False otherwise. diff --git a/tensorflow/lite/kernels/internal/quantization_util.h b/tensorflow/lite/kernels/internal/quantization_util.h index eb4e84013e1..e960697956f 100644 --- a/tensorflow/lite/kernels/internal/quantization_util.h +++ b/tensorflow/lite/kernels/internal/quantization_util.h @@ -111,7 +111,7 @@ QuantizationParams ChooseQuantizationParams(double rmin, double rmax) { // static_cast would cause undefined behavior for the following cases, which // have well-defined behavior for this function: // -// 1. If x is NaN, the result is zero. +// 1. If x is NaN, the result is nan_result, which defaults to zero. // // 2. If the truncated form of x is above the representable range of IntOut, // the result is std::numeric_limits::max(). @@ -125,47 +125,36 @@ QuantizationParams ChooseQuantizationParams(double rmin, double rmax) { // the results are undefined. // TODO(sfeuz): Replace by absl::SafeCast once available. template -IntOut SafeCast(FloatIn x) { +IntOut SafeCast(FloatIn x, IntOut nan_result = IntOut{0}) { static_assert(!std::numeric_limits::is_integer, "FloatIn is integer"); static_assert(std::numeric_limits::is_integer, "IntOut is not integer"); static_assert(std::numeric_limits::radix == 2, "IntOut is base 2"); - - // Special case NaN, for which the logic below doesn't work. + static_assert(std::numeric_limits::max_exponent > + std::numeric_limits::digits, + "FloatIn cannot represent IntOut's exclusive upper bound"); if (std::isnan(x)) { - return 0; - } - - // Negative values all clip to zero for unsigned results. - if (!std::numeric_limits::is_signed && x < 0) { - return 0; + return nan_result; } - // Handle infinities. - if (std::isinf(x)) { - return x < 0 ? std::numeric_limits::min() - : std::numeric_limits::max(); + const FloatIn min_value = + static_cast(std::numeric_limits::min()); + if (x < min_value) { + return std::numeric_limits::min(); } - // Set exp such that x == f * 2^exp for some f with |f| in [0.5, 1.0), - // unless x is zero in which case exp == 0. Note that this implies that the - // magnitude of x is strictly less than 2^exp. - int exp = 0; - std::frexp(x, &exp); - - // Let N be the number of non-sign bits in the representation of IntOut. If - // the magnitude of x is strictly less than 2^N, the truncated version of x - // is representable as IntOut. The only representable integer for which this - // is not the case is kMin for signed types (i.e. -2^N), but that is covered - // by the fall-through below. - if (exp <= std::numeric_limits::digits) { - return x; + // IntOut's exclusive upper bound is one greater than its maximum. Compute + // it from a representable power of two so that no rounded integer-to-float + // conversion can make an out-of-range value appear safe. + constexpr FloatIn kMaxExclusive = + static_cast(std::numeric_limits::max() / 2 + 1) * + FloatIn{2}; + if (x >= kMaxExclusive) { + return std::numeric_limits::max(); } - // Handle numbers with magnitude >= 2^N. - return x < 0 ? std::numeric_limits::min() - : std::numeric_limits::max(); + return static_cast(x); } // LINT.ThenChange(//tensorflow/compiler/mlir/lite/kernels/internal/quantization_util.h) diff --git a/tensorflow/lite/kernels/internal/reference/add.h b/tensorflow/lite/kernels/internal/reference/add.h index 395a623a028..198d576c486 100644 --- a/tensorflow/lite/kernels/internal/reference/add.h +++ b/tensorflow/lite/kernels/internal/reference/add.h @@ -204,7 +204,8 @@ inline void AddBroadcast(const T* input_data, const T* broadcast_data, T activation_max) { for (size_t c = 0; c < size; ++c) { output_data[c] = ActivationFunctionWithMinMax( - input_data[c] + broadcast_data[0], activation_min, activation_max); + WrappingAdd(input_data[c], broadcast_data[0]), activation_min, + activation_max); } } @@ -229,7 +230,8 @@ inline void AddBroadcast(const int32_t* input_data, #endif for (; c < size; ++c) { output_data[c] = ActivationFunctionWithMinMax( - input_data[c] + broadcast_data[0], activation_min, activation_max); + WrappingAdd(input_data[c], broadcast_data[0]), activation_min, + activation_max); } } diff --git a/tensorflow/lite/kernels/internal/reference/arg_min_max.h b/tensorflow/lite/kernels/internal/reference/arg_min_max.h index 8154fbf71e3..6664c7428a4 100644 --- a/tensorflow/lite/kernels/internal/reference/arg_min_max.h +++ b/tensorflow/lite/kernels/internal/reference/arg_min_max.h @@ -15,7 +15,9 @@ limitations under the License. #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_ARG_MIN_MAX_H_ #define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_ARG_MIN_MAX_H_ +#include #include +#include #include "tensorflow/lite/kernels/internal/types.h" @@ -23,8 +25,11 @@ namespace tflite { namespace reference_ops { +// Default comparator for non-floating-point types (no NaN possible). template -std::function GetComparefunction(bool is_arg_max) { +typename std::enable_if::value, + std::function>::type +GetComparefunction(bool is_arg_max) { if (is_arg_max) { return std::greater(); } else { @@ -32,6 +37,28 @@ std::function GetComparefunction(bool is_arg_max) { } } +// NaN-aware comparator for floating-point types. +// Matches TensorFlow eager semantics: NaN is treated as "less than any finite +// value" for ArgMax and "greater than any finite value" for ArgMin. A NaN +// candidate never replaces anything; a finite candidate always replaces a NaN +// accumulator. For all-NaN inputs the first index is returned. +template +typename std::enable_if::value, + std::function>::type +GetComparefunction(bool is_arg_max) { + if (is_arg_max) { + return [](T candidate, T current) { + return !std::isnan(candidate) && + (std::isnan(current) || candidate > current); + }; + } else { + return [](T candidate, T current) { + return !std::isnan(candidate) && + (std::isnan(current) || candidate < current); + }; + } +} + template void ArgMinMax(const RuntimeShape& input1_shape, const T1* input1_data, const T3* input2_data, const RuntimeShape& output_shape, diff --git a/tensorflow/lite/kernels/internal/reference/cumsum.h b/tensorflow/lite/kernels/internal/reference/cumsum.h index 7cbc87c0883..abf5204229b 100644 --- a/tensorflow/lite/kernels/internal/reference/cumsum.h +++ b/tensorflow/lite/kernels/internal/reference/cumsum.h @@ -71,9 +71,11 @@ inline void CumSum(const T* input_data, const RuntimeShape& shape, int32_t axis, if (exclusive) { output_data[index] = accumulator; - accumulator += input_data[index]; + accumulator = AddTensorValuesWithExpectedOverflow(accumulator, + input_data[index]); } else { - accumulator += input_data[index]; + accumulator = AddTensorValuesWithExpectedOverflow(accumulator, + input_data[index]); output_data[index] = accumulator; } } diff --git a/tensorflow/lite/kernels/internal/reference/mul.h b/tensorflow/lite/kernels/internal/reference/mul.h index ea66dae226f..56db0166f15 100644 --- a/tensorflow/lite/kernels/internal/reference/mul.h +++ b/tensorflow/lite/kernels/internal/reference/mul.h @@ -61,7 +61,7 @@ inline void Mul(const ArithmeticParams& params, MatchingExtendedShapeFlatSize(input1_shape, input2_shape, output_shape); for (int i = 0; i < flat_size; ++i) { output_data[i] = ActivationFunctionWithMinMax( - input1_data[i] * input2_data[i], output_activation_min, + WrappingMul(input1_data[i], input2_data[i]), output_activation_min, output_activation_max); } } @@ -132,8 +132,8 @@ BroadcastMul6DSlow(const ArithmeticParams& params, T output_activation_max; GetActivationParams(params, &output_activation_min, &output_activation_max); auto op = [output_activation_min, output_activation_max](T a, T b) { - return ActivationFunctionWithMinMax(a * b, output_activation_min, - output_activation_max); + return ActivationFunctionWithMinMax( + WrappingMul(a, b), output_activation_min, output_activation_max); }; BroadcastBinaryOpSimple(unextended_input1_shape, input1_data, unextended_input2_shape, input2_data, diff --git a/tensorflow/lite/kernels/internal/reference/reduce.h b/tensorflow/lite/kernels/internal/reference/reduce.h index 5b795ea8ff0..ff6281b9194 100644 --- a/tensorflow/lite/kernels/internal/reference/reduce.h +++ b/tensorflow/lite/kernels/internal/reference/reduce.h @@ -15,14 +15,20 @@ limitations under the License. #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_REDUCE_H_ #define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_REDUCE_H_ +#include + #include +#include +#include +#include "absl/types/span.h" #include "ruy/profiler/instrumentation.h" // from @ruy #include "tensorflow/lite/kernels/internal/common.h" #include "tensorflow/lite/kernels/internal/cppmath.h" #include "tensorflow/lite/kernels/internal/max.h" #include "tensorflow/lite/kernels/internal/min.h" #include "tensorflow/lite/kernels/internal/quantization_util.h" +#include "tensorflow/lite/kernels/internal/reduce_common.h" #include "tensorflow/lite/kernels/internal/types.h" // Check if the reduction at index is the first one along the dimensions given @@ -148,7 +154,7 @@ inline bool ReduceSumImpl(const In* input_data, const int* input_dims, Out* output_data) { auto reducer = [](const Out current, const In in) -> Out { const Out actual_in = static_cast(in); - return current + actual_in; + return AddTensorValuesWithExpectedOverflow(current, actual_in); }; return Reduce(input_data, input_dims, output_dims, input_num_dims, output_num_dims, axis, num_axis, input_iter, reducer, @@ -158,15 +164,14 @@ inline bool ReduceSumImpl(const In* input_data, const int* input_dims, template inline bool InitTensorDataForReduce(const int* dims, const int num_dims, const T init_value, T* data) { - size_t num_elements = 1; - for (int idx = 0; idx < num_dims; ++idx) { - size_t current = static_cast(dims[idx]); - // Overflow prevention. - if (current > 0 && - num_elements > std::numeric_limits::max() / current) { - return false; - } - num_elements *= current; + if (num_dims < 0 || (dims == nullptr && num_dims != 0)) { + return false; + } + size_t num_elements = 0; + if (!reduce_utils::CheckedElementCount( + absl::MakeConstSpan(dims, static_cast(num_dims)), + num_elements)) { + return false; } for (size_t idx = 0; idx < num_elements; ++idx) { data[idx] = init_value; @@ -221,14 +226,15 @@ inline bool Mean(const T* input_data, const int* input_dims, int* temp_index, int* resolved_axis, U* temp_sum) { ruy::profiler::ScopeLabel label("Mean"); // Reset output data. - size_t num_outputs = 1; - for (int idx = 0; idx < output_num_dims; ++idx) { - size_t current = static_cast(output_dims[idx]); - // Overflow prevention. - if (num_outputs > std::numeric_limits::max() / current) { - return false; - } - num_outputs *= current; + if (output_num_dims < 0 || (output_dims == nullptr && output_num_dims != 0)) { + return false; + } + size_t num_outputs = 0; + if (!reduce_utils::CheckedElementCount( + absl::MakeConstSpan(output_dims, + static_cast(output_num_dims)), + num_outputs)) { + return false; } for (size_t idx = 0; idx < num_outputs; ++idx) { output_data[idx] = T(); @@ -249,14 +255,17 @@ inline bool Mean(const T* input_data, const int* input_dims, } // 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(input_dims[resolved_axis[idx]]); - // Overflow prevention. - if (current > (std::numeric_limits::max() / num_elements_in_axis)) { - return false; - } - num_elements_in_axis *= current; + if (input_num_dims < 0 || (input_dims == nullptr && input_num_dims != 0) || + (resolved_axis == nullptr && num_resolved_axis != 0)) { + return false; + } + size_t num_elements_in_axis = 0; + if (!reduce_utils::CheckedReducedElementCount( + absl::MakeConstSpan(input_dims, static_cast(input_num_dims)), + absl::MakeConstSpan(resolved_axis, + static_cast(num_resolved_axis)), + num_elements_in_axis)) { + return false; } if (num_elements_in_axis > 0) { @@ -337,14 +346,15 @@ inline bool QuantizedMeanOrSum(const T* input_data, int32_t input_zero_point, ruy::profiler::ScopeLabel label(compute_sum ? "Sum/Int8" : "Mean/Int8"); } // Reset output data. - size_t num_outputs = 1; - for (int idx = 0; idx < output_num_dims; ++idx) { - size_t current = static_cast(output_dims[idx]); - // Overflow prevention. - if (num_outputs > std::numeric_limits::max() / current) { - return false; - } - num_outputs *= current; + if (output_num_dims < 0 || (output_dims == nullptr && output_num_dims != 0)) { + return false; + } + size_t num_outputs = 0; + if (!reduce_utils::CheckedElementCount( + absl::MakeConstSpan(output_dims, + static_cast(output_num_dims)), + num_outputs)) { + return false; } for (size_t idx = 0; idx < num_outputs; ++idx) { output_data[idx] = T(); @@ -373,15 +383,17 @@ inline bool QuantizedMeanOrSum(const T* input_data, int32_t input_zero_point, } // Calculate mean by dividing output_data by num of aggregated element. - int64_t num_elements_in_axis = 1; - for (int idx = 0; idx < num_resolved_axis; ++idx) { - size_t current = static_cast(input_dims[resolved_axis[idx]]); - // Overflow prevention. - if (current > static_cast(std::numeric_limits::max() / - num_elements_in_axis)) { - return false; - } - num_elements_in_axis *= current; + if (input_num_dims < 0 || (input_dims == nullptr && input_num_dims != 0) || + (resolved_axis == nullptr && num_resolved_axis != 0)) { + return false; + } + int64_t num_elements_in_axis = 0; + if (!reduce_utils::CheckedReducedElementCount( + absl::MakeConstSpan(input_dims, static_cast(input_num_dims)), + absl::MakeConstSpan(resolved_axis, + static_cast(num_resolved_axis)), + num_elements_in_axis)) { + return false; } if (num_elements_in_axis == 0) { diff --git a/tensorflow/lite/kernels/internal/reference/resize_nearest_neighbor.h b/tensorflow/lite/kernels/internal/reference/resize_nearest_neighbor.h index bf0b757e916..d9e0cb28ddb 100644 --- a/tensorflow/lite/kernels/internal/reference/resize_nearest_neighbor.h +++ b/tensorflow/lite/kernels/internal/reference/resize_nearest_neighbor.h @@ -71,9 +71,11 @@ inline void ResizeNearestNeighbor( int32_t output_height = output_size_data[0]; int32_t output_width = output_size_data[1]; - const int col_offset = input_shape.Dims(3); - const int row_offset = input_shape.Dims(2) * col_offset; - const int batch_offset = input_shape.Dims(1) * row_offset; + const int64_t col_offset = input_shape.Dims(3); + const int64_t row_offset = + static_cast(input_shape.Dims(2)) * col_offset; + const int64_t batch_offset = + static_cast(input_shape.Dims(1)) * row_offset; const T* input_ptr = input_data; T* output_ptr = output_data; diff --git a/tensorflow/lite/kernels/internal/reference/select.h b/tensorflow/lite/kernels/internal/reference/select.h index 4939d067433..934fcf56dff 100644 --- a/tensorflow/lite/kernels/internal/reference/select.h +++ b/tensorflow/lite/kernels/internal/reference/select.h @@ -237,10 +237,10 @@ void BroadcastSelect5DSlow(const RuntimeShape& input_condition_shape, const T* input_y_data, const RuntimeShape& output_shape, T* output_data) { ruy::profiler::ScopeLabel label("Select/BroadcastSelectSlow"); - TFLITE_DCHECK_LE(input_condition_shape.DimensionsCount(), 5); - TFLITE_DCHECK_LE(input_x_shape.DimensionsCount(), 5); - TFLITE_DCHECK_LE(input_y_shape.DimensionsCount(), 5); - TFLITE_DCHECK_LE(output_shape.DimensionsCount(), 5); + TFLITE_DCHECK_LE(input_condition_shape.DimensionsCount(), 8); + TFLITE_DCHECK_LE(input_x_shape.DimensionsCount(), 8); + TFLITE_DCHECK_LE(input_y_shape.DimensionsCount(), 8); + TFLITE_DCHECK_LE(output_shape.DimensionsCount(), 8); BroadcastSelectSimple(input_condition_shape, input_condition_data, input_x_shape, input_x_data, input_y_shape, diff --git a/tensorflow/lite/kernels/kernel_util.cc b/tensorflow/lite/kernels/kernel_util.cc index e2911c9440f..36ad5178158 100644 --- a/tensorflow/lite/kernels/kernel_util.cc +++ b/tensorflow/lite/kernels/kernel_util.cc @@ -599,8 +599,18 @@ bool HasUnspecifiedDimension(const TfLiteTensor* tensor) { TfLiteStatus CheckedShapeProduct(TfLiteContext* context, std::initializer_list dims, const char* error_message, size_t& product) { + return CheckedShapeProduct(context, dims.begin(), dims.size(), error_message, + product); +} + +TfLiteStatus CheckedShapeProduct(TfLiteContext* context, const int* dims, + int count, const char* error_message, + size_t& product) { + TF_LITE_ENSURE(context, count >= 0); + TF_LITE_ENSURE(context, dims != nullptr || count == 0); size_t checked_count = 1; - for (const int d : dims) { + for (int i = 0; i < count; ++i) { + const int d = dims[i]; TF_LITE_ENSURE_MSG(context, d >= 0, "Encountered a negative dimension."); TF_LITE_ENSURE_MSG( context, diff --git a/tensorflow/lite/kernels/kernel_util.h b/tensorflow/lite/kernels/kernel_util.h index eacd1bad38c..e107ab23177 100644 --- a/tensorflow/lite/kernels/kernel_util.h +++ b/tensorflow/lite/kernels/kernel_util.h @@ -26,9 +26,6 @@ limitations under the License. #include "tensorflow/lite/core/c/builtin_op_data.h" #include "tensorflow/lite/core/c/common.h" -#ifndef NDEBUG -#include "tensorflow/lite/kernels/op_macros.h" -#endif namespace tflite { @@ -173,17 +170,6 @@ inline int NumIntermediates(const TfLiteNode* node) { inline int64_t NumElements(const int* dims, int num_dims) { int64_t count = 1; for (int i = 0; i < num_dims; ++i) { -#ifndef NDEBUG - if (count <= 0) { - break; - } - // Check that number of elements can fit in 32 bit int. Most of tflite - // assumes the result of `NumElements` is < MAX_INT and static or implicit - // casts to `int32_t` without any checks. It is more meaningful to check - // that the result fits into 32 bits than for standard overflow on 64 bit - // type. - TF_LITE_ASSERT(dims[i] < std::numeric_limits::max() / count); -#endif count *= dims[i]; } return count; @@ -360,6 +346,20 @@ TfLiteStatus CheckedShapeProduct(TfLiteContext* context, * the dimensions is negative or if the product overflows. * @param context The context to use for error reporting. * @param dims The dimensions to multiply. + * @param count The length of the dims array. + * @param error_message The error message to use if an error is encountered. + * @param product The output parameter to store the product. + */ +TfLiteStatus CheckedShapeProduct(TfLiteContext* context, const int* dims, + int count, const char* error_message, + size_t& product); + +/** + * Calculates the product of the given dimensions. Returns an error if any of + * the dimensions is negative or if the product overflows. (Same as above + * function with dims built on the fly) + * @param context The context to use for error reporting. + * @param dims The dimensions to multiply. * @param error_message The error message to use if an error is encountered. * @param product The output parameter to store the product. */ diff --git a/tensorflow/lite/kernels/padding.h b/tensorflow/lite/kernels/padding.h index cc9d596f1a5..e51d683b70d 100644 --- a/tensorflow/lite/kernels/padding.h +++ b/tensorflow/lite/kernels/padding.h @@ -15,16 +15,63 @@ limitations under the License. #ifndef TENSORFLOW_LITE_KERNELS_PADDING_H_ #define TENSORFLOW_LITE_KERNELS_PADDING_H_ +#include +#include + #include "tensorflow/lite/core/c/builtin_op_data.h" +#include "tensorflow/lite/core/c/common.h" #include "tensorflow/lite/kernels/internal/types.h" namespace tflite { +inline TfLiteStatus CheckedNarrowPaddingValue(int64_t value, int* result) { + if (result == nullptr || value > std::numeric_limits::max() || + value < std::numeric_limits::min()) { + return kTfLiteError; + } + *result = static_cast(value); + return kTfLiteOk; +} + +inline int64_t ComputeEffectiveFilterSize(int filter_size, int dilation_rate) { + return (static_cast(filter_size) - 1) * dilation_rate + 1; +} + +inline TfLiteStatus ValidatePaddingArguments(TfLitePadding padding, + int image_size, int filter_size, + int stride, int dilation_rate) { + if ((padding != kTfLitePaddingSame && padding != kTfLitePaddingValid) || + image_size < 0 || filter_size <= 0 || stride <= 0 || dilation_rate <= 0) { + return kTfLiteError; + } + return kTfLiteOk; +} + +inline TfLiteStatus ComputePaddingWithOffsetChecked( + int stride, int dilation_rate, int in_size, int filter_size, int out_size, + int* offset, int* padding) { + if (offset == nullptr || padding == nullptr || in_size < 0 || + filter_size <= 0 || out_size < 0 || stride <= 0 || dilation_rate <= 0) { + return kTfLiteError; + } + const int64_t effective_filter_size = + ComputeEffectiveFilterSize(filter_size, dilation_rate); + int64_t total_padding = ((static_cast(out_size) - 1) * stride + + effective_filter_size - in_size); + total_padding = total_padding > 0 ? total_padding : 0; + *offset = static_cast(total_padding % 2); + return CheckedNarrowPaddingValue(total_padding / 2, padding); +} + inline int ComputePadding(int stride, int dilation_rate, int in_size, int filter_size, int out_size) { - int effective_filter_size = (filter_size - 1) * dilation_rate + 1; - int padding = ((out_size - 1) * stride + effective_filter_size - in_size) / 2; - return padding > 0 ? padding : 0; + int offset = 0; + int padding = 0; + return ComputePaddingWithOffsetChecked(stride, dilation_rate, in_size, + filter_size, out_size, &offset, + &padding) == kTfLiteOk + ? padding + : 0; } // It's not guaranteed that padding is symmetric. It's important to keep @@ -32,52 +79,97 @@ inline int ComputePadding(int stride, int dilation_rate, int in_size, inline int ComputePaddingWithOffset(int stride, int dilation_rate, int in_size, int filter_size, int out_size, int* offset) { - int effective_filter_size = (filter_size - 1) * dilation_rate + 1; - int total_padding = - ((out_size - 1) * stride + effective_filter_size - in_size); - total_padding = total_padding > 0 ? total_padding : 0; - *offset = total_padding % 2; - return total_padding / 2; + int padding = 0; + if (ComputePaddingWithOffsetChecked(stride, dilation_rate, in_size, + filter_size, out_size, offset, + &padding) != kTfLiteOk) { + if (offset != nullptr) *offset = 0; + return 0; + } + return padding; } // Matching GetWindowedOutputSize in TensorFlow. -inline int ComputeOutSize(TfLitePadding padding, int image_size, - int filter_size, int stride, int dilation_rate = 1) { - int effective_filter_size = (filter_size - 1) * dilation_rate + 1; - - // TODO(b/186448822): This uses 0 since the function has no other way to - // report error case - if (stride == 0) return 0; +inline TfLiteStatus ComputeOutSizeChecked(TfLitePadding padding, int image_size, + int filter_size, int stride, + int dilation_rate, int* out_size) { + if (out_size == nullptr || + ValidatePaddingArguments(padding, image_size, filter_size, stride, + dilation_rate) != kTfLiteOk) { + return kTfLiteError; + } + const int64_t effective_filter_size = + ComputeEffectiveFilterSize(filter_size, dilation_rate); + int64_t value = 0; switch (padding) { case kTfLitePaddingSame: - return (image_size + stride - 1) / stride; + value = (static_cast(image_size) + stride - 1) / stride; + break; case kTfLitePaddingValid: - return (image_size + stride - effective_filter_size) / stride; + value = + (static_cast(image_size) + stride - effective_filter_size) / + stride; + break; default: - return 0; + return kTfLiteError; } + if (value < 0) return kTfLiteError; + return CheckedNarrowPaddingValue(value, out_size); +} + +inline int ComputeOutSize(TfLitePadding padding, int image_size, + int filter_size, int stride, int dilation_rate = 1) { + int out_size = 0; + return ComputeOutSizeChecked(padding, image_size, filter_size, stride, + dilation_rate, &out_size) == kTfLiteOk + ? out_size + : 0; +} + +inline TfLiteStatus ComputePaddingHeightWidthChecked( + int stride_height, int stride_width, int dilation_rate_height, + int dilation_rate_width, int in_height, int in_width, int filter_height, + int filter_width, TfLitePadding padding, int* out_height, int* out_width, + TfLitePaddingValues* padding_values) { + if (out_height == nullptr || out_width == nullptr || + padding_values == nullptr) { + return kTfLiteError; + } + TF_LITE_ENSURE_STATUS(ComputeOutSizeChecked(padding, in_width, filter_width, + stride_width, dilation_rate_width, + out_width)); + TF_LITE_ENSURE_STATUS( + ComputeOutSizeChecked(padding, in_height, filter_height, stride_height, + dilation_rate_height, out_height)); + + int offset = 0; + TF_LITE_ENSURE_STATUS(ComputePaddingWithOffsetChecked( + stride_height, dilation_rate_height, in_height, filter_height, + *out_height, &offset, &padding_values->height)); + padding_values->height_offset = offset; + TF_LITE_ENSURE_STATUS(ComputePaddingWithOffsetChecked( + stride_width, dilation_rate_width, in_width, filter_width, *out_width, + &offset, &padding_values->width)); + padding_values->width_offset = offset; + return kTfLiteOk; } inline TfLitePaddingValues ComputePaddingHeightWidth( int stride_height, int stride_width, int dilation_rate_height, int dilation_rate_width, int in_height, int in_width, int filter_height, int filter_width, TfLitePadding padding, int* out_height, int* out_width) { - *out_width = ComputeOutSize(padding, in_width, filter_width, stride_width, - dilation_rate_width); - *out_height = ComputeOutSize(padding, in_height, filter_height, stride_height, - dilation_rate_height); - TfLitePaddingValues padding_values; - int offset = 0; - padding_values.height = - ComputePaddingWithOffset(stride_height, dilation_rate_height, in_height, - filter_height, *out_height, &offset); - padding_values.height_offset = offset; - padding_values.width = - ComputePaddingWithOffset(stride_width, dilation_rate_width, in_width, - filter_width, *out_width, &offset); - padding_values.width_offset = offset; + if (out_height != nullptr) *out_height = 0; + if (out_width != nullptr) *out_width = 0; + padding_values.height = 0; + padding_values.height_offset = 0; + padding_values.width = 0; + padding_values.width_offset = 0; + ComputePaddingHeightWidthChecked( + stride_height, stride_width, dilation_rate_height, dilation_rate_width, + in_height, in_width, filter_height, filter_width, padding, out_height, + out_width, &padding_values); return padding_values; } diff --git a/tensorflow/lite/tools/flatbuffer_utils_test.py b/tensorflow/lite/tools/flatbuffer_utils_test.py index b604d178f3a..ae715a7dd52 100644 --- a/tensorflow/lite/tools/flatbuffer_utils_test.py +++ b/tensorflow/lite/tools/flatbuffer_utils_test.py @@ -18,9 +18,9 @@ import subprocess import sys -from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.python import schema_py_generated as schema # pylint:disable=g-direct-tensorflow-import -from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import flatbuffer_utils -from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import test_utils +from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.python import schema_py_generated as schema # pylint:disable=g-direct-tensorflow-import +from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import flatbuffer_utils +from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import test_utils from tensorflow.python.framework import test_util from tensorflow.python.platform import test diff --git a/tensorflow/lite/tools/test_utils.py b/tensorflow/lite/tools/test_utils.py index d71267bb401..281baee11b2 100644 --- a/tensorflow/lite/tools/test_utils.py +++ b/tensorflow/lite/tools/test_utils.py @@ -18,7 +18,7 @@ """ import flatbuffers -from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.python import schema_py_generated as schema_fb +from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.python import schema_py_generated as schema_fb TFLITE_SCHEMA_VERSION = 3 diff --git a/tensorflow/lite/tools/visualize_test.py b/tensorflow/lite/tools/visualize_test.py index 5f00a189272..d3404c9785b 100644 --- a/tensorflow/lite/tools/visualize_test.py +++ b/tensorflow/lite/tools/visualize_test.py @@ -16,8 +16,8 @@ import os import re -from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import test_utils -from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import visualize +from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import test_utils +from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import visualize from tensorflow.python.framework import test_util from tensorflow.python.platform import test