1# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Starting point for writing scripts to integrate TFLM with external IDEs.
16
17This script can be used to output a tree containing only the sources and headers
18needed to use TFLM for a specific configuration (e.g. target and
19optimized_kernel_implementation). This should serve as a starting
20point to integrate TFLM with external IDEs.
21
22The goal is for this script to be an interface that is maintained by the TFLM
23team and any additional scripting needed for integration with a particular IDE
24should be written external to the TFLM repository and built to work on top of
25the output tree generated with this script.
26
27We will add more documentation for a desired end-to-end integration workflow as
28we get further along in our prototyping. See this github issue for more details:
29  https://github.com/tensorflow/tensorflow/issues/47413
30"""
31
32import argparse
33import fileinput
34import os
35import re
36import shutil
37import subprocess
38
39
40def _get_dirs(file_list):
41  dirs = set()
42  for filepath in file_list:
43    dirs.add(os.path.dirname(filepath))
44  return dirs
45
46
47def _get_file_list(key, makefile_options):
48  params_list = [
49      "make", "-f", "tensorflow/lite/micro/tools/make/Makefile", key
50  ] + makefile_options.split()
51  process = subprocess.Popen(params_list,
52                             stdout=subprocess.PIPE,
53                             stderr=subprocess.PIPE)
54  stdout, stderr = process.communicate()
55
56  if process.returncode != 0:
57    raise RuntimeError("%s failed with \n\n %s" %
58                       (" ".join(params_list), stderr.decode()))
59
60  return [bytepath.decode() for bytepath in stdout.split()]
61
62
63def _third_party_src_and_dest_files(prefix_dir, makefile_options):
64  src_files = []
65  src_files.extend(_get_file_list("list_third_party_sources",
66                                  makefile_options))
67  src_files.extend(_get_file_list("list_third_party_headers",
68                                  makefile_options))
69
70  # The list_third_party_* rules give path relative to the root of the git repo.
71  # However, in the output tree, we would like for the third_party code to be a
72  # tree under prefix_dir/third_party, with the path to the tflm_download
73  # directory removed. The path manipulation logic that follows removes the
74  # downloads directory prefix, and adds the third_party prefix to create a
75  # list of destination directories for each of the third party files.
76  tflm_download_path = "tensorflow/lite/micro/tools/make/downloads"
77  dest_files = [
78      os.path.join(prefix_dir, "third_party",
79                   os.path.relpath(f, tflm_download_path)) for f in src_files
80  ]
81
82  return src_files, dest_files
83
84
85def _tflm_src_and_dest_files(prefix_dir, makefile_options):
86  src_files = []
87  src_files.extend(_get_file_list("list_library_sources", makefile_options))
88  src_files.extend(_get_file_list("list_library_headers", makefile_options))
89  dest_files = [os.path.join(prefix_dir, src) for src in src_files]
90  return src_files, dest_files
91
92
93def _get_src_and_dest_files(prefix_dir, makefile_options):
94  tflm_src_files, tflm_dest_files = _tflm_src_and_dest_files(
95      prefix_dir, makefile_options)
96  third_party_srcs, third_party_dests = _third_party_src_and_dest_files(
97      prefix_dir, makefile_options)
98
99  all_src_files = tflm_src_files + third_party_srcs
100  all_dest_files = tflm_dest_files + third_party_dests
101  return all_src_files, all_dest_files
102
103
104def _copy(src_files, dest_files):
105  for dirname in _get_dirs(dest_files):
106    os.makedirs(dirname, exist_ok=True)
107
108  for src, dst in zip(src_files, dest_files):
109    shutil.copy(src, dst)
110
111
112# For examples, we are explicitly making a deicision to not have any source
113# specialization based on the TARGET and OPTIMIZED_KERNEL_DIR. The thinking
114# here is that any target-specific sources should not be part of the TFLM
115# tree. Rather, this function will return an examples directory structure for
116# x86 and it will be the responsibility of the target-specific examples
117# repository to provide all the additional sources (and remove the unnecessary
118# sources) for the examples to run on that specific target.
119def _create_examples_tree(prefix_dir, examples_list):
120  files = []
121  for e in examples_list:
122    files.extend(_get_file_list("list_%s_example_sources" % (e), ""))
123    files.extend(_get_file_list("list_%s_example_headers" % (e), ""))
124
125  # The get_file_list gives path relative to the root of the git repo (where the
126  # examples are in tensorflow/lite/micro/examples). However, in the output
127  # tree, we would like for the examples to be under prefix_dir/examples.
128  tflm_examples_path = "tensorflow/lite/micro/examples"
129  tflm_downloads_path = "tensorflow/lite/micro/tools/make/downloads"
130
131  # Some non-example source and headers will be in the {files} list. They need
132  # special handling or they will end up outside the {prefix_dir} tree.
133  dest_file_list = []
134  for f in files:
135    if tflm_examples_path in f:
136      # file is in examples tree
137      relative_path = os.path.relpath(f, tflm_examples_path)
138      full_filename = os.path.join(prefix_dir, "examples", relative_path)
139    elif tflm_downloads_path in f:
140      # is third-party file
141      relative_path = os.path.relpath(f, tflm_downloads_path)
142      full_filename = os.path.join(prefix_dir, "third_party", relative_path)
143    else:
144      # not third-party and not examples, don't modify file name
145      # ex. tensorflow/lite/experimental/microfrontend
146      full_filename = os.path.join(prefix_dir, f)
147    dest_file_list.append(full_filename)
148
149  for dest_file, filepath in zip(dest_file_list, files):
150    dest_dir = os.path.dirname(dest_file)
151    os.makedirs(dest_dir, exist_ok=True)
152    shutil.copy(filepath, dest_dir)
153
154  # Since we are changing the directory structure for the examples, we will also
155  # need to modify the paths in the code.
156  for filepath in dest_file_list:
157    with fileinput.FileInput(filepath, inplace=True) as f:
158      for line in f:
159        include_match = re.match(
160            r'.*#include.*"' + tflm_examples_path + r'/([^/]+)/.*"', line)
161        if include_match:
162          # We need a trailing forward slash because what we care about is
163          # replacing the include paths.
164          text_to_replace = os.path.join(tflm_examples_path,
165                                         include_match.group(1)) + "/"
166          line = line.replace(text_to_replace, "")
167        # end="" prevents an extra newline from getting added as part of the
168        # in-place find and replace.
169        print(line, end="")
170
171
172def main():
173  parser = argparse.ArgumentParser(
174      description="Starting script for TFLM project generation")
175  parser.add_argument("output_dir",
176                      help="Output directory for generated TFLM tree")
177  parser.add_argument("--no_copy",
178                      action="store_true",
179                      help="Do not copy files to output directory")
180  parser.add_argument(
181      "--no_download",
182      action="store_true",
183      help="Do not download the TFLM third_party dependencies.")
184  parser.add_argument("--print_src_files",
185                      action="store_true",
186                      help="Print the src files (i.e. files in the TFLM tree)")
187  parser.add_argument(
188      "--print_dest_files",
189      action="store_true",
190      help="Print the dest files (i.e. files in the output tree)")
191  parser.add_argument("--makefile_options",
192                      default="",
193                      help="Additional TFLM Makefile options. For example: "
194                      "--makefile_options=\"TARGET=<target> "
195                      "OPTIMIZED_KERNEL_DIR=<optimized_kernel_dir> "
196                      "TARGET_ARCH=corex-m4\"")
197  parser.add_argument("--examples",
198                      "-e",
199                      action="append",
200                      help="Examples to add to the output tree. For example: "
201                      "-e hello_world -e micro_speech")
202  args = parser.parse_args()
203
204  makefile_options = args.makefile_options
205  if args.no_download:
206    makefile_options += " DISABLE_DOWNLOADS=true"
207  else:
208    # TODO(b/143904317): Explicitly call make third_party_downloads. This will
209    # no longer be needed once all the downloads are switched over to bash
210    # scripts.
211    params_list = [
212        "make", "-f", "tensorflow/lite/micro/tools/make/Makefile",
213        "third_party_downloads"
214    ] + makefile_options.split()
215    process = subprocess.Popen(params_list,
216                               stdout=subprocess.PIPE,
217                               stderr=subprocess.PIPE)
218    _, stderr = process.communicate()
219    if process.returncode != 0:
220      raise RuntimeError("%s failed with \n\n %s" %
221                         (" ".join(params_list), stderr.decode()))
222
223  src_files, dest_files = _get_src_and_dest_files(args.output_dir,
224                                                  makefile_options)
225
226  if args.print_src_files:
227    print(" ".join(src_files))
228
229  if args.print_dest_files:
230    print(" ".join(dest_files))
231
232  if args.no_copy is False:
233    _copy(src_files, dest_files)
234
235  if args.examples is not None:
236    _create_examples_tree(args.output_dir, args.examples)
237
238
239if __name__ == "__main__":
240  main()
241