Update on "[Executorch] Introduce caching cpu memory allocator"

Meant to use this for temp allocator for kernels. Specifically for sdpa, it seems that on iOS there is a significant overhead coming from allocations

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

[ghstack-poisoned]

Author: kimishpatel

extension/memory_allocator/cpu_caching_malloc_allocator.cpp

@@ -6,7 +6,7 @@ namespace executorch::extension {
 
 namespace {
 size_t get_alignment_adjusted_size(size_t size, size_t alignment) {
-  alignment = std::max(alignment, kDefaultAlignment);
+  alignment = std::max(alignment, kCachingAllocatorDefaultAlignment);
   if (size % alignment != 0) {
     // Adjust size to the next multiple of alignment
     // This is needed for aligned_alloc to work
@@ -66,6 +66,7 @@ void CPUCachingAllocator::free_cached() {
     }
   }
   available_map_.clear();
+  current_size_ = 0;
 }
 
 void CPUCachingAllocator::reset() {
@@ -75,7 +76,6 @@ void CPUCachingAllocator::reset() {
     size_t alloc_size = it.second;
     // Cache the memory
     available_map_[alloc_size].push_back(ptr);
-    current_size_ -= alloc_size;
   }
   allocation_map_.clear();
 }

extension/memory_allocator/cpu_caching_malloc_allocator.h

@@ -41,7 +41,7 @@ template <typename KeyType, typename ValueType>
 using FlatHashMap = std::unordered_map<KeyType, ValueType>;
 #endif
 
-constexpr size_t kDefaultAlignment = 64;
+constexpr size_t kCachingAllocatorDefaultAlignment = 64;
 class CPUCachingAllocator : public executorch::runtime::MemoryAllocator {
   /*
    * What it does:
@@ -73,7 +73,7 @@ class CPUCachingAllocator : public executorch::runtime::MemoryAllocator {
   // Checks the cache to see if allocation of size bytes can be found.
   // If so return cached memory, else
   // allocates memory, records it for caching and returns.
-  void* allocate(size_t size, size_t alignment = kDefaultAlignment) override;
+  void* allocate(size_t size, size_t alignment = kCachingAllocatorDefaultAlignment) override;
   void reset() override;
   ~CPUCachingAllocator();
 };

extension/memory_allocator/test/cpu_caching_malloc_allocator_test.cpp

@@ -12,7 +12,7 @@
 using namespace ::testing;
 using executorch::extension::CPUCachingAllocator;
 
-constexpr auto kDefaultAlignment = executorch::extension::kDefaultAlignment;
+constexpr auto kDefaultAlignment = executorch::extension::kCachingAllocatorDefaultAlignment;
 
 class CPUCachingAllocatorTest : public ::testing::Test {
  protected:
@@ -25,7 +25,7 @@ class CPUCachingAllocatorTest : public ::testing::Test {
 
 bool is_aligned(const void* ptr, size_t alignment) {
   uintptr_t addr = reinterpret_cast<uintptr_t>(ptr);
-  return addr % alignment == 0;
+  return (addr & (alignment - 1)) == 0;
 }
 
 #define EXPECT_ALIGNED(ptr, alignment)        \
@@ -297,5 +297,6 @@ TEST_F(CPUCachingAllocatorTest, ResetMultipleTimes) {
 
     auto p2 = allocator.allocate(512);
     EXPECT_EQ(p, p2);
+    allocator.reset();
   }
 }