diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc index 881b9b9abb0..588e0674241 100644 --- a/tensorflow/lite/micro/micro_allocator.cc +++ b/tensorflow/lite/micro/micro_allocator.cc @@ -242,6 +242,18 @@ TfLiteStatus AllocationInfoBuilder::AddTensors(const SubGraph* subgraph, for (size_t n = 0; n < op->inputs()->size(); ++n) { const int tensor_index = op->inputs()->Get(n); AllocationInfo* current = &info_[tensor_index]; + + // In case operator input are not in subgraph inputs initialize them. + if (current->first_created == 0) { + for (size_t op_input = 0; op_input < op->inputs()->size(); ++op_input) { + const int op_tensor_index = op->inputs()->Get(op_input); + AllocationInfo* op_current = &info_[op_tensor_index]; + if (op_current->needs_allocating && op_current->first_created == -1) { + op_current->first_created = i; + } + } + } + if (((current->last_used == -1) || (current->last_used < i))) { current->last_used = i; } diff --git a/tensorflow/lite/micro/micro_allocator_test.cc b/tensorflow/lite/micro/micro_allocator_test.cc index 32d52a994d9..5e8f1b5df31 100644 --- a/tensorflow/lite/micro/micro_allocator_test.cc +++ b/tensorflow/lite/micro/micro_allocator_test.cc @@ -735,4 +735,64 @@ TF_LITE_MICRO_TEST(TestAllocateTfLiteTensorWithReset) { TF_LITE_MICRO_EXPECT(tensor2 == tensor1); } +TF_LITE_MICRO_TEST(TestOperatorInputsNotInSubgraphInputs) { + constexpr int nbr_tensors = 5; + tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver(); + tflite::NodeAndRegistration* node_and_registration; + const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize + + nbr_tensors] = { + 1, 0, nbr_tensors, // header: version, subgraph, nbr tensors + // memory offsets: + 0, // t0 + 0, // t1 + 0, // t2 + 48, // t3 + -1}; // t4 + + int t0 = 0; + int t1 = 1; + int t2 = 2; + int t3 = 3; + int t4 = 4; + + int num_conns = 2; + tflite::testing::NodeConnection node_list[2] = { + { + {t0, t1, t2}, // t0: input (actual input part of subgraph inputs as well as + // operator inputs) + // t1: scratch1 (only in operator inputs) + // t2: scratch2 (only in operator inputs) + {t3} // output + }, + { + {t3}, // input + {t4} // output + }, + }; + + const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning( + nbr_tensors, metadata_buffer, node_list, num_conns, + 1/* only first tensor (t0) is in subgraph input list*/); + + TfLiteEvalTensor* eval_tensors = nullptr; + constexpr size_t arena_size = 4096; + uint8_t arena[arena_size]; + tflite::MicroAllocator* allocator = + tflite::MicroAllocator::Create(arena, arena_size, micro_test::reporter); + + TF_LITE_MICRO_EXPECT_EQ( + kTfLiteOk, + allocator->StartModelAllocation(model, op_resolver, + &node_and_registration, &eval_tensors)); + TF_LITE_MICRO_EXPECT_EQ( + kTfLiteOk, allocator->FinishModelAllocation(model, eval_tensors)); + + uint8_t* start = eval_tensors[0].data.uint8; + TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[0].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[1].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[2].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(48, eval_tensors[3].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[4].data.uint8 - start); +} + TF_LITE_MICRO_TESTS_END diff --git a/tensorflow/lite/micro/test_helpers.cc b/tensorflow/lite/micro/test_helpers.cc index 23c7ca96408..f3450f05c2f 100644 --- a/tensorflow/lite/micro/test_helpers.cc +++ b/tensorflow/lite/micro/test_helpers.cc @@ -108,7 +108,8 @@ class ModelBuilder { // Constructs the flatbuffer model using `builder_` and return a pointer to // it. The returned model has the same lifetime as `builder_`. const Model* BuildModel(std::initializer_list inputs, - std::initializer_list outputs); + std::initializer_list outputs, + size_t num_subgraph_inputs = 0); private: // Adds a tensor to the model. @@ -179,7 +180,8 @@ void ModelBuilder::AddMetadata(const char* description_string, const Model* ModelBuilder::BuildModel( std::initializer_list inputs, - std::initializer_list outputs) { + std::initializer_list outputs, + size_t num_subgraph_inputs) { // Model schema requires an empty buffer at idx 0. size_t buffer_size = 1 + ModelBuilder::nbr_of_metadata_buffers_; flatbuffers::Offset buffers[kMaxMetadataBuffers]; @@ -193,10 +195,17 @@ const Model* ModelBuilder::BuildModel( // TFLM only supports single subgraph. constexpr size_t subgraphs_size = 1; + + // Find out number of subgraph inputs. + int num_tensors_in_subgraph_inputs = inputs.size(); // Default case = all inputs + if (num_subgraph_inputs > 0 && num_subgraph_inputs < inputs.size()) { + num_tensors_in_subgraph_inputs = num_subgraph_inputs; + } + const flatbuffers::Offset subgraphs[subgraphs_size] = { tflite::CreateSubGraph( *builder_, builder_->CreateVector(tensors_, next_tensor_id_), - builder_->CreateVector(inputs.begin(), inputs.size()), + builder_->CreateVector(inputs.begin(), num_tensors_in_subgraph_inputs), builder_->CreateVector(outputs.begin(), outputs.size()), builder_->CreateVector(operators_, next_operator_id_), builder_->CreateString("test_subgraph"))}; @@ -301,7 +310,8 @@ const Model* BuildSimpleModelWithBranch() { const Model* BuildModelWithOfflinePlanning(int number_of_tensors, const int32_t* metadata_buffer, NodeConnection* node_conn, - int num_conns) { + int num_conns, + int num_subgraph_inputs) { using flatbuffers::Offset; flatbuffers::FlatBufferBuilder* fb_builder = BuilderInstance(); @@ -324,7 +334,8 @@ const Model* BuildModelWithOfflinePlanning(int number_of_tensors, number_of_tensors + tflite::testing::kOfflinePlannerHeaderSize); return model_builder.BuildModel(node_conn[0].input, - node_conn[num_conns - 1].output); + node_conn[num_conns - 1].output, + num_subgraph_inputs); } const Model* BuildSimpleMockModel() { @@ -710,9 +721,10 @@ const Model* GetSimpleModelWithBranch() { const Model* GetModelWithOfflinePlanning(int num_tensors, const int32_t* metadata_buffer, NodeConnection* node_conn, - int num_conns) { + int num_conns, + int num_subgraph_inputs) { const Model* model = BuildModelWithOfflinePlanning( - num_tensors, metadata_buffer, node_conn, num_conns); + num_tensors, metadata_buffer, node_conn, num_conns, num_subgraph_inputs); return model; } diff --git a/tensorflow/lite/micro/test_helpers.h b/tensorflow/lite/micro/test_helpers.h index a7897145d26..7b88a68084a 100644 --- a/tensorflow/lite/micro/test_helpers.h +++ b/tensorflow/lite/micro/test_helpers.h @@ -91,7 +91,8 @@ const Model* GetSimpleModelWithBranch(); const Model* GetModelWithOfflinePlanning(int num_tensors, const int32_t* metadata_buffer, NodeConnection* node_conn, - int num_conns); + int num_conns, + int num_subgraph_inputs = 0/*0 means all*/); // Returns a flatbuffer model with `simple_stateful_op` const Model* GetSimpleStatefulModel();