[Executorch] Make module constructors uniform across

Existing constructors dont compose well such that if you want data loader or data files constructor then you cannot get to override memory allocator.
Fix that.

Differential Revision: [D86120037](https://our.internmc.facebook.com/intern/diff/D86120037/)

[ghstack-poisoned]

Author: kimishpatel

extension/module/module.cpp

@@ -78,11 +78,17 @@ runtime::Result<std::unique_ptr<runtime::DataLoader>> make_data_loader(
 Module::Module(
     const std::string& file_path,
     const LoadMode load_mode,
+    std::unique_ptr<runtime::MemoryAllocator> memory_allocator,
+    std::unique_ptr<runtime::MemoryAllocator> temp_allocator,
     std::unique_ptr<runtime::EventTracer> event_tracer)
     : file_path_(file_path),
       load_mode_(load_mode),
-      memory_allocator_(std::make_unique<MallocMemoryAllocator>()),
-      temp_allocator_(std::make_unique<MallocMemoryAllocator>()),
+      memory_allocator_(
+          memory_allocator ? std::move(memory_allocator)
+                           : std::make_unique<MallocMemoryAllocator>()),
+      temp_allocator_(
+          temp_allocator ? std::move(temp_allocator)
+                         : std::make_unique<MallocMemoryAllocator>()),
       event_tracer_(std::move(event_tracer)) {
   runtime::runtime_init();
 }
@@ -91,11 +97,17 @@ Module::Module(
     const std::string& file_path,
     const std::string& data_map_path,
     const LoadMode load_mode,
+    std::unique_ptr<runtime::MemoryAllocator> memory_allocator,
+    std::unique_ptr<runtime::MemoryAllocator> temp_allocator,
     std::unique_ptr<runtime::EventTracer> event_tracer)
     : file_path_(file_path),
       load_mode_(load_mode),
-      memory_allocator_(std::make_unique<MallocMemoryAllocator>()),
-      temp_allocator_(std::make_unique<MallocMemoryAllocator>()),
+      memory_allocator_(
+          memory_allocator ? std::move(memory_allocator)
+                           : std::make_unique<MallocMemoryAllocator>()),
+      temp_allocator_(
+          temp_allocator ? std::move(temp_allocator)
+                         : std::make_unique<MallocMemoryAllocator>()),
       event_tracer_(std::move(event_tracer)) {
   if (!data_map_path.empty()) {
     data_files_.push_back(data_map_path);
@@ -107,12 +119,18 @@ Module::Module(
     const std::string& file_path,
     std::vector<std::string> data_files,
     const LoadMode load_mode,
+    std::unique_ptr<runtime::MemoryAllocator> memory_allocator,
+    std::unique_ptr<runtime::MemoryAllocator> temp_allocator,
     std::unique_ptr<runtime::EventTracer> event_tracer)
     : file_path_(file_path),
       data_files_(std::move(data_files)),
       load_mode_(load_mode),
-      memory_allocator_(std::make_unique<MallocMemoryAllocator>()),
-      temp_allocator_(std::make_unique<MallocMemoryAllocator>()),
+      memory_allocator_(
+          memory_allocator ? std::move(memory_allocator)
+                           : std::make_unique<MallocMemoryAllocator>()),
+      temp_allocator_(
+          temp_allocator ? std::move(temp_allocator)
+                         : std::make_unique<MallocMemoryAllocator>()),
       event_tracer_(std::move(event_tracer)) {
   runtime::runtime_init();
 }

extension/module/module.h

@@ -63,6 +63,8 @@ class Module {
   explicit Module(
       const std::string& file_path,
       const LoadMode load_mode = LoadMode::File,
+      std::unique_ptr<runtime::MemoryAllocator> memory_allocator = nullptr,
+      std::unique_ptr<runtime::MemoryAllocator> temp_allocator = nullptr,
       std::unique_ptr<runtime::EventTracer> event_tracer = nullptr);
 
   /**
@@ -78,6 +80,8 @@ class Module {
       const std::string& file_path,
       const std::string& data_map_path,
       const LoadMode load_mode = LoadMode::File,
+      std::unique_ptr<runtime::MemoryAllocator> memory_allocator = nullptr,
+      std::unique_ptr<runtime::MemoryAllocator> temp_allocator = nullptr,
       std::unique_ptr<runtime::EventTracer> event_tracer = nullptr);
 
   /**
@@ -93,6 +97,8 @@ class Module {
       const std::string& file_path,
       std::vector<std::string> data_files,
       const LoadMode load_mode = LoadMode::File,
+      std::unique_ptr<runtime::MemoryAllocator> memory_allocator = nullptr,
+      std::unique_ptr<runtime::MemoryAllocator> temp_allocator = nullptr,
       std::unique_ptr<runtime::EventTracer> event_tracer = nullptr);
 
   /**