Prepare compiler partition option and add flag parse/unparse helper for tooling usage.

LiteRT-PiperOrigin-RevId: 828691730

Author: YunanAZ

litert/c/options/litert_compiler_options.cc

@@ -23,6 +23,8 @@
 #include "litert/core/cache/hash_util.h"
 
 struct LiteRtCompilerOptionsT {
+  LiteRtCompilerOptionsPartitionStrategy partition_strategy =
+      kLiteRtCompilerOptionsPartitionStrategyDefault;
   bool dummy_option = false;
 };
 
@@ -66,6 +68,26 @@ LiteRtStatus LiteRtFindCompilerOptions(
 
 const char* LiteRtGetCompilerOptionsIdentifier() { return "litert_compiler"; }
 
+LiteRtStatus LiteRtSetCompilerOptionsPartitionStrategy(
+    LiteRtCompilerOptions options,
+    LiteRtCompilerOptionsPartitionStrategy partition_strategy) {
+  if (options == nullptr) {
+    return kLiteRtStatusErrorInvalidArgument;
+  }
+  options->partition_strategy = partition_strategy;
+  return kLiteRtStatusOk;
+}
+
+LiteRtStatus LiteRtGetCompilerOptionsPartitionStrategy(
+    LiteRtCompilerOptionsConst options,
+    LiteRtCompilerOptionsPartitionStrategy* partition_strategy) {
+  if (options == nullptr || partition_strategy == nullptr) {
+    return kLiteRtStatusErrorInvalidArgument;
+  }
+  *partition_strategy = options->partition_strategy;
+  return kLiteRtStatusOk;
+}
+
 LiteRtStatus LiteRtSetDummyCompilerOptions(LiteRtCompilerOptions options,
                                            bool dummy_option) {
   if (options == nullptr) {

litert/c/options/litert_compiler_options.h

@@ -26,6 +26,22 @@ extern "C" {
 
 LITERT_DEFINE_HANDLE(LiteRtCompilerOptions);
 
+// Partition strategy for the compiler.
+//
+// Default:
+//   Default partition strategy is to partition the graph based on the naive
+//   cutting algorithm.
+//
+// Weakly connected (experimental feature):
+//   Partitions the graph into weakly connected components, this may result in
+//   incorrect partition where there exists a path from a selected to another,
+//   where there is an unselected node in between. This feature is currently in
+//   experimental stage.
+typedef enum LiteRtCompilerOptionsPartitionStrategy {
+  kLiteRtCompilerOptionsPartitionStrategyDefault = 0,
+  kLiteRtCompilerOptionsPartitionStrategyWeaklyConnected = 1,
+} LiteRtCompilerOptionsPartitionStrategy;
+
 // Creates an opaque options object holding Compiler options.
 LiteRtStatus LiteRtCreateCompilerOptions(LiteRtOpaqueOptions* options);
 
@@ -36,6 +52,15 @@ LiteRtStatus LiteRtFindCompilerOptions(LiteRtOpaqueOptions opaque_options,
 // Gets the identifier for Compiler options stored in opaque options.
 const char* LiteRtGetCompilerOptionsIdentifier();
 
+// Sets the partition strategy for the compiler.
+LiteRtStatus LiteRtSetCompilerOptionsPartitionStrategy(
+    LiteRtCompilerOptions options,
+    LiteRtCompilerOptionsPartitionStrategy partition_strategy);
+
+LiteRtStatus LiteRtGetCompilerOptionsPartitionStrategy(
+    LiteRtCompilerOptionsConst options,
+    LiteRtCompilerOptionsPartitionStrategy* partition_strategy);
+
 // Dummy options for testing.
 LiteRtStatus LiteRtSetDummyCompilerOptions(LiteRtCompilerOptions options,
                                            bool dummy_option);

litert/c/options/litert_compiler_options_test.cc

@@ -57,6 +57,38 @@ TEST(LiteRtCompilerOptionsTest, DummyOptions) {
   LiteRtDestroyOpaqueOptions(options);
 }
 
+TEST(LiteRtCompilerOptionsTest, GetDefaultPartitionStrategy) {
+  LiteRtOpaqueOptions options;
+  LITERT_ASSERT_OK(LiteRtCreateCompilerOptions(&options));
+  LiteRtCompilerOptions compiler_options;
+  LITERT_ASSERT_OK(LiteRtFindCompilerOptions(options, &compiler_options));
+
+  LiteRtCompilerOptionsPartitionStrategy partition_strategy;
+  LITERT_ASSERT_OK(LiteRtGetCompilerOptionsPartitionStrategy(
+      compiler_options, &partition_strategy));
+  EXPECT_EQ(partition_strategy, kLiteRtCompilerOptionsPartitionStrategyDefault);
+
+  LiteRtDestroyOpaqueOptions(options);
+}
+
+TEST(LiteRtCompilerOptionsTest, SetAndGetWCPartitionStrategy) {
+  LiteRtOpaqueOptions options;
+  LITERT_ASSERT_OK(LiteRtCreateCompilerOptions(&options));
+  LiteRtCompilerOptions compiler_options;
+  LITERT_ASSERT_OK(LiteRtFindCompilerOptions(options, &compiler_options));
+
+  LiteRtCompilerOptionsPartitionStrategy partition_strategy;
+  LITERT_ASSERT_OK(LiteRtSetCompilerOptionsPartitionStrategy(
+      compiler_options,
+      kLiteRtCompilerOptionsPartitionStrategyWeaklyConnected));
+  LITERT_ASSERT_OK(LiteRtGetCompilerOptionsPartitionStrategy(
+      compiler_options, &partition_strategy));
+  EXPECT_EQ(partition_strategy,
+            kLiteRtCompilerOptionsPartitionStrategyWeaklyConnected);
+
+  LiteRtDestroyOpaqueOptions(options);
+}
+
 TEST(LiteRtCompilerOptionsTest, Hash) {
   LiteRtOpaqueOptions options1;
   LITERT_ASSERT_OK(LiteRtCreateCompilerOptions(&options1));
@@ -71,6 +103,9 @@ TEST(LiteRtCompilerOptionsTest, Hash) {
   LiteRtCompilerOptions compiler_options;
   LITERT_ASSERT_OK(LiteRtFindCompilerOptions(options1, &compiler_options));
   LITERT_ASSERT_OK(LiteRtSetDummyCompilerOptions(compiler_options, true));
+  LITERT_ASSERT_OK(LiteRtSetCompilerOptionsPartitionStrategy(
+      compiler_options,
+      kLiteRtCompilerOptionsPartitionStrategyWeaklyConnected));
   LITERT_ASSERT_OK(LiteRtGetOpaqueOptionsHash(options1, &hash1));
   LITERT_ASSERT_OK(LiteRtGetOpaqueOptionsHash(options2, &hash2));
   EXPECT_NE(hash1, hash2);

litert/cc/options/BUILD

@@ -80,6 +80,7 @@ cc_test(
     srcs = ["compiler_options_test.cc"],
     deps = [
         ":compiler_options",
+        "//litert/c/options:litert_compiler_options",
         "//litert/test:matchers",
         "@com_google_googletest//:gtest_main",
     ],

litert/cc/options/compiler_options.cc

@@ -35,12 +35,31 @@ Expected<CompilerOptions> CompilerOptions::Create() {
 
 Expected<CompilerOptions> CompilerOptions::Create(OpaqueOptions& original) {
   const auto id = original.GetIdentifier();
-  if (!id || *id != Identifier()) {
+  if (!id || *id != Discriminator()) {
     return Error(kLiteRtStatusErrorInvalidArgument);
   }
   return CompilerOptions(original.Get(), OwnHandle::kNo);
 }
 
+Expected<void> CompilerOptions::SetPartitionStrategy(
+    LiteRtCompilerOptionsPartitionStrategy partition_strategy) {
+  LiteRtCompilerOptions compiler_options;
+  LITERT_RETURN_IF_ERROR(LiteRtFindCompilerOptions(Get(), &compiler_options));
+  LITERT_RETURN_IF_ERROR(LiteRtSetCompilerOptionsPartitionStrategy(
+      compiler_options, partition_strategy));
+  return {};
+}
+
+Expected<LiteRtCompilerOptionsPartitionStrategy>
+CompilerOptions::GetPartitionStrategy() const {
+  LiteRtCompilerOptions compiler_options;
+  LITERT_RETURN_IF_ERROR(LiteRtFindCompilerOptions(Get(), &compiler_options));
+  LiteRtCompilerOptionsPartitionStrategy partition_strategy;
+  LITERT_RETURN_IF_ERROR(LiteRtGetCompilerOptionsPartitionStrategy(
+      compiler_options, &partition_strategy));
+  return partition_strategy;
+}
+
 Expected<void> CompilerOptions::SetDummyOption(bool dummy_option) {
   LiteRtCompilerOptions compiler_options;
   LITERT_RETURN_IF_ERROR(LiteRtFindCompilerOptions(Get(), &compiler_options));

litert/cc/options/compiler_options.h

@@ -26,13 +26,17 @@ class CompilerOptions : public OpaqueOptions {
  public:
   using OpaqueOptions::OpaqueOptions;
 
-  static absl::string_view Identifier() {
+  static absl::string_view Discriminator() {
     return LiteRtGetCompilerOptionsIdentifier();
   };
 
   static Expected<CompilerOptions> Create();
   static Expected<CompilerOptions> Create(OpaqueOptions& original);
 
+  Expected<void> SetPartitionStrategy(
+      LiteRtCompilerOptionsPartitionStrategy partition_strategy);
+  Expected<LiteRtCompilerOptionsPartitionStrategy> GetPartitionStrategy() const;
+
   Expected<void> SetDummyOption(bool dummy_option);
   Expected<bool> GetDummyOption() const;
 };

litert/cc/options/compiler_options_test.cc

@@ -17,6 +17,7 @@
 #include "litert/cc/options/compiler_options.h"
 
 #include <gtest/gtest.h>
+#include "litert/c/options/litert_compiler_options.h"
 #include "litert/test/matchers.h"
 
 namespace litert {
@@ -29,5 +30,16 @@ TEST(CompilerOptionsTest, CreateSetAndGetDummyOptionWorks) {
   EXPECT_TRUE(options.GetDummyOption());
 }
 
+TEST(CompilerOptionsTest, SetAndGetPartitionStrategyReturnsSetValue) {
+  LITERT_ASSERT_OK_AND_ASSIGN(::litert::CompilerOptions options,
+                              ::litert::CompilerOptions::Create());
+  LITERT_EXPECT_OK(options.SetPartitionStrategy(
+      kLiteRtCompilerOptionsPartitionStrategyWeaklyConnected));
+  LITERT_ASSERT_OK_AND_ASSIGN(auto partition_strategy,
+                              options.GetPartitionStrategy());
+  EXPECT_EQ(partition_strategy,
+            kLiteRtCompilerOptionsPartitionStrategyWeaklyConnected);
+}
+
 }  // namespace
 }  // namespace litert

litert/tools/BUILD

@@ -135,6 +135,7 @@ litert_bin(
         ":apply_plugin_runtimecapi",
         ":outstream",
         "//litert/cc:litert_api_with_dynamic_runtime",
+        "//litert/cc/options:compiler_options",
         "//litert/core:build_stamp",
         "//litert/tools/flags:apply_plugin_flags",
         "//litert/tools/flags:common_flags",
@@ -145,6 +146,7 @@ litert_bin(
         "//litert/tools/flags/vendors:qualcomm_flags_with_dynamic_runtime",
         "@com_google_absl//absl/flags:flag",
         "@com_google_absl//absl/flags:parse",
+        "@com_google_absl//absl/functional:function_ref",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/strings:string_view",
     ],
@@ -174,6 +176,7 @@ litert_bin(
         ":apply_plugin_runtimecapi",
         ":outstream",
         "//litert/cc:litert_api_with_dynamic_runtime",
+        "//litert/cc/options:compiler_options",
         "//litert/core:build_stamp",
         "//litert/tools/flags:apply_plugin_flags",
         "//litert/tools/flags:common_flags",
@@ -184,6 +187,7 @@ litert_bin(
         "//litert/tools/flags/vendors:qualcomm_flags_with_dynamic_runtime",
         "@com_google_absl//absl/flags:flag",
         "@com_google_absl//absl/flags:parse",
+        "@com_google_absl//absl/functional:function_ref",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/strings:string_view",
     ],
@@ -257,6 +261,7 @@ cc_binary(
         ":apply_plugin_runtimecapi",
         ":outstream",
         "//litert/cc:litert_api_with_dynamic_runtime",
+        "//litert/cc/options:compiler_options",
         "//litert/tools/flags:apply_plugin_flags",
         "//litert/tools/flags:common_flags",
         "//litert/tools/flags:flag_types",
@@ -266,6 +271,7 @@ cc_binary(
         "//litert/tools/flags/vendors:qualcomm_flags_with_dynamic_runtime",
         "@com_google_absl//absl/flags:flag",
         "@com_google_absl//absl/flags:parse",
+        "@com_google_absl//absl/functional:function_ref",
         "@com_google_absl//absl/strings:str_format",
     ],
 )

litert/tools/apply_plugin_main.cc

@@ -8,7 +8,7 @@
 //
 // 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 expruns or implied.
+// 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.
 
@@ -24,8 +24,16 @@
 
 #include "absl/flags/flag.h"  // from @com_google_absl
 #include "absl/flags/parse.h"  // from @com_google_absl
+#include "absl/functional/function_ref.h"  // from @com_google_absl
 #include "absl/strings/str_format.h"  // from @com_google_absl
+#include "litert/cc/litert_expected.h"
+#include "litert/cc/litert_opaque_options.h"
 #include "litert/cc/litert_options.h"
+#include "litert/cc/options/compiler_options.h"
+#include "litert/cc/options/litert_google_tensor_options.h"
+#include "litert/cc/options/litert_intel_openvino_options.h"
+#include "litert/cc/options/litert_mediatek_options.h"
+#include "litert/cc/options/litert_qualcomm_options.h"
 #include "litert/tools/apply_plugin.h"
 #include "litert/tools/flags/apply_plugin_flags.h"
 #include "litert/tools/flags/common_flags.h"
@@ -36,6 +44,8 @@
 #include "litert/tools/flags/vendors/qualcomm_flags.h"  // IWYU pragma: keep
 #include "litert/tools/outstream.h"
 
+namespace {
+
 using ::litert::tools::ApplyPlugin;
 using ::litert::tools::ApplyPluginRun;
 using ::litert::tools::IntList;
@@ -80,6 +90,28 @@ ApplyPluginRun::Ptr ParseFlags() {
   return res;
 }
 
+template <typename T>
+bool ParseAndAddOptions(
+    litert::Options& opts, ApplyPluginRun::Ptr& run,
+    absl::FunctionRef<litert::Expected<T>()> options_from_flags) {
+  static_assert(std::is_base_of_v<litert::OpaqueOptions, T>);
+  litert::Expected<T> options = options_from_flags();
+  if (!options) {
+    run->dump_out.Get().get()
+        << "Failed to create " << T::Discriminator()
+        << " options, Error:" << options.Error().Message() << "\n";
+    return false;
+  }
+  if (!opts.AddOpaqueOptions(std::move(*options))) {
+    run->dump_out.Get().get()
+        << "Failed to add " << T::Discriminator() << " options to list\n";
+    return false;
+  }
+  return true;
+}
+
+}  // namespace
+
 int main(int argc, char* argv[]) {
   absl::ParseCommandLine(argc, argv);
 
@@ -110,62 +142,20 @@ int main(int argc, char* argv[]) {
     return 1;
   }
 
-  {
-    auto qnn_opts = litert::qualcomm::QualcommOptionsFromFlags();
-    if (!qnn_opts) {
-      run->dump_out.Get().get() << "Failed to create Qualcomm options\n";
-      return 1;
-    }
-
-    if (!opts->AddOpaqueOptions(std::move(*qnn_opts))) {
-      run->dump_out.Get().get() << "Failed to add Qualcomm options to list\n";
-      return 1;
-    }
-  }
-
-  {
-    auto google_tensor_opts =
-        litert::google_tensor::GoogleTensorOptionsFromFlags();
-    if (!google_tensor_opts) {
-      run->dump_out.Get().get() << "Failed to create Google Tensor options\n";
-      return 1;
-    }
-
-    if (!opts->AddOpaqueOptions(std::move(*google_tensor_opts))) {
-      run->dump_out.Get().get()
-          << "Failed to add google tensor options to list\n";
-      return 1;
-    }
-  }
-
-  {
-    auto intel_openvino_opts =
-        litert::intel_openvino::IntelOpenVinoOptionsFromFlags();
-    if (!intel_openvino_opts) {
-      run->dump_out.Get().get() << "Failed to create Intel OpenVINO options\n";
-      return 1;
-    }
-
-    if (!opts->AddOpaqueOptions(std::move(*intel_openvino_opts))) {
-      run->dump_out.Get().get()
-          << "Failed to add Intel OpenVINO options to list\n";
-      return 1;
-    }
-  }
-
-  {
-    auto mediatek_opts =
-        litert::mediatek::MediatekOptionsFromFlags();
-    if (!mediatek_opts) {
-      run->dump_out.Get().get() << "Failed to create Mediatek options\n";
-      return 1;
-    }
-
-    if (!opts->AddOpaqueOptions(std::move(*mediatek_opts))) {
-      run->dump_out.Get().get()
-          << "Failed to add Mediatek options to list\n";
-      return 1;
-    }
+  bool all_flags_parsed =
+      (ParseAndAddOptions<litert::CompilerOptions>(
+           *opts, run, litert::CompilerOptionsFromFlags) ||
+       ParseAndAddOptions<litert::qualcomm::QualcommOptions>(
+           *opts, run, litert::qualcomm::QualcommOptionsFromFlags) ||
+       ParseAndAddOptions<litert::google_tensor::GoogleTensorOptions>(
+           *opts, run, litert::google_tensor::GoogleTensorOptionsFromFlags) ||
+       ParseAndAddOptions<litert::intel_openvino::IntelOpenVinoOptions>(
+           *opts, run, litert::intel_openvino::IntelOpenVinoOptionsFromFlags) ||
+       ParseAndAddOptions<litert::mediatek::MediatekOptions>(
+           *opts, run, litert::mediatek::MediatekOptionsFromFlags));
+
+  if (!all_flags_parsed) {
+    return 1;
   }
 
   run->options = std::move(*opts);