WHILE operator input/output copy fix (#3633)

@tensorflow/micro

Remove extraneous tensor copy operation after first invocation of condition subgraph.

Move copy of operator inputs to outputs, such that it occurs before the first invocation of the condition subgraph. This preserves the operator inputs when one or more of them is the output of DECODE, and alternate decompression memory is in use. This is because the output of DECODE is for immediate consumption by the next operator in the graph (WHILE), yet it is possible for the WHILE subgraph invocations to share memory with the original DECODE output.

Update the unit test for multiple invocations of the condition and body subgraphs.

When copying tensors between operator inputs/outputs and subgraph inputs/outputs, check if the source and destination tensors share memory.

bug=fixes #3632

Author: ddavis-2015

tensorflow/lite/micro/kernels/kernel_util.cc

@@ -166,7 +166,10 @@ TfLiteStatus CopyOpInputsToOpOutputs(TfLiteContext* context, TfLiteNode* node) {
     TfLiteEvalTensor* output = tflite::micro::GetEvalOutput(context, node, i);
     int bytes = ValidateAndGetTensorSizes(input, output);
     TF_LITE_ENSURE(context, bytes >= 0);
-    memcpy(output->data.raw, input->data.raw, bytes);
+    // Don't copy if the tensors share memory
+    if (output->data.data != input->data.data) {
+      memcpy(output->data.data, input->data.data, bytes);
+    }
   }
   return kTfLiteOk;
 }
@@ -226,7 +229,10 @@ TfLiteStatus CopyOpInputsToSubgraphInputs(TfLiteContext* context,
         graph_info->GetSubgraphInput(subgraph_idx, i);
     int bytes = ValidateAndGetTensorSizes(input, subgraph_input);
     TF_LITE_ENSURE(context, bytes >= 0);
-    memcpy(subgraph_input->data.raw, input->data.raw, bytes);
+    // Don't copy if the tensors share memory
+    if (subgraph_input->data.data != input->data.data) {
+      memcpy(subgraph_input->data.data, input->data.data, bytes);
+    }
   }
   return kTfLiteOk;
 }
@@ -243,7 +249,10 @@ TfLiteStatus CopyOpOutputsToSubgraphInputs(TfLiteContext* context,
         graph_info->GetSubgraphInput(subgraph_idx, i);
     int bytes = ValidateAndGetTensorSizes(output, subgraph_input);
     TF_LITE_ENSURE(context, bytes >= 0);
-    memcpy(subgraph_input->data.raw, output->data.raw, bytes);
+    // Don't copy if the tensors share memory
+    if (subgraph_input->data.data != output->data.data) {
+      memcpy(subgraph_input->data.data, output->data.data, bytes);
+    }
   }
   return kTfLiteOk;
 }
@@ -261,7 +270,10 @@ TfLiteStatus CopySubgraphOutputsToOpOutputs(TfLiteContext* context,
         graph_info->GetSubgraphOutput(subgraph_idx, i);
     int bytes = ValidateAndGetTensorSizes(output, subgraph_output);
     TF_LITE_ENSURE(context, bytes >= 0);
-    memcpy(output->data.raw, subgraph_output->data.raw, bytes);
+    // Don't copy if the tensors share memory
+    if (output->data.data != subgraph_output->data.data) {
+      memcpy(output->data.data, subgraph_output->data.data, bytes);
+    }
   }
   return kTfLiteOk;
 }

tensorflow/lite/micro/kernels/while.cc

@@ -85,20 +85,20 @@ TfLiteStatus WhileEval(TfLiteContext* context, TfLiteNode* node) {
                         context, node, graph_info, op_data->cond_subgraph_index,
                         /*first_tensor_idx=*/0));
 
+  // Preserve the op inputs by copying to op outputs, prior to invoking
+  // a subgraph which could invalidate the memory of one or more op inputs.
+  // This is possible when using the output of the DECODE operator as input to
+  // the WHILE op, in the presence of alternate decompression memory.
+  TF_LITE_ENSURE_OK(context,
+                    tflite::micro::CopyOpInputsToOpOutputs(context, node));
+
   TF_LITE_ENSURE_OK(context,
                     graph_info->InvokeSubgraph(op_data->cond_subgraph_index));
 
   TfLiteEvalTensor* cond_subgraph_output = graph_info->GetSubgraphOutput(
       op_data->cond_subgraph_index, /*tensor_idx=*/0);
   bool cond_value = cond_subgraph_output->data.b[0];
 
-  TF_LITE_ENSURE_OK(context,
-                    tflite::micro::CopyOpInputsToSubgraphInputs(
-                        context, node, graph_info, op_data->body_subgraph_index,
-                        /*first_tensor_idx=*/0));
-  TF_LITE_ENSURE_OK(context,
-                    tflite::micro::CopyOpInputsToOpOutputs(context, node));
-
   while (cond_value == true) {
     // Copy output of this iteration back to the body input.
     TF_LITE_ENSURE_OK(

tensorflow/lite/micro/kernels/while_test.cc

@@ -73,4 +73,29 @@ TEST(WhileTest, WhileShouldInvokeOnce) {
   EXPECT_EQ(output1->data.f[0], 3.0f);
 }
 
+TEST(WhileTest, WhileShouldInvokeMultiple) {
+  constexpr int kArenaSize = 5000;
+  uint8_t arena[kArenaSize];
+
+  const tflite::Model* model =
+      tflite::testing::GetSimpleModelWithSubgraphsAndWhile();
+  tflite::MicroMutableOpResolver<3> resolver;
+  resolver.AddWhile();
+  resolver.AddAdd();
+  resolver.AddLess();
+  tflite::MicroInterpreter interpreter(model, resolver, arena, kArenaSize);
+  EXPECT_EQ(kTfLiteOk, interpreter.AllocateTensors());
+  TfLiteTensor* input0 = interpreter.input(0);
+  TfLiteTensor* input1 = interpreter.input(1);
+  TfLiteTensor* output0 = interpreter.output(0);
+  TfLiteTensor* output1 = interpreter.output(1);
+  input0->data.f[0] = -5.0f;
+  input1->data.f[0] = 3.0f;
+
+  interpreter.Invoke();
+
+  EXPECT_EQ(output0->data.f[0], 4.0f);
+  EXPECT_EQ(output1->data.f[0], 3.0f);
+}
+
 TF_LITE_MICRO_TESTS_MAIN