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

@@ -19,12 +19,12 @@ void* CPUCachingAllocator::allocate(size_t size, size_t alignment) {
     return nullptr;
   }
   alignment = std::max(alignment, kCachingAllocatorDefaultAlignment);
-  size_t adjusted_size =
+  auto adjusted_size_value =
       executorch::extension::utils::get_aligned_size(size, alignment);
-  if (adjusted_size == 0) {
+  if (!adjusted_size_value.ok()) {
     return nullptr;
   }
-  size = adjusted_size;
+  size = adjusted_size_value.get();
 
   std::lock_guard<std::mutex> guard(mutex_);
   const auto& it = available_map_.find(size);

extension/memory_allocator/malloc_memory_allocator.h

@@ -52,7 +52,12 @@ class MallocMemoryAllocator : public executorch::runtime::MemoryAllocator {
       return nullptr;
     }
 
-    size = executorch::extension::utils::get_aligned_size(size, alignment);
+    auto adjusted_size_value =
+        executorch::extension::utils::get_aligned_size(size, alignment);
+    if (!adjusted_size_value.ok()) {
+      return nullptr;
+    }
+    size = adjusted_size_value.get();
     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

@@ -13,12 +13,15 @@
 #include <cstdlib>
 
 #include <executorch/runtime/core/error.h>
+#include <executorch/runtime/core/result.h>
 #include <executorch/runtime/platform/compiler.h>
 
+using executorch::runtime::Error;
+using executorch::runtime::Result;
 namespace executorch::extension::utils {
 
 // Util to get alighment adjusted allocation size
-inline size_t get_aligned_size(size_t size, size_t alignment) {
+inline Result<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) {
@@ -28,7 +31,7 @@ inline size_t get_aligned_size(size_t size, size_t alignment) {
     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;
+      return Result<size_t>(Error::InvalidArgument);
     }
     size += extra;
   }