Automated sync from github.com/tensorflow/tensorflow

tensorflow/compiler/mlir/lite/core/c/tflite_types.h

@@ -66,6 +66,8 @@ typedef enum {
   kTfLiteBFloat16 = 19,
   kTfLiteInt2 = 20,
   kTfLiteUInt4 = 21,
+  kTfLiteFloat8E4M3FN = 22,
+  kTfLiteFloat8E5M2 = 23,
 } TfLiteType;
 // LINT.ThenChange(//tensorflow/lite/profiling/proto/model_runtime_info.proto:EdgeDataType)
 

tensorflow/compiler/mlir/lite/schema/schema.fbs

@@ -63,6 +63,8 @@ enum TensorType : byte {
   BFLOAT16 = 18,
   INT2 = 19,
   UINT4 = 20,
+  FLOAT8_E4M3FN = 21,
+  FLOAT8_E5M2 = 22,
 }
 
 // Custom quantization parameters for experimenting with new quantization

tensorflow/lite/core/api/flatbuffer_conversions.cc

@@ -1101,6 +1101,12 @@ TfLiteStatus ConvertTensorType(TensorType tensor_type, TfLiteType* type,
     case TensorType_UINT4:
       *type = kTfLiteUInt4;
       return kTfLiteOk;
+    case TensorType_FLOAT8_E4M3FN:
+      *type = kTfLiteFloat8E4M3FN;
+      return kTfLiteOk;
+    case TensorType_FLOAT8_E5M2:
+      *type = kTfLiteFloat8E5M2;
+      return kTfLiteOk;
     default:
       *type = kTfLiteNoType;
       TF_LITE_REPORT_ERROR(error_reporter,

tensorflow/lite/core/c/common.cc

@@ -20,6 +20,7 @@ limitations under the License.
 #endif  // TF_LITE_STATIC_MEMORY
 
 #include <cstring>
+#include <new>
 #include <type_traits>
 #include <utility>
 
@@ -111,6 +112,7 @@ TfLiteSparsity TfLiteSparsityClone(const TfLiteSparsity& src) {
   if (src.dim_metadata) {
     dst.dim_metadata = reinterpret_cast<TfLiteDimensionMetadata*>(
         calloc(1, sizeof(TfLiteDimensionMetadata) * src.dim_metadata_size));
+    if (src.dim_metadata_size > 0 && !dst.dim_metadata) return TfLiteSparsity();
     for (int i = 0; i < src.dim_metadata_size; ++i) {
       dst.dim_metadata[i] = src.dim_metadata[i];
       dst.dim_metadata[i].array_segments =
@@ -129,6 +131,7 @@ TfLiteSparsity* TfLiteSparsityClone(const TfLiteSparsity* const src) {
   }
   TfLiteSparsity* dst =
       reinterpret_cast<TfLiteSparsity*>(calloc(1, sizeof(TfLiteSparsity)));
+  if (!dst) return nullptr;
   *dst = TfLiteSparsityClone(*src);
   return dst;
 }
@@ -147,6 +150,7 @@ TfLiteQuantization TfLiteQuantizationClone(const TfLiteQuantization& src) {
       break;
     case kTfLiteAffineQuantization: {
       dst.params = calloc(1, sizeof(TfLiteAffineQuantization));
+      if (!dst.params) return TfLiteQuantization();
       const TfLiteAffineQuantization* const src_params =
           reinterpret_cast<TfLiteAffineQuantization*>(src.params);
       TfLiteAffineQuantization* const dst_params =
@@ -158,6 +162,7 @@ TfLiteQuantization TfLiteQuantizationClone(const TfLiteQuantization& src) {
     }
     case kTfLiteBlockwiseQuantization: {
       dst.params = calloc(1, sizeof(TfLiteBlockwiseQuantization));
+      if (!dst.params) return TfLiteQuantization();
       const TfLiteBlockwiseQuantization* const src_params =
           (TfLiteBlockwiseQuantization*)(src.params);
       TfLiteBlockwiseQuantization* const dst_params =
@@ -219,6 +224,9 @@ TfLiteFloatArray* TfLiteFloatArrayCopy(const TfLiteFloatArray* src) {
 void TfLiteFloatArrayFree(TfLiteFloatArray* a) { TfLiteVarArrayFree(a); }
 
 void TfLiteTensorDataFree(TfLiteTensor* t) {
+  if (t == nullptr) {
+    return;
+  }
   if (t->allocation_type == kTfLiteVariantObject && t->data.data) {
     delete static_cast<VariantData*>(t->data.data);
   } else if (t->allocation_type == kTfLiteDynamic ||
@@ -238,6 +246,9 @@ void TfLiteTensorDataFree(TfLiteTensor* t) {
 }
 
 void TfLiteQuantizationFree(TfLiteQuantization* quantization) {
+  if (quantization == nullptr) {
+    return;
+  }
   if (quantization->type == kTfLiteAffineQuantization) {
     TfLiteAffineQuantization* q_params =
         reinterpret_cast<TfLiteAffineQuantization*>(quantization->params);
@@ -294,6 +305,9 @@ void TfLiteSparsityFree(TfLiteSparsity* sparsity) {
 }
 
 void TfLiteTensorFree(TfLiteTensor* t) {
+  if (t == nullptr) {
+    return;
+  }
   TfLiteTensorDataFree(t);
   if (t->dims) TfLiteIntArrayFree(t->dims);
   t->dims = nullptr;
@@ -308,7 +322,7 @@ void TfLiteTensorFree(TfLiteTensor* t) {
   t->sparsity = nullptr;
 }
 
-TfLiteTensor TfLiteTensorClone(const TfLiteTensor src) {
+TfLiteTensor TfLiteTensorClone(TfLiteTensor src) {
   // We copy all of the source data first, then we clone the fields that can't
   // be shared between two tensor instances.
   TfLiteTensor dst = src;
@@ -335,16 +349,18 @@ TfLiteTensor TfLiteTensorClone(const TfLiteTensor src) {
         break;
       case kTfLiteAllocationStrategyMalloc:
         dst.data.data = malloc(src.bytes);
+        if (src.bytes > 0 && !dst.data.data) return TfLiteTensor();
         std::memcpy(dst.data.data, src.data.data, src.bytes);
         break;
       case kTfLiteAllocationStrategyNew:
         // Special case for variant objects. They are allocated using new/delete
         // but require using the `CloneTo` function.
         if (src.allocation_type == kTfLiteVariantObject) {
-          dst.data.data = reinterpret_cast<const VariantData*>(src.data.data)
-                              ->CloneTo(nullptr);
+          dst.data.data =
+              static_cast<const VariantData*>(src.data.data)->CloneTo(nullptr);
         } else {
-          dst.data.data = new char[src.bytes];
+          dst.data.data = new (std::nothrow) char[src.bytes];
+          if (src.bytes > 0 && !dst.data.data) return TfLiteTensor();
           std::memcpy(dst.data.data, src.data.data, src.bytes);
         }
         break;
@@ -394,13 +410,21 @@ TfLiteStatus TfLiteTensorCopy(const TfLiteTensor* src, TfLiteTensor* dst) {
     }
     auto* dst_vd = static_cast<VariantData*>(dst->data.data);
     auto* src_vd = static_cast<VariantData*>(src->data.data);
+    if (!src_vd) return kTfLiteError;
 
     // `CloneTo` will handle the case when `dst_vd` is nullptr, so it is safe
     // to `CloneTo` something which was "freed". Also, returning from `CloneTo`
     // will implicitly cast to `VariantData`; don't need static cast here.
     dst->data.data = src_vd->CloneTo(dst_vd);
   } else {
-    memcpy(dst->data.raw, src->data.raw, src->bytes);
+    if (dst->allocation_type == kTfLiteVariantObject) {
+      TfLiteTensorDataFree(dst);
+      dst->allocation_type = src->allocation_type;
+    }
+    if (src->bytes > 0) {
+      if (!dst->data.raw || !src->data.raw) return kTfLiteError;
+      memcpy(dst->data.raw, src->data.raw, src->bytes);
+    }
   }
   dst->buffer_handle = src->buffer_handle;
   dst->data_is_stale = src->data_is_stale;
@@ -513,6 +537,10 @@ const char* TfLiteTypeGetName(TfLiteType type) {
       return "INT2";
     case kTfLiteUInt4:
       return "UINT4";
+    case kTfLiteFloat8E4M3FN:
+      return "FLOAT8_E4M3FN";
+    case kTfLiteFloat8E5M2:
+      return "FLOAT8_E5M2";
   }
   return "Unknown type";
 }

tensorflow/lite/core/c/common.h

@@ -56,6 +56,7 @@ limitations under the License.
 #include <stdbool.h>
 #include <stddef.h>
 #include <stdint.h>
+#include <stdio.h>
 
 #include "tensorflow/lite/core/c/c_api_types.h"  // IWYU pragma: export
 
@@ -277,13 +278,34 @@ void TfLiteFloatArrayFree(TfLiteFloatArray* a);
     }                                                                        \
   } while (0)
 
-#define TF_LITE_ENSURE_OK(context, status) \
-  do {                                     \
-    const TfLiteStatus s = (status);       \
-    if ((s) != kTfLiteOk) {                \
-      return s;                            \
-    }                                      \
+#ifndef TF_LITE_STRIP_ERROR_STRINGS
+#define TF_LITE_VAR_ARG_HEAD(FIRST, ...) FIRST
+#define TF_LITE_STRINGIFY_HELPER(x) #x
+#define TF_LITE_STRINGIFY(x) TF_LITE_STRINGIFY_HELPER(x)
+// Checks that `status` evaluates to `kTfLiteOk`.
+//
+// Can take a printf style log message and its parameters after the status. The
+// message will be printed using `TF_LITE_KERNEL_LOG` in case of error.
+#define TF_LITE_ENSURE_OK(context, status, ...)                              \
+  do {                                                                       \
+    const TfLiteStatus s = (status);                                         \
+    if (s != kTfLiteOk) {                                                    \
+      if (sizeof(TF_LITE_VAR_ARG_HEAD("" __VA_ARGS__)) > sizeof("")) {       \
+        TF_LITE_MAYBE_KERNEL_LOG((context), __FILE__ ":" TF_LITE_STRINGIFY(  \
+                                                __LINE__) ": " __VA_ARGS__); \
+      }                                                                      \
+      return s;                                                              \
+    }                                                                        \
   } while (0)
+#else
+#define TF_LITE_ENSURE_OK(context, status, ...) \
+  do {                                          \
+    const TfLiteStatus s = (status);            \
+    if ((s) != kTfLiteOk) {                     \
+      return s;                                 \
+    }                                           \
+  } while (0)
+#endif
 
 // `std::unreachable` not available until CC23.
 #ifdef __GNUC__  // GCC, Clang, ICC
@@ -1060,6 +1082,13 @@ typedef struct TfLiteContext {
   /// WARNING: This is an experimental interface that is subject to change.
   TfLiteStatus (*ReleaseSubgraphContext)(struct TfLiteContext* context,
                                          int subgraph_index);
+#if defined(_WIN32)
+  /// Create a array of a given `size` (uninitialized entries).
+  TfLiteIntArray* (*TfLiteIntArrayCreate)(int size);  // NOLINT
+
+  /// Free memory of array `a`.
+  void (*TfLiteIntArrayFree)(TfLiteIntArray* a);  // NOLINT
+#endif                                            // defined(_WIN32)
 } TfLiteContext;
 
 /// `TfLiteOperator` is an external version of `TfLiteRegistration`

tensorflow/lite/kernels/internal/common.h

@@ -78,15 +78,23 @@ bool ReduceDimensionsForBroadcast(const RuntimeShape& input1_shape,
       if (!broadcast_input1) {
         broadcast_input1 = true;
         broadcast_input2 = false;
+        if (num_compressed_dims >= MAX_DIM) return false;
         num_compressed_dims++;
+        if (num_compressed_dims > MAX_DIM) {
+          return false;
+        }
       }
       compressed_input2_shape[num_compressed_dims - 1] *= input2_dim;
       compressed_output_shape[num_compressed_dims - 1] *= input2_dim;
     } else if (input2_dim == 1) {
       if (!broadcast_input2) {
         broadcast_input1 = false;
         broadcast_input2 = true;
+        if (num_compressed_dims >= MAX_DIM) return false;
         num_compressed_dims++;
+        if (num_compressed_dims > MAX_DIM) {
+          return false;
+        }
       }
       compressed_input1_shape[num_compressed_dims - 1] *= input1_dim;
       compressed_output_shape[num_compressed_dims - 1] *= input1_dim;
@@ -95,7 +103,11 @@ bool ReduceDimensionsForBroadcast(const RuntimeShape& input1_shape,
       if (broadcast_input1 || broadcast_input2 || first_nonunit) {
         broadcast_input1 = false;
         broadcast_input2 = false;
+        if (num_compressed_dims >= MAX_DIM) return false;
         num_compressed_dims++;
+        if (num_compressed_dims > MAX_DIM) {
+          return false;
+        }
       }
       compressed_input1_shape[num_compressed_dims - 1] *= input1_dim;
       compressed_input2_shape[num_compressed_dims - 1] *= input1_dim;
@@ -105,7 +117,11 @@ bool ReduceDimensionsForBroadcast(const RuntimeShape& input1_shape,
   }
   if (num_input1_dims > num_input2_dims) {
     if (!broadcast_input2) {
+      if (num_compressed_dims >= MAX_DIM) return false;
       num_compressed_dims++;
+      if (num_compressed_dims > MAX_DIM) {
+        return false;
+      }
     }
     for (size_t i = 0; i < num_input1_dims - num_input2_dims; i++) {
       const size_t input1_dim = input1_dims[i];
@@ -117,7 +133,11 @@ bool ReduceDimensionsForBroadcast(const RuntimeShape& input1_shape,
     }
   } else if (num_input2_dims > num_input1_dims) {
     if (!broadcast_input1) {
+      if (num_compressed_dims >= MAX_DIM) return false;
       num_compressed_dims++;
+      if (num_compressed_dims > MAX_DIM) {
+        return false;
+      }
     }
     for (size_t i = 0; i < num_input2_dims - num_input1_dims; i++) {
       const size_t input2_dim = input2_dims[i];