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.cpp

@@ -1,6 +1,7 @@
 #include <cstdlib>
 
 #include <executorch/extension/memory_allocator/cpu_caching_malloc_allocator.h>
+#include <executorch/extension/memory_allocator/memory_allocator_utils.h>
 
 namespace executorch::extension {
 
@@ -18,6 +19,11 @@ void* CPUCachingAllocator::allocate(size_t size, size_t alignment) {
     return nullptr;
   }
   alignment = std::max(alignment, kCachingAllocatorDefaultAlignment);
+  size_t adjusted_size = executorch::extension::utils::get_aligned_size(size, alignment);
+  if (adjusted_size == 0) {
+    return nullptr;
+  }
+  size = adjusted_size;
 
   std::lock_guard<std::mutex> guard(mutex_);
   const auto& it = available_map_.find(size);
@@ -26,14 +32,7 @@ void* CPUCachingAllocator::allocate(size_t size, size_t alignment) {
   // 2. Allocate new memory
   // 2 can lead to current_size > max_size_
   if (it == available_map_.end() || it->second.empty()) {
-    void* ptr = nullptr;
-#if defined(__ANDROID__)
-    ptr = memalign(alignment, size);
-#elif defined(_MSC_VER)
-    ptr = _aligned_malloc(size, alignment);
-#else
-    ptr = std::aligned_alloc(alignment, size);
-#endif
+    void* ptr = std::malloc(size);
     if (ptr == nullptr) {
       ET_LOG(Error, "Failed to allocate memory");
       return nullptr;

extension/memory_allocator/cpu_caching_malloc_allocator.h

@@ -70,7 +70,13 @@ class CPUCachingAllocator : public executorch::runtime::MemoryAllocator {
   /*
     max_size: Maximum size of memory to cache. Never cache more than that.
   */
-  CPUCachingAllocator(uint32_t max_size);
+  explicit CPUCachingAllocator(uint32_t max_size);
+  // No copies allowed
+  CPUCachingAllocator(const CPUCachingAllocator&) = delete;
+  CPUCachingAllocator& operator=(const CPUCachingAllocator&) = delete;
+  // No moves allowed
+  CPUCachingAllocator(CPUCachingAllocator&&) = delete;
+  CPUCachingAllocator& operator=(CPUCachingAllocator&&) = delete;
   // 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.

extension/memory_allocator/malloc_memory_allocator.h

@@ -13,6 +13,7 @@
 #include <cstdlib>
 #include <vector>
 
+#include <executorch/extension/memory_allocator/memory_allocator_utils.h>
 #include <executorch/runtime/core/memory_allocator.h>
 
 namespace executorch {
@@ -51,20 +52,11 @@ class MallocMemoryAllocator : public executorch::runtime::MemoryAllocator {
       return nullptr;
     }
 
-    // The minimum alignment that malloc() is guaranteed to provide.
-    static constexpr size_t kMallocAlignment = alignof(std::max_align_t);
-    if (alignment > kMallocAlignment) {
-      // To get higher alignments, allocate extra and then align the returned
-      // pointer. This will waste an extra `alignment - 1` bytes every time, but
-      // this is the only portable way to get aligned memory from the heap.
-      const size_t extra = alignment - 1;
-      if ET_UNLIKELY (extra >= SIZE_MAX - size) {
-        ET_LOG(
-            Error, "Malloc size overflow: size=%zu + extra=%zu", size, extra);
-        return nullptr;
-      }
-      size += extra;
+    size_t adjusted_size = executorch::extension::utils::get_aligned_size(size, alignment);
+    if (adjusted_size == 0) {
+      return nullptr;
     }
+    size = adjusted_size;
     void* mem_ptr = std::malloc(size);
     if (!mem_ptr) {
       ET_LOG(Error, "Malloc failed to allocate %zu bytes", size);

extension/memory_allocator/memory_allocator_utils.h

@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <cstdlib>
+
+#include <executorch/runtime/core/error.h>
+#include <executorch/runtime/platform/compiler.h>
+
+namespace executorch::extension::utils {
+
+// Util to get alighment adjusted allocation size
+inline size_t get_aligned_size(size_t size, size_t alignment) {
+  // The minimum alignment that malloc() is guaranteed to provide.
+  static constexpr size_t kMallocAlignment = alignof(std::max_align_t);
+  if (alignment > kMallocAlignment) {
+    // To get higher alignments, allocate extra and then align the returned
+    // pointer. This will waste an extra `alignment - 1` bytes every time, but
+    // this is the only portable way to get aligned memory from the heap.
+    const size_t extra = alignment - 1;
+    if ET_UNLIKELY (extra >= SIZE_MAX - size) {
+      ET_LOG(
+          Error, "Malloc size overflow: size=%zu + extra=%zu", size, extra);
+      return 0;
+    }
+    size += extra;
+  }
+  return size;
+}
+
+} // namespace executorch::extension::utils

extension/memory_allocator/targets.bzl

@@ -9,6 +9,9 @@ def define_common_targets():
 
     runtime.cxx_library(
         name = "malloc_memory_allocator",
+        headers = [
+            "memory_allocator_utils.h",
+        ],
         exported_headers = [
             "malloc_memory_allocator.h",
         ],
@@ -23,6 +26,9 @@ def define_common_targets():
 
     runtime.cxx_library(
         name = "cpu_caching_allocator",
+        headers = [
+            "memory_allocator_utils.h",
+        ],
         srcs = [
             "cpu_caching_malloc_allocator.cpp",
         ],

extension/memory_allocator/test/cpu_caching_malloc_allocator_test.cpp

@@ -241,6 +241,7 @@ TEST_F(CPUCachingAllocatorTest, ThreadSafety) {
   };
 
   // Create threads
+  threads.reserve(num_threads);
   for (int i = 0; i < num_threads; ++i) {
     threads.emplace_back(thread_work);
   }
@@ -284,7 +285,7 @@ TEST_F(CPUCachingAllocatorTest, SizeAlignmentAdjustment) {
   EXPECT_NE(p1, nullptr);
   EXPECT_ALIGNED(p1, 256);
 
-  auto p2 = allocator.allocate(100, 256);
+  allocator.allocate(100, 256);
   // Should not get cached pointer since size was adjusted during first
   // allocation
   allocator.reset();