Sync from upstream TF.

Author: TFLM-bot

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<int64_t>(
-        /*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<int64_t>(
-        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<int64_t>(
-        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<int64_t>(
+          /*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<int64_t>(
+          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<int64_t>(
+          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<int64_t>(
-        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<int64_t>(
+          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 {

tensorflow/lite/core/c/common.cc

@@ -20,6 +20,7 @@ limitations under the License.
 #endif  // TF_LITE_STATIC_MEMORY
 
 #include <cstring>
+#include <limits>
 #include <new>
 #include <type_traits>
 #include <utility>
@@ -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<size_t>::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);

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_

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<int32_t>(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:

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 <typename T>
+inline T WrappingAdd(T a, T b) {
+  if constexpr (std::is_integral_v<T> && !std::is_same_v<T, bool>) {
+    using U = std::make_unsigned_t<T>;
+    using P = std::common_type_t<U, unsigned int>;
+    return static_cast<T>(static_cast<U>(static_cast<P>(static_cast<U>(a)) +
+                                         static_cast<P>(static_cast<U>(b))));
+  } else {
+    return a + b;
+  }
+}
+
+template <typename T>
+inline T WrappingMul(T a, T b) {
+  if constexpr (std::is_integral_v<T> && !std::is_same_v<T, bool>) {
+    using U = std::make_unsigned_t<T>;
+    using P = std::common_type_t<U, unsigned int>;
+    return static_cast<T>(static_cast<U>(static_cast<P>(static_cast<U>(a)) *
+                                         static_cast<P>(static_cast<U>(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 <typename T>
+TFLITE_NO_SANITIZE_INTEGER_OVERFLOW inline T
+AddTensorValuesWithExpectedOverflow(T a, T b) {
+  return a + b;
+}
+
+template <typename T>
+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.

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<IntOut>::max().
@@ -125,47 +125,36 @@ QuantizationParams ChooseQuantizationParams(double rmin, double rmax) {
 // the results are undefined.
 // TODO(sfeuz): Replace by absl::SafeCast once available.
 template <class IntOut, class FloatIn>
-IntOut SafeCast(FloatIn x) {
+IntOut SafeCast(FloatIn x, IntOut nan_result = IntOut{0}) {
   static_assert(!std::numeric_limits<FloatIn>::is_integer,
                 "FloatIn is integer");
   static_assert(std::numeric_limits<IntOut>::is_integer,
                 "IntOut is not integer");
   static_assert(std::numeric_limits<IntOut>::radix == 2, "IntOut is base 2");
-
-  // Special case NaN, for which the logic below doesn't work.
+  static_assert(std::numeric_limits<FloatIn>::max_exponent >
+                    std::numeric_limits<IntOut>::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<IntOut>::is_signed && x < 0) {
-    return 0;
+    return nan_result;
   }
 
-  // Handle infinities.
-  if (std::isinf(x)) {
-    return x < 0 ? std::numeric_limits<IntOut>::min()
-                 : std::numeric_limits<IntOut>::max();
+  const FloatIn min_value =
+      static_cast<FloatIn>(std::numeric_limits<IntOut>::min());
+  if (x < min_value) {
+    return std::numeric_limits<IntOut>::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<IntOut>::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<FloatIn>(std::numeric_limits<IntOut>::max() / 2 + 1) *
+      FloatIn{2};
+  if (x >= kMaxExclusive) {
+    return std::numeric_limits<IntOut>::max();
   }
 
-  // Handle numbers with magnitude >= 2^N.
-  return x < 0 ? std::numeric_limits<IntOut>::min()
-               : std::numeric_limits<IntOut>::max();
+  return static_cast<IntOut>(x);
 }
 // LINT.ThenChange(//tensorflow/compiler/mlir/lite/kernels/internal/quantization_util.h)
 

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<T>(
-        input_data[c] + broadcast_data[0], activation_min, activation_max);
+        WrappingAdd<T>(input_data[c], broadcast_data[0]), activation_min,
+        activation_max);
   }
 }
 
@@ -229,7 +230,8 @@ inline void AddBroadcast<int32_t>(const int32_t* input_data,
 #endif
   for (; c < size; ++c) {
     output_data[c] = ActivationFunctionWithMinMax<int32_t>(
-        input_data[c] + broadcast_data[0], activation_min, activation_max);
+        WrappingAdd<int32_t>(input_data[c], broadcast_data[0]), activation_min,
+        activation_max);
   }
 }
 

tensorflow/lite/kernels/internal/reference/arg_min_max.h

@@ -15,23 +15,50 @@ 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 <cmath>
 #include <functional>
+#include <type_traits>
 
 #include "tensorflow/lite/kernels/internal/types.h"
 
 namespace tflite {
 
 namespace reference_ops {
 
+// Default comparator for non-floating-point types (no NaN possible).
 template <typename T>
-std::function<bool(T, T)> GetComparefunction(bool is_arg_max) {
+typename std::enable_if<!std::is_floating_point<T>::value,
+                        std::function<bool(T, T)>>::type
+GetComparefunction(bool is_arg_max) {
   if (is_arg_max) {
     return std::greater<T>();
   } else {
     return std::less<T>();
   }
 }
 
+// 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 T>
+typename std::enable_if<std::is_floating_point<T>::value,
+                        std::function<bool(T, T)>>::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 <typename T1, typename T2, typename T3, typename Cmp>
 void ArgMinMax(const RuntimeShape& input1_shape, const T1* input1_data,
                const T3* input2_data, const RuntimeShape& output_shape,

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;
         }
       }