From 799c34bf674e7964d6c75d61e041e3f71c981155 Mon Sep 17 00:00:00 2001
From: Kimish Patel <kimishpatel@fb.com>
Date: Mon, 10 Nov 2025 09:18:14 -0800
Subject: [PATCH] [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]
---
 extension/module/module.cpp | 30 ++++++++++++++++++++++++------
 extension/module/module.h   |  6 ++++++
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/extension/module/module.cpp b/extension/module/module.cpp
index 35228d06729..b95f40b9f40 100644
--- a/extension/module/module.cpp
+++ b/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();
 }
diff --git a/extension/module/module.h b/extension/module/module.h
index e523f163317..10fc366cb04 100644
--- a/extension/module/module.h
+++ b/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);
 
   /**