patch/0015-mtk-preemption.patch

From 1c747fd6365504125688579d75e74f58ff3b4ec0 Mon Sep 17 00:00:00 2001
From: Tommy Chiang <ototot@google.com>
Date: Tue, 22 Oct 2024 01:47:18 +0000
Subject: [PATCH] Add QoS/Preemption options to MtkNeuronSettings

PiperOrigin-RevId: 687135593
PATCH_NAME=mtk-preemption
---
 .../configuration/configuration.proto         | 11 +++++
 .../configuration/configuration_generated.h   | 46 ++++++++++++++++---
 .../configuration/flatbuffer_to_proto.cc      |  4 ++
 .../configuration/flatbuffer_to_proto_test.cc |  4 ++
 .../configuration/proto_to_flatbuffer.cc      |  3 +-
 .../configuration/proto_to_flatbuffer_test.cc |  8 ++++
 6 files changed, 69 insertions(+), 7 deletions(-)

diff --git a/tensorflow/lite/acceleration/configuration/configuration.proto b/tensorflow/lite/acceleration/configuration/configuration.proto
index 657ae5e1..33bc6d95 100644
--- a/tensorflow/lite/acceleration/configuration/configuration.proto
+++ b/tensorflow/lite/acceleration/configuration/configuration.proto
@@ -787,6 +787,17 @@ message MtkNeuronSettings {
   // Optional path to the platform-dependent Neuron configuration file.
   // See docs at https://neuropilot.mediatek.com/ for more details.
   optional string neuron_config_path = 10;
+
+  // The deadline time duration (in ms) of the inference (waiting + execution).
+  // The scheduler would adjust scheduling based on this value. Note that
+  // setting this value to zero implies no deadline requirement.
+  optional int32 inference_deadline_ms = 11;
+
+  // The maximum inference (waiting + execution) time duration (in ms). The
+  // scheduler would abort the inference if the inference time dutation exceed
+  // the time specified. Note that setting this value to zero implies no abort
+  // time requirement.
+  optional int32 inference_abort_time_ms = 12;
 }
 
 // How to configure TFLite.
diff --git a/tensorflow/lite/acceleration/configuration/configuration_generated.h b/tensorflow/lite/acceleration/configuration/configuration_generated.h
index 2820f7c4..ef43aa6f 100644
--- a/tensorflow/lite/acceleration/configuration/configuration_generated.h
+++ b/tensorflow/lite/acceleration/configuration/configuration_generated.h
@@ -2835,6 +2835,8 @@ struct MtkNeuronSettingsT : public ::flatbuffers::NativeTable {
   std::vector<std::string> compile_options{};
   std::vector<std::string> accelerator_names{};
   std::string neuron_config_path{};
+  int32_t inference_deadline_ms = 0;
+  int32_t inference_abort_time_ms = 0;
 };
 
 struct MtkNeuronSettings FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
@@ -2850,7 +2852,9 @@ struct MtkNeuronSettings FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table
     VT_USE_CACHEABLE_BUFFER = 16,
     VT_COMPILE_OPTIONS = 18,
     VT_ACCELERATOR_NAMES = 20,
