From df7525018e329de897044365deaf0752a3e678be Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 8 Jan 2025 14:36:27 -0800 Subject: [PATCH] Fix application of JIT compiler plugins 1) Restore some key logic lost when landing cl/707770943, in compiled_model.cpp:122 2) Don't abort CompiledModel creation if the runtime fails to apply compiler plugins, rather issue warnings 3) Log the list of compiler plugins that were successfully applied PiperOrigin-RevId: 713420589 --- .../litert/c/litert_compiled_model_test.cc | 2 +- .../litert/cc/litert_compiled_model.h | 2 +- .../experimental/litert/compiler/plugin/BUILD | 2 + .../litert/compiler/plugin/compiler_plugin.cc | 58 +++++++++++++++---- .../litert/compiler/plugin/compiler_plugin.h | 19 +++++- .../litert/runtime/compiled_model.cc | 28 ++++++--- .../litert/runtime/compiled_model_test.cc | 4 +- .../litert_dispatch_invocation_context.cc | 1 - 8 files changed, 89 insertions(+), 27 deletions(-) diff --git a/tensorflow/lite/experimental/litert/c/litert_compiled_model_test.cc b/tensorflow/lite/experimental/litert/c/litert_compiled_model_test.cc index f7d2bad73b1..705be3d5ddb 100644 --- a/tensorflow/lite/experimental/litert/c/litert_compiled_model_test.cc +++ b/tensorflow/lite/experimental/litert/c/litert_compiled_model_test.cc @@ -45,7 +45,7 @@ TEST(CompiledModelTest, Basic) { LiteRtCompiledModel compiled_model; ASSERT_EQ( - LiteRtCreateCompiledModel(model, kLiteRtHwAccelatorNone, &compiled_model), + LiteRtCreateCompiledModel(model, kLiteRtHwAccelatorCpu, &compiled_model), kLiteRtStatusOk); LiteRtSubgraph subgraph; diff --git a/tensorflow/lite/experimental/litert/cc/litert_compiled_model.h b/tensorflow/lite/experimental/litert/cc/litert_compiled_model.h index 9d41ca5db68..9b9499eef65 100644 --- a/tensorflow/lite/experimental/litert/cc/litert_compiled_model.h +++ b/tensorflow/lite/experimental/litert/cc/litert_compiled_model.h @@ -69,7 +69,7 @@ class CompiledModel // returned object. static Expected Create( litert::Model& model, - LiteRtCompilationOptions compilation_options = kLiteRtHwAccelatorNone) { + LiteRtCompilationOptions compilation_options = kLiteRtHwAccelatorCpu) { LiteRtCompiledModel compiled_model; if (auto status = LiteRtCreateCompiledModel( model.Get(), compilation_options, &compiled_model); diff --git a/tensorflow/lite/experimental/litert/compiler/plugin/BUILD b/tensorflow/lite/experimental/litert/compiler/plugin/BUILD index 1967f4d3f4a..6ed275db1d6 100644 --- a/tensorflow/lite/experimental/litert/compiler/plugin/BUILD +++ b/tensorflow/lite/experimental/litert/compiler/plugin/BUILD @@ -41,6 +41,8 @@ cc_library( "//tensorflow/lite/experimental/litert/vendors/c:litert_compiler_plugin", "//tensorflow/lite/experimental/litert/vendors/c:litert_compiler_plugin_api", "@com_google_absl//absl/log:absl_check", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/types:span", ], diff --git a/tensorflow/lite/experimental/litert/compiler/plugin/compiler_plugin.cc b/tensorflow/lite/experimental/litert/compiler/plugin/compiler_plugin.cc index 5a825403a7c..57b078ab8aa 100644 --- a/tensorflow/lite/experimental/litert/compiler/plugin/compiler_plugin.cc +++ b/tensorflow/lite/experimental/litert/compiler/plugin/compiler_plugin.cc @@ -24,6 +24,9 @@ #include #include "absl/log/absl_check.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_format.h" +#include "absl/strings/str_join.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "tensorflow/lite/experimental/litert/c/litert_any.h" @@ -291,6 +294,16 @@ CompilerPlugin::~CompilerPlugin() { } } +std::string CompilerPlugin::DebugString() const { + std::string version_str = "?"; + if (auto version = ApiVersion(); version) { + version_str = absl::StrFormat("%d.%d.%d", version->major, version->minor, + version->patch); + } + return absl::StrFormat("%s compiler plugin (ver %s)", SocManufacturer(), + version_str); +} + Expected CompilerPlugin::ApiVersion() const { LiteRtApiVersion api_version; LITERT_EXPECT_OK(plugin_api_.get_compiler_plugin_version(&api_version)); @@ -426,7 +439,7 @@ Expected ApplyPlugin(CompilerPlugin& compiler_plugin, LiteRtModelT& model, return {}; } -Expected> ApplyPlugins( +Expected ApplyPlugins( LiteRtModel model, LiteRtHwAccelerators selected_hw_accelerators) { auto environment = litert::internal::Environment::Instance(); if (!environment) { @@ -448,13 +461,25 @@ Expected> ApplyPlugins( if (!compiler_plugins) { return compiler_plugins.Error(); } + if (compiler_plugins->empty()) { + return litert::Error(kLiteRtStatusErrorRuntimeFailure, + "No compiler plugin found"); + } - std::optional> new_flatbuffer; + OwningBufferRef new_flatbuffer; + std::vector success_messages; + std::vector error_messages; + ApplyPluginsResult result; + result.num_applied_plugins = 0; for (auto& compiler_plugin : *compiler_plugins) { + auto plugin_name = compiler_plugin.DebugString(); + auto plugin_supported_hardware = compiler_plugin.SupportedHardware(); if (!plugin_supported_hardware) { - return plugin_supported_hardware.Error(); + error_messages.push_back(absl::StrCat( + plugin_name, " ", plugin_supported_hardware.Error().Message())); + continue; } if (*plugin_supported_hardware & selected_hw_accelerators) { @@ -462,28 +487,39 @@ Expected> ApplyPlugins( // shouldn't be needing to serialize a model to then read it again from // the serialized buffer when applying a compiler plugin. if (auto status = ApplyPlugin(compiler_plugin, *model); !status) { - return status.Error(); + error_messages.push_back( + absl::StrCat(plugin_name, " ", status.Error().Message())); + continue; } + auto serialized_model = litert::internal::SerializeModel(std::move(*model)); if (!serialized_model) { - return serialized_model.Error(); + error_messages.push_back( + absl::StrCat(plugin_name, " ", serialized_model.Error().Message())); + continue; } + auto new_model = litert::Model::CreateFromBuffer(*serialized_model); if (!new_model) { - return new_model.Error(); + error_messages.push_back( + absl::StrCat(plugin_name, " ", new_model.Error().Message())); + continue; } + new_flatbuffer = std::move(*serialized_model); *model = std::move(*new_model->Get()); + + success_messages.push_back(absl::StrCat(plugin_name)); + result.num_applied_plugins++; } } - if (!new_flatbuffer.has_value()) { - return litert::Error(kLiteRtStatusErrorRuntimeFailure, - "No applicable compiler plugin found"); - } + result.new_flatbuffer = std::move(new_flatbuffer); + result.success_message = absl::StrJoin(success_messages, ", "); + result.error_message = absl::StrJoin(error_messages, ", "); - return *new_flatbuffer; + return result; } } // namespace litert::internal diff --git a/tensorflow/lite/experimental/litert/compiler/plugin/compiler_plugin.h b/tensorflow/lite/experimental/litert/compiler/plugin/compiler_plugin.h index fb856294354..f3b93293a60 100644 --- a/tensorflow/lite/experimental/litert/compiler/plugin/compiler_plugin.h +++ b/tensorflow/lite/experimental/litert/compiler/plugin/compiler_plugin.h @@ -16,6 +16,7 @@ #define TENSORFLOW_LITE_EXPERIMENTAL_LITERT_COMPILER_PLUGIN_COMPILER_PLUGIN_H_ #include +#include #include #include "absl/strings/string_view.h" @@ -72,6 +73,8 @@ class CompiledResult { // Wraps vendor compiler plugin. class CompilerPlugin { public: + std::string DebugString() const; + // Get the compiler plugin's API version. Expected ApiVersion() const; @@ -150,9 +153,19 @@ Expected ApplyPlugin( Serialization serialization = Serialization::kAppend); // Apply all available plugins providing the selected HW accelerators to the -// given model, modify the model accordingly, and return a new flatbuffer -// backing the modified model. -Expected> ApplyPlugins( +// given model, modify the model accordingly, and return (1) the number of +// compiler plugins succesfully applied, (2) a new flatbuffer backing the +// modified model, (3) a string listing the compiler plugins that were +// succesfully applied, and (4) a string listing the compiler plugins that +// failed to apply with an associated error message. +struct ApplyPluginsResult { + size_t num_applied_plugins; + OwningBufferRef new_flatbuffer; + std::string success_message; + std::string error_message; +}; + +Expected ApplyPlugins( LiteRtModel model, LiteRtHwAccelerators selected_hw_accelerators); } // namespace litert::internal diff --git a/tensorflow/lite/experimental/litert/runtime/compiled_model.cc b/tensorflow/lite/experimental/litert/runtime/compiled_model.cc index 7779ad66529..a3b24dffb00 100644 --- a/tensorflow/lite/experimental/litert/runtime/compiled_model.cc +++ b/tensorflow/lite/experimental/litert/runtime/compiled_model.cc @@ -94,14 +94,23 @@ Expected LiteRtCompiledModelT::Create( std::optional> new_flatbuffer; // TODO: b/379317134 - Support other delegates with compilation options. if (compilation_options != kLiteRtHwAccelatorNone) { - LITERT_LOG(LITERT_INFO, "Applying compiler plugins"); - if (auto flatbuffer = + LITERT_LOG(LITERT_INFO, "Applying compiler plugins..."); + if (auto result = litert::internal::ApplyPlugins(model, compilation_options); - !flatbuffer) { - LITERT_LOG(LITERT_ERROR, "Failed to applying compiler plugins"); - return flatbuffer.Error(); + !result) { + LITERT_LOG(LITERT_WARNING, "Failed to apply compiler plugins: %s", + result.Error().Message().data()); } else { - new_flatbuffer = *flatbuffer; + if (result->num_applied_plugins > 0) { + LITERT_LOG(LITERT_INFO, "Successfully applied %d compiler plugins: %s", + result->num_applied_plugins, + result->success_message.c_str()); + new_flatbuffer = std::move(result->new_flatbuffer); + } + if (!result->error_message.empty()) { + LITERT_LOG(LITERT_WARNING, "Some compiler plugins failed to apply: %s", + result->error_message.c_str()); + } } } @@ -109,8 +118,11 @@ Expected LiteRtCompiledModelT::Create( size_t model_buffer_size = 0; // The following code gets the original FB pointer from LiteRtModel. // TODO b/383120429 - Use a better way of getting the FB pointer. - auto init_model_buffer = detail::GetTflInitFlatbuffer(*model); - if (init_model_buffer.Size() != 0) { + if (new_flatbuffer) { + model_buffer = reinterpret_cast(new_flatbuffer->Data()); + model_buffer_size = new_flatbuffer->Size(); + } else if (auto init_model_buffer = detail::GetTflInitFlatbuffer(*model); + init_model_buffer.Size() != 0) { // Use the saved the original FB pointer when the LiteRtModel was created // from a buffer. model_buffer = init_model_buffer.StrData(); diff --git a/tensorflow/lite/experimental/litert/runtime/compiled_model_test.cc b/tensorflow/lite/experimental/litert/runtime/compiled_model_test.cc index 7e4b4d4924a..45730efb511 100644 --- a/tensorflow/lite/experimental/litert/runtime/compiled_model_test.cc +++ b/tensorflow/lite/experimental/litert/runtime/compiled_model_test.cc @@ -137,7 +137,7 @@ TEST(CompiledModelTest, Basic) { ASSERT_EQ(LiteRtCreateModelFromFile(path.c_str(), &model), kLiteRtStatusOk); auto res_compiled_model = - LiteRtCompiledModelT::Create(model, kLiteRtHwAccelatorNone); + LiteRtCompiledModelT::Create(model, kLiteRtHwAccelatorCpu); ASSERT_TRUE(res_compiled_model) << "Failed to initialize CompiledModel: " << res_compiled_model.Error().Message(); auto& compiled_model = **res_compiled_model; @@ -216,7 +216,7 @@ TEST(CompiledModelTest, UseAhwbBuffer) { ASSERT_EQ(LiteRtCreateModelFromFile(path.c_str(), &model), kLiteRtStatusOk); auto res_compiled_model = - LiteRtCompiledModelT::Create(model, kLiteRtHwAccelatorNone); + LiteRtCompiledModelT::Create(model, kLiteRtHwAccelatorCpu); ASSERT_TRUE(res_compiled_model) << "Failed to initialize CompiledModel"; auto& compiled_model = **res_compiled_model; diff --git a/tensorflow/lite/experimental/litert/vendors/mediatek/dispatch/litert_dispatch_invocation_context.cc b/tensorflow/lite/experimental/litert/vendors/mediatek/dispatch/litert_dispatch_invocation_context.cc index d740bd224c5..33b6d1b9d54 100644 --- a/tensorflow/lite/experimental/litert/vendors/mediatek/dispatch/litert_dispatch_invocation_context.cc +++ b/tensorflow/lite/experimental/litert/vendors/mediatek/dispatch/litert_dispatch_invocation_context.cc @@ -82,7 +82,6 @@ Expected> LoadFromDlaBytecode( const litert::mediatek::NeuronAdapter& neuron_adapter, const void* bytecode_addr, size_t bytecode_size, int num_inputs, int num_outputs) { - LITERT_LOG(LITERT_INFO, "Creating model..."); Expected model = neuron_adapter.CreateModel(); if (!model) { return model.Error();