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.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();