Adds GetCurrentOperatorIndex to micro_interpreter (#2605)
In this way an AOT can add metadata that could then be accessed with the couple GetCurrentSubgraphIndex, GetCurrentOperatorIndex.
These can be accessed from tflite::GetMicroContext(context)->graph()
Refers to #2593
BUG=#2593
diff --git a/tensorflow/lite/micro/micro_interpreter_graph.cc b/tensorflow/lite/micro/micro_interpreter_graph.cc
index 0d18fe7..7f096ae 100644
--- a/tensorflow/lite/micro/micro_interpreter_graph.cc
+++ b/tensorflow/lite/micro/micro_interpreter_graph.cc
@@ -44,6 +44,7 @@
model_(model),
allocator_(allocator),
current_subgraph_index_(0),
+ current_operator_index_(0),
resource_variables_(resource_variables) {
if (model != nullptr) {
subgraphs_ = model->subgraphs();
@@ -54,17 +55,21 @@
TfLiteStatus MicroInterpreterGraph::InitSubgraphs() {
int previous_subgraph_idx = current_subgraph_index_;
+ uint32_t previous_operator_idx = current_operator_index_;
for (size_t subgraph_idx = 0; subgraph_idx < subgraphs_->size();
subgraph_idx++) {
current_subgraph_index_ = subgraph_idx;
uint32_t operators_size = NumSubgraphOperators(model_, subgraph_idx);
- for (size_t i = 0; i < operators_size; ++i) {
- TfLiteNode* node =
- &(subgraph_allocations_[subgraph_idx].node_and_registrations[i].node);
- const TFLMRegistration* registration = subgraph_allocations_[subgraph_idx]
- .node_and_registrations[i]
- .registration;
+ for (current_operator_index_ = 0; current_operator_index_ < operators_size;
+ ++current_operator_index_) {
+ TfLiteNode* node = &(subgraph_allocations_[subgraph_idx]
+ .node_and_registrations[current_operator_index_]
+ .node);
+ const TFLMRegistration* registration =
+ subgraph_allocations_[subgraph_idx]
+ .node_and_registrations[current_operator_index_]
+ .registration;
size_t init_data_size;
const char* init_data;
if (registration->builtin_code == BuiltinOperator_CUSTOM) {
@@ -81,52 +86,62 @@
}
}
current_subgraph_index_ = previous_subgraph_idx;
+ current_operator_index_ = previous_operator_idx;
return kTfLiteOk;
}
TfLiteStatus MicroInterpreterGraph::PrepareSubgraphs() {
int previous_subgraph_idx = current_subgraph_index_;
-
+ uint32_t previous_operator_idx = current_operator_index_;
for (size_t subgraph_idx = 0; subgraph_idx < subgraphs_->size();
subgraph_idx++) {
current_subgraph_index_ = subgraph_idx;
uint32_t operators_size = NumSubgraphOperators(model_, subgraph_idx);
- for (size_t i = 0; i < operators_size; ++i) {
- TfLiteNode* node =
- &(subgraph_allocations_[subgraph_idx].node_and_registrations[i].node);
- const TFLMRegistration* registration = subgraph_allocations_[subgraph_idx]
- .node_and_registrations[i]
- .registration;
+ for (current_operator_index_ = 0; current_operator_index_ < operators_size;
+ ++current_operator_index_) {
+ TfLiteNode* node = &(subgraph_allocations_[subgraph_idx]
+ .node_and_registrations[current_operator_index_]
+ .node);
+ const TFLMRegistration* registration =
+ subgraph_allocations_[subgraph_idx]
+ .node_and_registrations[current_operator_index_]
+ .registration;
if (registration->prepare != nullptr) {
TfLiteStatus prepare_status = registration->prepare(context_, node);
if (prepare_status != kTfLiteOk) {
MicroPrintf("Node %s (number %df) failed to prepare with status %d",
- OpNameFromRegistration(registration), i, prepare_status);
+ OpNameFromRegistration(registration),
+ current_operator_index_, prepare_status);
return kTfLiteError;
}
}
- allocator_->FinishPrepareNodeAllocations(/*node_id=*/i);
+ allocator_->FinishPrepareNodeAllocations(
+ /*node_id=*/current_operator_index_);
}
}
current_subgraph_index_ = previous_subgraph_idx;
-
+ current_operator_index_ = previous_operator_idx;
return kTfLiteOk;
}
TfLiteStatus MicroInterpreterGraph::ResetSubgraphs() {
int previous_subgraph_idx = current_subgraph_index_;
+ uint32_t previous_operator_idx = current_operator_index_;
for (size_t subgraph_idx = 0; subgraph_idx < subgraphs_->size();
subgraph_idx++) {
current_subgraph_index_ = subgraph_idx;
uint32_t operators_size = NumSubgraphOperators(model_, subgraph_idx);
- for (size_t i = 0; i < operators_size; ++i) {
- TfLiteNode* node =
- &(subgraph_allocations_[subgraph_idx].node_and_registrations[i].node);
- const TFLMRegistration* registration = subgraph_allocations_[subgraph_idx]
- .node_and_registrations[i]
- .registration;
+ for (current_operator_index_ = 0; current_operator_index_ < operators_size;
+ ++current_operator_index_) {
+ TfLiteNode* node = &(subgraph_allocations_[subgraph_idx]
+ .node_and_registrations[current_operator_index_]
+ .node);
+ const TFLMRegistration* registration =
+ subgraph_allocations_[subgraph_idx]
+ .node_and_registrations[current_operator_index_]
+ .registration;
// registration is allocated outside the interpreter, so double check to
// make sure it's not nullptr;
if (registration != nullptr && registration->reset != nullptr) {
@@ -135,23 +150,28 @@
}
}
current_subgraph_index_ = previous_subgraph_idx;
+ current_operator_index_ = previous_operator_idx;
return kTfLiteOk;
}
TfLiteStatus MicroInterpreterGraph::FreeSubgraphs() {
int previous_subgraph_idx = current_subgraph_index_;
+ uint32_t previous_operator_idx = current_operator_index_;
for (size_t subgraph_idx = 0; subgraph_idx < subgraphs_->size();
subgraph_idx++) {
current_subgraph_index_ = subgraph_idx;
uint32_t operators_size = NumSubgraphOperators(model_, subgraph_idx);
- for (size_t i = 0; i < operators_size; ++i) {
- TfLiteNode* node =
- &(subgraph_allocations_[subgraph_idx].node_and_registrations[i].node);
- const TFLMRegistration* registration = subgraph_allocations_[subgraph_idx]
- .node_and_registrations[i]
- .registration;
+ for (current_operator_index_ = 0; current_operator_index_ < operators_size;
+ ++current_operator_index_) {
+ TfLiteNode* node = &(subgraph_allocations_[subgraph_idx]
+ .node_and_registrations[current_operator_index_]
+ .node);
+ const TFLMRegistration* registration =
+ subgraph_allocations_[subgraph_idx]
+ .node_and_registrations[current_operator_index_]
+ .registration;
// registration is allocated outside the interpreter, so double check to
// make sure it's not nullptr;
if (registration != nullptr && registration->free != nullptr) {
@@ -160,12 +180,14 @@
}
}
current_subgraph_index_ = previous_subgraph_idx;
+ current_operator_index_ = previous_operator_idx;
return kTfLiteOk;
}
TfLiteStatus MicroInterpreterGraph::InvokeSubgraph(int subgraph_idx) {
int previous_subgraph_idx = current_subgraph_index_;
+ uint32_t previous_operator_idx = current_operator_index_;
current_subgraph_index_ = subgraph_idx;
if (static_cast<size_t>(subgraph_idx) >= subgraphs_->size()) {
@@ -174,12 +196,15 @@
return kTfLiteError;
}
uint32_t operators_size = NumSubgraphOperators(model_, subgraph_idx);
- for (size_t i = 0; i < operators_size; ++i) {
- TfLiteNode* node =
- &(subgraph_allocations_[subgraph_idx].node_and_registrations[i].node);
- const TFLMRegistration* registration = subgraph_allocations_[subgraph_idx]
- .node_and_registrations[i]
- .registration;
+ for (current_operator_index_ = 0; current_operator_index_ < operators_size;
+ ++current_operator_index_) {
+ TfLiteNode* node = &(subgraph_allocations_[subgraph_idx]
+ .node_and_registrations[current_operator_index_]
+ .node);
+ const TFLMRegistration* registration =
+ subgraph_allocations_[subgraph_idx]
+ .node_and_registrations[current_operator_index_]
+ .registration;
// This ifdef is needed (even though ScopedMicroProfiler itself is a no-op with
// -DTF_LITE_STRIP_ERROR_STRINGS) because the function OpNameFromRegistration is
@@ -201,13 +226,15 @@
if (invoke_status == kTfLiteError) {
MicroPrintf("Node %s (number %d) failed to invoke with status %d",
- OpNameFromRegistration(registration), i, invoke_status);
+ OpNameFromRegistration(registration), current_operator_index_,
+ invoke_status);
return kTfLiteError;
} else if (invoke_status != kTfLiteOk) {
return invoke_status;
}
}
current_subgraph_index_ = previous_subgraph_idx;
+ current_operator_index_ = previous_operator_idx;
return kTfLiteOk;
}
diff --git a/tensorflow/lite/micro/micro_interpreter_graph.h b/tensorflow/lite/micro/micro_interpreter_graph.h
index 5c2121a..7ae0794 100644
--- a/tensorflow/lite/micro/micro_interpreter_graph.h
+++ b/tensorflow/lite/micro/micro_interpreter_graph.h
@@ -86,7 +86,12 @@
// to be the subgraph of that operator.
int GetCurrentSubgraphIndex() { return current_subgraph_index_; }
- // Gets the list of alloctions for each subgraph. This is the source of truth
+ // Get the current operator index inside a subgraph.
+ // The couple GetCurrentSubgraphIndex GetCurrentSubgraphIndex creates a unique
+ // identifier of the operator inside the subgraph
+ int GetCurrentOperatorIndex() { return current_operator_index_; }
+
+ // Gets the list of allocations for each subgraph. This is the source of truth
// for all per-subgraph allocation data.
SubgraphAllocations* GetAllocations() { return subgraph_allocations_; }
@@ -99,6 +104,7 @@
MicroAllocator* allocator_;
SubgraphAllocations* subgraph_allocations_ = nullptr;
int current_subgraph_index_;
+ uint32_t current_operator_index_;
MicroResourceVariables* resource_variables_;
const flatbuffers::Vector<flatbuffers::Offset<SubGraph>>* subgraphs_;
