Use caching allocator for runner (#15730)

Summary:

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.
ghstack-source-id: 328001114
exported-using-ghexport

Reviewed By: navsud, derekdixu

Differential Revision: D86120038

Author: kimishpatel

CMakeLists.txt

@@ -1095,6 +1095,8 @@ if(EXECUTORCH_BUILD_EXTENSION_TRAINING)
 endif()
 
 if(EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/memory_allocator)
+  list(APPEND _executorch_extensions extension_memory_allocator)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/runner)
   list(APPEND _executorch_extensions extension_llm_runner)
 endif()

extension/llm/runner/CMakeLists.txt

@@ -40,7 +40,7 @@ add_subdirectory(
 )
 
 set(runner_deps executorch_core extension_module extension_tensor
-                extension_llm_sampler tokenizers::tokenizers
+                extension_llm_sampler extension_memory_allocator tokenizers::tokenizers
 )
 
 # depend on arange_utils

extension/llm/runner/llm_runner_helper.cpp

@@ -17,6 +17,7 @@
 #include <executorch/extension/llm/runner/text_llm_runner.h>
 #include <executorch/extension/llm/runner/text_prefiller.h>
 #include <executorch/extension/llm/runner/text_token_generator.h>
+#include <executorch/extension/memory_allocator/cpu_caching_malloc_allocator.h>
 #include <executorch/runtime/core/result.h>
 #include <executorch/runtime/platform/runtime.h>
 #include <pytorch/tokenizers/hf_tokenizer.h>
@@ -223,12 +224,25 @@ std::unique_ptr<TextLLMRunner> create_text_llm_runner(
 
   // Create the Module
   std::unique_ptr<Module> module;
+  uint32_t max_cached_memory_size_bytes_ = 1024 * 1024 * 10; // 10MB
   if (data_files.size() > 0) {
     module = std::make_unique<Module>(
-        model_path, data_files, load_mode, std::move(event_tracer));
+        model_path, data_files, load_mode, std::move(event_tracer),
+        nullptr, // memory allocator
+        std::make_unique<
+            executorch::extension::CPUCachingAllocator>( // temp memory
+                                                         // allocator
+            max_cached_memory_size_bytes_));
   } else {
     module = std::make_unique<Module>(
-        model_path, load_mode, std::move(event_tracer));
+        model_path,
+        load_mode,
+        std::move(event_tracer), // event tracer
+        nullptr, // memory allocator
+        std::make_unique<
+            executorch::extension::CPUCachingAllocator>( // temp memory
+                                                         // allocator
+            max_cached_memory_size_bytes_));
   }
 
   // Get metadata from Module

extension/llm/runner/targets.bzl

@@ -132,6 +132,7 @@ def define_common_targets():
                 ":text_prefiller" + aten_suffix,
                 ":text_token_generator" + aten_suffix,
                 "//executorch/extension/llm/runner/io_manager:io_manager" + aten_suffix,
+                "//executorch/extension/memory_allocator:cpu_caching_allocator",
                 "//pytorch/tokenizers:hf_tokenizer",
                 "//pytorch/tokenizers:llama2c_tokenizer",
                 "//pytorch/tokenizers:sentencepiece",