Implement MicroContext/Graph for codegen

codegen/examples/hello_world/Makefile.inc

@@ -1,9 +1,13 @@
+# TODO(rjascani): The codegen runtime files (ie, in runtime subdir) should be a
+# separate library.
 CODEGEN_HELLO_WORLD_SRCS := \
 $(TENSORFLOW_ROOT)codegen/examples/hello_world/hello_world.cc \
-$(TENSORFLOW_ROOT)codegen/examples/hello_world/hello_world_model.cc
+$(TENSORFLOW_ROOT)codegen/examples/hello_world/hello_world_model.cc \
+$(TENSORFLOW_ROOT)codegen/runtime/micro_codegen_context.cc
 
 CODEGEN_HELLO_WORLD_HDRS := \
-$(TENSORFLOW_ROOT)codegen/examples/hello_world/hello_world_model.h
+$(TENSORFLOW_ROOT)codegen/examples/hello_world/hello_world_model.h \
+$(TENSORFLOW_ROOT)codegen/runtime/micro_codegen_context.h
 
 # Builds a standalone binary.
 $(eval $(call microlite_test,codegen_hello_world,\

codegen/examples/hello_world/hello_world_model.cc

@@ -17,11 +17,14 @@ limitations under the License.
 
 #include "hello_world_model.h"
 
+#include "codegen/runtime/micro_codegen_context.h"
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/c_api_types.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/micro/kernels/micro_ops.h"
 #include "tensorflow/lite/micro/micro_common.h"
+#include "tensorflow/lite/micro/micro_context.h"
 
 namespace hello_world_model {
 namespace {
@@ -101,6 +104,10 @@ alignas(16) uint8_t buffer_7[16] = {
 
 // buffer_10 is located in the arena
 
+constexpr size_t kSubgraph0Inputs[1] = {0};
+
+constexpr size_t kSubgraph0Outputs[1] = {9};
+
 struct Node0_0 {
   struct Inputs {
     int size = 3;
@@ -202,13 +209,34 @@ struct Tensor0_9Dims {
   int data[2] = {1, 1};
 } tensor0_9_dims;
 
+TfLiteStatus InvokeSubgraph0(TfLiteContext* context,
+                             tflite::Span<TfLiteNode> nodes) {
+  TFLITE_DCHECK(nodes.size() == 3);
+  TF_LITE_ENSURE_OK(
+      context, op_table[OpCode::kFullyConnected].invoke(context, &nodes[0]));
+  TF_LITE_ENSURE_OK(
+      context, op_table[OpCode::kFullyConnected].invoke(context, &nodes[1]));
+  TF_LITE_ENSURE_OK(
+      context, op_table[OpCode::kFullyConnected].invoke(context, &nodes[2]));
+
+  return kTfLiteOk;
+}
+
 }  // namespace
 
-Model::Model() {
-  context_.impl_ = nullptr;
+Model::Model()
+  : subgraphs_{
+      {.inputs = {&kSubgraph0Inputs[0], 1},
+       .outputs = {&kSubgraph0Outputs[0], 1},
+       .nodes = {&subgraph0_nodes_[0], 3},
+       .tensors = {&subgraph0_tensors_[0], 10},
+       .invoke = &InvokeSubgraph0},
+    },
+    micro_context_{&context_, {&subgraphs_[0], 1}} {
+  context_.impl_ = static_cast<void*>(&micro_context_);
   context_.ReportError = nullptr;
   context_.GetTensor = nullptr;
-  context_.GetEvalTensor = nullptr;
+  context_.GetEvalTensor = tflite::MicroContextGetEvalTensor;
   context_.profiler = nullptr;
   context_.GetExternalContext = nullptr;
   context_.GetScratchBuffer = nullptr;
@@ -280,17 +308,6 @@ Model::Model() {
       .type = kTfLiteInt8};
 }
 
-TfLiteStatus Model::Invoke() { return InvokeSubgraph0(); }
-
-TfLiteStatus Model::InvokeSubgraph0() {
-  TF_LITE_ENSURE_OK(context_, op_table[OpCode::kFullyConnected].invoke(
-                                  &context_, &subgraph0_nodes_[0]));
-  TF_LITE_ENSURE_OK(context_, op_table[OpCode::kFullyConnected].invoke(
-                                  &context_, &subgraph0_nodes_[1]));
-  TF_LITE_ENSURE_OK(context_, op_table[OpCode::kFullyConnected].invoke(
-                                  &context_, &subgraph0_nodes_[2]));
-
-  return kTfLiteOk;
-}
+TfLiteStatus Model::Invoke() { return micro_context_.InvokeSubgraph(0); }
 
 }  // namespace hello_world_model

codegen/examples/hello_world/hello_world_model.h

@@ -17,6 +17,7 @@ limitations under the License.
 
 #pragma once
 
+#include "codegen/runtime/micro_codegen_context.h"
 #include "tensorflow/lite/c/c_api_types.h"
 #include "tensorflow/lite/c/common.h"
 
@@ -29,9 +30,9 @@ class Model {
   TfLiteStatus Invoke();
 
  private:
-  TfLiteStatus InvokeSubgraph0();
-
   TfLiteContext context_ = {};
+  tflite::Subgraph subgraphs_[1];
+  tflite::MicroCodegenContext micro_context_;
   TfLiteNode subgraph0_nodes_[3] = {};
   TfLiteEvalTensor subgraph0_tensors_[10] = {};
 };

codegen/graph.py

@@ -73,6 +73,14 @@ def __init__(self, model: schema_fb.ModelT, buffers: Sequence[tensor.Buffer],
   def index(self) -> int:
     return self._subgraph_idx
 
+  @property
+  def inputs(self) -> Sequence[int]:
+    return self._subgraph.inputs
+
+  @property
+  def outputs(self) -> Sequence[int]:
+    return self._subgraph.outputs
+
   @property
   def operators(self) -> Sequence[operator.Operator]:
     return self._operators
@@ -85,6 +93,18 @@ def tensors(self) -> Sequence[tensor.Tensor]:
   def needs_zero_length_int_array(self) -> bool:
     return any(t.needs_zero_length_int_array for t in self.tensors)
 
+  @property
+  def invoke_fn_name(self) -> str:
+    return f"InvokeSubgraph{self.index}"
+
+  @property
+  def inputs_array_name(self) -> str:
+    return f"kSubgraph{self.index}Inputs"
+
+  @property
+  def outputs_array_name(self) -> str:
+    return f"kSubgraph{self.index}Outputs"
+
   @property
   def nodes_array(self) -> str:
     return f"subgraph{self.index}_nodes_"
@@ -116,16 +136,54 @@ def generate_c_node_init(self, indent: str) -> str:
     return textwrap.indent("\n".join(node_init_strs), indent)
 
   def generate_c_invoke(self, indent: str) -> str:
-    invoke_template = string.Template(
-        "TF_LITE_ENSURE_OK(context_, op_table[${op_code}].invoke(\n"
-        "                                &context_, &${node}));\n")
+    function_template = string.Template(
+        "TfLiteStatus ${function_name}(TfLiteContext* context,\n"
+        "                             tflite::Span<TfLiteNode> nodes) {\n"
+        "  TFLITE_DCHECK(nodes.size() == ${num_nodes});\n"
+        "${body}\n"
+        "  return kTfLiteOk;\n"
+        "}")
+
+    body_template = string.Template(
+        "  TF_LITE_ENSURE_OK(\n"
+        "      context, op_table[${op_code}].invoke(context, &${node}));\n")
     invoke_strs: List[str] = []
     for op_idx, op in enumerate(self.operators):
       invoke_strs.append(
-          invoke_template.substitute(
+          body_template.substitute(
               op_code=self._op_codes[op.op_code_index].full_enum_name,
-              node=self.nodes_element(op_idx)))
-    return textwrap.indent("".join(invoke_strs), indent)
+              node=f"nodes[{op_idx}]"))
+
+    invoke = function_template.substitute(function_name=self.invoke_fn_name,
+                                          num_nodes=len(self.operators),
+                                          body="".join(invoke_strs))
+    return textwrap.indent(invoke, indent)
+
+  def generate_c_input_array(self, indent: str) -> str:
+    return utils.generate_c_int_array(indent, "size_t", self.inputs_array_name,
+                                      self.inputs)
+
+  def generate_c_output_array(self, indent: str) -> str:
+    return utils.generate_c_int_array(indent, "size_t",
+                                      self.outputs_array_name, self.outputs)
+
+  def generate_c_subgraph_init(self, indent: str) -> str:
+    init_template = string.Template(
+        "{.inputs = {&${input_array}[0], ${input_size}},\n"
+        " .outputs = {&${output_array}[0], ${output_size}},\n"
+        " .nodes = {&${node_array}[0], ${node_size}},\n"
+        " .tensors = {&${tensor_array}[0], ${tensor_size}},\n"
+        " .invoke = &${invoke}},")
+    return textwrap.indent(
+        init_template.substitute(input_array=self.inputs_array_name,
+                                 input_size=len(self.inputs),
+                                 output_array=self.outputs_array_name,
+                                 output_size=len(self.outputs),
+                                 node_array=self.nodes_array,
+                                 node_size=len(self.operators),
+                                 tensor_array=self.tensors_array,
+                                 tensor_size=len(self.tensors),
+                                 invoke=self.invoke_fn_name), indent)
 
   @property
   def tensors_array(self) -> str:

codegen/runtime/micro_codegen_context.cc

@@ -0,0 +1,139 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "codegen/runtime/micro_codegen_context.h"
+
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/op_macros.h"
+#include "tensorflow/lite/micro/micro_log.h"
+
+namespace tflite {
+
+MicroCodegenContext::MicroCodegenContext(TfLiteContext* context,
+                                         Span<Subgraph> subgraphs)
+    : context_(context), subgraphs_(subgraphs) {}
+
+void* MicroCodegenContext::GetScratchBuffer(int buffer_idx) {
+  // TODO(rjascani): Implement scratch buffers
+  return nullptr;
+}
+
+TfLiteEvalTensor* MicroCodegenContext::GetEvalTensor(int tensor_idx) {
+  TFLITE_DCHECK(static_cast<size_t>(tensor_idx) <
+                subgraphs_[current_subgraph_idx_].tensors.size());
+  return &subgraphs_[current_subgraph_idx_].tensors[tensor_idx];
+}
+
+TfLiteStatus MicroCodegenContext::set_external_context(
+    void* external_context_payload) {
+  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* MicroCodegenContext::external_context() {
+  return external_context_payload_;
+}
+
+MicroGraph& MicroCodegenContext::graph() { return *this; }
+
+void* MicroCodegenContext::AllocatePersistentBuffer(size_t) {
+  // Not allowed at Eval
+  TFLITE_ABORT;
+  return nullptr;
+}
+
+TfLiteStatus MicroCodegenContext::RequestScratchBufferInArena(size_t, int*) {
+  // Not allowed at Eval
+  TFLITE_ABORT;
+  return kTfLiteError;
+}
+
+TfLiteTensor* MicroCodegenContext::AllocateTempTfLiteTensor(int) {
+  // Not allowed at Eval
+  TFLITE_ABORT;
+  return nullptr;
+}
+
+void MicroCodegenContext::DeallocateTempTfLiteTensor(TfLiteTensor*) {
+  // Not allowed at Eval
+  TFLITE_ABORT;
+}
+
+uint8_t* MicroCodegenContext::AllocateTempBuffer(size_t, size_t) {
+  // Not allowed at Eval
+  TFLITE_ABORT;
+  return nullptr;
+}
+
+void MicroCodegenContext::DeallocateTempBuffer(uint8_t*) {
+  // Not allowed at Eval
+  TFLITE_ABORT;
+}
+
+TfLiteStatus MicroCodegenContext::InvokeSubgraph(int subgraph_idx) {
+  TF_LITE_ENSURE(context_,
+                 static_cast<size_t>(subgraph_idx) < subgraphs_.size());
+  size_t previous_subgraph_idx = current_subgraph_idx_;
+  current_subgraph_idx_ = subgraph_idx;
+  TfLiteStatus status =
+      subgraphs_[subgraph_idx].invoke(context_, subgraphs_[subgraph_idx].nodes);
+  current_subgraph_idx_ = previous_subgraph_idx;
+  return status;
+}
+
+size_t MicroCodegenContext::NumSubgraphInputs(int subgraph_idx) {
+  TFLITE_DCHECK(static_cast<size_t>(subgraph_idx) < subgraphs_.size());
+  return subgraphs_[subgraph_idx].inputs.size();
+}
+
+TfLiteEvalTensor* MicroCodegenContext::GetSubgraphInput(int subgraph_idx,
+                                                        int input_idx) {
+  TFLITE_DCHECK(static_cast<size_t>(subgraph_idx) < subgraphs_.size());
+  TFLITE_DCHECK(static_cast<size_t>(input_idx) <
+                subgraphs_[subgraph_idx].inputs.size());
+  const size_t tensor_idx = subgraphs_[subgraph_idx].inputs[input_idx];
+  return &subgraphs_[subgraph_idx].tensors[tensor_idx];
+}
+
+size_t MicroCodegenContext::NumSubgraphOutputs(int subgraph_idx) {
+  TFLITE_DCHECK(static_cast<size_t>(subgraph_idx) < subgraphs_.size());
+  return subgraphs_[subgraph_idx].outputs.size();
+}
+
+TfLiteEvalTensor* MicroCodegenContext::GetSubgraphOutput(int subgraph_idx,
+                                                         int output_idx) {
+  TFLITE_DCHECK(static_cast<size_t>(subgraph_idx) < subgraphs_.size());
+  TFLITE_DCHECK(static_cast<size_t>(output_idx) <
+                subgraphs_[subgraph_idx].outputs.size());
+  const size_t tensor_idx = subgraphs_[subgraph_idx].outputs[output_idx];
+  return &subgraphs_[subgraph_idx].tensors[tensor_idx];
+}
+
+int MicroCodegenContext::NumSubgraphs() { return subgraphs_.size(); }
+
+MicroResourceVariables* MicroCodegenContext::GetResourceVariables() {
+  return nullptr;
+}
+
+}  // namespace tflite