Update cmsis prepare functions to exit earlier (#2630)
* Adds additional checks to various cmsis-nn prepare functions
BUG=1121
Authored-by: Ryan O'Shea <ryan.oshea3@arm.com>
Change-Id: Ic6481873f064a94a4dd0b4a49790842180d73dd9
diff --git a/tensorflow/lite/micro/kernels/cmsis_nn/add.cc b/tensorflow/lite/micro/kernels/cmsis_nn/add.cc
index 898410a..fb166a1 100644
--- a/tensorflow/lite/micro/kernels/cmsis_nn/add.cc
+++ b/tensorflow/lite/micro/kernels/cmsis_nn/add.cc
@@ -1,4 +1,4 @@
-/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -301,6 +301,15 @@
micro_context->AllocateTempOutputTensor(node, kOutputTensor);
TF_LITE_ENSURE(context, output != nullptr);
+ TF_LITE_ENSURE_EQ(context, input1->type, output->type);
+ TF_LITE_ENSURE_MSG(
+ context,
+ input1->type == kTfLiteFloat32 || input1->type == kTfLiteInt32 ||
+ input1->type == kTfLiteInt16 || input1->type == kTfLiteInt8,
+ "Input data type not supported");
+ TF_LITE_ENSURE_MSG(context, input1->type == input2->type,
+ "Hybrid models are not supported on TFLite Micro.");
+
if (input1->type == kTfLiteInt16) {
TF_LITE_ENSURE_EQ(context, input1->params.zero_point, 0);
TF_LITE_ENSURE_EQ(context, input2->params.zero_point, 0);
diff --git a/tensorflow/lite/micro/kernels/cmsis_nn/conv.cc b/tensorflow/lite/micro/kernels/cmsis_nn/conv.cc
index 4c35970..cae68c7 100644
--- a/tensorflow/lite/micro/kernels/cmsis_nn/conv.cc
+++ b/tensorflow/lite/micro/kernels/cmsis_nn/conv.cc
@@ -67,11 +67,17 @@
TfLiteType bias_type = bias != nullptr ? bias->type : kTfLiteNoType;
TF_LITE_ENSURE_EQ(context, input->type, output->type);
+ TF_LITE_ENSURE_MSG(context,
+ input->type == kTfLiteFloat32 ||
+ input->type == kTfLiteInt16 ||
+ input->type == kTfLiteInt8,
+ "Input data type not supported");
TF_LITE_ENSURE_MSG(
context,
- input->type == filter->type ||
+ (input->type == kTfLiteFloat32 && filter->type == kTfLiteFloat32) ||
(input->type == kTfLiteInt16 && filter->type == kTfLiteInt8) ||
- (input->type == kTfLiteInt8 && filter->type == kTfLiteInt4),
+ (input->type == kTfLiteInt8 &&
+ (filter->type == kTfLiteInt4 || filter->type == kTfLiteInt8)),
"Hybrid models are not supported on TFLite Micro.");
// Consistency check tensor dims
diff --git a/tensorflow/lite/micro/kernels/cmsis_nn/depthwise_conv.cc b/tensorflow/lite/micro/kernels/cmsis_nn/depthwise_conv.cc
index f30a952..7183a28 100644
--- a/tensorflow/lite/micro/kernels/cmsis_nn/depthwise_conv.cc
+++ b/tensorflow/lite/micro/kernels/cmsis_nn/depthwise_conv.cc
@@ -1,4 +1,4 @@
-/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -75,6 +75,20 @@
micro_context->AllocateTempOutputTensor(node, kDepthwiseConvOutputTensor);
TF_LITE_ENSURE(context, output != nullptr);
+ TF_LITE_ENSURE_EQ(context, input->type, output->type);
+ TF_LITE_ENSURE_MSG(context,
+ input->type == kTfLiteFloat32 ||
+ input->type == kTfLiteInt16 ||
+ input->type == kTfLiteInt8,
+ "Input data type not supported");
+ TF_LITE_ENSURE_MSG(
+ context,
+ (input->type == kTfLiteFloat32 && filter->type == kTfLiteFloat32) ||
+ (input->type == kTfLiteInt16 && filter->type == kTfLiteInt8) ||
+ (input->type == kTfLiteInt8 &&
+ (filter->type == kTfLiteInt4 || filter->type == kTfLiteInt8)),
+ "Hybrid models are not supported on TFLite Micro.");
+
const TfLiteType data_type = input->type;
int input_width = SizeOfDimension(input, 2);
int input_height = SizeOfDimension(input, 1);
diff --git a/tensorflow/lite/micro/kernels/cmsis_nn/fully_connected.cc b/tensorflow/lite/micro/kernels/cmsis_nn/fully_connected.cc
index dc7b78c..7c373b5 100644
--- a/tensorflow/lite/micro/kernels/cmsis_nn/fully_connected.cc
+++ b/tensorflow/lite/micro/kernels/cmsis_nn/fully_connected.cc
@@ -76,7 +76,19 @@
node, kFullyConnectedOutputTensor);
TF_LITE_ENSURE(context, output != nullptr);
- TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
+ TF_LITE_ENSURE_EQ(context, input->type, output->type);
+ TF_LITE_ENSURE_MSG(context,
+ input->type == kTfLiteFloat32 ||
+ input->type == kTfLiteInt16 ||
+ input->type == kTfLiteInt8,
+ "Input data type not supported");
+ TF_LITE_ENSURE_MSG(
+ context,
+ (input->type == kTfLiteFloat32 && filter->type == kTfLiteFloat32) ||
+ (input->type == kTfLiteInt16 && filter->type == kTfLiteInt8) ||
+ (input->type == kTfLiteInt8 &&
+ (filter->type == kTfLiteInt4 || filter->type == kTfLiteInt8)),
+ "Hybrid models are not supported on TFLite Micro.");
const RuntimeShape filter_shape = GetTensorShape(filter);
const RuntimeShape output_shape = GetTensorShape(output);
@@ -125,7 +137,7 @@
input_dims.c = data->accum_depth;
buf_size = arm_convolve_1x1_s8_fast_get_buffer_size(&input_dims);
- } else {
+ } else if (input->type == kTfLiteInt8) {
buf_size = arm_fully_connected_s8_get_buffer_size(&filter_dims);
int8_t* filter_data = GetTensorData<int8_t>(filter);
diff --git a/tensorflow/lite/micro/kernels/cmsis_nn/softmax.cc b/tensorflow/lite/micro/kernels/cmsis_nn/softmax.cc
index f83a090..6515691 100644
--- a/tensorflow/lite/micro/kernels/cmsis_nn/softmax.cc
+++ b/tensorflow/lite/micro/kernels/cmsis_nn/softmax.cc
@@ -1,4 +1,4 @@
-/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -52,6 +52,17 @@
TfLiteTensor* output = micro_context->AllocateTempOutputTensor(node, 0);
TF_LITE_ENSURE(context, output != nullptr);
+ TF_LITE_ENSURE_MSG(
+ context,
+ input->type == output->type ||
+ (input->type == kTfLiteInt8 && output->type == kTfLiteInt16),
+ "Input and output data types are not supported together.");
+ TF_LITE_ENSURE_MSG(context,
+ input->type == kTfLiteFloat32 ||
+ input->type == kTfLiteInt16 ||
+ input->type == kTfLiteInt8,
+ "Input data type not supported");
+
TF_LITE_ENSURE(context, node->user_data != nullptr);
CMSISNNSoftmaxParams* op_data =
static_cast<CMSISNNSoftmaxParams*>(node->user_data);
diff --git a/tensorflow/lite/micro/kernels/cmsis_nn/transpose_conv.cc b/tensorflow/lite/micro/kernels/cmsis_nn/transpose_conv.cc
index 06305bc..20cf0e1 100644
--- a/tensorflow/lite/micro/kernels/cmsis_nn/transpose_conv.cc
+++ b/tensorflow/lite/micro/kernels/cmsis_nn/transpose_conv.cc
@@ -174,10 +174,17 @@
micro_context->AllocateTempInputTensor(node, kFilterTensor);
TF_LITE_ENSURE(context, filter != nullptr);
+ TF_LITE_ENSURE_EQ(context, input->type, output->type);
+ TF_LITE_ENSURE_MSG(context,
+ input->type == kTfLiteFloat32 ||
+ input->type == kTfLiteInt16 ||
+ input->type == kTfLiteInt8,
+ "Input data type not supported");
TF_LITE_ENSURE_MSG(
context,
- input->type == filter->type ||
- (input->type == kTfLiteInt16 && filter->type == kTfLiteInt8),
+ (input->type == kTfLiteFloat32 && filter->type == kTfLiteFloat32) ||
+ (input->type == kTfLiteInt16 && filter->type == kTfLiteInt8) ||
+ (input->type == kTfLiteInt8 && filter->type == kTfLiteInt8),
"Hybrid models are not supported on TFLite Micro.");
// Get height and width of the output.
