Run external weight testing with numeric check to avoid numeric drifting

litert/cc/litert_options.cc

@@ -40,9 +40,7 @@
 #include "litert/cc/options/litert_mediatek_options.h"
 #include "litert/cc/options/litert_qualcomm_options.h"
 #include "litert/cc/options/litert_runtime_options.h"
-#if defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
 #include "litert/core/options.h"
-#endif  // defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
 
 namespace litert {
 
@@ -146,7 +144,6 @@ Expected<void> Options::Build() {
 
 Expected<void> Options::SetExternalWeightScopedFile(
     ScopedFile& scoped_file, ScopedWeightSectionMap sections) {
-#if defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
   if (!scoped_file.IsValid()) {
     return Unexpected(kLiteRtStatusErrorInvalidArgument,
                       "Scoped file handle must be valid");
@@ -169,10 +166,6 @@ Expected<void> Options::SetExternalWeightScopedFile(
   options_impl->scoped_weight_source = std::make_unique<ScopedWeightSource>(
       std::move(scoped_file), std::move(sections));
   return {};
-#else
-  return Unexpected(kLiteRtStatusErrorInvalidArgument,
-                    "LiteRT was built without external weight loader support");
-#endif
 }
 
 }  // namespace litert

litert/cc/litert_options_test.cc

@@ -22,12 +22,10 @@
 
 #include <gtest/gtest.h>
 #include "litert/cc/litert_common.h"
-#if defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
 #include "absl/strings/str_cat.h"  // from @com_google_absl
 #include "litert/cc/internal/scoped_file.h"
 #include "litert/cc/internal/scoped_weight_source.h"
 #include "litert/core/options.h"
-#endif  // defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
 #include "litert/test/matchers.h"
 
 namespace litert {
@@ -45,7 +43,6 @@ TEST(OptionsTest, SetHardwareAccelerators) {
       HwAccelerators::kCpu | HwAccelerators::kGpu));
 }
 
-#if defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
 TEST(OptionsTest, SetExternalWeightScopedFileStoresMetadata) {
   LITERT_ASSERT_OK_AND_ASSIGN(auto options, Options::Create());
 
@@ -76,6 +73,5 @@ TEST(OptionsTest, SetExternalWeightScopedFileStoresMetadata) {
   EXPECT_EQ(it->second.length, 8);
   std::remove(path.c_str());
 }
-#endif
 }  // namespace
 }  // namespace litert

litert/core/options.h

@@ -16,17 +16,13 @@
 #define ODML_LITERT_LITERT_CORE_COMPILATION_OPTIONS_H_
 
 #include <cstddef>
-#if defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
 #include <memory>
-#endif  // defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
 #include <string>
 #include <vector>
 
 #include "litert/c/litert_common.h"
 #include "litert/c/litert_custom_op_kernel.h"
-#if defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
 #include "litert/cc/internal/scoped_weight_source.h"
-#endif  // defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
 
 namespace weight_loader {
 class WeightLoader;
@@ -67,13 +63,11 @@ struct LiteRtOptionsT {
   LiteRtOpaqueOptions options = nullptr;
   std::vector<CustomOpOption> custom_op_options;
   std::vector<LiteRtExternalTensorBinding> external_tensor_bindings;
-#if defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
   // Non-owning pointer used to expose the runtime's WeightLoader to delegates.
   weight_loader::WeightLoader* weight_loader = nullptr;
   // Optional scoped weight source when external weights are packed into a
   // single file with group sections.
   std::unique_ptr<litert::ScopedWeightSource> scoped_weight_source;
-#endif  // defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
 };
 
 #endif  // ODML_LITERT_LITERT_CORE_COMPILATION_OPTIONS_H_

litert/runtime/compiled_model.cc

@@ -95,9 +95,7 @@
 #include "litert/runtime/tensor_buffer_requirements.h"
 #include "litert/runtime/tensor_identifier.h"
 #include "litert/runtime/tfl_utils.h"
-#if defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
 #include "weight_loader/external_weight_loader_litert.h"
-#endif  // defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
 #include "tflite/converter/allocation.h"
 #include "tflite/builtin_ops.h"
 #include "tflite/core/api/profiler.h"
@@ -405,19 +403,25 @@ Expected<void> LiteRtCompiledModelT::InitializeRuntime(
       std::make_unique<LiteRtExternalLiteRtBufferContextT>(env, get_tensor_id);
   interp_->SetExternalContext(kTfLiteLiteRtBufferContext,
                               buffer_context_.get());
-#if defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
+
   std::unique_ptr<litert::ScopedWeightSource> scoped_weight_source;
   auto* options_impl =
     reinterpret_cast<LiteRtOptionsT*>(jit_compilation_options);
   if (options_impl != nullptr) {
+    options_impl->weight_loader = nullptr;
     scoped_weight_source = std::move(options_impl->scoped_weight_source);
   }
   weight_loader_ = weight_loader::CreateLiteRtWeightLoader(
       fb_model_->GetModel(), model_directory_, std::move(scoped_weight_source));
-  if (options_impl != nullptr) {
-    options_impl->weight_loader = weight_loader_.get();
-  }
+  has_external_weights_ = false;
   if (weight_loader_) {
+    auto weight_infos = weight_loader_->GetWeightInfo();
+    has_external_weights_ = !weight_infos.empty();
+    if (!has_external_weights_) {
+      LITERT_LOG(LITERT_DEBUG,
+                 "External weight loader: no external weight tensors found");
+      return {};
+    }
     weight_loader::WeightAccessRequest request;
     request.cpu = true;
     // TODO(b/456318365): Handle weight access request to support multiple
@@ -428,8 +432,12 @@ Expected<void> LiteRtCompiledModelT::InitializeRuntime(
       return litert::Unexpected(kLiteRtStatusErrorRuntimeFailure,
                                 std::string(prepare_status.message()));
     }
+
+    if (options_impl != nullptr) {
+      options_impl->weight_loader = weight_loader_.get();
+    }
+
     // Inspect the weight infos to log the available weights for GPU delegates.
-    auto weight_infos = weight_loader_->GetWeightInfo();
     LITERT_LOG(LITERT_DEBUG,
                "External weight loader: %zu weight tensors available for GPU "
                "delegates",
@@ -447,13 +455,11 @@ Expected<void> LiteRtCompiledModelT::InitializeRuntime(
       }
     }
   }
-#endif  // defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
   return {};
 }
 
-#if defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
 Expected<void> LiteRtCompiledModelT::RestoreExternalWeightsForCpu() {
-  if (!weight_loader_) {
+  if (!weight_loader_ || !has_external_weights_) {
     return {};
   }
 
@@ -518,7 +524,6 @@ Expected<void> LiteRtCompiledModelT::RestoreExternalWeightsForCpu() {
   }
   return {};
 }
-#endif  // defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
 
 namespace {
 
@@ -741,14 +746,12 @@ Expected<LiteRtCompiledModelT::Ptr> LiteRtCompiledModelT::Create(
     LITERT_RETURN_IF_ERROR(scoped_modifier.Append(std::move(dispatch_options)));
   }
 
-#if defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
   // Load and restore external weights for CPU execution before delegates are
   // applied. This ensures that XNNPack and other CPU delegates can see the
   // weight data.
   if (hardware_accelerators & kLiteRtHwAcceleratorCpu) {
     LITERT_RETURN_IF_ERROR(compiled_model->RestoreExternalWeightsForCpu());
   }
-#endif  // defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
 
   // Apply accelerators matching the requested hardware support to the
   // model in the order they were registered.

litert/runtime/compiled_model.h

@@ -36,9 +36,7 @@
 #include "litert/cc/litert_buffer_ref.h"
 #include "litert/cc/litert_expected.h"
 #include "litert/cc/litert_macros.h"
-#if defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
 #include "weight_loader/external_weight_loader_litert.h"
-#endif  // defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
 #if !defined(LITERT_DISABLE_NPU)
 #include "litert/core/cache/compilation_cache.h"
 #endif  // !defined(LITERT_DISABLE_NPU)
@@ -403,12 +401,10 @@ class LiteRtCompiledModelT {
   litert::Expected<bool> SignatureNeedsAllocation(
       const tflite::SignatureRunner* runner) const;
 
-#if defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
   // Restores external weights into tensor for CPU execution.
   // This is called before delegates are applied so that XNNPack and other
   // CPU delegates can see the weight data in the tensors.
   litert::Expected<void> RestoreExternalWeightsForCpu();
-#endif  // defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
 
 #if !defined(LITERT_DISABLE_NPU)
   // Applies the plugins to the model and caches the compiled model if
@@ -489,10 +485,11 @@ class LiteRtCompiledModelT {
   // `SetSchedulingInfo`.
   std::string model_debug_feature_id_;
 
-#if defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
   // The loader that manages external weight metadata and bindings.
   std::unique_ptr<weight_loader::WeightLoader> weight_loader_;
-#endif  // defined(LITERT_WITH_EXTERNAL_WEIGHT_LOADER)
+
+  // Indicates whether this model actually uses external weights.
+  bool has_external_weights_ = false;
 
   // File system hints about the originating model location.
   std::optional<std::string> model_directory_;

litert/runtime/compiled_model_test.cc

@@ -1276,6 +1276,33 @@ TEST(CompiledModelTest, BindExternalWeightBuffer) {
   LiteRtDestroyEnvironment(env_ptr);
 }
 
+TEST(CompiledModelTest, NonExternalModelKeepsWeightLoaderNull) {
+  LITERT_ASSERT_OK_AND_ASSIGN(LiteRtEnvironmentT::Ptr env,
+                              LiteRtEnvironmentT::CreateWithOptions({}));
+  LiteRtEnvironmentT* env_ptr = env.release();
+
+  std::string path = testing::GetTestFilePath(kModelFileName);
+  LiteRtModel model;
+  ASSERT_EQ(LiteRtCreateModelFromFile(path.c_str(), &model), kLiteRtStatusOk);
+
+  LiteRtOptions options;
+  ASSERT_EQ(LiteRtCreateOptions(&options), kLiteRtStatusOk);
+  ASSERT_EQ(
+      LiteRtSetOptionsHardwareAccelerators(options, kLiteRtHwAcceleratorCpu),
+      kLiteRtStatusOk);
+  auto* options_impl = reinterpret_cast<LiteRtOptionsT*>(options);
+  ASSERT_NE(options_impl, nullptr);
+
+  LITERT_ASSERT_OK_AND_ASSIGN(
+      LiteRtCompiledModelT::Ptr compiled_model,
+      LiteRtCompiledModelT::Create(env_ptr, model, options));
+  EXPECT_EQ(options_impl->weight_loader, nullptr);
+
+  LiteRtDestroyOptions(options);
+  LiteRtDestroyModel(model);
+  LiteRtDestroyEnvironment(env_ptr);
+}
+
 TEST(CompiledModelTest, GetInterpreter) {
   // Environment setup.
   LITERT_ASSERT_OK_AND_ASSIGN(LiteRtEnvironmentT::Ptr env,