Update Vision adapter GPU options to be in sync with the encoder

Package.swift

@@ -0,0 +1,78 @@
+// swift-tools-version: 5.9
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// https://www.apache.org/licenses/LICENSE-2.0
+//
+// 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 express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import PackageDescription
+
+let package = Package(
+  name: "LiteRTLM",
+  platforms: [
+    .iOS(.v15),
+    .macOS(.v12),
+  ],
+  products: [
+    .library(
+      name: "LiteRTLM",
+      targets: ["LiteRTLM"]
+    )
+  ],
+  targets: [
+    // 1. The Prebuilt Binary Target
+    .binaryTarget(
+      name: "CLiteRTLM",
+      path: "prebuilt/CLiteRTLM.xcframework"
+    ),
+    // 2. The Swift Wrapper Target
+    .target(
+      name: "LiteRTLM",
+      dependencies: ["CLiteRTLM"],
+      path: "swift",
+      exclude: [
+        "CapabilitiesTests.swift",
+        "EngineTests.swift",
+        "ConversationTests.swift",
+        "ToolTests.swift",
+        "BUILD",
+      ],
+      linkerSettings: [
+        .unsafeFlags(["-Xlinker", "-all_load"])
+      ]
+    ),
+    // Separate test targets for each file to avoid naming conflicts:
+    .testTarget(
+      name: "CapabilitiesTests",
+      dependencies: ["LiteRTLM"],
+      path: "swift",
+      sources: ["CapabilitiesTests.swift"]
+    ),
+    .testTarget(
+      name: "ConversationTests",
+      dependencies: ["LiteRTLM"],
+      path: "swift",
+      sources: ["ConversationTests.swift"]
+    ),
+    .testTarget(
+      name: "ToolTests",
+      dependencies: ["LiteRTLM"],
+      path: "swift",
+      sources: ["ToolTests.swift"]
+    ),
+    .testTarget(
+      name: "EngineTests",
+      dependencies: ["LiteRTLM"],
+      path: "swift",
+      sources: ["EngineTests.swift"]
+    ),
+  ]
+)

c/BUILD

@@ -14,6 +14,8 @@
 
 # [Google-internal load of `cc_library`]
 # [Google-internal load of `cc_test`]
