mirror of
https://github.com/vee1e/tflite-micro.git
synced 2026-09-02 10:17:54 +00:00
3323a41d69
The TFLM kernels use the MicroContext and MicroGraph interfaces to fetch eval tensors and access other parts of the graph. Since codegen does not have the MicroInterpreter and related objects to serve those, this PR introduces a new MicroCodegenContext class that serves as both the MicroContext and MicroGraph. The MicroCodegenContext is configured with a span of Subgraphs, each of which includes the inputs, outputs, nodes, tensors and an invocation function. The code generator will create the data and functions needed for each Subgraph and initialize the MicroCodegenContext with it. By having the re-usable MicroCodegenContext, the code generator won't have to generate nearly as much code. BUG=b/295174086
262 lines
9.1 KiB
Python
262 lines
9.1 KiB
Python
# 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.
|
|
# ==============================================================================
|
|
""" Provides object representation for the model that is conducive to code
|
|
generation using templates. """
|
|
|
|
from typing import Dict, List, Optional, Sequence
|
|
import string
|
|
import textwrap
|
|
|
|
from tflite_micro.codegen.operators import factory
|
|
from tflite_micro.codegen.operators import operator
|
|
from tflite_micro.codegen import tensor
|
|
from tflite_micro.codegen import utils
|
|
from tflite_micro.tensorflow.lite.python import schema_py_generated as schema_fb
|
|
from tflite_micro.tensorflow.lite.tools import visualize
|
|
|
|
|
|
class OpCode(object):
|
|
|
|
def __init__(self, op_code: schema_fb.OperatorCodeT):
|
|
self._op_code: schema_fb.OperatorCodeT = op_code
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
if self._op_code.customCode:
|
|
return self._op_code.customCode
|
|
return visualize.BuiltinCodeToName(self._op_code.builtinCode)
|
|
|
|
@property
|
|
def register_function(self) -> str:
|
|
return "tflite::RegisterInference_{}".format(self.name)
|
|
|
|
@property
|
|
def enum_name(self) -> str:
|
|
return "k{}".format(utils.to_pascal_case(self.name))
|
|
|
|
@property
|
|
def full_enum_name(self) -> str:
|
|
return "OpCode::" + self.enum_name
|
|
|
|
|
|
class Subgraph(object):
|
|
|
|
def __init__(self, model: schema_fb.ModelT, buffers: Sequence[tensor.Buffer],
|
|
subgraph_idx: int, subgraph: schema_fb.SubGraphT):
|
|
self._subgraph_idx: int = subgraph_idx
|
|
self._subgraph: schema_fb.SubGraphT = subgraph
|
|
self._op_codes: List[OpCode] = [
|
|
OpCode(op_code) for op_code in model.operatorCodes
|
|
]
|
|
self._tensors: List[Tensor] = []
|
|
for t in subgraph.tensors:
|
|
self._tensors.append(tensor.Tensor(buffers[t.buffer], t))
|
|
|
|
self._operators: List[operator.Operator] = []
|
|
for op in subgraph.operators:
|
|
op_code = model.operatorCodes[op.opcodeIndex]
|
|
self._operators.append(factory.create_operator(op_code, op))
|
|
|
|
@property
|
|
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
|
|
|
|
@property
|
|
def tensors(self) -> Sequence[tensor.Tensor]:
|
|
return self._tensors
|
|
|
|
@property
|
|
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_"
|
|
|
|
def nodes_element(self, operator_idx: int) -> str:
|
|
return self.nodes_array + f"[{operator_idx}]"
|
|
|
|
def node_data_type(self, operator_idx: int) -> str:
|
|
return f"Node{self.index}_{operator_idx}"
|
|
|
|
def node_data_name(self, operator_idx: int) -> str:
|
|
return f"node_{self.index}_{operator_idx}"
|
|
|
|
def generate_c_node_data(self, indent: str) -> str:
|
|
node_data_strs: List[str] = []
|
|
for op_idx, op in enumerate(self.operators):
|
|
type_name = self.node_data_type(op_idx)
|
|
node_name = self.node_data_name(op_idx)
|
|
node_data_strs.append(op.generate_c_node_data(type_name, node_name))
|
|
return textwrap.indent("\n\n".join(node_data_strs), indent)
|
|
|
|
def generate_c_node_init(self, indent: str) -> str:
|
|
node_init_strs: List[str] = []
|
|
for op_idx, op in enumerate(self.operators):
|
|
tflite_node_name = self.nodes_element(op_idx)
|
|
node_data_name = self.node_data_name(op_idx)
|
|
node_init_strs.append(
|
|
op.generate_c_node_init(tflite_node_name, node_data_name))
|
|
return textwrap.indent("\n".join(node_init_strs), indent)
|
|
|
|
def generate_c_invoke(self, indent: str) -> str:
|
|
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(
|
|
body_template.substitute(
|
|
op_code=self._op_codes[op.op_code_index].full_enum_name,
|
|
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:
|
|
return f"subgraph{self.index}_tensors_"
|
|
|
|
def tensors_element(self, tensor_idx: int) -> str:
|
|
return self.tensors_array + f"[{tensor_idx}]"
|
|
|
|
def tensor_data_type(self, tensor_idx: int) -> str:
|
|
return f"Tensor{self.index}_{tensor_idx}"
|
|
|
|
def tensor_data_name(self, tensor_idx: int) -> str:
|
|
return f"tensor{self.index}_{tensor_idx}"
|
|
|
|
def generate_c_tensor_data(self, indent: str) -> str:
|
|
tensor_dims_strs: List[str] = []
|
|
for tensor_idx, tensor in enumerate(self.tensors):
|
|
type_name = self.tensor_data_type(tensor_idx)
|
|
tensor_name = self.tensor_data_name(tensor_idx)
|
|
tensor_dims_strs.append(
|
|
tensor.generate_c_tensor_dims(type_name, tensor_name))
|
|
return textwrap.indent("\n\n".join(tensor_dims_strs), indent)
|
|
|
|
def generate_c_tensor_init(self, indent: str) -> str:
|
|
tensor_init_strs: List[str] = []
|
|
for tensor_idx, tensor in enumerate(self.tensors):
|
|
tflite_tensor_name = self.tensors_element(tensor_idx)
|
|
tensor_data_name = self.tensor_data_name(tensor_idx)
|
|
tensor_init_strs.append(
|
|
tensor.generate_c_tensor_init(tflite_tensor_name, tensor_data_name))
|
|
return textwrap.indent("\n".join(tensor_init_strs), indent)
|
|
|
|
|
|
class Graph(object):
|
|
|
|
def __init__(self, model: schema_fb.ModelT):
|
|
buffers: List[tensor.Buffer] = [
|
|
tensor.Buffer("buffer_{}".format(idx), buffer)
|
|
for idx, buffer in enumerate(model.buffers)
|
|
]
|
|
self._subgraphs: List[SubGraph] = [
|
|
Subgraph(model, buffers, idx, subgraph)
|
|
for idx, subgraph in enumerate(model.subgraphs)
|
|
]
|
|
|
|
@property
|
|
def subgraphs(self) -> Sequence[Subgraph]:
|
|
return self._subgraphs
|
|
|
|
@property
|
|
def buffers(self) -> Sequence[tensor.Buffer]:
|
|
buffers: List[tensor.Buffer] = []
|
|
for subgraph in self.subgraphs:
|
|
for t in subgraph.tensors:
|
|
buffers.append(t.buffer)
|
|
return buffers
|
|
|
|
@property
|
|
def needs_zero_length_int_array(self) -> bool:
|
|
return any(subgraph.needs_zero_length_int_array
|
|
for subgraph in self.subgraphs)
|
|
|
|
|
|
class OpCodeTable(object):
|
|
|
|
def __init__(self, models: Sequence[schema_fb.ModelT]):
|
|
op_codes = []
|
|
for model in models:
|
|
for op_code in model.operatorCodes:
|
|
op_codes.append(OpCode(op_code))
|
|
|
|
self._op_codes: List([OpCode]) = list(set(op_codes))
|
|
|
|
@property
|
|
def op_codes(self) -> Sequence[OpCode]:
|
|
return self._op_codes
|