Create codegen_preprocessor

This PR creates the initial scaffolding for the codegen preprocessor.
The preprocessor is a target-specific binary that will take a model,
load it into a TFLM Interpreter, perform the Init & Prepare stages, then
serialize the resulting data structures to an output file. Currently,
all this binary does is load the model file and write an output file
that simply contains the source model path in it. This will be expanded
as we expose the data.

BUG=b/295076056

codegen/BUILD

@@ -64,6 +64,7 @@ py_binary(
     deps = [
         ":graph",
         ":inference_generator",
+        "//codegen/preprocessor:preprocessor_schema_py",
         "//tensorflow/lite/tools:flatbuffer_utils",
         "@absl_py//absl:app",
         "@absl_py//absl/flags",

codegen/code_generator.py

@@ -22,18 +22,27 @@
 
 from tflite_micro.codegen import inference_generator
 from tflite_micro.codegen import graph
+from tflite_micro.codegen.preprocessor import preprocessor_schema_py_generated as preprocessor_fb
 from tflite_micro.tensorflow.lite.tools import flatbuffer_utils
 
 # Usage information:
 # Default:
-#   `bazel run codegen:code_generator -- --model=</path/to/my_model.tflite>`
+#   `bazel run codegen:code_generator -- \
+#        --model=</path/to/my_model.tflite> \
+#        --preprocessed_data=</path/to/preprocesser_output>`
 # Output will be located at: /path/to/my_model.h|cc
 
 _MODEL_PATH = flags.DEFINE_string(name="model",
                                   default=None,
                                   help="Path to the TFLite model file.",
                                   required=True)
 
+_PREPROCESSED_DATA_PATH = flags.DEFINE_string(
+    name="preprocessed_data",
+    default=None,
+    help="Path to output of codegen_preprocessor.",
+    required=True)
+
 _OUTPUT_DIR = flags.DEFINE_string(
     name="output_dir",
     default=None,
@@ -48,12 +57,27 @@
     required=False)
 
 
+def _read_preprocessed_data(
+    preprocessed_data_file: str) -> preprocessor_fb.DataT:
+  with open(preprocessed_data_file, 'rb') as file:
+    data_byte_array = bytearray(file.read())
+  return preprocessor_fb.DataT.InitFromObj(
+      preprocessor_fb.Data.GetRootAs(data_byte_array, 0))
+
+
 def main(argv: Sequence[str]) -> None:
   output_dir = _OUTPUT_DIR.value or os.path.dirname(_MODEL_PATH.value)
   output_name = _OUTPUT_NAME.value or os.path.splitext(
       os.path.basename(_MODEL_PATH.value))[0]
 
   model = flatbuffer_utils.read_model(_MODEL_PATH.value)
+  preprocessed_data = _read_preprocessed_data(_PREPROCESSED_DATA_PATH.value)
+
+  print("Generating inference code for model:\n"
+        "  model: {}\n"
+        "  preprocessed_model: {}\n".format(
+            _MODEL_PATH.value,
+            preprocessed_data.inputModelPath.decode('utf-8')))
 
   inference_generator.generate(output_dir, output_name,
                                graph.OpCodeTable([model]), graph.Graph(model))

codegen/examples/hello_world/README.md

@@ -1,18 +1,51 @@
 # Codegen Hello World Example
 
-This is a code-generated example of the hello world model.
+This is a code-generated example of the hello world model. The process is
+currently somewhat involved:
+
+## Build the preprocessor for your target
+
+This creates a target-specific preprocessor binary capable of performing the
+init and prepare stages of the Interpreter and serializing the output. This
+binary can be re-used for multiple models.
+
+### x86
+```
+make -f tensorflow/lite/micro/tools/make/Makefile codegen_preprocessor
+```
+
+## Run the preprocessor
+
+The preprocessor will take the provided model, create a TFLM Interpreter, and
+allocate tensors. It will then capture and serialize the resulting data
+structures needed for inference. For embedded targets, this should be run under
+simulation.
+
+### x86
+```
+./gen/linux_x86_64_default/bin/codegen_preprocessor \
+  $(pwd)/tensorflow/lite/micro/examples/hello_world/models/hello_world_int8.tflite \
+  $(pwd)/gen/linux_86_64_default/genfiles/hello_world_int8.ppd
+```
+
+## Generate the inference code
 
 To generate the inference code at `codegen/example/hello_world_model.h/.cc`:
 
+### x86
 ```
 bazel run codegen:code_generator -- \
   --model $(pwd)/tensorflow/lite/micro/examples/hello_world/models/hello_world_int8.tflite \
+  --preprocessed_data $(pwd)/gen/linux_86_64_default/genfiles/hello_world_int8.ppd \
   --output_dir $(pwd)/codegen/examples/hello_world \
   --output_name hello_world_model
 ```
 
-To compile the generated source, you can use the Makefile:
+## Compile the generated inference code
+
+ To compile the generated source, you can use the Makefile:
 
+### x86
 ```
 make -f tensorflow/lite/micro/tools/make/Makefile codegen_hello_world
 ```

codegen/preprocessor/Makefile.inc

@@ -0,0 +1,17 @@
+# 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.
+# ==============================================================================
+
+CODEGEN_PREPROCESSOR_SRCS := \
+$(TENSORFLOW_ROOT)codegen/preprocessor/main.cc

codegen/preprocessor/main.cc

@@ -0,0 +1,87 @@
+/* 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 <cstdlib>
+#include <fstream>
+#include <iostream>
+#include <memory>
+
+#include "codegen/preprocessor/preprocessor_schema_generated.h"
+#include "flatbuffers/flatbuffers.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace {
+
+std::unique_ptr<char[]> ReadModelFile(const char* model_file_name) {
+  std::ifstream model_file(model_file_name, std::ios::binary);
+  if (!model_file.is_open()) {
+    std::cerr << "codegen_preprocessor: could not open model file: "
+              << model_file_name << std::endl;
+    return nullptr;
+  }
+
+  model_file.seekg(0, std::ios::end);
+  size_t num_bytes = model_file.tellg();
+  std::unique_ptr<char[]> model_data(new char[num_bytes]);
+  model_file.read(model_data.get(), num_bytes);
+
+  return model_data;
+}
+
+int WriteOutputFile(const char* output_file_name,
+                    flatbuffers::span<uint8_t> output) {
+  std::ofstream output_file(output_file_name, std::ios::trunc);
+  if (!output_file.is_open()) {
+    std::cerr << "codegen_preprocessor: could not open output file: "
+              << output_file_name << std::endl;
+    return EXIT_FAILURE;
+  }
+
+  output_file.write(reinterpret_cast<char*>(output.data()), output.size());
+  return 0;
+}
+
+}  // namespace
+
+int main(int argc, char* argv[]) {
+  if (argc < 2) {
+    std::cerr << "codegen_preprocessor: invalid usage!" << std::endl;
+    std::cerr << "usage: codegen_preprocessor <tflite_model> <output_file>"
+              << std::endl;
+    return EXIT_FAILURE;
+  }
+
+  const char* model_file_name = argv[1];
+  const char* output_file_name = argv[2];
+
+  const auto model_data = ReadModelFile(model_file_name);
+  if (!model_data) {
+    return EXIT_FAILURE;
+  }
+
+  // We have to create our own allocator, as the typical TFLM runtime disables
+  // its use (to avoid dynamic allocation).
+  flatbuffers::DefaultAllocator allocator;
+  flatbuffers::FlatBufferBuilder builder{2048, &allocator};
+  const auto input_model_path = builder.CreateString(model_file_name);
+
+  // Do the preprocess work.
+
+  tflm::codegen::preprocessor::DataBuilder data_builder(builder);
+  data_builder.add_input_model_path(input_model_path);
+  builder.Finish(data_builder.Finish());
+
+  return WriteOutputFile(output_file_name, builder.GetBufferSpan());
+}

tensorflow/lite/micro/tools/make/Makefile

@@ -294,6 +294,8 @@ MICRO_LITE_BENCHMARKS := $(wildcard $(TENSORFLOW_ROOT)tensorflow/lite/micro/tool
 MICROLITE_BENCHMARK_SRCS := \
 $(wildcard $(TENSORFLOW_ROOT)tensorflow/lite/micro/tools/benchmarking/*benchmark.cc)
 
+MICRO_LITE_CODEGEN_PREPROCESSOR := $(TENSORFLOW_ROOT)codegen/preprocessor/Makefile.inc
+
 MICRO_LITE_CODEGEN_EXAMPLES := $(shell find $(TENSORFLOW_ROOT)codegen/examples/ -name Makefile.inc)
 
 MICROLITE_TEST_SRCS := \
@@ -704,6 +706,9 @@ include $(MICRO_LITE_BENCHMARKS)
 # Load custom kernel tests.
 include $(MAKEFILE_DIR)/additional_tests.inc
 
+# Load codegen preprocessor rules
+include $(MICRO_LITE_CODEGEN_PREPROCESSOR)
+
 # Create rules for downloading third-party dependencies.
 THIRD_PARTY_TARGETS :=
 $(foreach DOWNLOAD,$(THIRD_PARTY_DOWNLOADS),$(eval $(call create_download_rule,$(DOWNLOAD))))
@@ -863,6 +868,16 @@ integration_tests: $(MICROLITE_INTEGRATION_TEST_TARGETS)
 generated_micro_mutable_op_resolver: $(MICROLITE_GEN_OP_RESOLVER_TEST_TARGETS)
 endif
 
+CODEGEN_PREPROCESSOR_PATH := $(BINDIR)codegen_preprocessor
+
+codegen_preprocessor: $(CODEGEN_PREPROCESSOR_PATH)
+
+$(CODEGEN_PREPROCESSOR_PATH): $(CODEGEN_PREPROCESSOR_SRCS) $(MICROLITE_LIB_PATH)
+	@mkdir -p $(dir $@)
+	$(CXX) $(CXXFLAGS) $(INCLUDES) \
+	-o $@ $< \
+	$(MICROLITE_LIB_PATH) $(LDFLAGS) $(MICROLITE_LIBS)
+
 # Just build the test targets
 build: $(MICROLITE_BUILD_TARGETS)