Make MicroContext and MicroGraph abstract (#2202)

MicroContext and MicroGraph serve as the primary API surface between the interpreter and the kernels. For interpreter-less execution, we need to be able to substitute different MicroContext and MicroGraph implementations. This PR makes MicroContext and MicroGraph abstract and moves their existing implementations to MicroInterpreterContext and MicroInterpreterGraph, respectively. This will allow us to substitute an inference-only implementation of each for codegen.

FakeMicroContext and MockMicroGraph were also updated to implement the base classes.

BUG=b/295174086

tensorflow/lite/micro/BUILD

@@ -39,8 +39,8 @@ cc_library(
     deps = [
         ":memory_helpers",
         ":micro_allocator",
-        ":micro_context",
-        ":micro_graph",
+        ":micro_interpreter_context",
+        ":micro_interpreter_graph",
         ":micro_profiler_interface",
         ":op_resolvers",
         "//tensorflow/lite:type_to_tflitetype",
@@ -62,10 +62,28 @@ cc_library(
         "micro_context.h",
     ],
     copts = micro_copts(),
+    deps = [
+        ":micro_common",
+        ":micro_graph",
+        ":micro_log",
+        "//tensorflow/lite/c:common",
+    ],
+)
+
+cc_library(
+    name = "micro_interpreter_context",
+    srcs = [
+        "micro_interpreter_context.cc",
+    ],
+    hdrs = [
+        "micro_interpreter_context.h",
+    ],
+    copts = micro_copts(),
     deps = [
         ":memory_helpers",
         ":micro_allocator",
-        ":micro_graph",
+        ":micro_context",
+        ":micro_interpreter_graph",
         ":micro_log",
         ":micro_profiler_interface",
         "//tensorflow/lite/c:common",
@@ -94,22 +112,34 @@ cc_library(
     copts = micro_copts(),
     deps = [
         ":memory_helpers",
-        ":micro_allocator",
+        ":micro_arena_constants",
         ":micro_context",
         ":micro_log",
         ":mock_micro_graph",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/micro/arena_allocator:simple_memory_allocator",
     ],
 )
 
 cc_library(
     name = "micro_graph",
-    srcs = ["micro_graph.cc"],
     hdrs = ["micro_graph.h"],
+    deps = [
+        ":micro_common",
+        ":micro_resource_variable",
+        "//tensorflow/lite/kernels/internal:compatibility",
+    ],
+)
+
+cc_library(
+    name = "micro_interpreter_graph",
+    srcs = ["micro_interpreter_graph.cc"],
+    hdrs = ["micro_interpreter_graph.h"],
     deps = [
         ":memory_helpers",
         ":micro_allocator",
         ":micro_common",
+        ":micro_graph",
         ":micro_log",
         ":micro_profiler",
         ":micro_resource_variable",
@@ -401,13 +431,14 @@ cc_test(
 )
 
 cc_test(
-    name = "micro_context_test",
+    name = "micro_interpreter_context_test",
     srcs = [
-        "micro_context_test.cc",
+        "micro_interpreter_context_test.cc",
     ],
     deps = [
         ":micro_allocator",
-        ":micro_context",
+        ":micro_interpreter_context",
+        ":micro_interpreter_graph",
         ":test_helpers",
         "//tensorflow/lite/micro/testing:micro_test",
     ],
@@ -421,6 +452,7 @@ cc_test(
     deps = [
         ":fake_micro_context",
         ":micro_allocator",
+        ":mock_micro_graph",
         ":test_helpers",
         "//tensorflow/lite/micro/testing:micro_test",
     ],

tensorflow/lite/micro/fake_micro_context.cc

@@ -15,28 +15,18 @@ limitations under the License.
 
 #include "tensorflow/lite/micro/fake_micro_context.h"
 
+#include "tensorflow/lite/c/c_api_types.h"
 #include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/micro/arena_allocator/single_arena_buffer_allocator.h"
-#include "tensorflow/lite/micro/micro_allocator.h"
 #include "tensorflow/lite/micro/micro_arena_constants.h"
 #include "tensorflow/lite/micro/micro_log.h"
 
 namespace tflite {
-namespace {
-// Dummy static variables to allow creation of dummy MicroAllocator.
-// All tests are guarateed to run serially.
-static constexpr int KDummyTensorArenaSize = 256;
-static uint8_t dummy_tensor_arena[KDummyTensorArenaSize];
-}  // namespace
 
 FakeMicroContext::FakeMicroContext(TfLiteTensor* tensors,
                                    SingleArenaBufferAllocator* allocator,
                                    MicroGraph* micro_graph)
-    : MicroContext(
-          MicroAllocator::Create(dummy_tensor_arena, KDummyTensorArenaSize),
-          nullptr, micro_graph),
-      tensors_(tensors),
-      allocator_(allocator) {}
+    : graph_(*micro_graph), tensors_(tensors), allocator_(allocator) {}
 
 TfLiteTensor* FakeMicroContext::AllocateTempTfLiteTensor(int tensor_index) {
   allocated_temp_count_++;
@@ -113,4 +103,13 @@ void* FakeMicroContext::GetScratchBuffer(int buffer_index) {
   return scratch_buffers_[buffer_index];
 }
 
+TfLiteStatus FakeMicroContext::set_external_context(
+    void* external_context_payload) {
+  return kTfLiteError;
+}
+
+void* FakeMicroContext::external_context() { return nullptr; }
+
+MicroGraph& FakeMicroContext::graph() { return graph_; }
+
 }  // namespace tflite

tensorflow/lite/micro/fake_micro_context.h

@@ -27,6 +27,8 @@ namespace tflite {
 
 class FakeMicroContext : public MicroContext {
  public:
+  ~FakeMicroContext() = default;
+
   FakeMicroContext(TfLiteTensor* tensors, SingleArenaBufferAllocator* allocator,
                    MicroGraph* micro_graph);
 
@@ -44,9 +46,14 @@ class FakeMicroContext : public MicroContext {
 
   TfLiteEvalTensor* GetEvalTensor(int tensor_index) override;
 
+  TfLiteStatus set_external_context(void* external_context_payload) override;
+  void* external_context() override;
+  MicroGraph& graph() override;
+
  private:
   static constexpr int kNumScratchBuffers_ = 12;
 
+  MicroGraph& graph_;
   int scratch_buffer_count_ = 0;
   uint8_t* scratch_buffers_[kNumScratchBuffers_];
 

tensorflow/lite/micro/fake_micro_context_test.cc

@@ -18,6 +18,7 @@ limitations under the License.
 
 #include "tensorflow/lite/micro/arena_allocator/single_arena_buffer_allocator.h"
 #include "tensorflow/lite/micro/micro_allocator.h"
+#include "tensorflow/lite/micro/mock_micro_graph.h"
 #include "tensorflow/lite/micro/test_helpers.h"
 #include "tensorflow/lite/micro/testing/micro_test.h"
 
@@ -58,7 +59,7 @@ TF_LITE_MICRO_TEST(TestGetBeforeRequestScratchBufferWouldReturnNull) {
   uint8_t arena_buffer[kArenaSize];
   tflite::SingleArenaBufferAllocator simple_memory_allocator(arena_buffer,
                                                              kArenaSize);
-  tflite::MicroGraph dummy_micro_graph(nullptr, nullptr, nullptr, nullptr);
+  tflite::MockMicroGraph dummy_micro_graph{&simple_memory_allocator};
 
   tflite::FakeMicroContext micro_context = tflite::CreateFakeMicroContext(
       &simple_memory_allocator, &dummy_micro_graph);
@@ -71,7 +72,7 @@ TF_LITE_MICRO_TEST(TestRequestScratchBufferAndThenGetShouldSucceed) {
   uint8_t arena_buffer[kArenaSize];
   tflite::SingleArenaBufferAllocator simple_memory_allocator(arena_buffer,
                                                              kArenaSize);
-  tflite::MicroGraph dummy_micro_graph(nullptr, nullptr, nullptr, nullptr);
+  tflite::MockMicroGraph dummy_micro_graph{&simple_memory_allocator};
 
   tflite::FakeMicroContext micro_context = tflite::CreateFakeMicroContext(
       &simple_memory_allocator, &dummy_micro_graph);

tensorflow/lite/micro/micro_context.cc

@@ -17,48 +17,14 @@ limitations under the License.
 
 #include <cstdarg>
 #include <cstddef>
-#include <cstdint>
 
-#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/micro/micro_common.h"
 #include "tensorflow/lite/micro/micro_log.h"
 
 namespace tflite {
-MicroContext::MicroContext(MicroAllocator* allocator, const Model* model,
-                           MicroGraph* graph)
-    : allocator_(*allocator),
-      graph_(*graph),
-      model_(model),
-      state_(InterpreterState::kInit) {}
+namespace {
 
-MicroContext::~MicroContext() {}
-
-void* MicroContext::AllocatePersistentBuffer(size_t bytes) {
-  TFLITE_DCHECK(state_ == InterpreterState::kPrepare ||
-                state_ == InterpreterState::kInit);
-  return allocator_.AllocatePersistentBuffer(bytes);
-}
-
-TfLiteStatus MicroContext::RequestScratchBufferInArena(size_t bytes,
-                                                       int* buffer_idx) {
-  TFLITE_DCHECK(state_ == InterpreterState::kPrepare);
-  return allocator_.RequestScratchBufferInArena(
-      bytes, graph_.GetCurrentSubgraphIndex(), buffer_idx);
-}
-
-void* MicroContext::GetScratchBuffer(int buffer_idx) {
-  TFLITE_DCHECK(state_ == InterpreterState::kInvoke);
-  ScratchBufferHandle* handle = scratch_buffer_handles_ + buffer_idx;
-  return handle->data;
-}
-
-TfLiteTensor* MicroContext::AllocateTempTfLiteTensor(int tensor_idx) {
-  return allocator_.AllocateTempTfLiteTensor(model_, graph_.GetAllocations(),
-                                             tensor_idx,
-                                             graph_.GetCurrentSubgraphIndex());
-}
-
-int MicroContext::GetTensorIndex(int index, int max_size,
-                                 const int* tensor_indices) {
+int GetTensorIndex(int index, int max_size, const int* tensor_indices) {
   if (index >= 0 && index < max_size) {
     const int tensor_index = tensor_indices[index];
     if (tensor_index != kTfLiteOptionalTensor) {
@@ -68,6 +34,8 @@ int MicroContext::GetTensorIndex(int index, int max_size,
   return -1;
 }
 
+}  // namespace
+
 TfLiteTensor* MicroContext::AllocateTempInputTensor(const TfLiteNode* node,
                                                     int index) {
   const int tensor_index =
@@ -98,46 +66,6 @@ TfLiteTensor* MicroContext::AllocateTempIntermediateTensor(
   return AllocateTempTfLiteTensor(tensor_index);
 }
 
-void MicroContext::DeallocateTempTfLiteTensor(TfLiteTensor* tensor) {
-  return allocator_.DeallocateTempTfLiteTensor(tensor);
-}
-
-uint8_t* MicroContext::AllocateTempBuffer(size_t size, size_t alignment) {
-  TFLITE_DCHECK(state_ == InterpreterState::kPrepare);
-  return allocator_.AllocateTempBuffer(size, alignment);
-}
-
-void MicroContext::DeallocateTempBuffer(uint8_t* buffer) {
-  TFLITE_DCHECK(state_ == InterpreterState::kPrepare);
-  allocator_.DeallocateTempBuffer(buffer);
-}
-
-TfLiteEvalTensor* MicroContext::GetEvalTensor(int tensor_idx) {
-  return &graph_.GetAllocations()[graph_.GetCurrentSubgraphIndex()]
-              .tensors[tensor_idx];
-}
-
-void MicroContext::SetScratchBufferHandles(
-    ScratchBufferHandle* scratch_buffer_handles) {
-  scratch_buffer_handles_ = scratch_buffer_handles;
-}
-
-TfLiteStatus MicroContext::set_external_context(
-    void* external_context_payload) {
-  TFLITE_DCHECK(state_ == InterpreterState::kPrepare ||
-                state_ == InterpreterState::kInvoke);
-  if (external_context_payload == nullptr ||
-      external_context_payload_ != nullptr) {
-    MicroPrintf(
-        "Attempting to set external context to %x but it was %x already",
-        external_context_payload, external_context_payload_);
-    return kTfLiteError;
-  }
-
-  external_context_payload_ = external_context_payload;
-  return kTfLiteOk;
-}
-
 void MicroContextReportOpError(struct TfLiteContext* context,
                                const char* format, ...) {
   va_list args;
@@ -146,12 +74,4 @@ void MicroContextReportOpError(struct TfLiteContext* context,
   va_end(args);
 }
 
-void MicroContext::SetInterpreterState(MicroContext::InterpreterState state) {
-  state_ = state;
-}
-
-MicroContext::InterpreterState MicroContext::GetInterpreterState() const {
-  return state_;
-}
-
 }  // namespace tflite