Internal changes and clean up.

LiteRT-LM-PiperOrigin-RevId: 890153576

Author: ai-edge-bot

runtime/components/embedding_lookup/BUILD

@@ -63,6 +63,7 @@ cc_library(
             "@litert//litert/cc:litert_options",
             "@litert//litert/cc:litert_ranked_tensor_type",
             "@litert//litert/cc:litert_tensor_buffer",
+            "@litert//litert/cc/options:litert_qualcomm_options",
         ],
     }),
 )
@@ -127,6 +128,8 @@ cc_library(
             "@litert//litert/cc:litert_model",
             "@litert//litert/cc:litert_options",
             "@litert//litert/cc:litert_tensor_buffer",
+            "@litert//litert/cc/internal:litert_handle",
+            "@litert//litert/cc/options:litert_qualcomm_options",
         ],
     }),
 )
@@ -158,6 +161,7 @@ cc_library(
             "@litert//litert/cc:litert_api_with_dynamic_runtime",
         ],
         "//conditions:default": [
+            "@litert//litert/cc:litert_environment",
             "@litert//litert/cc:litert_model",
             "@litert//litert/cc:litert_tensor_buffer",
         ],

runtime/components/embedding_lookup/embedding_lookup_end_of_multi_modal.cc

@@ -27,6 +27,7 @@
 #include "absl/status/statusor.h"  // from @com_google_absl
 #include "absl/strings/str_cat.h"  // from @com_google_absl
 #include "absl/types/span.h"  // from @com_google_absl
+#include "litert/cc/internal/litert_handle.h"  // from @litert
 #include "litert/cc/litert_common.h"  // from @litert
 #include "litert/cc/litert_compiled_model.h"  // from @litert
 #include "litert/cc/litert_element_type.h"  // from @litert
@@ -36,6 +37,9 @@
 #include "litert/cc/litert_options.h"  // from @litert
 #include "litert/cc/litert_tensor_buffer.h"  // from @litert
 #include "runtime/util/status_macros.h"  //NOLINT
