Update base for Update on "[Executorch] Use temp allocator for allocating scratch memory"

This allows us to leverage temp memory allocator and if that allocator is caching allocator it reduces the allocaiton overhead.

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

[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