Additional const-tensor checks (CMSIS_NN) (#3229)

@tensorflow/micro

Add additional const-tensor checks to the following CMSIS_NN kernels that attempt to use tensor data during the Prepare phase:

FULLY_CONNECTED
SVDF
UNIDIRECTIONAL_SEQUENCE_LSTM

Unit tests updated.

bug=fixes #3228

Co-authored-by: Esun Kim <[email protected]>

Author: ddavis-2015

tensorflow/lite/micro/kernels/cmsis_nn/fully_connected.cc

@@ -151,9 +151,9 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
       data->kernel_sums = nullptr;
 
 #if defined(KERNELS_OPTIMIZED_FOR_SPEED)
-      const int8_t* filter_data = GetTensorData<const int8_t>(filter);
-
-      if (buf_size > 0 && filter_data != nullptr) {
+      if (buf_size > 0 && IsConstantTensor(filter) &&
+          (bias == nullptr || IsConstantTensor(bias))) {
+        const int8_t* filter_data = GetTensorData<const int8_t>(filter);
         const int32_t input_offset = -data->reference_op_data.input_zero_point;
         const int32_t filter_offset =
             -data->reference_op_data.filter_zero_point;

tensorflow/lite/micro/kernels/cmsis_nn/svdf.cc

@@ -193,6 +193,7 @@ TfLiteStatus CmsisNnPrepareSvdf(TfLiteContext* context, TfLiteNode* node) {
 
     if (buf_size > 0) {
 #if defined(KERNELS_OPTIMIZED_FOR_SPEED)
+      TF_LITE_ENSURE(context, IsConstantTensor(weights_feature));
       data->kernel_sums = static_cast<int32_t*>(
           context->AllocatePersistentBuffer(context, buf_size));
 
@@ -210,6 +211,12 @@ TfLiteStatus CmsisNnPrepareSvdf(TfLiteContext* context, TfLiteNode* node) {
           "Either KERNELS_OPTIMIZED_FOR_SIZE or KERNELS_OPTIMIZED_FOR_SPEED "
           "must be defined");
       return kTfLiteError;
+#endif
+    } else {
+      // safety first!
+      data->kernel_sums = nullptr;
+#if defined(KERNELS_OPTIMIZED_FOR_SIZE)
+      data->scratch_weight_tensor_index = -1;
 #endif
     }
 
@@ -310,6 +317,7 @@ TfLiteStatus EvalIntegerSVDF(TfLiteContext* context, TfLiteNode* node,
 #if defined(KERNELS_OPTIMIZED_FOR_SPEED)
       ctx.buf = data.kernel_sums;
 #elif defined(KERNELS_OPTIMIZED_FOR_SIZE)
+      TF_LITE_ENSURE(context, data.scratch_weight_tensor_index != -1);
       ctx.buf = static_cast<int32_t*>(
           context->GetScratchBuffer(context, data.scratch_weight_tensor_index));
 

tensorflow/lite/micro/kernels/cmsis_nn/unidirectional_sequence_lstm.cc

@@ -344,6 +344,18 @@ TfLiteStatus UnidirectionalSequenceLstmPrepare(TfLiteContext* context,
   // All TempTfLiteTensors will be deallocated through the destructor.
   LstmTensors lstm_tensors(context, node);
   TF_LITE_ENSURE_OK(context, lstm_tensors.ValidateTensorStatus(context));
+  // Additional validation of weights and biases.
+  // ValidateTensorStatus() ensures no tensor is <nullptr>.
+  for (size_t i = 1; i < 9; i++) {
+    // check weight
+    TF_LITE_ENSURE(context,
+                   IsConstantTensor(lstm_tensors.GetInternalTensor(i)));
+  }
+  for (size_t i = 12; i < 16; i++) {
+    // check bias
+    TF_LITE_ENSURE(context,
+                   IsConstantTensor(lstm_tensors.GetInternalTensor(i)));
+  }
 
   op_data_lstm->cell_gate_nonlinear_type = builtin_data->activation;
   op_data_lstm->size_info =

tensorflow/lite/micro/kernels/fully_connected_test.cc

@@ -307,6 +307,11 @@ TfLiteStatus ValidateFullyConnectedGoldens(
   TfLiteIntArray* inputs_array = IntArrayFromInts(inputs_array_data);
   TfLiteIntArray* outputs_array = IntArrayFromInts(outputs_array_data);
 
+  tensors[1].allocation_type = kTfLiteMmapRo;
+  if (!null_bias) {
+    tensors[2].allocation_type = kTfLiteMmapRo;
+  }
+
 #ifdef USE_TFLM_COMPRESSION
 
   TestCompressedList<kMaxTensors> tcl;

tensorflow/lite/micro/kernels/svdf_test.cc

@@ -498,6 +498,8 @@ void ValidateSVDFGoldens(const int batch_size, const int num_units,
   int outputs_array_data[] = {1, 5};
   TfLiteIntArray* outputs_array = IntArrayFromInts(outputs_array_data);
 
+  tensors[1].allocation_type = kTfLiteMmapRo;
+
   const TFLMRegistration registration = Register_SVDF();
   micro::KernelRunner runner(registration, tensors, tensor_count, inputs_array,
                              outputs_array, &params);

tensorflow/lite/micro/kernels/unidirectional_sequence_lstm_test.cc

@@ -29,6 +29,18 @@ namespace {
 
 constexpr int kLstmMaxNumInputOutputTensors = 24 + 1;
 
+// Set weights and biases to be const-tensors
+void SetConstTensors(TfLiteTensor* tensors) {
+  for (size_t i = 1; i < 9; i++) {
+    // weights
+    tensors[i].allocation_type = kTfLiteMmapRo;
+  }
+  for (size_t i = 12; i < 16; i++) {
+    // biases
+    tensors[i].allocation_type = kTfLiteMmapRo;
+  }
+}
+
 // Validate the output result array with golden values
 template <typename T>
 void ValidateResultGoldens(const T* golden, const T* output_data,
@@ -49,6 +61,7 @@ void TestUnidirectionalLSTMInteger(
     LstmNodeContent<ActivationType, WeightType, BiasType, CellType, batch_size,
                     time_steps, input_dimension, state_dimension>&
         node_contents) {
+  SetConstTensors(node_contents.GetTensors());
   const TFLMRegistration registration = Register_UNIDIRECTIONAL_SEQUENCE_LSTM();
   auto buildin_data = node_contents.BuiltinData();
   micro::KernelRunner runner(
@@ -101,6 +114,7 @@ void TestUnidirectionalLSTMFloat(
     const float hidden_state_tolerance, const float cell_state_tolerance,
     LstmNodeContent<float, float, float, float, batch_size, time_steps,
                     input_dimension, state_dimension>& node_contents) {
+  SetConstTensors(node_contents.GetTensors());
   const TFLMRegistration registration = Register_UNIDIRECTIONAL_SEQUENCE_LSTM();
   auto buildin_data = node_contents.BuiltinData();
   micro::KernelRunner runner(