Add a streaming-aware version of LitertLmLoader

LitertLmStreamingLoader provides utility functions for reading a .litertlm model from a DataStream.

LiteRT-LM-PiperOrigin-RevId: 888363584

Author: ai-edge-bot

runtime/util/BUILD

@@ -87,7 +87,6 @@ cc_test(
         ":test_utils",
         "@com_google_googletest//:gtest_main",
         "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:statusor",
     ],
 )
 
@@ -524,6 +523,43 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "litert_lm_streaming_loader",
+    srcs = ["litert_lm_streaming_loader.cc"],
+    hdrs = ["litert_lm_streaming_loader.h"],
+    deps = [
+        ":data_stream",
+        ":litert_lm_loader",
+        ":litert_status_util",
+        "@com_google_absl//absl/log:absl_check",
+        "@com_google_absl//absl/log:absl_log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@litert//litert/cc:litert_buffer_ref",
+        "//runtime/components:model_resources",
+        "//schema/core:litertlm_header_schema",
+        "//schema/core:litertlm_read",
+    ],
+)
+
+cc_test(
+    name = "litert_lm_streaming_loader_test",
+    srcs = ["litert_lm_streaming_loader_test.cc"],
+    data = ["//runtime/testdata"],
+    deps = [
+        ":file_data_stream",
+        ":litert_lm_streaming_loader",
+        ":test_utils",
+        "@com_google_googletest//:gtest_main",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "//runtime/util:data_stream",
+        "//schema/core:litertlm_header_schema",
+        "//schema/core:litertlm_read",
+    ],
+)
+
 cc_library(
     name = "metadata_util",
     srcs = ["metadata_util.cc"],

runtime/util/data_stream.cc

@@ -27,36 +27,36 @@
 namespace litert::lm {
 
 SubStream::~SubStream() {
-  if (auto parent = parent_.lock()) {
+  if (parent_) {
     // Note: We ignore errors here as we are in a destructor.
-    (void)parent->Discard(offset_, size_);
+    (void)parent_->Discard(offset_, size_);
   }
 }
 
 absl::Status SubStream::ReadAndDiscard(void* buffer, uint64_t offset,
                                        uint64_t size) {
   RETURN_IF_ERROR(CheckBounds(offset, size));
-  if (auto parent = parent_.lock()) {
-    return parent->ReadAndDiscard(buffer, offset_ + offset, size);
+  if (parent_) {
+    return parent_->ReadAndDiscard(buffer, offset_ + offset, size);
   }
-  return absl::FailedPreconditionError("Parent stream is expired");
+  return absl::FailedPreconditionError("Parent stream is null");
 }
 
 absl::Status SubStream::ReadAndPreserve(void* buffer, uint64_t offset,
                                         uint64_t size) {
   RETURN_IF_ERROR(CheckBounds(offset, size));
-  if (auto parent = parent_.lock()) {
-    return parent->ReadAndPreserve(buffer, offset_ + offset, size);
+  if (parent_) {
+    return parent_->ReadAndPreserve(buffer, offset_ + offset, size);
   }
-  return absl::FailedPreconditionError("Parent stream is expired");
+  return absl::FailedPreconditionError("Parent stream is null");
 }
 
 absl::Status SubStream::Discard(uint64_t offset, uint64_t size) {
   RETURN_IF_ERROR(CheckBounds(offset, size));
-  if (auto parent = parent_.lock()) {
-    return parent->Discard(offset_ + offset, size);
+  if (parent_) {
+    return parent_->Discard(offset_ + offset, size);
   }
-  return absl::FailedPreconditionError("Parent stream is expired");
+  return absl::FailedPreconditionError("Parent stream is null");
 }
 
 absl::Status SubStream::CheckBounds(uint64_t offset, uint64_t size) const {
@@ -69,7 +69,7 @@ absl::Status SubStream::CheckBounds(uint64_t offset, uint64_t size) const {
   return absl::OkStatus();
 }
 
-absl::StatusOr<std::shared_ptr<DataStream>> SubStream::OpenSubStream(
+absl::StatusOr<std::unique_ptr<DataStream>> SubStream::OpenSubStream(
     uint64_t offset, uint64_t size) {
   // Check if the requested substream fits within this SubStream's bounds.
   // Note that the parent DataStream::OpenSubStream method doesn't do this for
@@ -80,7 +80,7 @@ absl::StatusOr<std::shared_ptr<DataStream>> SubStream::OpenSubStream(
   return DataStream::OpenSubStream(offset, size);
 }
 
-absl::StatusOr<std::shared_ptr<DataStream>> DataStream::OpenSubStream(
+absl::StatusOr<std::unique_ptr<DataStream>> DataStream::OpenSubStream(
     uint64_t offset, uint64_t size) {
   for (const auto& region : locked_regions_) {
     // Check for overlap: Is [offset, offset + size) overlapping with
@@ -94,7 +94,7 @@ absl::StatusOr<std::shared_ptr<DataStream>> DataStream::OpenSubStream(
     }
   }
   locked_regions_.emplace_back(offset, size);
-  return std::make_shared<SubStream>(shared_from_this(), offset, size);
+  return std::make_unique<SubStream>(this, offset, size);
 }
 
 }  // namespace litert::lm

runtime/util/data_stream.h

@@ -32,7 +32,7 @@ namespace litert::lm {
 // A `SubStream` is a view into a parent `DataStream`. It holds a weak pointer
 // to its parent, so it does not keep the parent alive. The user must ensure
 // that the parent `DataStream` outlives all of its `SubStream`s.
-class DataStream : public std::enable_shared_from_this<DataStream> {
+class DataStream {
  public:
   virtual ~DataStream() = default;
 
@@ -55,7 +55,7 @@ class DataStream : public std::enable_shared_from_this<DataStream> {
   // `offset` + `size`).
   // Note that substreams cannot overlap regions in the parent stream, even if
   // the first substream is destroyed.
-  virtual absl::StatusOr<std::shared_ptr<DataStream>> OpenSubStream(
+  virtual absl::StatusOr<std::unique_ptr<DataStream>> OpenSubStream(
       uint64_t offset, uint64_t size);
 
  private:
@@ -68,7 +68,7 @@ class DataStream : public std::enable_shared_from_this<DataStream> {
 // that the parent `DataStream` outlives all of its `SubStream`s.
 class SubStream : public DataStream {
  public:
-  SubStream(std::shared_ptr<DataStream> parent, uint64_t offset, uint64_t size)
+  SubStream(DataStream* parent, uint64_t offset, uint64_t size)
       : parent_(parent), offset_(offset), size_(size) {}
 
   ~SubStream() override;
@@ -81,11 +81,11 @@ class SubStream : public DataStream {
 
   absl::Status Discard(uint64_t offset, uint64_t size) override;
 
-  absl::StatusOr<std::shared_ptr<DataStream>> OpenSubStream(
+  absl::StatusOr<std::unique_ptr<DataStream>> OpenSubStream(
       uint64_t offset, uint64_t size) override;
 
  private:
-  std::weak_ptr<DataStream> parent_;
+  DataStream* parent_;
   uint64_t offset_;
   uint64_t size_;
 

runtime/util/data_stream_test.cc

@@ -104,20 +104,6 @@ TEST(DataStreamTest, SubStreamOfSubStream) {
   EXPECT_EQ(magic_number, "ITERTLM");
 }
 
-TEST(DataStreamTest, SubStreamReadAfterParentDestroyed) {
-  std::shared_ptr<DataStream> sub_stream;
-  {
-    ASSERT_OK_AND_ASSIGN(auto stream,
-                         FileDataStream::Create(GetTestModelPath()));
-    ASSERT_OK_AND_ASSIGN(sub_stream, stream->OpenSubStream(0, 8));
-  }
-
-  std::vector<char> buffer(1);
-  absl::Status status = sub_stream->ReadAndPreserve(buffer.data(), 0, 1);
-  EXPECT_EQ(status.code(), absl::StatusCode::kFailedPrecondition);
-  EXPECT_THAT(status.message(), HasSubstr("Parent stream is expired"));
-}
-
 TEST(DataStreamTest, SubStreamCannotOverlap) {
   ASSERT_OK_AND_ASSIGN(auto stream, FileDataStream::Create(GetTestModelPath()));
   ASSERT_OK_AND_ASSIGN(auto sub_stream1, stream->OpenSubStream(0, 10));

runtime/util/litert_lm_loader.cc

@@ -59,10 +59,48 @@ absl::StatusOr<std::unique_ptr<MemoryMappedFile>> CreateMemoryMapFromScopedFile(
                                               "whole");
 }
 
-constexpr uint64_t kLitertLmHeaderMaxSize = 16 * 1024;
-
 }  // namespace
 
+absl::StatusOr<std::pair<BufferKey, std::optional<std::string>>>
+ExtractBufferKeyAndBackendConstraint(const schema::SectionObject* section) {
+  auto items = section->items();
+  BufferKey buffer_key(section->data_type());
+  std::optional<std::string> backend_constraint;
+  // Extract the specific model type from the section items KeyValuePairs.
+  if (section->data_type() == schema::AnySectionDataType_TFLiteModel ||
+      section->data_type() == schema::AnySectionDataType_TFLiteWeights) {
+    bool found_model_type = false;
+    std::string model_type;
+    for (size_t j = 0; j < items->size(); ++j) {
+      auto item = items->Get(j);
+      if (item->key() &&
+          absl::AsciiStrToLower(item->key()->str()) == "model_type" &&
+          item->value()) {
+        found_model_type = true;
+        model_type = *(item->value_as_StringValue()->value());
+      }
+      if (item->key() &&
+          absl::AsciiStrToLower(item->key()->str()) == "backend_constraint" &&
+          item->value()) {
+        backend_constraint = *(item->value_as_StringValue()->value());
+      }
+    }
+    if (found_model_type) {
+      ABSL_LOG(INFO) << "model_type: " << model_type;
+      ASSIGN_OR_RETURN(ModelType model_type_enum,
+                       StringToModelType(model_type));
+      buffer_key = BufferKey(section->data_type(), model_type_enum);
+    } else {
+      ABSL_LOG(WARNING) << "model_type not found, use kTfLitePrefillDecode";
+      // For backward compatibility, we will use the default model type if
+      // model_type is not found.
+      buffer_key =
+          BufferKey(section->data_type(), ModelType::kTfLitePrefillDecode);
+    }
+  }
+  return std::make_pair(buffer_key, backend_constraint);
+}
+
 absl::Status LitertLmLoader::MapSection(BufferKey buffer_key,
                                         uint64_t begin_offset,
                                         uint64_t end_offset) {
@@ -149,7 +187,6 @@ absl::Status LitertLmLoader::Initialize() {
   ABSL_LOG(INFO) << "mmap_status is ok";
 
   // Read the header information.
-  schema::LitertlmHeader header;
   absl::Status status =
       ReadHeaderFromLiteRTLM(header_data, header_size, &header_);
   ABSL_LOG(INFO) << "status: " << status;
@@ -164,44 +201,14 @@ absl::Status LitertLmLoader::Initialize() {
   auto sections = header_.metadata->section_metadata()->objects();
   for (size_t i = 0; i < sections->size(); ++i) {
     const schema::SectionObject* section = sections->Get(i);
-    auto items = section->items();
-    BufferKey buffer_key(section->data_type());
-    // Extract the specific model type from the section items KeyValuePairs.
-    if (section->data_type() == schema::AnySectionDataType_TFLiteModel ||
-        section->data_type() == schema::AnySectionDataType_TFLiteWeights) {
-      bool found_model_type = false;
-      std::string model_type;
-      std::string backend_constraint;
-      for (size_t j = 0; j < items->size(); ++j) {
-        auto item = items->Get(j);
-        if (item->key() &&
-            absl::AsciiStrToLower(item->key()->str()) == "model_type" &&
-            item->value()) {
-          found_model_type = true;
-          model_type = *(item->value_as_StringValue()->value());
-        }
-        if (item->key() &&
-            absl::AsciiStrToLower(item->key()->str()) == "backend_constraint" &&
-            item->value()) {
-          backend_constraint = *(item->value_as_StringValue()->value());
-        }
-      }
-      if (found_model_type) {
-        ABSL_LOG(INFO) << "model_type: " << model_type;
-        ASSIGN_OR_RETURN(ModelType model_type_enum,
-                         StringToModelType(model_type));
-        buffer_key = BufferKey(section->data_type(), model_type_enum);
-      } else {
-        ABSL_LOG(WARNING) << "model_type not found, use kTfLitePrefillDecode";
-        // For backward compatibility, we will use the default model type if
-        // model_type is not found.
-        buffer_key =
-            BufferKey(section->data_type(), ModelType::kTfLitePrefillDecode);
-      }
-      if (!backend_constraint.empty()) {
-        section_backend_constraint_[buffer_key] = backend_constraint;
-        ABSL_LOG(INFO) << "section_backend_constraint: " << backend_constraint;
-      }
+    ASSIGN_OR_RETURN(auto key_and_constraint,
+                     ExtractBufferKeyAndBackendConstraint(section));
+    BufferKey buffer_key = key_and_constraint.first;
+    if (key_and_constraint.second.has_value() &&
+        !key_and_constraint.second->empty()) {
+      section_backend_constraint_[buffer_key] = *key_and_constraint.second;
+      ABSL_LOG(INFO) << "section_backend_constraint: "
+                     << *key_and_constraint.second;
     }
     section_locations_[buffer_key] =
         std::make_pair(section->begin_offset(), section->end_offset());

runtime/util/litert_lm_loader.h

@@ -40,6 +40,8 @@
 
 namespace litert::lm {
 
+inline constexpr uint64_t kLitertLmHeaderMaxSize = 16 * 1024;
+
 // Each buffer is keyed by the data type as the major key and the model type
 // as the optional secondary key when the data type is TFLiteModel or
 // TFLiteWeights.
@@ -68,6 +70,10 @@ struct BufferKey {
   }
 };
 
+// Extracts the BufferKey and backend constraint from the section metadata.
+absl::StatusOr<std::pair<BufferKey, std::optional<std::string>>>
+ExtractBufferKeyAndBackendConstraint(const schema::SectionObject* section);
+
 // Hash function for BufferKey
 struct BufferKeyHash {
   size_t operator()(const BufferKey& k) const {