Skip to content

Commit 72af1b2

Browse files
committed
Sync from upstream TF.
1 parent 348eed0 commit 72af1b2

31 files changed

Lines changed: 698 additions & 1071 deletions

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,8 @@ typedef enum {
6666
kTfLiteBFloat16 = 19,
6767
kTfLiteInt2 = 20,
6868
kTfLiteUInt4 = 21,
69+
kTfLiteFloat8E4M3FN = 22,
70+
kTfLiteFloat8E5M2 = 23,
6971
} TfLiteType;
7072
// LINT.ThenChange(//tensorflow/lite/profiling/proto/model_runtime_info.proto:EdgeDataType)
7173

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,8 @@ enum TensorType : byte {
6363
BFLOAT16 = 18,
6464
INT2 = 19,
6565
UINT4 = 20,
66+
FLOAT8_E4M3FN = 21,
67+
FLOAT8_E5M2 = 22,
6668
}
6769

6870
// Custom quantization parameters for experimenting with new quantization

tensorflow/lite/core/api/flatbuffer_conversions.cc

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1101,6 +1101,12 @@ TfLiteStatus ConvertTensorType(TensorType tensor_type, TfLiteType* type,
11011101
case TensorType_UINT4:
11021102
*type = kTfLiteUInt4;
11031103
return kTfLiteOk;
1104+
case TensorType_FLOAT8_E4M3FN:
1105+
*type = kTfLiteFloat8E4M3FN;
1106+
return kTfLiteOk;
1107+
case TensorType_FLOAT8_E5M2:
1108+
*type = kTfLiteFloat8E5M2;
1109+
return kTfLiteOk;
11041110
default:
11051111
*type = kTfLiteNoType;
11061112
TF_LITE_REPORT_ERROR(error_reporter,

tensorflow/lite/core/c/common.cc

Lines changed: 33 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ limitations under the License.
2020
#endif // TF_LITE_STATIC_MEMORY
2121

2222
#include <cstring>
23+
#include <new>
2324
#include <type_traits>
2425
#include <utility>
2526

@@ -111,6 +112,7 @@ TfLiteSparsity TfLiteSparsityClone(const TfLiteSparsity& src) {
111112
if (src.dim_metadata) {
112113
dst.dim_metadata = reinterpret_cast<TfLiteDimensionMetadata*>(
113114
calloc(1, sizeof(TfLiteDimensionMetadata) * src.dim_metadata_size));
115+
if (src.dim_metadata_size > 0 && !dst.dim_metadata) return TfLiteSparsity();
114116
for (int i = 0; i < src.dim_metadata_size; ++i) {
115117
dst.dim_metadata[i] = src.dim_metadata[i];
116118
dst.dim_metadata[i].array_segments =
@@ -129,6 +131,7 @@ TfLiteSparsity* TfLiteSparsityClone(const TfLiteSparsity* const src) {
129131
}
130132
TfLiteSparsity* dst =
131133
reinterpret_cast<TfLiteSparsity*>(calloc(1, sizeof(TfLiteSparsity)));
134+
if (!dst) return nullptr;
132135
*dst = TfLiteSparsityClone(*src);
133136
return dst;
134137
}
@@ -147,6 +150,7 @@ TfLiteQuantization TfLiteQuantizationClone(const TfLiteQuantization& src) {
147150
break;
148151
case kTfLiteAffineQuantization: {
149152
dst.params = calloc(1, sizeof(TfLiteAffineQuantization));
153+
if (!dst.params) return TfLiteQuantization();
150154
const TfLiteAffineQuantization* const src_params =
151155
reinterpret_cast<TfLiteAffineQuantization*>(src.params);
152156
TfLiteAffineQuantization* const dst_params =
@@ -158,6 +162,7 @@ TfLiteQuantization TfLiteQuantizationClone(const TfLiteQuantization& src) {
158162
}
159163
case kTfLiteBlockwiseQuantization: {
160164
dst.params = calloc(1, sizeof(TfLiteBlockwiseQuantization));
165+
if (!dst.params) return TfLiteQuantization();
161166
const TfLiteBlockwiseQuantization* const src_params =
162167
(TfLiteBlockwiseQuantization*)(src.params);
163168
TfLiteBlockwiseQuantization* const dst_params =
@@ -219,6 +224,9 @@ TfLiteFloatArray* TfLiteFloatArrayCopy(const TfLiteFloatArray* src) {
219224
void TfLiteFloatArrayFree(TfLiteFloatArray* a) { TfLiteVarArrayFree(a); }
220225

221226
void TfLiteTensorDataFree(TfLiteTensor* t) {
227+
if (t == nullptr) {
228+
return;
229+
}
222230
if (t->allocation_type == kTfLiteVariantObject && t->data.data) {
223231
delete static_cast<VariantData*>(t->data.data);
224232
} else if (t->allocation_type == kTfLiteDynamic ||
@@ -238,6 +246,9 @@ void TfLiteTensorDataFree(TfLiteTensor* t) {
238246
}
239247

240248
void TfLiteQuantizationFree(TfLiteQuantization* quantization) {
249+
if (quantization == nullptr) {
250+
return;
251+
}
241252
if (quantization->type == kTfLiteAffineQuantization) {
242253
TfLiteAffineQuantization* q_params =
243254
reinterpret_cast<TfLiteAffineQuantization*>(quantization->params);
@@ -294,6 +305,9 @@ void TfLiteSparsityFree(TfLiteSparsity* sparsity) {
294305
}
295306

296307
void TfLiteTensorFree(TfLiteTensor* t) {
308+
if (t == nullptr) {
309+
return;
310+
}
297311
TfLiteTensorDataFree(t);
298312
if (t->dims) TfLiteIntArrayFree(t->dims);
299313
t->dims = nullptr;
@@ -308,7 +322,7 @@ void TfLiteTensorFree(TfLiteTensor* t) {
308322
t->sparsity = nullptr;
309323
}
310324

311-
TfLiteTensor TfLiteTensorClone(const TfLiteTensor src) {
325+
TfLiteTensor TfLiteTensorClone(TfLiteTensor src) {
312326
// We copy all of the source data first, then we clone the fields that can't
313327
// be shared between two tensor instances.
314328
TfLiteTensor dst = src;
@@ -335,16 +349,18 @@ TfLiteTensor TfLiteTensorClone(const TfLiteTensor src) {
335349
break;
336350
case kTfLiteAllocationStrategyMalloc:
337351
dst.data.data = malloc(src.bytes);
352+
if (src.bytes > 0 && !dst.data.data) return TfLiteTensor();
338353
std::memcpy(dst.data.data, src.data.data, src.bytes);
339354
break;
340355
case kTfLiteAllocationStrategyNew:
341356
// Special case for variant objects. They are allocated using new/delete
342357
// but require using the `CloneTo` function.
343358
if (src.allocation_type == kTfLiteVariantObject) {
344-
dst.data.data = reinterpret_cast<const VariantData*>(src.data.data)
345-
->CloneTo(nullptr);
359+
dst.data.data =
360+
static_cast<const VariantData*>(src.data.data)->CloneTo(nullptr);
346361
} else {
347-
dst.data.data = new char[src.bytes];
362+
dst.data.data = new (std::nothrow) char[src.bytes];
363+
if (src.bytes > 0 && !dst.data.data) return TfLiteTensor();
348364
std::memcpy(dst.data.data, src.data.data, src.bytes);
349365
}
350366
break;
@@ -394,13 +410,21 @@ TfLiteStatus TfLiteTensorCopy(const TfLiteTensor* src, TfLiteTensor* dst) {
394410
}
395411
auto* dst_vd = static_cast<VariantData*>(dst->data.data);
396412
auto* src_vd = static_cast<VariantData*>(src->data.data);
413+
if (!src_vd) return kTfLiteError;
397414

398415
// `CloneTo` will handle the case when `dst_vd` is nullptr, so it is safe
399416
// to `CloneTo` something which was "freed". Also, returning from `CloneTo`
400417
// will implicitly cast to `VariantData`; don't need static cast here.
401418
dst->data.data = src_vd->CloneTo(dst_vd);
402419
} else {
403-
memcpy(dst->data.raw, src->data.raw, src->bytes);
420+
if (dst->allocation_type == kTfLiteVariantObject) {
421+
TfLiteTensorDataFree(dst);
422+
dst->allocation_type = src->allocation_type;
423+
}
424+
if (src->bytes > 0) {
425+
if (!dst->data.raw || !src->data.raw) return kTfLiteError;
426+
memcpy(dst->data.raw, src->data.raw, src->bytes);
427+
}
404428
}
405429
dst->buffer_handle = src->buffer_handle;
406430
dst->data_is_stale = src->data_is_stale;
@@ -513,6 +537,10 @@ const char* TfLiteTypeGetName(TfLiteType type) {
513537
return "INT2";
514538
case kTfLiteUInt4:
515539
return "UINT4";
540+
case kTfLiteFloat8E4M3FN:
541+
return "FLOAT8_E4M3FN";
542+
case kTfLiteFloat8E5M2:
543+
return "FLOAT8_E5M2";
516544
}
517545
return "Unknown type";
518546
}

tensorflow/lite/core/c/common.h

Lines changed: 35 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ limitations under the License.
5656
#include <stdbool.h>
5757
#include <stddef.h>
5858
#include <stdint.h>
59+
#include <stdio.h>
5960

6061
#include "tensorflow/lite/core/c/c_api_types.h" // IWYU pragma: export
6162

@@ -277,13 +278,34 @@ void TfLiteFloatArrayFree(TfLiteFloatArray* a);
277278
} \
278279
} while (0)
279280

280-
#define TF_LITE_ENSURE_OK(context, status) \
281-
do { \
282-
const TfLiteStatus s = (status); \
283-
if ((s) != kTfLiteOk) { \
284-
return s; \
285-
} \
281+
#ifndef TF_LITE_STRIP_ERROR_STRINGS
282+
#define TF_LITE_VAR_ARG_HEAD(FIRST, ...) FIRST
283+
#define TF_LITE_STRINGIFY_HELPER(x) #x
284+
#define TF_LITE_STRINGIFY(x) TF_LITE_STRINGIFY_HELPER(x)
285+
// Checks that `status` evaluates to `kTfLiteOk`.
286+
//
287+
// Can take a printf style log message and its parameters after the status. The
288+
// message will be printed using `TF_LITE_KERNEL_LOG` in case of error.
289+
#define TF_LITE_ENSURE_OK(context, status, ...) \
290+
do { \
291+
const TfLiteStatus s = (status); \
292+
if (s != kTfLiteOk) { \
293+
if (sizeof(TF_LITE_VAR_ARG_HEAD("" __VA_ARGS__)) > sizeof("")) { \
294+
TF_LITE_MAYBE_KERNEL_LOG((context), __FILE__ ":" TF_LITE_STRINGIFY( \
295+
__LINE__) ": " __VA_ARGS__); \
296+
} \
297+
return s; \
298+
} \
286299
} while (0)
300+
#else
301+
#define TF_LITE_ENSURE_OK(context, status, ...) \
302+
do { \
303+
const TfLiteStatus s = (status); \
304+
if ((s) != kTfLiteOk) { \
305+
return s; \
306+
} \
307+
} while (0)
308+
#endif
287309

288310
// `std::unreachable` not available until CC23.
289311
#ifdef __GNUC__ // GCC, Clang, ICC
@@ -1060,6 +1082,13 @@ typedef struct TfLiteContext {
10601082
/// WARNING: This is an experimental interface that is subject to change.
10611083
TfLiteStatus (*ReleaseSubgraphContext)(struct TfLiteContext* context,
10621084
int subgraph_index);
1085+
#if defined(_WIN32)
1086+
/// Create a array of a given `size` (uninitialized entries).
1087+
TfLiteIntArray* (*TfLiteIntArrayCreate)(int size); // NOLINT
1088+
1089+
/// Free memory of array `a`.
1090+
void (*TfLiteIntArrayFree)(TfLiteIntArray* a); // NOLINT
1091+
#endif // defined(_WIN32)
10631092
} TfLiteContext;
10641093

10651094
/// `TfLiteOperator` is an external version of `TfLiteRegistration`

tensorflow/lite/kernels/internal/common.h

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,15 +78,23 @@ bool ReduceDimensionsForBroadcast(const RuntimeShape& input1_shape,
7878
if (!broadcast_input1) {
7979
broadcast_input1 = true;
8080
broadcast_input2 = false;
81+
if (num_compressed_dims >= MAX_DIM) return false;
8182
num_compressed_dims++;
83+
if (num_compressed_dims > MAX_DIM) {
84+
return false;
85+
}
8286
}
8387
compressed_input2_shape[num_compressed_dims - 1] *= input2_dim;
8488
compressed_output_shape[num_compressed_dims - 1] *= input2_dim;
8589
} else if (input2_dim == 1) {
8690
if (!broadcast_input2) {
8791
broadcast_input1 = false;
8892
broadcast_input2 = true;
93+
if (num_compressed_dims >= MAX_DIM) return false;
8994
num_compressed_dims++;
95+
if (num_compressed_dims > MAX_DIM) {
96+
return false;
97+
}
9098
}
9199
compressed_input1_shape[num_compressed_dims - 1] *= input1_dim;
92100
compressed_output_shape[num_compressed_dims - 1] *= input1_dim;
@@ -95,7 +103,11 @@ bool ReduceDimensionsForBroadcast(const RuntimeShape& input1_shape,
95103
if (broadcast_input1 || broadcast_input2 || first_nonunit) {
96104
broadcast_input1 = false;
97105
broadcast_input2 = false;
106+
if (num_compressed_dims >= MAX_DIM) return false;
98107
num_compressed_dims++;
108+
if (num_compressed_dims > MAX_DIM) {
109+
return false;
110+
}
99111
}
100112
compressed_input1_shape[num_compressed_dims - 1] *= input1_dim;
101113
compressed_input2_shape[num_compressed_dims - 1] *= input1_dim;
@@ -105,7 +117,11 @@ bool ReduceDimensionsForBroadcast(const RuntimeShape& input1_shape,
105117
}
106118
if (num_input1_dims > num_input2_dims) {
107119
if (!broadcast_input2) {
120+
if (num_compressed_dims >= MAX_DIM) return false;
108121
num_compressed_dims++;
122+
if (num_compressed_dims > MAX_DIM) {
123+
return false;
124+
}
109125
}
110126
for (size_t i = 0; i < num_input1_dims - num_input2_dims; i++) {
111127
const size_t input1_dim = input1_dims[i];
@@ -117,7 +133,11 @@ bool ReduceDimensionsForBroadcast(const RuntimeShape& input1_shape,
117133
}
118134
} else if (num_input2_dims > num_input1_dims) {
119135
if (!broadcast_input1) {
136+
if (num_compressed_dims >= MAX_DIM) return false;
120137
num_compressed_dims++;
138+
if (num_compressed_dims > MAX_DIM) {
139+
return false;
140+
}
121141
}
122142
for (size_t i = 0; i < num_input2_dims - num_input1_dims; i++) {
123143
const size_t input2_dim = input2_dims[i];

0 commit comments

Comments
 (0)