Merge pull request from mansnils:scratch_tensors

PiperOrigin-RevId: 338767807
Change-Id: I2e1962f773872cb94abbe21881b57a1144cdf645
This commit is contained in:
TensorFlower Gardener 2020-10-23 16:55:40 -07:00
commit b35dceb805
3 changed files with 142 additions and 130 deletions

View File

@ -229,21 +229,6 @@ 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];
// TODO(b/166484865): Figure out a more general solution.
// This workaround is needed to handle situations where subgraph input !=
// operator input.
// In case operator input(s) 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;
}
@ -257,16 +242,15 @@ TfLiteStatus AllocationInfoBuilder::AddTensors(const SubGraph* subgraph,
}
}
// Work out which tensors need to be allocated.
// Sanity check for valid tensor lifetime.
for (size_t i = 0; i < tensor_count_; ++i) {
AllocationInfo* current = &info_[i];
const bool is_read_only =
// Even though tensor appears to be read only it may still need to be
// allocated.
const bool appears_read_only =
(current->first_created == -1) && (current->last_used != -1);
if (is_read_only) {
current->needs_allocating = false;
}
const bool has_partial_lifetime =
!is_read_only &&
!appears_read_only &&
((current->first_created == -1) || (current->last_used == -1));
if (has_partial_lifetime && current->needs_allocating) {
TF_LITE_REPORT_ERROR(

View File

@ -28,6 +28,12 @@ namespace testing {
namespace {
constexpr int kExpectedAlignment = 4;
constexpr int t0 = 0;
constexpr int t1 = 1;
constexpr int t2 = 2;
constexpr int t3 = 3;
constexpr int t4 = 4;
constexpr int t5 = 5;
void VerifyMockTfLiteTensor(TfLiteTensor* tensor, bool is_variable = false) {
TF_LITE_MICRO_EXPECT_EQ(kTfLiteInt32, tensor->type);
@ -446,18 +452,18 @@ TF_LITE_MICRO_TEST(TestAllocationForComplexModelAllocation) {
TF_LITE_MICRO_TEST(OfflinePlannerBranchesAllOnline) {
int version = 1;
int subgraph = 0;
constexpr int nbr_tensors = 4;
constexpr int number_tensors = 4;
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
tflite::NodeAndRegistration* node_and_registration;
const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize +
nbr_tensors] = {version, subgraph,
nbr_tensors, // header
// memory offsets:
-1, -1, -1, -1};
number_tensors] = {version, subgraph,
number_tensors, // header
// memory offsets:
-1, -1, -1, -1};
// The structure is identical to the one in
// TestAllocationForModelsWithBranches
int num_conns = 3;
int number_connections = 3;
tflite::testing::NodeConnection node_list[3] = {{
{0}, // input
{1} // output
@ -472,7 +478,7 @@ TF_LITE_MICRO_TEST(OfflinePlannerBranchesAllOnline) {
}};
const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning(
nbr_tensors, metadata_buffer, node_list, num_conns);
number_tensors, metadata_buffer, node_list, number_connections);
TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
@ -506,37 +512,26 @@ TF_LITE_MICRO_TEST(OfflinePlannerBranchesAllOnline) {
}
TF_LITE_MICRO_TEST(OfflinePlannerBasic) {
constexpr int nbr_tensors = 4;
constexpr int number_tensors = 4;
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,
0, // t0
48, // t1
0, // t2
48}; // t3
int t0 = 0;
int t1 = 1;
int t2 = 2;
int t3 = 3;
int num_conns = 3;
tflite::testing::NodeConnection node_list[3] = {{
{t0}, // input
{t1} // output
},
{
{t1}, // input
{t2} // output
},
{
{t2}, // input
{t3} // output
}};
number_tensors] = {1, 0, number_tensors,
/*t0=*/0,
/*t1=*/48,
/*t2=*/0,
/*t3=*/48};
constexpr int number_connections = 3;
tflite::testing::NodeConnection node_list[number_connections] = {
{/*input=*/{tflite::testing::t0},
/*output=*/{tflite::testing::t1}},
{/*input=*/{tflite::testing::t1},
/*output=*/{tflite::testing::t2}},
{/*input=*/{tflite::testing::t2},
/*output=*/{tflite::testing::t3}}};
const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning(
nbr_tensors, metadata_buffer, node_list, num_conns);
number_tensors, metadata_buffer, node_list, number_connections);
TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
@ -561,37 +556,28 @@ TF_LITE_MICRO_TEST(OfflinePlannerBasic) {
}
TF_LITE_MICRO_TEST(OfflinePlannerOverlappingAllocation) {
constexpr int nbr_tensors = 4;
constexpr int number_tensors = 4;
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
48, // t2
-1}; // t3
number_tensors] = {/*version=*/1,
/*subgraph=*/0,
number_tensors,
/*t0=*/0,
/*t1=*/0,
/*t2=*/48,
/*t3=*/-1};
int t0 = 0;
int t1 = 1;
int t2 = 2;
int t3 = 3;
int num_conns = 2;
int number_connections = 2;
tflite::testing::NodeConnection node_list[2] = {
{
{t0, t1}, // input, scratch
{t2} // output
},
{
{t2}, // input
{t3} // output
},
{/*input, scratch=*/{tflite::testing::t0, tflite::testing::t1},
/*output=*/{tflite::testing::t2}},
{/*input=*/{tflite::testing::t2},
/*output=*/{tflite::testing::t3}},
};
const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning(
nbr_tensors, metadata_buffer, node_list, num_conns);
number_tensors, metadata_buffer, node_list, number_connections);
TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
@ -617,39 +603,33 @@ TF_LITE_MICRO_TEST(OfflinePlannerOverlappingAllocation) {
}
TF_LITE_MICRO_TEST(OfflinePlannerOfflineOnline) {
constexpr int nbr_tensors = 5;
constexpr int number_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
48, // t1
-1, // t2
0, // t3
-1}; // t4
number_tensors] = {/*version=*/1,
/*subgraph=*/0,
number_tensors,
/*t0=*/0,
/*t1=*/48,
/*t2=*/-1,
/*t3=*/0,
/*t4=*/-1};
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] = {
constexpr int number_connections = 2;
tflite::testing::NodeConnection node_list[number_connections] = {
{
{t0, t1}, // input, scratch
{t2}, // output
/*input, scratch=*/{tflite::testing::t0, tflite::testing::t1},
/*output=*/{tflite::testing::t2},
},
{
{t2}, // input
{t3, t4}, // output1, output2
/*input=*/{tflite::testing::t2},
/*output1, output2=*/{tflite::testing::t3, tflite::testing::t4},
},
};
const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning(
nbr_tensors, metadata_buffer, node_list, num_conns);
number_tensors, metadata_buffer, node_list, number_connections);
TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
@ -757,43 +737,34 @@ TF_LITE_MICRO_TEST(TestAllocateTfLiteTensorWithReset) {
}
TF_LITE_MICRO_TEST(TestOperatorInputsNotInSubgraphInputs) {
constexpr int nbr_tensors = 5;
constexpr int number_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
number_tensors] = {/*version=*/1,
/*subgraph=*/0,
number_tensors,
/*t0=*/0,
/*t1=*/0,
/*t2=*/0,
/*t3=*/48,
/*t4=*/-1};
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
},
constexpr int number_connections = 2;
tflite::testing::NodeConnection node_list[number_connections] = {
{// 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)
{tflite::testing::t0, tflite::testing::t1, tflite::testing::t2},
/*t3: output=*/{tflite::testing::t3}},
{/*t3: input=*/{tflite::testing::t3},
/*t4: output=*/{tflite::testing::t4}},
};
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*/);
number_tensors, metadata_buffer, node_list, number_connections,
/*Only first tensor (t0) is in subgraph input list=*/1);
TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
@ -818,4 +789,61 @@ TF_LITE_MICRO_TEST(TestOperatorInputsNotInSubgraphInputs) {
TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[4].data.uint8 - start);
}
TF_LITE_MICRO_TEST(TestTypicalFirstOpAndSecondOpWithScratchTensors) {
constexpr int number_tensors = 6;
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
tflite::NodeAndRegistration* node_and_registration;
const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize +
number_tensors] = {/*version=*/1,
/*subgraph=*/0,
number_tensors,
/*t0=*/0,
/*t1=*/0,
/*t2=*/0,
/*t3=*/0,
/*t4=*/48,
/*t5=*/-1};
constexpr int number_connections = 3;
tflite::testing::NodeConnection node_list[number_connections] = {
{/*t0: input (subgraph and operator input)=*/{tflite::testing::t0},
/*t1: output=*/{tflite::testing::t1}},
{// t1: input
// t2: scratch1 (only in operator inputs)
// t3: scratch2 (only in operator inputs)
{tflite::testing::t1, tflite::testing::t2, tflite::testing::t3},
/*t4: output=*/{tflite::testing::t4}},
{/*t4: input=*/{tflite::testing::t4},
/*t5: output=*/{tflite::testing::t5}},
};
const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning(
number_tensors, metadata_buffer, node_list, number_connections,
/*Only first tensor (t0) is in subgraph input list=*/1);
TfLiteEvalTensor* eval_tensors = nullptr;
tflite::ScratchBufferHandle* scratch_buffer_handles = 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,
&scratch_buffer_handles));
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(0, eval_tensors[3].data.uint8 - start);
TF_LITE_MICRO_EXPECT_EQ(48, eval_tensors[4].data.uint8 - start);
TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[5].data.uint8 - start);
}
TF_LITE_MICRO_TESTS_END

View File

@ -206,7 +206,7 @@ const Model* ModelBuilder::BuildModel(
} else {
// A non-zero value of num_subgraph_inputs means that some of
// the operator input tensors are not subgraph inputs.
TFLITE_DCHECK(num_subgraph_inputs < inputs.size());
TFLITE_DCHECK(num_subgraph_inputs <= inputs.size());
}
const flatbuffers::Offset<SubGraph> subgraphs[subgraphs_size] = {