+# [Google-internal load of `objc_library`]
+load("@build_bazel_rules_apple//apple:apple_xcframework.bzl", "apple_xcframework")
 
 package(
     default_hdrs_check = "strict",
@@ -71,6 +73,7 @@ cc_library(
     deps = ENGINE_COMMON_DEPS + [
         "//runtime/core:engine_impl",
     ],
+    alwayslink = True,
 )
 
 cc_library(
@@ -115,3 +118,37 @@ cc_test(
         "//runtime/executor:llm_executor_settings",
     ],
 )
+
+objc_library(
+    name = "CLiteRTLM_objc",
+    deps = [
+        ":engine",
+        "//schema/capabilities:capabilities_c",
+    ],
+)
+
+apple_xcframework(
+    name = "CLiteRTLM",
+    bundle_id = "com.google.odml.litertlm.CLiteRTLM",
+    features = ["exported_symbols"],
+    infoplists = ["Info.plist"],
+    ios = {
+        "simulator": [
+            "arm64",
+            "x86_64",
+        ],
+        "device": [
+            "arm64",
+        ],
+    },
+    minimum_os_versions = {
+        "ios": "15.0",
+    },
+    public_hdrs = [
+        "engine.h",
+        "//schema/capabilities:capabilities_c.h",
+    ],
+    deps = [
+        ":CLiteRTLM_objc",
+    ],
+)

c/Info.plist

@@ -0,0 +1,22 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>CFBundleDevelopmentRegion</key>
+	<string>en</string>
+	<key>CFBundleExecutable</key>
+	<string>CLiteRTLM</string>
+	<key>CFBundleIdentifier</key>
+	<string>com.google.odml.litertlm.CLiteRTLM</string>
+	<key>CFBundleInfoDictionaryVersion</key>
+	<string>6.0</string>
+	<key>CFBundleName</key>
+	<string>CLiteRTLM</string>
+	<key>CFBundlePackageType</key>
+	<string>FMWK</string>
+	<key>CFBundleShortVersionString</key>
+	<string>1.0</string>
+	<key>CFBundleVersion</key>
+	<string>1</string>
+</dict>
+</plist>

runtime/components/embedding_lookup/BUILD

@@ -106,6 +106,9 @@ cc_library(
     name = "embedding_lookup_end_of_multi_modal",
     srcs = ["embedding_lookup_end_of_multi_modal.cc"],
     hdrs = ["embedding_lookup_end_of_multi_modal.h"],
+    visibility = [
+        "//:__subpackages__",
+    ],
     deps = [
         ":embedding_lookup",
         "@com_google_absl//absl/base:nullability",

runtime/executor/llm_litert_npu_compiled_model_executor.cc

@@ -2740,8 +2740,9 @@ absl::StatusOr<int> LlmLiteRtNpuCompiledModelExecutor::GetVocabSize() {
       llm_inference_context_
           .decode_output_buffers[LlmSignatures::kDecodeLogitsOutput]
           .TensorType());
-  RET_CHECK_EQ(logits_tensor_type.Layout().Dimensions().size(), 2);
-  return logits_tensor_type.Layout().Dimensions()[1];
+  const auto rank = logits_tensor_type.Layout().Dimensions().size();
+  RET_CHECK(rank == 2 || rank == 3) << "Logits must be a 2D or 3D tensor.";
+  return logits_tensor_type.Layout().Dimensions()[rank - 1];
 }
 
 const LlmLiteRtNpuCompiledModelExecutor::LatencyStats&

runtime/executor/llm_litert_npu_compiled_model_executor_utils.cc

@@ -972,7 +972,13 @@ absl::Status DequantizeLogits(const ::litert::TensorBuffer& src,
       dst_ptr[i] = scale * (static_cast<float>(src_ptr[i]) -
                             static_cast<float>(zero_point));
     }
-  } else {
+  } else if (src_elem_type == ::litert::ElementType::Float32) {
+      // This is for dealing with unquantized float 32 logits.
+      const float* src_ptr = static_cast<const float*>(src_raw_ptr);
+      for (size_t i = 0; i < num_elements; ++i) {
+        dst_ptr[i] = src_ptr[i];
+      }
+    } else {
     return absl::InvalidArgumentError(absl::StrCat(
         "Unsupported source type for dequantization: ", (int)src_elem_type));
   }

runtime/executor/vision_litert_compiled_model_executor.cc

@@ -313,14 +313,9 @@ absl::Status VisionLiteRtCompiledModelExecutor::VisionAdapter::Initialize() {
       break;
     }
     case Backend::GPU: {
-      // TODO: b/403132820 - Add accelerator compilation options for ML_DRIFT.
       LITERT_ASSIGN_OR_RETURN(auto& gpu_options, options.GetGpuOptions());
-      gpu_options.EnableConstantTensorSharing(true);
-      gpu_options.EnableAllowSrcQuantizedFcConvOps(true);
-
-      gpu_options.SetPrecision(GpuOptions::Precision::kFp16);
-      gpu_options.SetPreferTextureWeights(true);
-      options.SetHardwareAccelerators(litert::HwAccelerators::kGpu);
+      LITERT_RETURN_IF_ERROR(
+          SetGpuOptions(vision_executor_settings_, gpu_options));
       break;
     }
 #if !defined(LITERT_DISABLE_NPU)

runtime/util/litert_util.cc

@@ -15,15 +15,14 @@
 #include "runtime/util/litert_util.h"
 
 #include <cstdint>
-#include <filesystem>  // NOLINT: Required for path manipulation.
 #include <string>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
 #include "absl/base/no_destructor.h"  // from @com_google_absl
 #include "absl/log/absl_log.h"  // from @com_google_absl
 #include "absl/status/statusor.h"  // from @com_google_absl
-#include "absl/strings/string_view.h"  // from @com_google_absl
 #include "litert/cc/litert_environment.h"  // from @litert
 #include "litert/cc/litert_environment_options.h"  // from @litert
 #include "litert/cc/litert_macros.h"  // from @litert
@@ -41,77 +40,85 @@ absl::StatusOr<Environment&> GetEnvironment(EngineSettings& engine_settings,
   // called. Since env is used multiple times, it should also be static.
   static absl::NoDestructor<MagicNumberConfigsHelper> helper;
 
-  static absl::NoDestructor<absl::StatusOr<Environment>> kEnvironment(
-      [&]() -> absl::StatusOr<Environment> {
-        const auto& main_executor_settings =
-            engine_settings.GetMainExecutorSettings();
-        std::vector<EnvironmentOptions::Option> env_options;
+  const auto& main_executor_settings =
+      engine_settings.GetMainExecutorSettings();
+  Backend backend = main_executor_settings.GetBackend();
 
-        if (model_resources != nullptr &&
-            (main_executor_settings.GetBackend() == Backend::CPU ||
-             main_executor_settings.GetBackend() == Backend::GPU)) {
-          if (!main_executor_settings.GetAdvancedSettings() ||
-              main_executor_settings.GetAdvancedSettings()
-                  ->configure_magic_numbers) {
-            env_options = helper->GetLiteRtEnvOptions(*model_resources,
-                                                      main_executor_settings);
-          }
+  static absl::NoDestructor<
+      std::unordered_map<Backend, absl::StatusOr<Environment>>>
+      kEnvironments;
+
+  auto it = kEnvironments->find(backend);
+  if (it == kEnvironments->end()) {
+    auto env_res = [&]() -> absl::StatusOr<Environment> {
+      std::vector<EnvironmentOptions::Option> env_options;
+
+      if (model_resources != nullptr &&
+          (backend == Backend::CPU || backend == Backend::GPU)) {
+        if (!main_executor_settings.GetAdvancedSettings() ||
+            main_executor_settings.GetAdvancedSettings()
+                ->configure_magic_numbers) {
+          env_options = helper->GetLiteRtEnvOptions(*model_resources,
+                                                    main_executor_settings);
         }
+      }
 
 #if !defined(LITERT_DISABLE_NPU)
-        if (!main_executor_settings.GetLitertDispatchLibDir().empty()) {
-          // If the dispatch library directory is provided, use it.
-          env_options.push_back(::litert::EnvironmentOptions::Option{
-              ::litert::EnvironmentOptions::Tag::kDispatchLibraryDir,
-              main_executor_settings.GetLitertDispatchLibDir()});
-          ABSL_LOG(INFO) << "Setting dispatch library path from "
-                            "main_executor_settings: "
-                           << main_executor_settings.GetLitertDispatchLibDir();
-        } else {
+      if (!main_executor_settings.GetLitertDispatchLibDir().empty()) {
+        // If the dispatch library directory is provided, use it.
+        env_options.push_back(::litert::EnvironmentOptions::Option{
+            ::litert::EnvironmentOptions::Tag::kDispatchLibraryDir,
+            main_executor_settings.GetLitertDispatchLibDir()});
+        ABSL_LOG(INFO) << "Setting dispatch library path from "
+                          "main_executor_settings: "
+                       << main_executor_settings.GetLitertDispatchLibDir();
+      } else {
 #if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
-          // Otherwise, use the directory of the model file.
-          std::string model_path(
-              main_executor_settings.GetModelAssets().GetPath().value_or(""));
-          std::filesystem::path path(model_path);
-          // Note: Existence check for path was here, but it's better to check
-          // before calling this function if needed.
-          std::string dispatch_library_path = path.parent_path().string();
-          // In WASM, the parent path is often just "/" which is usually not
-          // what we want for dispatch libraries.
+        // Otherwise, use the directory of the model file.
+        std::string model_path(
+            main_executor_settings.GetModelAssets().GetPath().value_or(""));
+        std::filesystem::path path(model_path);
+        // Note: Existence check for path was here, but it's better to check
+        // before calling this function if needed.
+        std::string dispatch_library_path = path.parent_path().string();
+        // In WASM, the parent path is often just "/" which is usually not
+        // what we want for dispatch libraries.
 #ifdef __EMSCRIPTEN__
-          bool should_set_path =
-              !dispatch_library_path.empty() && dispatch_library_path != "/";
+        bool should_set_path =
+            !dispatch_library_path.empty() && dispatch_library_path != "/";
 #else
-          bool should_set_path = !dispatch_library_path.empty();
+        bool should_set_path = !dispatch_library_path.empty();
 #endif
-          if (should_set_path) {
-            ABSL_LOG(INFO) << "Setting dispatch library path: "
-                           << dispatch_library_path;
-            env_options.push_back(::litert::EnvironmentOptions::Option{
-                ::litert::EnvironmentOptions::Tag::kDispatchLibraryDir,
-                absl::string_view(dispatch_library_path)});
-          } else {
-            ABSL_LOG(INFO) << "No dispatch library path provided.";
-          }
-#endif  // defined(__ANDROID__) || defined(__EMSCRIPTEN__)
+        if (should_set_path) {
+          ABSL_LOG(INFO) << "Setting dispatch library path: "
+                         << dispatch_library_path;
+          env_options.push_back(::litert::EnvironmentOptions::Option{
+              ::litert::EnvironmentOptions::Tag::kDispatchLibraryDir,
+              absl::string_view(dispatch_library_path)});
+        } else {
+          ABSL_LOG(INFO) << "No dispatch library path provided.";
         }
+#endif  // defined(__ANDROID__) || defined(__EMSCRIPTEN__)
+      }
 #endif  // defined(LITERT_DISABLE_NPU)
 
-        if (auto severity = GetMinLogSeverity()) {
-          env_options.push_back(::litert::EnvironmentOptions::Option{
-              ::litert::EnvironmentOptions::Tag::kMinLoggerSeverity,
-              static_cast<int64_t>(ToLiteRtLogSeverityInt8(*severity))});
-        }
+      if (auto severity = GetMinLogSeverity()) {
+        env_options.push_back(::litert::EnvironmentOptions::Option{
+            ::litert::EnvironmentOptions::Tag::kMinLoggerSeverity,
+            static_cast<int64_t>(ToLiteRtLogSeverityInt8(*severity))});
+      }
 
-        LITERT_ASSIGN_OR_RETURN(
-            auto env, Environment::Create(EnvironmentOptions(env_options)));
-        return std::move(env);
-      }());
+      LITERT_ASSIGN_OR_RETURN(
+          auto env, Environment::Create(EnvironmentOptions(env_options)));
+      return std::move(env);
+    }();
+    it = kEnvironments->emplace(backend, std::move(env_res)).first;
+  }
 
-  if (!kEnvironment->ok()) {
-    return kEnvironment->status();
+  if (!it->second.ok()) {
+    return it->second.status();
   }
-  return **kEnvironment;
+  return *it->second;
 }
 
 }  // namespace litert::lm

schema/capabilities/BUILD

@@ -24,6 +24,8 @@ package(
 
 licenses(["notice"])
 
+exports_files(["capabilities_c.h"])
+
 cc_library(
     name = "capabilities",
     hdrs = ["capabilities.h"],
@@ -37,6 +39,7 @@ cc_library(
     deps = [
         ":capabilities",
     ],
+    alwayslink = True,
 )
 
 cc_library(