Use unique_ptr and shared_ptr properly for runner components

Summary:
The ownership of these components need some clarification.

* `Module` should be solely owned by `TextDecoderRunner`
* `Tokenizer` should be owned by the `Runner` as well as `TextTokenGenerator`

Differential Revision: D73399600

Author: larryliu0820

examples/models/llama/runner/runner.cpp

@@ -99,7 +99,7 @@ Error Runner::load() {
         "Failed to load %s as a Tiktoken artifact, trying BPE tokenizer",
         tokenizer_path_.c_str());
     tokenizer_.reset();
-    tokenizer_ = std::make_unique<::tokenizers::Llama2cTokenizer>();
+    tokenizer_ = std::make_shared<::tokenizers::Llama2cTokenizer>();
     err = tokenizer_->load(tokenizer_path_);
     ET_CHECK_TK_OK_OR_RETURN_ERROR(
         err,
@@ -143,17 +143,17 @@ Error Runner::load() {
     }
   }
   // @lint-ignore CLANGTIDY facebook-hte-Deprecated
-  text_decoder_runner_ = std::make_unique<llm::TextDecoderRunner>(
-      module_.get(), metadata_.at(kUseKVCache));
+  text_decoder_runner_ = std::make_shared<llm::TextDecoderRunner>(
+      std::move(module_), metadata_.at(kUseKVCache));
   text_prefiller_ = std::make_unique<llm::TextPrefiller>(
-      text_decoder_runner_.get(),
+      text_decoder_runner_,
       metadata_.at(kUseKVCache),
       metadata_.at(kEnableDynamicShape),
       metadata_.at(kMaxSeqLen));
 
   text_token_generator_ = std::make_unique<llm::TextTokenGenerator>(
-      tokenizer_.get(),
-      text_decoder_runner_.get(),
+      tokenizer_,
+      text_decoder_runner_,
       metadata_.at(kUseKVCache),
       std::move(eos_ids),
       &stats_);

examples/models/llama/runner/runner.h