+#if defined(__ANDROID__)
+#include "litert/cc/options/litert_qualcomm_options.h"  // from @litert
+#endif
 
 namespace litert::lm {
 
@@ -149,18 +153,42 @@ absl::Status EndOfMultiModalEmbedding::LookupPrefill(
 
 absl::StatusOr<std::unique_ptr<EndOfMultiModalEmbedding>>
 EndOfMultiModalEmbedding::Create(const litert::Model* absl_nonnull model,
-                                 int special_token) {
-  LITERT_ASSIGN_OR_RETURN(auto env, ::litert::Environment::Create({}));
-  auto handler = std::unique_ptr<EndOfMultiModalEmbedding>(
-      new EndOfMultiModalEmbedding(std::move(env), model, special_token));
+                                 int special_token, litert::Environment* env) {
+  if (env == nullptr) {
+    LITERT_ASSIGN_OR_RETURN(auto local_env, ::litert::Environment::Create({}));
+    auto handler =
+        std::unique_ptr<EndOfMultiModalEmbedding>(new EndOfMultiModalEmbedding(
+            std::move(local_env), model, special_token));
+    RETURN_IF_ERROR(handler->Initialize());
+    return handler;
+  }
+  auto handler =
+      std::unique_ptr<EndOfMultiModalEmbedding>(new EndOfMultiModalEmbedding(
+          ::litert::Environment::WrapCObject(
+              env->GetHolder(),
+              ::litert::OwnHandle::kNo),  // NOLINT(build/include_what_you_use)
+          model, special_token));
   RETURN_IF_ERROR(  // IWYU pragma: keep as is included by status_macros.h
       handler->Initialize());
   return handler;
 }
 
 absl::Status EndOfMultiModalEmbedding::Initialize() {
   LITERT_ASSIGN_OR_RETURN(auto options, Options::Create());
+#if defined(__ANDROID__)
+  options.SetHardwareAccelerators(litert::HwAccelerators::kNpu |
+                                  litert::HwAccelerators::kCpu);
+#else
   options.SetHardwareAccelerators(litert::HwAccelerators::kCpu);
+#endif
+#if defined(__ANDROID__)
+  LITERT_ASSIGN_OR_RETURN(::litert::qualcomm::QualcommOptions & qnn_opts,
+                          options.GetQualcommOptions());
+  qnn_opts.SetLogLevel(::litert::qualcomm::QualcommOptions::LogLevel::kOff);
+  qnn_opts.SetHtpPerformanceMode(
+      ::litert::qualcomm::QualcommOptions::HtpPerformanceMode::
+          kSustainedHighPerformance);
+#endif
 
   LITERT_ASSIGN_OR_RETURN(
       litert::CompiledModel compiled_model,

runtime/components/embedding_lookup/embedding_lookup_end_of_multi_modal.h

@@ -29,6 +29,7 @@
 #include "litert/cc/litert_environment.h"  // from @litert
 #include "litert/cc/litert_layout.h"  // from @litert
 #include "litert/cc/litert_model.h"  // from @litert
+#include "litert/cc/litert_options.h"  // from @litert
 #include "litert/cc/litert_tensor_buffer.h"  // from @litert
 #include "runtime/components/embedding_lookup/embedding_lookup.h"
 
@@ -45,7 +46,8 @@ class EndOfMultiModalEmbedding : public EmbeddingLookup {
   // multi-modal embedding. If the special token is not found in the tokens,
   // the end of multi-modal embedding will not be inserted.
   static absl::StatusOr<std::unique_ptr<EndOfMultiModalEmbedding>> Create(
-      const litert::Model* absl_nonnull model, int special_token);
+      const litert::Model* absl_nonnull model, int special_token,
+      litert::Environment* env = nullptr);
 
   // Multimodal embeddings are not supported during decode.
   absl::Status LookupDecode(int token,

runtime/components/embedding_lookup/embedding_lookup_manager.cc

@@ -29,6 +29,7 @@
 #include "absl/status/statusor.h"  // from @com_google_absl
 #include "absl/strings/string_view.h"  // from @com_google_absl
 #include "absl/types/span.h"  // from @com_google_absl
+#include "litert/cc/litert_environment.h"  // from @litert
 #include "litert/cc/litert_macros.h"  // from @litert
 #include "litert/cc/litert_model.h"  // from @litert
 #include "litert/cc/litert_tensor_buffer.h"  // from @litert
@@ -45,22 +46,24 @@ EmbeddingLookupManager::Create(
     const litert::Model* absl_nonnull text_embedding_model,
     absl::flat_hash_map<int, const litert::Model*>&
         end_of_multi_modal_embedding_models,
-    bool fully_supports_multi_modal, std::optional<std::string> signature_key) {
+    bool fully_supports_multi_modal, std::optional<std::string> signature_key,
+    litert::Environment* env) {
   auto embedding_lookup_manager = std::make_unique<EmbeddingLookupManager>();
   RETURN_IF_ERROR(embedding_lookup_manager->Initialize(
       text_embedding_model, end_of_multi_modal_embedding_models,
-      fully_supports_multi_modal, signature_key));
+      fully_supports_multi_modal, signature_key, env));
   return std::move(embedding_lookup_manager);
 }
 
 absl::StatusOr<std::unique_ptr<EmbeddingLookupManager>>
 EmbeddingLookupManager::Create(
     const litert::Model* absl_nonnull text_embedding_model,
-    bool fully_supports_multi_modal, std::optional<std::string> signature_key) {
+    bool fully_supports_multi_modal, std::optional<std::string> signature_key,
+    litert::Environment* env) {
   absl::flat_hash_map<int, const litert::Model*>
       end_of_multi_modal_embedding_models;
   return Create(text_embedding_model, end_of_multi_modal_embedding_models,
-                fully_supports_multi_modal, signature_key);
+                fully_supports_multi_modal, signature_key, env);
 }
 
 absl::Status EmbeddingLookupManager::UpdateMultiModalEmbeddings(
@@ -239,7 +242,8 @@ absl::Status EmbeddingLookupManager::Initialize(
     const litert::Model* absl_nonnull text_embedding_model,
     absl::flat_hash_map<int, const litert::Model*>&
         end_of_multi_modal_embedding_models,
-    bool fully_supports_multi_modal, std::optional<std::string> signature_key) {
+    bool fully_supports_multi_modal, std::optional<std::string> signature_key,
+    litert::Environment* env) {
   if (!fully_supports_multi_modal &&
       !end_of_multi_modal_embedding_models.empty()) {
     return absl::InvalidArgumentError(
@@ -249,12 +253,12 @@ absl::Status EmbeddingLookupManager::Initialize(
   fully_supports_multi_modal_ = fully_supports_multi_modal;
   ASSIGN_OR_RETURN(text_embedding_lookup_,
                    EmbeddingLookupText::Create(std::move(text_embedding_model),
-                                               signature_key));
+                                               signature_key, env));
   for (const auto& [special_token, embedding_model] :
        end_of_multi_modal_embedding_models) {
     ASSIGN_OR_RETURN(auto end_of_multi_modal_embedding_lookup,
                      EndOfMultiModalEmbedding::Create(
-                         std::move(embedding_model), special_token));
+                         std::move(embedding_model), special_token, env));
     end_of_multi_modal_embedding_lookups_.push_back(
         std::move(end_of_multi_modal_embedding_lookup));
   }

runtime/components/embedding_lookup/embedding_lookup_manager.h

@@ -56,12 +56,14 @@ class EmbeddingLookupManager {
       absl::flat_hash_map<int, const litert::Model*>&
           end_of_multi_modal_embedding_models,
       bool fully_supports_multi_modal = true,
-      std::optional<std::string> signature_key = std::nullopt);
+      std::optional<std::string> signature_key = std::nullopt,
+      litert::Environment* env = nullptr);
 
   static absl::StatusOr<std::unique_ptr<EmbeddingLookupManager>> Create(
       const litert::Model* absl_nonnull text_embedding_model,
       bool fully_supports_multi_modal = true,
-      std::optional<std::string> signature_key = std::nullopt);
+      std::optional<std::string> signature_key = std::nullopt,
+      litert::Environment* env = nullptr);
 
   // Updates the multimodal embeddings for the given ExecutorInputs.
   // Intended to be called at the beginning of the prefill pass.
@@ -118,8 +120,8 @@ class EmbeddingLookupManager {
       const litert::Model* absl_nonnull text_embedding_model,
       absl::flat_hash_map<int, const litert::Model*>&
           end_of_multi_modal_embedding_models,
-      bool fully_supports_multi_modal,
-      std::optional<std::string> signature_key);
+      bool fully_supports_multi_modal, std::optional<std::string> signature_key,
+      litert::Environment* env = nullptr);
 
   std::unique_ptr<EmbeddingLookupText> text_embedding_lookup_;
   std::vector<std::unique_ptr<EmbeddingLookupMultiModal>>

