codegen/graph.py - 3p/tensorflow/tflite-micro - Git at Google


            
            # 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
	# 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