-    VT_NEURON_CONFIG_PATH = 22
+    VT_NEURON_CONFIG_PATH = 22,
+    VT_INFERENCE_DEADLINE_MS = 24,
+    VT_INFERENCE_ABORT_TIME_MS = 26
   };
   tflite::MtkNeuronSettings_::ExecutionPreference execution_preference() const {
     return static_cast<tflite::MtkNeuronSettings_::ExecutionPreference>(GetField<int32_t>(VT_EXECUTION_PREFERENCE, 0));
@@ -2882,6 +2886,12 @@ struct MtkNeuronSettings FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table
   const ::flatbuffers::String *neuron_config_path() const {
     return GetPointer<const ::flatbuffers::String *>(VT_NEURON_CONFIG_PATH);
   }
+  int32_t inference_deadline_ms() const {
+    return GetField<int32_t>(VT_INFERENCE_DEADLINE_MS, 0);
+  }
+  int32_t inference_abort_time_ms() const {
+    return GetField<int32_t>(VT_INFERENCE_ABORT_TIME_MS, 0);
+  }
   bool Verify(::flatbuffers::Verifier &verifier) const {
     return VerifyTableStart(verifier) &&
            VerifyField<int32_t>(verifier, VT_EXECUTION_PREFERENCE, 4) &&
@@ -2900,6 +2910,8 @@ struct MtkNeuronSettings FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table
            verifier.VerifyVectorOfStrings(accelerator_names()) &&
            VerifyOffset(verifier, VT_NEURON_CONFIG_PATH) &&
            verifier.VerifyString(neuron_config_path()) &&
+           VerifyField<int32_t>(verifier, VT_INFERENCE_DEADLINE_MS, 4) &&
+           VerifyField<int32_t>(verifier, VT_INFERENCE_ABORT_TIME_MS, 4) &&
            verifier.EndTable();
   }
   MtkNeuronSettingsT *UnPack(const ::flatbuffers::resolver_function_t *_resolver = nullptr) const;
@@ -2941,6 +2953,12 @@ struct MtkNeuronSettingsBuilder {
   void add_neuron_config_path(::flatbuffers::Offset<::flatbuffers::String> neuron_config_path) {
     fbb_.AddOffset(MtkNeuronSettings::VT_NEURON_CONFIG_PATH, neuron_config_path);
   }
+  void add_inference_deadline_ms(int32_t inference_deadline_ms) {
+    fbb_.AddElement<int32_t>(MtkNeuronSettings::VT_INFERENCE_DEADLINE_MS, inference_deadline_ms, 0);
+  }
+  void add_inference_abort_time_ms(int32_t inference_abort_time_ms) {
+    fbb_.AddElement<int32_t>(MtkNeuronSettings::VT_INFERENCE_ABORT_TIME_MS, inference_abort_time_ms, 0);
+  }
   explicit MtkNeuronSettingsBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
         : fbb_(_fbb) {
     start_ = fbb_.StartTable();
@@ -2963,8 +2981,12 @@ inline ::flatbuffers::Offset<MtkNeuronSettings> CreateMtkNeuronSettings(
     bool use_cacheable_buffer = true,
     ::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<::flatbuffers::String>>> compile_options = 0,
     ::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<::flatbuffers::String>>> accelerator_names = 0,
-    ::flatbuffers::Offset<::flatbuffers::String> neuron_config_path = 0) {
+    ::flatbuffers::Offset<::flatbuffers::String> neuron_config_path = 0,
+    int32_t inference_deadline_ms = 0,
+    int32_t inference_abort_time_ms = 0) {
   MtkNeuronSettingsBuilder builder_(_fbb);
+  builder_.add_inference_abort_time_ms(inference_abort_time_ms);
+  builder_.add_inference_deadline_ms(inference_deadline_ms);
   builder_.add_neuron_config_path(neuron_config_path);
   builder_.add_accelerator_names(accelerator_names);
   builder_.add_compile_options(compile_options);
@@ -2989,7 +3011,9 @@ inline ::flatbuffers::Offset<MtkNeuronSettings> CreateMtkNeuronSettingsDirect(
     bool use_cacheable_buffer = true,
     const std::vector<::flatbuffers::Offset<::flatbuffers::String>> *compile_options = nullptr,
     const std::vector<::flatbuffers::Offset<::flatbuffers::String>> *accelerator_names = nullptr,
-    const char *neuron_config_path = nullptr) {
+    const char *neuron_config_path = nullptr,
+    int32_t inference_deadline_ms = 0,
+    int32_t inference_abort_time_ms = 0) {
   auto optimization_hints__ = optimization_hints ? _fbb.CreateVector<int32_t>(*optimization_hints) : 0;
   auto compile_options__ = compile_options ? _fbb.CreateVector<::flatbuffers::Offset<::flatbuffers::String>>(*compile_options) : 0;
   auto accelerator_names__ = accelerator_names ? _fbb.CreateVector<::flatbuffers::Offset<::flatbuffers::String>>(*accelerator_names) : 0;
@@ -3005,7 +3029,9 @@ inline ::flatbuffers::Offset<MtkNeuronSettings> CreateMtkNeuronSettingsDirect(
       use_cacheable_buffer,
       compile_options__,
       accelerator_names__,
-      neuron_config_path__);
+      neuron_config_path__,
+      inference_deadline_ms,
+      inference_abort_time_ms);
 }
 
 ::flatbuffers::Offset<MtkNeuronSettings> CreateMtkNeuronSettings(::flatbuffers::FlatBufferBuilder &_fbb, const MtkNeuronSettingsT *_o, const ::flatbuffers::rehasher_function_t *_rehasher = nullptr);
@@ -5534,7 +5560,9 @@ inline bool operator==(const MtkNeuronSettingsT &lhs, const MtkNeuronSettingsT &
       (lhs.use_cacheable_buffer == rhs.use_cacheable_buffer) &&
       (lhs.compile_options == rhs.compile_options) &&
       (lhs.accelerator_names == rhs.accelerator_names) &&
-      (lhs.neuron_config_path == rhs.neuron_config_path);
+      (lhs.neuron_config_path == rhs.neuron_config_path) &&
+      (lhs.inference_deadline_ms == rhs.inference_deadline_ms) &&
+      (lhs.inference_abort_time_ms == rhs.inference_abort_time_ms);
 }
 
 inline bool operator!=(const MtkNeuronSettingsT &lhs, const MtkNeuronSettingsT &rhs) {
@@ -5561,6 +5589,8 @@ inline void MtkNeuronSettings::UnPackTo(MtkNeuronSettingsT *_o, const ::flatbuff
   { auto _e = compile_options(); if (_e) { _o->compile_options.resize(_e->size()); for (::flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->compile_options[_i] = _e->Get(_i)->str(); } } else { _o->compile_options.resize(0); } }
   { auto _e = accelerator_names(); if (_e) { _o->accelerator_names.resize(_e->size()); for (::flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->accelerator_names[_i] = _e->Get(_i)->str(); } } else { _o->accelerator_names.resize(0); } }
   { auto _e = neuron_config_path(); if (_e) _o->neuron_config_path = _e->str(); }
+  { auto _e = inference_deadline_ms(); _o->inference_deadline_ms = _e; }
+  { auto _e = inference_abort_time_ms(); _o->inference_abort_time_ms = _e; }
 }
 
 inline ::flatbuffers::Offset<MtkNeuronSettings> MtkNeuronSettings::Pack(::flatbuffers::FlatBufferBuilder &_fbb, const MtkNeuronSettingsT* _o, const ::flatbuffers::rehasher_function_t *_rehasher) {
@@ -5581,6 +5611,8 @@ inline ::flatbuffers::Offset<MtkNeuronSettings> CreateMtkNeuronSettings(::flatbu
   auto _compile_options = _o->compile_options.size() ? _fbb.CreateVectorOfStrings(_o->compile_options) : 0;
   auto _accelerator_names = _o->accelerator_names.size() ? _fbb.CreateVectorOfStrings(_o->accelerator_names) : 0;
   auto _neuron_config_path = _o->neuron_config_path.empty() ? 0 : _fbb.CreateString(_o->neuron_config_path);
+  auto _inference_deadline_ms = _o->inference_deadline_ms;
+  auto _inference_abort_time_ms = _o->inference_abort_time_ms;
   return tflite::CreateMtkNeuronSettings(
       _fbb,
       _execution_preference,
@@ -5592,7 +5624,9 @@ inline ::flatbuffers::Offset<MtkNeuronSettings> CreateMtkNeuronSettings(::flatbu
       _use_cacheable_buffer,
       _compile_options,
       _accelerator_names,
-      _neuron_config_path);
+      _neuron_config_path,
+      _inference_deadline_ms,
+      _inference_abort_time_ms);
 }
 
 
diff --git a/tensorflow/lite/acceleration/configuration/flatbuffer_to_proto.cc b/tensorflow/lite/acceleration/configuration/flatbuffer_to_proto.cc
index 87f822cf..8ed43c97 100644
--- a/tensorflow/lite/acceleration/configuration/flatbuffer_to_proto.cc
+++ b/tensorflow/lite/acceleration/configuration/flatbuffer_to_proto.cc
@@ -469,6 +469,10 @@ proto::MtkNeuronSettings ConvertMtkNeuronSettings(
     proto_settings.set_neuron_config_path(settings.neuron_config_path()->str());
   }
 
+  proto_settings.set_inference_deadline_ms(settings.inference_deadline_ms());
+  proto_settings.set_inference_abort_time_ms(
+      settings.inference_abort_time_ms());
+
   return proto_settings;
 }
 
diff --git a/tensorflow/lite/acceleration/configuration/flatbuffer_to_proto_test.cc b/tensorflow/lite/acceleration/configuration/flatbuffer_to_proto_test.cc
index bd5ef446..7ac336c8 100644
--- a/tensorflow/lite/acceleration/configuration/flatbuffer_to_proto_test.cc
+++ b/tensorflow/lite/acceleration/configuration/flatbuffer_to_proto_test.cc
@@ -572,6 +572,8 @@ TEST_F(ConversionTest, MtkNeuronSettings) {
   input_settings->compile_options = {"TEST_COMPILE_OPTIONS"};
   input_settings->accelerator_names = {"TEST_ACCELERATOR_NAME"};
   input_settings->neuron_config_path = "TEST_NEURON_CONFIG_PATH";
+  input_settings->inference_deadline_ms = 1337;
+  input_settings->inference_abort_time_ms = 42;
 
   const proto::ComputeSettings compute = ConvertFromFlatbuffer(settings_);
   const proto::MtkNeuronSettings& output_settings =
@@ -596,6 +598,8 @@ TEST_F(ConversionTest, MtkNeuronSettings) {
   EXPECT_EQ(output_settings.accelerator_names().size(), 1);
   EXPECT_EQ(output_settings.accelerator_names().at(0), "TEST_ACCELERATOR_NAME");
   EXPECT_EQ(output_settings.neuron_config_path(), "TEST_NEURON_CONFIG_PATH");
+  EXPECT_EQ(output_settings.inference_deadline_ms(), 1337);
+  EXPECT_EQ(output_settings.inference_abort_time_ms(), 42);
 }
 
 TEST_F(ConversionTest, MiniBenchmarkSettings) {
diff --git a/tensorflow/lite/acceleration/configuration/proto_to_flatbuffer.cc b/tensorflow/lite/acceleration/configuration/proto_to_flatbuffer.cc
index 366869ba..a35c259e 100644
--- a/tensorflow/lite/acceleration/configuration/proto_to_flatbuffer.cc
+++ b/tensorflow/lite/acceleration/configuration/proto_to_flatbuffer.cc
@@ -428,7 +428,8 @@ Offset<MtkNeuronSettings> ConvertMtkNeuronSettings(
                                     settings.compile_options().end()),
       builder.CreateVectorOfStrings(settings.accelerator_names().begin(),
                                     settings.accelerator_names().end()),
-      builder.CreateString(settings.neuron_config_path()));
+      builder.CreateString(settings.neuron_config_path()),
+      settings.inference_deadline_ms(), settings.inference_abort_time_ms());
 }
 
 Offset<CoralSettings> ConvertCoralSettings(const proto::CoralSettings& settings,
diff --git a/tensorflow/lite/acceleration/configuration/proto_to_flatbuffer_test.cc b/tensorflow/lite/acceleration/configuration/proto_to_flatbuffer_test.cc
index e504f09c..c0db23fa 100644
--- a/tensorflow/lite/acceleration/configuration/proto_to_flatbuffer_test.cc
+++ b/tensorflow/lite/acceleration/configuration/proto_to_flatbuffer_test.cc
@@ -184,6 +184,8 @@ TEST(ConversionTest, MtkNeuronSettings) {
   const std::string kCompileOptions = "TEST_COMPILE_OPTIONS";
   const std::string kAcceleratorName = "TEST_ACCELERATOR_NAME";
   const std::string kNeuronConfigPath = "TEST_NEURON_CONFIG_PATH";
+  const int32_t kInferenceDeadlineMs = 1337;
+  const int32_t kInferenceAbortTimeMs = 42;
 
   // Create the proto settings.
   proto::TFLiteSettings input_settings;
@@ -198,6 +200,8 @@ TEST(ConversionTest, MtkNeuronSettings) {
   mtk_neuron_settings->add_compile_options(kCompileOptions);
   mtk_neuron_settings->add_accelerator_names(kAcceleratorName);
   mtk_neuron_settings->set_neuron_config_path(kNeuronConfigPath);
+  mtk_neuron_settings->set_inference_deadline_ms(kInferenceDeadlineMs);
+  mtk_neuron_settings->set_inference_abort_time_ms(kInferenceAbortTimeMs);
   flatbuffers::FlatBufferBuilder flatbuffers_builder;
 
   // Convert.
@@ -231,6 +235,10 @@ TEST(ConversionTest, MtkNeuronSettings) {
             kAcceleratorName);
   EXPECT_EQ(output_mtk_neuron_settings->neuron_config_path()->str(),
             kNeuronConfigPath);
+  EXPECT_EQ(output_mtk_neuron_settings->inference_deadline_ms(),
+            kInferenceDeadlineMs);
+  EXPECT_EQ(output_mtk_neuron_settings->inference_abort_time_ms(),
+            kInferenceAbortTimeMs);
 }
 
 }  // namespace