mirror of
https://github.com/vee1e/tflite-micro.git
synced 2026-09-01 17:57:27 +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
101 lines
3.8 KiB
Python
101 lines
3.8 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.
|
|
# ==============================================================================
|
|
""" Utility functions and classes for code generation. """
|
|
|
|
from typing import Any, Generator, Iterable, List, Optional, Sequence, Tuple
|
|
import string
|
|
import textwrap
|
|
import itertools
|
|
|
|
|
|
def to_pascal_case(s: str) -> str:
|
|
""" Basic function for converting snake_case to PascalCase. """
|
|
# This isn't perfect, as there might be some cases where we want underscores
|
|
# to remain if they are used as number separators.
|
|
return s.title().replace('_', '')
|
|
|
|
|
|
def bool_to_c_str(b: bool) -> str:
|
|
""" Convert a python bool value to a C bool string. Ie, False -> 'false' """
|
|
return str(b).lower()
|
|
|
|
|
|
def split_into_chunks(
|
|
data: Iterable[Any],
|
|
chunk_size: int) -> Generator[Tuple[Any, ...], None, None]:
|
|
"""Splits an iterable into chunks of a given size."""
|
|
data_iterator = iter(data)
|
|
while True:
|
|
chunk = tuple(itertools.islice(data_iterator, chunk_size))
|
|
if not chunk:
|
|
break
|
|
yield chunk
|
|
|
|
|
|
def generate_c_int_array(indent: str, int_type: str, name: str,
|
|
ints: Sequence[int]) -> str:
|
|
int_strs = ['{}'.format(i) for i in ints]
|
|
|
|
# Try to do it on a single line first
|
|
single_line_array_template = string.Template(
|
|
"constexpr ${int_type} ${name}[${size}] = {${data}};")
|
|
single_line = textwrap.indent(
|
|
single_line_array_template.substitute(int_type=int_type,
|
|
name=name,
|
|
size=len(int_strs),
|
|
data=', '.join(int_strs)), indent)
|
|
|
|
if len(single_line) < 81:
|
|
return single_line
|
|
|
|
# Couldn't fit, so split it across multiple lines
|
|
multi_line_array_template = string.Template(
|
|
"constexpr ${int_type} ${name}[${size}] = {\n"
|
|
"${body}\n"
|
|
"};\n")
|
|
|
|
lines = []
|
|
for int_strs_for_line in split_into_chunks(int_strs, 12):
|
|
ints_segment = ', '.join(int_strs_for_line)
|
|
lines.append(f' {ints_segment},')
|
|
|
|
return textwrap.indent(
|
|
multi_line_array_template.substitute(int_type=int_type,
|
|
name=name,
|
|
size=len(ints),
|
|
body='\n'.join(lines)), indent)
|
|
|
|
|
|
class IntArray(object):
|
|
""" A helper class for generating int arrays that can be used to provide the
|
|
backing storage for a TfLiteIntArray. """
|
|
|
|
def __init__(self, data: List[int]):
|
|
self._data = data
|
|
|
|
def generate_c_struct(self, type_name: str,
|
|
variable_name: Optional[str]) -> str:
|
|
struct_template = string.Template("struct ${type_name} {\n"
|
|
" int size = ${size};\n"
|
|
" int data[${size}] = {${data}};\n"
|
|
"}")
|
|
# TODO(rjascani): Make this pretty print in multi-line chunks
|
|
int_strs = ['{}'.format(i) for i in self._data]
|
|
c_struct_str = struct_template.substitute(type_name=type_name,
|
|
size=len(int_strs),
|
|
data=', '.join(int_strs))
|
|
if variable_name:
|
|
return c_struct_str + " {};".format(variable_name)
|
|
return c_struct_str + ";"
|