chore: import upstream snapshot with attribution

tensorflow/lite/experimental/acceleration/mini_benchmark/benchmark_result_evaluator.cc

@@ -0,0 +1,36 @@
+/* Copyright 2022 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/benchmark_result_evaluator.h"
+
+#include <memory>
+
+#include "tensorflow/lite/acceleration/configuration/configuration_generated.h"
+
+namespace tflite {
+namespace acceleration {
+
+EmbeddedResultEvaluator* EmbeddedResultEvaluator::GetInstance() {
+  static EmbeddedResultEvaluator* const instance =
+      new EmbeddedResultEvaluator();
+  return instance;
+}
+
+bool EmbeddedResultEvaluator::HasPassedAccuracyCheck(
+    const BenchmarkResult& result) {
+  return result.ok();
+}
+
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/benchmark_result_evaluator.h

@@ -0,0 +1,55 @@
+/* Copyright 2022 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_BENCHMARK_RESULT_EVALUATOR_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_BENCHMARK_RESULT_EVALUATOR_H_
+
+#include "tensorflow/lite/acceleration/configuration/configuration_generated.h"
+
+namespace tflite {
+namespace acceleration {
+
+// Evaluates the BenchmarkEvent output from validator.
+class AbstractBenchmarkResultEvaluator {
+ public:
+  virtual ~AbstractBenchmarkResultEvaluator() = default;
+
+  // Returns whether this event means the validation test has passed. It checks
+  // that the test has finished successfully, and the test result passed
+  // accuracy checks.
+  bool IsValidationSuccessEvent(const BenchmarkEvent& event) {
+    return event.event_type() == BenchmarkEventType_END && event.result() &&
+           HasPassedAccuracyCheck(*event.result());
+  }
+
+  // Returns whether this BenchmarkResult should pass the accuracy check.
+  virtual bool HasPassedAccuracyCheck(const BenchmarkResult& result) = 0;
+};
+
+// Evaluator for embedded validation scenario.
+class EmbeddedResultEvaluator : public AbstractBenchmarkResultEvaluator {
+ public:
+  static EmbeddedResultEvaluator* GetInstance();
+
+  bool HasPassedAccuracyCheck(const BenchmarkResult& result) override;
+
+ private:
+  EmbeddedResultEvaluator() = default;
+  ~EmbeddedResultEvaluator() override = default;
+};
+
+}  // namespace acceleration
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_BENCHMARK_RESULT_EVALUATOR_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/big_little_affinity.cc

@@ -0,0 +1,123 @@
+/* Copyright 2021 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/big_little_affinity.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <map>
+#include <set>
+
+#include "include/cpuinfo.h"
+
+namespace tflite {
+namespace acceleration {
+
+namespace {
+bool IsInOrderArch(cpuinfo_uarch arch) {
+  switch (arch) {
+    case cpuinfo_uarch_cortex_a53:
+    case cpuinfo_uarch_cortex_a55r0:
+    case cpuinfo_uarch_cortex_a55:
+    case cpuinfo_uarch_cortex_a57:
+      return true;
+    default:
+      return false;
+  }
+  return false;
+}
+}  // namespace
+
+BigLittleAffinity GetAffinity() {
+  BigLittleAffinity affinity;
+  if (!cpuinfo_initialize()) {
+    return affinity;
+  }
+  std::map<uint32_t, uint64_t> cluster_to_max_frequency;
+  uint64_t smallest_max_frequency = UINT64_MAX;
+  uint64_t largest_max_frequency = 0;
+  uint64_t processors_count = cpuinfo_get_processors_count();
+  for (auto i = 0; i < processors_count; i++) {
+    const struct cpuinfo_processor* processor = cpuinfo_get_processor(i);
+    if (processor->core->frequency > 0) {
+      cluster_to_max_frequency[processor->cluster->cluster_id] =
+          processor->core->frequency;
+      smallest_max_frequency =
+          std::min(smallest_max_frequency, processor->core->frequency);
+      largest_max_frequency =
+          std::max(largest_max_frequency, processor->core->frequency);
+    }
+  }
+
+  int count_of_processors_with_largest_max_frequency = 0;
+  for (auto i = 0; i < cpuinfo_get_processors_count(); i++) {
+    const struct cpuinfo_processor* processor = cpuinfo_get_processor(i);
+    uint64_t max_frequency =
+        cluster_to_max_frequency[processor->cluster->cluster_id];
+    if (max_frequency == largest_max_frequency) {
+      ++count_of_processors_with_largest_max_frequency;
+    }
+  }
+  std::set<cpuinfo_uarch> archs;
+
+  // Three variants for detecting the big/little split:
+  // - all cores have the same frequency, check the uarch for in-order (on
+  //   big.LITTLE, the big cores are typically out-of-order and the LITTLE
+  //   cores in-order)
+  // - if there are 2 cores with largest max frequency, those are counted as big
+  // - otherwise the cores with smallest max frequency are counted as LITTLE
+  for (auto i = 0; i < cpuinfo_get_processors_count(); i++) {
+    const struct cpuinfo_processor* processor = cpuinfo_get_processor(i);
+    uint64_t max_frequency =
+        cluster_to_max_frequency[processor->cluster->cluster_id];
+    bool is_little;
+    archs.insert(processor->core->uarch);
+    if (count_of_processors_with_largest_max_frequency ==
+        cpuinfo_get_processors_count()) {
+      is_little = IsInOrderArch(processor->core->uarch);
+    } else if (count_of_processors_with_largest_max_frequency == 2) {
+      is_little = (max_frequency != largest_max_frequency);
+    } else {
+      is_little = (max_frequency == smallest_max_frequency);
+    }
+#ifdef __ANDROID__
+    // On desktop linux there are easily more processors than bits in an int, so
+    // skip this code. It's still convenient to enable the rest of the code on
+    // non-Android for quicker testing.
+    if (is_little) {
+      affinity.little_core_affinity |= (0x1 << processor->linux_id);
+    } else {
+      affinity.big_core_affinity |= (0x1 << processor->linux_id);
+    }
+#endif  // __ANDROID__
+  }
+  // After the detection we may have determined that all cores are big or
+  // LITTLE. This is ok if there is only one cluster or if all the cores are the
+  // same, and in that case we return the same for both masks.
+  if (cluster_to_max_frequency.size() == 1) {
+    // Only one cluster.
+    affinity.big_core_affinity = affinity.little_core_affinity =
+        std::max(affinity.big_core_affinity, affinity.little_core_affinity);
+  } else if (count_of_processors_with_largest_max_frequency ==
+                 cpuinfo_get_processors_count() &&
+             archs.size() == 1) {
+    // All cores have same uarch and frequency.
+    affinity.big_core_affinity = affinity.little_core_affinity =
+        std::max(affinity.big_core_affinity, affinity.little_core_affinity);
+  }
+  return affinity;
+}
+
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/big_little_affinity.h

@@ -0,0 +1,33 @@
+/* Copyright 2021 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_BIG_LITTLE_AFFINITY_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_BIG_LITTLE_AFFINITY_H_
+
+#include <cstdint>
+
+namespace tflite {
+namespace acceleration {
+
+struct BigLittleAffinity {
+  uint16_t big_core_affinity = 0;
+  uint16_t little_core_affinity = 0;
+};
+
+BigLittleAffinity GetAffinity();
+
+}  // namespace acceleration
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_BIG_LITTLE_AFFINITY_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/big_little_affinity_test.cc

@@ -0,0 +1,70 @@
+/* Copyright 2021 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/big_little_affinity.h"
+
+#include <cstdint>
+#include <map>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "include/cpuinfo.h"
+#include "tensorflow/lite/experimental/acceleration/compatibility/android_info.h"
+
+namespace tflite {
+namespace acceleration {
+namespace {
+
+TEST(BigLittle, CheckBasics) {
+  ASSERT_TRUE(cpuinfo_initialize());
+  auto processors_count = cpuinfo_get_processors_count();
+  ASSERT_GT(processors_count, 0);
+#if defined(__ANDROID__)
+  AndroidInfo android_info;
+  auto status = RequestAndroidInfo(&android_info);
+  if (android_info.is_emulator) {
+    std::cout << "Running on emulator\n";
+    return;
+  } else {
+    std::cout << "Running on hardware\n";
+  }
+  ASSERT_TRUE(status.ok());
+  std::map<uint32_t, uint64_t> cluster_to_max_frequency;
+  for (auto i = 0; i < cpuinfo_get_processors_count(); i++) {
+    const struct cpuinfo_processor* processor = cpuinfo_get_processor(i);
+    if (processor->core->frequency > 0) {
+      cluster_to_max_frequency[processor->cluster->cluster_id] =
+          processor->core->frequency;
+    }
+  }
+  EXPECT_GT(cluster_to_max_frequency.size(), 0);
+  EXPECT_LE(cluster_to_max_frequency.size(), 3);
+  for (auto i = 0; i < cpuinfo_get_processors_count(); i++) {
+    const struct cpuinfo_processor* processor = cpuinfo_get_processor(i);
+    EXPECT_TRUE(cluster_to_max_frequency.find(processor->cluster->cluster_id) !=
+                cluster_to_max_frequency.end());
+  }
+  BigLittleAffinity affinity = GetAffinity();
+  EXPECT_GT(affinity.little_core_affinity, 0);
+  EXPECT_GT(affinity.big_core_affinity, 0);
+  std::cout << "Little core affinity: " << std::hex
+            << affinity.little_core_affinity << std::endl;
+  std::cout << "Big core affinity: " << std::hex << affinity.big_core_affinity
+            << std::endl;
+#endif
+}
+
+}  // namespace
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/blocking_validator_runner.cc

@@ -0,0 +1,171 @@
+/* Copyright 2022 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/blocking_validator_runner.h"
+
+#include <algorithm>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/strings/str_cat.h"
+#include "absl/time/time.h"
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "tensorflow/lite/acceleration/configuration/configuration_generated.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/status_codes.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/validator.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/validator_runner_options.h"
+#include "tensorflow/lite/logger.h"
+#include "tensorflow/lite/minimal_logging.h"
+
+namespace tflite {
+namespace acceleration {
+namespace {
+
+using ::flatbuffers::FlatBufferBuilder;
+using ::flatbuffers::GetRoot;
+
+// Wait time between each query to the test result file, defined in
+// microseconds.
+constexpr absl::Duration kWaitBetweenRefresh = absl::Milliseconds(20);
+
+// Generate a string of 10 chars.
+std::string GenerateRandomString() {
+  static const char charset[] =
+      "0123456789"
+      "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+      "abcdefghijklmnopqrstuvwxyz";
+  const int size = 10;
+  std::string result;
+  result.resize(size);
+  for (int i = 0; i < size; ++i) {
+    result[i] = charset[rand() % (sizeof(charset) - 1)];
+  }
+
+  return result;
+}
+
+}  // namespace
+
+BlockingValidatorRunner::BlockingValidatorRunner(
+    const ValidatorRunnerOptions& options)
+    : per_test_timeout_ms_(options.per_test_timeout_ms),
+      storage_path_base_(options.storage_path) {
+  validator_runner_impl_ = std::make_unique<ValidatorRunnerImpl>(
+      CreateModelLoaderPath(options), options.storage_path,
+      options.data_directory_path, options.per_test_timeout_ms,
+      options.custom_input_data.empty()
+          ? nullptr
+          : std::make_unique<CustomValidationEmbedder>(
+                options.custom_input_batch_size, options.custom_input_data,
+                options.error_reporter),
+      options.error_reporter, options.nnapi_sl, options.gpu_plugin_handle,
+      options.validation_entrypoint_name, options.benchmark_result_evaluator);
+}
+
+MinibenchmarkStatus BlockingValidatorRunner::Init() {
+  return validator_runner_impl_->Init();
+}
+
+std::vector<FlatBufferBuilder> BlockingValidatorRunner::TriggerValidation(
+    const std::vector<const TFLiteSettings*>& for_settings) {
+  if (for_settings.empty()) {
+    return {};
+  }
+
+  // Create a unique storage_path.
+  std::string storage_path =
+      absl::StrCat(storage_path_base_, ".", GenerateRandomString());
+  TFLITE_LOG_PROD(TFLITE_LOG_INFO, "Validation storage path: %s",
+                  storage_path.c_str());
+  std::vector<flatbuffers::FlatBufferBuilder> to_be_run;
+  std::vector<TFLiteSettingsT> for_settings_obj;
+  for_settings_obj.reserve(for_settings.size());
+  for (auto settings : for_settings) {
+    TFLiteSettingsT tflite_settings;
+    settings->UnPackTo(&tflite_settings);
+    flatbuffers::FlatBufferBuilder copy;
+    copy.Finish(CreateTFLiteSettings(copy, &tflite_settings));
+    to_be_run.emplace_back(std::move(copy));
+    for_settings_obj.emplace_back(tflite_settings);
+  }
+  validator_runner_impl_->TriggerValidationAsync(std::move(to_be_run),
+                                                 storage_path);
+
+  // The underlying process runner should ensure each test finishes on time or
+  // timed out. deadline_us is added here as an extra safety guard.
+  int64_t total_timeout_ms = per_test_timeout_ms_ * (1 + for_settings.size());
+  int64_t deadline_us = Validator::BootTimeMicros() + total_timeout_ms * 1000;
+
+  bool within_timeout = true;
+  // TODO(b/249274787): GetNumCompletedResults() loads the file from disk each
+  // time when called. We should find a way to optimize the FlatbufferStorage to
+  // reduce the I/O and remove the sleep().
+  while ((validator_runner_impl_->GetNumCompletedResults()) <
+             for_settings.size() &&
+         (within_timeout = Validator::BootTimeMicros() < deadline_us)) {
+    usleep(absl::ToInt64Microseconds(kWaitBetweenRefresh));
+  }
+
+  std::vector<FlatBufferBuilder> results =
+      validator_runner_impl_->GetCompletedResults();
+  if (!within_timeout) {
+    TFLITE_LOG_PROD(
+        TFLITE_LOG_WARNING,
+        "Validation timed out after %ld ms. Return before all tests finished.",
+        total_timeout_ms);
+  } else if (for_settings.size() != results.size()) {
+    TFLITE_LOG_PROD(TFLITE_LOG_WARNING,
+                    "Validation completed.Started benchmarking for %d "
+                    "TFLiteSettings, received %d results.",
+                    for_settings.size(), results.size());
+  }
+
+  // If there are any for_settings missing from results, add an error event.
+  std::vector<TFLiteSettingsT> result_settings;
+  result_settings.reserve(results.size());
+  for (auto& result : results) {
+    const BenchmarkEvent* event =
+        GetRoot<BenchmarkEvent>(result.GetBufferPointer());
+    TFLiteSettingsT event_settings;
+    event->tflite_settings()->UnPackTo(&event_settings);
+    result_settings.emplace_back(std::move(event_settings));
+  }
+  for (auto& settings_obj : for_settings_obj) {
+    auto result_it =
+        std::find(result_settings.begin(), result_settings.end(), settings_obj);
+    if (result_it == result_settings.end()) {
+      FlatBufferBuilder fbb;
+      fbb.Finish(CreateBenchmarkEvent(
+          fbb, CreateTFLiteSettings(fbb, &settings_obj),
+          BenchmarkEventType_ERROR, /* result */ 0,
+          CreateBenchmarkError(fbb, BenchmarkStage_UNKNOWN,
+                               /* exit_code */ 0, /* signal */ 0,
+                               /* error_code */ 0,
+                               /* mini_benchmark_error_code */
+                               kMinibenchmarkCompletionEventMissing),
+          Validator::BootTimeMicros(), Validator::WallTimeMicros()));
+      results.emplace_back(std::move(fbb));
+    }
+  }
+  // Delete storage_file before returning. In case of test timeout, the child
+  // thread or process may create and continue to write to the storage_path. In
+  // this case we cannot delete the file.
+  (void)unlink(storage_path.c_str());
+  return results;
+}
+
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/blocking_validator_runner.h

@@ -0,0 +1,59 @@
+/* Copyright 2022 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_BLOCKING_VALIDATOR_RUNNER_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_BLOCKING_VALIDATOR_RUNNER_H_
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "tensorflow/lite/acceleration/configuration/configuration_generated.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/status_codes.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/validator_runner_impl.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/validator_runner_options.h"
+
+namespace tflite {
+namespace acceleration {
+
+// Class that runs mini-benchmark validation in a separate process and gives
+// access to the results. This class provides a synchronous API for the callers
+// to wait until the all the tests have finished.
+//
+// This class is thread-safe when using different storage_path_. When
+// storage_path_ is shared between multiple runners, they will interfere with
+// each other.
+class BlockingValidatorRunner {
+ public:
+  explicit BlockingValidatorRunner(const ValidatorRunnerOptions& options);
+
+  MinibenchmarkStatus Init();
+
+  // Trigger the validation tests with for_settings, and return the test result.
+  // Each for_settings will have a corresponding result. The result is of schema
+  // BenchmarkEvent.
+  std::vector<flatbuffers::FlatBufferBuilder> TriggerValidation(
+      const std::vector<const TFLiteSettings*>& for_settings);
+
+ private:
+  int per_test_timeout_ms_ = 0;
+  const std::string storage_path_base_;
+  std::unique_ptr<ValidatorRunnerImpl> validator_runner_impl_;
+};
+
+}  // namespace acceleration
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_BLOCKING_VALIDATOR_RUNNER_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/blocking_validator_runner_test.cc

@@ -0,0 +1,262 @@
+/* Copyright 2022 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/blocking_validator_runner.h"
+
+#include <fcntl.h>
+
+#include <iostream>
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/strings/str_cat.h"
+#include "flatbuffers/buffer.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "tensorflow/lite/acceleration/configuration/configuration_generated.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/benchmark_result_evaluator.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/embedded_mobilenet_model.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/embedded_mobilenet_validation_model.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/mini_benchmark_test_helper.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/status_codes.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/validator_runner_options.h"
+
+namespace tflite {
+namespace acceleration {
+namespace {
+
+using ::flatbuffers::FlatBufferBuilder;
+using ::flatbuffers::GetRoot;
+
+class CustomResultEvaluator : public AbstractBenchmarkResultEvaluator {
+ public:
+  bool HasPassedAccuracyCheck(const BenchmarkResult& result) override {
+    return true;
+  }
+};
+
+class BlockingValidatorRunnerTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    MiniBenchmarkTestHelper helper;
+    should_perform_test_ = helper.should_perform_test();
+    options_.model_path = helper.DumpToTempFile(
+        "mobilenet_quant_with_validation.tflite",
+        g_tflite_acceleration_embedded_mobilenet_validation_model,
+        g_tflite_acceleration_embedded_mobilenet_validation_model_len);
+    ASSERT_TRUE(!options_.model_path.empty());
+
+    options_.data_directory_path = ::testing::TempDir();
+    options_.storage_path =
+        absl::StrCat(::testing::TempDir(), "storage_path.fb.1");
+    options_.per_test_timeout_ms = 5000;
+
+    plain_model_path_ = MiniBenchmarkTestHelper::DumpToTempFile(
+        "mobilenet_quant.tflite",
+        g_tflite_acceleration_embedded_mobilenet_model,
+        g_tflite_acceleration_embedded_mobilenet_model_len);
+  }
+
+  std::string plain_model_path_;
+  ValidatorRunnerOptions options_;
+  bool should_perform_test_ = true;
+};
+
+TEST_F(BlockingValidatorRunnerTest, SucceedWithEmbeddedValidation) {
+  if (!should_perform_test_) {
+    std::cerr << "Skipping test";
+    return;
+  }
+
+  BlockingValidatorRunner runner(options_);
+  ASSERT_EQ(runner.Init(), kMinibenchmarkSuccess);
+  FlatBufferBuilder fbb;
+#ifdef __ANDROID__
+  fbb.Finish(CreateTFLiteSettings(fbb, Delegate_GPU));
+#else
+  fbb.Finish(CreateTFLiteSettings(fbb));
+#endif  // __ANDROID__
+
+  std::vector<FlatBufferBuilder> results = runner.TriggerValidation(
+      {flatbuffers::GetRoot<TFLiteSettings>(fbb.GetBufferPointer())});
+  EXPECT_THAT(results, testing::Not(testing::IsEmpty()));
+  for (auto& result : results) {
+    const BenchmarkEvent* event =
+        GetRoot<BenchmarkEvent>(result.GetBufferPointer());
+    EXPECT_EQ(event->event_type(), BenchmarkEventType_END);
+    EXPECT_TRUE(event->result()->ok());
+  }
+}
+
+TEST_F(BlockingValidatorRunnerTest, SucceedWithFdCloexecEmbeddedValidation) {
+  if (!should_perform_test_) {
+    std::cerr << "Skipping test";
+    return;
+  }
+
+  options_.model_fd = open(options_.model_path.c_str(), O_RDONLY | O_CLOEXEC);
+  ASSERT_GE(options_.model_fd, 0);
+  struct stat stat_buf = {0};
+  ASSERT_EQ(fstat(options_.model_fd, &stat_buf), 0);
+  options_.model_size = stat_buf.st_size;
+  options_.model_offset = 0;
+  options_.model_path.clear();
+
+  BlockingValidatorRunner runner(options_);
+  ASSERT_EQ(runner.Init(), kMinibenchmarkSuccess);
+  FlatBufferBuilder fbb;
+#ifdef __ANDROID__
+  fbb.Finish(CreateTFLiteSettings(fbb, Delegate_GPU));
+#else
+  fbb.Finish(CreateTFLiteSettings(fbb));
+#endif  // __ANDROID__
+
+  std::vector<FlatBufferBuilder> results = runner.TriggerValidation(
+      {flatbuffers::GetRoot<TFLiteSettings>(fbb.GetBufferPointer())});
+  EXPECT_THAT(results, testing::Not(testing::IsEmpty()));
+  for (auto& result : results) {
+    const BenchmarkEvent* event =
+        GetRoot<BenchmarkEvent>(result.GetBufferPointer());
+    EXPECT_EQ(event->event_type(), BenchmarkEventType_END);
+    EXPECT_TRUE(event->result()->ok());
+  }
+}
+
+TEST_F(BlockingValidatorRunnerTest, SucceedWithBufferModel) {
+  if (!should_perform_test_) {
+    std::cerr << "Skipping test";
+    return;
+  }
+
+  options_.model_buffer =
+      g_tflite_acceleration_embedded_mobilenet_validation_model;
+  options_.model_size =
+      g_tflite_acceleration_embedded_mobilenet_validation_model_len;
+  options_.model_path.clear();
+
+  BlockingValidatorRunner runner(options_);
+  ASSERT_EQ(runner.Init(), kMinibenchmarkSuccess);
+  FlatBufferBuilder fbb;
+  fbb.Finish(CreateTFLiteSettings(fbb));
+
+  std::vector<FlatBufferBuilder> results = runner.TriggerValidation(
+      {flatbuffers::GetRoot<TFLiteSettings>(fbb.GetBufferPointer())});
+  EXPECT_THAT(results, testing::Not(testing::IsEmpty()));
+  for (auto& result : results) {
+    const BenchmarkEvent* event =
+        GetRoot<BenchmarkEvent>(result.GetBufferPointer());
+    EXPECT_EQ(event->event_type(), BenchmarkEventType_END);
+    EXPECT_TRUE(event->result()->ok());
+  }
+}
+
+TEST_F(BlockingValidatorRunnerTest, SucceedWithFdModelCustomValidation) {
+  if (!should_perform_test_) {
+    std::cerr << "Skipping test";
+    return;
+  }
+
+  options_.model_path.clear();
+  options_.model_fd = open(plain_model_path_.c_str(), O_RDONLY);
+  ASSERT_GE(options_.model_fd, 0);
+  struct stat stat_buf = {0};
+  ASSERT_EQ(fstat(options_.model_fd, &stat_buf), 0);
+  options_.model_size = stat_buf.st_size;
+  options_.model_offset = 0;
+  options_.custom_input_batch_size = 3;
+  options_.custom_input_data = {std::vector<uint8_t>(3 * 224 * 224 * 3, 1)};
+  CustomResultEvaluator evaluator;
+  options_.benchmark_result_evaluator = &evaluator;
+
+  BlockingValidatorRunner runner(options_);
+  ASSERT_EQ(runner.Init(), kMinibenchmarkSuccess);
+  FlatBufferBuilder fbb;
+#ifdef __ANDROID__
+  fbb.Finish(CreateTFLiteSettings(fbb, Delegate_XNNPACK));
+#else
+  fbb.Finish(CreateTFLiteSettings(fbb));
+#endif  // __ANDROID__
+
+  std::vector<FlatBufferBuilder> results = runner.TriggerValidation(
+      {flatbuffers::GetRoot<TFLiteSettings>(fbb.GetBufferPointer())});
+  EXPECT_THAT(results, testing::Not(testing::IsEmpty()));
+  for (auto& result : results) {
+    const BenchmarkEvent* event =
+        GetRoot<BenchmarkEvent>(result.GetBufferPointer());
+    EXPECT_EQ(event->event_type(), BenchmarkEventType_END);
+  }
+}
+
+TEST_F(BlockingValidatorRunnerTest, SucceedWhenRunningMultipleTimes) {
+  if (!should_perform_test_) {
+    std::cerr << "Skipping test";
+    return;
+  }
+
+  BlockingValidatorRunner runner(options_);
+  ASSERT_EQ(runner.Init(), kMinibenchmarkSuccess);
+  FlatBufferBuilder fbb;
+  fbb.Finish(CreateTFLiteSettings(fbb));
+
+  int num_runs = 3;
+  for (int i = 0; i < num_runs; i++) {
+    std::vector<FlatBufferBuilder> results = runner.TriggerValidation(
+        {flatbuffers::GetRoot<TFLiteSettings>(fbb.GetBufferPointer()),
+         flatbuffers::GetRoot<TFLiteSettings>(fbb.GetBufferPointer())});
+    EXPECT_THAT(results, testing::Not(testing::IsEmpty()));
+    for (auto& result : results) {
+      const BenchmarkEvent* event =
+          GetRoot<BenchmarkEvent>(result.GetBufferPointer());
+      EXPECT_EQ(event->event_type(), BenchmarkEventType_END);
+      EXPECT_TRUE(event->result()->ok());
+    }
+  }
+}
+
+TEST_F(BlockingValidatorRunnerTest, ReturnErrorWhenTimedOut) {
+  if (!should_perform_test_) {
+    std::cerr << "Skipping test";
+    return;
+  }
+  options_.per_test_timeout_ms = 50;
+  BlockingValidatorRunner runner(options_);
+  ASSERT_EQ(runner.Init(), kMinibenchmarkSuccess);
+  FlatBufferBuilder fbb;
+  fbb.Finish(CreateTFLiteSettings(fbb));
+
+  std::vector<FlatBufferBuilder> results = runner.TriggerValidation(
+      {flatbuffers::GetRoot<TFLiteSettings>(fbb.GetBufferPointer())});
+  EXPECT_THAT(results, testing::SizeIs(1));
+  for (auto& result : results) {
+    const BenchmarkEvent* event =
+        GetRoot<BenchmarkEvent>(result.GetBufferPointer());
+    EXPECT_EQ(event->event_type(), BenchmarkEventType_ERROR);
+    ASSERT_NE(nullptr, event->error());
+    // The timeout can result in two different behaviors:
+    // 1. The popen() subprocess got killed by the detached thread because the
+    // timeout has reached, and the thread wrote error code
+    // kMinibenchmarkCommandTimedOut, or
+    // 2. The thread didn't respond the main process in time, and the main
+    // process returned after the timeout, with error code
+    // kMinibenchmarkCompletionEventMissing.
+    EXPECT_THAT(event->error()->mini_benchmark_error_code(),
+                testing::AnyOf(kMinibenchmarkCommandTimedOut,
+                               kMinibenchmarkCompletionEventMissing));
+  }
+}
+
+}  // namespace
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/build_defs.bzl

@@ -0,0 +1,262 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Helpers for mini-benchmark build rules."""
+
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+load(
+    "//tensorflow:tensorflow.bzl",
+    "clean_dep",
+)
+load("//tensorflow/lite/core/shims:cc_library_with_tflite.bzl", "add_suffix")
+load("//tensorflow/lite/experimental/acceleration/mini_benchmark:special_rules.bzl", "libjpeg_handle_deps")
+
+def _concat(lists):
+    """Concatenate a list of lists, without requiring the inner lists to be iterable.
+
+    This allows the inner lists to be obtained by calls to select().
+    """
+    result = []
+    for selected_list in lists:
+        result = result + selected_list
+    return result
+
+def embedded_binary(name, binary, array_variable_name, testonly = False, exec_properties = None):
+    """Create a cc_library that embeds a binary as constant data.
+
+    Args:
+        name: name for the generated cc_library target, and the base name for
+              generated header file
+        binary: binary file to be embedded
+        array_variable_name: name of the constant array for the data.
+    """
+    cc_name = "%s.cc" % name
+    h_name = "%s.h" % name
+    native.genrule(
+        name = name + "_src",
+        srcs = [binary],
+        outs = [
+            cc_name,
+            h_name,
+        ],
+        cmd = """
+      $(location //tensorflow/lite/experimental/acceleration/compatibility:convert_binary_to_cc_source) \
+          --input_binary_file $(location %s) \
+          --output_header_file $(location :%s) \
+          --output_source_file $(location :%s) \
+          --array_variable_name %s
+      """ % (binary, h_name, cc_name, array_variable_name),
+        tools = ["//tensorflow/lite/experimental/acceleration/compatibility:convert_binary_to_cc_source"],
+        testonly = testonly,
+    )
+    cc_library(
+        name = name,
+        srcs = [cc_name],
+        hdrs = [h_name],
+        testonly = testonly,
+        exec_properties = exec_properties,
+    )
+
+def validation_model(
+        name,
+        main_model,
+        metrics_model,
+        jpegs,
+        scale = "",
+        zeropoint = "",
+        use_ondevice_cpu_for_golden = False,
+        testonly = 0):
+    """Create a tflite model with embedded validation.
+
+    Args:
+        name: name of the target. A file called 'name'.tflite is generated
+        main_model: main tflite model target
+        metrics_model: metrics tflite model target
+        jpegs: target with 1 or more jpeg files
+        scale: the input (de)quantization scale parameter for float models
+        zeropoint: the input (de)quantization zeropoint parameter for float models
+        use_ondevice_cpu_for_golden: use on-device CPU for golden data (rather than embedding)
+        testonly: whether target is marked testonly
+    """
+    if use_ondevice_cpu_for_golden:
+        use_ondevice_cpu_for_golden = "true"
+    else:
+        use_ondevice_cpu_for_golden = "false"
+    scale_arg = ""
+    zeropoint_arg = ""
+    if scale:
+        scale_arg = "--scale=" + scale
+        zeropoint_arg = "--zero_point=" + zeropoint
+    schema_location = "//tensorflow/compiler/mlir/lite/schema:schema.fbs"
+    native.genrule(
+        name = name,
+        testonly = testonly,
+        srcs = [
+            main_model,
+            jpegs,
+            schema_location,
+            metrics_model,
+        ],
+        outs = [name + ".tflite"],
+        cmd = """
+          JPEGS='$(locations %s)'
+          JPEGS=$${JPEGS// /,}
+          $(location //tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier:embedder_cmdline) \
+              --schema=$(location %s) \
+              --main_model=$(location %s) \
+              --metrics_model=$(location %s) \
+              %s %s \
+              --jpegs=$$JPEGS \
+              --use_ondevice_cpu_for_golden=%s \
+              --output='$(@D)/%s.tflite.tmp'
+          $(location //tensorflow/lite/experimental/acceleration/mini_benchmark:copy_associated_files) \
+              '$(@D)/%s.tflite.tmp' \
+              $(location %s) \
+              $(location %s.tflite)
+          rm '$(@D)/%s.tflite.tmp'
+        """ % (
+            jpegs,
+            schema_location,
+            main_model,
+            metrics_model,
+            scale_arg,
+            zeropoint_arg,
+            use_ondevice_cpu_for_golden,
+            name,
+            name,
+            main_model,
+            name,
+            name,
+        ),
+        tools = [
+            "//tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier:embedder_cmdline",
+            "//tensorflow/lite/experimental/acceleration/mini_benchmark:copy_associated_files",
+        ],
+    )
+
+def validation_test(name, validation_model, tags = [], copts = [], deps = []):
+    """Create a test binary for the given model with validation.
+
+    Args:
+        name: name of the target.
+        validation_model: tflite model with validation target.
+        tags: to be passed to cc_test.
+        copts: to be passed to cc_test.
+        deps: to be passed to cc_test.
+    """
+    embed_name = name + "_embed_model"
+    embedded_binary(
+        embed_name,
+        binary = validation_model,
+        array_variable_name = "g_tflite_acceleration_" + name + "_model",
+    )
+    cc_test(
+        name = name,
+        srcs = ["//tensorflow/lite/experimental/acceleration/mini_benchmark:model_validation_test.cc"],
+        tags = tags + ["no_mac", "no_windows", "tflite_not_portable_ios"],
+        copts = copts + [
+            "-DTENSORFLOW_ACCELERATION_MODEL_DATA_VARIABLE=\"g_tflite_acceleration_%s_model\"" % name,
+            "-DTENSORFLOW_ACCELERATION_MODEL_LENGTH_VARIABLE=\"g_tflite_acceleration_%s_model_len\"" % name,
+        ],
+        deps = deps + [
+            embed_name,
+            "@com_google_googletest//:gtest_main",
+            "@flatbuffers",
+            "//tensorflow/lite/acceleration/configuration:configuration_fbs",
+            "//tensorflow/lite/acceleration/configuration:nnapi_plugin",
+            "//tensorflow/lite/experimental/acceleration/compatibility:android_info",
+            "//tensorflow/lite/experimental/acceleration/mini_benchmark:big_little_affinity",
+            "//tensorflow/lite/experimental/acceleration/mini_benchmark:status_codes",
+            "//tensorflow/lite/experimental/acceleration/mini_benchmark:validator",
+            "//tensorflow/lite/tools:model_loader",
+        ] + select({
+            clean_dep("//tensorflow:android"): [
+                "//tensorflow/lite/acceleration/configuration:gpu_plugin",
+            ],
+            "//conditions:default": [],
+        }) + libjpeg_handle_deps(),
+        linkstatic = 1,
+    )
+
+def cc_library_with_forced_in_process_benchmark_variant(
+        name,
+        deps = [],
+        forced_in_process_deps = [],
+        in_process_deps = [],
+        non_in_process_deps_selects = [],
+        **kwargs):
+    """Defines a cc_library that optionally forces benchmark runs in process.
+
+    This generates two cc_library target. The first one runs the benchmark in a
+    separate process on Android, while it runs the benchmark in process on all
+    other platforms. It doesn't have TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS
+    defined.
+    The second one, which has "_in_process" appended to the name, forces
+    benchmark runs in process on all platforms. It has
+    TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS defined.
+
+    The default option for MiniBenchmark is to run the benchmark in a separate
+    process on Android, as this is safer than running the benchmark in the app
+    process. However, forcing the benchmark to run in-process on Android allows
+    the benchmark to reuse the same TF Lite runtime that is initialized in the
+    application process. These two variants may use different dependencies.
+    For example, the in-process variant uses the statically linked libjpeg
+    handle, while the other variant uses the dynamically linked libjpeg handle
+    on Android to minimize binary size.
+
+    This build rule ensures that the dependencies listed in
+    "forced_in_process_deps" are added only when
+    TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS is defined, that the dependencies
+    listed in "non_in_process_deps_selects" are added only when
+    TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS is NOT defined, and that
+    TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS is defined automatically when
+    using the "_in_process" target.
+
+
+    Args:
+      name: determines the name used for the generated cc_library targets.
+      forced_in_process_deps: dependencies that will be enabled only when the
+        benchmark is forced to run in-process on all platforms. This should be
+        used for dependencies arising from code inside
+        '#ifdef TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS'.
+      deps: dependencies that will be unconditionally included in the deps of
+        the generated cc_library targets.
+      in_process_deps: dependencies on rules that are themselves defined using
+        'cc_library_with_forced_in_process_benchmark_variant'. Must be
+        iterable, so cannot be computed by calling 'select'.
+      non_in_process_deps_selects: A list of dictionaries that will be
+        converted to dependencies with select on rules. The dependencies will
+        be enabled only when the benchmark runs in a separate process on
+        Android. This should be used for dependencies arising from code inside
+        '#ifndef TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS'.
+      **kwargs:
+        Additional cc_library parameters.
+    """
+    cc_library(
+        name = name,
+        deps = deps + in_process_deps + _concat([select(map) for map in non_in_process_deps_selects]) + [
+            "//tensorflow/lite/experimental/acceleration/mini_benchmark:tflite_acceleration_in_process_default",
+        ],
+        **kwargs
+    )
+
+    in_process_deps_renamed = [add_suffix(in_process_dep, "_in_process") for in_process_dep in in_process_deps]
+    cc_library(
+        name = name + "_in_process",
+        deps = deps + in_process_deps_renamed + forced_in_process_deps + [
+            "//tensorflow/lite/experimental/acceleration/mini_benchmark:tflite_acceleration_in_process_enable",
+        ],
+        **kwargs
+    )

tensorflow/lite/experimental/acceleration/mini_benchmark/c/BUILD

@@ -0,0 +1,44 @@
+# 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.
+# ==============================================================================
+load("//tensorflow/lite:build_def.bzl", "tflite_cc_library_with_c_headers_test")
+load("//tensorflow/lite/core/shims:cc_library_with_tflite.bzl", "cc_library_with_tflite_with_c_headers_test")
+load("//tensorflow/lite/experimental/acceleration/mini_benchmark:special_rules.bzl", "minibenchmark_visibility_allowlist")
+
+default_visibility_group = [
+    "//tensorflow/lite/experimental/acceleration/mini_benchmark:__subpackages__",
+] + minibenchmark_visibility_allowlist()
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = default_visibility_group,
+    licenses = ["notice"],
+)
+
+# This target runs MiniBenchmark in a separate process on Android, while it runs MiniBenchmark
+# in-process on all other platforms.
+cc_library_with_tflite_with_c_headers_test(
+    name = "c_api",
+    hdrs = ["c_api.h"],
+    deps = ["//tensorflow/lite/core/experimental/acceleration/mini_benchmark/c:c_api"],
+)
+
+# This target forces MiniBenchmark to run in-process on all platforms including Android.
+tflite_cc_library_with_c_headers_test(
+    name = "c_api_in_process",
+    hdrs = ["c_api.h"],
+    deps = [
+        "//tensorflow/lite/core/experimental/acceleration/mini_benchmark/c:c_api_in_process",
+    ],
+)

tensorflow/lite/experimental/acceleration/mini_benchmark/c/c_api.h

@@ -0,0 +1,23 @@
+/* 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_C_C_API_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_C_C_API_H_
+
+/// For documentation, see
+/// third_party/tensorflow/lite/core/experimental/acceleration/mini_benchmark/c/c_api.h
+
+#include "tensorflow/lite/core/experimental/acceleration/mini_benchmark/c/c_api.h"  // IWYU pragma: export
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_C_C_API_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/call.cc

@@ -0,0 +1,288 @@
+/* Copyright 2020 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 <stddef.h>
+
+#include <cstring>
+#include <sstream>
+#include <string>
+#include <vector>
+
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/context_util.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/call_register.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/op_macros.h"
+
+namespace tflite {
+namespace acceleration {
+namespace ops {
+namespace call_kernel {
+
+namespace {
+
+bool MatchDimensionsExceptBatchSize(TfLiteTensor* a, TfLiteTensor* b) {
+  if (a->dims->size != b->dims->size) {
+    return false;
+  }
+  // First dimension is the batch size.
+  for (int i = 1; i < a->dims->size; ++i) {
+    if (a->dims->data[i] != b->dims->data[i]) {
+      return false;
+    }
+  }
+  return true;
+}
+
+// Verifies that number, shape and type of inputs match for subgraph and CALL
+// node. If shape of inputs is unset for subgraph, it is inferred.
+TfLiteStatus ValidateAndResizeInputsIfNeeded(TfLiteContext* context,
+                                             TfLiteNode* node,
+                                             Subgraph* subgraph,
+                                             int loop_count) {
+  // Match number of inputs for subgraph and CALL node.
+  TF_LITE_ENSURE_EQ(context, subgraph->inputs().size(), node->inputs->size);
+
+  for (int i = 0; i < node->inputs->size; ++i) {
+    TfLiteTensor* node_input = context->tensors + node->inputs->data[i];
+    TfLiteTensor* subgraph_input = subgraph->tensor(subgraph->inputs()[i]);
+    // Match input types.
+    TF_LITE_ENSURE_TYPES_EQ(context, node_input->type, subgraph_input->type);
+    TF_LITE_ENSURE_MSG(
+        context, node_input->dims->size > 0,
+        "Dimensions of all of call node's inputs should be non-zero.");
+    // Ensure batch size of CALL node's input is same as loop size.
+    TF_LITE_ENSURE_EQ(context, node_input->dims->data[0], loop_count);
+    if (!subgraph_input->dims->size) {
+      // Subgraph input dimensions unset and will be inferred.
+      std::vector<int> new_dims;
+      new_dims.reserve(node_input->dims->size);
+      new_dims.push_back(1);  // Batch size is fixed as 1 for subgraph.
+      new_dims.insert(new_dims.end(), node_input->dims->data + 1,
+                      node_input->dims->data + node_input->dims->size);
+      subgraph->ResizeInputTensor(subgraph->inputs()[i], new_dims);
+    } else {
+      // Dimensions already set for subgraph, match input dimensions.
+      if (!MatchDimensionsExceptBatchSize(node_input, subgraph_input)) {
+        std::stringstream node_input_dims, subgraph_input_dims;
+        for (int i = 0; i < node_input->dims->size; i++) {
+          node_input_dims << node_input->dims->data[i] << " ";
+          subgraph_input_dims << subgraph_input->dims->data[i] << " ";
+        }
+        TF_LITE_KERNEL_LOG(
+            context,
+            "%s:%d: All dimensions except the batch size should match for call "
+            "node and the subgraph to invoke (input tensor %s[ %s], subgraph "
+            "tensor %s[ %s])",
+            __FILE__, __LINE__, node_input->name, node_input_dims.str().c_str(),
+            subgraph_input->name, subgraph_input_dims.str().c_str());
+        return kTfLiteError;
+      }
+      // Batch size of subgraph's input should be 1.
+      TF_LITE_ENSURE_EQ(context, subgraph_input->dims->data[0], 1);
+    }
+  }
+
+  return kTfLiteOk;
+}
+
+// Verifies that number of outputs match for subgraph and CALL node. Infer the
+// shape and type of outputs for CALL node and resize accordingly.
+TfLiteStatus ValidateAndResizeOutputs(TfLiteContext* context, TfLiteNode* node,
+                                      Subgraph* subgraph, int loop_count) {
+  // Match number of outputs for subgraph and CALL node.
+  TF_LITE_ENSURE_EQ(context, subgraph->outputs().size(), node->outputs->size);
+
+  // Infer output shape for the CALL node.
+  for (int i = 0; i < node->outputs->size; ++i) {
+    const TfLiteTensor* subgraph_output =
+        subgraph->tensor(subgraph->outputs()[i]);
+    TfLiteTensor* node_output = context->tensors + node->outputs->data[i];
+
+    TF_LITE_ASSERT(subgraph_output->dims->size > 0);
+    TfLiteIntArray* new_dims_array = TfLiteIntArrayCopy(subgraph_output->dims);
+    // Batch size is fixed as `loop_count` for CALL node.
+    new_dims_array->data[0] = loop_count;
+    TF_LITE_ENSURE_OK(
+        context, context->ResizeTensor(context, node_output, new_dims_array));
+    node_output->type = subgraph_output->type;
+  }
+  return kTfLiteOk;
+}
+
+// Copy input tensor data from CALL node's inputs to subgraph's inputs.
+TfLiteStatus CopyInputTensorsData(TfLiteContext* context, TfLiteNode* node,
+                                  Subgraph* dst_subgraph, int loop_index,
+                                  int loop_count) {
+  const std::vector<int>& dst_tensor_indices = dst_subgraph->inputs();
+  TF_LITE_ENSURE_EQ(context, node->inputs->size, dst_tensor_indices.size());
+  for (int i = 0; i < dst_tensor_indices.size(); ++i) {
+    TfLiteTensor* src_tensor = context->tensors + node->inputs->data[i];
+    TfLiteTensor* dst_tensor = dst_subgraph->tensor(dst_tensor_indices[i]);
+    size_t offset = src_tensor->bytes / loop_count * loop_index;
+    TF_LITE_ENSURE_EQ(context, src_tensor->bytes / loop_count,
+                      dst_tensor->bytes);
+    memcpy(dst_tensor->data.raw, src_tensor->data.raw + offset,
+           src_tensor->bytes / loop_count);
+  }
+  return kTfLiteOk;
+}
+
+// Copy tensor data from subgraph's outputs to CALL node's outputs.
+TfLiteStatus CopyOutputTensorsData(TfLiteContext* context,
+                                   Subgraph* src_subgraph, TfLiteNode* node,
+                                   int loop_index, int loop_count) {
+  const std::vector<int>& src_tensor_indices = src_subgraph->outputs();
+  TF_LITE_ENSURE_EQ(context, src_tensor_indices.size(), node->outputs->size);
+  for (int i = 0; i < src_tensor_indices.size(); ++i) {
+    const TfLiteTensor* src_tensor =
+        src_subgraph->tensor(src_tensor_indices[i]);
+    TfLiteTensor* dst_tensor = context->tensors + node->outputs->data[i];
+    size_t offset = dst_tensor->bytes / loop_count * loop_index;
+    TF_LITE_ENSURE_EQ(context, src_tensor->bytes,
+                      dst_tensor->bytes / loop_count);
+    memcpy(dst_tensor->data.raw + offset, src_tensor->data.raw,
+           src_tensor->bytes);
+  }
+  return kTfLiteOk;
+}
+
+}  // namespace
+
+struct OpData {
+  // Index of the subgraph that needs to be invoked.
+  // Subgraph should have batch size 1.
+  int subgraph_index;
+  // The number of times the CALL op should call the subgraph.
+  // The inputs to the call op are expected to have this value as their batch
+  // size.
+  int loop_count;
+};
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  if (!buffer) {
+    return nullptr;
+  }
+  auto* op_data = new OpData;
+  const uint8_t* buffer_fixed_width = reinterpret_cast<const uint8_t*>(buffer);
+  const flexbuffers::Map& map =
+      flexbuffers::GetRoot(buffer_fixed_width, length).AsMap();
+  // Note: The values below will be set as 0 if the parsing fails or if the
+  // values have been unset.
+  op_data->subgraph_index = map["subgraph_index"].AsInt32();
+  op_data->loop_count = map["loop_count"].AsInt32();
+  return op_data;
+}
+
+void Free(TfLiteContext* context, void* buffer) {
+  delete reinterpret_cast<OpData*>(buffer);
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
+  TF_LITE_ENSURE(context, op_data);
+
+  // Check subgraph index and get subgraph.
+  Subgraph* this_subgraph = reinterpret_cast<Subgraph*>(context->impl_);
+  auto* subgraphs = this_subgraph->GetSubgraphs();
+  TF_LITE_ENSURE_MSG(context,
+                     (op_data->subgraph_index < subgraphs->size()) &&
+                         (op_data->subgraph_index >= 0),
+                     "Index of subgraph to be invoked is invalid.");
+  Subgraph* subgraph = (*subgraphs)[op_data->subgraph_index].get();
+  TF_LITE_ENSURE_MSG(
+      context, subgraph != this_subgraph,
+      "Subgraph to invoke must be different from the invoking graph.");
+  int loop_count = op_data->loop_count;
+  TF_LITE_ENSURE_MSG(context, loop_count >= 0, "Loop count must be positive. ");
+
+  // Check if the inputs and outputs of the CALL node and the subgraph have
+  // proper shapes and types.
+  TF_LITE_ENSURE_OK(context, ValidateAndResizeInputsIfNeeded(
+                                 context, node, subgraph, loop_count));
+  TF_LITE_ENSURE_OK(context, subgraph->AllocateTensors());
+  TF_LITE_ENSURE_OK(
+      context, ValidateAndResizeOutputs(context, node, subgraph, loop_count));
+
+  // Since delegates don't support acceleration of models with dynamic outputs,
+  // this op doesn't support it either.
+  TF_LITE_ENSURE(context, !subgraph->HasDynamicTensors());
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  const OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
+  Subgraph* this_subgraph = reinterpret_cast<Subgraph*>(context->impl_);
+  auto* subgraphs = this_subgraph->GetSubgraphs();
+  Subgraph* subgraph = (*subgraphs)[op_data->subgraph_index].get();
+
+  // The following graph illustrates the current implementation.
+  //
+  // This Subgraph        Subgraph to invoke
+  // +-----------+   (1)   +------------+
+  // |   CALL    |-------->|  SUBGRAPH  |
+  // |   INPUT   |         |   INPUT    |
+  // +-----------+         +------------+
+  //                            |
+  //               (3)          | (2)
+  //                            v
+  // +-----------+        +------------+
+  // |   CALL    |<-------|  SUBGRAPH  |
+  // |   OUTPUT  |        |   OUTPUT   |
+  // +-----------+        +------------+
+  // For every ith loop iteration:
+  // (1) Copy the ith input of CALL op to the inputs of subgraph.
+  // (2) Invoke subgraph.
+  // (3) Copy the outputs of subgraph to the ith output of CALL op.
+  //
+  // Requires the subgraph to have a batch size of 1.
+  // Requires the CALL node's inputs and outputs to have a batch size equal to
+  // `loop_count`.
+  //
+  //
+  // TODO(b/120234921): Optimize and avoid copying tensors between subgraphs.
+
+  for (int loop_index = 0; loop_index < op_data->loop_count; loop_index++) {
+    // Copy inputs needed for this iteration.
+    TF_LITE_ENSURE_OK(context,
+                      CopyInputTensorsData(context, node, subgraph, loop_index,
+                                           op_data->loop_count));
+    // Invoke subgraph for this iteration.
+    TF_LITE_ENSURE_OK(context, subgraph->Invoke());
+
+    for (int tensor_index : subgraph->outputs()) {
+      subgraph->EnsureTensorDataIsReadable(tensor_index);
+    }
+
+    TF_LITE_ENSURE_OK(context,
+                      CopyOutputTensorsData(context, subgraph, node, loop_index,
+                                            op_data->loop_count));
+  }
+  return kTfLiteOk;
+}
+
+}  // namespace call_kernel
+
+TfLiteRegistration* Register_CALL() {
+  static TfLiteRegistration r = {call_kernel::Init, call_kernel::Free,
+                                 call_kernel::Prepare, call_kernel::Eval};
+  return &r;
+}
+
+}  // namespace ops
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/call_register.h

@@ -0,0 +1,37 @@
+/* Copyright 2020 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_CALL_REGISTER_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_CALL_REGISTER_H_
+#include "tensorflow/lite/core/c/common.h"
+namespace tflite {
+namespace acceleration {
+namespace ops {
+// CALL op can be used to invoke a subgraph a given number of times.
+TfLiteRegistration* Register_CALL();
+
+typedef struct {
+  // Index of the subgraph that needs to be invoked.
+  // Subgraph should have batch size 1.
+  int subgraph_index;
+  // The number of times the CALL op should call the subgraph.
+  // The inputs to the call op are expected to have this value as their batch
+  // size.
+  int loop_count;
+} TfLiteCallParams;
+
+}  // namespace ops
+}  // namespace acceleration
+}  // namespace tflite
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_CALL_REGISTER_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/call_test.cc

@@ -0,0 +1,420 @@
+/* Copyright 2020 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 <cstddef>
+#include <memory>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/builtin_ops.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/core/kernels/builtin_op_kernels.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/call_register.h"
+#include "tensorflow/lite/interpreter_test_util.h"
+#include "tensorflow/lite/kernels/subgraph_test_util.h"
+#include "tensorflow/lite/testing/util.h"
+
+namespace tflite {
+
+namespace {
+
+class CallTest : public subgraph_test_util::ControlFlowOpTest {
+ public:
+  CallTest() { interpreter_ = std::make_unique<Interpreter>(&error_reporter_); }
+  ~CallTest() override = default;
+  void SetupTensor(Subgraph* subgraph, int tensor_index, TfLiteType type) {
+    ASSERT_EQ(subgraph->SetTensorParametersReadWrite(tensor_index, type, "", 0,
+                                                     nullptr, {}, false),
+              kTfLiteOk);
+  }
+  void BuildCallSubgraph(Subgraph* subgraph, std::vector<uint8_t> params_buffer,
+                         std::vector<int> inputs, std::vector<int> outputs,
+                         int expected_node_index, bool single_node_subgraph) {
+    if (single_node_subgraph) {
+      int first_new_tensor_index;
+      ASSERT_EQ(subgraph->AddTensors(inputs.size() + outputs.size(),
+                                     &first_new_tensor_index),
+                kTfLiteOk);
+      ASSERT_EQ(first_new_tensor_index, 0);
+      ASSERT_EQ(subgraph->SetInputs(inputs), kTfLiteOk);
+      ASSERT_EQ(subgraph->SetOutputs(outputs), kTfLiteOk);
+    }
+    for (const int& idx : inputs) {
+      SetupTensor(subgraph, idx, kTfLiteInt32);
+    }
+    for (const int& idx : outputs) {
+      SetupTensor(subgraph, idx, kTfLiteInt32);
+    }
+    int node_index;
+    subgraph->AddNodeWithParameters(
+        inputs, outputs, {},
+        reinterpret_cast<const char*>(params_buffer.data()),
+        params_buffer.size(), nullptr, acceleration::ops::Register_CALL(),
+        &node_index);
+    ASSERT_EQ(node_index, expected_node_index);
+  }
+  void BuildCallSubgraph(Subgraph* subgraph, int index, int loop_count,
+                         std::vector<int> inputs, std::vector<int> outputs,
+                         int expected_node_index = 0,
+                         bool single_node_subgraph = true) {
+    flexbuffers::Builder fbb;
+    fbb.Map([&] {
+      fbb.Int("subgraph_index", index);
+      fbb.Int("loop_count", loop_count);
+    });
+    fbb.Finish();
+    BuildCallSubgraph(subgraph, fbb.GetBuffer(), inputs, outputs,
+                      expected_node_index, single_node_subgraph);
+  }
+
+  void BuildGraphWithMultipleOutputs(Subgraph* subgraph) {
+    const int kInput1 = 0;
+    const int kInput2 = 1;
+    const int kMulOutput = 2;
+    const int kAddOutput = 3;
+    const int kTensorCount = 4;
+    // kInput1(0) --> +---+
+    //                |MUL| --> kOutput(2)
+    // kInput2(1) --> +---+
+    //
+    // kInput1(0) --> +---+
+    //                |ADD| --> kOutput(3)
+    // kInput2(1) --> +---+
+
+    int first_new_tensor_index;
+    ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
+              kTfLiteOk);
+    ASSERT_EQ(first_new_tensor_index, 0);
+    ASSERT_EQ(subgraph->SetInputs({kInput1, kInput2}), kTfLiteOk);
+    ASSERT_EQ(subgraph->SetOutputs({kMulOutput, kAddOutput}), kTfLiteOk);
+
+    SetupTensor(subgraph, kInput1, kTfLiteInt32);
+    SetupTensor(subgraph, kInput2, kTfLiteInt32);
+    SetupTensor(subgraph, kMulOutput, kTfLiteInt32);
+    SetupTensor(subgraph, kAddOutput, kTfLiteInt32);
+    TfLiteMulParams* params_mul =
+        reinterpret_cast<TfLiteMulParams*>(malloc(sizeof(TfLiteMulParams)));
+    params_mul->activation = kTfLiteActNone;
+    int node_index;
+    subgraph->AddNodeWithParameters(
+        {kInput1, kInput2}, {kMulOutput}, {}, nullptr, 0, params_mul,
+        ::tflite::ops::builtin::Register_MUL(), &node_index);
+    TfLiteAddParams* params_add =
+        reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
+    params_add->activation = kTfLiteActNone;
+    subgraph->AddNodeWithParameters(
+        {kInput1, kInput2}, {kAddOutput}, {}, nullptr, 0, params_add,
+        ::tflite::ops::builtin::Register_ADD(), &node_index);
+  }
+  void BuildMultiNodeGraph(Subgraph* this_subgraph) {
+    // kIn1(0)----------------
+    //                       |
+    //                       |        +----+
+    //            +---+      -------->|    |              +---+
+    // kIn2(1)--> |PAD|-->kOut1(4)--->|CALL|-->kOut2(5)-->|MUL|-->kOut3(6)
+    // kIn3(2)--> |   |               |    |              |   |
+    //            +---+               +----+         ---->|   |
+    //                                               |    +---+
+    //                                               |
+    //                                               |
+    // kIn4(3)----------------------------------------
+    const int kInput1 = 0, kInput2 = 1, kInput3 = 2, kInput4 = 3;
+    const int kOutput1 = 4, kOutput2 = 5, kOutput3 = 6;
+    const int kTensorCount = 7;
+    int first_new_tensor_index;
+    ASSERT_EQ(this_subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
+              kTfLiteOk);
+    ASSERT_EQ(first_new_tensor_index, 0);
+    std::vector<int> inputs = {kInput1, kInput2, kInput3, kInput4};
+    std::vector<int> outputs = {kOutput3};
+    ASSERT_EQ(this_subgraph->SetInputs(inputs), kTfLiteOk);
+    ASSERT_EQ(this_subgraph->SetOutputs({kOutput3}), kTfLiteOk);
+    for (int idx = 0; idx < kTensorCount; ++idx) {
+      SetupTensor(this_subgraph, idx, kTfLiteInt32);
+    }
+    int expected_node_index = 0, node_index;
+    // Node 1: Pad op.
+    auto* pad_reg = ops::builtin::Register_PAD();
+    pad_reg->builtin_code = kTfLiteBuiltinPad;
+    this_subgraph->AddNodeWithParameters(
+        {kInput2, kInput3}, {kOutput1}, {}, nullptr, 0,
+        reinterpret_cast<TfLitePadParams*>(malloc(sizeof(TfLitePadParams))),
+        pad_reg, &node_index);
+    ASSERT_EQ(node_index, expected_node_index++);
+    // Node 2: Call op, calls subgraph that contains Add op.
+    AddSubgraphs(1);
+    const int kLoopCount = 1;
+    const int kSubgraphIndex = 1;
+    builder_->BuildAddSubgraph(interpreter_->subgraph(1));
+    CallTest::BuildCallSubgraph(this_subgraph, kSubgraphIndex, kLoopCount,
+                                {kInput1, kOutput1}, {kOutput2},
+                                expected_node_index++, false);
+    // Node 3: Mul op.
+    TfLiteMulParams* mul_params =
+        reinterpret_cast<TfLiteMulParams*>(malloc(sizeof(TfLiteMulParams)));
+    mul_params->activation = kTfLiteActNone;
+    auto* mul_reg = ops::builtin::Register_MUL();
+    mul_reg->builtin_code = kTfLiteBuiltinMul;
+    this_subgraph->AddNodeWithParameters({kInput4, kOutput2}, {kOutput3}, {},
+                                         nullptr, 0, mul_params, mul_reg,
+                                         &node_index);
+    ASSERT_EQ(node_index, expected_node_index++);
+  }
+  TestErrorReporter error_reporter_;
+};
+
+/** Tests the happy path for `call` op. **/
+TEST_F(CallTest, SubgraphMultipleInputsSingleOutput) {
+  std::vector<std::vector<int>> test_shapes = {
+      {3, 2}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}};
+  // Will loop over and will be fed to the subgraph as {1,2}, {1,3}, {1,1,3},
+  // {1,3,1,2}.
+  for (size_t i = 0; i < test_shapes.size(); ++i) {
+    interpreter_ = std::make_unique<Interpreter>();
+    AddSubgraphs(1);
+    int loop_count = test_shapes[i][0];
+    builder_->BuildMulSubgraph(interpreter_->subgraph(1));
+    CallTest::BuildCallSubgraph(&interpreter_->primary_subgraph(), 1,
+                                loop_count, {0, 1}, {2});
+    interpreter_->ResizeInputTensor(interpreter_->inputs()[0], test_shapes[i]);
+    interpreter_->ResizeInputTensor(interpreter_->inputs()[1], test_shapes[i]);
+    ASSERT_EQ(interpreter_->subgraph(1)->AllocateTensors(), kTfLiteOk);
+    ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
+
+    subgraph_test_util::FillIntTensor(
+        interpreter_->tensor(interpreter_->inputs()[0]), {-1, 2, -3, 4, -5, 6});
+    subgraph_test_util::FillIntTensor(
+        interpreter_->tensor(interpreter_->inputs()[1]), {-1, 2, -3, 4, -5, 6});
+    ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
+
+    TfLiteTensor* output = interpreter_->tensor(interpreter_->outputs()[0]);
+    subgraph_test_util::CheckIntTensor(output, test_shapes[i],
+                                       {1, 4, 9, 16, 25, 36});
+  }
+}
+
+TEST_F(CallTest, ShouldBeANoOpWhenLoopCountIsZero) {
+  AddSubgraphs(1);
+  int loop_count = 0;
+  builder_->BuildMulSubgraph(interpreter_->subgraph(1));
+  CallTest::BuildCallSubgraph(&interpreter_->primary_subgraph(), 1, loop_count,
+                              {0, 1}, {2});
+  interpreter_->ResizeInputTensor(interpreter_->inputs()[0], {0, 3});
+  interpreter_->ResizeInputTensor(interpreter_->inputs()[1], {0, 3});
+  ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
+  ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
+  TfLiteTensor* output = interpreter_->tensor(interpreter_->outputs()[0]);
+  subgraph_test_util::CheckIntTensor(output, {0, 3}, {});
+}
+
+TEST_F(CallTest, SubgraphWithFixedInputShapes) {
+  AddSubgraphs(1);
+  const int kLoopCount = 2;
+  const int kBatchSizeSubgraph = 1;
+  const int kFixedInputLen = 3;
+  const std::vector<int> kCallOpInputShape = {kLoopCount, kFixedInputLen};
+  const std::vector<int> kSubgraphInputShape = {kBatchSizeSubgraph,
+                                                kFixedInputLen};
+
+  Subgraph* subgraph = interpreter_->subgraph(1);
+  builder_->BuildMulSubgraph(subgraph);
+  CallTest::BuildCallSubgraph(&interpreter_->primary_subgraph(), 1, kLoopCount,
+                              {0, 1}, {2});
+  interpreter_->ResizeInputTensor(interpreter_->inputs()[0], kCallOpInputShape);
+  interpreter_->ResizeInputTensor(interpreter_->inputs()[1], kCallOpInputShape);
+
+  subgraph->ResizeInputTensor(subgraph->inputs()[0], kSubgraphInputShape);
+  subgraph->ResizeInputTensor(subgraph->inputs()[1], kSubgraphInputShape);
+
+  ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
+  subgraph_test_util::FillIntTensor(
+      interpreter_->tensor(interpreter_->inputs()[0]), {-1, 2, -3, 4, -5, 6});
+  subgraph_test_util::FillIntTensor(
+      interpreter_->tensor(interpreter_->inputs()[1]), {-1, 2, -3, 4, -5, 6});
+  ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
+
+  TfLiteTensor* output = interpreter_->tensor(interpreter_->outputs()[0]);
+  subgraph_test_util::CheckIntTensor(output, kCallOpInputShape,
+                                     {1, 4, 9, 16, 25, 36});
+}
+
+TEST_F(CallTest, SubgraphWithMultipleInputsAndOutputs) {
+  std::vector<std::vector<int>> test_shapes = {
+      {3, 2, 1}, {1, 2, 3}, {2, 1, 3}, {2, 3, 1, 1}, {2, 3}};
+  for (size_t i = 0; i < test_shapes.size(); ++i) {
+    interpreter_ = std::make_unique<Interpreter>();
+    AddSubgraphs(1);
+    int loop_count = test_shapes[i][0];
+    CallTest::BuildGraphWithMultipleOutputs(interpreter_->subgraph(1));
+    CallTest::BuildCallSubgraph(&interpreter_->primary_subgraph(), 1,
+                                loop_count, {0, 1}, {2, 3});
+    interpreter_->ResizeInputTensor(interpreter_->inputs()[0], test_shapes[i]);
+    interpreter_->ResizeInputTensor(interpreter_->inputs()[1], test_shapes[i]);
+    ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
+
+    subgraph_test_util::FillIntTensor(
+        interpreter_->tensor(interpreter_->inputs()[0]), {-1, 2, -3, 4, -5, 6});
+    subgraph_test_util::FillIntTensor(
+        interpreter_->tensor(interpreter_->inputs()[1]), {-1, 2, -3, 4, -5, 6});
+    ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
+
+    TfLiteTensor* output_mul = interpreter_->tensor(interpreter_->outputs()[0]);
+    subgraph_test_util::CheckIntTensor(output_mul, test_shapes[i],
+                                       {1, 4, 9, 16, 25, 36});
+    TfLiteTensor* output_add = interpreter_->tensor(interpreter_->outputs()[1]);
+    subgraph_test_util::CheckIntTensor(output_add, test_shapes[i],
+                                       {-2, 4, -6, 8, -10, 12});
+  }
+}
+
+TEST_F(CallTest, ShouldHandleInvalidParamsAndSetToDefault) {
+  flexbuffers::Builder fbb;
+  fbb.Vector([&]() {
+    fbb.String("hi");
+    fbb.String("hello");
+  });
+  fbb.Finish();
+  AddSubgraphs(1);
+
+  CallTest::BuildCallSubgraph(&interpreter_->primary_subgraph(),
+                              fbb.GetBuffer(), {0}, {1}, 0, true);
+  const int kNodeIndex = 0;
+  const TfLiteNode* call_node = &interpreter_->primary_subgraph()
+                                     .nodes_and_registration()[kNodeIndex]
+                                     .first;
+  tflite::acceleration::ops::TfLiteCallParams* op_data =
+      reinterpret_cast<tflite::acceleration::ops::TfLiteCallParams*>(
+          call_node->user_data);
+
+  EXPECT_EQ(op_data->subgraph_index, 0);
+  EXPECT_EQ(op_data->loop_count, 0);
+}
+TEST_F(CallTest, MultiNodeGraph) {
+  CallTest::BuildMultiNodeGraph(&interpreter_->primary_subgraph());
+  interpreter_->ResizeInputTensor(interpreter_->inputs()[0], {1, 4, 4, 1});
+  interpreter_->ResizeInputTensor(interpreter_->inputs()[1], {1, 2, 2, 1});
+  interpreter_->ResizeInputTensor(interpreter_->inputs()[2], {4, 2});
+  interpreter_->ResizeInputTensor(interpreter_->inputs()[3], {1, 4, 4, 1});
+
+  ASSERT_EQ(interpreter_->subgraph(1)->AllocateTensors(), kTfLiteOk);
+  ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
+
+  subgraph_test_util::FillIntTensor(
+      interpreter_->tensor(interpreter_->inputs()[0]), std::vector<int>(16, 1));
+  subgraph_test_util::FillIntTensor(
+      interpreter_->tensor(interpreter_->inputs()[1]), {1, 2, 3, 4});
+  subgraph_test_util::FillIntTensor(
+      interpreter_->tensor(interpreter_->inputs()[2]),
+      {0, 0, 1, 1, 1, 1, 0, 0});
+  subgraph_test_util::FillIntTensor(
+      interpreter_->tensor(interpreter_->inputs()[3]), std::vector<int>(16, 2));
+
+  ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
+  TfLiteTensor* output = interpreter_->tensor(interpreter_->outputs()[0]);
+  subgraph_test_util::CheckIntTensor(
+      output, {1, 4, 4, 1}, {2, 2, 2, 2, 2, 4, 6, 2, 2, 8, 10, 2, 2, 2, 2, 2});
+}
+
+// Note: For the tests below the error messages returned by the error reporter
+// are of the following format:
+// "<filename>:<line number> <error message>. Node <number name> failed to
+// prepare.\n"
+// It's sufficient to test whether the string returned by error reporter
+// contains the expected error message.
+TEST_F(CallTest, ShouldFailWith0DInputs) {
+  AddSubgraphs(1);
+  int loop_count = 5;
+  builder_->BuildMulSubgraph(interpreter_->subgraph(1));
+  interpreter_->subgraph(1)->ResizeInputTensor(0, {});
+  interpreter_->subgraph(1)->ResizeInputTensor(1, {});
+  CallTest::BuildCallSubgraph(&interpreter_->primary_subgraph(), 1, loop_count,
+                              {0, 1}, {2});
+
+  ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteError);
+
+  EXPECT_THAT(
+      error_reporter_.error_messages(),
+      testing::HasSubstr(
+          "Dimensions of all of call node's inputs should be non-zero."));
+}
+
+TEST_F(CallTest, ShouldFailWhenLoopCountDoesNotMatchBatchSize) {
+  AddSubgraphs(1);
+  int loop_count = 7;
+  builder_->BuildMulSubgraph(interpreter_->subgraph(1));
+  CallTest::BuildCallSubgraph(&interpreter_->primary_subgraph(), 1, loop_count,
+                              {0, 1}, {2});
+  interpreter_->ResizeInputTensor(interpreter_->inputs()[0], {5, 3});
+  interpreter_->ResizeInputTensor(interpreter_->inputs()[1], {5, 3});
+
+  ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteError);
+  EXPECT_THAT(
+      error_reporter_.error_messages(),
+      testing::HasSubstr("node_input->dims->data[0] != loop_count (5 != 7)"));
+}
+
+TEST_F(CallTest, ShouldFailForSubgraphWithIncompatibleInputShapes) {
+  AddSubgraphs(1);
+  const int kLoopCount = 5;
+  const int kBatchSizeSubgraph = 1;
+  std::vector<int> call_op_input = {kLoopCount, 3};
+  std::vector<int> subgraph_input = {kBatchSizeSubgraph, 7};
+  Subgraph* subgraph = interpreter_->subgraph(1);
+  builder_->BuildMulSubgraph(subgraph);
+  CallTest::BuildCallSubgraph(&interpreter_->primary_subgraph(), 1, kLoopCount,
+                              {0, 1}, {2});
+  interpreter_->ResizeInputTensor(interpreter_->inputs()[0], call_op_input);
+  interpreter_->ResizeInputTensor(interpreter_->inputs()[1], call_op_input);
+  subgraph->ResizeInputTensor(subgraph->inputs()[0], subgraph_input);
+  subgraph->ResizeInputTensor(subgraph->inputs()[1], subgraph_input);
+
+  ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteError);
+
+  EXPECT_THAT(
+      error_reporter_.error_messages(),
+      testing::HasSubstr("All dimensions except the batch size should match "
+                         "for call node and the subgraph to invoke"));
+}
+
+TEST_F(CallTest, ShouldFailWhenSubgraphIndexMatchesInvokedSubgraph) {
+  const int kPrimarySubgraphIndex = 0;
+  CallTest::BuildCallSubgraph(&interpreter_->primary_subgraph(),
+                              kPrimarySubgraphIndex, 1, {0}, {1});
+
+  ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteError);
+
+  EXPECT_THAT(
+      error_reporter_.error_messages(),
+      testing::HasSubstr(
+          "Subgraph to invoke must be different from the invoking graph."));
+}
+
+TEST_F(CallTest, ShouldFailWithNegativeLoopCount) {
+  AddSubgraphs(1);
+  CallTest::BuildCallSubgraph(&interpreter_->primary_subgraph(), 1, -1, {0},
+                              {1});
+
+  ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteError);
+
+  EXPECT_THAT(error_reporter_.error_messages(),
+              testing::HasSubstr("Loop count must be positive."));
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/constants.h

@@ -0,0 +1,32 @@
+/* Copyright 2022 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_CONSTANTS_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_CONSTANTS_H_
+
+namespace tflite {
+namespace acceleration {
+
+// Model modification.
+inline constexpr int kModelSchemaVersion = 3;
+inline constexpr char kValidationGraphName[] = "VALIDATION:main";
+
+// Process Execution.
+// Number of bytes used to store process id.
+inline constexpr int kPidBufferLength = 20;
+
+}  // namespace acceleration
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_CONSTANTS_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/copy_associated_files.py

@@ -0,0 +1,62 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""TFLite model associated files copier.
+
+A TFLite model with metadata may have 'associated files': extra data stored as a
+zip file appended to the model. This script copies such files from one model to
+another. This is used as part of constructing validation models.
+
+See https://www.tensorflow.org/lite/convert/metadata for description of models
+with metadata.
+
+If there are no associated files in the provided file, the model is output
+as-is.
+"""
+
+import argparse
+import sys
+import zipfile
+
+parser = argparse.ArgumentParser(
+    description='Script to generate a metrics model for mobilenet v1.')
+parser.add_argument('model', help='Input model filepath')
+parser.add_argument(
+    'copy_associated_files_from',
+    help='Model with potential associated files filepath')
+parser.add_argument('output', help='Output filepath')
+
+
+def main(model_path, associated_files_path, output_path):
+  with open(model_path, 'rb') as input_file:
+    with open(output_path, 'wb') as output_file:
+      output_file.write(input_file.read())
+  if zipfile.is_zipfile(associated_files_path):
+    zip_src = zipfile.ZipFile(associated_files_path, 'r')
+    zip_tgt = zipfile.ZipFile(output_path, 'a')
+    for info in zip_src.infolist():
+      zip_tgt.writestr(info, zip_src.read(info))
+
+
+if __name__ == '__main__':
+  flags, unparsed = parser.parse_known_args()
+
+  if unparsed:
+    parser.print_usage()
+    sys.stderr.write('\nGot the following unparsed args, %r please fix.\n' %
+                     unparsed)
+    exit(1)
+  else:
+    main(flags.model, flags.copy_associated_files_from, flags.output)
+    exit(0)

tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg.cc

@@ -0,0 +1,185 @@
+/* Copyright 2021 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 <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg_register.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg_status.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/string_util.h"
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  if (!buffer) {
+    return nullptr;
+  }
+#define RET_ENSURE(context, condition)                                  \
+  do {                                                                  \
+    if (!(condition)) {                                                 \
+      TF_LITE_KERNEL_LOG((context), "%s:%d %s was not true.", __FILE__, \
+                         __LINE__, #condition);                         \
+      return nullptr;                                                   \
+    }                                                                   \
+  } while (0)
+  const uint8_t* buffer_t = reinterpret_cast<const uint8_t*>(buffer);
+  const flexbuffers::Map m = flexbuffers::GetRoot(buffer_t, length).AsMap();
+  RET_ENSURE(context, m["height"].IsInt());
+  RET_ENSURE(context, m["width"].IsInt());
+  RET_ENSURE(context, m["num_images"].IsInt());
+  RET_ENSURE(context, m["channels"].IsInt());
+  OpData* op_data = new OpData();
+  op_data->height = m["height"].AsInt32();
+  op_data->width = m["width"].AsInt32();
+  op_data->num_images = m["num_images"].AsInt32();
+  op_data->channels = m["channels"].AsInt32();
+  return op_data;
+#undef RET_ENSURE
+}
+
+void Free(TfLiteContext* context, void* buffer) {
+  delete reinterpret_cast<OpData*>(buffer);
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
+  TF_LITE_ENSURE(context, op_data);
+  TF_LITE_ENSURE(context, op_data->height > 0);
+  TF_LITE_ENSURE(context, op_data->width > 0);
+  TF_LITE_ENSURE(context, op_data->num_images > 0);
+  // TODO(b/172544567): Support grayscale images.
+  TF_LITE_ENSURE(context, op_data->channels == 3 || op_data->channels == 4);
+
+  TF_LITE_ENSURE_EQ(context, node->inputs->size, 1);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);
+
+  const TfLiteTensor* input_buffer;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, /*index=*/0, &input_buffer));
+
+  TfLiteTensor* output_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, /*index=*/0, &output_tensor));
+
+  TF_LITE_ENSURE_TYPES_EQ(context, input_buffer->type, kTfLiteString);
+  TF_LITE_ENSURE_TYPES_EQ(context, output_tensor->type, kTfLiteUInt8);
+
+  TF_LITE_ENSURE_EQ(context, NumDimensions(input_buffer), 1);
+  TF_LITE_ENSURE_EQ(context, input_buffer->dims->data[0], op_data->num_images);
+
+  // Resize output.
+  // Output shape is determined as {num_images, height, width, channels}.
+  TfLiteIntArray* new_dims = TfLiteIntArrayCreate(4);
+  new_dims->data[0] = op_data->num_images;
+  new_dims->data[1] = op_data->height;
+  new_dims->data[2] = op_data->width;
+  new_dims->data[3] = op_data->channels;
+  output_tensor->type = kTfLiteUInt8;
+  TF_LITE_ENSURE_OK(context,
+                    context->ResizeTensor(context, output_tensor, new_dims));
+  return kTfLiteOk;
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  // Decodes a batch of JPEG images.
+
+  OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
+
+  const TfLiteTensor* input_buffer;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, /*index=*/0, &input_buffer));
+  TF_LITE_ENSURE(context, input_buffer);
+  TF_LITE_ENSURE(context, input_buffer->data.raw);
+  const int channels = op_data->channels;
+  // TODO(b/172544567): Support grayscale images.
+  const int decode_channels = 3;
+  TfLiteTensor* output_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, /*index=*/0, &output_tensor));
+  // kTfliteUInt8 corresponds to unsigned char as shown in
+  // "tensorflow/lite/portable_type_to_tflitetype.h".
+  unsigned char* output_arr = GetTensorData<unsigned char>(output_tensor);
+  Status decoder_status;
+  std::unique_ptr<LibjpegDecoder> decoder =
+      LibjpegDecoder::Create(decoder_status);
+  if (decoder_status.code != kTfLiteOk) {
+    TF_LITE_KERNEL_LOG(context, "%s", decoder_status.error_message.c_str());
+    return kTfLiteError;
+  }
+
+  const int kDecodedImageSize =
+      op_data->width * op_data->height * decode_channels;
+  const int kOutputImageSize = op_data->width * op_data->height * channels;
+
+  int output_array_offset = 0;
+  for (int img = 0; img < op_data->num_images; ++img) {
+    tflite::StringRef inputref =
+        tflite::GetString(input_buffer, /*string_index=*/img);
+    unsigned char* decoded = output_arr + output_array_offset;
+
+    Status decode_status = decoder->DecodeImage(
+        inputref, {op_data->height, op_data->width, decode_channels}, decoded,
+        kDecodedImageSize);
+
+    if (channels == 4) {
+      // Reorganize the decoded buffer from 3 channels to 4 channels.
+      size_t height = op_data->height;
+      size_t src_offset = kDecodedImageSize;
+      size_t dst_offset = kOutputImageSize;
+      while (height--) {
+        size_t width = op_data->width;
+        while (width--) {
+          src_offset -= decode_channels;
+          dst_offset -= channels;
+          std::copy_n(decoded + src_offset, decode_channels,
+                      decoded + dst_offset);
+          // Add an alpha channel value of 255 (fully opaque) to the
+          // current pixel if the target output channels is provided as 4. This
+          // is a workaround to allow jpeg decoder to work with 4 channel input
+          // models.
+          decoded[dst_offset + 3] = 255;
+        }
+      }
+    }
+
+    output_array_offset += kOutputImageSize;
+
+    if (decode_status.code != kTfLiteOk) {
+      TF_LITE_KERNEL_LOG(context, "%s", decode_status.error_message.c_str());
+      return kTfLiteError;
+    }
+  }
+  return kTfLiteOk;
+}
+
+TfLiteRegistration* Register_DECODE_JPEG() {
+  static TfLiteRegistration r = {
+      decode_jpeg_kernel::Init, decode_jpeg_kernel::Free,
+      decode_jpeg_kernel::Prepare, decode_jpeg_kernel::Eval};
+  return &r;
+}
+
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg_register.h

@@ -0,0 +1,54 @@
+/* Copyright 2021 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_DECODE_JPEG_REGISTER_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_DECODE_JPEG_REGISTER_H_
+
+#include <cstdint>
+
+#include "tensorflow/lite/core/c/common.h"
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+
+// DECODE_JPEG can be used to decode a batch of JPEG images on Android.
+// TODO(b/172544567): Support iOS.
+// TODO(b/172544567): Support greyscale images.
+// Expects single 1D input of the shape {num_images} and type string.
+// Single output containing the decoded images with shape {num_images, height,
+// width, channels}. All input images are required to have 3 channels. The
+// decoded images can have 3 or 4 channels depending on the shape of the
+// target model input. This op will add an alpha channel value of 255 (fully
+// opaque) if the target model accepts input images with 4 channels. This op
+// will eventually be included in mainline Tflite as a built-in/custom op once
+// it supports both Android and iOS.
+TfLiteRegistration* Register_DECODE_JPEG();
+
+struct OpData {
+  // Number of images to decode.
+  int32_t num_images;
+  // All images should have the same height and width.
+  // Height of images after decoding.
+  int32_t height;
+  // Width of images after decoding.
+  int32_t width;
+  // Number of channels to be decoded. Accepted values are 3 (RGB) and 4 (RGBA).
+  int32_t channels;
+};
+
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_DECODE_JPEG_REGISTER_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg_status.h

@@ -0,0 +1,35 @@
+/* Copyright 2021 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_DECODE_JPEG_STATUS_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_DECODE_JPEG_STATUS_H_
+
+#include <string>
+
+#include "tensorflow/lite/core/c/c_api_types.h"
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+
+struct Status {
+  TfLiteStatus code = kTfLiteOk;
+  std::string error_message;
+};
+
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_DECODE_JPEG_STATUS_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg_test.cc

@@ -0,0 +1,107 @@
+/* Copyright 2021 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 <cstdint>
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg_register.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/embedded_chessboard_jpeg.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/embedded_test_card_jpeg.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder_test_helper.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+
+namespace {
+
+using testing::ElementsAre;
+
+const int kHeight = 300, kWidth = 250, kChannels = 3;
+const int kDecodedSize = kHeight * kWidth * kChannels;
+
+class DecodeJPEGOpModel : public SingleOpModel {
+ public:
+  DecodeJPEGOpModel(const TensorData& input, const TensorData& output,
+                    int num_images, int height, int width, int channels) {
+    input_id_ = AddInput(input);
+    output_id_ = AddOutput(output);
+    flexbuffers::Builder fbb;
+    fbb.Map([&] {
+      fbb.Int("num_images", num_images);
+      fbb.Int("height", height);
+      fbb.Int("width", width);
+      fbb.Int("channels", channels);
+    });
+    fbb.Finish();
+    SetCustomOp("DECODE_JPEG", fbb.GetBuffer(),
+                tflite::acceleration::decode_jpeg_kernel::Register_DECODE_JPEG);
+    BuildInterpreter({GetShape(input_id_)});
+  }
+
+  int input_buffer_id() { return input_id_; }
+  int output_id() { return output_id_; }
+  std::vector<uint8_t> GetOutput() {
+    return ExtractVector<uint8_t>(output_id_);
+  }
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_id_); }
+
+ protected:
+  int input_id_;
+  int shapes_id_;
+  int output_id_;
+};
+
+// TODO(b/172544567): Add more tests to verify that invalid shapes, types and
+// params are handled gracefully by the op.
+
+TEST(DecodeJpegTest, TestMultipleJPEGImages) {
+  // Set up model and populate the input.
+  std::string chessboard_image(
+      reinterpret_cast<const char*>(g_tflite_acceleration_chessboard_jpeg),
+      g_tflite_acceleration_chessboard_jpeg_len);
+  std::string test_card_image(
+      reinterpret_cast<const char*>(g_tflite_acceleration_test_card_jpeg),
+      g_tflite_acceleration_test_card_jpeg_len);
+  const int kNumImages = 2;
+  DecodeJPEGOpModel model({TensorType_STRING, {kNumImages}},
+                          {TensorType_UINT8, {}}, kNumImages, kHeight, kWidth,
+                          kChannels);
+  model.PopulateStringTensor(model.input_buffer_id(),
+                             {chessboard_image, test_card_image});
+
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  // Check output values and shape.
+  ASSERT_THAT(model.GetOutputShape(),
+              ElementsAre(kNumImages, kHeight, kWidth, kChannels));
+  std::vector<uint8_t> output_flattened = model.GetOutput();
+  std::vector<uint8_t> img1(output_flattened.begin(),
+                            output_flattened.begin() + kDecodedSize);
+  EXPECT_THAT(img1, HasChessboardPatternWithTolerance(12));
+  std::vector<uint8_t> img2(output_flattened.begin() + kDecodedSize,
+                            output_flattened.end());
+  EXPECT_THAT(img2, HasRainbowPatternWithTolerance(5));
+}
+
+}  // namespace
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/fb_storage.cc

@@ -0,0 +1,157 @@
+/* Copyright 2021 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/fb_storage.h"
+
+#include <fcntl.h>
+#include <string.h>
+#ifndef _WIN32
+#include <sys/file.h>
+#include <unistd.h>
+#endif
+
+#include <fstream>
+#include <sstream>
+#include <string>
+
+#include "absl/base/attributes.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/lite/core/c/c_api_types.h"
+
+// We only really care about Android, but we want the code to be portable for
+// ease of testing. See also discussion in cl/224174491.
+#ifndef TEMP_FAILURE_RETRY
+#ifdef __ANDROID__
+#error "TEMP_FAILURE_RETRY not set although on Android"
+#else  // ! defined(__ANDROID__)
+#define TEMP_FAILURE_RETRY(exp) exp
+#endif  // defined(__ANDROID__)
+#endif  // defined(TEMP_FAILURE_RETRY)
+
+namespace tflite {
+namespace acceleration {
+FileStorage::FileStorage(absl::string_view path, ErrorReporter* error_reporter)
+    : path_(path), error_reporter_(error_reporter) {}
+
+MinibenchmarkStatus FileStorage::ReadFileIntoBuffer() {
+#ifndef _WIN32
+  buffer_.clear();
+  // O_CLOEXEC is needed for correctness, as another thread may call
+  // popen() and the callee inherit the lock if it's not O_CLOEXEC.
+  int fd = TEMP_FAILURE_RETRY(open(path_.c_str(), O_RDONLY | O_CLOEXEC, 0600));
+  int open_error_no = errno;
+  if (fd < 0) {
+    // Try to create if it doesn't exist.
+    int fd = TEMP_FAILURE_RETRY(
+        open(path_.c_str(), O_WRONLY | O_APPEND | O_CREAT | O_CLOEXEC, 0600));
+    if (fd >= 0) {
+      // Successfully created file, all good.
+      close(fd);
+      return kMinibenchmarkSuccess;
+    }
+    int create_error_no = errno;
+    TF_LITE_REPORT_ERROR(
+        error_reporter_,
+        "Could not open %s for reading: %s, creating failed as well: %s",
+        path_.c_str(), std::strerror(open_error_no),
+        std::strerror(create_error_no));
+    return kMinibenchmarkCantCreateStorageFile;
+  }
+  int lock_status = flock(fd, LOCK_EX);
+  int lock_error_no = errno;
+  if (lock_status < 0) {
+    close(fd);
+    TF_LITE_REPORT_ERROR(error_reporter_, "Could not flock %s: %s",
+                         path_.c_str(), std::strerror(lock_error_no));
+    return kMinibenchmarkFlockingStorageFileFailed;
+  }
+  char buffer[512];
+  while (true) {
+    int bytes_read = TEMP_FAILURE_RETRY(read(fd, buffer, 512));
+    int read_error_no = errno;
+    if (bytes_read == 0) {
+      // EOF
+      close(fd);
+      return kMinibenchmarkSuccess;
+    } else if (bytes_read < 0) {
+      close(fd);
+      TF_LITE_REPORT_ERROR(error_reporter_, "Error reading %s: %s",
+                           path_.c_str(), std::strerror(read_error_no));
+      return kMinibenchmarkErrorReadingStorageFile;
+    } else {
+      buffer_.append(buffer, bytes_read);
+    }
+  }
+#else  // _WIN32
+  return kMinibenchmarkUnsupportedPlatform;
+#endif
+}
+
+MinibenchmarkStatus FileStorage::AppendDataToFile(absl::string_view data) {
+#ifndef _WIN32
+  // We use a file descriptor (as opposed to FILE* or C++ streams) for writing
+  // because we want to be able to use fsync and flock.
+  // O_CLOEXEC is needed for correctness, as another thread may call
+  // popen() and the callee inherit the lock if it's not O_CLOEXEC.
+  int fd = TEMP_FAILURE_RETRY(
+      open(path_.c_str(), O_WRONLY | O_APPEND | O_CREAT | O_CLOEXEC, 0600));
+  if (fd < 0) {
+    int error_no = errno;
+    TF_LITE_REPORT_ERROR(error_reporter_, "Could not open %s for writing: %s",
+                         path_.c_str(), std::strerror(error_no));
+    return kMinibenchmarkFailedToOpenStorageFileForWriting;
+  }
+  int lock_status = flock(fd, LOCK_EX);
+  int lock_error_no = errno;
+  if (lock_status < 0) {
+    close(fd);
+    TF_LITE_REPORT_ERROR(error_reporter_, "Could not flock %s: %s",
+                         path_.c_str(), std::strerror(lock_error_no));
+    return kMinibenchmarkFlockingStorageFileFailed;
+  }
+  absl::string_view bytes = data;
+  while (!bytes.empty()) {
+    ssize_t bytes_written =
+        TEMP_FAILURE_RETRY(write(fd, bytes.data(), bytes.size()));
+    if (bytes_written < 0) {
+      int error_no = errno;
+      close(fd);
+      TF_LITE_REPORT_ERROR(error_reporter_, "Could not write to %s: %s",
+                           path_.c_str(), std::strerror(error_no));
+      return kMinibenchmarkErrorWritingStorageFile;
+    }
+    bytes.remove_prefix(bytes_written);
+  }
+  if (TEMP_FAILURE_RETRY(fsync(fd)) < 0) {
+    int error_no = errno;
+    close(fd);
+    TF_LITE_REPORT_ERROR(error_reporter_, "Failed to fsync %s: %s",
+                         path_.c_str(), std::strerror(error_no));
+    return kMinibenchmarkErrorFsyncingStorageFile;
+  }
+  if (TEMP_FAILURE_RETRY(close(fd)) < 0) {
+    int error_no = errno;
+    TF_LITE_REPORT_ERROR(error_reporter_, "Failed to close %s: %s",
+                         path_.c_str(), std::strerror(error_no));
+    return kMinibenchmarkErrorClosingStorageFile;
+  }
+  return kMinibenchmarkSuccess;
+#else   // _WIN32
+  return kMinibenchmarkUnsupportedPlatform;
+#endif  // !_WIN32
+}
+
+ABSL_CONST_INIT const char kFlatbufferStorageIdentifier[] = "STO1";
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/fb_storage.h

@@ -0,0 +1,161 @@
+/* Copyright 2021 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_FB_STORAGE_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_FB_STORAGE_H_
+
+#include <errno.h>
+
+#include <cstring>
+#include <string>
+#include <vector>
+
+#include "absl/base/attributes.h"
+#include "absl/strings/string_view.h"
+#include "flatbuffers/base.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/status_codes.h"
+#include "tensorflow/lite/stderr_reporter.h"
+namespace tflite {
+namespace acceleration {
+
+// FileStorage wraps storage of data in a file with locking and error handling.
+// Locking makes appends and reads atomic, using flock(2).
+//
+// The locking in this class is not meant for general purpose multiple
+// reader/writer support, but primarily for the case where a previous instance
+// of a program has not finished and we'd like to not corrupt the file
+// unnecessarily.
+class FileStorage {
+ public:
+  FileStorage(absl::string_view path, ErrorReporter* error_reporter);
+  // Read contents into buffer_. Returns an error if file exists but cannot be
+  // read.
+  MinibenchmarkStatus ReadFileIntoBuffer();
+  // Append data to file. Resets the in-memory items and returns an error if
+  // writing fails in any way.
+  //
+  // This calls fsync() on the file to guarantee persistence and is hence quite
+  // expensive. The assumption is that this is not done often or in a critical
+  // path.
+  MinibenchmarkStatus AppendDataToFile(absl::string_view data);
+
+ protected:
+  std::string path_;
+  ErrorReporter* error_reporter_;
+  std::string buffer_;
+};
+
+// FlatbufferStorage stores several flatbuffer objects in a file. The primary
+// usage is for storing mini benchmark results.
+//
+// Flatbuffers are not designed for easy mutation. This class is append-only.
+// The intended usage is to store a log of events like 'start benchmark with
+// configuration X', 'benchmark results for X' / 'crash observed with X' that
+// are then parsed to make decisions about how to configure TFLite.
+//
+// The data is stored as consecutive length-prefixed flatbuffers with identifier
+// "STO1".
+ABSL_CONST_INIT extern const char kFlatbufferStorageIdentifier[];
+template <typename T>
+class FlatbufferStorage : protected FileStorage {
+ public:
+  explicit FlatbufferStorage(
+      absl::string_view path,
+      ErrorReporter* error_reporter = DefaultErrorReporter())
+      : FileStorage(path, error_reporter) {}
+  // Reads current contents. Returns an error if file is inaccessible or
+  // contents are corrupt. The file not existing is not an error.
+  MinibenchmarkStatus Read();
+  // Get count of objects stored.
+  size_t Count() { return contents_.size(); }
+  // Get object at index i, i < Count();
+  const T* Get(size_t i) { return contents_[i]; }
+
+  // Append a new object to storage and write out to disk. Returns an error if
+  // disk write or re-read fails.
+  MinibenchmarkStatus Append(flatbuffers::FlatBufferBuilder* fbb,
+                             flatbuffers::Offset<T> object);
+
+ private:
+  std::vector<const T*> contents_;
+};
+
+template <typename T>
+MinibenchmarkStatus FlatbufferStorage<T>::Read() {
+  contents_.clear();
+  MinibenchmarkStatus status = ReadFileIntoBuffer();
+  if (status != kMinibenchmarkSuccess) {
+    return status;
+  }
+  size_t remaining_size = buffer_.size();
+  const uint8_t* current_ptr =
+      reinterpret_cast<const uint8_t*>(buffer_.c_str());
+  while (remaining_size != 0) {
+    if (remaining_size < sizeof(flatbuffers::uoffset_t)) {
+      TF_LITE_REPORT_ERROR(
+          error_reporter_,
+          "Corrupt size-prefixed flatbuffer file %s (remaining size less than "
+          "size of uoffset_t)",
+          path_.c_str());
+      return kMinibenchmarkCorruptSizePrefixedFlatbufferFile;
+    }
+    flatbuffers::uoffset_t current_size =
+        flatbuffers::ReadScalar<flatbuffers::uoffset_t>(current_ptr);
+    flatbuffers::Verifier verifier(
+        current_ptr, sizeof(flatbuffers::uoffset_t) + current_size);
+    if (!verifier.VerifySizePrefixedBuffer<T>(kFlatbufferStorageIdentifier)) {
+      TF_LITE_REPORT_ERROR(
+          error_reporter_,
+          "Corrupt size-prefixed flatbuffer file %s (verifier returned false)",
+          path_.c_str());
+      return kMinibenchmarkCorruptSizePrefixedFlatbufferFile;
+    }
+    contents_.push_back(flatbuffers::GetSizePrefixedRoot<T>(current_ptr));
+    size_t consumed = sizeof(flatbuffers::uoffset_t) + current_size;
+    if (remaining_size < consumed) {
+      TF_LITE_REPORT_ERROR(
+          error_reporter_,
+          "Corrupt size-prefixed flatbuffer file %s (mismatched size "
+          "calculation)",
+          path_.c_str());
+      return kMinibenchmarkCorruptSizePrefixedFlatbufferFile;
+    }
+    remaining_size -= consumed;
+    current_ptr += consumed;
+  }
+  return kMinibenchmarkSuccess;
+}
+
+template <typename T>
+MinibenchmarkStatus FlatbufferStorage<T>::Append(
+    flatbuffers::FlatBufferBuilder* fbb, flatbuffers::Offset<T> object) {
+  contents_.clear();
+  fbb->FinishSizePrefixed(object, kFlatbufferStorageIdentifier);
+  const char* data = reinterpret_cast<const char*>(fbb->GetBufferPointer());
+  size_t size = fbb->GetSize();
+  MinibenchmarkStatus status = AppendDataToFile({data, size});
+  if (status != kMinibenchmarkSuccess) {
+    return status;
+  }
+  return Read();
+}
+
+}  // namespace acceleration
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_FB_STORAGE_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/fb_storage_test.cc

@@ -0,0 +1,150 @@
+/* Copyright 2021 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/fb_storage.h"
+
+#include <algorithm>
+#include <string>
+#include <thread>  // NOLINT - only production use is on Android, where std::thread is allowed
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/acceleration/configuration/configuration_generated.h"
+#include "tensorflow/lite/core/c/c_api_types.h"
+
+namespace tflite {
+namespace acceleration {
+namespace {
+
+std::string GetTemporaryDirectory() {
+#ifdef __ANDROID__
+  return "/data/local/tmp";
+#else
+  if (getenv("TEST_TMPDIR")) {
+    return getenv("TEST_TMPDIR");
+  }
+  if (getenv("TEMP")) {
+    return getenv("TEMP");
+  }
+  return ".";
+#endif
+}
+
+std::string GetStoragePath() {
+  std::string path = GetTemporaryDirectory() + "/storage.fb";
+  unlink(path.c_str());
+  return path;
+}
+
+TEST(FlatbufferStorageTest, AppendAndReadOneItem) {
+  std::string path = GetStoragePath();
+  flatbuffers::FlatBufferBuilder fbb;
+  flatbuffers::Offset<BenchmarkEvent> o =
+      CreateBenchmarkEvent(fbb, 0, BenchmarkEventType_START);
+
+  FlatbufferStorage<BenchmarkEvent> storage(path);
+  EXPECT_EQ(storage.Read(), kMinibenchmarkSuccess);
+  EXPECT_EQ(storage.Count(), 0);
+
+  EXPECT_EQ(storage.Append(&fbb, o), kMinibenchmarkSuccess);
+  ASSERT_EQ(storage.Count(), 1);
+  EXPECT_EQ(storage.Get(0)->event_type(), BenchmarkEventType_START);
+
+  storage = FlatbufferStorage<BenchmarkEvent>(path);
+  EXPECT_EQ(storage.Read(), kMinibenchmarkSuccess);
+  ASSERT_EQ(storage.Count(), 1);
+  EXPECT_EQ(storage.Get(0)->event_type(), BenchmarkEventType_START);
+}
+
+TEST(FlatbufferStorageTest, AppendAndReadThreeItems) {
+  std::string path = GetStoragePath();
+  FlatbufferStorage<BenchmarkEvent> storage(path);
+  EXPECT_EQ(storage.Read(), kMinibenchmarkSuccess);
+  EXPECT_EQ(storage.Count(), 0);
+
+  for (auto event : {BenchmarkEventType_START, BenchmarkEventType_ERROR,
+                     BenchmarkEventType_END}) {
+    flatbuffers::FlatBufferBuilder fbb;
+    flatbuffers::Offset<BenchmarkEvent> object =
+        CreateBenchmarkEvent(fbb, 0, event);
+    EXPECT_EQ(storage.Append(&fbb, object), kMinibenchmarkSuccess);
+  }
+
+  ASSERT_EQ(storage.Count(), 3);
+  EXPECT_EQ(storage.Get(0)->event_type(), BenchmarkEventType_START);
+  EXPECT_EQ(storage.Get(1)->event_type(), BenchmarkEventType_ERROR);
+  EXPECT_EQ(storage.Get(2)->event_type(), BenchmarkEventType_END);
+
+  storage = FlatbufferStorage<BenchmarkEvent>(path);
+  EXPECT_EQ(storage.Read(), kMinibenchmarkSuccess);
+  ASSERT_EQ(storage.Count(), 3);
+  EXPECT_EQ(storage.Get(0)->event_type(), BenchmarkEventType_START);
+  EXPECT_EQ(storage.Get(1)->event_type(), BenchmarkEventType_ERROR);
+  EXPECT_EQ(storage.Get(2)->event_type(), BenchmarkEventType_END);
+}
+
+TEST(FlatbufferStorageTest, PathDoesntExist) {
+  std::string path = GetTemporaryDirectory() + "/nosuchdirectory/storage.pb";
+  FlatbufferStorage<BenchmarkEvent> storage(path);
+  EXPECT_EQ(storage.Read(), kMinibenchmarkCantCreateStorageFile);
+}
+
+#ifndef __ANDROID__
+// chmod(0444) doesn't block writing on Android.
+TEST(FlatbufferStorageTest, WriteFailureResetsStorage) {
+  std::string path = GetStoragePath();
+  flatbuffers::FlatBufferBuilder fbb;
+  flatbuffers::Offset<BenchmarkEvent> o =
+      CreateBenchmarkEvent(fbb, 0, BenchmarkEventType_START);
+
+  FlatbufferStorage<BenchmarkEvent> storage(path);
+  EXPECT_EQ(storage.Append(&fbb, o), kMinibenchmarkSuccess);
+  ASSERT_EQ(storage.Count(), 1);
+
+  chmod(path.c_str(), 0444);
+  EXPECT_EQ(storage.Append(&fbb, o),
+            kMinibenchmarkFailedToOpenStorageFileForWriting);
+  ASSERT_EQ(storage.Count(), 0);
+}
+#endif  // !__ANDROID__
+
+TEST(FlatbufferStorageTest, Locking) {
+  std::string path = GetStoragePath();
+
+  std::vector<std::thread> threads;
+  const int kNumThreads = 4;
+  const int kIterations = 10;
+  threads.reserve(kNumThreads);
+  for (int i = 0; i < kNumThreads; i++) {
+    threads.push_back(std::thread([path]() {
+      for (int j = 0; j < kIterations; j++) {
+        FlatbufferStorage<BenchmarkEvent> storage(path);
+        flatbuffers::FlatBufferBuilder fbb;
+        flatbuffers::Offset<BenchmarkEvent> o =
+            CreateBenchmarkEvent(fbb, 0, BenchmarkEventType_START);
+        EXPECT_EQ(storage.Append(&fbb, o), kMinibenchmarkSuccess);
+      }
+    }));
+  }
+  std::for_each(threads.begin(), threads.end(),
+                [](std::thread& t) { t.join(); });
+  FlatbufferStorage<BenchmarkEvent> storage(path);
+  EXPECT_EQ(storage.Read(), kMinibenchmarkSuccess);
+  EXPECT_EQ(storage.Count(), kNumThreads * kIterations);
+}
+
+}  // namespace
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/file_lock.cc

@@ -0,0 +1,48 @@
+/* Copyright 2022 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/file_lock.h"
+
+#ifndef _WIN32
+#include <fcntl.h>
+#include <sys/file.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <unistd.h>
+#endif  // !_WIN32
+
+#include <string>
+
+namespace tflite {
+namespace acceleration {
+
+bool FileLock::TryLock() {
+#ifndef _WIN32
+  if (fd_ < 0) {
+    // O_CLOEXEC is needed for correctness, as another thread may call
+    // popen() and the callee would then inherit the lock if it's not O_CLOEXEC.
+    fd_ = open(path_.c_str(), O_WRONLY | O_CREAT | O_CLOEXEC, 0600);
+  }
+  if (fd_ < 0) {
+    return false;
+  }
+  if (flock(fd_, LOCK_EX | LOCK_NB) == 0) {
+    return true;
+  }
+#endif  // !_WIN32
+  return false;
+}
+
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/file_lock.h

@@ -0,0 +1,57 @@
+/* Copyright 2022 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_FILE_LOCK_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_FILE_LOCK_H_
+
+#ifndef _WIN32
+#include <unistd.h>
+#endif  // !_WIN32
+
+#include <string>
+
+namespace tflite {
+namespace acceleration {
+
+// A simple mutex lock implemented with file descriptor. Not supported in
+// Windows. This lock will release safely when the calling thread / process
+// crashes.
+class FileLock {
+ public:
+  explicit FileLock(const std::string& path) : path_(path) {}
+
+  // Move only.
+  FileLock(FileLock&& other) = default;
+  FileLock& operator=(FileLock&& other) = default;
+
+  ~FileLock() {
+#ifndef _WIN32
+    if (fd_ >= 0) {
+      close(fd_);
+    }
+#endif  // !_WIN32
+  }
+
+  // Returns whether the lock is acquired successfully.
+  bool TryLock();
+
+ private:
+  std::string path_;
+  int fd_ = -1;
+};
+
+}  // namespace acceleration
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_FILE_LOCK_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/file_lock_test.cc

@@ -0,0 +1,66 @@
+/* Copyright 2022 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/file_lock.h"
+
+#include <csignal>
+#include <iostream>
+#include <string>
+#include <utility>
+
+#include <gtest/gtest.h>
+
+namespace tflite {
+namespace acceleration {
+namespace {
+
+class FileLockTest : public ::testing::Test {
+ protected:
+  void SetUp() override { file_path_ = ::testing::TempDir() + "/file_lock"; }
+
+  std::string file_path_;
+};
+
+TEST_F(FileLockTest, CanLock) { EXPECT_TRUE(FileLock(file_path_).TryLock()); }
+
+TEST_F(FileLockTest, FailIfLockMoreThanOnce) {
+  FileLock lock_one(file_path_);
+  FileLock lock_two(file_path_);
+  ASSERT_TRUE(lock_one.TryLock());
+  EXPECT_FALSE(lock_two.TryLock());
+}
+
+TEST_F(FileLockTest, LockReleasedWhenThreadCrash) {
+  pid_t pid = fork();
+  if (pid == 0) {
+    // Child process crashed after TryLock().
+    FileLock lock(file_path_);
+    if (!lock.TryLock()) {
+      _exit(1);
+    }
+    std::cout << "Lock acquired successfully.";
+    kill(getpid(), SIGKILL);
+  }
+  int wstatus;
+  int w = waitpid(pid, &wstatus, WUNTRACED);
+  ASSERT_NE(w, -1);
+
+  // Lock again from main process.
+  FileLock lock_two(file_path_);
+  EXPECT_TRUE(lock_two.TryLock());
+}
+
+}  // namespace
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/gpu_module_plugin.cc

@@ -0,0 +1,110 @@
+/* Copyright 2020 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/gpu_module_plugin.h"
+
+#include <dlfcn.h>
+
+#include <memory>
+#include <string>
+
+#include "flatbuffers/buffer.h"  // from @flatbuffers
+#include "tensorflow/lite/acceleration/configuration/c/delegate_plugin.h"
+#include "tensorflow/lite/acceleration/configuration/configuration_generated.h"
+#include "tensorflow/lite/core/acceleration/configuration/delegate_registry.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/status_codes.h"
+#include "tensorflow/lite/logger.h"
+#include "tensorflow/lite/minimal_logging.h"
+
+namespace tflite {
+namespace acceleration {
+
+using ::tflite::delegates::TfLiteDelegatePtr;
+using SymbolFunc = const TfLiteDelegatePlugin*();
+
+// Function name used to get a pointer to GpuDelegatePlugin.
+constexpr char kPluginGetterSymbolName[] = "TfLiteGpuDelegatePluginCApi";
+
+std::unique_ptr<delegates::DelegatePluginInterface> GpuModulePlugin::New(
+    const TFLiteSettings& acceleration) {
+  return std::unique_ptr<GpuModulePlugin>(new GpuModulePlugin(acceleration));
+}
+
+int GpuModulePlugin::GetDelegateErrno(TfLiteDelegate* from_delegate) {
+  if (!plugin_handle_) {
+    return error_code_;
+  }
+  return plugin_handle_->get_delegate_errno(from_delegate);
+}
+
+TfLiteDelegatePtr GpuModulePlugin::Create() {
+  if (!plugin_handle_) {
+    return TfLiteDelegatePtr(nullptr, [](TfLiteDelegate*) {});
+  }
+  return TfLiteDelegatePtr(plugin_handle_->create(tflite_settings_),
+                           (plugin_handle_->destroy));
+}
+
+// In case GPU acceleration is not supported for this platform, we still need to
+// construct an empty object so that Create() can later be called on it.
+GpuModulePlugin::GpuModulePlugin(const TFLiteSettings& tflite_settings) {
+  TFLiteSettingsT settings_obj;
+  tflite_settings.UnPackTo(&settings_obj);
+  fbb_.Finish(CreateTFLiteSettings(fbb_, &settings_obj));
+  tflite_settings_ =
+      flatbuffers::GetRoot<TFLiteSettings>(fbb_.GetBufferPointer());
+
+  module_ = dlopen(tflite_settings_->stable_delegate_loader_settings()
+                       ->delegate_path()
+                       ->c_str(),
+                   RTLD_NOW | RTLD_LOCAL);
+  if (!module_) {
+    TFLITE_LOG_PROD(TFLITE_LOG_WARNING, "Failed to load Gpu Module from %s",
+                    tflite_settings_->stable_delegate_loader_settings()
+                        ->delegate_path()
+                        ->c_str());
+    error_code_ = kMinibenchmarkCannotLoadGpuModule;
+    return;
+  }
+  void* sym = dlsym(module_, kPluginGetterSymbolName);
+  if (!sym) {
+    TFLITE_LOG_PROD(TFLITE_LOG_WARNING, "Failed to create symbol '%s'",
+                    kPluginGetterSymbolName);
+    error_code_ = kMinibenchmarkCannotLoadGpuModule;
+    return;
+  }
+
+  plugin_handle_ = reinterpret_cast<SymbolFunc*>(sym)();
+  if (!plugin_handle_) {
+    TFLITE_LOG_PROD(
+        TFLITE_LOG_WARNING,
+        "GPU Module loaded successfully from %s, but plugin handle is null.",
+        tflite_settings_->stable_delegate_loader_settings()
+            ->delegate_path()
+            ->c_str());
+    error_code_ = kMinibenchmarkDelegatePluginNotFound;
+  }
+}
+
+GpuModulePlugin::~GpuModulePlugin() {
+  if (module_) {
+    dlclose(module_);
+  }
+}
+
+static auto* g_delegate_plugin_GpuModulePlugin =
+    new tflite::delegates::DelegatePluginRegistry::Register(
+        "GpuModulePlugin", GpuModulePlugin ::New);
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/gpu_module_plugin.h

@@ -0,0 +1,63 @@
+/* Copyright 2021 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_GPU_MODULE_PLUGIN_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_GPU_MODULE_PLUGIN_H_
+
+// This file provides the GpuPlugin class, which implements the
+// TFLite Delegate Plugin for the GPU Delegate.
+
+#include <memory>
+#include <string>
+
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "tensorflow/lite/acceleration/configuration/c/delegate_plugin.h"
+#include "tensorflow/lite/acceleration/configuration/configuration_generated.h"
+#include "tensorflow/lite/core/acceleration/configuration/delegate_registry.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/status_codes.h"
+
+namespace tflite {
+namespace acceleration {
+
+// A DelegatePlugin that uses external library to create GPU Plugin.
+class GpuModulePlugin : public delegates::DelegatePluginInterface {
+ public:
+  static std::unique_ptr<DelegatePluginInterface> New(
+      const TFLiteSettings& acceleration);
+
+  // Move only.
+  GpuModulePlugin(GpuModulePlugin&& other) = default;
+  GpuModulePlugin& operator=(GpuModulePlugin&& other) = default;
+  ~GpuModulePlugin() override;
+
+  delegates::TfLiteDelegatePtr Create() override;
+  int GetDelegateErrno(TfLiteDelegate* from_delegate) override;
+
+ private:
+  explicit GpuModulePlugin(const TFLiteSettings& tflite_settings);
+
+  // The handle to the loaded external library.
+  void* module_ = nullptr;
+  const TfLiteDelegatePlugin* plugin_handle_ = nullptr;
+  // A copy of the input tflite_settings.
+  flatbuffers::FlatBufferBuilder fbb_;
+  // A pointer to the data in fbb_.
+  const TFLiteSettings* tflite_settings_;
+  MinibenchmarkStatus error_code_ = kMinibenchmarkSuccess;
+};
+
+}  // namespace acceleration
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_GPU_MODULE_PLUGIN_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_common.h

@@ -0,0 +1,35 @@
+/* Copyright 2021 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_JPEG_COMMON_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_JPEG_COMMON_H_
+
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg.h"
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+
+struct JpegHeader {
+  int height;
+  int width;
+  int channels;
+  int bits_per_sample = BITS_IN_JSAMPLE;
+};
+
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_JPEG_COMMON_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_decompress_buffered_struct.h

@@ -0,0 +1,78 @@
+/* Copyright 2021 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_JPEG_DECOMPRESS_BUFFERED_STRUCT_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_JPEG_DECOMPRESS_BUFFERED_STRUCT_H_
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdlib>
+#include <vector>
+
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg.h"
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+
+// May provide an extra buffer of characters beyond the `jpeg_decompress_struct`
+// for some builds of Libjpeg Dynamic Library on Android that expect a larger
+// struct than we were compiled with. Zeroes out any allocated bytes beyond
+// sizeof(jpeg_decompress_struct). This class is exclusively used by
+// decode_jpeg.cc to resize `jpeg_decompress_struct`. This is to fix a struct
+// mismatch problem. See go/libjpeg-android for more details.
+class JpegDecompressBufferedStruct {
+ public:
+  explicit JpegDecompressBufferedStruct(std::size_t expected_size)
+      : resized_size_(std::max(sizeof(jpeg_decompress_struct), expected_size)),
+        buffer_(reinterpret_cast<char*>(malloc(resized_size_))) {
+    // Note: Malloc guarantees alignment for 8 bytes. Hence, using malloc
+    // instead of aligned_alloc.
+    // https://www.gnu.org/software/libc/manual/html_node/Aligned-Memory-Blocks.html
+    // alignof(jpeg_decompress_struct) is 8 bytes both on 32 and 64 bit.
+    // It's safe to align the buffered struct as
+    // alignof(jpeg_decompress_struct). This is because we only access the
+    // `jpeg_common_fields` fields of `jpeg_decompress_struct`, all of which are
+    // pointers. The alignment of these pointer fields is 8 and 4 bytes for 64
+    // bit and 32 bit platforms respectively. Since
+    // alignof(jpeg_decompress_struct) is 8 bytes on both platforms, accessing
+    // these fields shouldn't be a problem.
+    // Zero out any excess bytes. Zero-initialization is safe for the bytes
+    // beyond sizeof(jpeg_decompress_struct) because both the dynamic library
+    // and the implementation in decode_jpeg.cc limit their access only to
+    // `jpeg_common_fields` in `jpeg_decompress_struct`.
+    while (--expected_size >= sizeof(jpeg_decompress_struct)) {
+      buffer_[expected_size] = 0;
+    }
+  }
+  ~JpegDecompressBufferedStruct() { std::free(buffer_); }
+  JpegDecompressBufferedStruct(const JpegDecompressBufferedStruct&) = delete;
+  JpegDecompressBufferedStruct& operator=(const JpegDecompressBufferedStruct&) =
+      delete;
+  jpeg_decompress_struct* get() const {
+    return reinterpret_cast<jpeg_decompress_struct*>(buffer_);
+  }
+  int const size() { return resized_size_; }
+  const char* buffer() { return buffer_; }
+
+ private:
+  int resized_size_;
+  char* const buffer_;
+};
+
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_JPEG_DECOMPRESS_BUFFERED_STRUCT_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_decompress_buffered_struct_test.cc

@@ -0,0 +1,56 @@
+/* Copyright 2021 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_decompress_buffered_struct.h"
+
+#include <gtest/gtest.h>
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+namespace {
+
+const int kSizeOfJpegDecompressStruct = sizeof(jpeg_decompress_struct);
+
+TEST(JpegDecompressBufferedStructTest,
+     ExpectInitializationSizeMatchesStructSize) {
+  JpegDecompressBufferedStruct buffered_struct(kSizeOfJpegDecompressStruct);
+  EXPECT_EQ(buffered_struct.size(), kSizeOfJpegDecompressStruct);
+}
+
+TEST(JpegDecompressBufferedStructTest,
+     StructWithSizeGreaterThanCompiledStruct) {
+  int excess_bytes = 16;
+  JpegDecompressBufferedStruct buffered_struct(kSizeOfJpegDecompressStruct +
+                                               excess_bytes);
+  EXPECT_EQ(buffered_struct.size(), kSizeOfJpegDecompressStruct + excess_bytes);
+  const char* buffer = buffered_struct.buffer();
+  ASSERT_NE(buffer, nullptr);
+  while (excess_bytes--) {
+    EXPECT_EQ(
+        (unsigned char)(buffer[kSizeOfJpegDecompressStruct + excess_bytes]),
+        '\0');
+  }
+}
+
+TEST(JpegDecompressBufferedStructTest, StructWithSizeLessThanCompiledStruct) {
+  JpegDecompressBufferedStruct buffered_struct(kSizeOfJpegDecompressStruct -
+                                               16);
+  EXPECT_EQ(buffered_struct.size(), kSizeOfJpegDecompressStruct);
+}
+
+}  // namespace
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_header_parser.cc

@@ -0,0 +1,337 @@
+/* Copyright 2021 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_header_parser.h"
+
+#include <cstdint>
+#include <string>
+
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/minimal_logging.h"
+#include "tensorflow/lite/string_util.h"
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+
+/*
+
+JPEG file overall file structure
+
+SOI Marker                 FFD8
+Marker XX size=SSSS        FFXX	SSSS	DDDD......
+Marker YY size=TTTT        FFYY	TTTT	DDDD......
+SOFn marker with the info we want
+SOS Marker size=UUUU       FFDA	UUUU	DDDD....
+Image stream               I I I I....
+EOI Marker                 FFD9
+
+The first marker is either APP0 (JFIF format) or APP1 (EXIF format)
+
+We support only JFIF images
+*/
+
+namespace {
+
+using MarkerId = uint16_t;
+
+void AsWord(int value, char* msb, char* lsb) {
+  *msb = static_cast<char>(value >> 8);
+  *lsb = static_cast<char>(value);
+}
+
+// JFIF spec at
+// https://www.ecma-international.org/publications-and-standards/technical-reports/ecma-tr-98/
+// Marker definition summary at
+// http://lad.dsc.ufcg.edu.br/multimidia/jpegmarker.pdf
+// Overall JPEG File structure with discussion of the supported number of
+// channels per format
+// https://docs.oracle.com/javase/8/docs/api/javax/imageio/metadata/doc-files/jpeg_metadata.html
+//
+
+class JfifHeaderParser {
+ public:
+  explicit JfifHeaderParser(const tflite::StringRef& jpeg_image_data)
+      : jpeg_image_data_(jpeg_image_data), offset_(0) {
+    if (!IsJpegImage(jpeg_image_data_)) {
+      is_valid_image_buffer_ = false;
+      validation_error_message_ = "Not a valid JPEG image.";
+    } else if (!IsJfifImage(jpeg_image_data_)) {
+      is_valid_image_buffer_ = false;
+      validation_error_message_ = "Image is not in JFIF format.";
+      return;
+    } else {
+      is_valid_image_buffer_ = true;
+    }
+  }
+
+#define ENSURE_READ_STATUS(a)                           \
+  do {                                                  \
+    const TfLiteStatus s = (a);                         \
+    if (s != kTfLiteOk) {                               \
+      return {s, "Error trying to parse JPEG header."}; \
+    }                                                   \
+  } while (0)
+
+  Status ReadJpegHeader(JpegHeader* result) {
+    if (!is_valid_image_buffer_) {
+      return {kTfLiteError, validation_error_message_};
+    }
+
+    Status move_to_sof_status = MoveToStartOfFrameMarker();
+    if (move_to_sof_status.code != kTfLiteOk) {
+      return move_to_sof_status;
+    }
+
+    ENSURE_READ_STATUS(SkipBytes(2));  // skipping marker length
+    char precision;
+    ENSURE_READ_STATUS(ReadByte(&precision));
+    uint16_t height;
+    ENSURE_READ_STATUS(ReadWord(&height));
+    uint16_t width;
+    ENSURE_READ_STATUS(ReadWord(&width));
+    char num_of_components;
+    ENSURE_READ_STATUS(ReadByte(&num_of_components));
+
+    if (num_of_components != 1 && num_of_components != 3) {
+      return {kTfLiteError,
+              "A JFIF image without App14 marker doesn't support a number of "
+              "components = " +
+                  std::to_string(static_cast<int>(num_of_components))};
+    }
+
+    result->width = width;
+    result->height = height;
+    result->channels = num_of_components;
+    result->bits_per_sample = precision;
+
+    return {kTfLiteOk, ""};
+  }
+
+  Status ApplyHeaderToImage(const JpegHeader& new_header,
+                            std::string& write_to) {
+    if (!is_valid_image_buffer_) {
+      return {kTfLiteError, validation_error_message_};
+    }
+
+    Status move_to_sof_status = MoveToStartOfFrameMarker();
+    if (move_to_sof_status.code != kTfLiteOk) {
+      return move_to_sof_status;
+    }
+    ENSURE_READ_STATUS(SkipBytes(2));  // skipping marker length
+
+    if (!HasData(6)) {
+      return {kTfLiteError,
+              "Invalid SOF marker, image buffer ends before end of marker"};
+    }
+
+    char header[6];
+    header[0] = static_cast<char>(new_header.bits_per_sample);
+    AsWord(new_header.height, header + 1, header + 2);
+    AsWord(new_header.width, header + 3, header + 4);
+    header[5] = static_cast<char>(new_header.channels);
+
+    write_to.clear();
+    write_to.append(jpeg_image_data_.str, offset_);
+    write_to.append(header, 6);
+
+    ENSURE_READ_STATUS(SkipBytes(6));
+    if (HasData()) {
+      write_to.append(jpeg_image_data_.str + offset_,
+                      jpeg_image_data_.len - offset_);
+    }
+
+    return {kTfLiteOk, ""};
+  }
+
+ private:
+  const tflite::StringRef jpeg_image_data_;
+  // Using int for consistency with the size in StringRef
+  int offset_;
+  bool is_valid_image_buffer_;
+  std::string validation_error_message_;
+
+  // Moves to the begin of the first SOF marker
+  Status MoveToStartOfFrameMarker() {
+    const MarkerId kStartOfStreamMarkerId = 0xFFDA;  // Start of image data
+
+    offset_ = 0;
+    ENSURE_READ_STATUS(SkipBytes(4));  // skipping SOI and APP0 marker IDs
+    ENSURE_READ_STATUS(SkipCurrentMarker());  // skipping APP0
+    MarkerId curr_marker_id;
+    // We need at least 2 bytes for the marker ID and 2 for the length
+    while (HasData(/*min_data_size=*/4)) {
+      ENSURE_READ_STATUS(ReadWord(&curr_marker_id));
+      // We are breaking at the first met SOF marker. This won't generate
+      // results inconsistent with LibJPEG because only
+      // image with a single SOF marker are successfully parsed by it.
+      // LibJPEG fails if more than one marker is found in the header (see
+      // https://github.com/libjpeg-turbo/libjpeg-turbo/blob/main/jerror.h#L121
+      // and
+      // https://github.com/libjpeg-turbo/libjpeg-turbo/blob/main/jdmarker.c#L264-L265
+      if (IsStartOfFrameMarkerId(curr_marker_id)) {
+        break;
+      }
+      if (curr_marker_id == kStartOfStreamMarkerId) {
+        return {kTfLiteError, "Error trying to parse JPEG header."};
+      }
+      ENSURE_READ_STATUS(SkipCurrentMarker());
+    }
+
+    return {kTfLiteOk, ""};
+  }
+
+#undef ENSURE_READ_STATUS
+
+  bool HasData(int min_data_size = 1) {
+    return offset_ <= jpeg_image_data_.len - min_data_size;
+  }
+
+  TfLiteStatus SkipBytes(int bytes) {
+    if (!HasData(bytes)) {
+      TFLITE_LOG(TFLITE_LOG_WARNING,
+                 "Trying to move out of image boundaries from offset %d, "
+                 "skipping %d bytes",
+                 offset_, bytes);
+      return kTfLiteError;
+    }
+
+    offset_ += bytes;
+
+    return kTfLiteOk;
+  }
+
+  TfLiteStatus ReadByte(char* result) {
+    if (!HasData()) {
+      return kTfLiteError;
+    }
+
+    *result = jpeg_image_data_.str[offset_];
+
+    return SkipBytes(1);
+  }
+
+  TfLiteStatus ReadWord(uint16_t* result) {
+    TF_LITE_ENSURE_STATUS(ReadWordAt(jpeg_image_data_, offset_, result));
+    return SkipBytes(2);
+  }
+
+  TfLiteStatus SkipCurrentMarker() {
+    // We just read the marker ID so we are on top of the marker len
+    uint16_t full_marker_len;
+    TF_LITE_ENSURE_STATUS(ReadWord(&full_marker_len));
+    if (full_marker_len <= 2) {
+      TFLITE_LOG(TFLITE_LOG_WARNING,
+                 "Invalid marker length %d read at offset %X", full_marker_len,
+                 offset_);
+      return kTfLiteError;
+    }
+
+    // The marker len includes the 2 bytes of marker length
+    return SkipBytes(full_marker_len - 2);
+  }
+
+  static TfLiteStatus ReadWordAt(const tflite::StringRef& jpeg_image_data,
+                                 int read_offset, uint16_t* result) {
+    if (read_offset < 0 || read_offset + 2 > jpeg_image_data.len) {
+      return kTfLiteError;
+    }
+    // Cast to unsigned since char can be signed.
+    const unsigned char* buf =
+        reinterpret_cast<const unsigned char*>(jpeg_image_data.str);
+
+    *result = (buf[read_offset] << 8) + buf[read_offset + 1];
+
+    return kTfLiteOk;
+  }
+
+  static bool IsJpegImage(const tflite::StringRef& jpeg_image_data) {
+    const MarkerId kStartOfImageMarkerId = 0xFFD8;
+    const MarkerId kEndOfImageMarkerId = 0xFFD9;
+
+    MarkerId soi_marker_id;
+    MarkerId eoi_marker_id;
+    if (ReadWordAt(jpeg_image_data, 0, &soi_marker_id) != kTfLiteOk) {
+      return false;
+    }
+    if (ReadWordAt(jpeg_image_data, jpeg_image_data.len - 2, &eoi_marker_id) !=
+        kTfLiteOk) {
+      return false;
+    }
+
+    return (soi_marker_id == kStartOfImageMarkerId) &&
+           (eoi_marker_id == kEndOfImageMarkerId);
+  }
+
+  static bool IsJfifImage(const tflite::StringRef& jpeg_image_data) {
+    const MarkerId kApp0MarkerId = 0xFFE0;  // First marker in JIFF image
+
+    MarkerId app_marker_id;
+    if ((ReadWordAt(jpeg_image_data, 2, &app_marker_id) != kTfLiteOk) ||
+        (app_marker_id != kApp0MarkerId)) {
+      return false;
+    }
+
+    // Checking Jfif identifier string "JFIF\0" in APP0 Marker
+    const std::string kJfifIdString{"JFIF\0", 5};
+
+    // The ID starts after SOI (2 bytes), APP0 marker IDs (2 bytes) and 2 other
+    // bytes with APP0 marker length
+    const int KJfifIdStringStartOffset = 6;
+
+    if (KJfifIdStringStartOffset + kJfifIdString.size() >=
+        jpeg_image_data.len) {
+      TFLITE_LOG(TFLITE_LOG_WARNING,
+                 "Invalid image, reached end of data at offset while "
+                 "parsing APP0 header");
+      return false;
+    }
+
+    const std::string actualImgId(
+        jpeg_image_data.str + KJfifIdStringStartOffset, kJfifIdString.size());
+    if (kJfifIdString != actualImgId) {
+      TFLITE_LOG(TFLITE_LOG_WARNING, "Invalid image, invalid APP0 header");
+
+      return false;
+    }
+
+    return true;
+  }
+
+  static bool IsStartOfFrameMarkerId(MarkerId marker_id) {
+    return 0xFFC0 <= marker_id && marker_id < 0xFFCF;
+  }
+};
+
+}  // namespace
+Status ReadJpegHeader(const tflite::StringRef& jpeg_image_data,
+                      JpegHeader* header) {
+  JfifHeaderParser parser(jpeg_image_data);
+
+  return parser.ReadJpegHeader(header);
+}
+
+Status BuildImageWithNewHeader(const tflite::StringRef& orig_jpeg_image_data,
+                               const JpegHeader& new_header,
+                               std::string& new_image_data) {
+  JfifHeaderParser parser(orig_jpeg_image_data);
+
+  return parser.ApplyHeaderToImage(new_header, new_image_data);
+}
+
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_header_parser.h

@@ -0,0 +1,47 @@
+/* Copyright 2021 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_JPEG_HEADER_PARSER_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_JPEG_HEADER_PARSER_H_
+
+#include <string>
+#include <tuple>
+
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg_status.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_common.h"
+#include "tensorflow/lite/string_type.h"
+#include "tensorflow/lite/string_util.h"
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+
+// Extract the info in JpegHeader from the given buffer.
+// Fails if the buffer doesn't contain a valid JPEG image in JFIF format.
+Status ReadJpegHeader(const tflite::StringRef& jpeg_image_data,
+                      JpegHeader* header);
+
+// Writes into the given string the content of the JPEG image altered with
+// the content of new_header.
+// This is intented to be used in tests to forge existing images.
+Status BuildImageWithNewHeader(const tflite::StringRef& orig_jpeg_image_data,
+                               const JpegHeader& new_header,
+                               std::string& new_image_data);
+
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_JPEG_HEADER_PARSER_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_header_parser_test.cc

@@ -0,0 +1,124 @@
+/* Copyright 2021 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_header_parser.h"
+
+#include <cstddef>
+#include <ostream>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/embedded_chessboard_jpeg.h"
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+
+void PrintTo(const Status& status, std::ostream* os) {
+  *os << "{ code: " + std::to_string(status.code) + ", error_message: '" +
+             status.error_message + "'}";
+}
+
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite
+
+namespace {
+
+using ::testing::AllOf;
+using ::testing::Eq;
+using ::testing::Field;
+using ::testing::Matcher;
+
+using tflite::acceleration::decode_jpeg_kernel::JpegHeader;
+using tflite::acceleration::decode_jpeg_kernel::ReadJpegHeader;
+
+Matcher<JpegHeader> JpegHeaderEq(const JpegHeader& expected) {
+  return AllOf(
+      Field(&JpegHeader::channels, Eq(expected.channels)),
+      Field(&JpegHeader::height, Eq(expected.height)),
+      Field(&JpegHeader::width, Eq(expected.width)),
+      Field(&JpegHeader::bits_per_sample, Eq(expected.bits_per_sample)));
+}
+
+using tflite::acceleration::decode_jpeg_kernel::Status;
+
+Matcher<Status> StatusEq(const Status& expected) {
+  return AllOf(Field(&Status::code, Eq(expected.code)),
+               Field(&Status::error_message, Eq(expected.error_message)));
+}
+
+const int kChessboardImgHeight = 300;
+const int kChessboardImgWidth = 250;
+const int kChessboardImgChannels = 3;
+
+TEST(ReadJpegHeader, ShouldParseValidJpgImage) {
+  const tflite::StringRef chessboard_image{
+      reinterpret_cast<const char*>(g_tflite_acceleration_chessboard_jpeg),
+      static_cast<size_t>(g_tflite_acceleration_chessboard_jpeg_len)};
+  ASSERT_GT(chessboard_image.len, 4);
+
+  JpegHeader header;
+
+  ASSERT_THAT(ReadJpegHeader(chessboard_image, &header),
+              StatusEq({kTfLiteOk, ""}));
+  EXPECT_THAT(header, JpegHeaderEq({kChessboardImgHeight, kChessboardImgWidth,
+                                    kChessboardImgChannels}));
+}
+
+TEST(ReadJpegHeader, ShouldFailForInvalidJpegImage) {
+  const std::string invalid_image = "invalid image content";
+  const tflite::StringRef invalid_image_ref{invalid_image.c_str(),
+                                            invalid_image.size()};
+
+  JpegHeader header;
+
+  EXPECT_THAT(ReadJpegHeader(invalid_image_ref, &header),
+              StatusEq({kTfLiteError, "Not a valid JPEG image."}));
+}
+
+TEST(ReadJpegHeader, ShouldFailForEmptyJpegImage) {
+  const tflite::StringRef invalid_image_ref{"", 0};
+
+  JpegHeader header;
+
+  EXPECT_THAT(ReadJpegHeader(invalid_image_ref, &header),
+              StatusEq({kTfLiteError, "Not a valid JPEG image."}));
+}
+
+TEST(ApplyHeaderToImage, ReturnsNewImageWithDifferentHeader) {
+  const tflite::StringRef chessboard_image{
+      reinterpret_cast<const char*>(g_tflite_acceleration_chessboard_jpeg),
+      static_cast<size_t>(g_tflite_acceleration_chessboard_jpeg_len)};
+
+  JpegHeader new_header{
+      .height = 20, .width = 30, .channels = 1, .bits_per_sample = 3};
+
+  std::string new_image_data;
+
+  ASSERT_THAT(
+      BuildImageWithNewHeader(chessboard_image, new_header, new_image_data),
+      StatusEq({kTfLiteOk, ""}));
+
+  const tflite::StringRef altered_image{new_image_data.c_str(),
+                                        new_image_data.size()};
+  JpegHeader header;
+  ASSERT_THAT(ReadJpegHeader(altered_image, &header),
+              StatusEq({kTfLiteOk, ""}));
+  EXPECT_THAT(header, JpegHeaderEq(new_header));
+}
+
+}  // namespace

tensorflow/lite/experimental/acceleration/mini_benchmark/libc_handle.cc

@@ -0,0 +1,62 @@
+/* Copyright 2021 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/libc_handle.h"
+
+#ifdef __ANDROID__
+#include <dlfcn.h>
+#endif
+#include <stdio.h>
+
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg_status.h"
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+
+LibCHandle LibCHandle::Create(Status &status) {
+#ifndef __ANDROID__
+#ifndef _WIN32
+  // Use the statically linked C lib.
+  return LibCHandle(nullptr, ::fmemopen);
+#else   // _WIN32
+  status = {kTfLiteError, "Windows not supported."};
+  return LibCHandle(nullptr, nullptr);
+#endif  // !_WIN32
+#else   // __ANDROID__
+  void *libc = nullptr;
+  FmemopenPtr fmemopen_ptr = nullptr;
+  if (!(libc = dlopen("libc.so", RTLD_NOW | RTLD_LOCAL))) {
+    status = {kTfLiteError,
+              "Failed to load the libc dynamic shared object library."};
+    return LibCHandle(nullptr, nullptr);
+  }
+
+  if (!(fmemopen_ptr =
+            reinterpret_cast<FmemopenPtr>(dlsym(libc, "fmemopen")))) {
+    status = {kTfLiteError, "Failed to dynamically load the method: fmemopen"};
+    return LibCHandle(nullptr, nullptr);
+  }
+  status = {kTfLiteOk, ""};
+  return LibCHandle(libc, fmemopen_ptr);
+#endif  // !__ANDROID__
+}
+
+FILE *LibCHandle::fmemopen(void *buf, size_t size, const char *mode) const {
+  return fmemopen_(buf, size, mode);
+}
+
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/libc_handle.h

@@ -0,0 +1,81 @@
+/* Copyright 2021 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_LIBC_HANDLE_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_LIBC_HANDLE_H_
+
+#ifdef __ANDROID__
+#include <dlfcn.h>
+#endif
+
+#include <cstdio>
+#include <memory>
+#include <utility>
+
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg_status.h"
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+
+// This class offers a handle to C Standard Library (LibC) shared object
+// library on Android. It offers pointers to functions in LibC.
+// Fmemopen is available as native API from Android SDK 23 onwards. In order to
+// support Android devices from SDK 21 onwards, we load fmemopen dynamically
+// from the libc shared object library.
+// TODO(b/172544567): Support Apple.
+class LibCHandle {
+ public:
+  // Factory to get an initialised instance of LibCHandle.
+  // Loads the libc dynamic library and gets handle to all the
+  // required functions. Stores the initialisation status in `status`.
+  static LibCHandle Create(Status& status);
+  LibCHandle(LibCHandle const&) = delete;
+  LibCHandle& operator=(const LibCHandle&) = delete;
+  LibCHandle(LibCHandle&& other)
+      : libc_(std::exchange(other.libc_, nullptr)),
+        fmemopen_(std::exchange(other.fmemopen_, nullptr)) {}
+  LibCHandle& operator=(LibCHandle&& other) {
+    if (&other != this) {
+      CloseHandle();
+      libc_ = std::exchange(other.libc_, nullptr);
+      fmemopen_ = std::exchange(other.fmemopen_, nullptr);
+    }
+    return *this;
+  }
+  ~LibCHandle() { CloseHandle(); }
+  // Definition can be found here
+  // https://man7.org/linux/man-pages/man3/fmemopen.3.html
+  FILE* fmemopen(void* buf, size_t size, const char* mode) const;
+
+ private:
+  using FmemopenPtr = FILE* (*)(void*, size_t, const char*);
+  LibCHandle(void* libc, FmemopenPtr ptr) : libc_(libc), fmemopen_(ptr) {}
+  // Closes the dynamic library loaded in libc_.
+  void CloseHandle() {
+#ifdef __ANDROID__
+    if (libc_ != nullptr) {
+      dlclose(libc_);
+    }
+#endif
+  }
+  void* libc_ = nullptr;
+  FmemopenPtr fmemopen_ = nullptr;
+};
+
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_LIBC_HANDLE_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/libc_handle_test.cc

@@ -0,0 +1,34 @@
+/* Copyright 2021 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/libc_handle.h"
+
+#include <gtest/gtest.h>
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+namespace {
+
+TEST(LibCHandleTest, LoadingSucceedsAndroidPlatforms) {
+  Status status;
+  LibCHandle handle = LibCHandle::Create(status);
+  EXPECT_EQ(status.error_message, "");
+  EXPECT_EQ(status.code, kTfLiteOk);
+}
+
+}  // namespace
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg.h

@@ -0,0 +1,20 @@
+/* Copyright 2021 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_LIBJPEG_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_LIBJPEG_H_
+
+#include "tensorflow/core/platform/jpeg.h"
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_LIBJPEG_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder.cc

@@ -0,0 +1,299 @@
+/* Copyright 2021 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder.h"
+
+#include <setjmp.h>
+
+#include <algorithm>
+#include <cctype>
+#include <cstddef>
+#include <cstdio>
+#include <cstring>
+#include <functional>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/strings/match.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg_status.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_decompress_buffered_struct.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_header_parser.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_handle.h"
+#include "tensorflow/lite/minimal_logging.h"
+#include "tensorflow/lite/string_type.h"
+#include "tensorflow/lite/string_util.h"
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+
+// Limiting max image size to 10,000x10,000x3
+// This size would fit on 32 bit systems.
+// static
+const size_t LibjpegDecoder::kMaxImageHeight = 10000;
+// static
+const size_t LibjpegDecoder::kMaxImageWidth = 10000;
+
+constexpr char kSizeMismatchError[] =
+    "JPEG parameter struct mismatch: library thinks size is ";
+
+LibjpegDecoder::Impl::Impl(size_t decompress_struct_size,
+                           const LibjpegHandle* handle)
+    : decompress_struct_size_(decompress_struct_size),
+      handle_(handle),
+      cinfo_(decompress_struct_size) {
+  cinfo_.get()->err = handle->jpeg_std_error_(&jerr_);
+  jerr_.error_exit = ErrorExit;
+  cinfo_.get()->client_data = this;
+}
+
+void LibjpegDecoder::Impl::ErrorExit(j_common_ptr cinfo) {
+  Impl* const impl = reinterpret_cast<Impl*>(cinfo->client_data);
+  char message[JMSG_LENGTH_MAX];
+  cinfo->err->format_message(cinfo, message);
+  impl->status_.code = kTfLiteError;
+  impl->status_.error_message = message;
+  // Libjpeg aborts the program in case of any errors by using longjmp and then
+  // calling exit(). The only way to avoid this, is to transfer the control flow
+  // to the caller by using setjmp/longjmp.
+  // Note: Ensure that function containing the corresponding setjmp() is
+  // guaranteed not to have completed execution.
+  // https://wiki.sei.cmu.edu/confluence/display/c/MSC22-C.+Use+the+setjmp%28%29%2C+longjmp%28%29+facility+securely
+  longjmp(impl->env_, 1);
+}
+
+Status ExtractSizeFromErrorMessage(const std::string& error_message,
+                                   size_t& expected_size) {
+  Status status;
+  // Special error handling for struct mismatch issues.
+  // If there's a mismatch, set `expected_size` with the expected
+  // size. Error messages are like this: "JPEG parameter struct
+  // mismatch: library thinks size is 480, caller expects 464".
+  static const int kExpLengthStart = strlen(kSizeMismatchError);
+  int end = kExpLengthStart;
+  while (end < error_message.length() && std::isdigit(error_message[end])) {
+    end++;
+  }
+  if (end > kExpLengthStart) {
+    expected_size = std::stoi(error_message.substr(kExpLengthStart, end));
+  } else {
+    status.code = kTfLiteError;
+    status.error_message =
+        "Couldn't parse the size from message: \'" + error_message + "\'";
+  }
+  return status;
+}
+
+std::unique_ptr<LibjpegDecoder> LibjpegDecoder::Create(Status& status) {
+  std::unique_ptr<LibjpegDecoder> decoder(
+      new LibjpegDecoder(LibCHandle::Create(status)));
+  if (status.code != kTfLiteOk) {
+    return nullptr;
+  }
+  decoder->libjpeg_handle_ = LibjpegHandle::Create(status);
+  if (decoder->libjpeg_handle_ == nullptr) {
+    return nullptr;
+  }
+
+  // Tries to probe the libjpeg library to get the expected size of
+  // `jpeg_decompress_struct`.
+  Impl impl(sizeof(jpeg_decompress_struct), decoder->libjpeg_handle_.get());
+  impl.jpeg_CreateDecompress(LibjpegHandle::kLibjpegVersion,
+                             sizeof(jpeg_decompress_struct));
+  status = impl.status();
+  if (status.code == kTfLiteOk) {
+    decoder->expected_size_for_decompress_struct_ =
+        sizeof(jpeg_decompress_struct);
+    return decoder;
+  }
+  if (!absl::StrContains(status.error_message, kSizeMismatchError)) {
+    return nullptr;
+  }
+  status = ExtractSizeFromErrorMessage(
+      status.error_message, decoder->expected_size_for_decompress_struct_);
+  if (status.code != kTfLiteOk) {
+    return nullptr;
+  }
+  return decoder;
+}
+
+namespace {
+
+std::string JpegHeaderToString(const JpegHeader& header) {
+  return "(" + std::to_string(header.height) + ", " +
+         std::to_string(header.width) + ", " + std::to_string(header.channels) +
+         ", " + std::to_string(header.bits_per_sample) + ")";
+}
+
+}  // namespace
+
+Status LibjpegDecoder::DecodeImage(const tflite::StringRef& encoded,
+                                   const JpegHeader& expected_image_dimensions,
+                                   unsigned char* decoded,
+                                   const size_t& decoded_size) const {
+  if (expected_image_dimensions.bits_per_sample != 8) {
+    return {kTfLiteError, "Supporting only images with 8 bits per sample"};
+  }
+  if (expected_image_dimensions.channels != 1 &&
+      expected_image_dimensions.channels != 3) {
+    return {kTfLiteError, "Supporting only images with 1 or 3 channels"};
+  }
+  if (expected_image_dimensions.width > kMaxImageWidth ||
+      expected_image_dimensions.height > kMaxImageHeight) {
+    return {kTfLiteError, "Image is too big, dimensions (" +
+                              std::to_string(expected_image_dimensions.width) +
+                              "," +
+                              std::to_string(expected_image_dimensions.width) +
+                              ") larger than the maximum allowed (" +
+                              std::to_string(kMaxImageWidth) + ", " +
+                              std::to_string(kMaxImageHeight) + ")"};
+  }
+  // We match the buffer size and the expected size of the decoded image from
+  // the header to prevent buffer overflows.
+  JpegHeader header;
+  Status read_header_status = ReadJpegHeader(encoded, &header);
+  if (read_header_status.code != kTfLiteOk) {
+    return read_header_status;
+  }
+
+  if (expected_image_dimensions.channels != header.channels ||
+      expected_image_dimensions.width != header.width ||
+      expected_image_dimensions.height != header.height ||
+      expected_image_dimensions.bits_per_sample != header.bits_per_sample) {
+    return {kTfLiteError, "Decoded image size " + JpegHeaderToString(header) +
+                              " is different from provided image size " +
+                              JpegHeaderToString(expected_image_dimensions)};
+  }
+
+  size_t header_image_size = static_cast<size_t>(header.width) *
+                             static_cast<size_t>(header.height) *
+                             static_cast<size_t>(header.channels);
+
+  if (header_image_size != decoded_size) {
+    return {kTfLiteError, "Size of buffer(" + std::to_string(decoded_size) +
+                              ") for storing decoded image must be equal to "
+                              "the size of decoded image(" +
+                              std::to_string(header_image_size) + ")."};
+  }
+
+  // Dropping constness as fmemopen requires non-const buffers.
+  char* image_buffer = const_cast<char*>(encoded.str);
+  size_t image_size = encoded.len;
+  std::unique_ptr<FILE, std::function<void(FILE*)>> file(
+      libc_handle_.fmemopen(image_buffer, image_size, "r"),
+      [](FILE* f) { fclose(f); });
+  if (file == nullptr) {
+    return {kTfLiteError, "Fmemopen failed."};
+  }
+
+  Impl impl(expected_size_for_decompress_struct_, libjpeg_handle_.get());
+  if (impl.jpeg_CreateDecompress(LibjpegHandle::kLibjpegVersion,
+                                 expected_size_for_decompress_struct_)) {
+    return impl.status();
+  }
+  if (impl.jpeg_stdio_src(file.get())) {
+    return impl.status();
+  }
+  // jpeg_read_header() must be called before calling jpeg_start_decompress().
+  // It initialises decompression parameters to default values.
+  // jpeg_read_header() should not be relied upon for getting image information
+  // (width, height) etc. Fields populated by jpeg_read_header() such as
+  // `image_width` and `image_height` come after the `jpeg_common_fields` and
+  // these may have been shifted in the struct on some builds of libjpeg.
+  // See go/libjpeg-android.
+  int read_header_result = 0;
+  if (impl.jpeg_read_header(read_header_result, true) != kTfLiteOk) {
+    return impl.status();
+  }
+  if (read_header_result != JPEG_HEADER_OK) {
+    return {kTfLiteError, "Failed call jpeg_read_header"};
+  }
+  boolean start_decompress_result = false;
+  if (impl.jpeg_start_decompress(start_decompress_result) != kTfLiteOk) {
+    return impl.status();
+  }
+  if (!start_decompress_result) {
+    return {kTfLiteError, "Failed call jpeg_start_decompress_"};
+  }
+
+  size_t height = header.height;
+  size_t row_stride = header.width * header.channels;
+
+  // Decoding the data in a buffer as large as the largest allowed JPEG
+  // image row to avoid overflows in case  we are reading a wrong value for the
+  // image size in the header or we are receiving an image with an header
+  // deliberately incorrect to cause a buffer overflow.
+  // Using 4 channels because the decode color process can handle 3 or 4
+  // channels:
+  // https://github.com/libjpeg-turbo/libjpeg-turbo/blob/main/jdcolor.c#L383
+  const size_t kMaxImageSize = JPEG_MAX_DIMENSION * 4;
+  // Initializing the buffer in case we are trying to read more data than
+  // actually available to avoid having access to uninitialized memory.
+  // Libjpeg turbo actually fills the unread bytes with zeros
+  // https://github.com/libjpeg-turbo/libjpeg-turbo/blob/main/jdhuff.c#L360
+  // but we don't know what the library on the target system would do.
+  // Output tensor data stored as RGBRGBRGB..., row-wise for all images.
+  std::vector<unsigned char> decode_buffer(kMaxImageSize);
+  // Do not rely on fields such as `output_scanline` from
+  // `jpeg_decompress_struct` as these would have been shifted. See
+  // go/libjpeg-android.
+  unsigned char* buffer_array[1];
+  buffer_array[0] = decode_buffer.data();
+  size_t decoded_offset = 0;
+  while (height--) {
+    // According to the documentation, jpeg_read_scanlines returns the number
+    // of lines read
+    // https://android.googlesource.com/platform/external/jpeg/+/c6859b743e7248b9f401264aac939a5af0d63799/libjpeg.doc#655
+    // In case of premature ending of the image, the implementation of
+    // jpeg_read_scanlines in the version of JPEG Turbo we are using to test
+    // emits a warning ("Corrupt JPEG data: premature end of data segment")
+    // but doesn't fail and consider the line as successfully read.
+    // See test
+    // LibjpegDecoderTest::DoesNotFailDecodingAnImageWithLessDataThanDeclaredInJpegHeader
+    unsigned int num_of_scanlines_read = 0;
+    if (impl.jpeg_read_scanlines(num_of_scanlines_read, buffer_array, 1) !=
+        kTfLiteOk) {
+      return impl.status();
+    }
+
+    if (num_of_scanlines_read != 1) {
+      return {kTfLiteError, "Expected " + std::to_string(header.height) +
+                                " lines but found only " +
+                                std::to_string(header.height - height) +
+                                " read scanlines is " +
+                                std::to_string(num_of_scanlines_read)};
+    }
+
+    std::copy_n(buffer_array[0], row_stride, decoded + decoded_offset);
+
+    decoded_offset += row_stride;
+  }
+  boolean finish_decompress_result = false;
+  if (impl.jpeg_finish_decompress(finish_decompress_result) != kTfLiteOk) {
+    return impl.status();
+  }
+  if (!finish_decompress_result) {
+    return {kTfLiteError, "Failed call jpeg_finish_decompress_"};
+  }
+  if (impl.jpeg_destroy_decompress() != kTfLiteOk) {
+    return impl.status();
+  }
+  return impl.status();
+}
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder.h

@@ -0,0 +1,221 @@
+/* Copyright 2021 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_LIBJPEG_DECODER_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_LIBJPEG_DECODER_H_
+
+#include <memory.h>
+
+#include <csetjmp>
+#include <cstddef>
+#include <memory>
+#include <string>
+#include <type_traits>
+#include <utility>
+
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg_status.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_common.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_decompress_buffered_struct.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/libc_handle.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_handle.h"
+#include "tensorflow/lite/string_type.h"
+#include "tensorflow/lite/string_util.h"
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+
+// Extracts the expected size of `jpeg_decompress_struct` from the "struct
+// mismatch" error message and stores it in `expected_size`. Returns status code
+// kTfLiteOk if the extraction was successful, error otherwise.
+Status ExtractSizeFromErrorMessage(const std::string& error_message,
+                                   size_t& expected_size);
+
+class LibjpegDecoder {
+ public:
+  // The maximum height allowed for the decoded image. Any attempt to call
+  // DecodeImage for an image with height or width over the allowed limits will
+  // fail.
+  // The size is define to 10,000 lines.
+  static const size_t kMaxImageHeight;
+  // The maximum width allowed for the decoded image. Any attempt to call
+  // DecodeImage for an image with height or width over the allowed limits will
+  // fail.
+  // The size is define to 10,000 pixels per line.
+  static const size_t kMaxImageWidth;
+
+  // Creates and initialises the decoder.
+  // Dynamically loads libjpeg (into handle_) and sets the expected size for
+  // `jpeg_decompress_struct` (in expected_size_for_decompress_struct_). Returns
+  // an initalised instance of decoder if successful, else returns nullptr.
+  // Stores initialisation status in status.
+  static std::unique_ptr<LibjpegDecoder> Create(Status& status);
+  Status DecodeImage(const tflite::StringRef& encoded,
+                     const JpegHeader& expected_image_dimensions,
+                     unsigned char* decoded, const size_t& decoded_size) const;
+
+ private:
+  explicit LibjpegDecoder(LibCHandle libc_handle)
+      : libc_handle_(std::move(libc_handle)) {}
+  // Wraps all objects required for using the libjpeg library.
+  // This is to avoid stack-allocating these variables in the function that
+  // calls setjmp().
+  class Impl {
+   public:
+    explicit Impl(size_t decompress_struct_size, const LibjpegHandle* handle);
+    ~Impl() { jpeg_destroy_decompress(); }
+    Impl(const Impl&) = delete;
+    Impl& operator=(const Impl&) = delete;
+    Impl(Impl&& other) = delete;
+    Impl& operator=(Impl&& other) = delete;
+
+    // Wrapping calls to LibjpegHandle functions in Run and RunAndSetStatus.
+    TfLiteStatus jpeg_CreateDecompress(int version, size_t struct_size) {
+      // Note: It is safe to call jpeg_destroy_decompress even if the
+      // corresponding call to create_jpeg_decompress fails. See the end of
+      // section "Compression details" in
+      // https://www.freedesktop.org/wiki/Software/libjpeg/.
+      safe_to_invoke_destroy_decompress_ = true;
+      return Run(&LibjpegHandle::jpeg_create_decompress_, version, struct_size);
+    }
+    TfLiteStatus jpeg_stdio_src(FILE* infile) {
+      return Run(&LibjpegHandle::jpeg_stdio_src_, infile);
+    }
+
+    TfLiteStatus jpeg_read_header(int& read_header_result,
+                                  boolean require_image) {
+      return RunAndSetResult(&LibjpegHandle::jpeg_read_header_,
+                             &read_header_result, require_image);
+    }
+
+    TfLiteStatus jpeg_start_decompress(boolean& start_decompress_result) {
+      return RunAndSetResult(&LibjpegHandle::jpeg_start_decompress_,
+                             &start_decompress_result);
+    }
+    TfLiteStatus jpeg_read_scanlines(unsigned int& read_scanlines_result,
+                                     JSAMPARRAY scanlines,
+                                     JDIMENSION max_lines) {
+      return RunAndSetResult(&LibjpegHandle::jpeg_read_scanlines_,
+                             &read_scanlines_result, scanlines, max_lines);
+    }
+    TfLiteStatus jpeg_finish_decompress(boolean& finish_decompress_result) {
+      return RunAndSetResult(&LibjpegHandle::jpeg_finish_decompress_,
+                             &finish_decompress_result);
+    }
+    TfLiteStatus jpeg_destroy_decompress() {
+      if (safe_to_invoke_destroy_decompress_) {
+        safe_to_invoke_destroy_decompress_ = false;
+        return Run(&LibjpegHandle::jpeg_destroy_decompress_);
+      }
+      return kTfLiteOk;
+    }
+
+    // Status from the libjpeg layer that is to be returned to the caller.
+    Status status() { return status_; }
+
+   private:
+    // Delegates to one of the LibjpegHandle::jpeg_* methods.
+    // This is to restrict the call to setjmp() to a stack frame free from
+    // stack allocated C++ variables. The type of f is T
+    // (LibjpegHandle::*f)(Args...), for some T. All args must be
+    // pointers/references/primitive types. Since we use a
+    // non-suspending JPEG encoded data source, return value from a
+    // LibjpegHandle::jpeg_* methods is not required by client and hence
+    // discarded. Returns an Ok status if the execution was successful, error
+    // otherwise.
+    template <typename Fn, typename... Args>
+    TfLiteStatus Run(Fn f, Args... args) {
+      // Note(1): C++ variables local to this function should not be stack
+      // allocated
+      // and should be passed as pointers or references. Using setjmp/longjmp
+      // with stack allocated C++ objects that have non-trivial destructors can
+      // lead to undefined behaviour.
+      // https://wiki.sei.cmu.edu/confluence/pages/viewpage.action?pageId=88046492
+      // All such variables whose scope contains calls to libjpeg (that may do a
+      // longjmp) should be passed in as arguments.
+      //
+      // Note(2): All other variables local to this function that need to be
+      // accessed after longjmp() returns control to this function, should be
+      // volatile-qualified.
+      // After invoking longjump(), non-volatile local variables should not be
+      // accessed for two reasons:
+      // - their values may be indeterminate. According to the C standard, if
+      // the variable's value has changed between setjmp() and longjmp(), their
+      // value is considered indeterminate, and accessing them is undefined
+      // behaviour.
+      // https://wiki.sei.cmu.edu/confluence/display/c/MSC22-C.+Use+the+setjmp%28%29%2C+longjmp%28%29+facility+securely
+      // - the register storing such variables might be clobbered. Even if the
+      // variable remains unchanged between setjmp() and longjmp(), the stack
+      // slot for the variable may get incorrectly clobbered. This is a known
+      // LLVM bug: https://bugs.llvm.org/show_bug.cgi?id=21183
+      if (setjmp(env_)) return kTfLiteError;
+      (handle_->*f)(cinfo_.get(), args...);
+      return kTfLiteOk;
+    }
+    // Extension of the Run method for non-void JPEG calls when we need to
+    // collect the returned value.
+    // See Run comments above for details.
+    template <
+        typename Fn, typename... Args,
+        typename ResultType = typename std::result_of_t<Fn>,
+        typename = typename std::enable_if<!std::is_void<ResultType>::value> >
+    TfLiteStatus RunAndSetResult(Fn f, ResultType* result, Args... args) {
+      if (setjmp(env_)) return kTfLiteError;
+      *result = (handle_->*f)(cinfo_.get(), args...);
+      return kTfLiteOk;
+    }
+    // Size of `jpeg_decompress_struct` as expected by libjpeg library.
+    size_t decompress_struct_size_;
+    const LibjpegHandle* handle_;
+    // Using a buffered struct for `jpeg_decompress_struct` as the size expected
+    // by libjpeg can be different from the size of the compiled struct. See
+    // go/libjpeg-android. Note: Since we resize the struct, accessing some of
+    // the fields of this struct may lead to undefined behaviour. For
+    // decompression, only the fields within `jpeg_common_fields` are required
+    // viz. error manager(`err`) and client data(`client_data`). This code
+    // limits its usage to these two fields and we recommend future contributors
+    // to not access fields beyond `jpeg_common_fields`.
+    JpegDecompressBufferedStruct cinfo_;
+    struct jpeg_error_mgr jerr_;
+    // Stores the information of the calling environment which can be restored
+    // later. Libjpeg aborts the program in case of any errors by using longjmp
+    // and then calling exit(). The only way to avoid this, is to transfer the
+    // control flow to the caller by using setjmp/longjmp.
+    jmp_buf env_;
+    static void ErrorExit(j_common_ptr cinfo);
+    // Calls to jpeg_create_decompress and jpeg_destroy_decompress need to be
+    // paired. This flag indicates if it's safe to invoke
+    // jpeg_destroy_decompress.
+    bool safe_to_invoke_destroy_decompress_ = false;
+    // Status of the most recent execution of a LibjpegHandle::jpeg_* method
+    // invoked using Run or RunAndSetResult.
+    Status status_;
+  };
+  // Size of `jpeg_decompress_struct` as expected by the libjpeg dynamic
+  // library. The expected size is different from the size of the compiled
+  // struct on some Android Devices. See go/libjpeg-android.
+  size_t expected_size_for_decompress_struct_;
+  // Handle to the Libjpeg dynamic library.
+  std::unique_ptr<LibjpegHandle> libjpeg_handle_;
+  // Handle to the LibC dynamic library.
+  LibCHandle libc_handle_;
+};
+
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_LIBJPEG_DECODER_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder_test.cc

@@ -0,0 +1,344 @@
+/* Copyright 2021 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder.h"
+
+#include <stddef.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg_status.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/embedded_chessboard_jpeg.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/embedded_snow_jpeg.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/embedded_test_card_jpeg.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_header_parser.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder_test_helper.h"
+#include "tensorflow/lite/minimal_logging.h"
+#include "tensorflow/lite/string_util.h"
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+namespace {
+
+using testing::IsEmpty;
+using testing::NotNull;
+
+constexpr JpegHeader kExpectedImageDimensions{
+    .height = 300, .width = 250, .channels = 3};
+
+constexpr int kDecodedSize = kExpectedImageDimensions.height *
+                             kExpectedImageDimensions.width *
+                             kExpectedImageDimensions.channels;
+
+TEST(LibjpegDecoderTest, InitShouldSucceedOnAndroid) {
+  Status status;
+  std::unique_ptr<LibjpegDecoder> decoder = LibjpegDecoder::Create(status);
+  EXPECT_EQ(status.code, kTfLiteOk);
+  EXPECT_THAT(status.error_message, IsEmpty());
+}
+
+TEST(LibjpegDecoderTest, DecodingChessboardShouldSucceedOnAndroid) {
+  Status status;
+  std::unique_ptr<LibjpegDecoder> decoder = LibjpegDecoder::Create(status);
+  ASSERT_THAT(decoder, NotNull());
+  tflite::StringRef string_ref = {
+      reinterpret_cast<const char*>(g_tflite_acceleration_chessboard_jpeg),
+      static_cast<size_t>(g_tflite_acceleration_chessboard_jpeg_len)};
+  unsigned char decoded[kDecodedSize];
+
+  status = decoder->DecodeImage(string_ref, kExpectedImageDimensions, decoded,
+                                kDecodedSize);
+
+  ASSERT_EQ(status.error_message, "");
+  ASSERT_EQ(status.code, kTfLiteOk);
+  std::vector<uint8_t> decoded_vec(decoded, decoded + kDecodedSize);
+  EXPECT_THAT(decoded_vec, HasChessboardPatternWithTolerance(12));
+}
+
+TEST(LibjpegDecoderTest, DecodingRainbowTestCardShouldSucceedOnAndroid) {
+  Status status;
+  std::unique_ptr<LibjpegDecoder> decoder = LibjpegDecoder::Create(status);
+  EXPECT_THAT(status.error_message, IsEmpty());
+  EXPECT_EQ(status.code, kTfLiteOk);
+  ASSERT_THAT(decoder, NotNull());
+  std::string encoded(
+      reinterpret_cast<const char*>(g_tflite_acceleration_test_card_jpeg),
+      g_tflite_acceleration_test_card_jpeg_len);
+  tflite::StringRef string_ref = {encoded.c_str(), encoded.length()};
+  unsigned char decoded[kDecodedSize];
+
+  status = decoder->DecodeImage(string_ref, kExpectedImageDimensions, decoded,
+                                kDecodedSize);
+
+  ASSERT_EQ(status.error_message, "");
+  ASSERT_EQ(status.code, kTfLiteOk);
+  std::vector<uint8_t> decoded_vec(decoded, decoded + kDecodedSize);
+  EXPECT_THAT(decoded_vec, HasRainbowPatternWithTolerance(5));
+}
+
+TEST(LibjpegDecoderTest, ErrorsFromJpegLayerAreReturnedToCaller) {
+  Status status;
+  std::unique_ptr<LibjpegDecoder> decoder = LibjpegDecoder::Create(status);
+  EXPECT_THAT(status.error_message, IsEmpty());
+  EXPECT_EQ(status.code, kTfLiteOk);
+  ASSERT_THAT(decoder, NotNull());
+  std::string str = "this is not a jpeg image";
+  tflite::StringRef encoded = {str.c_str(), str.length()};
+  unsigned char decoded_image[12];
+
+  status = decoder->DecodeImage(encoded, kExpectedImageDimensions,
+                                decoded_image, 12);
+
+  EXPECT_EQ(status.code, kTfLiteError);
+  EXPECT_EQ(status.error_message, "Not a valid JPEG image.");
+}
+
+TEST(LibjpegDecoderTest, DecodingFailsWhenDecodeBufferIsSmall) {
+  Status status;
+  std::unique_ptr<LibjpegDecoder> decoder = LibjpegDecoder::Create(status);
+  EXPECT_THAT(status.error_message, IsEmpty());
+  EXPECT_EQ(status.code, kTfLiteOk);
+  ASSERT_THAT(decoder, NotNull());
+  std::string encoded(
+      reinterpret_cast<const char*>(g_tflite_acceleration_test_card_jpeg),
+      g_tflite_acceleration_test_card_jpeg_len);
+  tflite::StringRef string_ref = {encoded.c_str(), encoded.length()};
+  const int decoded_size = 100;
+  unsigned char decoded[decoded_size];
+  status = decoder->DecodeImage(string_ref, kExpectedImageDimensions, decoded,
+                                decoded_size);
+  EXPECT_EQ(status.code, kTfLiteError);
+  EXPECT_EQ(status.error_message,
+            "Size of buffer(100) for storing decoded image must be equal to "
+            "the size of decoded image(225000).");
+}
+
+TEST(LibjpegDecoderTest, DecodingFailsWhenImageDimensionsDifferFromExpected) {
+  Status status;
+  std::unique_ptr<LibjpegDecoder> decoder = LibjpegDecoder::Create(status);
+  EXPECT_THAT(status.error_message, IsEmpty());
+  EXPECT_EQ(status.code, kTfLiteOk);
+  ASSERT_THAT(decoder, NotNull());
+  std::string encoded(
+      reinterpret_cast<const char*>(g_tflite_acceleration_test_card_jpeg),
+      g_tflite_acceleration_test_card_jpeg_len);
+  tflite::StringRef string_ref = {encoded.c_str(), encoded.length()};
+  unsigned char decoded[kDecodedSize];
+
+  status = decoder->DecodeImage(string_ref,
+                                {.height = 300, .width = 250, .channels = 1},
+                                decoded, kDecodedSize);
+  EXPECT_EQ(status.code, kTfLiteError);
+  EXPECT_EQ(
+      status.error_message,
+      "Decoded image size (300, 250, 3, 8) is different from provided image "
+      "size (300, 250, 1, 8)");
+}
+
+TEST(LibjpegDecoderTest, DecodingFailsWhenImageDimensionsAreOverThreshold) {
+  Status status;
+  std::unique_ptr<LibjpegDecoder> decoder = LibjpegDecoder::Create(status);
+  EXPECT_THAT(status.error_message, IsEmpty());
+  EXPECT_EQ(status.code, kTfLiteOk);
+  ASSERT_THAT(decoder, NotNull());
+  std::string encoded(
+      reinterpret_cast<const char*>(g_tflite_acceleration_test_card_jpeg),
+      g_tflite_acceleration_test_card_jpeg_len);
+  tflite::StringRef origin_string_ref = {encoded.c_str(), encoded.length()};
+
+  const JpegHeader kHeader{
+      .height = static_cast<int>(LibjpegDecoder::kMaxImageHeight + 1),
+      .width = static_cast<int>(LibjpegDecoder::kMaxImageWidth + 1),
+      .channels = 3};
+
+  const size_t decoded_size = static_cast<size_t>(kHeader.height) *
+                              static_cast<size_t>(kHeader.width) *
+                              static_cast<size_t>(kHeader.channels);
+
+  std::string altered_image;
+  Status alter_header_status =
+      BuildImageWithNewHeader(origin_string_ref, kHeader, altered_image);
+  ASSERT_EQ(alter_header_status.code, kTfLiteOk);
+
+  tflite::StringRef altered_string_ref = {altered_image.c_str(),
+                                          altered_image.length()};
+
+  std::vector<unsigned char> decoded(decoded_size);
+  status = decoder->DecodeImage(altered_string_ref, kHeader, decoded.data(),
+                                decoded_size);
+  EXPECT_EQ(status.code, kTfLiteError);
+  EXPECT_EQ(status.error_message,
+            "Image is too big, dimensions (" + std::to_string(kHeader.width) +
+                "," + std::to_string(kHeader.width) +
+                ") larger than the maximum allowed (" +
+                std::to_string(LibjpegDecoder::kMaxImageHeight) + ", " +
+                std::to_string(LibjpegDecoder::kMaxImageWidth) + ")");
+}
+
+TEST(LibjpegDecoderTest, DecodingFailsWhenImageHasUnsupportedNumberOfChannels) {
+  Status status;
+  std::unique_ptr<LibjpegDecoder> decoder = LibjpegDecoder::Create(status);
+  EXPECT_THAT(status.error_message, IsEmpty());
+  EXPECT_EQ(status.code, kTfLiteOk);
+  ASSERT_THAT(decoder, NotNull());
+  std::string encoded(
+      reinterpret_cast<const char*>(g_tflite_acceleration_test_card_jpeg),
+      g_tflite_acceleration_test_card_jpeg_len);
+  tflite::StringRef string_ref = {encoded.c_str(), encoded.length()};
+  unsigned char decoded[300 * 250 * 4];
+
+  const JpegHeader kHeader{.height = 300, .width = 250, .channels = 4};
+  status = decoder->DecodeImage(string_ref, kHeader, decoded, kDecodedSize);
+  EXPECT_EQ(status.code, kTfLiteError);
+  EXPECT_EQ(status.error_message,
+            "Supporting only images with 1 or 3 channels");
+}
+
+TEST(LibjpegDecoderTest, DecodingFailsWhenExpectedBitPerSampleIsNot8) {
+  Status status;
+  std::unique_ptr<LibjpegDecoder> decoder = LibjpegDecoder::Create(status);
+  EXPECT_THAT(status.error_message, IsEmpty());
+  EXPECT_EQ(status.code, kTfLiteOk);
+  ASSERT_THAT(decoder, NotNull());
+  std::string encoded(
+      reinterpret_cast<const char*>(g_tflite_acceleration_test_card_jpeg),
+      g_tflite_acceleration_test_card_jpeg_len);
+  tflite::StringRef string_ref = {encoded.c_str(), encoded.length()};
+  unsigned char decoded[kDecodedSize];
+
+  status = decoder->DecodeImage(
+      string_ref,
+      {.height = 300, .width = 250, .channels = 3, .bits_per_sample = 4},
+      decoded, kDecodedSize);
+  EXPECT_EQ(status.code, kTfLiteError);
+  EXPECT_EQ(status.error_message,
+            "Supporting only images with 8 bits per sample");
+}
+
+TEST(LibjpegDecoderTest, DoesNotDecodeBeyondWhatIsSpecifiedInHeader) {
+  Status status;
+  std::unique_ptr<LibjpegDecoder> decoder = LibjpegDecoder::Create(status);
+  EXPECT_THAT(status.error_message, IsEmpty());
+  EXPECT_EQ(status.code, kTfLiteOk);
+  ASSERT_THAT(decoder, NotNull());
+  std::string origin_encoded_img(
+      reinterpret_cast<const char*>(g_tflite_acceleration_test_card_jpeg),
+      g_tflite_acceleration_test_card_jpeg_len);
+  tflite::StringRef origin_string_ref = {origin_encoded_img.c_str(),
+                                         origin_encoded_img.length()};
+
+  JpegHeader undersized_image_header = {
+      .height = kExpectedImageDimensions.height / 2,
+      .width = kExpectedImageDimensions.width / 2,
+      .channels = kExpectedImageDimensions.channels};
+  std::string altered_image;
+  Status alter_header_status = BuildImageWithNewHeader(
+      origin_string_ref, undersized_image_header, altered_image);
+  ASSERT_EQ(alter_header_status.code, kTfLiteOk);
+
+  tflite::StringRef altered_string_ref{altered_image.c_str(),
+                                       altered_image.length()};
+
+  unsigned char decoded[kDecodedSize / 4];
+
+  status = decoder->DecodeImage(altered_string_ref, undersized_image_header,
+                                decoded, kDecodedSize / 4);
+  EXPECT_EQ(status.code, kTfLiteOk);
+}
+
+TEST(LibjpegDecoderTest, CanReadImagesWithVeryLargeRows) {
+  Status status;
+  std::unique_ptr<LibjpegDecoder> decoder = LibjpegDecoder::Create(status);
+  EXPECT_THAT(status.error_message, IsEmpty());
+  EXPECT_EQ(status.code, kTfLiteOk);
+  ASSERT_THAT(decoder, NotNull());
+  std::string origin_encoded_img(
+      reinterpret_cast<const char*>(g_tflite_acceleration_snow_jpeg),
+      g_tflite_acceleration_snow_jpeg_len);
+  tflite::StringRef origin_string_ref = {origin_encoded_img.c_str(),
+                                         origin_encoded_img.length()};
+
+  JpegHeader one_long_row_image_header = {
+      .height = 1,
+      .width = static_cast<int>(LibjpegDecoder::kMaxImageWidth),
+      .channels = kExpectedImageDimensions.channels};
+  std::string altered_image;
+  Status alter_header_status = BuildImageWithNewHeader(
+      origin_string_ref, one_long_row_image_header, altered_image);
+  ASSERT_EQ(alter_header_status.code, kTfLiteOk);
+
+  tflite::StringRef altered_string_ref = {altered_image.c_str(),
+                                          altered_image.length()};
+
+  const size_t kImageSize = LibjpegDecoder::kMaxImageWidth * 3;
+  std::vector<unsigned char> decoded(kImageSize);
+
+  status = decoder->DecodeImage(altered_string_ref, one_long_row_image_header,
+                                decoded.data(), kImageSize);
+  EXPECT_EQ(status.code, kTfLiteOk);
+}
+
+TEST(LibjpegDecoderTest, FailDecodingAnImageWithUnexpectedEofInDataStream) {
+  Status status;
+  std::unique_ptr<LibjpegDecoder> decoder = LibjpegDecoder::Create(status);
+  EXPECT_THAT(status.error_message, IsEmpty());
+  EXPECT_EQ(status.code, kTfLiteOk);
+  ASSERT_THAT(decoder, NotNull());
+  std::string img(
+      reinterpret_cast<const char*>(g_tflite_acceleration_test_card_jpeg),
+      g_tflite_acceleration_test_card_jpeg_len);
+  tflite::StringRef truncated_image_ref = {img.c_str(), img.length() - 100};
+
+  unsigned char decoded[kDecodedSize];
+
+  status = decoder->DecodeImage(truncated_image_ref, kExpectedImageDimensions,
+                                decoded, kDecodedSize);
+  EXPECT_EQ(status.code, kTfLiteError);
+  // The image is not even read because is recognised as invalid JPEG image
+  EXPECT_EQ(status.error_message, "Not a valid JPEG image.");
+}
+
+TEST(LibjpegDecoderTest, JpegErrorMsgParsingForValidMsg) {
+  size_t extracted_size;
+  Status status = ExtractSizeFromErrorMessage(
+      "JPEG parameter struct mismatch: library thinks size is 480, caller "
+      "expects 464.",
+      extracted_size);
+  ASSERT_EQ(status.code, kTfLiteOk);
+  EXPECT_EQ(extracted_size, 480);
+}
+
+TEST(LibjpegDecoderTest, JpegErrorMsgParsingForMaformedMsg) {
+  size_t extracted_size;
+  std::string err_msg =
+      "JPEG parameter struct mismatch: library thinks size is abcde, caller "
+      "expects 464.";
+  Status status = ExtractSizeFromErrorMessage(err_msg, extracted_size);
+  EXPECT_EQ(status.code, kTfLiteError);
+  EXPECT_EQ(status.error_message,
+            "Couldn't parse the size from message: '" + err_msg + "'");
+}
+
+}  // namespace
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder_test_helper.h

@@ -0,0 +1,124 @@
+/* Copyright 2021 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_LIBJPEG_DECODER_TEST_HELPER_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_LIBJPEG_DECODER_TEST_HELPER_H_
+
+#include <cmath>
+#include <cstdint>
+#include <fstream>
+#include <sstream>
+#include <string>
+#include <tuple>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/strings/string_view.h"
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+
+// Checks if the values are almost equal and their absolute difference doesn't
+// exceed tolerance.
+// Params: int tolerance. arg is int.
+MATCHER_P(AreAlmostEqualWithTolerance, tolerance, "") {
+  int a = static_cast<int>(testing::get<0>(arg));
+  int b = static_cast<int>(testing::get<1>(arg));
+  int diff = std::abs(b - a);
+  if (result_listener) {
+    std::ostringstream os;
+    os << "difference between " << a << " and " << b << " is " << diff
+       << " which is greater than " << tolerance << ". ";
+    *result_listener << os.str();
+  }
+  return diff <= tolerance;
+}
+
+// Matcher for matching the 13x13 yellow-white chessboard pattern in an image.
+// Parameters: int tolerance. arg is vector<uint8_t> image.
+// Checks that relative difference between pixel values is within tolerance.
+MATCHER_P(HasChessboardPatternWithTolerance, tolerance, "") {
+  const std::vector<uint8_t> image = static_cast<std::vector<uint8_t>>(arg);
+  const std::vector<uint8_t> kYellow = {253, 242, 0};
+  const std::vector<uint8_t> kWhite = {255, 255, 255};
+  const int kHeightRect = 23;
+  const int kWidthRect = 19;
+  const int kChannels = 3;
+  // Rectangles on the chessboard have fixed height and width.
+  // Check color at centre pixel for every rectangle.
+  // There are 13x13 rectangles in all.
+  const int row_stride = kChannels * 250;
+  int rect = 0;
+
+  for (int i = 0; i < 13; i++)
+    for (int j = 0; j < 13; j++) {
+      int row = i * kHeightRect + (kHeightRect / 2);
+      int col = j * kWidthRect + (kWidthRect / 2);
+      int pixel = row * row_stride + col * kChannels;
+      std::vector<uint8_t> decoded_color = {image[pixel], image[pixel + 1],
+                                            image[pixel + 2]};
+      if (!testing::ExplainMatchResult(
+              testing::Pointwise(AreAlmostEqualWithTolerance(tolerance),
+                                 (rect & 1) ? kWhite : kYellow),
+              decoded_color, result_listener)) {
+        *result_listener << "Pixel values at row#" << row << ", col#" << col
+                         << " don't match.";
+
+        return false;
+      }
+      rect++;
+    }
+  return true;
+}
+
+// Matcher for matching the
+// SMPTE(https://en.wikipedia.org/wiki/SMPTE_color_bars) rainbow color bars
+// pattern in an image. Parameters: int tolerance. arg is vector<uint8_t> image.
+// Verifies that relative difference between pixel values is within tolerance.
+MATCHER_P(HasRainbowPatternWithTolerance, tolerance, "") {
+  const std::vector<uint8_t> image = static_cast<std::vector<uint8_t>>(arg);
+
+  const int kWidth = 250;
+  const int kRow = 150;
+  const int kChannels = 3;
+  const int kBandSize = 35;
+  std::vector<std::vector<uint8_t>> colors = {
+      {192, 192, 192}, {192, 192, 0}, {0, 193, 192}, {0, 192, 0},
+      {192, 0, 192},   {193, 0, 0},   {0, 0, 193}};
+  // Match pixel colors at 7 locations in a fixed row.
+  for (int i = 0; i < 7; ++i) {
+    int col = (i * kBandSize + kBandSize / 2);
+    int pixel = kRow * kWidth * kChannels + col * kChannels;
+    std::vector<uint8_t> decoded_color = {image[pixel], image[pixel + 1],
+                                          image[pixel + 2]};
+    if (!testing::ExplainMatchResult(
+            testing::Pointwise(AreAlmostEqualWithTolerance(tolerance),
+                               colors[i]),
+            decoded_color, result_listener)) {
+      *result_listener << "Pixel values at row#" << kRow << ", col#" << col
+                       << " don't match.";
+
+      return false;
+    }
+  }
+  return true;
+}
+
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_LIBJPEG_DECODER_TEST_HELPER_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_handle.h

@@ -0,0 +1,73 @@
+/* Copyright 2021 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_LIBJPEG_HANDLE_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_LIBJPEG_HANDLE_H_
+
+#include <stddef.h>
+#include <stdio.h>
+
+#include <memory>
+#include <string>
+
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg_status.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg.h"
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+
+// This class offers a handle to Libjpeg shared object library on Android.
+// It offers pointers to functions in Libjpeg that are required for decoding
+// JPEG images.
+// TODO(b/172544567): Support Apple.
+class LibjpegHandle {
+ public:
+  // Factory for creating an initialised instance of LibjpegHandle.
+  // Loads the libjpeg dynamic library and gets handle to all the functions
+  // required for decompressing JPEGs. Returns an initialised instance of
+  // LibjpegHandle if successful, else nullptr. Stores initialisation status in
+  // `status`.
+  static std::unique_ptr<LibjpegHandle> Create(Status& status);
+  // Closes the dynamic library loaded in libjpeg_.
+  ~LibjpegHandle();
+  LibjpegHandle(LibjpegHandle const&) = delete;
+  LibjpegHandle& operator=(const LibjpegHandle&) = delete;
+  LibjpegHandle(LibjpegHandle&& LibjpegHandle) = delete;
+  LibjpegHandle& operator=(LibjpegHandle&& other) = delete;
+  // Based on our analysis of Android devices in the ODML lab, it is reasonable
+  // to expect 62 (6b) as the version of libjpeg on all Android devices from SDK
+  // 22 onwards.
+  static const int kLibjpegVersion = 62;
+  // Definitions of the functions below can be found in
+  // third_party/libjpeg_turbo/src/jpeglib.h
+  struct jpeg_error_mgr* (*jpeg_std_error_)(struct jpeg_error_mgr*);
+  void (*jpeg_destroy_decompress_)(j_decompress_ptr);
+  void (*jpeg_create_decompress_)(j_decompress_ptr, int, size_t);
+  void (*jpeg_stdio_src_)(j_decompress_ptr, FILE*);
+  int (*jpeg_read_header_)(j_decompress_ptr, boolean);
+  boolean (*jpeg_start_decompress_)(j_decompress_ptr);
+  unsigned int (*jpeg_read_scanlines_)(j_decompress_ptr, JSAMPARRAY,
+                                       JDIMENSION);
+  boolean (*jpeg_finish_decompress_)(j_decompress_ptr);
+
+ private:
+  LibjpegHandle() {}
+  void* libjpeg_ = nullptr;
+};
+
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_LIBJPEG_HANDLE_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_handle_dynamic_link.cc

@@ -0,0 +1,78 @@
+/* Copyright 2021 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 <dlfcn.h>
+#include <stddef.h>
+#include <stdio.h>
+
+#include <memory>
+#include <type_traits>
+
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg_status.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_handle.h"
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+
+std::unique_ptr<LibjpegHandle> LibjpegHandle::Create(Status &status) {
+  std::unique_ptr<LibjpegHandle> handle(new LibjpegHandle());
+
+  if (!(handle->libjpeg_ = dlopen("libjpeg.so", RTLD_NOW | RTLD_LOCAL))) {
+    status = {kTfLiteError, "Failed to load dynamic library."};
+    return nullptr;
+  }
+
+  // On Android S, the system libjpeg symbols have been prefixed
+  // with 'chromium_' to avoid collisions.
+#define LOAD(variable, symbol_name)                                      \
+  do {                                                                   \
+    static_assert(                                                       \
+        std::is_same<decltype(variable), decltype(&symbol_name)>::value, \
+        "Mismatched types");                                             \
+    void *symbol = dlsym(handle->libjpeg_, #symbol_name);                \
+    if (!symbol) {                                                       \
+      symbol = dlsym(handle->libjpeg_, "chromium_" #symbol_name);        \
+    }                                                                    \
+    if (!symbol) {                                                       \
+      status = {kTfLiteError,                                            \
+                "Failed to dynamically load the method: " #symbol_name}; \
+      return nullptr;                                                    \
+    }                                                                    \
+    variable = reinterpret_cast<decltype(variable)>(symbol);             \
+  } while (0)
+
+  LOAD(handle->jpeg_std_error_, jpeg_std_error);
+  LOAD(handle->jpeg_destroy_decompress_, jpeg_destroy_decompress);
+  LOAD(handle->jpeg_create_decompress_, jpeg_CreateDecompress);
+  LOAD(handle->jpeg_stdio_src_, jpeg_stdio_src);
+  LOAD(handle->jpeg_read_header_, jpeg_read_header);
+  LOAD(handle->jpeg_start_decompress_, jpeg_start_decompress);
+  LOAD(handle->jpeg_read_scanlines_, jpeg_read_scanlines);
+  LOAD(handle->jpeg_finish_decompress_, jpeg_finish_decompress);
+#undef LOAD
+  status = {kTfLiteOk, ""};
+  return handle;
+}
+
+LibjpegHandle::~LibjpegHandle() {
+  if (libjpeg_) {
+    dlclose(libjpeg_);
+  }
+}
+
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_handle_static_link.cc

@@ -0,0 +1,46 @@
+/* Copyright 2021 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_handle.h"
+
+#include <memory>
+
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg_status.h"
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+
+std::unique_ptr<LibjpegHandle> LibjpegHandle::Create(Status &status) {
+  std::unique_ptr<LibjpegHandle> handle(new LibjpegHandle());
+
+  // Use statically linked implementation unless otherwise configured.
+  handle->jpeg_std_error_ = jpeg_std_error;
+  handle->jpeg_destroy_decompress_ = jpeg_destroy_decompress;
+  handle->jpeg_create_decompress_ = jpeg_CreateDecompress;
+  handle->jpeg_stdio_src_ = jpeg_stdio_src;
+  handle->jpeg_read_header_ = jpeg_read_header;
+  handle->jpeg_start_decompress_ = jpeg_start_decompress;
+  handle->jpeg_read_scanlines_ = jpeg_read_scanlines;
+  handle->jpeg_finish_decompress_ = jpeg_finish_decompress;
+  status = {kTfLiteOk, ""};
+  return handle;
+}
+
+LibjpegHandle::~LibjpegHandle() = default;
+
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_handle_test.cc

@@ -0,0 +1,39 @@
+/* Copyright 2021 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_handle.h"
+
+#include <memory>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg_status.h"
+
+namespace tflite {
+namespace acceleration {
+namespace decode_jpeg_kernel {
+namespace {
+
+TEST(LibjpegHandleTest, LoadingSucceeds) {
+  Status status;
+  std::unique_ptr<LibjpegHandle> handle = LibjpegHandle::Create(status);
+  EXPECT_TRUE(handle != nullptr);
+  EXPECT_EQ(status.error_message, "");
+  EXPECT_EQ(status.code, kTfLiteOk);
+}
+
+}  // namespace
+}  // namespace decode_jpeg_kernel
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/metrics/BUILD

@@ -0,0 +1,128 @@
+# Copyright 2021 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.
+# ==============================================================================
+
+load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow/core/platform:build_config_root.bzl", "if_pywrap")
+
+# Building blocks for generating accuracy validation metric graphs.
+load("//tensorflow/lite/experimental/acceleration/mini_benchmark:build_defs.bzl", "validation_model", "validation_test")
+load("//tensorflow/lite/experimental/acceleration/mini_benchmark:special_rules.bzl", "minibenchmark_visibility_allowlist")
+
+default_visibility_group = [
+    "//tensorflow/lite/experimental/acceleration/mini_benchmark:__subpackages__",
+    "//tensorflow/lite/tools/benchmark:__subpackages__",
+] + minibenchmark_visibility_allowlist()
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = default_visibility_group,
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "kl_divergence",
+    srcs = [
+        "kl_divergence.py",
+    ],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow:tensorflow_py",
+    ],
+)
+
+py_binary(
+    name = "mobilenet",
+    srcs = ["mobilenet.py"],
+    strict_deps = True,
+    deps = [
+        ":kl_divergence",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/lite/python:lite",
+        "//tensorflow/lite/tools:flatbuffer_utils",
+    ] + if_pywrap(
+        if_true = ["//tensorflow/python:_pywrap_tensorflow"],
+    ),
+)
+
+genrule(
+    name = "mobilenet_metrics_tflite",
+    outs = ["mobilenet_metrics.tflite"],
+    cmd = "$(location :mobilenet) $(location :mobilenet_metrics.tflite)",
+    tools = [":mobilenet"],
+)
+
+validation_model(
+    name = "mobilenet_quant_with_validation",
+    testonly = 1,
+    jpegs = "//tensorflow/lite/experimental/acceleration/mini_benchmark:odt_classifier_testfiles",
+    main_model = "//tensorflow/lite/experimental/acceleration/mini_benchmark/models:mobilenet_v1_1.0_224_quant.tflite",
+    metrics_model = ":mobilenet_metrics_tflite",
+)
+
+validation_test(
+    name = "mobilenet_quant_validation_test",
+    validation_model = ":mobilenet_quant_with_validation.tflite",
+)
+
+validation_model(
+    name = "mobilenet_float_with_validation",
+    testonly = 1,
+    jpegs = "//tensorflow/lite/experimental/acceleration/mini_benchmark:odt_classifier_testfiles",
+    main_model = "//tensorflow/lite/experimental/acceleration/mini_benchmark/models:mobilenet_v1_1.0_224.tflite",
+    metrics_model = ":mobilenet_metrics_tflite",
+)
+
+validation_test(
+    name = "mobilenet_float_validation_test",
+    validation_model = ":mobilenet_float_with_validation.tflite",
+)
+
+py_binary(
+    name = "blazeface_metrics",
+    srcs = ["blazeface_metrics.py"],
+    strict_deps = True,
+    deps = [
+        ":kl_divergence",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/lite/python:lite",
+        "//tensorflow/lite/tools:flatbuffer_utils",
+    ] + if_pywrap(
+        if_true = ["//tensorflow/python:_pywrap_tensorflow"],
+    ),
+)
+
+genrule(
+    name = "blazeface_metrics_tflite",
+    outs = ["blazeface_metrics.tflite"],
+    cmd = "$(location :blazeface_metrics) $(location :blazeface_metrics.tflite)",
+    tools = [":blazeface_metrics"],
+)
+
+validation_model(
+    name = "blazeface_mlkit_v1_with_validation",
+    testonly = 1,
+    jpegs = "//tensorflow/lite/experimental/acceleration/mini_benchmark:blazeface_testfiles",
+    main_model = "//tensorflow/lite/experimental/acceleration/mini_benchmark/models:blazeface_mlkit_v1.tfl",
+    metrics_model = ":blazeface_metrics_tflite",
+    scale = "0.007812",
+    zeropoint = "128",
+)
+
+validation_test(
+    name = "blazeface_mlkit_v1_validation_regression_test",
+    validation_model = ":blazeface_mlkit_v1_with_validation.tflite",
+)

tensorflow/lite/experimental/acceleration/mini_benchmark/metrics/README.md

@@ -0,0 +1,4 @@
+# Code for generating accuracy metric tflite models
+
+This package provides building blocks for generating accuracy metrics used in
+the mini-benchmark.

tensorflow/lite/experimental/acceleration/mini_benchmark/metrics/blazeface_metrics.py

@@ -0,0 +1,134 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Metrics model generator for Blazeface.
+
+The produced model is to be used as part of the mini-benchmark, combined into
+the same flatbuffer with the main model.
+
+The blazeface model is described in
+https://tfhub.dev/tensorflow/tfjs-model/blazeface/1/default/1
+
+The metrics are roughly equivalent to the training time loss function for SSD
+(https://arxiv.org/abs/1512.02325): localization loss and classification loss.
+
+The localization loss is MSE (L2-norm) of box encodings over high-probability
+boxes. A box encoding contains the size and location difference between the
+prediction and the prototype box (see section 2 in the linked paper).
+
+The classification loss is symmetric KL-divergence over classification scores
+squashed to 0..1.
+
+This follows the general rationale of the mini-benchmark: use as much of the
+model outputs as possible for metrics, so that less example data is needed.
+"""
+import argparse
+import sys
+# TODO(b/152872335): (re-)port to tf v2 after output names are kept during
+# conversion in v2.
+import tensorflow.compat.v1 as tf
+from tensorflow.lite.experimental.acceleration.mini_benchmark.metrics import kl_divergence
+from tensorflow.lite.python import lite
+from tensorflow.lite.tools import flatbuffer_utils
+
+parser = argparse.ArgumentParser(
+    description='Script to generate a metrics model for the Blazeface.')
+parser.add_argument('output', help='Output filepath')
+
+
+@tf.function
+def metrics(expected_box_encodings, expected_scores, actual_box_encodings,
+            actual_scores):
+  """Calculate metrics from expected and actual blazeface outputs.
+
+  Args:
+    expected_box_encodings: box encodings from model
+    expected_scores: classifications from model
+    actual_box_encodings: golden box encodings
+    actual_scores: golden classifications
+
+  Returns:
+    two-item list with classification error and localization error
+  """
+  squashed_expected_scores = tf.math.divide(1.0,
+                                            1.0 + tf.math.exp(-expected_scores))
+  squashed_actual_scores = tf.math.divide(1.0,
+                                          1.0 + tf.math.exp(-actual_scores))
+  kld_metric = kl_divergence.symmetric_kl_divergence(expected_scores,
+                                                     actual_scores)
+  # ML Kit uses 0.5 as the threshold. We use
+  # 0.1 to use more possible boxes based on experimentation with the model.
+  high_scoring_indices = tf.math.logical_or(
+      tf.math.greater(squashed_expected_scores, 0.1),
+      tf.math.greater(squashed_actual_scores, 0.1))
+
+  high_scoring_actual_boxes = tf.where(
+      condition=tf.broadcast_to(
+          input=high_scoring_indices, shape=tf.shape(actual_box_encodings)),
+      x=actual_box_encodings,
+      y=expected_box_encodings)
+  box_diff = high_scoring_actual_boxes - expected_box_encodings
+  box_squared_diff = tf.math.pow(box_diff, 2)
+  # MSE is calculated over the high-scoring boxes.
+  box_mse = tf.divide(
+      tf.math.reduce_sum(box_squared_diff),
+      tf.math.maximum(
+          tf.math.count_nonzero(high_scoring_indices, dtype=tf.float32), 1.0))
+  # Thresholds were determined experimentally by running validation on a variety
+  # of devices. Known good devices give KLD ~10-e7 and MSE ~10-e12. A buggy
+  # NNAPI implementation gives KLD > 200 and MSE > 100.
+  ok = tf.logical_and(kld_metric < 0.1, box_mse < 0.01)
+
+  return [kld_metric, box_mse, ok]
+
+
+def main(output_path):
+  tf.reset_default_graph()
+  with tf.Graph().as_default():
+    expected_box_encodings = tf.placeholder(
+        dtype=tf.float32, shape=[1, 564, 16])
+    expected_scores = tf.placeholder(dtype=tf.float32, shape=[1, 564, 1])
+    actual_box_encodings = tf.placeholder(dtype=tf.float32, shape=[1, 564, 16])
+    actual_scores = tf.placeholder(dtype=tf.float32, shape=[1, 564, 1])
+
+    [kld_metric, box_mse, ok] = metrics(expected_box_encodings, expected_scores,
+                                        actual_box_encodings, actual_scores)
+    ok = tf.reshape(ok, [1], name='ok')
+    kld_metric = tf.reshape(kld_metric, [1], name='symmetric_kl_divergence')
+    box_mse = tf.reshape(box_mse, [1], name='box_mse')
+    sess = tf.compat.v1.Session()
+    converter = lite.TFLiteConverter.from_session(sess, [
+        expected_box_encodings, expected_scores, actual_box_encodings,
+        actual_scores
+    ], [kld_metric, box_mse, ok])
+    converter.experimental_new_converter = True
+    tflite_model = converter.convert()
+    if sys.byteorder == 'big':
+      tflite_model = flatbuffer_utils.byte_swap_tflite_buffer(
+          tflite_model, 'big', 'little'
+      )
+    open(output_path, 'wb').write(tflite_model)
+
+
+if __name__ == '__main__':
+  flags, unparsed = parser.parse_known_args()
+
+  if unparsed:
+    parser.print_usage()
+    sys.stderr.write('\nGot the following unparsed args, %r please fix.\n' %
+                     unparsed)
+    exit(1)
+  else:
+    main(flags.output)
+    exit(0)

tensorflow/lite/experimental/acceleration/mini_benchmark/metrics/kl_divergence.py

@@ -0,0 +1,47 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""KL-divergence metrics calculation."""
+
+# TODO(b/152872335): (re-)port to tf v2 after output names are kept during
+# conversion in v2.
+import tensorflow.compat.v1 as tf
+
+
+def symmetric_kl_divergence(predicted, actual):
+  """Calculate symmetric KL-divergence over two classification tensors.
+
+  Note that here the classifications do not form a probability distribution.
+  They are, however normalized to 0..1 and calculating a KL-divergence over them
+  gives reasonable numerical results.
+
+  Shape of the two inputs must be the same at inference time but is otherwise
+  unconstrained.
+
+  Args:
+    predicted: classification outputs from model
+    actual: golden classification outputs
+
+  Returns:
+    Single scalar tensor with symmetric KL-divergence between predicted and
+    actual.
+  """
+  epsilon = tf.constant(1e-7, dtype=tf.float32, name='epsilon')
+  p = tf.math.maximum(predicted, epsilon)
+  q = tf.math.maximum(actual, epsilon)
+  kld_1 = tf.math.reduce_sum(
+      tf.math.multiply(p, tf.math.log(tf.math.divide(p, q))))
+  kld_2 = tf.math.reduce_sum(
+      tf.math.multiply(q, tf.math.log(tf.math.divide(q, p))))
+  return tf.add(kld_1, kld_2)

tensorflow/lite/experimental/acceleration/mini_benchmark/metrics/mobilenet.py

@@ -0,0 +1,77 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Metrics model generator for mobilenet v1.
+
+The produced model is to be used as part of the mini-benchmark, combined into
+the same flatbuffer with the main model.
+
+Mobilenet v1 is described in
+https://tfhub.dev/tensorflow/coral-model/mobilenet_v1_1.0_224_quantized/1/default/1
+
+The metrics used are symmetric KL-divergence and MSE.
+"""
+import argparse
+import sys
+# TODO(b/152872335): (re-)port to tf v2 after output names are kept during
+# conversion in v2.
+import tensorflow.compat.v1 as tf
+from tensorflow.lite.experimental.acceleration.mini_benchmark.metrics import kl_divergence
+from tensorflow.lite.python import lite
+from tensorflow.lite.tools import flatbuffer_utils
+parser = argparse.ArgumentParser(
+    description='Script to generate a metrics model for mobilenet v1.')
+parser.add_argument('output', help='Output filepath')
+
+
+def main(output_path):
+  tf.reset_default_graph()
+  with tf.Graph().as_default():
+    expected_scores = tf.placeholder(dtype=tf.float32, shape=[1, 1001])
+    actual_scores = tf.placeholder(dtype=tf.float32, shape=[1, 1001])
+    mse = tf.reshape(
+        tf.math.reduce_mean((expected_scores - actual_scores)**2), [1],
+        name='mse')
+    kld_metric = kl_divergence.symmetric_kl_divergence(expected_scores,
+                                                       actual_scores)
+    kld_metric = tf.reshape(kld_metric, [1], name='symmetric_kl_divergence')
+    # Thresholds chosen by comparing NNAPI top-k accuracy on MLTS on devices
+    # with top-k accuracy within 1%-p of tflite CPU and with a 5-%p drop.
+    ok = tf.reshape(
+        tf.logical_and(kld_metric < 5.5, mse < 0.003), [1], name='ok')
+    sess = tf.compat.v1.Session()
+    converter = lite.TFLiteConverter.from_session(sess, [
+        expected_scores,
+        actual_scores,
+    ], [kld_metric, mse, ok])
+    converter.experimental_new_converter = True
+    tflite_model = converter.convert()
+    if sys.byteorder == 'big':
+      tflite_model = flatbuffer_utils.byte_swap_tflite_buffer(
+          tflite_model, 'big', 'little'
+      )
+    open(output_path, 'wb').write(tflite_model)
+
+
+if __name__ == '__main__':
+  flags, unparsed = parser.parse_known_args()
+
+  if unparsed:
+    parser.print_usage()
+    sys.stderr.write('\nGot the following unparsed args, %r please fix.\n' %
+                     unparsed)
+    exit(1)
+  else:
+    main(flags.output)
+    exit(0)

tensorflow/lite/experimental/acceleration/mini_benchmark/mini_benchmark.cc

@@ -0,0 +1,95 @@
+/* Copyright 2021 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/mini_benchmark.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "absl/synchronization/mutex.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+
+namespace tflite {
+namespace acceleration {
+
+namespace {
+class NoopMiniBenchmark : public MiniBenchmark {
+ public:
+  ComputeSettingsT GetBestAcceleration() override { return ComputeSettingsT(); }
+  void TriggerMiniBenchmark() override {}
+  void SetEventTimeoutForTesting(int64_t) override {}
+  std::vector<MiniBenchmarkEventT> MarkAndGetEventsToLog() override {
+    return {};
+  }
+  // We return -1 as this no-op instance doesn't have the overall
+  // mini-benchmark-related setup properly initialized.
+  int NumRemainingAccelerationTests() override { return -1; }
+};
+}  // namespace
+
+std::unique_ptr<MiniBenchmark> CreateMiniBenchmark(
+    const MinibenchmarkSettings& settings, const std::string& model_namespace,
+    const std::string& model_id) {
+  absl::StatusOr<std::unique_ptr<MiniBenchmark>> s_or_mb =
+      MinibenchmarkImplementationRegistry::CreateByName(
+          "Impl", settings, model_namespace, model_id);
+  if (!s_or_mb.ok()) {
+    return std::unique_ptr<MiniBenchmark>(new NoopMiniBenchmark());
+  } else {
+    return std::move(*s_or_mb);
+  }
+}
+
+void MinibenchmarkImplementationRegistry::RegisterImpl(
+    const std::string& name, CreatorFunction creator_function) {
+  absl::MutexLock lock(mutex_);
+  factories_[name] = creator_function;
+}
+
+std::unique_ptr<MiniBenchmark> MinibenchmarkImplementationRegistry::CreateImpl(
+    const std::string& name, const MinibenchmarkSettings& settings,
+    const std::string& model_namespace, const std::string& model_id) {
+  absl::MutexLock lock(mutex_);
+  auto it = factories_.find(name);
+  return (it != factories_.end())
+             ? it->second(settings, model_namespace, model_id)
+             : nullptr;
+}
+
+MinibenchmarkImplementationRegistry*
+MinibenchmarkImplementationRegistry::GetSingleton() {
+  static auto* instance = new MinibenchmarkImplementationRegistry();
+  return instance;
+}
+
+std::unique_ptr<MiniBenchmark>
+MinibenchmarkImplementationRegistry::CreateByName(
+    const std::string& name, const MinibenchmarkSettings& settings,
+    const std::string& model_namespace, const std::string& model_id) {
+  auto* const instance = MinibenchmarkImplementationRegistry::GetSingleton();
+  return instance->CreateImpl(name, settings, model_namespace, model_id);
+}
+
+MinibenchmarkImplementationRegistry::Register::Register(
+    const std::string& name, CreatorFunction creator_function) {
+  auto* const instance = MinibenchmarkImplementationRegistry::GetSingleton();
+  instance->RegisterImpl(name, creator_function);
+}
+
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/mini_benchmark.h

@@ -0,0 +1,127 @@
+/* Copyright 2021 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MINI_BENCHMARK_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MINI_BENCHMARK_H_
+
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "absl/base/thread_annotations.h"
+#include "absl/synchronization/mutex.h"
+#include "tensorflow/lite/acceleration/configuration/configuration_generated.h"
+
+namespace tflite {
+namespace acceleration {
+// Instances are thread-compatible, access from multiple threads must be guarded
+// by a mutex.
+//
+// Caution: The mini-benchmark runs silently in-process on non-Android, rather
+// than in a separate process.
+class MiniBenchmark {
+ public:
+  // Get best acceleration based on tests done so far. If no successful tests
+  // are found, the best settings are on CPU or if the settings do not contain
+  // configurations to test or not all relevant fields are present, the returned
+  // ComputeSettingsT will be an object created by the default constructor.
+  // Note: if we have successful mini-benchmark run events, the best
+  // acceleration configuration will be persisted on the local storage as a new
+  // mini-benchmark event.
+  virtual ComputeSettingsT GetBestAcceleration() = 0;
+
+  // Trigger the running of tests in the background in a separate thread on
+  // Linux, but a separate process on Android. If triggering fails, errors are
+  // stored internally.
+  //
+  // Does nothing if the settings do not contain configurations to test or not
+  // all relevant fields are present.
+  virtual void TriggerMiniBenchmark() = 0;
+
+  virtual void SetEventTimeoutForTesting(int64_t timeout_us) = 0;
+
+  // Mark mini-benchmark events that have not yet been marked as to be logged,
+  // and return these events. This can include errors in triggering the
+  // mini-benchmark.
+  virtual std::vector<MiniBenchmarkEventT> MarkAndGetEventsToLog() = 0;
+
+  // Get the number of remaining tests that still need to be completed.
+  // Note that this method should be only called after calling
+  // TriggerMiniBenchmark or GetBestAcceleration where additional
+  // mini-benchmark-related setup could be initialized. In short, -1 is returned
+  // if the overall mini-benchmark-related setup isn't properly initialized.
+  virtual int NumRemainingAccelerationTests() = 0;
+
+  MiniBenchmark() {}
+  virtual ~MiniBenchmark() {}
+
+  MiniBenchmark(MiniBenchmark&) = delete;
+  MiniBenchmark& operator=(const MiniBenchmark&) = delete;
+  MiniBenchmark(MiniBenchmark&&) = delete;
+  MiniBenchmark& operator=(const MiniBenchmark&&) = delete;
+};
+
+// Instantiate a mini-benchmark. This will return a no-op implementation unless
+// the :mini_benchmark_implementation target has been linked in.
+std::unique_ptr<MiniBenchmark> CreateMiniBenchmark(
+    const MinibenchmarkSettings& settings, const std::string& model_namespace,
+    const std::string& model_id);
+
+// A simple registry that allows different mini-benchmark implementations to be
+// created by name.
+//
+// Limitations:
+// - Doesn't allow deregistration.
+// - Doesn't check for duplication registration.
+//
+class MinibenchmarkImplementationRegistry {
+ public:
+  using CreatorFunction = std::function<std::unique_ptr<MiniBenchmark>(
+      const MinibenchmarkSettings& /*settings*/,
+      const std::string& /*model_namespace*/, const std::string& /*model_id*/)>;
+
+  // Returns a MiniBenchmark registered with `name` or nullptr if no matching
+  // mini-benchmark implementation found.
+  static std::unique_ptr<MiniBenchmark> CreateByName(
+      const std::string& name, const MinibenchmarkSettings& settings,
+      const std::string& model_namespace, const std::string& model_id);
+
+  // Struct to be statically allocated for registration.
+  struct Register {
+    Register(const std::string& name, CreatorFunction creator_function);
+  };
+
+ private:
+  void RegisterImpl(const std::string& name, CreatorFunction creator_function);
+  std::unique_ptr<MiniBenchmark> CreateImpl(
+      const std::string& name, const MinibenchmarkSettings& settings,
+      const std::string& model_namespace, const std::string& model_id);
+  static MinibenchmarkImplementationRegistry* GetSingleton();
+
+  absl::Mutex mutex_;
+  std::unordered_map<std::string, CreatorFunction> factories_
+      ABSL_GUARDED_BY(mutex_);
+};
+
+}  // namespace acceleration
+}  // namespace tflite
+
+#define TFLITE_REGISTER_MINI_BENCHMARK_FACTORY_FUNCTION(name, f) \
+  static auto* g_tflite_mini_benchmark_##name##_ =               \
+      new MinibenchmarkImplementationRegistry::Register(#name, f);
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MINI_BENCHMARK_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/mini_benchmark_implementation.cc

@@ -0,0 +1,551 @@
+/* Copyright 2021 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 <algorithm>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/acceleration/configuration/configuration_generated.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/fb_storage.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/mini_benchmark.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/status_codes.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/validator_runner.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/validator_runner_options.h"
+#include "tensorflow/lite/minimal_logging.h"
+#include "tensorflow/lite/nnapi/sl/include/SupportLibrary.h"
+
+namespace tflite {
+namespace acceleration {
+
+// This class is used to store the results of a GetBestAcceleration and
+// information on the events used to take the decision.
+// The class is thread-compatible as MiniBenchmark.
+class MemoizedBestAccelerationSelector {
+ public:
+  // Note 'settings' has to outlast this instance.
+  MemoizedBestAccelerationSelector(const MinibenchmarkSettings& settings,
+                                   const std::string model_namespace,
+                                   const std::string& model_id,
+                                   const std::string& storage_path)
+      : settings_(settings),
+        model_namespace_(model_namespace),
+        model_id_(model_id),
+        number_of_events_in_memoized_call_(0),
+        memoised_result_(nullptr),
+        storage_(storage_path, tflite::DefaultErrorReporter()) {
+    storage_.Read();
+
+    TFLITE_LOG_PROD_ONCE(TFLITE_LOG_INFO,
+                         "Initializing BestAccelerationSelector for model (%s, "
+                         "%s) and storage path %s. Storage has %zu events.\n",
+                         model_namespace_.c_str(), model_id_.c_str(),
+                         storage_path.c_str(), storage_.Count());
+
+    // Read saved status from storage_.
+    for (int i = storage_.Count() - 1; i >= 0; i--) {
+      const MiniBenchmarkEvent* event = storage_.Get(i);
+      if (event == nullptr || event->best_acceleration_decision() == nullptr) {
+        continue;
+      }
+
+      const auto* best_decision = event->best_acceleration_decision();
+      const TFLiteSettings* acceleration_settings =
+          CreateAccelerationFromBenchmark(
+              best_decision->min_latency_event(),
+              best_decision->min_inference_time_us());
+      Memoize(acceleration_settings, best_decision->number_of_source_events());
+      TFLITE_LOG_PROD_ONCE(
+          TFLITE_LOG_INFO,
+          "Rebuilding memoised best acceleration from storage. It has been "
+          "generated based on %d events.\n",
+          number_of_events_in_memoized_call_);
+      break;
+    }
+  }
+
+  ComputeSettingsT GetBestAcceleration(
+      const std::vector<const BenchmarkEvent*>& events) {
+    ComputeSettingsT result;
+    if (events.empty()) {
+      TFLITE_LOG_PROD_ONCE(
+          TFLITE_LOG_INFO,
+          "No completed events are available to calculate best "
+          "acceleration result for model (%s, %s).\n",
+          model_namespace_.c_str(), model_id_.c_str());
+      return result;
+    }
+    if (memoised_result_ != nullptr &&
+        (events.size() == number_of_events_in_memoized_call_)) {
+      TFLITE_LOG_PROD_ONCE(TFLITE_LOG_INFO,
+                           "Returning memoized best acceleration result for "
+                           "model (%s, %s) based on %d events.\n",
+                           model_namespace_.c_str(), model_id_.c_str(),
+                           number_of_events_in_memoized_call_);
+      memoised_result_->UnPackTo(&result);
+      return result;
+    }
+
+    TFLITE_LOG_PROD_ONCE(TFLITE_LOG_INFO,
+                         "Calculating best acceleration result for model (%s, "
+                         "%s) based on %zu events.\n",
+                         model_namespace_.c_str(), model_id_.c_str(),
+                         events.size());
+
+    int64_t min_latency = -1;
+    auto min_latency_event = FindMinLatencyEvent(events, min_latency);
+    if (min_latency_event == nullptr) {
+      // We won't memoise a decision on no events.
+      return result;
+    }
+
+    const TFLiteSettings* acceleration_settings =
+        CreateAccelerationFromBenchmark(min_latency_event, min_latency);
+    Memoize(acceleration_settings, static_cast<int>(events.size()));
+    StoreBestAcceleration(min_latency_event, min_latency);
+    memoised_result_->UnPackTo(&result);
+    return result;
+  }
+
+  // Note that events used here are all benchmark-run-ok events, i.e. those
+  // indicating the acceleration configuration has run successfully and produced
+  // correct results.
+  int NumEventsUsedInBestAcceleration() const {
+    return number_of_events_in_memoized_call_;
+  }
+
+ private:
+  void Memoize(const TFLiteSettings* acceleration_settings, int num_events) {
+    number_of_events_in_memoized_call_ = num_events;
+    memoised_result_buffer_.Clear();
+    flatbuffers::Offset<tflite::TFLiteSettings> tflite_setting_offset = 0;
+    if (acceleration_settings != nullptr) {
+      TFLiteSettingsT copy;
+      acceleration_settings->UnPackTo(&copy);
+      tflite_setting_offset =
+          TFLiteSettings::Pack(memoised_result_buffer_, &copy);
+    }
+    auto compute_settings = CreateComputeSettings(
+        memoised_result_buffer_, tflite::ExecutionPreference_ANY,
+        tflite_setting_offset,
+        memoised_result_buffer_.CreateString(model_namespace_),
+        memoised_result_buffer_.CreateString(model_id_));
+    memoised_result_buffer_.Finish(compute_settings);
+    memoised_result_ = flatbuffers::GetRoot<ComputeSettings>(
+        memoised_result_buffer_.GetBufferPointer());
+  }
+
+  // Stores the BestAcceleration to persistent storage to handle restart.
+  void StoreBestAcceleration(const BenchmarkEvent* min_latency_event,
+                             int64_t min_latency) {
+    flatbuffers::FlatBufferBuilder fbb;
+    tflite::BenchmarkEventT min_latency_ev_copy;
+    min_latency_event->UnPackTo(&min_latency_ev_copy);
+    auto best_acceleration_decision = tflite::CreateBestAccelerationDecision(
+        fbb, number_of_events_in_memoized_call_,
+        CreateBenchmarkEvent(fbb, &min_latency_ev_copy), min_latency);
+    storage_.Append(
+        &fbb, CreateMiniBenchmarkEvent(fbb, /*is_log_flushing_event=*/false,
+                                       best_acceleration_decision));
+  }
+
+  const BenchmarkEvent* FindMinLatencyEvent(
+      const std::vector<const BenchmarkEvent*>& events, int64_t& min_latency) {
+    const BenchmarkEvent* min_latency_event = nullptr;
+    min_latency = -1;
+    for (const BenchmarkEvent* ev : events) {
+      for (int i = 0; i < ev->result()->inference_time_us()->size(); i++) {
+        int64_t latency = ev->result()->inference_time_us()->Get(i);
+        if (latency < 0) {
+          continue;
+        }
+        if ((min_latency < 0) || (latency < min_latency)) {
+          min_latency = latency;
+          min_latency_event = ev;
+        }
+      }
+    }
+
+    return min_latency_event;
+  }
+
+  std::string GetDelegateFromBenchmarkEvent(const BenchmarkEvent* event) const {
+    if (event->tflite_settings()->delegate() == Delegate_NNAPI) {
+      return "NNAPI";
+    }
+    if (event->tflite_settings()->delegate() == Delegate_GPU) {
+      return "GPU";
+    }
+    if (event->tflite_settings()->delegate() == Delegate_XNNPACK) {
+      return "XNNPACK";
+    }
+
+    return "CPU";
+  }
+
+  // Returns the acceleration settings in the list of settings_to_test
+  // corresponding to the tflite settings of the provided min_latency_event.
+  // Returns nullptr if no matching settings is found.
+  const TFLiteSettings* FindAccelerationToTestFromMiniBenchmarkEvent(
+      const BenchmarkEvent* min_latency_event) const {
+    TFLiteSettingsT event_tflite_settings;
+    min_latency_event->tflite_settings()->UnPackTo(&event_tflite_settings);
+    for (int i = 0; i < settings_.settings_to_test()->size(); i++) {
+      auto one_setting = settings_.settings_to_test()->Get(i);
+      TFLiteSettingsT to_test_tflite_settings;
+      one_setting->UnPackTo(&to_test_tflite_settings);
+      if (to_test_tflite_settings == event_tflite_settings) {
+        return one_setting;
+      }
+    }
+    TFLITE_LOG_PROD_ONCE(
+        TFLITE_LOG_WARNING,
+        "Couldn't find  setting to test matching the best latency event for "
+        "model %s, returning no acceleration.\n",
+        model_id_.c_str());
+    return nullptr;
+  }
+
+  const TFLiteSettings* CreateAccelerationFromBenchmark(
+      const BenchmarkEvent* min_latency_event, int64_t min_latency) const {
+    std::string delegate = GetDelegateFromBenchmarkEvent(min_latency_event);
+    TFLITE_LOG_PROD_ONCE(
+        TFLITE_LOG_INFO,
+        "Found best latency for %s with delegate %s ( %ld us).\n",
+        model_id_.c_str(), delegate.c_str(), min_latency);
+
+    if (min_latency_event->tflite_settings()->delegate() == Delegate_NONE) {
+      TFLITE_LOG_PROD_ONCE(
+          TFLITE_LOG_INFO,
+          "Best latency for %s is without a delegate, not overring defaults.\n",
+          model_id_.c_str());
+      return nullptr;
+    }
+
+    return FindAccelerationToTestFromMiniBenchmarkEvent(min_latency_event);
+  }
+
+  const MinibenchmarkSettings& settings_;
+  std::string model_namespace_;
+  std::string model_id_;
+  // The number of events we are basing our best acceleration decision on.
+  // We are assuming that the events passed to this object will only increase
+  // and never change.
+  int number_of_events_in_memoized_call_ = 0;
+  flatbuffers::FlatBufferBuilder memoised_result_buffer_;
+  // Pointer to the 'memoised_result_buffer_' for convenience.
+  const ComputeSettings* memoised_result_;
+
+  FlatbufferStorage<MiniBenchmarkEvent> storage_;
+};
+
+class MiniBenchmarkImpl : public MiniBenchmark {
+ public:
+  MiniBenchmarkImpl(const MinibenchmarkSettings& settings,
+                    const std::string& model_namespace,
+                    const std::string& model_id)
+      : model_namespace_(model_namespace), model_id_(model_id) {
+    // Keep a copy of the passed in 'settings' to simplify the memory
+    // management. Otherwise, the 'settings' has to outlast this instance.
+    MinibenchmarkSettingsT copy;
+    settings.UnPackTo(&copy);
+    settings_buffer_.Finish(
+        MinibenchmarkSettings::Pack(settings_buffer_, &copy));
+    settings_ = flatbuffers::GetRoot<MinibenchmarkSettings>(
+        settings_buffer_.GetBufferPointer());
+
+    is_enabled_ = BenchmarkIsEnabled();
+    if (!is_enabled_) return;
+
+    is_cpu_validation_specified_ = false;
+    total_validation_tests_ = settings_->settings_to_test()->size();
+    for (int i = 0; i < settings_->settings_to_test()->size(); i++) {
+      auto one_setting = settings_->settings_to_test()->Get(i);
+      if (one_setting->delegate() == Delegate_NONE) {
+        is_cpu_validation_specified_ = true;
+      }
+    }
+    // By default, will always add a cpu testing even if it's not requested.
+    if (total_validation_tests_ != 0 && !is_cpu_validation_specified_) {
+      total_validation_tests_ += 1;
+    }
+
+    const std::string local_event_fp = LocalEventStorageFileName(settings);
+    best_acceleration_selector_ =
+        std::make_unique<MemoizedBestAccelerationSelector>(
+            *settings_, model_namespace, model_id, local_event_fp);
+
+    storage_ = std::make_unique<FlatbufferStorage<MiniBenchmarkEvent>>(
+        local_event_fp, tflite::DefaultErrorReporter());
+    storage_->Read();
+    for (int i = storage_->Count() - 1; i >= 0; i--) {
+      auto* event = storage_->Get(i);
+      if (event != nullptr && event->initialization_failure()) {
+        initialization_failure_logged_ = true;
+        break;
+      }
+    }
+  }
+
+  ComputeSettingsT GetBestAcceleration() override {
+    if (!is_enabled_) return ComputeSettingsT();
+    CreateValidatorIfNececessary();
+    if (!validator_initialized_) return ComputeSettingsT();
+
+    std::vector<const BenchmarkEvent*> events =
+        validator_->GetSuccessfulResults();
+    TFLITE_LOG_PROD_ONCE(TFLITE_LOG_INFO,
+                         "Got %zu successful minibenchmark events for %s.\n",
+                         events.size(), model_id_.c_str());
+
+    return best_acceleration_selector_->GetBestAcceleration(events);
+  }
+
+  void TriggerMiniBenchmark() override {
+    if (!is_enabled_) return;
+    CreateValidatorIfNececessary();
+    if (!validator_initialized_) return;
+
+    std::vector<const TFLiteSettings*> settings;
+    for (int i = 0; i < settings_->settings_to_test()->size(); i++) {
+      auto one_setting = settings_->settings_to_test()->Get(i);
+      settings.push_back(one_setting);
+    }
+    // By default, always add a cpu testing even if it's not requested.
+    flatbuffers::FlatBufferBuilder cpu_fbb;
+    if (!settings.empty() && !is_cpu_validation_specified_) {
+      cpu_fbb.Finish(CreateTFLiteSettings(cpu_fbb));
+      settings.push_back(
+          flatbuffers::GetRoot<TFLiteSettings>(cpu_fbb.GetBufferPointer()));
+    }
+    int triggered = validator_->TriggerMissingValidation(settings);
+    if (triggered > 0) {
+      TFLITE_LOG_PROD(TFLITE_LOG_INFO,
+                      "Triggered mini benchmark for %s with %d possibilities "
+                      "(including CPU).\n",
+                      model_id_.c_str(), triggered);
+    }
+  }
+
+  std::vector<MiniBenchmarkEventT> MarkAndGetEventsToLog() override {
+    if (!is_enabled_) return {};
+
+    std::vector<tflite::MiniBenchmarkEventT> result;
+
+    // Internal MiniBenchmarkEvents.
+    storage_->Read();
+    int storage_size_pre_flush = storage_->Count();
+    if (storage_size_pre_flush != 0) {
+      // Marking current events as read. Adding a MiniBenchmarkEvent
+      // with the is_log_flushing_event flag set to true.
+      flatbuffers::FlatBufferBuilder fbb;
+      storage_->Append(&fbb, ::tflite::CreateMiniBenchmarkEvent(
+                                 fbb, /*is_log_flushing_event=*/true));
+
+      storage_->Read();
+      for (int i = storage_size_pre_flush - 1; i >= 0; i--) {
+        const MiniBenchmarkEvent* event = storage_->Get(i);
+        if (event == nullptr || event->is_log_flushing_event()) {
+          break;
+        }
+        tflite::MiniBenchmarkEventT mini_benchmark_event;
+        event->UnPackTo(&mini_benchmark_event);
+        result.push_back(std::move(mini_benchmark_event));
+      }
+    }
+    std::reverse(result.begin(), result.end());
+
+    // BenchmarkEvents from the validaton runner.
+    std::vector<const BenchmarkEvent*> event_ptrs =
+        validator_->GetAndFlushEventsToLog(event_timeout_us_);
+    for (const auto* event_ptr : event_ptrs) {
+      tflite::MiniBenchmarkEventT mini_benchmark_event;
+      mini_benchmark_event.benchmark_event =
+          std::unique_ptr<BenchmarkEventT>(event_ptr->UnPack());
+      result.push_back(std::move(mini_benchmark_event));
+    }
+
+    return result;
+  }
+
+  void SetEventTimeoutForTesting(int64_t timeout_us) override {
+    event_timeout_us_ = (timeout_us == -1)
+                            ? ValidatorRunner::kDefaultEventTimeoutUs
+                            : timeout_us;
+  }
+
+  int NumRemainingAccelerationTests() override {
+    // We return -1 when the overall mini-benchmark-related setup isn't properly
+    // initialized.
+    if (!is_enabled_ || !validator_initialized_) return -1;
+
+    // We first check if there has been a previous execution of the runner
+    // already using some of the validation results to decide the best
+    // acceleration.
+    const int to_complete_tests =
+        total_validation_tests_ -
+        best_acceleration_selector_->NumEventsUsedInBestAcceleration();
+    // No remaining tests, skip reading the validation events log and just
+    // return.
+    if (to_complete_tests == 0) return 0;
+
+    // Read the whole validation events log to find the number of completed
+    // runs.
+    return total_validation_tests_ - validator_->GetNumCompletedResults();
+  }
+
+ private:
+  static std::string LocalEventStorageFileName(
+      const MinibenchmarkSettings& settings) {
+    if (settings.storage_paths() == nullptr ||
+        settings.storage_paths()->storage_file_path() == nullptr) {
+      return "mini_benchmark.default.extra.fb";
+    }
+    return settings.storage_paths()->storage_file_path()->str() + ".extra.fb";
+  }
+
+  bool BenchmarkIsEnabled() const {
+    if (settings_->settings_to_test() == nullptr) return false;
+    if (settings_->settings_to_test()->size() <= 0) return false;
+
+    const auto* storage_paths = settings_->storage_paths();
+    if (storage_paths == nullptr) return false;
+    if (storage_paths->storage_file_path() == nullptr ||
+        storage_paths->storage_file_path()->str().empty()) {
+      TFLITE_LOG_PROD_ONCE(
+          TFLITE_LOG_ERROR,
+          "Minibenchmark requested for %s but storage_file_path not set.\n",
+          model_id_.c_str());
+      return false;
+    } else if (storage_paths->data_directory_path() == nullptr ||
+               storage_paths->data_directory_path()->str().empty()) {
+      TFLITE_LOG_PROD_ONCE(
+          TFLITE_LOG_ERROR,
+          "Minibenchmark requested for %s but data_directory_path not set.\n",
+          model_id_.c_str());
+      return false;
+    }
+
+    const auto* model_file = settings_->model_file();
+    if (model_file == nullptr) return false;
+    if (model_file->fd() <= 0 && (model_file->filename() == nullptr ||
+                                  model_file->filename()->str().empty())) {
+      TFLITE_LOG_PROD_ONCE(
+          TFLITE_LOG_ERROR,
+          "Minibenchmark requested for %s but model_file not set.\n",
+          model_id_.c_str());
+      return false;
+    }
+
+    return true;
+  }
+
+  MinibenchmarkStatus GetNnApiSlPointerIfPresent(
+      const NnApiSLDriverImplFL5** nnapi_sl) {
+    *nnapi_sl = nullptr;
+    const auto& settings_to_test = *settings_->settings_to_test();
+    for (const auto* setting_to_test : settings_to_test) {
+      // Check that there are not two different NNAPI-with-SL configurations
+      // with different support library instances.
+      if (setting_to_test->delegate() == Delegate_NNAPI &&
+          setting_to_test->nnapi_settings() &&
+          setting_to_test->nnapi_settings()->support_library_handle()) {
+        const NnApiSLDriverImplFL5* curr_nnapi_sl_handle =
+            reinterpret_cast<const NnApiSLDriverImplFL5*>(
+                setting_to_test->nnapi_settings()->support_library_handle());
+
+        if (*nnapi_sl != nullptr && *nnapi_sl != curr_nnapi_sl_handle) {
+          return kMiniBenchmarkInvalidSupportLibraryConfiguration;
+        }
+
+        *nnapi_sl = curr_nnapi_sl_handle;
+      }
+    }
+    return kMinibenchmarkSuccess;
+  }
+
+  void LogInitializationFailure(MinibenchmarkStatus status) {
+    if (!initialization_failure_logged_) {
+      flatbuffers::FlatBufferBuilder fbb;
+      storage_->Append(&fbb, CreateMiniBenchmarkEvent(
+                                 fbb, /*is_log_flushing_event=*/false,
+                                 /*best_acceleration_decision=*/0,
+                                 ::tflite::CreateBenchmarkInitializationFailure(
+                                     fbb, status)));
+      initialization_failure_logged_ = true;
+    }
+  }
+
+  void CreateValidatorIfNececessary() {
+    if (validator_) return;
+
+    ValidatorRunnerOptions options =
+        CreateValidatorRunnerOptionsFrom(*settings_);
+    MinibenchmarkStatus get_nnapi_sl_status =
+        GetNnApiSlPointerIfPresent(&options.nnapi_sl);
+    if (get_nnapi_sl_status != kMinibenchmarkSuccess) {
+      LogInitializationFailure(get_nnapi_sl_status);
+      return;
+    }
+    validator_ = std::make_unique<ValidatorRunner>(options);
+
+    MinibenchmarkStatus status = validator_->Init();
+    if (status == kMinibenchmarkValidationEntrypointSymbolNotFound) {
+      TFLITE_LOG_PROD_ONCE(TFLITE_LOG_ERROR,
+                           "Model %s does not contain a validation subgraph.",
+                           model_id_.c_str());
+    } else if (status != kMinibenchmarkSuccess) {
+      TFLITE_LOG_PROD_ONCE(TFLITE_LOG_ERROR,
+                           "ValidatorRunner::Init() failed for model %s.",
+                           model_id_.c_str());
+    } else {
+      validator_initialized_ = true;
+    }
+    if (status != kMinibenchmarkSuccess) {
+      LogInitializationFailure(status);
+    }
+  }
+
+  flatbuffers::FlatBufferBuilder settings_buffer_;
+  // Just a pointer to the 'settings_buffer_' for convenience.
+  const MinibenchmarkSettings* settings_ = nullptr;
+  bool is_enabled_ = false;
+  int total_validation_tests_ = 0;
+  bool is_cpu_validation_specified_ = false;
+
+  std::unique_ptr<ValidatorRunner> validator_ = nullptr;
+  bool validator_initialized_ = false;
+  std::string model_namespace_;
+  std::string model_id_;
+  int64_t event_timeout_us_ = ValidatorRunner::kDefaultEventTimeoutUs;
+  std::unique_ptr<MemoizedBestAccelerationSelector>
+      best_acceleration_selector_ = nullptr;
+
+  std::unique_ptr<FlatbufferStorage<MiniBenchmarkEvent>> storage_ = nullptr;
+  bool initialization_failure_logged_ = false;
+};
+
+std::unique_ptr<MiniBenchmark> CreateMiniBenchmarkImpl(
+    const MinibenchmarkSettings& settings, const std::string& model_namespace,
+    const std::string& model_id) {
+  return std::unique_ptr<MiniBenchmark>(
+      new MiniBenchmarkImpl(settings, model_namespace, model_id));
+}
+
+TFLITE_REGISTER_MINI_BENCHMARK_FACTORY_FUNCTION(Impl, CreateMiniBenchmarkImpl);
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/mini_benchmark_test.cc

@@ -0,0 +1,320 @@
+/* Copyright 2021 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/mini_benchmark.h"
+
+#include <unistd.h>
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "absl/time/clock.h"
+#include "absl/time/time.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/acceleration/configuration/configuration.pb.h"
+#include "tensorflow/lite/acceleration/configuration/configuration_generated.h"
+#include "tensorflow/lite/acceleration/configuration/proto_to_flatbuffer.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/embedded_mobilenet_float_validation_model.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/embedded_nnapi_sl_fake_impl.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/mini_benchmark_test_helper.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/nnapi_sl_fake_impl.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/status_codes.h"
+#include "tensorflow/lite/nnapi/sl/include/SupportLibrary.h"
+
+namespace tflite {
+namespace acceleration {
+namespace {
+
+TEST(BasicMiniBenchmarkTest, EmptySettings) {
+  proto::MinibenchmarkSettings settings_proto;
+  flatbuffers::FlatBufferBuilder empty_settings_buffer_;
+  const MinibenchmarkSettings* empty_settings =
+      ConvertFromProto(settings_proto, &empty_settings_buffer_);
+  std::unique_ptr<MiniBenchmark> mb(
+      CreateMiniBenchmark(*empty_settings, "ns", "id"));
+  mb->TriggerMiniBenchmark();
+  const ComputeSettingsT acceleration = mb->GetBestAcceleration();
+  EXPECT_EQ(nullptr, acceleration.tflite_settings);
+  EXPECT_TRUE(mb->MarkAndGetEventsToLog().empty());
+  EXPECT_EQ(-1, mb->NumRemainingAccelerationTests());
+}
+
+class MiniBenchmarkTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    MiniBenchmarkTestHelper helper;
+    should_perform_test_ = helper.should_perform_test();
+
+    if (should_perform_test_) {
+      mobilenet_model_path_ = MiniBenchmarkTestHelper::DumpToTempFile(
+          "mobilenet_float_with_validation.tflite",
+          g_tflite_acceleration_embedded_mobilenet_float_validation_model,
+          g_tflite_acceleration_embedded_mobilenet_float_validation_model_len);
+    }
+  }
+
+  void SetupBenchmark(proto::Delegate delegate, const std::string& model_path,
+                      bool reset_storage = true,
+                      const nnapi::NnApiSupportLibrary* nnapi_sl = nullptr) {
+    proto::MinibenchmarkSettings settings;
+    proto::TFLiteSettings* tflite_settings = settings.add_settings_to_test();
+    tflite_settings->set_delegate(delegate);
+    if ((delegate == proto::Delegate::NNAPI) && nnapi_sl) {
+      std::cerr << "Using NNAPI SL\n";
+      tflite_settings->mutable_nnapi_settings()->set_support_library_handle(
+          reinterpret_cast<int64_t>(nnapi_sl->getFL5()));
+    }
+    proto::ModelFile* file = settings.mutable_model_file();
+    file->set_filename(model_path);
+    proto::BenchmarkStoragePaths* paths = settings.mutable_storage_paths();
+    paths->set_storage_file_path(::testing::TempDir() + "/storage.fb");
+    if (reset_storage) {
+      (void)unlink(paths->storage_file_path().c_str());
+      // The suffix needs to be same as the one used in
+      // MiniBenchmarkImpl::LocalEventStorageFileName
+      (void)unlink((paths->storage_file_path() + ".extra.fb").c_str());
+    }
+    paths->set_data_directory_path(::testing::TempDir());
+
+    if (delegate != proto::Delegate::NONE) {
+      // Some of the tests rely on XNNPack beating CPU - need to not apply
+      // XNNPack for the CPU variant for the tests to pass.
+      proto::TFLiteSettings* cpu_tflite_settings =
+          settings.add_settings_to_test();
+      cpu_tflite_settings->set_disable_default_delegates(false);
+    }
+
+    settings_ = ConvertFromProto(settings, &settings_buffer_);
+
+    mb_ = CreateMiniBenchmark(*settings_, ns_, model_id_);
+  }
+
+  void TriggerBenchmark(proto::Delegate delegate, const std::string& model_path,
+                        bool reset_storage = true,
+                        const nnapi::NnApiSupportLibrary* nnapi_sl = nullptr) {
+    SetupBenchmark(delegate, model_path, reset_storage, nnapi_sl);
+    mb_->TriggerMiniBenchmark();
+  }
+
+  void WaitForValidationCompletion(
+      absl::Duration timeout = absl::Seconds(300)) {
+    absl::Time deadline = absl::Now() + timeout;
+    while (absl::Now() < deadline) {
+      if (mb_->NumRemainingAccelerationTests() == 0) return;
+      absl::SleepFor(absl::Milliseconds(200));
+    }
+
+    // We reach here only when the timeout has been reached w/o validation
+    // completing.
+    ASSERT_NE(0, mb_->NumRemainingAccelerationTests());
+  }
+
+  const std::string ns_ = "org.tensorflow.lite.mini_benchmark.test";
+  const std::string model_id_ = "test_minibenchmark_model";
+
+  bool should_perform_test_ = true;
+
+  std::unique_ptr<MiniBenchmark> mb_;
+  std::string mobilenet_model_path_;
+
+  flatbuffers::FlatBufferBuilder settings_buffer_;
+  // Simply a reference to settings_buffer_ for convenience.
+  const MinibenchmarkSettings* settings_;
+};
+
+TEST_F(MiniBenchmarkTest, OnlyCPUSettings) {
+  if (!should_perform_test_) return;
+  SetupBenchmark(proto::Delegate::NONE, mobilenet_model_path_);
+  EXPECT_EQ(-1, mb_->NumRemainingAccelerationTests());
+  // We haven't triggered the benchmark yet, so a default ComputeSettingsT is
+  // expected.
+  ComputeSettingsT acceleration = mb_->GetBestAcceleration();
+  EXPECT_EQ(nullptr, acceleration.tflite_settings);
+  EXPECT_EQ(1, mb_->NumRemainingAccelerationTests());
+
+  mb_->TriggerMiniBenchmark();
+  // We just have 1 acceleration test to complete as the default CPU execution.
+  WaitForValidationCompletion();
+  acceleration = mb_->GetBestAcceleration();
+  // As the best is the default CPU execution, the returned acceleration above
+  // is still a default ComputeSettingsT.
+  EXPECT_EQ(nullptr, acceleration.tflite_settings);
+}
+
+TEST_F(MiniBenchmarkTest, RunSuccessfully) {
+  if (!should_perform_test_) return;
+
+  TriggerBenchmark(proto::Delegate::XNNPACK, mobilenet_model_path_);
+
+  // We have 2 acceleration tests to complete: one for the default CPU execution
+  // and the other for XNNPACK delegate execution.
+  WaitForValidationCompletion();
+  // Mark existing events to be logged to simplify the logic of checking the
+  // best decision event later.
+  mb_->MarkAndGetEventsToLog();
+
+  const ComputeSettingsT acceleration1 = mb_->GetBestAcceleration();
+
+  // The 2nd call should return the same acceleration settings.
+  const ComputeSettingsT acceleration2 = mb_->GetBestAcceleration();
+  EXPECT_EQ(acceleration1, acceleration2);
+#ifndef ADDRESS_SANITIZER  // XNNPack is slower under Asan.
+  // As we choose mobilenet-v1 float model, XNNPACK delegate should be the best
+  // on CPU.
+  ASSERT_NE(nullptr, acceleration1.tflite_settings);
+  EXPECT_EQ(tflite::Delegate_XNNPACK, acceleration1.tflite_settings->delegate);
+#endif  // !ADDRESS_SANITIZER
+
+  EXPECT_EQ(model_id_, acceleration1.model_identifier_for_statistics);
+  EXPECT_EQ(ns_, acceleration1.model_namespace_for_statistics);
+
+  // As the best decision event has not been marked as to-be-logged, we should
+  // get one best decision event after the call.
+  auto events = mb_->MarkAndGetEventsToLog();
+  EXPECT_EQ(1, events.size());
+  const auto& decision = events.front().best_acceleration_decision;
+  EXPECT_NE(nullptr, decision);
+#ifndef ADDRESS_SANITIZER  // XNNPack is slower under Asan.
+  EXPECT_EQ(tflite::Delegate_XNNPACK,
+            decision->min_latency_event->tflite_settings->delegate);
+#endif  // !ADDRESS_SANITIZER
+}
+
+TEST_F(MiniBenchmarkTest, BestAccelerationEventIsMarkedLoggedAfterRestart) {
+  if (!should_perform_test_) return;
+
+  TriggerBenchmark(proto::Delegate::XNNPACK, mobilenet_model_path_);
+  // We have 2 acceleration tests to complete: one for the default CPU execution
+  // and the other for XNNPACK delegate execution.
+  WaitForValidationCompletion();
+  // Mark existing events to be logged to simplify the logic of checking the
+  // best decision event later.
+  mb_->MarkAndGetEventsToLog();
+  mb_->GetBestAcceleration();
+
+  // The best acceleration decision event was already persisted to the storage
+  // above. So, we could retrieve the best acceleration immediately.
+  TriggerBenchmark(proto::Delegate::XNNPACK, mobilenet_model_path_,
+                   /*reset_storage=*/false);
+  // As all acceleration tests have completed before, we expect no remaining
+  // tests to be performed.
+  EXPECT_EQ(0, mb_->NumRemainingAccelerationTests());
+
+  const ComputeSettingsT acceleration = mb_->GetBestAcceleration();
+#ifndef ADDRESS_SANITIZER  // XNNPack is slower under Asan.
+  ASSERT_NE(nullptr, acceleration.tflite_settings);
+  // As we choose mobilenet-v1 float model, XNNPACK delegate should be the best
+  // on CPU.
+  EXPECT_EQ(tflite::Delegate_XNNPACK, acceleration.tflite_settings->delegate);
+#endif  // !ADDRESS_SANITIZER
+  EXPECT_EQ(model_id_, acceleration.model_identifier_for_statistics);
+  EXPECT_EQ(ns_, acceleration.model_namespace_for_statistics);
+
+  // Similar to "RunSuccessfully' test above, the best decision event has not
+  // been marked as to-be-logged, we should get one best decision event after
+  // the call.
+  auto events = mb_->MarkAndGetEventsToLog();
+  EXPECT_EQ(1, events.size());
+}
+
+TEST_F(MiniBenchmarkTest,
+       BestAccelerationEventIsNotReMarkedLoggedAfterRestart) {
+  if (!should_perform_test_) return;
+
+  TriggerBenchmark(proto::Delegate::XNNPACK, mobilenet_model_path_);
+  // We have 2 acceleration tests to complete: one for the default CPU execution
+  // and the other for XNNPACK delegate execution.
+  WaitForValidationCompletion();
+  mb_->GetBestAcceleration();
+
+  // As the GetBestAcceleration above generates events synchronously, the event
+  // will be persisted to the storage after the call, thus no waiting is needed
+  // here.
+  mb_->MarkAndGetEventsToLog();
+
+  // The best acceleration decision event was already collected above. So, we
+  // could retrieve the best acceleration immediately.
+  TriggerBenchmark(proto::Delegate::XNNPACK, mobilenet_model_path_,
+                   /*reset_storage=*/false);
+  mb_->GetBestAcceleration();
+
+  // As we have marked mini-benchmark events to be logged, we will expect
+  // empty to-log events.
+  EXPECT_TRUE(mb_->MarkAndGetEventsToLog().empty());
+}
+
+TEST_F(MiniBenchmarkTest, DelegatePluginNotSupported) {
+  if (!should_perform_test_) return;
+
+  // As Hexagon delegate plugin isn't supported in mini-benchmark, we will
+  // expect a delegate plugin not-supported error.
+  // Also, note that if a supported delegate plugin isn't linked to the this
+  // test itself or the ":validator_runner_so_for_tests" target on Android,
+  // one will expect a delegate plugin not-found error.
+  TriggerBenchmark(proto::Delegate::HEXAGON, mobilenet_model_path_);
+
+  // We have 2 acceleration tests to complete: one for the default CPU execution
+  // and the other for HEXAGON delegate not supported.
+  WaitForValidationCompletion();
+
+  const ComputeSettingsT acceleration = mb_->GetBestAcceleration();
+  // As the best performance is achieved on the default CPU, there's no
+  // acceleration settings.
+  EXPECT_EQ(nullptr, acceleration.tflite_settings);
+  EXPECT_EQ(model_id_, acceleration.model_identifier_for_statistics);
+  EXPECT_EQ(ns_, acceleration.model_namespace_for_statistics);
+
+  // Check there is a Hexagon-delegate-not-supported event.
+  const auto events = mb_->MarkAndGetEventsToLog();
+  bool is_found = false;
+  for (const auto& event : events) {
+    const auto& t = event.benchmark_event;
+    if (t == nullptr) continue;
+    if (t->event_type == tflite::BenchmarkEventType_ERROR &&
+        t->error->mini_benchmark_error_code ==
+            tflite::acceleration::kMinibenchmarkDelegateNotSupported) {
+      is_found = true;
+      break;
+    }
+  }
+  EXPECT_TRUE(is_found);
+}
+
+#ifdef __ANDROID__  // Loading shared libraries only works on Android.
+TEST_F(MiniBenchmarkTest, UseNnApiSl) {
+  if (!should_perform_test_) return;
+
+  std::string nnapi_sl_path_ = MiniBenchmarkTestHelper::DumpToTempFile(
+      "libnnapi_fake.so", g_nnapi_sl_fake_impl, g_nnapi_sl_fake_impl_len);
+
+  std::unique_ptr<const ::tflite::nnapi::NnApiSupportLibrary> nnapi_sl =
+      ::tflite::nnapi::loadNnApiSupportLibrary(nnapi_sl_path_);
+
+  ASSERT_TRUE(nnapi_sl);
+
+  TriggerBenchmark(proto::Delegate::NNAPI, mobilenet_model_path_,
+                   /*reset_storage=*/true, nnapi_sl.get());
+
+  WaitForValidationCompletion();
+
+  EXPECT_TRUE(tflite::acceleration::WasNnApiSlInvoked());
+}
+#endif  // __ANDROID__
+
+}  // namespace
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/mini_benchmark_test_helper.cc

@@ -0,0 +1,111 @@
+/* Copyright 2021 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/mini_benchmark_test_helper.h"
+
+#include <fcntl.h>
+
+#include <string>
+#ifndef _WIN32
+#include <dlfcn.h>
+#include <signal.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#endif  // !_WIN32
+
+#include <fstream>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/acceleration/compatibility/android_info.h"
+#include "tensorflow/lite/tools/logging.h"
+
+#ifdef __ANDROID__
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/embedded_runner_executable.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/embedded_validator_runner_entrypoint.h"
+#endif  // __ANDROID__
+
+namespace tflite {
+namespace acceleration {
+namespace {
+#ifdef __ANDROID__
+void* LoadEntryPointModule(const std::string& module_path) {
+  void* module =
+      dlopen(module_path.c_str(), RTLD_NOW | RTLD_GLOBAL | RTLD_NODELETE);
+  if (module == nullptr) {
+    TFLITE_LOG(FATAL) << dlerror();
+  }
+  return module;
+}
+#endif  // __ANDROID__
+
+std::string JoinPath(const std::string& path1, const std::string& path2) {
+  if (path1.empty()) return path2;
+  if (path2.empty()) return path1;
+  if (path1.back() == '/') {
+    if (path2.front() == '/') return path1 + path2.substr(1);
+  } else {
+    if (path2.front() != '/') return path1 + "/" + path2;
+  }
+  return path1 + path2;
+}
+}  // namespace
+
+MiniBenchmarkTestHelper::MiniBenchmarkTestHelper(
+    bool should_load_entrypoint_dynamically)
+    : should_perform_test_(true) {
+#ifdef __ANDROID__
+  AndroidInfo android_info;
+  const auto status = RequestAndroidInfo(&android_info);
+  if (!status.ok() || android_info.is_emulator) {
+    should_perform_test_ = false;
+    return;
+  }
+
+  if (!should_load_entrypoint_dynamically) {
+    return;
+  }
+  DumpToTempFile("librunner_main.so", g_tflite_acceleration_embedded_runner,
+                 g_tflite_acceleration_embedded_runner_len);
+
+  // We extract the test files here as that's the only way to get the right
+  // architecture when building tests for multiple architectures.
+  std::string validator_runner_so_path = DumpToTempFile(
+      "libvalidator_runner_entrypoint.so",
+      g_tflite_acceleration_embedded_validator_runner_entrypoint,
+      g_tflite_acceleration_embedded_validator_runner_entrypoint_len);
+
+  // Load this library here because it contains the validation entry point
+  // "Java_org_tensorflow_lite_acceleration_validation_entrypoint" that is then
+  // found using dlsym (using RTLD_DEFAULT hence not needing the handle) in the
+  // mini-benchmark code.
+  LoadEntryPointModule(validator_runner_so_path);
+#endif  // __ANDROID__
+}
+
+std::string MiniBenchmarkTestHelper::DumpToTempFile(const std::string& filename,
+                                                    const unsigned char* data,
+                                                    size_t length) {
+  std::string path = JoinPath(::testing::TempDir(), filename);
+  (void)unlink(path.c_str());
+  std::string contents(reinterpret_cast<const char*>(data), length);
+  std::ofstream f(path, std::ios::binary);
+  EXPECT_TRUE(f.is_open());
+  f << contents;
+  f.close();
+  EXPECT_EQ(0, chmod(path.c_str(), 0500));
+  return path;
+}
+
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/mini_benchmark_test_helper.h

@@ -0,0 +1,50 @@
+/* Copyright 2021 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MINI_BENCHMARK_TEST_HELPER_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MINI_BENCHMARK_TEST_HELPER_H_
+
+#include <string>
+
+namespace tflite {
+namespace acceleration {
+
+class MiniBenchmarkTestHelper {
+ public:
+  // Dump the in-memory binary data stream to the testing temporary directory w/
+  // a file name as 'filename'.
+  // It retruns the full file path of the dumped file.
+  static std::string DumpToTempFile(const std::string& filename,
+                                    const unsigned char* data, size_t length);
+
+  // The constructor will check whether the testing environment supports to run
+  // the mini benchmark. If yes, it will do additional testing setup
+  // accordingly.
+  explicit MiniBenchmarkTestHelper(
+#ifdef __ANDROID__
+      bool should_load_entrypoint_dynamically = true
+#else   // !__ANDROID__
+      bool should_load_entrypoint_dynamically = false
+#endif  // __ANDROID__
+  );
+  ~MiniBenchmarkTestHelper() = default;
+  bool should_perform_test() const { return should_perform_test_; }
+
+ private:
+  bool should_perform_test_;
+};
+
+}  // namespace acceleration
+}  // namespace tflite
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MINI_BENCHMARK_TEST_HELPER_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/BUILD

@@ -0,0 +1,137 @@
+load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+
+# Copyright 2022 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.
+# ==============================================================================
+load("//tensorflow/lite/experimental/acceleration/mini_benchmark:special_rules.bzl", "libjpeg_handle_deps", "register_selected_ops_deps")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = [
+        "//tensorflow/lite/experimental/acceleration/mini_benchmark:__subpackages__",
+    ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "grafter",
+    srcs = ["grafter.cc"],
+    hdrs = ["grafter.h"],
+    deps = [
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs_with_reflection",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:str_format",
+        "@flatbuffers",
+    ],
+)
+
+cc_library(
+    name = "embedder",
+    srcs = ["embedder.cc"],
+    hdrs = ["embedder.h"],
+    deps = [
+        ":validation_graph_builder",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:str_format",
+        "@flatbuffers",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:minimal_logging",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/kernels:builtin_ops",
+        "//tensorflow/lite/experimental/acceleration/mini_benchmark:decode_jpeg_status",
+        "//tensorflow/lite/experimental/acceleration/mini_benchmark:jpeg_common",
+        "//tensorflow/lite/experimental/acceleration/mini_benchmark:jpeg_header_parser",
+        # TODO(bekzhan): Remove duplicate dependency when only one of the two schemas is used.
+        "//tensorflow/lite/schema:schema_fbs",
+        "//tensorflow/lite/schema:schema_fbs_with_reflection",
+    ],
+)
+
+cc_library(
+    name = "validation_graph_builder",
+    srcs = ["validation_graph_builder.cc"],
+    hdrs = ["validation_graph_builder.h"],
+    deps = [
+        ":grafter",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:string_util",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:str_format",
+        "@flatbuffers",
+    ],
+)
+
+cc_binary(
+    name = "embedder_cmdline",
+    srcs = ["embedder_main.cc"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":embedder",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core:model_builder",
+        "//tensorflow/lite/core/kernels:builtin_ops",
+        "//tensorflow/lite/experimental/acceleration/mini_benchmark:call",
+        "//tensorflow/lite/experimental/acceleration/mini_benchmark:decode_jpeg",
+        "//tensorflow/lite/schema:schema_fbs",
+        "//tensorflow/lite/schema:schema_fbs_with_reflection",
+        "//tensorflow/lite/tools:command_line_flags",
+        "@com_google_absl//absl/strings",
+        "@flatbuffers",
+    ] + libjpeg_handle_deps() + register_selected_ops_deps(),
+)
+
+cc_library(
+    name = "custom_validation_embedder",
+    srcs = ["custom_validation_embedder.cc"],
+    hdrs = ["custom_validation_embedder.h"],
+    deps = [
+        "//tensorflow/lite:stderr_reporter",
+        "//tensorflow/lite/core/api:error_reporter",
+        "//tensorflow/lite/core/tools:verifier",
+        "//tensorflow/lite/experimental/acceleration/mini_benchmark:constants",
+        "//tensorflow/lite/experimental/acceleration/mini_benchmark:status_codes",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@flatbuffers",
+    ],
+)
+
+cc_test(
+    name = "custom_validation_embedder_test",
+    srcs = ["custom_validation_embedder_test.cc"],
+    deps = [
+        ":custom_validation_embedder",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/api:error_reporter",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/core/kernels:builtin_ops",
+        "//tensorflow/lite/experimental/acceleration/mini_benchmark:call",
+        "//tensorflow/lite/experimental/acceleration/mini_benchmark:embedded_mobilenet_model",
+        "//tensorflow/lite/experimental/acceleration/mini_benchmark:mini_benchmark_test_helper",
+        "//tensorflow/lite/experimental/acceleration/mini_benchmark:status_codes",
+        "//tensorflow/lite/kernels/internal:tensor",
+        "//tensorflow/lite/schema:schema_fbs",
+        "//tensorflow/lite/tools:model_loader",
+        "@com_google_googletest//:gtest_main",
+        "@flatbuffers//:runtime_cc",
+    ] + libjpeg_handle_deps() + register_selected_ops_deps(),
+)

tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/custom_validation_embedder.cc

@@ -0,0 +1,172 @@
+/* Copyright 2022 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/custom_validation_embedder.h"
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "flatbuffers/buffer.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "flatbuffers/vector.h"  // from @flatbuffers
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/core/tools/verifier.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/constants.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/status_codes.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace acceleration {
+namespace {
+
+using ::flatbuffers::FlatBufferBuilder;
+
+// Create options for call operator.
+flatbuffers::Offset<flatbuffers::Vector<uint8_t>> CallOpCustomOptions(
+    int primary_subgraph_index, int batch_size, FlatBufferBuilder& output) {
+  flexbuffers::Builder flexbuffer_builder;
+  flexbuffer_builder.Map([&] {
+    flexbuffer_builder.Int("subgraph_index", primary_subgraph_index);
+    flexbuffer_builder.Int("loop_count", batch_size);
+  });
+  flexbuffer_builder.Finish();
+  return output.CreateVector(flexbuffer_builder.GetBuffer());
+}
+
+}  // namespace
+
+void CustomValidationEmbedder::CreateTensorsFrom(
+    const SubGraph& from_subgraph, const std::vector<int>& from_indexes,
+    std::vector<std::vector<uint8_t>>* buffer_content,
+    flatbuffers::FlatBufferBuilder& fbb, std::vector<int>& new_indexes,
+    std::vector<flatbuffers::Offset<Buffer>>& buffers,
+    std::vector<flatbuffers::Offset<Tensor>>& tensors) {
+  int tensor_index_start = tensors.size();
+  for (int i = 0; i < from_indexes.size(); i++) {
+    TensorT base_tensor;
+    from_subgraph.tensors()->Get(from_indexes[i])->UnPackTo(&base_tensor);
+    if (!base_tensor.shape.empty() && base_tensor.shape[0] == 1) {
+      base_tensor.shape[0] = batch_size_;
+    }
+    if (!base_tensor.shape_signature.empty() &&
+        base_tensor.shape_signature[0] == 1) {
+      base_tensor.shape_signature[0] = batch_size_;
+    }
+    // Set the index in buffer.
+    base_tensor.buffer = buffers.size();
+
+    tensors.push_back(CreateTensor(fbb, &base_tensor));
+    new_indexes.push_back(tensor_index_start + i);
+    // If buffer content is provided, embed the content. Otherwise create an
+    // empty buffer.
+    if (buffer_content && !(*buffer_content)[i].empty()) {
+      buffers.push_back(
+          CreateBuffer(fbb, fbb.CreateVector((*buffer_content)[i])));
+    } else {
+      buffers.push_back(CreateBuffer(fbb));
+    }
+  }
+}
+
+MinibenchmarkStatus CustomValidationEmbedder::BuildModel(
+    const Model& main_model, flatbuffers::FlatBufferBuilder& fbb) {
+  ModelT main_model_obj;
+  main_model.UnPackTo(&main_model_obj);
+  if (main_model_obj.subgraphs[0]->inputs.size() != custom_input_.size()) {
+    TF_LITE_REPORT_ERROR(
+        error_reporter_,
+        "Unexpected custom_input size. Expected: %d. Actual: %d.",
+        main_model_obj.subgraphs[0]->inputs.size(), custom_input_.size());
+    return kMinibenchmarkValidationSubgraphBuildFailed;
+  }
+
+  // Copy all the data from main_model.
+  std::vector<flatbuffers::Offset<Metadata>> metadata;
+  metadata.reserve(main_model_obj.metadata.size());
+  for (auto& iter : main_model_obj.metadata) {
+    metadata.push_back(CreateMetadata(fbb, iter.get()));
+  }
+
+  std::vector<flatbuffers::Offset<SignatureDef>> signature_defs;
+  signature_defs.reserve(main_model_obj.signature_defs.size());
+  for (auto& iter : main_model_obj.signature_defs) {
+    signature_defs.push_back(CreateSignatureDef(fbb, iter.get()));
+  }
+
+  std::vector<flatbuffers::Offset<SubGraph>> subgraphs;
+  subgraphs.reserve(main_model_obj.subgraphs.size());
+  for (auto& iter : main_model_obj.subgraphs) {
+    subgraphs.push_back(CreateSubGraph(fbb, iter.get()));
+  }
+
+  std::vector<flatbuffers::Offset<Buffer>> buffers;
+  buffers.reserve(main_model_obj.buffers.size());
+  for (auto& iter : main_model_obj.buffers) {
+    buffers.push_back(CreateBuffer(fbb, iter.get()));
+  }
+
+  std::vector<flatbuffers::Offset<OperatorCode>> operator_codes;
+  operator_codes.reserve(main_model_obj.operator_codes.size());
+  for (auto& iter : main_model_obj.operator_codes) {
+    operator_codes.push_back(CreateOperatorCode(fbb, iter.get()));
+  }
+
+  // Create validation subgraph.
+  operator_codes.push_back(CreateOperatorCode(
+      fbb, BuiltinOperator_CUSTOM, fbb.CreateString("validation/call")));
+  int operator_code_index = operator_codes.size() - 1;
+
+  // Input and output tensors.
+  std::vector<flatbuffers::Offset<Tensor>> tensors;
+  std::vector<int32_t> input;
+  CreateTensorsFrom(*main_model.subgraphs()->Get(0),
+                    main_model_obj.subgraphs[0]->inputs, &custom_input_, fbb,
+                    input, buffers, tensors);
+
+  std::vector<int32_t> output;
+  CreateTensorsFrom(*main_model.subgraphs()->Get(0),
+                    main_model_obj.subgraphs[0]->outputs, nullptr, fbb, output,
+                    buffers, tensors);
+  auto input_offset = fbb.CreateVector(input);
+  auto output_offset = fbb.CreateVector(output);
+  std::vector<flatbuffers::Offset<Operator>> operators{CreateOperator(
+      fbb, operator_code_index, input_offset, output_offset,
+      tflite::BuiltinOptions_NONE, 0,
+      CallOpCustomOptions(/*primary_graph_index*/ 0, batch_size_, fbb),
+      tflite::CustomOptionsFormat_FLEXBUFFERS)};
+  subgraphs.push_back(
+      CreateSubGraph(fbb, fbb.CreateVector(tensors), input_offset,
+                     output_offset, fbb.CreateVector(operators),
+                     fbb.CreateString(std::string(kValidationGraphName))));
+
+  fbb.Finish(
+      CreateModel(fbb, kModelSchemaVersion, fbb.CreateVector(operator_codes),
+                  fbb.CreateVector(subgraphs),
+                  fbb.CreateString(main_model_obj.description),
+                  fbb.CreateVector(buffers),
+                  /* metadata_buffer */ 0, fbb.CreateVector(metadata),
+                  fbb.CreateVector(signature_defs)),
+      "TFL3");
+
+  if (Verify(fbb.GetBufferPointer(), fbb.GetSize(), error_reporter_)) {
+    return kMinibenchmarkSuccess;
+  } else {
+    return kMinibenchmarkValidationSubgraphBuildFailed;
+  }
+}
+
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/custom_validation_embedder.h

@@ -0,0 +1,95 @@
+/* Copyright 2022 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MODEL_MODIFIER_CUSTOM_VALIDATION_EMBEDDER_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MODEL_MODIFIER_CUSTOM_VALIDATION_EMBEDDER_H_
+
+#include <utility>
+#include <vector>
+
+#include "flatbuffers/buffer.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/status_codes.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/stderr_reporter.h"
+
+namespace tflite {
+namespace acceleration {
+
+// Create a model with custom validation graph.
+//
+// 'validation model' (new subgraph)
+// input (batch_size)
+//           |
+// +-----------------------+
+// |'main_model' (0)       |
+// | +---------------+     |
+// | |input          +---+ |
+// | +---------------+   | |
+// |                     ~ |
+// | +---------------+   | |
+// | |outputs        +<--+ |
+// | +---------------+     |
+// |                       |
+// +-----------------------+
+//           |
+// output (batch_size)
+//
+// The new model contains all the information from main_model, with an extra
+// subgraph for validation purposes. The validation graph calls the primary
+// subgraph with batch_size. The input data is embedded to the validation graph.
+// custom_input should have the same order as the input in the main_model. E.g.
+// custom_input[i] will be mapped to main_model.input[i].
+class CustomValidationEmbedder {
+ public:
+  CustomValidationEmbedder(
+      int batch_size, std::vector<std::vector<uint8_t>> custom_input,
+      ErrorReporter* error_reporter = DefaultErrorReporter())
+      : batch_size_(batch_size),
+        custom_input_(std::move(custom_input)),
+        error_reporter_(error_reporter) {}
+
+  // Move only.
+  CustomValidationEmbedder(CustomValidationEmbedder&&) = default;
+  CustomValidationEmbedder& operator=(CustomValidationEmbedder&&) = default;
+
+  // Build the final model with main_model and validation subgraph.
+  MinibenchmarkStatus BuildModel(const Model& main_model,
+                                 flatbuffers::FlatBufferBuilder& fbb);
+
+ private:
+  // Helper function to create tensors in validation graph based on primary
+  // subgraph. This function creates new tensors and buffers based on the
+  // from_subgraphs.tensors[from_indexes]. The new tensors will have shape[0]
+  // set to batch_size_, and indexes stored in new_indexes.
+  // New buffers will be created for each of the new tensors, and buffer data is
+  // copied from the corresponding buffer_content.
+  void CreateTensorsFrom(const SubGraph& from_subgraph,
+                         const std::vector<int>& from_indexes,
+                         std::vector<std::vector<uint8_t>>* buffer_content,
+                         flatbuffers::FlatBufferBuilder& fbb,
+                         std::vector<int>& new_indexes,
+                         std::vector<flatbuffers::Offset<Buffer>>& buffers,
+                         std::vector<flatbuffers::Offset<Tensor>>& tensors);
+
+  int batch_size_;
+  std::vector<std::vector<uint8_t>> custom_input_;
+  ErrorReporter* error_reporter_;
+};
+
+}  // namespace acceleration
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MODEL_MODIFIER_CUSTOM_VALIDATION_EMBEDDER_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/custom_validation_embedder_test.cc

@@ -0,0 +1,111 @@
+/* Copyright 2022 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/custom_validation_embedder.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/core/interpreter_builder.h"
+#include "tensorflow/lite/core/kernels/register.h"
+#include "tensorflow/lite/core/model_builder.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/call_register.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/embedded_mobilenet_model.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/mini_benchmark_test_helper.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/status_codes.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/tools/model_loader.h"
+
+namespace tflite {
+namespace acceleration {
+namespace {
+using ::flatbuffers::FlatBufferBuilder;
+
+constexpr int kMobileNetModelInputByteSize = 1 * 224 * 224 * 3;
+
+class CustomValidationEmbedderTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    std::string plain_model_path = MiniBenchmarkTestHelper::DumpToTempFile(
+        "mobilenet_quant.tflite",
+        g_tflite_acceleration_embedded_mobilenet_model,
+        g_tflite_acceleration_embedded_mobilenet_model_len);
+    ASSERT_TRUE(!plain_model_path.empty());
+    plain_model_loader_ =
+        std::make_unique<tools::PathModelLoader>(plain_model_path);
+    ASSERT_TRUE(plain_model_loader_->Init());
+  }
+
+  std::unique_ptr<tools::ModelLoader> plain_model_loader_;
+};
+
+TEST_F(CustomValidationEmbedderTest, BuildValidationModelSucceed) {
+  int batch_size = 5;
+  std::vector<uint8_t> input_buffer(batch_size * kMobileNetModelInputByteSize);
+  CustomValidationEmbedder embedder(batch_size, {input_buffer});
+
+  FlatBufferBuilder fbb;
+  EXPECT_EQ(
+      embedder.BuildModel(*plain_model_loader_->GetModel()->GetModel(), fbb),
+      kMinibenchmarkSuccess);
+
+  // Verify validation graph can be invoked.
+  auto model =
+      FlatBufferModel::BuildFromModel(GetModel(fbb.GetBufferPointer()));
+  auto interpreter = std::make_unique<Interpreter>();
+  auto resolver = std::make_unique<
+      ::tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates>();
+  resolver->AddCustom("validation/call", ops::Register_CALL(), 1);
+  ASSERT_EQ(InterpreterBuilder(*model, *resolver)(&interpreter), kTfLiteOk);
+  ASSERT_NE(interpreter, nullptr);
+  Subgraph* validation_graph = interpreter->subgraph(1);
+  EXPECT_THAT(input_buffer, testing::ElementsAreArray(
+                                GetTensorData<uint8_t>(validation_graph->tensor(
+                                    validation_graph->inputs()[0])),
+                                input_buffer.size()));
+  EXPECT_EQ(validation_graph->AllocateTensors(), kTfLiteOk);
+  EXPECT_EQ(validation_graph->Invoke(), kTfLiteOk);
+}
+
+TEST_F(CustomValidationEmbedderTest, BuildValidationModelTooManyInput) {
+  int batch_size = 5;
+  CustomValidationEmbedder embedder(batch_size, {{}, {}});
+
+  FlatBufferBuilder fbb;
+  EXPECT_EQ(
+      embedder.BuildModel(*plain_model_loader_->GetModel()->GetModel(), fbb),
+      kMinibenchmarkValidationSubgraphBuildFailed);
+}
+
+TEST_F(CustomValidationEmbedderTest, BuildValidationModelInvalidBufferSize) {
+  CustomValidationEmbedder embedder(2, {std::vector<uint8_t>(2, 2)});
+
+  FlatBufferBuilder fbb;
+  EXPECT_EQ(
+      embedder.BuildModel(*plain_model_loader_->GetModel()->GetModel(), fbb),
+      kMinibenchmarkValidationSubgraphBuildFailed);
+}
+
+}  // namespace
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/embedder.cc

@@ -0,0 +1,237 @@
+/* Copyright 2022 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/embedder.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_format.h"
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "flatbuffers/reflection_generated.h"  // from @flatbuffers
+#include "flatbuffers/vector.h"  // from @flatbuffers
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/core/interpreter_builder.h"
+#include "tensorflow/lite/core/kernels/register.h"
+#include "tensorflow/lite/core/model_builder.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg_status.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_common.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_header_parser.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/validation_graph_builder.h"
+#include "tensorflow/lite/logger.h"
+#include "tensorflow/lite/minimal_logging.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace fb = flatbuffers;
+
+namespace tflite {
+namespace acceleration {
+namespace {
+
+constexpr char kMetricPrefix[] = "metrics/";
+
+std::string DescribeShape(const fb::Vector<int32_t>* shape) {
+  std::string desc = "[";
+  for (int i = 0; i < shape->size(); i++) {
+    if (i != 0) {
+      desc += ", ";
+    }
+    desc += absl::StrFormat("%d", shape->Get(i));
+  }
+  desc += "]";
+  return desc;
+}
+
+}  // namespace
+
+Embedder::Embedder(const Model* main_model,
+                   const std::vector<std::string>& jpeg_data, float scale,
+                   int64_t zero_point, const Model* validation_model,
+                   const reflection::Schema* schema,
+                   bool use_ondevice_cpu_for_golden)
+    : main_model_(main_model),
+      jpeg_data_(jpeg_data),
+      scale_(scale),
+      zero_point_(zero_point),
+      validation_model_(validation_model),
+      schema_(schema),
+      use_ondevice_cpu_for_golden_(use_ondevice_cpu_for_golden) {}
+
+absl::Status Embedder::ValidateInputs() {
+#define VALIDATE(condition, ...)                                     \
+  if (!(condition)) {                                                \
+    return absl::InvalidArgumentError(absl::StrFormat(__VA_ARGS__)); \
+  }
+  VALIDATE(main_model_, "main_model may not be null");
+  VALIDATE(main_model_->subgraphs()->size(), "main model must have subgraphs");
+  const SubGraph* main_subgraph = main_model_->subgraphs()->Get(0);
+  VALIDATE(main_subgraph->inputs()->size() == 1,
+           "main subgraph must have 1 input (got %d)",
+           main_subgraph->inputs()->size());
+  const auto* shape =
+      main_subgraph->tensors()->Get(main_subgraph->inputs()->Get(0))->shape();
+  VALIDATE(shape->size() == 4,
+           "main subgraph input must have 4 dimensions (got %d)",
+           shape->size());
+  jpeg_output_channels_ = shape->Get(3);
+  VALIDATE(shape->Get(0) == 1,
+           "main subgraph input must have batch size 1 (got %d)",
+           shape->Get(0));
+  VALIDATE(shape->Get(3) == 1 || shape->Get(3) == 3 || shape->Get(3) == 4,
+           "main subgraph input must have 1 or 3 or 4 channels (got %d)",
+           shape->Get(3));
+
+  VALIDATE(!jpeg_data_.empty(), "must have at least 1 jpeg input");
+  int jpeg_number = 0;
+  for (const std::string& jpeg_image_data : jpeg_data_) {
+    int width, height, components;
+    decode_jpeg_kernel::JpegHeader header{0};
+    auto status = decode_jpeg_kernel::ReadJpegHeader(
+        {jpeg_image_data.data(), jpeg_image_data.size()}, &header);
+    VALIDATE(status.code == kTfLiteOk,
+             "Failed to decompress jpeg data at index %d: %s", jpeg_number,
+             status.error_message.c_str());
+    width = header.width;
+    height = header.height;
+    components = header.channels;
+    VALIDATE(height == shape->Get(1) && width == shape->Get(2) &&
+                 (components == shape->Get(3) ||
+                  // A workaround to allow RGBA channels extracted from RGB
+                  // images with alpha channel as 255 (fully opaque) by default.
+                  components == 3 && shape->Get(3) == 4),
+             "Jpeg input at index %d has different size from input tensor "
+             "(jpeg h: %d, w: %d, c: %d; tensor h: %d, w: %d, c: %d)",
+             jpeg_number, height, width, components, shape->Get(1),
+             shape->Get(2), shape->Get(3));
+    if (components < shape->Get(3)) {
+      TFLITE_LOG_PROD(TFLITE_LOG_INFO,
+                      "Jpeg input at index %d has %d channels. Lower than the "
+                      "expected %d channels.",
+                      jpeg_number, components, shape->Get(3));
+    }
+    jpeg_number++;
+  }
+
+  int main_output_count = main_subgraph->outputs()->size();
+  VALIDATE(main_output_count > 0,
+           "main subgraph must have at least 1 output (got %d)",
+           main_output_count);
+  VALIDATE(validation_model_->subgraphs()->size(),
+           "validation model must have subgraphs");
+  const SubGraph* validation_subgraph = validation_model_->subgraphs()->Get(0);
+  int validation_input_count = validation_subgraph->inputs()->size();
+  VALIDATE(
+      validation_input_count == main_output_count * 2,
+      "validation subgraph input count must be 2 times main subgraph output "
+      "count (validation input count: %d, main subgraph output count: %d)",
+      validation_input_count, main_output_count);
+  for (int i = 0; i < main_output_count; i++) {
+    auto main_output_tensor =
+        main_subgraph->tensors()->Get(main_subgraph->outputs()->Get(i));
+    auto main_output_shape = DescribeShape(main_output_tensor->shape());
+    VALIDATE(main_output_shape != "[]",
+             "All main outputs must be tensors, %d is a scalar", i);
+    VALIDATE(main_output_tensor->name()->str().find(kMetricPrefix) != 0,
+             "Main output %d name %s clashes with metrics/ prefix", i,
+             main_output_tensor->name()->c_str());
+    auto validation_input_shape_1 =
+        DescribeShape(validation_subgraph->tensors()
+                          ->Get(validation_subgraph->inputs()->Get(i))
+                          ->shape());
+    auto validation_input_shape_2 = DescribeShape(
+        validation_subgraph->tensors()
+            ->Get(validation_subgraph->inputs()->Get(main_output_count + i))
+            ->shape());
+    VALIDATE(main_output_shape == validation_input_shape_1,
+             "Main output %d dimensions %s do not match validation input %d "
+             "dimensions %s",
+             i, main_output_shape, i, validation_input_shape_1);
+    VALIDATE(main_output_shape == validation_input_shape_2,
+             "Main output %d dimensions %s do not match validation input %d "
+             "dimensions %s",
+             i, main_output_shape, main_output_count + i,
+             validation_input_shape_2);
+  }
+  int validation_output_count = validation_subgraph->outputs()->size();
+  VALIDATE(validation_output_count >= 2,
+           "validation output count must be at least 2 (got "
+           "%d)",
+           validation_output_count);
+  bool seen_ok = false;
+  const std::string kOk = "ok";
+  const std::string kPrefixedOk = kMetricPrefix + kOk;
+  std::string names = "";
+  for (int i = 0; i < validation_output_count; i++) {
+    const Tensor* t = validation_subgraph->tensors()->Get(
+        validation_subgraph->outputs()->Get(i));
+    VALIDATE(t->shape()->size(),
+             "validation outputs must be tensors, %d is a scalar", i);
+    seen_ok = (seen_ok || (kOk == t->name()->str()) ||
+               (kPrefixedOk == t->name()->str()));
+    if (i != 0) {
+      names += ", ";
+    }
+    names += t->name()->str();
+  }
+  VALIDATE(seen_ok, "validation must have an output named 'ok' (saw %s)",
+           names);
+#undef VALIDATE
+  return absl::OkStatus();
+}
+
+absl::Status Embedder::CreateModelWithEmbeddedValidation(
+    fb::FlatBufferBuilder* fbb, ops::builtin::BuiltinOpResolver* resolver) {
+  auto status = ValidateInputs();
+  if (!status.ok()) {
+    return status;
+  }
+  fb::FlatBufferBuilder intermediate_fbb;
+  ValidationGraphBuilder builder(
+      kMetricPrefix, main_model_, jpeg_data_, jpeg_output_channels_, scale_,
+      zero_point_, validation_model_, schema_, use_ondevice_cpu_for_golden_);
+  status = builder.BuildIntermediateModel(&intermediate_fbb);
+  if (!status.ok()) {
+    return status;
+  }
+  auto intermediate_model = FlatBufferModel::VerifyAndBuildFromBuffer(
+      reinterpret_cast<const char*>(intermediate_fbb.GetBufferPointer()),
+      intermediate_fbb.GetSize());
+  if (!intermediate_model) {
+    return absl::InternalError("Failed to load intermediate model");
+  }
+  std::unique_ptr<Interpreter> interpreter;
+  InterpreterBuilder(*intermediate_model, *resolver)(&interpreter);
+  if (!interpreter) {
+    return absl::InternalError(
+        "Failed to build interpreter from intermediate model");
+  }
+  Subgraph* subgraph = interpreter->subgraph(1);
+  if (subgraph->AllocateTensors() != kTfLiteOk) {
+    return absl::InternalError(
+        "Failed to AllocateTensors() on validation subgraph of intermediate "
+        "model");
+  }
+  if (subgraph->Invoke() != kTfLiteOk) {
+    return absl::InternalError(
+        "Failed to Invoke() on validation subgraph of intermediate model");
+  }
+  return builder.BuildFinalModel(fbb, subgraph);
+}
+
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/embedder.h

@@ -0,0 +1,112 @@
+/* Copyright 2022 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MODEL_MODIFIER_EMBEDDER_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MODEL_MODIFIER_EMBEDDER_H_
+
+#include <stdint.h>
+
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "flatbuffers/reflection_generated.h"  // from @flatbuffers
+#include "tensorflow/lite/core/kernels/register.h"
+#include "tensorflow/lite/schema/reflection/schema_generated.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace acceleration {
+// Class to embed a mini-benchmark into a tflite file.
+//
+// The inputs are:
+// - 'main_model': the actual inference graph (e.g., mobilenet classifier)
+// - 'jpeg_data': jpeg images used as test data.
+// - 'validation_model': a graph that takes as input two sets of values (the
+// known-good main model output and the to-be-tested main model output) and
+// produces 2 or more outputs where one must be called 'ok' (whether the
+// results are good enough) and rest are metrics that were used to determine
+// 'ok' and can be used for debugging/telemetry.
+// (Known good outputs are produced inside this class, i.e. running TFLite CPU
+// on the build host).
+//
+// The output is:
+// - A new benchmark model which has 3 subgraphs. The 'main_model' subgraph, a
+// new 'validate' subgraph that invokes the other two subgraphs when required,
+// and the 'validation_model' subgraph.
+// - The model output is the output of 'validation_model' + output of
+// 'main_model'
+// - This model has additional buffers that store the 'jpeg_data' and the actual
+// outputs.
+// - The 'main_model' subgraph is fed the 'jpeg_data' and produces an output
+// which is used by the 'validation_model' with the known-good outputs to
+// evaluate the model.
+// - This entire process is handled end-to-end by the 'validate' subgraph using
+// two custom ops: 'validate/call' (implemented in :call in this directory) and
+// 'validate/decode_jpeg' (being implemented).
+//
+// Constraints on inputs:
+// - 'main_model' must have a single input of dimensions
+//   [1, height, width, 1 or 3]
+// - the images encoded in 'jpeg_data' must have same height, width and channels
+//   as 'main_model' input
+// - the 'validation_model' must have inputs equal to 'main_model' outputs
+//   duplicated (e.g, if 'main_model' has outputs with dimensions
+//   [1, 10] and [1, 20]; the 'validation_model' must have inputs with
+//   dimensions [1, 10], [1, 20], [1, 10], [1, 20]).
+// - the 'validation_model' must have 2 or more outputs, and one of them must be
+//   called 'ok'.
+// - all inputs and outputs must be tensors (not scalars).
+//
+// TODO(b/172541832):
+// - Mark the validation graph so that it's not delegated in the inference case.
+// - Allow known-good outputs to be given rather than always being calculated
+// inside this class.
+class Embedder {
+ public:
+  // Construct Embedder with inputs. The Model* inputs are owned by the caller
+  // and must outlive the Embedder. The `schema` must contain the tflite
+  // flatbuffer schema. If the model is quantized, scale and zero_point are
+  // ignored.
+  Embedder(const Model* main_model, const std::vector<std::string>& jpeg_data,
+           float scale, int64_t zero_point, const Model* validation_model,
+           const reflection::Schema* schema,
+           bool use_ondevice_cpu_for_golden = false);
+  // Construct the output model. Calls Finish() on 'fbb'.
+  // The 'resolver' must have the call and decode_jpeg ops from this directory
+  // registered as 'validation/call' and 'validation/decode_jpeg'.
+  absl::Status CreateModelWithEmbeddedValidation(
+      flatbuffers::FlatBufferBuilder* fbb,
+      ::tflite::ops::builtin::BuiltinOpResolver* resolver);
+  // Check that the inputs fulfill the constraints. Called automatically as part
+  // of CreateModelWithEmbeddedValidation.
+  absl::Status ValidateInputs();
+
+ private:
+  const Model* main_model_;
+  std::vector<std::string> jpeg_data_;
+  int32_t jpeg_output_channels_;
+  float scale_;
+  int64_t zero_point_;
+  const Model* validation_model_;
+  const reflection::Schema* schema_;
+  bool use_ondevice_cpu_for_golden_;
+};
+
+}  // namespace acceleration
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MODEL_MODIFIER_EMBEDDER_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/embedder_main.cc

@@ -0,0 +1,200 @@
+/* Copyright 2022 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.
+==============================================================================*/
+// Command line tool for embedding validation data in tflite models.
+#include <cstdint>
+#include <fstream>
+#include <iostream>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/strings/str_split.h"
+#include "flatbuffers/base.h"  // from @flatbuffers
+#include "flatbuffers/buffer.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "flatbuffers/idl.h"  // from @flatbuffers
+#include "flatbuffers/reflection_generated.h"  // from @flatbuffers
+#include "flatbuffers/util.h"  // from @flatbuffers
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/core/interpreter_builder.h"
+#include "tensorflow/lite/core/kernels/register.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#if FLATBUFFERS_LITTLEENDIAN == 0
+#include "tensorflow/lite/core/model_builder.h"
+#endif
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/call_register.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg_register.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/embedder.h"
+#include "tensorflow/lite/schema/reflection/schema_generated.h"
+#include "tensorflow/lite/tools/benchmark/register_custom_op.h"
+#include "tensorflow/lite/tools/command_line_flags.h"
+
+namespace tflite {
+namespace acceleration {
+struct EmbedderOptions {
+  std::string schema, main_model, metrics_model, output, jpegs_arg;
+  float scale = 0.;
+  int64_t zero_point = -1;
+  bool use_ondevice_cpu_for_golden = false;
+};
+
+int RunEmbedder(const EmbedderOptions& options) {
+  // Load schema.
+  std::string fbs_contents;
+  if (!flatbuffers::LoadFile(options.schema.c_str(), false, &fbs_contents)) {
+    std::cerr << "Unable to load schema file " << options.schema << std::endl;
+    return 1;
+  }
+  const char* include_directories[] = {nullptr};
+  flatbuffers::Parser schema_parser;
+  if (!schema_parser.Parse(fbs_contents.c_str(), include_directories)) {
+    std::cerr << "Unable to parse schema " << schema_parser.error_ << std::endl;
+    return 2;
+  }
+  schema_parser.Serialize();
+  const reflection::Schema* schema =
+      reflection::GetSchema(schema_parser.builder_.GetBufferPointer());
+
+  // Load main model.
+  std::string main_model_contents;
+  if (!flatbuffers::LoadFile(options.main_model.c_str(), false,
+                             &main_model_contents)) {
+    std::cerr << "Unable to load main model file " << options.main_model
+              << std::endl;
+    return 3;
+  }
+#if FLATBUFFERS_LITTLEENDIAN == 0
+  tflite::FlatBufferModel::ByteSwapSerializedModel(&main_model_contents, false);
+#endif
+  const Model* main_model =
+      flatbuffers::GetRoot<Model>(main_model_contents.data());
+
+  // Load metrics model.
+  std::string metrics_model_contents;
+  if (!flatbuffers::LoadFile(options.metrics_model.c_str(), false,
+                             &metrics_model_contents)) {
+    std::cerr << "Unable to load metrics model file " << options.metrics_model
+              << std::endl;
+    return 4;
+  }
+#if FLATBUFFERS_LITTLEENDIAN == 0
+  tflite::FlatBufferModel::ByteSwapSerializedModel(&metrics_model_contents,
+                                                   false);
+#endif
+  const Model* metrics_model =
+      flatbuffers::GetRoot<Model>(metrics_model_contents.data());
+
+  // Load sample images.
+  std::vector<std::string> jpeg_paths = absl::StrSplit(options.jpegs_arg, ',');
+  std::vector<std::string> jpeg_data;
+  for (const std::string& jpeg_path : jpeg_paths) {
+    std::string data;
+    if (!flatbuffers::LoadFile(jpeg_path.c_str(), false, &data)) {
+      std::cerr << "Unable to load jpeg file '" << jpeg_path << "'"
+                << std::endl;
+      return 5;
+    }
+    jpeg_data.push_back(data);
+  }
+
+  // Create model with embedded validation.
+  tflite::acceleration::Embedder embedder(
+      main_model, jpeg_data, options.scale, options.zero_point, metrics_model,
+      schema, options.use_ondevice_cpu_for_golden);
+  flatbuffers::FlatBufferBuilder fbb;
+  ::tflite::ops::builtin::BuiltinOpResolver resolver;
+  resolver.AddCustom("validation/call",
+                     ::tflite::acceleration::ops::Register_CALL(), 1);
+  resolver.AddCustom(
+      "validation/decode_jpeg",
+      ::tflite::acceleration::decode_jpeg_kernel::Register_DECODE_JPEG(), 1);
+
+  RegisterSelectedOps(&resolver);
+
+  auto status = embedder.CreateModelWithEmbeddedValidation(&fbb, &resolver);
+  if (!status.ok()) {
+    std::cerr << "Creating model with embedded validation failed: "
+              << status.ToString() << std::endl;
+    return 6;
+  }
+
+  // Write created model to output path.
+  std::string binary(reinterpret_cast<const char*>(fbb.GetBufferPointer()),
+                     fbb.GetSize());
+  std::ofstream f;
+  f.open(options.output);
+  if (!f.good()) {
+    std::cerr << "Opening " << options.output
+              << " for writing failed: " << strerror(errno) << std::endl;
+    return 7;
+  }
+#if FLATBUFFERS_LITTLEENDIAN == 0
+  tflite::FlatBufferModel::ByteSwapSerializedModel(&binary, true);
+#endif
+  f << binary;
+  f.close();
+  if (!f.good()) {
+    std::cerr << "Writing to " << options.output
+              << " failed: " << strerror(errno) << std::endl;
+    return 8;
+  }
+
+  const Model* model = flatbuffers::GetRoot<Model>(fbb.GetBufferPointer());
+  std::unique_ptr<Interpreter> interpreter;
+  if (InterpreterBuilder(model, resolver)(&interpreter) != kTfLiteOk) {
+    std::cerr << "Loading the created model failed" << std::endl;
+    return 9;
+  }
+
+  return 0;
+}
+
+}  // namespace acceleration
+}  // namespace tflite
+
+int main(int argc, char* argv[]) {
+  tflite::acceleration::EmbedderOptions options;
+  std::vector<tflite::Flag> flags = {
+      tflite::Flag::CreateFlag("schema", &options.schema,
+                               "Path to tflite schema.fbs"),
+      tflite::Flag::CreateFlag("main_model", &options.main_model,
+                               "Path to main inference tflite model"),
+      tflite::Flag::CreateFlag("metrics_model", &options.metrics_model,
+                               "Path to metrics tflite model"),
+      tflite::Flag::CreateFlag("output", &options.output,
+                               "Path to tflite output file"),
+      tflite::Flag::CreateFlag(
+          "jpegs", &options.jpegs_arg,
+          "Comma-separated list of jpeg files to use as input"),
+      tflite::Flag::CreateFlag("scale", &options.scale,
+                               "Scale to use when dequantizing input images"),
+      tflite::Flag::CreateFlag(
+          "zero_point", &options.zero_point,
+          "Zero-point to use when dequantizing input images"),
+      tflite::Flag::CreateFlag(
+          "use_ondevice_cpu_for_golden", &options.use_ondevice_cpu_for_golden,
+          "Use on-device CPU as golden data (rather than embedding golden "
+          "data)"),
+  };
+  if (!tflite::Flags::Parse(&argc, const_cast<const char**>(argv), flags) ||
+      options.schema.empty() || options.main_model.empty() ||
+      options.output.empty() || options.jpegs_arg.empty()) {
+    std::cerr << tflite::Flags::Usage("embedder_cmdline", flags);
+    return 1;
+  }
+  return tflite::acceleration::RunEmbedder(options);
+
+  return 0;
+}

tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/grafter.cc

@@ -0,0 +1,300 @@
+/* Copyright 2022 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/grafter.h"
+
+#include <stdint.h>
+
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "flatbuffers/buffer.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "flatbuffers/reflection.h"  // from @flatbuffers
+#include "flatbuffers/reflection_generated.h"  // from @flatbuffers
+#include "flatbuffers/table.h"  // from @flatbuffers
+#include "flatbuffers/vector.h"  // from @flatbuffers
+#include "tensorflow/lite/schema/reflection/schema_generated.h"
+
+namespace fb = flatbuffers;
+
+namespace tflite {
+namespace acceleration {
+
+namespace {
+
+class Combiner : FlatbufferHelper {
+ public:
+  Combiner(flatbuffers::FlatBufferBuilder* fbb,
+           std::vector<const Model*> models,
+           std::vector<std::string> subgraph_names,
+           const reflection::Schema* schema)
+      : FlatbufferHelper(fbb, schema),
+        fbb_(fbb),
+        models_(models),
+        subgraph_names_(subgraph_names),
+        schema_(schema) {}
+  absl::Status Combine() {
+    auto operator_codes = OperatorCodes();
+    if (!operator_codes.ok()) {
+      return operator_codes.status();
+    }
+    auto subgraphs = SubGraphs();
+    if (!subgraphs.ok()) {
+      return subgraphs.status();
+    }
+    auto buffers = Buffers();
+    if (!buffers.ok()) {
+      return buffers.status();
+    }
+    auto metadata = Metadatas();
+    if (!metadata.ok()) {
+      return metadata.status();
+    }
+    auto signature_defs = SignatureDefs();
+    if (!signature_defs.ok()) {
+      return signature_defs.status();
+    }
+    fb::Offset<Model> model = CreateModel(
+        *fbb_, 3, *operator_codes, *subgraphs,
+        fbb_->CreateString(models_[0]->description()->str()), *buffers,
+        /* metadata_buffer */ 0, *metadata, *signature_defs);
+    fbb_->Finish(model, "TFL3");
+    return absl::OkStatus();
+  }
+
+ private:
+  absl::StatusOr<fb::Offset<fb::Vector<fb::Offset<OperatorCode>>>>
+  OperatorCodes() {
+    std::vector<fb::Offset<OperatorCode>> codes;
+    for (const Model* model : models_) {
+      for (int i = 0; i < model->operator_codes()->size(); i++) {
+        auto status = CopyTableToVector(
+            "tflite.OperatorCode", model->operator_codes()->Get(i), &codes);
+        if (!status.ok()) {
+          return status;
+        }
+      }
+    }
+    return fbb_->CreateVector(codes);
+  }
+  absl::StatusOr<fb::Offset<fb::Vector<fb::Offset<SubGraph>>>> SubGraphs() {
+    std::vector<fb::Offset<SubGraph>> graphs;
+    int buffer_offset = 0;
+    int operator_code_offset = 0;
+    int subgraph_index = 0;
+    for (const Model* model : models_) {
+      if (model->subgraphs()->size() != 1) {
+        return absl::InvalidArgumentError(
+            "Every model to be combined must have a single subgraph.");
+      }
+      auto graph =
+          AdjustSubGraph(model->subgraphs()->Get(0), buffer_offset,
+                         operator_code_offset, subgraph_names_[subgraph_index]);
+      if (!graph.ok()) {
+        return graph.status();
+      }
+      graphs.push_back(*graph);
+      buffer_offset += model->buffers()->size();
+      operator_code_offset += model->operator_codes()->size();
+      ++subgraph_index;
+    }
+    return fbb_->CreateVector(graphs);
+  }
+  absl::StatusOr<fb::Offset<fb::Vector<fb::Offset<Buffer>>>> Buffers() {
+    std::vector<fb::Offset<Buffer>> buffers;
+    for (const Model* model : models_) {
+      for (int i = 0; i < model->buffers()->size(); i++) {
+        auto status = CopyTableToVector("tflite.Buffer",
+                                        model->buffers()->Get(i), &buffers);
+        if (!status.ok()) {
+          return status;
+        }
+      }
+    }
+    return fbb_->CreateVector(buffers);
+  }
+  absl::StatusOr<fb::Offset<fb::Vector<fb::Offset<Metadata>>>> Metadatas() {
+    std::vector<fb::Offset<Metadata>> metadatas;
+    int buffer_offset = 0;
+    for (const Model* model : models_) {
+      for (int i = 0; model->metadata() && i < model->metadata()->size(); i++) {
+        auto metadata =
+            AdjustMetadata(model->metadata()->Get(i), buffer_offset);
+        if (!metadata.ok()) {
+          return metadata.status();
+        }
+        metadatas.push_back(*metadata);
+        buffer_offset += model->buffers()->size();
+      }
+    }
+    return fbb_->CreateVector(metadatas);
+  }
+  absl::StatusOr<fb::Offset<fb::Vector<fb::Offset<SignatureDef>>>>
+  SignatureDefs() {
+    std::vector<fb::Offset<SignatureDef>> signature_defs;
+    const Model* model = models_[0];
+    for (int i = 0;
+         model->signature_defs() && i < model->signature_defs()->size(); i++) {
+      auto status =
+          CopyTableToVector("tflite.SignatureDef",
+                            model->signature_defs()->Get(i), &signature_defs);
+      if (!status.ok()) {
+        return status;
+      }
+    }
+    return fbb_->CreateVector(signature_defs);
+  }
+
+  absl::StatusOr<fb::Offset<SubGraph>> AdjustSubGraph(const SubGraph* graph,
+                                                      int buffer_offset,
+                                                      int operator_code_offset,
+                                                      const std::string& name) {
+    auto tensors = AdjustTensors(graph, buffer_offset);
+    if (!tensors.ok()) {
+      return tensors.status();
+    }
+    auto ops = AdjustOps(graph, operator_code_offset);
+    if (!ops.ok()) {
+      return ops.status();
+    }
+    return CreateSubGraph(*fbb_, fbb_->CreateVector(*tensors),
+                          CopyIntVector(graph->inputs()),
+                          CopyIntVector(graph->outputs()),
+                          fbb_->CreateVector(*ops), fbb_->CreateString(name));
+  }
+
+  absl::StatusOr<std::vector<fb::Offset<Operator>>> AdjustOps(
+      const SubGraph* graph, int operator_code_offset) {
+    std::vector<fb::Offset<Operator>> ops;
+    auto op_object = FindObject("tflite.Operator");
+    const reflection::Field* builtin_options_field = nullptr;
+    for (auto it = op_object->fields()->cbegin();
+         it != op_object->fields()->cend(); it++) {
+      auto candidate = *it;
+      if (candidate->name()->str() == "builtin_options") {
+        builtin_options_field = candidate;
+        break;
+      }
+    }
+    if (!builtin_options_field) {
+      return absl::UnknownError(
+          "Wasn't able to find the builtin_options field on tflite.Operator");
+    }
+    for (int i = 0; i < graph->operators()->size(); i++) {
+      const Operator* op = graph->operators()->Get(i);
+      fb::Offset<void> copied_builtin_options = 0;
+      if (op->builtin_options() != nullptr) {
+        const fb::Table* opt = (const fb::Table*)op;  // NOLINT
+        auto& builtin_options_object = fb::GetUnionType(
+            *schema_, *op_object, *builtin_options_field, *opt);
+        copied_builtin_options =
+            fb::CopyTable(*fbb_, *schema_, builtin_options_object,
+                          *fb::GetFieldT(*opt, *builtin_options_field))
+                .o;
+      }
+      ops.push_back(CreateOperator(
+          *fbb_, op->opcode_index() + operator_code_offset,
+          CopyIntVector(op->inputs()), CopyIntVector(op->outputs()),
+          op->builtin_options_type(), copied_builtin_options,
+          CopyIntVector(op->custom_options()), op->custom_options_format(),
+          CopyIntVector(op->mutating_variable_inputs()),
+          CopyIntVector(op->intermediates())));
+    }
+    return ops;
+  }
+
+  absl::StatusOr<std::vector<fb::Offset<Tensor>>> AdjustTensors(
+      const SubGraph* graph, int buffer_offset) {
+    std::vector<fb::Offset<Tensor>> tensors;
+    auto orig_tensors = graph->tensors();
+    for (auto iter = orig_tensors->cbegin(); iter != orig_tensors->cend();
+         iter++) {
+      auto i = *iter;
+      std::vector<int32_t> shape{i->shape()->cbegin(), i->shape()->cend()};
+      std::vector<int32_t> shape_signature;
+      if (i->shape_signature()) {
+        shape_signature.assign(i->shape_signature()->cbegin(),
+                               i->shape_signature()->cend());
+      }
+      auto quantization =
+          CopyTable("tflite.QuantizationParameters", i->quantization());
+      if (!quantization.ok()) {
+        return quantization.status();
+      }
+      auto sparsity = CopyTable("tflite.SparsityParameters", i->sparsity());
+      if (!sparsity.ok()) {
+        return sparsity.status();
+      }
+      tensors.push_back(CreateTensor(
+          *fbb_, fbb_->CreateVector(shape), i->type(),
+          i->buffer() + buffer_offset, fbb_->CreateString(i->name()->str()),
+          *quantization, i->is_variable(), *sparsity,
+          shape_signature.empty() ? 0 : fbb_->CreateVector(shape_signature)));
+    }
+    return tensors;
+  }
+
+  absl::StatusOr<fb::Offset<Metadata>> AdjustMetadata(const Metadata* metadata,
+                                                      int buffer_offset) {
+    return CreateMetadata(*fbb_,
+                          metadata->name()
+                              ? fbb_->CreateString(metadata->name()->str())
+                              : 0,
+                          metadata->buffer())
+        .o;
+  }
+
+  flatbuffers::FlatBufferBuilder* fbb_;
+  std::vector<const Model*> models_;
+  std::vector<std::string> subgraph_names_;
+  const reflection::Schema* schema_;
+};
+
+}  // namespace
+
+absl::Status CombineModels(flatbuffers::FlatBufferBuilder* fbb,
+                           std::vector<const Model*> models,
+                           std::vector<std::string> subgraph_names,
+                           const reflection::Schema* schema) {
+  if (!fbb || !schema) {
+    return absl::InvalidArgumentError(
+        "Must provide FlatBufferBuilder and Schema");
+  }
+  if (models.size() < 2) {
+    return absl::InvalidArgumentError("Must have 2+ models to combine");
+  }
+  Combiner combiner(fbb, models, subgraph_names, schema);
+  return combiner.Combine();
+}
+
+FlatbufferHelper::FlatbufferHelper(flatbuffers::FlatBufferBuilder* fbb,
+                                   const reflection::Schema* schema)
+    : fbb_(fbb), schema_(schema) {}
+
+const reflection::Object* FlatbufferHelper::FindObject(
+    const std::string& name) {
+  for (auto candidate = schema_->objects()->cbegin();
+       candidate != schema_->objects()->cend(); candidate++) {
+    if (candidate->name()->str() == name) {
+      return *candidate;
+    }
+  }
+  return nullptr;
+}
+
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/grafter.h

@@ -0,0 +1,117 @@
+/* Copyright 2022 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MODEL_MODIFIER_GRAFTER_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MODEL_MODIFIER_GRAFTER_H_
+
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_format.h"
+#include "flatbuffers/buffer.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "flatbuffers/idl.h"  // from @flatbuffers
+#include "flatbuffers/reflection.h"  // from @flatbuffers
+#include "flatbuffers/reflection_generated.h"  // from @flatbuffers
+#include "flatbuffers/table.h"  // from @flatbuffers
+#include "flatbuffers/vector.h"  // from @flatbuffers
+
+namespace tflite {
+struct Model;
+}  // namespace tflite
+
+namespace tflite {
+namespace acceleration {
+
+// Combines the given models into one, using the FlatBufferBuilder.
+//
+// This is useful for constructing models that contain validation data and
+// metrics.
+//
+// The model fields are handled as follows:
+// - version is set to 3
+// - operator codes are concatenated (no deduplication)
+// - subgraphs are concatenated in order, rewriting operator and buffer indices
+// to match the combined model. Subgraph names are set from 'subgraph_names'
+// - description is taken from first model
+// - buffers are concatenated
+// - metadata buffer is left unset
+// - metadata are concatenated
+// - signature_defs are taken from the first model (as they refer to the main
+// subgraph).
+absl::Status CombineModels(flatbuffers::FlatBufferBuilder* fbb,
+                           std::vector<const Model*> models,
+                           std::vector<std::string> subgraph_names,
+                           const reflection::Schema* schema);
+
+// Convenience methods for copying flatbuffer Tables and Vectors.
+//
+// These are used by CombineModels above, but also needed for constructing
+// validation subgraphs to be combined with models.
+class FlatbufferHelper {
+ public:
+  FlatbufferHelper(flatbuffers::FlatBufferBuilder* fbb,
+                   const reflection::Schema* schema);
+  template <typename T>
+  absl::Status CopyTableToVector(const std::string& name, const T* o,
+                                 std::vector<flatbuffers::Offset<T>>* v) {
+    auto copied = CopyTable(name, o);
+    if (!copied.ok()) {
+      return copied.status();
+    }
+    v->push_back(*copied);
+    return absl::OkStatus();
+  }
+  template <typename T>
+  absl::StatusOr<flatbuffers::Offset<T>> CopyTable(const std::string& name,
+                                                   const T* o) {
+    if (o == nullptr) return 0;
+    const reflection::Object* def = FindObject(name);
+    if (!def) {
+      return absl::NotFoundError(
+          absl::StrFormat("Type %s not found in schema", name));
+    }
+    // We want to use the general copying mechanisms that operate on
+    // flatbuffers::Table pointers. Flatbuffer types are not directly
+    // convertible to Table, as they inherit privately from table.
+    // For type* -> Table*, use reinterpret cast.
+    const flatbuffers::Table* ot =
+        reinterpret_cast<const flatbuffers::Table*>(o);
+    // For Offset<Table *> -> Offset<type>, rely on uoffset_t conversion to
+    // any flatbuffers::Offset<T>.
+    return flatbuffers::CopyTable(*fbb_, *schema_, *def, *ot).o;
+  }
+  template <typename int_type>
+  flatbuffers::Offset<flatbuffers::Vector<int_type>> CopyIntVector(
+      const flatbuffers::Vector<int_type>* from) {
+    if (from == nullptr) {
+      return 0;
+    }
+    std::vector<int_type> v{from->cbegin(), from->cend()};
+    return fbb_->CreateVector(v);
+  }
+  const reflection::Object* FindObject(const std::string& name);
+
+ private:
+  flatbuffers::FlatBufferBuilder* fbb_;
+  const reflection::Schema* schema_;
+};
+
+}  // namespace acceleration
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MODEL_MODIFIER_GRAFTER_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/validation_graph_builder.cc

@@ -0,0 +1,570 @@
+/* Copyright 2022 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/validation_graph_builder.h"
+
+#include <stdint.h>
+
+#include <functional>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_format.h"
+#include "flatbuffers/buffer.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "flatbuffers/vector.h"  // from @flatbuffers
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/grafter.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_util.h"
+
+namespace tflite {
+namespace acceleration {
+
+using ::flatbuffers::FlatBufferBuilder;
+using ::flatbuffers::GetRoot;
+
+absl::Status ValidationGraphBuilder::BuildIntermediateModel(
+    FlatBufferBuilder* fbb) {
+  fbb_.Reset();
+  auto model = MakeModel(/* intermediate_only */ true,
+                         /* subgraph_with_golden_outputs */ nullptr);
+  if (!model.ok()) {
+    return model.status();
+  }
+  fbb_.Finish(*model, "TFL3");
+  std::vector<const Model*> models{
+      main_model_, flatbuffers::GetRoot<Model>(fbb_.GetBufferPointer())};
+  std::vector<std::string> subgraph_names_not_important(2);
+  return CombineModels(fbb, models, subgraph_names_not_important, schema_);
+}
+
+absl::Status ValidationGraphBuilder::BuildFinalModel(
+    FlatBufferBuilder* fbb, Subgraph* subgraph_with_golden_outputs) {
+  fbb_.Reset();
+  auto model =
+      MakeModel(/* intermediate_only */ false, subgraph_with_golden_outputs);
+  if (!model.ok()) {
+    return model.status();
+  }
+  fbb_.Finish(*model, "TFL3");
+  std::vector<const Model*> models{
+      main_model_, GetRoot<Model>(fbb_.GetBufferPointer()), validation_model_};
+  std::vector<std::string> subgraph_names;
+  auto main_subgraph_name = main_model_->subgraphs()->Get(0)->name();
+  subgraph_names.push_back(main_subgraph_name ? main_subgraph_name->str() : "");
+  subgraph_names.push_back("VALIDATION:main");
+  subgraph_names.push_back("VALIDATION:metrics");
+  return CombineModels(fbb, models, subgraph_names, schema_);
+}
+
+absl::StatusOr<flatbuffers::Offset<Model>> ValidationGraphBuilder::MakeModel(
+    bool intermediate_only, Subgraph* subgraph_with_golden_outputs) {
+  TensorInfo tensor_info;
+  auto operator_codes = OperatorCodes();
+  if (!operator_codes.ok()) {
+    return operator_codes.status();
+  }
+  auto subgraphs = SubGraphs(intermediate_only, &tensor_info);
+  if (!subgraphs.ok()) {
+    return subgraphs.status();
+  }
+  auto buffers =
+      Buffers(intermediate_only, tensor_info, subgraph_with_golden_outputs);
+  if (!buffers.ok()) {
+    return buffers.status();
+  }
+  return CreateModel(fbb_, kModelVersion, *operator_codes, *subgraphs,
+                     fbb_.CreateString("validation"), *buffers,
+                     /* metadata_buffer */ 0, /* metadata */ 0,
+                     /* signature_defs */ 0);
+}
+
+absl::StatusOr<
+    flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<OperatorCode>>>>
+ValidationGraphBuilder::OperatorCodes() {
+#define RET_CHECK_INDEX(constant, code_index)                              \
+  do {                                                                     \
+    if ((constant) != (code_index)) {                                      \
+      return absl::InternalError(absl::StrFormat(                          \
+          "Operator code indexing mismatch %s (%d) != %s (%d)", #constant, \
+          (constant), #code_index, (code_index)));                         \
+    }                                                                      \
+  } while (0)
+  std::vector<flatbuffers::Offset<OperatorCode>> codes;
+  RET_CHECK_INDEX(kCallOperatorCode, codes.size());
+  codes.push_back(CreateOperatorCode(fbb_, BuiltinOperator_CUSTOM,
+                                     fbb_.CreateString("validation/call")));
+  RET_CHECK_INDEX(kDequantizeOperatorCode, codes.size());
+  codes.push_back(CreateOperatorCode(fbb_, BuiltinOperator_DEQUANTIZE));
+  RET_CHECK_INDEX(kDecodeJpegOperatorCode, codes.size());
+  codes.push_back(
+      CreateOperatorCode(fbb_, BuiltinOperator_CUSTOM,
+                         fbb_.CreateString("validation/decode_jpeg")));
+  return fbb_.CreateVector(codes);
+#undef RET_CHECK_INDEX
+}
+
+absl::StatusOr<
+    flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<Tensor>>>>
+ValidationGraphBuilder::Tensors(bool intermediate_only,
+                                TensorInfo* tensor_info) {
+  std::vector<flatbuffers::Offset<Tensor>> tensors;
+  int buffer_count = 0;
+
+  // Copy tensors from a source subgraph, overriding the batch_size where
+  // necessary (the called subgraph always uses batch size 1, the calling
+  // subgraph always uses batch size equal jpeg_data_.size()).
+  auto copy =
+      [&tensors, this, &buffer_count](
+          const SubGraph* from_subgraph,
+          const flatbuffers::Vector<int32_t>* indices,
+          std::vector<int32_t>* store_indices_into, int batch_size,
+          const std::string prefix = "",
+          std::function<absl::StatusOr<bool>(const Tensor*, int)> filter =
+              nullptr) -> absl::Status {
+    int counter = 0;
+    for (auto index = indices->cbegin(); index != indices->cend();
+         index++, counter++) {
+      const Tensor* tensor = from_subgraph->tensors()->Get(*index);
+      if (filter) {
+        auto statusor = filter(tensor, counter);
+        if (!statusor.ok()) {
+          return statusor.status();
+        } else if (!statusor.value()) {
+          store_indices_into->push_back(kSkippedIndex);
+          continue;
+        }
+      }
+      std::vector<int32_t> shape{tensor->shape()->cbegin(),
+                                 tensor->shape()->cend()};
+      if (shape.size() >= 2 && shape[0] == 1 && batch_size > 0) {
+        shape[0] = batch_size;
+      }
+      std::vector<int32_t> shape_signature;
+      if (tensor->shape_signature()) {
+        shape_signature.assign(tensor->shape_signature()->cbegin(),
+                               tensor->shape_signature()->cend());
+        if (shape_signature.size() >= 2 && shape_signature[0] == 1 &&
+            batch_size > 0) {
+          shape_signature[0] = batch_size;
+        }
+      }
+      auto quantization_parameters = helper_.CopyTable(
+          "tflite.QuantizationParameters", tensor->quantization());
+      if (!quantization_parameters.ok()) {
+        return quantization_parameters.status();
+      }
+      auto sparsity_parameters =
+          helper_.CopyTable("tflite.SparsityParameters", tensor->sparsity());
+      if (!sparsity_parameters.ok()) {
+        return sparsity_parameters.status();
+      }
+      store_indices_into->push_back(tensors.size());
+      std::string name = tensor->name()->str();
+      if (!prefix.empty() && name.find(prefix) != 0) {  // NOLINT
+        name = prefix + name;
+      }
+      tensors.push_back(CreateTensor(
+          fbb_, fbb_.CreateVector(shape), tensor->type(), buffer_count,
+          fbb_.CreateString(name), *quantization_parameters,
+          tensor->is_variable(), *sparsity_parameters,
+          shape_signature.empty() ? 0 : fbb_.CreateVector(shape_signature)));
+      buffer_count++;
+    }
+    return absl::OkStatus();
+  };
+  // Input image, jpeg data.
+  tensor_info->jpeg_images.push_back(tensors.size());
+  DynamicBuffer jpeg_buffer;
+  for (int i = 0; i < jpeg_data_.size(); i++) {
+    jpeg_buffer.AddString(jpeg_data_[i].data(), jpeg_data_[i].size());
+  }
+  tensor_info->jpeg_buffer_length =
+      jpeg_buffer.WriteToBuffer(&(tensor_info->jpeg_buffer_contents));
+  tensors.push_back(CreateTensor(fbb_,
+                                 fbb_.CreateVector(std::vector<int32_t>{
+                                     static_cast<int32_t>(jpeg_data_.size())}),
+                                 TensorType::TensorType_STRING, buffer_count,
+                                 fbb_.CreateString("call/jpeg_images")));
+  buffer_count++;
+
+  // Input image.
+  const SubGraph* main_subgraph = main_model_->subgraphs()->Get(0);
+  const Tensor* input_tensor =
+      main_subgraph->tensors()->Get(main_subgraph->inputs()->Get(0));
+  tensor_info->jpeg_height = input_tensor->shape()->Get(1);
+  tensor_info->jpeg_width = input_tensor->shape()->Get(2);
+  if (input_tensor->type() == TensorType_FLOAT32) {
+    // Quantized.
+    std::vector<int32_t> input_shape{input_tensor->shape()->cbegin(),
+                                     input_tensor->shape()->cend()};
+    input_shape[0] = static_cast<int32_t>(jpeg_data_.size());
+    tensor_info->quantized_images.push_back(tensors.size());
+    tensors.push_back(CreateTensor(
+        fbb_, fbb_.CreateVector(input_shape), TensorType::TensorType_UINT8,
+        buffer_count, fbb_.CreateString("call/quant_image"),
+        CreateQuantizationParameters(
+            fbb_, 0, 0, fbb_.CreateVector(std::vector<float>{scale_}),
+            fbb_.CreateVector(std::vector<int64_t>{zero_point_}))));
+    buffer_count++;
+    // Float.
+    tensor_info->float_images.push_back(tensors.size());
+    tensors.push_back(CreateTensor(fbb_, fbb_.CreateVector(input_shape),
+                                   TensorType::TensorType_FLOAT32, buffer_count,
+                                   fbb_.CreateString("call/float_image")));
+    buffer_count++;
+  } else {
+    // Quantized only.
+    auto status = copy(main_model_->subgraphs()->Get(0),
+                       main_model_->subgraphs()->Get(0)->inputs(),
+                       &tensor_info->quantized_images, jpeg_data_.size());
+    if (!status.ok()) {
+      return status;
+    }
+  }
+
+  // Validation inputs, actual.
+  auto status = copy(main_model_->subgraphs()->Get(0),
+                     main_model_->subgraphs()->Get(0)->outputs(),
+                     &tensor_info->main_outputs, jpeg_data_.size());
+  if (!status.ok()) {
+    return status;
+  }
+  if (intermediate_only) {
+    return fbb_.CreateVector(tensors);
+  }
+  // Validation inputs, golden.
+  tensor_info->validation_inputs = tensor_info->main_outputs;
+  status = copy(main_model_->subgraphs()->Get(0),
+                main_model_->subgraphs()->Get(0)->outputs(),
+                &tensor_info->validation_inputs, jpeg_data_.size());
+  if (!status.ok()) {
+    return status;
+  }
+  // Entrypoint inputs. Golden first (validator relies on this).
+  for (int i = tensor_info->validation_inputs.size() / 2;
+       i < tensor_info->validation_inputs.size(); i++) {
+    tensor_info->entrypoint_inputs.push_back(tensor_info->validation_inputs[i]);
+  }
+  tensor_info->entrypoint_inputs.push_back(tensor_info->jpeg_images[0]);
+  // Validation inputs, dequantized.
+  status = copy(
+      validation_model_->subgraphs()->Get(0),
+      validation_model_->subgraphs()->Get(0)->inputs(),
+      &tensor_info->dequantized_validation_inputs, jpeg_data_.size(), "",
+      [&tensors, &tensor_info, this](const Tensor* validation_model_input,
+                                     int i) -> absl::StatusOr<bool> {
+        // validation_model_input is the tensor for metrics calculation.
+        // validation_graph_input is the under-construction graph will be
+        // given to the metrics calculation but need to be dequantized
+        // first.
+        const Tensor* validation_graph_input = flatbuffers::GetTemporaryPointer(
+            fbb_, tensors[tensor_info->validation_inputs[i]]);
+        if (validation_model_input->type() == TensorType_FLOAT32 &&
+            (validation_graph_input->type() == TensorType_UINT8 ||
+             validation_graph_input->type() == TensorType_INT8)) {
+          return true;
+        } else if (validation_model_input->type() !=
+                   validation_graph_input->type()) {
+          const char* name = "(null)";
+          if (validation_model_input->name()) {
+            name = validation_model_input->name()->c_str();
+          }
+          return absl::InvalidArgumentError(
+              absl::StrFormat("Validation model input %s with type %d is "
+                              "incompatible with main model output type %d",
+                              name, validation_model_input->type(),
+                              validation_graph_input->type()));
+        } else {
+          return false;
+        }
+      });
+  if (!status.ok()) {
+    return status;
+  }
+  // Validation outputs.
+  status =
+      copy(validation_model_->subgraphs()->Get(0),
+           validation_model_->subgraphs()->Get(0)->outputs(),
+           &tensor_info->validation_outputs, jpeg_data_.size(), metric_prefix_);
+  if (!status.ok()) {
+    return status;
+  }
+  // Outputs from entrypoint graph
+  // Actuals first (validator relies on this);
+  for (int i = 0; i < tensor_info->validation_inputs.size() / 2; i++) {
+    tensor_info->entrypoint_outputs.push_back(
+        tensor_info->validation_inputs[i]);
+  }
+  // Metrics.
+  for (int i = 0; i < tensor_info->validation_outputs.size(); i++) {
+    tensor_info->entrypoint_outputs.push_back(
+        tensor_info->validation_outputs[i]);
+  }
+  return fbb_.CreateVector(tensors);
+}
+
+flatbuffers::Offset<flatbuffers::Vector<uint8_t>>
+ValidationGraphBuilder::CallOpCustomOptions(int subgraph) {
+  flexbuffers::Builder fbb;
+  fbb.Map([&] {
+    fbb.Int("subgraph_index", subgraph);
+    fbb.Int("loop_count", static_cast<int32_t>(jpeg_data_.size()));
+  });
+  fbb.Finish();
+  return fbb_.CreateVector(fbb.GetBuffer());
+}
+
+flatbuffers::Offset<flatbuffers::Vector<uint8_t>>
+ValidationGraphBuilder::JpegOpCustomOptions(int height, int width,
+                                            int channels) {
+  flexbuffers::Builder fbb;
+  fbb.Map([&] {
+    fbb.Int("height", height);
+    fbb.Int("width", width);
+    fbb.Int("channels", channels);
+    fbb.Int("num_images", jpeg_data_.size());
+  });
+  fbb.Finish();
+  return fbb_.CreateVector(fbb.GetBuffer());
+}
+
+flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<Operator>>>
+ValidationGraphBuilder::Operators(bool intermediate_only,
+                                  const TensorInfo& tensor_info) {
+  std::vector<flatbuffers::Offset<Operator>> ops;
+  // Jpeg decode.
+  ops.push_back(CreateOperator(
+      fbb_, kDecodeJpegOperatorCode, fbb_.CreateVector(tensor_info.jpeg_images),
+      fbb_.CreateVector(tensor_info.quantized_images),
+      tflite::BuiltinOptions_NONE, 0,
+      JpegOpCustomOptions(tensor_info.jpeg_height, tensor_info.jpeg_width,
+                          jpeg_output_channels_)));
+  if (!tensor_info.float_images.empty()) {
+    // Dequantize.
+    ops.push_back(
+        CreateOperator(fbb_, kDequantizeOperatorCode,
+                       fbb_.CreateVector(tensor_info.quantized_images),
+                       fbb_.CreateVector(tensor_info.float_images),
+                       BuiltinOptions_DequantizeOptions, 0));
+    // Call main model.
+    ops.push_back(CreateOperator(
+        fbb_, kCallOperatorCode, fbb_.CreateVector(tensor_info.float_images),
+        fbb_.CreateVector(tensor_info.main_outputs),
+        tflite::BuiltinOptions_NONE, 0, CallOpCustomOptions(kMainSubgraphIndex),
+        tflite::CustomOptionsFormat_FLEXBUFFERS));
+  } else {
+    // Call main model.
+    ops.push_back(CreateOperator(
+        fbb_, kCallOperatorCode,
+        fbb_.CreateVector(tensor_info.quantized_images),
+        fbb_.CreateVector(tensor_info.main_outputs),
+        tflite::BuiltinOptions_NONE, 0, CallOpCustomOptions(kMainSubgraphIndex),
+        tflite::CustomOptionsFormat_FLEXBUFFERS));
+  }
+  if (intermediate_only) {
+    return fbb_.CreateVector(ops);
+  }
+  // Call validation model.
+  std::vector<int32_t> validation_input_indices;
+  for (int i = 0; i < tensor_info.dequantized_validation_inputs.size(); i++) {
+    int32_t validation_input_index;
+    if (tensor_info.dequantized_validation_inputs[i] == kSkippedIndex) {
+      validation_input_index = tensor_info.validation_inputs[i];
+    } else {
+      validation_input_index = tensor_info.dequantized_validation_inputs[i];
+      std::vector<int32_t> dequantize_inputs{tensor_info.validation_inputs[i]};
+      std::vector<int32_t> dequantize_outputs{
+          tensor_info.dequantized_validation_inputs[i]};
+      ops.push_back(CreateOperator(fbb_, kDequantizeOperatorCode,
+                                   fbb_.CreateVector(dequantize_inputs),
+                                   fbb_.CreateVector(dequantize_outputs),
+                                   BuiltinOptions_DequantizeOptions, 0));
+    }
+    validation_input_indices.push_back(validation_input_index);
+  }
+  ops.push_back(CreateOperator(
+      fbb_, kCallOperatorCode, fbb_.CreateVector(validation_input_indices),
+      fbb_.CreateVector(tensor_info.validation_outputs),
+      tflite::BuiltinOptions_NONE, 0,
+      CallOpCustomOptions(kValidationSubgraphIndex),
+      tflite::CustomOptionsFormat_FLEXBUFFERS));
+  return fbb_.CreateVector(ops);
+}
+
+absl::StatusOr<
+    flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<SubGraph>>>>
+ValidationGraphBuilder::SubGraphs(bool intermediate_only,
+                                  TensorInfo* tensor_info) {
+  auto tensors = Tensors(intermediate_only, tensor_info);
+  if (!tensors.ok()) {
+    return tensors.status();
+  }
+  std::vector<flatbuffers::Offset<SubGraph>> graphs;
+  if (intermediate_only) {
+    graphs.push_back(CreateSubGraph(
+        fbb_, *tensors, fbb_.CreateVector(tensor_info->jpeg_images),
+        fbb_.CreateVector(tensor_info->main_outputs),
+        Operators(intermediate_only, *tensor_info), fbb_.CreateString("call")));
+  } else {
+    graphs.push_back(CreateSubGraph(
+        fbb_, *tensors, fbb_.CreateVector(tensor_info->entrypoint_inputs),
+        fbb_.CreateVector(tensor_info->entrypoint_outputs),
+        Operators(intermediate_only, *tensor_info), fbb_.CreateString("call")));
+  }
+  return fbb_.CreateVector(graphs);
+}
+
+absl::StatusOr<
+    flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<Buffer>>>>
+ValidationGraphBuilder::Buffers(bool intermediate_only,
+                                const TensorInfo& tensor_info,
+                                Subgraph* subgraph_with_golden_outputs) {
+  std::vector<flatbuffers::Offset<Buffer>> buffers;
+
+  // The buffers created in this method map 1:1 to the
+  // tensors created in Tensors() - a tensor at index X
+  // uses buffer at index X. The numbering is checked along
+  // the way using the RET_CHECK_INDEX macro below.
+#define RET_CHECK_INDEX(tensor_index, buffer_index)                         \
+  do {                                                                      \
+    if ((tensor_index) != (buffer_index)) {                                 \
+      return absl::InternalError(absl::StrFormat(                           \
+          "%s:%d, Tensor/buffer indexing mismatch %s (%d) != %s (%d)",      \
+          __FILE__, __LINE__, #tensor_index, (tensor_index), #buffer_index, \
+          (buffer_index)));                                                 \
+    }                                                                       \
+  } while (0)
+
+  // Jpeg input.
+  RET_CHECK_INDEX(tensor_info.jpeg_images.size(), 1);
+  RET_CHECK_INDEX(tensor_info.jpeg_images[0], buffers.size());
+  std::vector<uint8_t> jpeg_buffer_vec{
+      reinterpret_cast<const uint8_t*>(tensor_info.jpeg_buffer_contents),
+      reinterpret_cast<const uint8_t*>(tensor_info.jpeg_buffer_contents) +
+          tensor_info.jpeg_buffer_length};
+  buffers.push_back(CreateBuffer(fbb_, fbb_.CreateVector(jpeg_buffer_vec)));
+
+  // Decoded and dequantized image.
+  RET_CHECK_INDEX(tensor_info.quantized_images.size(), 1);
+  RET_CHECK_INDEX(tensor_info.quantized_images[0], buffers.size());
+  buffers.push_back(CreateBuffer(fbb_));
+  if (!tensor_info.float_images.empty()) {
+    RET_CHECK_INDEX(tensor_info.float_images.size(), 1);
+    RET_CHECK_INDEX(tensor_info.float_images[0], buffers.size());
+    buffers.push_back(CreateBuffer(fbb_));
+  }
+
+  // Main graph outputs / first half of validation inputs.
+  auto main_subgraph = main_model_->subgraphs()->Get(0);
+  RET_CHECK_INDEX(main_subgraph->outputs()->size(),
+                  tensor_info.main_outputs.size());
+  int main_output_index = 0;
+  int validation_graph_input_index = 0;
+  for (auto i = main_subgraph->outputs()->cbegin();
+       i != main_subgraph->outputs()->cend(); i++) {
+    RET_CHECK_INDEX(tensor_info.main_outputs[main_output_index],
+                    buffers.size());
+    main_output_index++;
+    if (!intermediate_only) {
+      RET_CHECK_INDEX(
+          tensor_info.validation_inputs[validation_graph_input_index],
+          buffers.size());
+      validation_graph_input_index++;
+    }
+    auto t = main_subgraph->tensors()->Get(*i);
+    auto status = helper_.CopyTableToVector(
+        "tflite.Buffer", main_model_->buffers()->Get(t->buffer()), &buffers);
+    if (!status.ok()) {
+      return status;
+    }
+  }
+  if (intermediate_only) {
+    return fbb_.CreateVector(buffers);
+  }
+
+  // Golden outputs / second half of validation inputs.
+  RET_CHECK_INDEX(tensor_info.validation_inputs.size(),
+                  validation_graph_input_index +
+                      subgraph_with_golden_outputs->outputs().size());
+  for (auto i : subgraph_with_golden_outputs->outputs()) {
+    RET_CHECK_INDEX(tensor_info.validation_inputs[validation_graph_input_index],
+                    buffers.size());
+    validation_graph_input_index++;
+    auto t = subgraph_with_golden_outputs->tensor(i);
+    if (!use_ondevice_cpu_for_golden_) {
+      std::vector<uint8_t> output_data{
+          reinterpret_cast<const uint8_t*>(t->data.raw),
+          reinterpret_cast<const uint8_t*>(t->data.raw + t->bytes)};
+      buffers.push_back(CreateBuffer(fbb_, fbb_.CreateVector(output_data)));
+    } else {
+      buffers.push_back(CreateBuffer(fbb_));
+    }
+  }
+
+  auto validation_model_subgraph = validation_model_->subgraphs()->Get(0);
+  // Dequantized validation inputs.
+  RET_CHECK_INDEX(tensor_info.dequantized_validation_inputs.size(),
+                  validation_model_subgraph->inputs()->size());
+  int validation_graph_dequantized_input_index = 0;
+  for (auto i = validation_model_subgraph->inputs()->cbegin();
+       i != validation_model_subgraph->inputs()->cend(); i++) {
+    if (tensor_info.dequantized_validation_inputs
+            [validation_graph_dequantized_input_index] == kSkippedIndex) {
+      validation_graph_dequantized_input_index++;
+      continue;
+    }
+    RET_CHECK_INDEX(tensor_info.dequantized_validation_inputs
+                        [validation_graph_dequantized_input_index],
+                    buffers.size());
+    validation_graph_dequantized_input_index++;
+    auto t = validation_model_subgraph->tensors()->Get(*i);
+    auto status = helper_.CopyTableToVector(
+        "tflite.Buffer", validation_model_->buffers()->Get(t->buffer()),
+        &buffers);
+    if (!status.ok()) {
+      return status;
+    }
+  }
+
+  // Validation outputs.
+  RET_CHECK_INDEX(tensor_info.validation_outputs.size(),
+                  validation_model_subgraph->outputs()->size());
+  int validation_graph_output_index = 0;
+  for (auto i = validation_model_subgraph->outputs()->cbegin();
+       i != validation_model_subgraph->outputs()->cend(); i++) {
+    RET_CHECK_INDEX(
+        tensor_info.validation_outputs[validation_graph_output_index],
+        buffers.size());
+    validation_graph_output_index++;
+    auto t = validation_model_subgraph->tensors()->Get(*i);
+    auto status = helper_.CopyTableToVector(
+        "tflite.Buffer", validation_model_->buffers()->Get(t->buffer()),
+        &buffers);
+    if (!status.ok()) {
+      return status;
+    }
+  }
+  return fbb_.CreateVector(buffers);
+#undef RET_CHECK_INDEX
+}
+
+}  // namespace acceleration
+
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/validation_graph_builder.h

@@ -0,0 +1,221 @@
+/* Copyright 2022 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MODEL_MODIFIER_VALIDATION_GRAPH_BUILDER_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MODEL_MODIFIER_VALIDATION_GRAPH_BUILDER_H_
+
+#include <stdint.h>
+#include <stdlib.h>
+
+#include <cstdlib>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "flatbuffers/buffer.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "flatbuffers/reflection_generated.h"  // from @flatbuffers
+#include "flatbuffers/vector.h"  // from @flatbuffers
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier/grafter.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace acceleration {
+
+// Class for building the validation entry-point graph that calls into the main
+// graph and a metrics graph. Like this (boxes are tensors with plural names
+// meaning possibly multiple tensors, arrows are ops and numbers in parentheses
+// are subgraph indices):
+// +--------------------------------------+
+// | Graph created by this class (1)      |
+// |                                      |
+// | +-----------input-+                  |
+// | |jpeg input       |                  |
+// | +-----+-----------+                  |
+// |       |                              |
+// |       | decode                       |
+// |       v                              |
+// | +-----+-----------+                  |
+// | |quantized image  |                  |
+// | +-----+-----------+                  |  +-----------------------+
+// |       |                              |  |'main_model' (0)       |
+// |       | dequantize (optional)        |  | +---------------+     |
+// |       v                              |  | |input          +---+ |
+// | +-----+-----------+                  |  | +---------------+   | |
+// | |float image      |                  |  |                     ~ |
+// | +-----+-----------+                  |  | +---------------+   | |
+// |       |  call                        |  | |outputs        +<--+ |
+// |       +<------------------------------->+ +---------------+     |
+// |       v                              |  |                       |
+// | +-----+-----output+ +---------input+ |  +-----------------------+
+// | |actual outputs   | |golden outputs| |
+// | +-----+-----------+ +-----------+--+ |
+// |       |                         |    |
+// |       | dequantize (optional)   |    |
+// |       |                         |    |
+// | +-----+-------------------------+-+  |
+// | | dequantized actual and golden   |  |
+// | | outputs (validation inputs)     |  |
+// | +-----+---------------------------+  |  +-----------------------+
+// |       |  call                        |  |'validation model' (2) |
+// |       +<------------------------------->+                       |
+// |       v                              |  | +---------------+     |
+// | +-----+-----output+                  |  | |inputs         +---+ |
+// | |results          |                  |  | +---------------+   | |
+// | +-----------------+                  |  |                     ~ |
+// |                                      |  | +---------------+   | |
+// |                                      |  | |outputs        +<--+ |
+// |                                      |  | +---------------+     |
+// |                                      |  |                       |
+// +--------------------------------------+  +-----------------------+
+//
+// It's important the 'main_model' has subgraph index 0 so that it is used as
+// the primary subgraph by the TFLite interpreter. The other indices are
+// arbitrary.
+// TODO(b/172541832): Handle a main model with more than one subgraph.
+//
+// Note that the jpeg input is marked as an input in this graph, as TFLite
+// graphs must have inputs. However, it will be pre-filled from the jpeg_data
+// and doesn't need to be filled by the user of the model.
+class ValidationGraphBuilder {
+ public:
+  ValidationGraphBuilder(const std::string& metric_prefix,
+                         const Model* main_model,
+                         std::vector<std::string> jpeg_data,
+                         int32_t jpeg_output_channels, float scale,
+                         int64_t zero_point, const Model* validation_model,
+                         const reflection::Schema* schema,
+                         bool use_ondevice_cpu_for_golden)
+      : metric_prefix_(metric_prefix),
+        main_model_(main_model),
+        jpeg_data_(jpeg_data),
+        jpeg_output_channels_(jpeg_output_channels),
+        scale_(scale),
+        zero_point_(zero_point),
+        validation_model_(validation_model),
+        schema_(schema),
+        helper_(&fbb_, schema_),
+        use_ondevice_cpu_for_golden_(use_ondevice_cpu_for_golden) {}
+
+  ValidationGraphBuilder(const ValidationGraphBuilder&) = delete;
+  ValidationGraphBuilder& operator=(const ValidationGraphBuilder&) = delete;
+
+  // Builds the part of the model drawn above until the call to the validation
+  // graph. The model is used to generate golden outputs. Calls Finish on the
+  // FlatbufferBuilder.
+  absl::Status BuildIntermediateModel(flatbuffers::FlatBufferBuilder* fbb);
+
+  // Builds the whole model as drawn above. The subgraph_with_golden_outputs
+  // should be the result of invoking subgraph 1 on the output of
+  // BuildIntermediateModel(). Calls Finish on the FlatbufferBuilder.
+  absl::Status BuildFinalModel(flatbuffers::FlatBufferBuilder* fbb,
+                               Subgraph* subgraph_with_golden_outputs);
+
+ private:
+  // Allocation of tensors, for communication between methods that create the
+  // tensors, the operations and the buffers.
+  // (Some of these vectors will always contain only one element, but using the
+  // same type for them simplifies the code a lot).
+  struct TensorInfo {
+    ~TensorInfo() { std::free(jpeg_buffer_contents); }
+
+    std::vector<int32_t> entrypoint_inputs;
+    std::vector<int32_t> entrypoint_outputs;
+    std::vector<int32_t> jpeg_images;
+
+    // With float main model, both quantized_images and float_images are set,
+    // and float_images is the same as main input. With a quantized model
+    // only quantized_images is set and it's the same as main input.
+    std::vector<int32_t> quantized_images;
+    std::vector<int32_t> float_images;
+
+    std::vector<int32_t> main_outputs;  // First half of validation_inputs.
+    std::vector<int32_t> validation_inputs;
+    // With a float model, validation_inputs is used directly. With a quantized
+    // model, the inputs are first dequantized.
+    // Some models have a mixture of quantized outputs that need to be
+    // dequantized to floats; and integer outputs. For integer outputs
+    // kSkippedIndex is used.
+    std::vector<int32_t> dequantized_validation_inputs;
+    std::vector<int32_t> validation_outputs;
+
+    char* jpeg_buffer_contents = nullptr;
+    int32_t jpeg_buffer_length = -1;
+    int32_t jpeg_height = -1;
+    int32_t jpeg_width = -1;
+  };
+
+  static constexpr int32_t kModelVersion = 3;
+  static constexpr int32_t kSkippedIndex = -1;
+  // Operator code numbering.
+  static constexpr int32_t kCallOperatorCode = 0;
+  static constexpr int32_t kDequantizeOperatorCode = 1;
+  static constexpr int32_t kDecodeJpegOperatorCode = 2;
+  // Subgraph numbering.
+  static constexpr int32_t kMainSubgraphIndex = 0;
+  static constexpr int32_t kValidationSubgraphIndex = 2;
+
+  absl::StatusOr<flatbuffers::Offset<Model>> MakeModel(
+      bool intermediate_only, Subgraph* subgraph_with_golden_outputs);
+
+  absl::StatusOr<flatbuffers::Offset<
+      flatbuffers::Vector<flatbuffers::Offset<OperatorCode>>>>
+  OperatorCodes();
+
+  absl::StatusOr<
+      flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<Tensor>>>>
+  Tensors(bool intermediate_only, TensorInfo* tensor_info);
+
+  // Create the options for the custom call op (see call.cc for the options
+  // format).
+  flatbuffers::Offset<flatbuffers::Vector<uint8_t>> CallOpCustomOptions(
+      int subgraph);
+
+  // Create the options for the custom jpeg op (see decode_jpeg.cc for the
+  // options format).
+  flatbuffers::Offset<flatbuffers::Vector<uint8_t>> JpegOpCustomOptions(
+      int height, int width, int channels);
+
+  flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<Operator>>>
+  Operators(bool intermediate_only, const TensorInfo& tensor_info);
+
+  absl::StatusOr<
+      flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<SubGraph>>>>
+  SubGraphs(bool intermediate_only, TensorInfo* tensor_info);
+
+  absl::StatusOr<
+      flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<Buffer>>>>
+  Buffers(bool intermediate_only, const TensorInfo& tensor_info,
+          Subgraph* subgraph_with_golden_outputs);
+
+  const std::string metric_prefix_;
+  const Model* main_model_;
+  std::vector<std::string> jpeg_data_;
+  int32_t jpeg_output_channels_;
+  float scale_;
+  int64_t zero_point_;
+  const Model* validation_model_;
+  const reflection::Schema* schema_;
+  flatbuffers::FlatBufferBuilder fbb_;
+  FlatbufferHelper helper_;
+  bool use_ondevice_cpu_for_golden_;
+};
+
+}  // namespace acceleration
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MODEL_MODIFIER_VALIDATION_GRAPH_BUILDER_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/model_validation_test.cc

@@ -0,0 +1,286 @@
+/* Copyright 2021 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 <errno.h>
+#include <fcntl.h>
+#include <sched.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+
+#include <fstream>
+#include <memory>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/acceleration/configuration/configuration_generated.h"
+#include "tensorflow/lite/experimental/acceleration/compatibility/android_info.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/big_little_affinity.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/validator.h"
+#include "tensorflow/lite/tools/model_loader.h"
+#ifdef ENABLE_NNAPI_SL_TEST
+#include "tensorflow/lite/nnapi/sl/include/SupportLibrary.h"
+#endif /* ENABLE_NNAPI_SL_TEST */
+
+extern const unsigned char TENSORFLOW_ACCELERATION_MODEL_DATA_VARIABLE[];
+extern const int TENSORFLOW_ACCELERATION_MODEL_LENGTH_VARIABLE;
+
+namespace tflite {
+namespace acceleration {
+namespace {
+
+class LocalizerValidationRegressionTest : public ::testing::Test {
+ protected:
+  static void SetUpTestSuite() {
+    int error = -1;
+#if defined(__ANDROID__)
+    AndroidInfo android_info;
+    auto status = RequestAndroidInfo(&android_info);
+    ASSERT_TRUE(status.ok());
+    if (android_info.is_emulator) {
+      std::cerr << "Running on an emulator, skipping processor affinity\n";
+    } else {
+      std::cerr << "Running on hardware, setting processor affinity\n";
+      BigLittleAffinity affinity = GetAffinity();
+      cpu_set_t set;
+      CPU_ZERO(&set);
+      for (int i = 0; i < 16; i++) {
+        if (affinity.big_core_affinity & (0x1 << i)) {
+          CPU_SET(i, &set);
+        }
+      }
+      error = sched_setaffinity(getpid(), sizeof(set), &set);
+      if (error == -1) {
+        perror("sched_setaffinity failed");
+      }
+    }
+#endif
+    std::string dir = GetTestTmpDir();
+    error = mkdir(dir.c_str(), 0777);
+    if (error == -1) {
+      if (errno != EEXIST) {
+        perror("mkdir failed");
+        ASSERT_TRUE(false);
+      }
+    }
+
+    std::string path = ModelPath();
+    (void)unlink(path.c_str());
+    std::string contents(reinterpret_cast<const char*>(
+                             TENSORFLOW_ACCELERATION_MODEL_DATA_VARIABLE),
+                         TENSORFLOW_ACCELERATION_MODEL_LENGTH_VARIABLE);
+    std::ofstream f(path, std::ios::binary);
+    ASSERT_TRUE(f.is_open());
+    f << contents;
+    f.close();
+    ASSERT_EQ(chmod(path.c_str(), 0500), 0);
+  }
+  static std::string ModelPath() { return GetTestTmpDir() + "/model.tflite"; }
+  static std::string GetTestTmpDir() {
+    const char* from_env = getenv("TEST_TMPDIR");
+    if (from_env) {
+      return from_env;
+    }
+#ifdef __ANDROID__
+    return "/data/local/tmp";
+#else
+    return "/tmp";
+#endif
+  }
+  void CheckValidation(const std::string& accelerator_name) {
+    std::string path = ModelPath();
+    const ComputeSettings* settings =
+        flatbuffers::GetRoot<ComputeSettings>(fbb_.GetBufferPointer());
+    int fd = open(path.c_str(), O_RDONLY);
+    ASSERT_GE(fd, 0);
+    struct stat stat_buf = {0};
+    ASSERT_EQ(fstat(fd, &stat_buf), 0);
+    auto validator =
+        std::make_unique<Validator>(std::make_unique<tools::MmapModelLoader>(
+                                        fd, /*offset=*/0, stat_buf.st_size),
+                                    settings);
+    close(fd);
+
+    Validator::Results results;
+    Validator::Status validation_run = validator->RunValidation(&results);
+    EXPECT_EQ(validation_run.status, kMinibenchmarkSuccess);
+    EXPECT_EQ(validation_run.stage, BenchmarkStage_UNKNOWN);
+    EXPECT_TRUE(results.ok);
+    EXPECT_EQ(results.delegate_error, 0);
+    if (accelerator_name != "CPU") {
+      // For any non-CPU delegate, we validate that model execution was at least
+      // partially delegated by expecting non-zero number of delegated kernels
+      EXPECT_NE(results.delegated_kernels, 0);
+    }
+
+    for (const auto& metric : results.metrics) {
+      int test_case = 0;
+      std::cerr << "Metric " << metric.first;
+      for (float v : metric.second) {
+        std::cerr << " " << v;
+        RecordProperty("[" + std::to_string(test_case++) + "] " +
+                           accelerator_name + " " + metric.first,
+                       std::to_string(v));
+      }
+      std::cerr << "\n";
+    }
+    std::cerr << "Delegate prep time us " << results.delegate_prep_time_us
+              << std::endl;
+    RecordProperty(accelerator_name + " Delegate prep time us",
+                   results.delegate_prep_time_us);
+    std::cerr << "Execution time us";
+    int test_case = 0;
+    int64_t total_execution_time_us = 0;
+    for (int64_t t : results.execution_time_us) {
+      std::cerr << " " << t;
+      RecordProperty("[" + std::to_string(test_case++) + "] " +
+                         accelerator_name + " Execution time us",
+                     t);
+      total_execution_time_us += t;
+    }
+    std::cerr << "\n";
+    int64_t average_execution_time_us = total_execution_time_us / test_case;
+    std::cerr << "Avg execution time us " << average_execution_time_us
+              << std::endl;
+    RecordProperty(accelerator_name + " Avg execution time us",
+                   average_execution_time_us);
+    std::cerr << std::endl;
+  }
+  flatbuffers::FlatBufferBuilder fbb_;
+};
+
+TEST_F(LocalizerValidationRegressionTest, Cpu) {
+  fbb_.Finish(CreateComputeSettings(fbb_, ExecutionPreference_ANY,
+                                    CreateTFLiteSettings(fbb_)));
+  CheckValidation("CPU");
+}
+
+#ifndef DISABLE_PLATFORM_NNAPI_TEST
+TEST_F(LocalizerValidationRegressionTest, Nnapi) {
+  fbb_.Finish(
+      CreateComputeSettings(fbb_, ExecutionPreference_ANY,
+                            CreateTFLiteSettings(fbb_, Delegate_NNAPI)));
+  AndroidInfo android_info;
+  auto status = RequestAndroidInfo(&android_info);
+  ASSERT_TRUE(status.ok());
+  if (android_info.android_sdk_version >= "28") {
+    CheckValidation("NNAPI");
+  }
+}
+#endif  // DISABLE_PLATFORM_NNAPI_TEST
+
+#ifdef ENABLE_NNAPI_SL_TEST
+TEST_F(LocalizerValidationRegressionTest, NnapiSl) {
+  const char* accelerator_name = getenv("TEST_ACCELERATOR_NAME");
+
+  auto nnapi_sl_handle = nnapi::loadNnApiSupportLibrary(
+      GetTestTmpDir() + "/libnnapi_sl_driver.so");
+
+  ASSERT_NE(nnapi_sl_handle.get(), nullptr);
+  int res;
+  uint32_t count;
+  res = nnapi_sl_handle->getFL5()->ANeuralNetworks_getDeviceCount(&count);
+  ASSERT_EQ(res, ANEURALNETWORKS_NO_ERROR);
+  ASSERT_GE(count, 1);
+
+  fbb_.Finish(CreateComputeSettings(
+      fbb_, ExecutionPreference_ANY,
+      CreateTFLiteSettings(
+          fbb_, Delegate_NNAPI,
+          CreateNNAPISettings(
+              fbb_, fbb_.CreateString(accelerator_name ? accelerator_name : ""),
+              /* cache_directory */ 0,
+              /* model_token */ 0, NNAPIExecutionPreference_UNDEFINED,
+              /* no_of_nnapi_instances_to_cache */ 0,
+              /* fallback_settings */ 0,
+              /* allow_nnapi_cpu_on_android_10_plus */ false,
+              /* execution_priority */
+              NNAPIExecutionPriority_NNAPI_PRIORITY_UNDEFINED,
+              /* allow_dynamic_dimensions */ false,
+              /* allow_fp16_precision_for_fp32 */ true,
+              /* use_burst_computation */ false,
+              reinterpret_cast<uint64_t>(nnapi_sl_handle->getFL5())),
+          /* gpu_settings */ 0,
+          /* hexagon_settings */ 0,
+          /* xnnpack_settings */ 0,
+          /* coreml_settings */ 0,
+          /* cpu_settings */ 0,
+          /* max_delegated_partitions */ 1)));
+  AndroidInfo android_info;
+  auto status = RequestAndroidInfo(&android_info);
+  ASSERT_TRUE(status.ok());
+  if (android_info.android_sdk_version >= "30") {
+    CheckValidation("NNAPISL");
+  }
+}
+#endif  // ENABLE_NNAPI_SL_TEST
+
+TEST_F(LocalizerValidationRegressionTest, GpuAny) {
+  AndroidInfo android_info;
+  auto status = RequestAndroidInfo(&android_info);
+  ASSERT_TRUE(status.ok());
+  if (android_info.is_emulator) {
+    std::cerr << "Skipping GPU on emulator\n";
+    return;
+  }
+  fbb_.Finish(CreateComputeSettings(fbb_, ExecutionPreference_ANY,
+                                    CreateTFLiteSettings(fbb_, Delegate_GPU)));
+#ifdef __ANDROID__
+  CheckValidation("GPUANY");
+#endif  // __ANDROID__
+}
+
+TEST_F(LocalizerValidationRegressionTest, GpuOpenGL) {
+  AndroidInfo android_info;
+  auto status = RequestAndroidInfo(&android_info);
+  ASSERT_TRUE(status.ok());
+  if (android_info.is_emulator) {
+    std::cerr << "Skipping GPU on emulator\n";
+    return;
+  }
+  fbb_.Finish(CreateComputeSettings(
+      fbb_, ExecutionPreference_ANY,
+      CreateTFLiteSettings(
+          fbb_, Delegate_GPU, 0,
+          CreateGPUSettings(fbb_, /* allow_precision_loss */ false,
+                            /* allow_quantized_inference */ true,
+                            GPUBackend_OPENGL))));
+#ifdef __ANDROID__
+  CheckValidation("GPUOPENGL");
+#endif  // __ANDROID__
+}
+
+TEST_F(LocalizerValidationRegressionTest, GpuOpenCL) {
+  AndroidInfo android_info;
+  auto status = RequestAndroidInfo(&android_info);
+  ASSERT_TRUE(status.ok());
+  if (android_info.is_emulator) {
+    std::cerr << "Skipping GPU on emulator\n";
+    return;
+  }
+  fbb_.Finish(CreateComputeSettings(
+      fbb_, ExecutionPreference_ANY,
+      CreateTFLiteSettings(
+          fbb_, Delegate_GPU, 0,
+          CreateGPUSettings(fbb_, /* allow_precision_loss */ false,
+                            /* allow_quantized_inference */ true,
+                            GPUBackend_OPENCL))));
+#ifdef __ANDROID__
+  CheckValidation("GPUOPENCL");
+#endif  // __ANDROID__
+}
+
+}  // namespace
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/models/BUILD

@@ -0,0 +1,48 @@
+# Copyright 2021 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.
+# ==============================================================================
+
+# Model files for mini-benchmark tests and examples.
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = [
+        "//tensorflow/lite/experimental/acceleration/mini_benchmark:__subpackages__",
+        "//tensorflow/lite/tools/benchmark:__subpackages__",
+    ],
+    licenses = ["notice"],
+)
+
+exports_files(["blazeface_mlkit_v1.tfl"])
+
+filegroup(
+    name = "add.bin",
+    srcs = [
+        "//tensorflow/lite:testdata/add.bin",
+    ],
+)
+
+filegroup(
+    name = "mobilenet_v1_1.0_224.tflite",
+    srcs = [
+        "@tflite_mobilenet_float//:mobilenet_v1_1.0_224.tflite",
+    ],
+)
+
+filegroup(
+    name = "mobilenet_v1_1.0_224_quant.tflite",
+    srcs = [
+        "@tflite_mobilenet_quant//:mobilenet_v1_1.0_224_quant.tflite",
+    ],
+)

tensorflow/lite/experimental/acceleration/mini_benchmark/nnapi_sl_fake_impl.cc

@@ -0,0 +1,306 @@
+/* Copyright 2021 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 <stdlib.h>
+#include <unistd.h>
+
+#include <cerrno>
+#include <cstdint>
+#include <cstdio>
+#include <fstream>
+#include <iostream>
+#include <memory>
+#include <string>
+
+#include "tensorflow/lite/nnapi/sl/include/SupportLibrary.h"
+#include "tensorflow/lite/nnapi/sl/public/NeuralNetworksSupportLibraryImpl.h"
+
+namespace {
+
+std::string GetTempDir() {
+  const char* temp_dir = getenv("TEST_TMPDIR");
+  if (temp_dir == nullptr || temp_dir[0] == '\0') {
+#ifdef __ANDROID__
+    return "/data/local/tmp";
+#else
+    return "/tmp";
+#endif
+  } else {
+    return temp_dir;
+  }
+}
+
+std::string CallCountFilePath() {
+  return GetTempDir() + "/nnapi_sl_fake_impl.out";
+}
+// We write a . in the trace file to allow a caller to count the number of
+// calls to NNAPI SL.
+void TraceCall() {
+  std::ofstream trace_file(CallCountFilePath().c_str(), std::ofstream::app);
+  if (trace_file) {
+    std::cerr << "Tracing call\n";
+    trace_file << '.';
+    if (!trace_file) {
+      std::cerr << "Error writing to '" << CallCountFilePath() << "'\n";
+    }
+  } else {
+    std::cerr << "FAKE_NNAPI_SL: UNABLE TO TRACE CALL\n";
+  }
+}
+
+template <int return_value, typename... Types>
+int TraceCallAndReturn(Types... args) {
+  TraceCall();
+  return return_value;
+}
+
+template <typename... Types>
+void JustTraceCall(Types... args) {
+  TraceCall();
+}
+
+const uint32_t kDefaultMemoryPaddingAndAlignment = 64;
+
+NnApiSLDriverImplFL5 GetNnApiSlDriverImpl() {
+  NnApiSLDriverImplFL5 sl_driver_impl;
+
+  sl_driver_impl.base = {ANEURALNETWORKS_FEATURE_LEVEL_5};
+  sl_driver_impl.ANeuralNetworksBurst_create =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksBurst_free = JustTraceCall;
+  sl_driver_impl.ANeuralNetworksCompilation_createForDevices =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksCompilation_finish =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksCompilation_free = JustTraceCall;
+  sl_driver_impl
+      .ANeuralNetworksCompilation_getPreferredMemoryAlignmentForInput =
+      [](const ANeuralNetworksCompilation* compilation, uint32_t index,
+         uint32_t* alignment) -> int {
+    TraceCall();
+    *alignment = kDefaultMemoryPaddingAndAlignment;
+    return ANEURALNETWORKS_NO_ERROR;
+  };
+  sl_driver_impl
+      .ANeuralNetworksCompilation_getPreferredMemoryAlignmentForOutput =
+      [](const ANeuralNetworksCompilation* compilation, uint32_t index,
+         uint32_t* alignment) -> int {
+    TraceCall();
+    *alignment = kDefaultMemoryPaddingAndAlignment;
+    return ANEURALNETWORKS_NO_ERROR;
+  };
+  sl_driver_impl.ANeuralNetworksCompilation_getPreferredMemoryPaddingForInput =
+      [](const ANeuralNetworksCompilation* compilation, uint32_t index,
+         uint32_t* padding) -> int {
+    TraceCall();
+    *padding = kDefaultMemoryPaddingAndAlignment;
+    return ANEURALNETWORKS_NO_ERROR;
+  };
+  sl_driver_impl.ANeuralNetworksCompilation_getPreferredMemoryPaddingForOutput =
+      [](const ANeuralNetworksCompilation* compilation, uint32_t index,
+         uint32_t* padding) -> int {
+    TraceCall();
+    *padding = kDefaultMemoryPaddingAndAlignment;
+    return ANEURALNETWORKS_NO_ERROR;
+  };
+  sl_driver_impl.ANeuralNetworksCompilation_setCaching =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksCompilation_setPreference =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksCompilation_setPriority =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksCompilation_setTimeout =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksDevice_getExtensionSupport =
+      [](const ANeuralNetworksDevice* device, const char* extensionName,
+         bool* isExtensionSupported) -> int {
+    *isExtensionSupported = false;
+    return ANEURALNETWORKS_NO_ERROR;
+  };
+  sl_driver_impl.ANeuralNetworksDevice_getFeatureLevel =
+      [](const ANeuralNetworksDevice* device, int64_t* featureLevel) -> int {
+    TraceCall();
+    *featureLevel = ANEURALNETWORKS_FEATURE_LEVEL_5;
+    return ANEURALNETWORKS_NO_ERROR;
+  };
+  sl_driver_impl.ANeuralNetworksDevice_getName =
+      [](const ANeuralNetworksDevice* device, const char** name) -> int {
+    TraceCall();
+    *name = "mockDevice";
+    return ANEURALNETWORKS_BAD_DATA;
+  };
+  sl_driver_impl.ANeuralNetworksDevice_getType =
+      [](const ANeuralNetworksDevice* device, int32_t* type) -> int {
+    *type = ANEURALNETWORKS_DEVICE_CPU;
+    return ANEURALNETWORKS_BAD_DATA;
+  };
+  sl_driver_impl.ANeuralNetworksDevice_getVersion =
+      [](const ANeuralNetworksDevice* device, const char** version) -> int {
+    TraceCall();
+    *version = "mock";
+    return ANEURALNETWORKS_NO_ERROR;
+  };
+  sl_driver_impl.ANeuralNetworksDevice_wait =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksEvent_createFromSyncFenceFd =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksEvent_free = JustTraceCall;
+  sl_driver_impl.ANeuralNetworksEvent_getSyncFenceFd =
+      TraceCallAndReturn<ANEURALNETWORKS_BAD_DATA>;
+  sl_driver_impl.ANeuralNetworksEvent_wait =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksExecution_burstCompute =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksExecution_compute =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksExecution_create =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksExecution_enableInputAndOutputPadding =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksExecution_free = JustTraceCall;
+  sl_driver_impl.ANeuralNetworksExecution_getDuration =
+      [](const ANeuralNetworksExecution* execution, int32_t durationCode,
+         uint64_t* duration) -> int {
+    TraceCall();
+    *duration = UINT64_MAX;
+    return ANEURALNETWORKS_NO_ERROR;
+  };
+  sl_driver_impl.ANeuralNetworksExecution_getOutputOperandDimensions =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksExecution_getOutputOperandRank =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksExecution_setInput =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksExecution_setInputFromMemory =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksExecution_setLoopTimeout =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksExecution_setMeasureTiming =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksExecution_setOutput =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksExecution_setOutputFromMemory =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksExecution_setReusable =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksExecution_setTimeout =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksExecution_startComputeWithDependencies =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksMemoryDesc_addInputRole =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksMemoryDesc_addOutputRole =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksMemoryDesc_create =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksMemoryDesc_finish =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksMemoryDesc_free = JustTraceCall;
+  sl_driver_impl.ANeuralNetworksMemoryDesc_setDimensions =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksMemory_copy =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksMemory_createFromAHardwareBuffer =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksMemory_createFromDesc =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksMemory_createFromFd =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksMemory_free = JustTraceCall;
+  sl_driver_impl.ANeuralNetworksModel_addOperand =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksModel_addOperation =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksModel_create =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksModel_finish =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksModel_free = JustTraceCall;
+  sl_driver_impl.ANeuralNetworksModel_getExtensionOperandType =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksModel_getExtensionOperationType =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksModel_getSupportedOperationsForDevices =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksModel_identifyInputsAndOutputs =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksModel_relaxComputationFloat32toFloat16 =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksModel_setOperandExtensionData =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksModel_setOperandSymmPerChannelQuantParams =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksModel_setOperandValue =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksModel_setOperandValueFromMemory =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworksModel_setOperandValueFromModel =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworks_getDefaultLoopTimeout = []() -> uint64_t {
+    TraceCall();
+    return UINT64_MAX;
+  };
+  sl_driver_impl.ANeuralNetworks_getDevice =
+      TraceCallAndReturn<ANEURALNETWORKS_NO_ERROR>;
+  sl_driver_impl.ANeuralNetworks_getDeviceCount =
+      [](uint32_t* num_devices) -> int {
+    TraceCall();
+    *num_devices = 0;
+    return ANEURALNETWORKS_NO_ERROR;
+  };
+  sl_driver_impl.ANeuralNetworks_getMaximumLoopTimeout = []() -> uint64_t {
+    TraceCall();
+    return UINT64_MAX;
+  };
+  sl_driver_impl.ANeuralNetworks_getRuntimeFeatureLevel = []() -> int64_t {
+    TraceCall();
+    return ANEURALNETWORKS_FEATURE_LEVEL_5;
+  };
+
+  return sl_driver_impl;
+}
+
+}  // namespace
+
+extern "C" NnApiSLDriverImpl* ANeuralNetworks_getSLDriverImpl() {
+  static NnApiSLDriverImplFL5 sl_driver_impl = GetNnApiSlDriverImpl();
+  return reinterpret_cast<NnApiSLDriverImpl*>(&sl_driver_impl);
+}
+
+namespace tflite {
+namespace acceleration {
+
+void InitNnApiSlInvocationStatus() { unlink(CallCountFilePath().c_str()); }
+
+bool WasNnApiSlInvoked() {
+  std::cerr << "Checking if file '" << CallCountFilePath() << "' exists.\n";
+  if (FILE* trace_file = fopen(CallCountFilePath().c_str(), "r")) {
+    fclose(trace_file);
+    return true;
+  } else {
+    return false;
+  }
+}
+
+int CountNnApiSlApiCalls() {
+  FILE* trace_file = fopen(CallCountFilePath().c_str(), "r");
+  int call_count = 0;
+  while (fgetc(trace_file) != EOF) {
+    call_count++;
+  }
+  return call_count;
+}
+
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/nnapi_sl_fake_impl.h

@@ -0,0 +1,33 @@
+/* Copyright 2021 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_NNAPI_SL_FAKE_IMPL_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_NNAPI_SL_FAKE_IMPL_H_
+
+namespace tflite {
+namespace acceleration {
+
+// Initialize the shared file used to check if NNAPI SL has been called
+// and count the number of API calls.
+void InitNnApiSlInvocationStatus();
+
+// Checks if any API call to the fake NNAPI SL has been done.
+bool WasNnApiSlInvoked();
+
+// Returns the number of calls to the fake NNAPI SL.
+int CountNnApiSlApiCalls();
+
+}  // namespace acceleration
+}  // namespace tflite
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_NNAPI_SL_FAKE_IMPL_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/runner.cc

@@ -0,0 +1,426 @@
+/* Copyright 2021 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/runner.h"
+
+#ifndef TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS
+#include <dlfcn.h>
+#endif  // !TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS
+#include <stdio.h>
+#include <stdlib.h>
+#include <sys/stat.h>
+#ifndef _WIN32
+#include <poll.h>
+#include <signal.h>
+#include <unistd.h>
+#endif
+
+#include <cstdlib>
+#include <fstream>
+#include <sstream>
+#include <string>
+#include <thread>  // NOLINT(build/c++11)
+#include <vector>
+
+#include "absl/strings/numbers.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/lite/allocation.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/constants.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/status_codes.h"
+#include "tensorflow/lite/logger.h"
+#include "tensorflow/lite/minimal_logging.h"
+
+#if defined(__ANDROID__) && !defined(TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS)
+#include "tensorflow/lite/experimental/acceleration/compatibility/android_info.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/embedded_runner_executable.h"
+#endif  // __ANDROID__ && !TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS
+
+// Implementation notes and rationale:
+//
+// This class's primary client is the mini-benchmark. The mini-benchmark tries
+// out different acceleration configurations for TFLite. The acceleration may
+// hang or crash due to driver bugs. By Default, the mini-benchmark on Android
+// runs in a separate process that is forked from the host application. This is
+// done to prevent the benchmark from crashing or hanging the host application.
+// All other platforms run the benchmark in the same process. If
+// TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS is defined, the mini-benchmark is
+// forced to run in process..
+//
+// The separate process is implemented in native code. The main alternative
+// would be to use a separate service process at the Android application
+// level. The most important drawbacks of that approach would have been
+// manifest merge issues in tests (see for example b/158142805) and the need for
+// all applications to explicitly integrate the mini-benchmark at the app level.
+//
+// The native code uses popen(3) for the separate process. This is one of the
+// few officially supported ways of safely using fork(2) on Android (See
+// b/129156634 for a discussion on popen use in Android studio device-side, see
+// https://groups.google.com/g/android-platform/c/80jr-_A-9bU for discussion on
+// fork in general).
+//
+// The native process executes a small helper binary that then dynamically
+// loads the shared library that tflite code lives in. This allows for minimal
+// code size as only one copy of the tflite code is needed.
+//
+// The 8kb helper binary is extracted into an app data folder On P- we execute
+// it directly.  On Q+ this is no longer allowed (see b/112357170) but we can
+// use /system/bin/linker(64) as a trampoline.  (See also Chrome's documentation
+// for a similar problem:
+// https://chromium.googlesource.com/chromium/src/+/master/docs/android_native_libraries.md#crashpad-packaging).
+// Using /system/bin/linker(64) does not work before Android Q (See
+// b/112050209).
+//
+// The shared library where the code to be called lives is a JNI library that
+// contains tflite. We detect the path to that library with dladdr(3). This way
+// differences in the way the JNI libraries are bundled, named or delivered is
+// automatically handled. The downside is that dladdr is broken on Android 23
+// for JNI libraries that are not extracted from the apk (See b/37069572). If
+// this becomes a serious issue we can try to parse /proc/self/maps and the apk
+// zip directory to figure out the name. The alternative where the caller needs
+// to pass in the library name requires more integration work at the packaging
+// (app or SDK) level.
+//
+// The methods in this class return detailed error codes, so that telemetry can
+// be used to detect issues in production without using strings which can be
+// hard to make privacy-compliant.
+
+namespace tflite {
+namespace acceleration {
+namespace {
+std::string ShellEscape(const std::string& src);
+}  // namespace
+
+MinibenchmarkStatus ProcessRunner::Init() {
+  if (!function_pointer_) {
+    return kMinibenchmarkPreconditionNotMet;
+  }
+#if !defined(__ANDROID__) || defined(TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS)
+  return kMinibenchmarkSuccess;
+#else  // __ANDROID__ && !TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS
+  tflite::acceleration::AndroidInfo android_info;
+  if (!tflite::acceleration::RequestAndroidInfo(&android_info).ok()) {
+    return kMinibenchmarkRequestAndroidInfoFailed;
+  }
+  if (android_info.android_sdk_version.length() < 2 ||
+      android_info.android_sdk_version < "23") {
+    // The codepaths have only been tested on 23+.
+    return kMinibenchmarkUnsupportedPlatform;
+  }
+
+  // Find name of this shared object.
+  std::string soname;
+  Dl_info dl_info;
+  int status = dladdr(function_pointer_, &dl_info);
+  if (status != 0) {
+    if (dl_info.dli_fname) {
+      soname = dl_info.dli_fname;
+    } else {
+      return kMinibenchmarkDliFnameWasNull;
+    }
+  } else {
+    return kMinibenchmarkDladdrReturnedZero;
+  }
+  // Check for b/37069572 - on SDK level 23 dli_fname may not contain the
+  // library name, only the apk. (See comment at top of file).
+  if (soname.size() >= 4 && soname.substr(soname.size() - 4) == ".apk") {
+    return kMinibenchmarkDliFnameHasApkNameOnly;
+  }
+
+  // Construct path to runner, extracting the helper binary if needed.
+  std::string runner_path;
+  // TODO(b/172541832): handle multiple concurrent callers.
+  runner_path = temporary_path_ + "/runner";
+  (void)unlink(runner_path.c_str());
+  std::string runner_contents(
+      reinterpret_cast<const char*>(g_tflite_acceleration_embedded_runner),
+      g_tflite_acceleration_embedded_runner_len);
+  std::ofstream f(runner_path, std::ios::binary);
+  if (!f.is_open()) {
+    return kMinibenchmarkCouldntOpenTemporaryFileForBinary;
+  }
+  f << runner_contents;
+  f.close();
+  if (chmod(runner_path.c_str(), 0500) != 0) {
+    return kMinibenchmarkCouldntChmodTemporaryFile;
+  }
+  runner_path = ShellEscape(runner_path);
+  if (android_info.android_sdk_version >= "29") {
+    // On 29+ we need to use /system/bin/linker to load the binary from the app,
+    // as exec()ing writable files was blocked for security. (See comment at top
+    // of file).
+#if defined(__arm__) || defined(__i386__)
+    std::string linker_path = "/system/bin/linker";
+#else
+    std::string linker_path = "/system/bin/linker64";
+#endif
+    runner_path = linker_path + " " + runner_path;
+  }
+
+  // Commit.
+  runner_path_ = runner_path;
+  soname_ = soname;
+  return kMinibenchmarkSuccess;
+#endif  // !__ANDROID__ || TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS
+}
+
+// TODO(b/245901066): Refactor the runner to separate Multi-process
+// implementation and in process implementors, and remove the ifdef guards.
+#ifndef _WIN32
+bool ProcessRunner::KillProcessWhenTimedOut(FILE* fstream) {
+  // The first fread() should get subprocess id. It's important to
+  // read the same number of bytes as on the write side, so that this fread()
+  // does not block.
+  const int array_length = 1 + kPidBufferLength;
+  char buffer[array_length];
+  memset(buffer, '\0', array_length);
+  ssize_t length = fread(buffer, 1, kPidBufferLength, fstream);
+  int pid;
+  if (length != kPidBufferLength || !absl::SimpleAtoi(buffer, &pid)) {
+    TF_LITE_REPORT_ERROR(error_reporter_,
+                         "Failed to get Validator subprocess id: %s", buffer);
+    return false;
+  }
+  struct pollfd pfd[1];
+  pfd[0].fd = fileno(fstream);
+  // Wait for the fstream to be closed.
+  pfd[0].events = POLLHUP;
+  int poll_ret = poll(pfd, 1, timeout_millisec_);
+  // Kill the subprocess if timed out.
+  if (poll_ret == 0) {
+    kill(pid, SIGKILL);
+    return true;
+  } else if (poll_ret < 0) {
+    TF_LITE_REPORT_ERROR(error_reporter_, "Validator timer failed: %s",
+                         strerror(errno));
+  }
+  return false;
+}
+#endif  // _WIN32
+
+MinibenchmarkStatus ProcessRunner::Run(const Allocation* model_allocation,
+                                       const std::vector<std::string>& args,
+                                       std::string* output, int* exitcode,
+                                       int* signal) {
+#ifdef _WIN32
+  return kMinibenchmarkUnsupportedPlatform;
+#else  // !_WIN32
+  if (!output || !exitcode) {
+    return kMinibenchmarkPreconditionNotMet;
+  }
+  int benchmark_process_status = 0;
+  MinibenchmarkStatus status = kMinibenchmarkCommandFailed;
+#ifdef TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS
+  if (function_pointer_) {
+    benchmark_process_status = RunInprocess(model_allocation, args);
+  } else {
+    return kMinibenchmarkPreconditionNotMet;
+  }
+#else   // !TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS
+  if (runner_path_.empty()) {
+    return kMinibenchmarkPreconditionNotMet;
+  }
+  // runner_path_ components are escaped earlier.
+  std::string cmd = runner_path_ + " " + ShellEscape(soname_) + " " +
+                    ShellEscape(function_name_);
+
+  // If model is not null, open a pipe() and add pipe model path as cmdline
+  // argv[3]. If model is null, argv[0] should be the model path.
+  int pipe_fds[2];
+  if (model_allocation != nullptr) {
+    if (pipe(pipe_fds) < 0) {
+      *exitcode = errno;
+      return kMinibenchmarkPipeFailed;
+    }
+    std::string pipe_model_path = absl::StrCat(
+        "pipe:", pipe_fds[0], ":", pipe_fds[1], ":", model_allocation->bytes());
+    cmd = cmd + " " + ShellEscape(pipe_model_path);
+  }
+
+  // Add the rest of the cmdline args.
+  for (const auto& arg : args) {
+    cmd = cmd + " " + ShellEscape(arg);
+  }
+
+  FILE* f = popen(cmd.c_str(), "r");
+  if (!f) {
+    *exitcode = errno;
+    return kMinibenchmarkPopenFailed;
+  }
+
+  // Write model to MiniBenchmark process.
+  if (model_allocation != nullptr) {
+    close(pipe_fds[0]);
+    int written_bytes = 0;
+    int remaining_bytes = model_allocation->bytes();
+    const uint8_t* current =
+        static_cast<const uint8_t*>(model_allocation->base());
+    while (remaining_bytes > 0 &&
+           (written_bytes = write(pipe_fds[1], current, remaining_bytes)) > 0) {
+      remaining_bytes -= written_bytes;
+      current += written_bytes;
+    }
+    close(pipe_fds[1]);
+    if (written_bytes <= 0 || remaining_bytes > 0) {
+      *exitcode = errno;
+      return kMinibenchmarkPipeFailed;
+    }
+  }
+
+  // Note: KillProcessWhenTimedOut() will block until f is closed or timeout has
+  // reached. It will cause issue if subprocess is blocked on writing to f.
+  if (timeout_millisec_ > 0 && KillProcessWhenTimedOut(f)) {
+    status = kMinibenchmarkCommandTimedOut;
+    TFLITE_LOG_PROD(
+        TFLITE_LOG_INFO,
+        "Validator did not finish after %dms. Tried to kill the test.",
+        timeout_millisec_);
+  }
+  std::vector<char> buffer(4 * 1024, 0);
+  ssize_t length;
+  std::string ret;
+  do {
+    length = fread(buffer.data(), 1, buffer.size(), f);
+    ret = ret + std::string(buffer.data(), length);
+  } while (length == buffer.size());
+  *output = ret;
+  benchmark_process_status = pclose(f);
+#endif  //  TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS
+  if (WIFEXITED(benchmark_process_status)) {
+    *exitcode = WEXITSTATUS(benchmark_process_status);
+    *signal = 0;
+    if (*exitcode == kMinibenchmarkSuccess) {
+      status = kMinibenchmarkSuccess;
+    }
+  } else if (WIFSIGNALED(benchmark_process_status)) {
+    *exitcode = 0;
+    *signal = WTERMSIG(benchmark_process_status);
+  }
+  return status;
+#endif  // _WIN32
+}
+
+#ifdef TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS
+#ifndef __W_EXITCODE  // Mac
+#define __W_EXITCODE(ret, sig) ((ret) << 8 | (sig))
+#endif
+
+int ProcessRunner::RunInprocess(const Allocation* model_allocation,
+                                const std::vector<std::string>& user_args) {
+  TFLITE_LOG_PROD(TFLITE_LOG_INFO, "Running Validator in-process.");
+  std::vector<std::string> args_string;
+  args_string.push_back("inprocess");
+  args_string.push_back("inprocess");
+  args_string.push_back(function_name_);
+
+  std::thread write_thread;
+  if (model_allocation != nullptr) {
+    int pipe_fds[2];
+    if (pipe(pipe_fds) < 0) {
+      return __W_EXITCODE(kMinibenchmarkPipeFailed, 0);
+    }
+
+    // Add pipe_model_path when model is not null.
+    // Model loader won't close the write file descriptor when it's -1.
+    args_string.push_back(
+        absl::StrCat("pipe:", pipe_fds[0], ":-1:", model_allocation->bytes()));
+
+    // When running MiniBenchmark in-process, we start a separate thread for
+    // writing to pipe.
+    write_thread = std::thread([pipe_fds, model_allocation,
+                                error_reporter = error_reporter_]() {
+      int written_bytes = 0;
+      int remaining_bytes = model_allocation->bytes();
+      const uint8_t* current =
+          static_cast<const uint8_t*>(model_allocation->base());
+      while (remaining_bytes > 0 &&
+             (written_bytes = write(pipe_fds[1], current, remaining_bytes)) >
+                 0) {
+        remaining_bytes -= written_bytes;
+        current += written_bytes;
+      }
+      close(pipe_fds[1]);
+      if (written_bytes < 0 || remaining_bytes > 0) {
+        TF_LITE_REPORT_ERROR(
+            error_reporter,
+            "Failed to write Model to pipe: %s. Expect to write %d "
+            "bytes, %d bytes written.",
+            strerror(errno), remaining_bytes, written_bytes);
+      }
+    });
+  }
+
+  for (int i = 0; i < user_args.size(); i++) {
+    args_string.push_back(user_args[i]);
+  }
+  std::vector<std::vector<char>> args_char(args_string.size());
+  std::vector<char*> argv(args_string.size());
+  for (int i = 0; i < args_string.size(); i++) {
+    args_char[i] = {args_string[i].begin(), args_string[i].end()};
+    // Compiler adds '\0' for std::string to indicate the end of string
+    // automatically. For char* string, '\0' needs to be add at the end of
+    // string manually.
+    args_char[i].push_back('\0');
+    argv[i] = args_char[i].data();
+  }
+
+  int (*function_pointer)(int, char**) =
+      reinterpret_cast<int (*)(int, char**)>(function_pointer_);
+  int exit_code = __W_EXITCODE(function_pointer(argv.size(), argv.data()), 0);
+  if (write_thread.joinable()) {
+    write_thread.join();
+  }
+  return exit_code;
+}
+#endif  // TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS
+
+namespace {
+
+// kDontNeedShellEscapeChars and ShellEscape are copied from absl, which
+// copied them from Python. Copied here because tflite core libraries should
+// not depend on absl (both for size reasons and for possible version skew):
+
+static const char kDontNeedShellEscapeChars[] =
+    "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+    "abcdefghijklmnopqrstuvwxyz"
+    "0123456789+-_.=/:,@";
+
+std::string ShellEscape(const std::string& src) {
+  if (!src.empty() &&  // empty string needs quotes
+      src.find_first_not_of(kDontNeedShellEscapeChars) == std::string::npos) {
+    // only contains chars that don't need quotes; it's fine
+    return src;
+  } else if (src.find('\'') == std::string::npos) {  // NOLINT
+    // no single quotes; just wrap it in single quotes
+    return "'" + src + "'";
+  } else {
+    // needs double quote escaping
+    std::string result = "\"";
+    for (const char c : src) {
+      switch (c) {
+        case '\\':
+        case '$':
+        case '"':
+        case '`':
+          result.push_back('\\');
+      }
+      result.push_back(c);
+    }
+    result.push_back('"');
+    return result;
+  }
+}
+}  // namespace
+
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/runner.h

@@ -0,0 +1,138 @@
+/* Copyright 2021 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_RUNNER_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_RUNNER_H_
+
+#include <cstdio>
+#include <string>
+#include <vector>
+
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/allocation.h"
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/status_codes.h"
+#include "tensorflow/lite/stderr_reporter.h"
+
+namespace tflite {
+namespace acceleration {
+
+// Class that runs a C-main-function -compatible exported symbol in a separate
+// process on Android. Support all Android 24+ devices and some 23 devices. May
+// also work on 22- but has not been tested.
+//
+// Requirements on the caller:
+// - The function to be called must reside in the same shared library as this
+// code and it must be exported (typically by being extern "C" and having a name
+// that starts with Java).
+// - The function to be called must have the same signature as C main:
+//     extern "C" int Java_foo(int argc, char** argv)
+// - The librunner_main.so shared object from this directory must
+// reside in the same location as the shared object above (typically by
+// depending on the runner_main_library_for_deps target from this
+// directory in a java_library or java_binary)
+//
+// The return values are meant to be detailed enough for telemetry.
+//
+// For reasons behind the above restrictions and requirements, see
+// implementation notes in runner.cc
+//
+// Warning: this class will just run the provided code in-process when compiled
+// for non-Android, and timeout is not enforced.
+class ProcessRunner {
+ public:
+  // Construct ProcessRunner. 'temporary_path' should be a suitable subdirectory
+  // of the app data path for extracting the helper binary on Android P-.
+  //
+  // Since the function will be called through popen() only return values
+  // between 0 and 127 are well-defined.
+  ProcessRunner(const std::string& temporary_path,
+                const std::string& function_name,
+                int (*function_pointer)(int argc, char** argv),
+                int timeout_millisec = 0,
+                ErrorReporter* error_reporter = tflite::DefaultErrorReporter())
+      : temporary_path_(temporary_path),
+        function_name_(function_name),
+        function_pointer_(reinterpret_cast<void*>(function_pointer)),
+        timeout_millisec_(timeout_millisec),
+        error_reporter_(error_reporter) {}
+
+  // Initialize runner.
+  MinibenchmarkStatus Init();
+
+  // Run function in separate process. Returns function's output to stdout and
+  // the shell exitcode. Stderr is discarded.
+  //
+  // The function will be called with argc and argv corresponding to a command
+  // line like:
+  //     helper_binary function_name (optional: model path) args
+  // If model_allocation is not null, runner will use pipe() to pass the model
+  // data to subprocess. Otherwise, args[0] should be a model path.
+  // The args are escaped for running through the shell.
+  //
+  // The 'output' and 'exitcode' and `signal` are set as follows based on the
+  // return value:
+  //   kMinibenchmarkUnknownStatus, kMinibenchmarkPreconditionNotMet: undefined
+  //   kMinibenchmarkPopenFailed:
+  //       *output is an empty string
+  //       *exitcode is errno after popen()
+  //       *signal is 0
+  //   kMinibenchmarkCommandFailed, kMinibenchmarkSuccess:
+  //       *output is stdout produced from function
+  //       *exitcode is:
+  //        - if the process terminated normally:
+  //          the return value of the benchmark function or, if function
+  //          loading fails one the MinibenchmarkStatus values listed under
+  //          'Runner main status codes' to describe the failure.
+  //        - if the process has been terminated by a signal: 0
+  //       *signal is:
+  //        - if the process has been terminated by a signal: the signal number
+  //        - 0 otherwise
+  //
+  // To be considered successful, the function must return
+  // kMinibenchmarkSuccess. This is because some GPU drivers call exit(0) as a
+  // bailout and we don't want to confuse that with a successful run.
+  MinibenchmarkStatus Run(const Allocation* model_allocation,
+                          const std::vector<std::string>& args,
+                          std::string* output, int* exitcode, int* signal);
+
+  ProcessRunner(ProcessRunner&) = delete;
+  ProcessRunner& operator=(const ProcessRunner&) = delete;
+
+ private:
+#ifdef TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS
+  int RunInprocess(const Allocation* model_allocation,
+                   const std::vector<std::string>& args);
+#endif  // TFLITE_ACCELERATION_BENCHMARK_IN_PROCESS
+
+#ifndef _WIN32
+  // This function first reads the subprocess id from fstream, and then block
+  // until timeout_millisec_ has passed, OR fstream is closed. If timeout is
+  // reached first, we will kill the subprocess. Returns whether kill() is
+  // triggered.
+  bool KillProcessWhenTimedOut(FILE* fstream);
+#endif  // !_WIN32
+
+  std::string temporary_path_;
+  std::string function_name_;
+  void* function_pointer_;
+  std::string runner_path_;
+  std::string soname_;
+  int timeout_millisec_;
+  ErrorReporter* error_reporter_;
+};
+
+}  // namespace acceleration
+}  // namespace tflite
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_RUNNER_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/runner_main.c

@@ -0,0 +1,67 @@
+/* Copyright 2021 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.
+==============================================================================*/
+
+// LINT.IfChange
+#define DLOPEN_FAILED 11
+#define SYMBOL_LOOKUP_FAILED 12
+#define TOO_FEW_ARGUMENTS 13
+#define UNSUPPORTED_PLATFORM 14
+// LINT.ThenChange(//tensorflow/lite/experimental/acceleration/mini_benchmark/status_codes.h)
+
+#ifndef _WIN32
+#include <dlfcn.h>
+#include <stdio.h>
+
+/**
+ * Mini-benchmark helper executable entry point.
+ *
+ * This is a minimal executable that loads the shared library that contains the
+ * actual logic as not to duplicate code between the mini-benchmark and
+ * applicaiton code that uses TFLite.
+ *
+ * The executable takes 2 or more arguments:
+ *   executable path_to_shared_library symbol_to_call ...
+ * It loads the path_to_shared_library and calls symbol_to_call with all the
+ * arguments (using original argc and argv).
+ *
+ * Exit codes:
+ * - Happy path exits with the return value from symbol_to_call().
+ * - If library loading fails, exits with 1.
+ * - If symbol lookup fails, exits with 2.
+ * - If less than 2 arguments are passed, exits with 3.
+ */
+int main(int argc, char** argv) {
+  if (argc < 3) {
+    return TOO_FEW_ARGUMENTS;
+  }
+  void* lib = dlopen(argv[1], RTLD_NOW | RTLD_LOCAL);
+  if (!lib) {
+    printf("dlopen failed for %s: %s\n", argv[1], dlerror());
+    return DLOPEN_FAILED;
+  }
+  void* sym = dlsym(lib, argv[2]);
+  if (!sym) {
+    printf("dlsym failed for %s on library %s with error: %s\n", argv[2],
+           argv[1], dlerror());
+    return SYMBOL_LOOKUP_FAILED;
+  }
+  int (*f)(int argc, char** argv) = (int (*)(int, char**))sym;
+  int exitcode = (*f)(argc, argv);
+  return exitcode;
+}
+
+#else   // _WIN32
+int main(int argc, char** argv) { return UNSUPPORTED_PLATFORM; }
+#endif  // !_WIN32

tensorflow/lite/experimental/acceleration/mini_benchmark/runner_test.cc

@@ -0,0 +1,258 @@
+/* Copyright 2021 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 "tensorflow/lite/experimental/acceleration/mini_benchmark/runner.h"
+
+#include <dlfcn.h>
+#include <signal.h>
+
+#include <cstddef>
+#include <fstream>
+#include <memory>
+#include <string>
+
+#include <gtest/gtest.h>
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/compiler/mlir/lite/schema/mutable/schema_generated.h"
+#include "tensorflow/lite/allocation.h"
+#include "tensorflow/lite/stderr_reporter.h"
+#ifdef __ANDROID__
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/embedded_runner_executable.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/embedded_runner_unit_test_entry_points.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/mini_benchmark_test_helper.h"
+#endif  // __ANDROID__
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/status_codes.h"
+
+extern "C" {
+int TfLiteFunctionInBinary(int argc, char** argv) { return 2; }
+}  // extern "C"
+
+namespace tflite {
+namespace acceleration {
+namespace {
+
+typedef int (*EntryPoint)(int, char**);
+using flatbuffers::FlatBufferBuilder;
+
+struct RunnerTest : ::testing::Test {
+ protected:
+  void* LoadEntryPointModule() {
+    void* module =
+        dlopen(entry_point_file.c_str(), RTLD_NOW | RTLD_LOCAL | RTLD_NODELETE);
+    EXPECT_TRUE(module) << dlerror();
+    return module;
+  }
+
+  EntryPoint Load(const char* name) {
+#ifdef __ANDROID__
+    void* module = LoadEntryPointModule();
+    if (!module) {
+      return nullptr;
+    }
+#else   // !__ANDROID__
+    auto module = RTLD_DEFAULT;
+#endif  // __ANDROID__
+    void* symbol = dlsym(module, name);
+    return reinterpret_cast<EntryPoint>(symbol);
+  }
+
+  void SetUp() override {
+#ifdef __ANDROID__
+    // We extract the test files here as that's the only way to get the right
+    // architecture when building tests for multiple architectures.
+    entry_point_file = MiniBenchmarkTestHelper::DumpToTempFile(
+        "librunner_unit_test_entry_points.so",
+        g_tflite_acceleration_embedded_runner_unit_test_entry_points,
+        g_tflite_acceleration_embedded_runner_unit_test_entry_points_len);
+    ASSERT_TRUE(!entry_point_file.empty());
+#endif  // __ANDROID__
+  }
+
+  void Init(const char* symbol_name) {
+    EntryPoint function = Load(symbol_name);
+    ASSERT_TRUE(function);
+    runner = std::make_unique<ProcessRunner>(::testing::TempDir(), symbol_name,
+                                             function, timeout_ms);
+    ASSERT_EQ(runner->Init(), kMinibenchmarkSuccess);
+  }
+
+  FlatBufferBuilder CreateTestModel() {
+    ModelT model;
+    model.description = "test";
+    flatbuffers::FlatBufferBuilder fbb;
+    FinishModelBuffer(fbb, CreateModel(fbb, &model));
+    return fbb;
+  }
+  int exitcode = 0;
+  int signal = 0;
+  int timeout_ms = 0;
+  std::string output;
+  std::unique_ptr<ProcessRunner> runner;
+  std::string entry_point_file;
+};
+
+// These tests are only for Android. They are also disabled on 64-bit arm
+// because the 64-bit arm emulator doesn't have a shell that works with popen().
+// These tests are run on x86 emulators.
+#if !defined(__aarch64__)
+TEST_F(RunnerTest, LoadSymbol) {
+  EntryPoint TfLiteJustReturnZero = Load("TfLiteJustReturnZero");
+  ASSERT_TRUE(TfLiteJustReturnZero);
+#ifdef __ANDROID__
+  Dl_info dl_info;
+  int status = dladdr(reinterpret_cast<void*>(TfLiteJustReturnZero), &dl_info);
+  ASSERT_TRUE(status) << dlerror();
+  ASSERT_TRUE(dl_info.dli_fname) << dlerror();
+
+  void* this_module =
+      dlopen(dl_info.dli_fname, RTLD_NOW | RTLD_LOCAL | RTLD_NODELETE);
+  ASSERT_TRUE(this_module);
+  void* symbol = dlsym(this_module, "TfLiteJustReturnZero");
+  EXPECT_TRUE(symbol);
+#endif  // __ANDROID__
+}
+
+TEST_F(RunnerTest, JustReturnZero) {
+  ASSERT_NO_FATAL_FAILURE(Init("TfLiteJustReturnZero"));
+  EXPECT_EQ(runner->Run(nullptr, {}, &output, &exitcode, &signal),
+            kMinibenchmarkCommandFailed);
+  EXPECT_EQ(exitcode, 0);
+  EXPECT_EQ(signal, 0);
+  EXPECT_EQ(output, "");
+}
+
+TEST_F(RunnerTest, ReturnOne) {
+  ASSERT_NO_FATAL_FAILURE(Init("TfLiteReturnOne"));
+  EXPECT_EQ(runner->Run(nullptr, {}, &output, &exitcode, &signal),
+            kMinibenchmarkCommandFailed);
+  EXPECT_EQ(exitcode, 1);
+  EXPECT_EQ(signal, 0);
+  EXPECT_EQ(output, "");
+}
+
+TEST_F(RunnerTest, ReturnSuccess) {
+  ASSERT_NO_FATAL_FAILURE(Init("TfLiteReturnSuccess"));
+  EXPECT_EQ(runner->Run(nullptr, {}, &output, &exitcode, &signal),
+            kMinibenchmarkSuccess);
+  EXPECT_EQ(exitcode, kMinibenchmarkSuccess);
+  EXPECT_EQ(signal, 0);
+  EXPECT_EQ(output, "");
+}
+
+TEST_F(RunnerTest, NullFunctionPointer) {
+  ProcessRunner runner("foo", "bar", nullptr);
+  EXPECT_EQ(runner.Init(), kMinibenchmarkPreconditionNotMet);
+  EXPECT_EQ(runner.Run(nullptr, {}, &output, &exitcode, &signal),
+            kMinibenchmarkPreconditionNotMet);
+}
+
+#ifdef __ANDROID__
+TEST_F(RunnerTest, SigKillSubprocess) {
+  ASSERT_NO_FATAL_FAILURE(Init("TfLiteSigKillSelf"));
+  EXPECT_EQ(runner->Run(nullptr, {}, &output, &exitcode, &signal),
+            kMinibenchmarkCommandFailed);
+  EXPECT_EQ(exitcode, 0);
+  EXPECT_EQ(signal, SIGKILL);
+  EXPECT_EQ(output, "");
+}
+
+TEST_F(RunnerTest, WriteOk) {
+  ASSERT_NO_FATAL_FAILURE(Init("TfLiteWriteOk"));
+  EXPECT_EQ(runner->Run(nullptr, {}, &output, &exitcode, &signal),
+            kMinibenchmarkSuccess);
+  EXPECT_EQ(exitcode, kMinibenchmarkSuccess);
+  EXPECT_EQ(signal, 0);
+  EXPECT_EQ(output, "ok\n");
+}
+
+TEST_F(RunnerTest, Write10kChars) {
+  ASSERT_NO_FATAL_FAILURE(Init("TfLiteWrite10kChars"));
+  EXPECT_EQ(runner->Run(nullptr, {}, &output, &exitcode, &signal),
+            kMinibenchmarkSuccess);
+  EXPECT_EQ(exitcode, kMinibenchmarkSuccess);
+  EXPECT_EQ(signal, 0);
+  EXPECT_EQ(output.size(), 10000);
+}
+
+TEST_F(RunnerTest, ArgsArePassed) {
+  ASSERT_NO_FATAL_FAILURE(Init("TfLiteWriteArgs"));
+  EXPECT_EQ(
+      runner->Run(nullptr, {"foo", "bar", "baz"}, &output, &exitcode, &signal),
+      kMinibenchmarkSuccess);
+  EXPECT_EQ(exitcode, kMinibenchmarkSuccess);
+  EXPECT_EQ(signal, 0);
+  EXPECT_EQ(output, "foo\nbar\nbaz\n");
+}
+
+TEST_F(RunnerTest, SymbolLookupFailed) {
+  ProcessRunner runner(::testing::TempDir(), "TfLiteFunctionInBinary",
+                       TfLiteFunctionInBinary);
+  EXPECT_EQ(runner.Init(), kMinibenchmarkSuccess);
+  EXPECT_EQ(runner.Run(nullptr, {}, &output, &exitcode, &signal),
+            kMinibenchmarkCommandFailed)
+      << output;
+  EXPECT_EQ(exitcode, kMinibenchmarkRunnerMainSymbolLookupFailed) << output;
+}
+
+TEST_F(RunnerTest, ReadModelFromPipe) {
+  ASSERT_NO_FATAL_FAILURE(Init("TfLiteReadFromPipe"));
+  FlatBufferBuilder model = CreateTestModel();
+  MemoryAllocation alloc(model.GetBufferPointer(), model.GetSize(),
+                         tflite::DefaultErrorReporter());
+  EXPECT_EQ(runner->Run(&alloc, {}, &output, &exitcode, &signal),
+            kMinibenchmarkSuccess);
+  EXPECT_EQ(exitcode, kMinibenchmarkSuccess);
+  EXPECT_EQ(signal, 0);
+  EXPECT_EQ(output,
+            std::string((char*)model.GetBufferPointer(), model.GetSize()));
+}
+
+TEST_F(RunnerTest, ProcessTimedOut) {
+  timeout_ms = 1000;
+  ASSERT_NO_FATAL_FAILURE(Init("TfLiteWritePidThenSleepNSec"));
+  // Process will sleep for 30 seconds.
+  EXPECT_EQ(runner->Run(nullptr, {"30"}, &output, &exitcode, &signal),
+            kMinibenchmarkCommandTimedOut);
+  EXPECT_EQ(signal, SIGKILL);
+  EXPECT_EQ(output, "");
+  EXPECT_EQ(exitcode, 0);
+}
+
+TEST_F(RunnerTest, ProcessDoNotTimedOut) {
+  timeout_ms = 3000;
+  ASSERT_NO_FATAL_FAILURE(Init("TfLiteWritePidThenSleepNSec"));
+  // Process will sleep for 1 second.
+  EXPECT_EQ(runner->Run(nullptr, {"1"}, &output, &exitcode, &signal),
+            kMinibenchmarkSuccess);
+  EXPECT_EQ(signal, 0);
+  EXPECT_EQ(output, "");
+  EXPECT_EQ(exitcode, kMinibenchmarkSuccess);
+}
+#else  // __ANDROID__
+
+TEST_F(RunnerTest, ReadModelFromPipeNonAndroid) {
+  ASSERT_NO_FATAL_FAILURE(Init("TfLiteReadFromPipeInProcess"));
+  FlatBufferBuilder model = CreateTestModel();
+  MemoryAllocation alloc(model.GetBufferPointer(), model.GetSize(),
+                         tflite::DefaultErrorReporter());
+  EXPECT_EQ(runner->Run(&alloc, {}, &output, &exitcode, &signal),
+            kMinibenchmarkSuccess);
+}
+
+#endif  // __ANDROID__
+#endif  // !__aarch64__
+
+}  // namespace
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/runner_test_entry_points.cc

@@ -0,0 +1,109 @@
+/* Copyright 2021 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 <signal.h>
+#include <stdio.h>
+#include <string.h>
+#include <sys/types.h>
+#include <unistd.h>
+
+#include <memory>
+#include <string>
+
+#include "absl/strings/numbers.h"
+#include "tensorflow/lite/allocation.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/constants.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/status_codes.h"
+#include "tensorflow/lite/tools/model_loader.h"
+
+extern "C" {
+
+constexpr int kStdOutFd = 1;
+
+int TfLiteJustReturnZero(int argc, char** argv) { return 0; }
+
+int TfLiteReturnOne(int argc, char** argv) { return 1; }
+
+int TfLiteReturnSuccess(int argc, char** argv) {
+  return ::tflite::acceleration::kMinibenchmarkSuccess;
+}
+
+int TfLiteSigKillSelf(int argc, char** argv) {
+  kill(getpid(), SIGKILL);
+  return 1;
+}
+
+int TfLiteWriteOk(int argc, char** argv) {
+  if (write(kStdOutFd, "ok\n", 3) == -1) {
+    return -1;
+  }
+  return ::tflite::acceleration::kMinibenchmarkSuccess;
+}
+
+// Write the pid to output stream and then sleep N seconds. N is parsed from
+// argv[3].
+int TfLiteWritePidThenSleepNSec(int argc, char** argv) {
+  std::string pid = std::to_string(getpid());
+  pid.resize(::tflite::acceleration::kPidBufferLength);
+  if (write(kStdOutFd, pid.data(), ::tflite::acceleration::kPidBufferLength) ==
+      -1) {
+    return 1;
+  }
+
+  int sleep_sec;
+  if (!absl::SimpleAtoi(argv[3], &sleep_sec)) {
+    return 1;
+  }
+  sleep(sleep_sec);
+  return ::tflite::acceleration::kMinibenchmarkSuccess;
+}
+
+int TfLiteWrite10kChars(int argc, char** argv) {
+  char buffer[10000];
+  memset(buffer, 'A', 10000);
+  return write(kStdOutFd, buffer, 10000) == 10000
+             ? ::tflite::acceleration::kMinibenchmarkSuccess
+             : 1;
+}
+
+int TfLiteWriteArgs(int argc, char** argv) {
+  for (int i = 3; i < argc; i++) {
+    if (write(1, argv[i], strlen(argv[i])) == -1 || write(1, "\n", 1) == -1) {
+      return 1;
+    }
+  }
+  return ::tflite::acceleration::kMinibenchmarkSuccess;
+}
+
+int TfLiteReadFromPipe(int argc, char** argv) {
+  std::unique_ptr<tflite::tools::ModelLoader> model_loader =
+      tflite::tools::CreateModelLoaderFromPath(argv[3]);
+  const tflite::Allocation* alloc;
+  if (!model_loader->Init() ||
+      !(alloc = model_loader->GetModel()->allocation()) || !alloc->base()) {
+    return 1;
+  }
+  return write(kStdOutFd, alloc->base(), alloc->bytes()) == alloc->bytes()
+             ? ::tflite::acceleration::kMinibenchmarkSuccess
+             : 1;
+}
+
+int TfLiteReadFromPipeInProcess(int argc, char** argv) {
+  std::unique_ptr<tflite::tools::ModelLoader> model_loader =
+      tflite::tools::CreateModelLoaderFromPath(argv[3]);
+  return model_loader->Init() ? ::tflite::acceleration::kMinibenchmarkSuccess
+                              : 1;
+}
+
+}  // extern "C"