@@ -62,9 +62,9 @@ class ET_EXPERIMENTAL Runner : public executorch::extension::llm::IRunner {
   // model
   std::unique_ptr<::executorch::extension::Module> module_;
   std::string tokenizer_path_;
-  std::unique_ptr<::tokenizers::Tokenizer> tokenizer_;
+  std::shared_ptr<::tokenizers::Tokenizer> tokenizer_;
   std::unordered_map<std::string, int64_t> metadata_;
-  std::unique_ptr<::executorch::extension::llm::TextDecoderRunner>
+  std::shared_ptr<::executorch::extension::llm::TextDecoderRunner>
       text_decoder_runner_;
   std::unique_ptr<::executorch::extension::llm::TextPrefiller> text_prefiller_;
   std::unique_ptr<::executorch::extension::llm::TextTokenGenerator>

examples/models/llama/tokenizer/llama_tiktoken.cpp

@@ -53,8 +53,8 @@ std::unique_ptr<std::vector<std::string>> _get_special_tokens(Version version) {
 
 } // namespace
 
-std::unique_ptr<Tiktoken> get_tiktoken_for_llama(Version version) {
-  return std::make_unique<Tiktoken>(
+std::shared_ptr<Tiktoken> get_tiktoken_for_llama(Version version) {
+  return std::make_shared<Tiktoken>(
       _get_special_tokens(version), kBOSTokenIndex, kEOSTokenIndex);
 }
 

examples/models/llama/tokenizer/llama_tiktoken.h

@@ -17,7 +17,7 @@ enum class Version {
   Multimodal,
 };
 
-std::unique_ptr<::tokenizers::Tiktoken> get_tiktoken_for_llama(
+std::shared_ptr<::tokenizers::Tiktoken> get_tiktoken_for_llama(
     Version version = Version::Default);
 
 std::unique_ptr<std::vector<std::string>> get_multimodal_special_tokens();

extension/llm/runner/text_decoder_runner.cpp

@@ -21,8 +21,8 @@ namespace llm {
 // NOTE: we observed ~2x loading performance increase on iPhone 15
 // and a ~5% improvement on Galaxy S22 by switching to
 // FileDataLoader instead of MmapDataLoader + UseMlockIgnoreErrors.
-TextDecoderRunner::TextDecoderRunner(Module* module, bool use_kv_cache)
-    : module_(module), use_kv_cache_(use_kv_cache) {}
+TextDecoderRunner::TextDecoderRunner(std::unique_ptr<Module> module, bool use_kv_cache)
+    : module_(std::move(module)), use_kv_cache_(use_kv_cache) {}
 
 // This function is functional, meaning it shouldn't modify any state of the
 // input. It should be safe to call multiple times with the same inputs. The

extension/llm/runner/text_decoder_runner.h

@@ -14,15 +14,14 @@
 #include <executorch/extension/module/module.h>
 #include <executorch/extension/tensor/tensor.h>
 #include <executorch/runtime/platform/compiler.h>
-#include <functional>
 
 namespace executorch {
 namespace extension {
 namespace llm {
 
 class ET_EXPERIMENTAL TextDecoderRunner {
  public:
-  TextDecoderRunner(Module* module, bool use_kv_cache);
+  TextDecoderRunner(std::unique_ptr<Module> module, bool use_kv_cache);
 
   virtual ~TextDecoderRunner() = default;
 
@@ -95,7 +94,7 @@ class ET_EXPERIMENTAL TextDecoderRunner {
 
  protected:
   // TODO: use shared_ptr for module
-  Module* module_;
+  std::unique_ptr<Module> module_;
   bool use_kv_cache_;
   bool should_stop_{false};
 };

extension/llm/runner/text_prefiller.cpp

@@ -17,7 +17,7 @@ namespace extension {
 namespace llm {
 
 TextPrefiller::TextPrefiller(
-    TextDecoderRunner* text_decoder_runner,
+    std::shared_ptr<TextDecoderRunner> text_decoder_runner,
     bool use_kv_cache,
     bool enable_parallel_prefill,
     int64_t max_seq_len)

extension/llm/runner/text_prefiller.h

@@ -20,10 +20,12 @@ namespace llm {
 class ET_EXPERIMENTAL TextPrefiller {
  public:
   TextPrefiller(
-      TextDecoderRunner* text_decoder_runner,
+      std::shared_ptr<TextDecoderRunner> text_decoder_runner,
       bool use_kv_cache_,
       bool enable_parallel_prefill,
       int64_t max_seq_len = 128);
+
+  virtual ~TextPrefiller() = default;
   /**
    * Prefill an LLM Module with the given text input.
    * @param prompt_tokens The text prompt tokens to the LLM Module. Encoded by
@@ -32,7 +34,7 @@ class ET_EXPERIMENTAL TextPrefiller {
    * Module.
    * @return The next token of the LLM Module after prefill.
    */
-  ::executorch::runtime::Result<uint64_t> prefill(
+  virtual ::executorch::runtime::Result<uint64_t> prefill(
       std::vector<uint64_t>& prompt_tokens,
       int64_t& start_pos);
 
@@ -47,8 +49,25 @@ class ET_EXPERIMENTAL TextPrefiller {
       std::vector<uint64_t>& prompt_tokens,
       int64_t& start_pos);
 
+  /**
+   * Load the necessary resources for the TextPrefiller.
+   * This method should be called before using the prefill methods.
+   */
+  ::executorch::runtime::Error load() {
+    return text_decoder_runner_->load();
+  }
+
+  /**
+   * Check if the TextPrefiller has been successfully loaded.
+   * @return True if the resources are loaded, false otherwise.
+   */
+  bool inline is_loaded() const {
+    // Implementation to check if resources are loaded
+    return text_decoder_runner_->is_method_loaded();
+  }
+
  private:
-  TextDecoderRunner* text_decoder_runner_;
+  std::shared_ptr<TextDecoderRunner> text_decoder_runner_;
   bool use_kv_cache_;
   bool enable_parallel_prefill_;
   int64_t max_seq_len_;

extension/llm/runner/text_token_generator.h

@@ -21,16 +21,18 @@ namespace llm {
 class ET_EXPERIMENTAL TextTokenGenerator {
  public:
   TextTokenGenerator(
-      ::tokenizers::Tokenizer* tokenizer,
-      TextDecoderRunner* text_decoder_runner,
+      std::shared_ptr<::tokenizers::Tokenizer> tokenizer,
+      std::shared_ptr<TextDecoderRunner> text_decoder_runner,
       bool use_kv_cache,
       std::unique_ptr<std::unordered_set<uint64_t>>&& eos_ids,
-      Stats* stats)
+      std::unique_ptr<Stats> stats)
       : tokenizer_(tokenizer),
         text_decoder_runner_(text_decoder_runner),
         eos_ids_(std::move(eos_ids)),
         use_kv_cache_(use_kv_cache),
-        stats_(stats) {}
+        stats_(std::move(stats)) {}
+
+  virtual ~TextTokenGenerator() = default;
 
   /**
    * Token generation loop.
@@ -135,17 +137,35 @@ class ET_EXPERIMENTAL TextTokenGenerator {
     should_stop_ = true;
   }
 
+  /**
+   * Load the necessary resources for TextTokenGenerator.
+   * This method should be called before using the generate() method.
+   */
+  ::executorch::runtime::Error load() {
+    return text_decoder_runner_->load();
+  }
+
+  /**
+   * Check if the TextTokenGenerator has been successfully loaded.
+   * @return True if the resources are loaded, false otherwise.
+   */
+  bool inline is_loaded() const {
+    // Implementation to check if resources are loaded
+    return tokenizer_->is_initialized() &&
+        text_decoder_runner_->is_method_loaded();
+  }
+
  private:
-  ::tokenizers::Tokenizer* tokenizer_;
-  TextDecoderRunner* text_decoder_runner_;
+  std::shared_ptr<::tokenizers::Tokenizer> tokenizer_;
+  std::shared_ptr<TextDecoderRunner> text_decoder_runner_;
   std::unique_ptr<std::unordered_set<uint64_t>> eos_ids_;
   bool use_kv_cache_;
 
   // state machine
   bool should_stop_ = false;
 
   // stats
-  Stats* stats_;
+  std::unique_ptr<Stats> stats_;
 };
 
 } // namespace llm