Change the compile function type to take a model rather than a list of subgraphs

PiperOrigin-RevId: 723715188

Author: LukeBoyer

tensorflow/lite/experimental/litert/compiler/plugin/BUILD

@@ -37,6 +37,7 @@ cc_library(
         "//tensorflow/lite/experimental/litert/core/model",
         "//tensorflow/lite/experimental/litert/core/model:ir_allocator",
         "//tensorflow/lite/experimental/litert/core/model:model_serialize",
+        "//tensorflow/lite/experimental/litert/core/util:flatbuffer_tools",
         "//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",

tensorflow/lite/experimental/litert/compiler/plugin/compiler_plugin.cc

@@ -18,6 +18,7 @@
 #include <array>
 #include <cstddef>
 #include <cstdint>
+#include <memory>
 #include <optional>
 #include <string>
 #include <utility>
@@ -46,6 +47,7 @@
 #include "tensorflow/lite/experimental/litert/core/model/ir_allocator.h"
 #include "tensorflow/lite/experimental/litert/core/model/model.h"
 #include "tensorflow/lite/experimental/litert/core/model/model_serialize.h"
+#include "tensorflow/lite/experimental/litert/core/util/flatbuffer_tools.h"
 #include "tensorflow/lite/experimental/litert/vendors/c/litert_compiler_plugin.h"
 #include "tensorflow/lite/experimental/litert/vendors/c/litert_compiler_plugin_api.h"
 
@@ -339,16 +341,16 @@ Expected<std::vector<LiteRtOp>> CompilerPlugin::Partition(
   return ops.Vec();
 }
 
-Expected<CompiledResult> CompilerPlugin::Compile(
-    absl::Span<LiteRtSubgraph> partitions, absl::string_view soc_model) {
+Expected<CompiledResult> CompilerPlugin::Compile(LiteRtModel partitions,
+                                                 absl::string_view soc_model) {
   CompiledResult result = MakeResult();
   // If the user has passed an soc_model, then we use it; otherwise we let the
   // backend pick the appropriate one by passing nullptr as soc_model. This is
   // important for on-device compilation, where the backend must determine the
   // SoC model based on the user device.
   const char* soc_model_str = !soc_model.empty() ? soc_model.data() : nullptr;
   LITERT_RETURN_IF_ERROR(plugin_api_.compiler_plugin_compile(
-      plugin_handle_, soc_model_str, partitions.data(), partitions.size(),
+      plugin_handle_, soc_model_str, partitions,
       &result.compiled_result_handle_));
   return result;
 }
@@ -408,9 +410,23 @@ Expected<void> ApplyPlugin(CompilerPlugin& compiler_plugin, LiteRtModelT& model,
   auto& dispatch_ops = partitions->first;
   auto& subgraphs = partitions->second;
 
+  // Wrap the partitioned subgraphs in a LiteRtModel.
+  LiteRtModelT sliced_model;
+  sliced_model.TransferSubgraphs(std::move(subgraphs));
+
+  // Copy op codes.
+  const auto& op_codes = detail::GetTflOpCodes(model);
+
+  LiteRtModelT::TflOpCodes codes;
+  codes.reserve(op_codes.size());
+  for (const auto& op_code : op_codes) {
+    codes.emplace_back(std::make_unique<TflOpCode>(*op_code));
+  }
+
+  detail::SetTflOpCodes(sliced_model, std::move(codes));
+
   // Pass sliced subgraphs to plugin for compilation.
-  auto compiled_result =
-      compiler_plugin.Compile(subgraphs.Elements(), soc_model);
+  auto compiled_result = compiler_plugin.Compile(&sliced_model, soc_model);
   if (!compiled_result) {
     return compiled_result.Error();
   }

tensorflow/lite/experimental/litert/compiler/plugin/compiler_plugin.h

@@ -103,7 +103,7 @@ class CompilerPlugin {
 
   // Compile given LiteRtSubgraphs. Result object must be outlived by
   // this CompilerPlugin.
-  Expected<CompiledResult> Compile(absl::Span<LiteRtSubgraph> partitions,
+  Expected<CompiledResult> Compile(LiteRtModel partitions,
                                    absl::string_view soc_model = "");
 
   // Search for shared library files with prefix "libLiteRtCompilerPlugin" in

tensorflow/lite/experimental/litert/compiler/plugin/compiler_plugin_test.cc

@@ -122,7 +122,7 @@ TEST(CompilerPluginTest, Compile) {
   auto model_wrap = testing::LoadTestFileModel("mul_simple.tflite");
   auto& model = *model_wrap.Get();
 
-  auto result = plugins->front().Compile(model.Subgraphs());
+  auto result = plugins->front().Compile(&model);
   ASSERT_TRUE(result);
 
   auto byte_code = result->ByteCode();

tensorflow/lite/experimental/litert/tools/apply_plugin.cc

@@ -361,8 +361,7 @@ LiteRtStatus Compile(Context& ctx) {
   ctx.Dump().Start("Compiling");
   DumpCompilationRequest(ctx.Dump(), ctx.SocModelTarget(),
                          model.NumSubgraphs());
-  auto compilation_result =
-      plugin->Compile(model.Subgraphs(), ctx.SocModelTarget());
+  auto compilation_result = plugin->Compile(&model, ctx.SocModelTarget());
   if (!compilation_result) {
     ctx.Dump().Fail();
     return compilation_result.Error().Status();

tensorflow/lite/experimental/litert/vendors/c/litert_compiler_plugin.h

@@ -65,14 +65,12 @@ LiteRtStatus LiteRtCompilerPluginPartition(LiteRtCompilerPlugin compiler_plugin,
                                            LiteRtSubgraph subgraph,
                                            LiteRtOpList selected_ops);
 
-// Prepare result to pass to the runtime for given partition and, optionally,
-// for a given SoC model (parameter `soc_model` can be NULL to specify a default
-// SoC model). The given subgraphs are valid sub-DAG within the ops selected in
-// partition step.
+// Prepare result to pass to the runtime for given model containing partitioned
+// subgraphs. Optionally, handles a SoC model (parameter `soc_model` can be NULL
+// to specify a default SoC model).
 LiteRtStatus LiteRtCompilerPluginCompile(LiteRtCompilerPlugin compiler_plugin,
                                          const char* soc_model,
-                                         LiteRtSubgraph* partitions,
-                                         LiteRtParamIndex num_partitions,
+                                         LiteRtModel partitions,
                                          LiteRtCompiledResult* compiled_result);
 
 //

tensorflow/lite/experimental/litert/vendors/c/litert_compiler_plugin_api.h

@@ -54,8 +54,8 @@ typedef LiteRtStatus (*LiteRtCompilerPluginPartitionT)(
     LiteRtCompilerPlugin, LiteRtSubgraph subgraph, LiteRtOpList selected_ops);
 
 typedef LiteRtStatus (*LiteRtCompilerPluginCompileT)(
-    LiteRtCompilerPlugin, const char* soc_model, LiteRtSubgraph* partitions,
-    LiteRtParamIndex num_partitions, LiteRtCompiledResult* compiled_result);
+    LiteRtCompilerPlugin, const char* soc_model, LiteRtModel partitions,
+    LiteRtCompiledResult* compiled_result);
 
 typedef void (*LiteRtDestroyCompiledResultT)(LiteRtCompiledResult);
 

tensorflow/lite/experimental/litert/vendors/examples/example_plugin.cc

@@ -88,12 +88,14 @@ LiteRtStatus CompileSinglePartition(LiteRtParamIndex partition_index,
 
 LiteRtStatus LiteRtCompilerPluginCompile(
     LiteRtCompilerPlugin compiler_plugin, const char* soc_model,
-    LiteRtSubgraph* partitions, LiteRtParamIndex num_partitions,
-    LiteRtCompiledResult* compiled_result) {
+    LiteRtModel partitions, LiteRtCompiledResult* compiled_result) {
   LiteRtCompiledResult result = new LiteRtCompiledResultT;
 
+  auto model = litert::Model::CreateFromNonOwnedHandle(partitions);
+  const auto num_partitions = model.NumSubgraphs();
   for (auto i = 0; i < num_partitions; ++i) {
-    LITERT_RETURN_IF_ERROR(CompileSinglePartition(i, partitions[i], *result));
+    LITERT_RETURN_IF_ERROR(
+        CompileSinglePartition(i, model.Subgraph(i)->Get(), *result));
   }
 
   *compiled_result = result;

tensorflow/lite/experimental/litert/vendors/examples/example_plugin_test.cc

@@ -65,12 +65,9 @@ TEST(TestCallDummyPlugin, CompileMulSubgraph) {
   auto plugin = CreatePlugin();
   auto model = testing::LoadTestFileModel("mul_simple.tflite");
 
-  auto main_subgraph = model.MainSubgraph();
-  LiteRtSubgraph litert_subgraph = main_subgraph->Get();
-
   LiteRtCompiledResult compiled;
   LITERT_ASSERT_OK(LiteRtCompilerPluginCompile(
-      plugin.get(), /*soc_model=*/nullptr, &litert_subgraph, 1, &compiled));
+      plugin.get(), /*soc_model=*/nullptr, model.Get(), &compiled));
 
   const void* byte_code;
   size_t byte_code_size;

tensorflow/lite/experimental/litert/vendors/examples/example_plugin_with_conversions.cc

@@ -118,15 +118,16 @@ LiteRtStatus CompileSinglePartition(
 // infrastructure.
 LiteRtStatus LiteRtCompilerPluginCompile(
     LiteRtCompilerPlugin compiler_plugin, const char* soc_model,
-    LiteRtSubgraph* partitions, LiteRtParamIndex num_partitions,
-    LiteRtCompiledResult* compiled_result) {
+    LiteRtModel partitions, LiteRtCompiledResult* compiled_result) {
   auto* result = new LiteRtCompiledResultT;
 
+  auto model = litert::Model::CreateFromNonOwnedHandle(partitions);
+  const auto num_partitions = model.NumSubgraphs();
   for (auto i = 0; i < num_partitions; ++i) {
     auto name = absl::StrFormat("partition_%lu", i);
     LITERT_RETURN_IF_ERROR(
         CompileSinglePartition(compiler_plugin->legalizations, std::move(name),
-                               partitions[i], *result));
+                               model.Subgraph(i)->Get(), *result));
   }
 
   *compiled_result = result;

tensorflow/lite/experimental/litert/vendors/examples/example_plugin_with_conversions_test.cc

@@ -69,13 +69,9 @@ TEST(ExamplePluginWithConvertTypesTest, CompileMulSubgraph) {
   auto plugin = CreatePlugin();
   auto model = litert::testing::LoadTestFileModel("mul_simple.tflite");
 
-  auto main_subgraph = model.MainSubgraph();
-  LiteRtSubgraph litert_subgraph = main_subgraph->Get();
-
   LiteRtCompiledResult compiled;
   LITERT_ASSERT_OK(LiteRtCompilerPluginCompile(
-      plugin.get(), /*soc_model=*/nullptr, &litert_subgraph,
-      /*num_partitions*/ 1, &compiled));
+      plugin.get(), /*soc_model=*/nullptr, model.Get(), &compiled));
 
   const void* byte_code;
   size_t byte_code_size;

tensorflow/lite/experimental/litert/vendors/google_tensor/compiler/google_tensor_compiler_plugin_test.cc

@@ -65,12 +65,9 @@ TEST(TestCallGoogleTensorPlugin, CompileMulSubgraph) {
   auto plugin = CreatePlugin();
   auto model = testing::LoadTestFileModel("mul_simple.tflite");
 
-  auto main_subgraph = model.MainSubgraph();
-  LiteRtSubgraph litert_subgraph = main_subgraph->Get();
-
   LiteRtCompiledResult compiled;
-  LITERT_ASSERT_OK(LiteRtCompilerPluginCompile(plugin.get(), "P25",
-                                               &litert_subgraph, 1, &compiled));
+  LITERT_ASSERT_OK(
+      LiteRtCompilerPluginCompile(plugin.get(), "P25", model.Get(), &compiled));
 
   LiteRtDestroyCompiledResult(compiled);
 }  // Todo(abhirs): activate this test once the compiler wrapper is updated

tensorflow/lite/experimental/litert/vendors/mediatek/compiler/compiler_plugin.cc

@@ -283,8 +283,10 @@ Expected<std::vector<uint8_t>> CompilePartition(
 
 LiteRtStatus LiteRtCompilerPluginCompile(
     LiteRtCompilerPlugin compiler_plugin, const char* soc_model,
-    LiteRtSubgraph* partitions, LiteRtParamIndex num_partitions,
-    LiteRtCompiledResult* compiled_result) {
+    LiteRtModel partitions, LiteRtCompiledResult* compiled_result) {
+  auto model = litert::Model::CreateFromNonOwnedHandle(partitions);
+  const auto num_partitions = model.NumSubgraphs();
+
   LITERT_LOG(LITERT_INFO,
              "Starting MediaTek Compilation for %d subgraphs, soc_model=%s",
              num_partitions, soc_model);
@@ -306,11 +308,11 @@ LiteRtStatus LiteRtCompilerPluginCompile(
   }
 
   auto result = std::make_unique<LiteRtCompiledResultT>();
+
   for (auto i = 0; i < num_partitions; ++i) {
-    auto partition = litert::Subgraph(partitions[i]);
     auto graph_name = absl::StrFormat("Partition_%d", i);
     auto bytecode =
-        CompilePartition(**api, partition, graph_name, opt_soc_model);
+        CompilePartition(**api, *model.Subgraph(i), graph_name, opt_soc_model);
     if (!bytecode) {
       LITERT_LOG(LITERT_INFO, "%s", bytecode.Error().Message().c_str());
       return bytecode.Error().Status();

tensorflow/lite/experimental/litert/vendors/mediatek/compiler/compiler_plugin_test.cc

@@ -25,7 +25,6 @@
 #include "tensorflow/lite/experimental/litert/cc/litert_macros.h"
 #include "tensorflow/lite/experimental/litert/core/model/model.h"
 #include "tensorflow/lite/experimental/litert/test/common.h"
-#include "tensorflow/lite/experimental/litert/test/matchers.h"
 #include "tensorflow/lite/experimental/litert/test/test_models.h"
 #include "tensorflow/lite/experimental/litert/vendors/c/litert_compiler_plugin.h"
 #include "tensorflow/lite/experimental/litert/vendors/cc/litert_compiler_plugin.h"
@@ -43,18 +42,24 @@ const auto kSupportedOps = Values(
 // clang-format on
 
 TEST(TestQnnPlugin, GetConfigInfo) {
+#ifndef __ANDROID__
+  GTEST_SKIP() << "Loading shared lib not currently supported on linux.";
+#endif  // __ANDROID__
+
   EXPECT_STREQ(LiteRtGetCompilerPluginSocManufacturer(), "MediaTek");
 
   auto plugin = CreatePlugin();
 
   LiteRtParamIndex num_supported_soc_models;
-  LITERT_ASSERT_OK(LiteRtGetNumCompilerPluginSupportedSocModels(
-      plugin.get(), &num_supported_soc_models));
+  ASSERT_EQ(LiteRtGetNumCompilerPluginSupportedSocModels(
+                plugin.get(), &num_supported_soc_models),
+            kLiteRtStatusOk);
   ASSERT_EQ(num_supported_soc_models, 12);
 
   const char* config_id;
-  LITERT_ASSERT_OK(
-      LiteRtGetCompilerPluginSupportedSocModel(plugin.get(), 0, &config_id));
+  ASSERT_EQ(
+      LiteRtGetCompilerPluginSupportedSocModel(plugin.get(), 0, &config_id),
+      kLiteRtStatusOk);
   EXPECT_STREQ(config_id, "mt6853");
 }
 
@@ -63,8 +68,9 @@ TEST(TestQnnPlugin, PartitionAdd) {
   auto model = testing::LoadTestFileModel("add_simple.tflite");
 
   LiteRtOpListT selected_op_list;
-  LITERT_ASSERT_OK(LiteRtCompilerPluginPartition(
-      plugin.get(), model.Subgraph(0)->Get(), &selected_op_list));
+  ASSERT_EQ(LiteRtCompilerPluginPartition(
+                plugin.get(), model.Subgraph(0)->Get(), &selected_op_list),
+            kLiteRtStatusOk);
   const auto selected_ops = selected_op_list.Vec();
 
   ASSERT_EQ(selected_ops.size(), 1);
@@ -77,27 +83,30 @@ class MtkPluginOpCompatibilityTest
     : public ::testing::TestWithParam<std::string> {};
 
 TEST_P(MtkPluginOpCompatibilityTest, SupportedOpsTest) {
+#ifndef __ANDROID__
+  GTEST_SKIP() << "Loading shared lib not currently supported on linux.";
+#endif  // __ANDROID__
+
   LITERT_LOG(LITERT_INFO, "Testing TFLite model: %s", GetParam().c_str());
   auto plugin = CreatePlugin();
   auto model = testing::LoadTestFileModel(GetParam());
 
-  const auto subgraph = model.MainSubgraph();
-  LiteRtSubgraph litert_subgraph = subgraph->Get();
-
   LiteRtCompiledResult compiled;
-  LITERT_ASSERT_OK(LiteRtCompilerPluginCompile(
-      plugin.get(), /*soc_model=*/nullptr, &litert_subgraph, 1, &compiled));
+  ASSERT_EQ(LiteRtCompilerPluginCompile(plugin.get(), /*soc_model=*/nullptr,
+                                        model.Get(), &compiled),
+            kLiteRtStatusOk);
 
   LiteRtParamIndex num_byte_code;
-  LITERT_ASSERT_OK(
-      LiteRtCompiledResultNumByteCodeModules(compiled, &num_byte_code));
+  ASSERT_EQ(LiteRtCompiledResultNumByteCodeModules(compiled, &num_byte_code),
+            kLiteRtStatusOk);
   ASSERT_EQ(num_byte_code, 1);
 
   const void* byte_code;
   size_t byte_code_size;
 
-  LITERT_ASSERT_OK(LiteRtGetCompiledResultByteCode(
-      compiled, /*byte_code_idx=*/0, &byte_code, &byte_code_size));
+  ASSERT_EQ(LiteRtGetCompiledResultByteCode(compiled, /*byte_code_idx=*/0,
+                                            &byte_code, &byte_code_size),
+            kLiteRtStatusOk);
 
   absl::string_view byte_code_string(reinterpret_cast<const char*>(byte_code),
                                      byte_code_size);
@@ -107,8 +116,9 @@ TEST_P(MtkPluginOpCompatibilityTest, SupportedOpsTest) {
   size_t op_data_size;
   LiteRtParamIndex byte_code_idx;
 
-  LITERT_ASSERT_OK(LiteRtGetCompiledResultCallInfo(
-      compiled, /*call_idx=*/0, &op_data, &op_data_size, &byte_code_idx));
+  ASSERT_EQ(LiteRtGetCompiledResultCallInfo(compiled, /*call_idx=*/0, &op_data,
+                                            &op_data_size, &byte_code_idx),
+            kLiteRtStatusOk);
 
   EXPECT_EQ(byte_code_idx, 0);
 

tensorflow/lite/experimental/litert/vendors/qualcomm/compiler/qnn_compiler_plugin.cc

@@ -257,8 +257,10 @@ LiteRtStatus LiteRtCompilerPluginPartition(LiteRtCompilerPlugin compiler_plugin,
 
 LiteRtStatus LiteRtCompilerPluginCompile(
     LiteRtCompilerPlugin compiler_plugin, const char* soc_model,
-    LiteRtSubgraph* partitions, LiteRtParamIndex num_partitions,
-    LiteRtCompiledResult* compiled_result) {
+    LiteRtModel partitions, LiteRtCompiledResult* compiled_result) {
+  auto model = litert::Model::CreateFromNonOwnedHandle(partitions);
+  const auto num_partitions = model.NumSubgraphs();
+
   LITERT_LOG(LITERT_INFO,
              "Starting QNN Compilation for %d subgraphs, soc_model=%s",
              num_partitions, soc_model);
@@ -304,8 +306,9 @@ LiteRtStatus LiteRtCompilerPluginCompile(
   {
     std::string& entry_point_name = result->graph_names.emplace_back();
     entry_point_name = "qnn_partition_0";
+    LiteRtSubgraph partition = model.Subgraph(0)->Get();
     LITERT_RETURN_IF_ERROR(litert::qnn::ComposeGraph(
-        **qnn_manager, context_handle->get(), partitions[0], entry_point_name));
+        **qnn_manager, context_handle->get(), partition, entry_point_name));
   }
   LITERT_LOG(LITERT_INFO, "%s", "Graph composed");
 

tensorflow/lite/experimental/litert/vendors/qualcomm/compiler/qnn_compiler_plugin_test.cc

@@ -123,12 +123,9 @@ TEST(TestQnnPlugin, CompileMulSubgraph) {
   auto plugin = CreatePlugin();
   auto model = testing::LoadTestFileModel("one_mul.tflite");
 
-  const auto subgraph = model.MainSubgraph();
-  LiteRtSubgraph litert_subgraph = subgraph->Get();
-
   LiteRtCompiledResult compiled;
-  LITERT_ASSERT_OK(LiteRtCompilerPluginCompile(plugin.get(), "V75",
-                                               &litert_subgraph, 1, &compiled));
+  LITERT_ASSERT_OK(
+      LiteRtCompilerPluginCompile(plugin.get(), "V75", model.Get(), &compiled));
 
   const void* byte_code;
   size_t byte_code_size;
@@ -254,12 +251,9 @@ TEST_P(QnnPluginOpCompatibilityTest, SupportedOpsTest) {
   auto plugin = CreatePlugin();
   auto model = testing::LoadTestFileModel(GetParam());
 
-  const auto subgraph = model.MainSubgraph();
-  LiteRtSubgraph litert_subgraph = subgraph->Get();
-
   LiteRtCompiledResult compiled;
-  LITERT_ASSERT_OK(LiteRtCompilerPluginCompile(plugin.get(), "V75",
-                                               &litert_subgraph, 1, &compiled));
+  LITERT_ASSERT_OK(
+      LiteRtCompilerPluginCompile(plugin.get(), "V75", model.Get(), &compiled));
 
   const void* byte_code;
   size_t byte_code_size;