Update base for Update on "[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/memory_allocator/cpu_caching_malloc_allocator.h

@@ -59,8 +59,9 @@ class CPUCachingAllocator : public executorch::runtime::MemoryAllocator {
   // New invariants must be written.
   FlatHashMap<size_t, SmallVector<void*, 16>> available_map_;
   FlatHashMap<void*, size_t> allocation_map_;
-  // Since allocation_map, which is a global instance, is mutated/read via
-  // all public APIs we need a global mutex.
+  // Since allocation_map_ and other member variables are mutated/read via
+  // all public APIs, we need a mutex to protect concurrent access to these
+  // instance members.
   std::mutex mutex_;
   size_t max_size_;
   size_t current_size_;

extension/memory_allocator/test/cpu_caching_malloc_allocator_test.cpp

@@ -280,7 +280,7 @@ TEST_F(CPUCachingAllocatorTest, SizeAlignmentAdjustment) {
   CPUCachingAllocator allocator(1024 * 1024); // 1MB max size
 
   // Test that allocation sizes get properly aligned
-  auto p1 = allocator.allocate(100, 256); // Size not aligned to 256
+  auto p1 = allocator.allocate(100, 256); // Size aligned to 256
   EXPECT_NE(p1, nullptr);
   EXPECT_ALIGNED(p1, 256);
 
@@ -289,9 +289,9 @@ TEST_F(CPUCachingAllocatorTest, SizeAlignmentAdjustment) {
   // allocation
   allocator.reset();
 
-  auto p3 = allocator.allocate(100, 256);
+  auto p3 = allocator.allocate(100, 512);
   // Should reuse p1 due to alignment adjustment
-  EXPECT_EQ(p1, p3);
+  EXPECT_NE(p1, p3);
 }
 
 TEST_F(CPUCachingAllocatorTest, ResetMultipleTimes) {
@@ -348,7 +348,8 @@ TEST_F(CPUCachingAllocatorTest, ResetCachesWhenUnderMaxSize) {
   EXPECT_NE(p1, nullptr);
   EXPECT_NE(p2, nullptr);
 
-  // Reset should cache the allocations since current_size (1024) <= max_size (2048)
+  // Reset should cache the allocations since current_size (1024) <= max_size
+  // (2048)
   allocator.reset();
 
   // Subsequent allocations should reuse the cached pointers