runtime/components/embedding_lookup/embedding_lookup_text.cc

@@ -40,6 +40,9 @@
 #include "litert/cc/litert_options.h"  // from @litert
 #include "litert/cc/litert_tensor_buffer.h"  // from @litert
 #include "runtime/util/status_macros.h"  //NOLINT
+#if defined(__ANDROID__)
+#include "litert/cc/options/litert_qualcomm_options.h"  // from @litert
+#endif
 
 namespace litert::lm {
 
@@ -243,17 +246,39 @@ absl::Status EmbeddingLookupText::LookupPrefill(absl::Span<const int> tokens,
 
 absl::StatusOr<std::unique_ptr<EmbeddingLookupText>>
 EmbeddingLookupText::Create(const litert::Model* absl_nonnull model,
-                            std::optional<std::string> signature_key) {
-  LITERT_ASSIGN_OR_RETURN(auto env, ::litert::Environment::Create({}));
-  auto handler = std::unique_ptr<EmbeddingLookupText>(
-      new EmbeddingLookupText(std::move(env), model, signature_key));
+                            std::optional<std::string> signature_key,
+                            litert::Environment* env) {
+  if (env == nullptr) {
+    LITERT_ASSIGN_OR_RETURN(auto local_env, ::litert::Environment::Create({}));
+    auto handler = std::unique_ptr<EmbeddingLookupText>(
+        new EmbeddingLookupText(std::move(local_env), model, signature_key));
+    RETURN_IF_ERROR(handler->Initialize());
+    return handler;
+  }
+  auto handler = std::unique_ptr<EmbeddingLookupText>(new EmbeddingLookupText(
+      ::litert::Environment::WrapCObject(env->GetHolder(),
+                                         ::litert::OwnHandle::kNo),  // NOLINT
+      model, signature_key));
   RETURN_IF_ERROR(handler->Initialize());
   return handler;
 }
 
 absl::Status EmbeddingLookupText::Initialize() {
   LITERT_ASSIGN_OR_RETURN(auto options, Options::Create());
+#if defined(__ANDROID__)
+  options.SetHardwareAccelerators(litert::HwAccelerators::kNpu |
+                                  litert::HwAccelerators::kCpu);
+#else
   options.SetHardwareAccelerators(litert::HwAccelerators::kCpu);
+#endif
+#if defined(__ANDROID__)
+  LITERT_ASSIGN_OR_RETURN(::litert::qualcomm::QualcommOptions & qnn_opts,
+                          options.GetQualcommOptions());
+  qnn_opts.SetLogLevel(::litert::qualcomm::QualcommOptions::LogLevel::kOff);
+  qnn_opts.SetHtpPerformanceMode(
+      ::litert::qualcomm::QualcommOptions::HtpPerformanceMode::
+          kSustainedHighPerformance);
+#endif
 
   LITERT_ASSIGN_OR_RETURN(compiled_model_, litert::CompiledModel::Create(
                                                env_, model_.Get(), options));
@@ -323,6 +348,11 @@ absl::Status EmbeddingLookupText::Initialize() {
     floats_per_token_output_ *= output_buffer_layout.Dimensions()[i];
   }
 
+  ABSL_LOG(INFO) << "EmbeddingLookupText initialized: "
+                 << "signature=" << signature_key_.value_or("default")
+                 << ", rank=" << output_buffer_layout.Rank()
+                 << ", floats_per_token=" << floats_per_token_output_;
+
   // Initialize the default embedding vector to be the embedding of token 0.
   default_embedding_vector_.resize(floats_per_token_output_);
   RETURN_IF_ERROR(LookupInternal(

runtime/components/embedding_lookup/embedding_lookup_text.h

@@ -32,6 +32,7 @@
 #include "litert/cc/litert_compiled_model.h"  // from @litert
 #include "litert/cc/litert_environment.h"  // from @litert
 #include "litert/cc/litert_model.h"  // from @litert
+#include "litert/cc/litert_options.h"  // from @litert
 #include "litert/cc/litert_ranked_tensor_type.h"  // from @litert
 #include "litert/cc/litert_tensor_buffer.h"  // from @litert
 #include "runtime/components/embedding_lookup/embedding_lookup.h"
@@ -55,7 +56,8 @@ class EmbeddingLookupText : public EmbeddingLookup {
   // signature_key is not provided, the first signature will be used by default.
   static absl::StatusOr<std::unique_ptr<EmbeddingLookupText>> Create(
       const litert::Model* absl_nonnull model,
-      std::optional<std::string> signature_key = std::nullopt);
+      std::optional<std::string> signature_key = std::nullopt,
+      litert::Environment* env = nullptr);
 
   // For a given token, looks up the embedding and stores it in the
   // provided vector. The caller is responsible for ensuring that the vector is

runtime/engine/litert_lm_advanced_main.cc

@@ -30,13 +30,10 @@
 #include <string>
 #include <vector>
 
-#include "absl/base/log_severity.h"  // from @com_google_absl
 #include "absl/flags/flag.h"  // from @com_google_absl
-#include "absl/flags/marshalling.h"  // from @com_google_absl
 #include "absl/flags/parse.h"  // from @com_google_absl
 #include "absl/log/absl_check.h"  // from @com_google_absl
 #include "absl/log/absl_log.h"  // from @com_google_absl
-#include "absl/log/globals.h"  // from @com_google_absl
 #include "absl/status/status.h"  // from @com_google_absl
 #include "absl/status/statusor.h"  // from @com_google_absl
 #include "absl/strings/numbers.h"  // from @com_google_absl
@@ -257,6 +254,7 @@ absl::Status MainHelper(int argc, char** argv) {
   settings.enable_speculative_decoding =
       absl::GetFlag(FLAGS_enable_speculative_decoding);
 
+
   // Adjust max_num_tokens and prefill_batch_size if not set on benchmark mode.
   if (settings.benchmark && settings.benchmark_prefill_tokens > 0) {
     if (settings.max_num_tokens == 0 && settings.benchmark_decode_tokens > 0) {

runtime/engine/litert_lm_lib.cc

@@ -45,6 +45,7 @@
 #include "absl/strings/str_format.h"  // from @com_google_absl
 #include "absl/strings/string_view.h"  // from @com_google_absl
 #include "absl/time/time.h"  // from @com_google_absl
+#include "runtime/util/status_macros.h"  // NOLINT
 #include "nlohmann/json.hpp"  // from @nlohmann_json
 #include "litert/cc/internal/scoped_file.h"  // from @litert
 #include "runtime/components/constrained_decoding/constraint.h"
@@ -61,7 +62,6 @@
 #include "runtime/executor/llm_executor_settings.h"
 #include "runtime/proto/sampler_params.pb.h"
 #include "runtime/util/scoped_file.h"
-#include "runtime/util/status_macros.h"  // IWYU pragma: keep
 #include "re2/re2.h"  // from @com_googlesource_code_re2
 #include "tflite/profiling/memory_info.h"  // from @litert
 #include "tflite/profiling/memory_usage_monitor.h"  // from @litert