Use __cpp_aligned_new instead of hand-rolling the implementation

Differential Revision: D75806635

Pull Request resolved: pytorch#11293

Author: tamird

extension/data_loader/file_data_loader.cpp

@@ -112,17 +112,29 @@ Result<FileDataLoader> FileDataLoader::from(
 }
 
 namespace {
+
+inline void* et_aligned_alloc(size_t size, std::align_val_t alignment) {
+  return ::operator new(size, alignment);
+}
+
+inline void et_aligned_free(void* ptr, std::align_val_t alignment) {
+  return ::operator delete(ptr, alignment);
+}
+
 /**
  * FreeableBuffer::FreeFn-compatible callback.
  *
- * `context` is the original buffer pointer. It is allocated with
- * ET_ALIGNED_ALLOC, and must be freed with ET_ALIGNED_FREE.
+ * `data` is the original buffer pointer.
+ * `context` is the original alignment.
  *
- * `data` and `size` are unused.
+ * `size` is unused.
  */
 void FreeSegment(void* context, void* data, ET_UNUSED size_t size) {
-  ET_ALIGNED_FREE(context);
+  et_aligned_free(
+      data,
+      static_cast<std::align_val_t>(reinterpret_cast<uintptr_t>(context)));
 }
+
 } // namespace
 
 Result<FreeableBuffer> FileDataLoader::load(
@@ -149,26 +161,31 @@ Result<FreeableBuffer> FileDataLoader::load(
   }
 
   // Allocate memory for the FreeableBuffer.
-  void* aligned_buffer = ET_ALIGNED_ALLOC(alignment_, size);
+  void* aligned_buffer = et_aligned_alloc(size, alignment_);
   if (aligned_buffer == nullptr) {
     ET_LOG(
         Error,
-        "Reading from %s at offset %zu: ET_ALIGNED_ALLOC(%zd, %zd) failed",
+        "Reading from %s at offset %zu: et_aligned_alloc(%zu, %zu) failed",
         file_name_,
         offset,
-        alignment_,
-        size);
+        size,
+        static_cast<size_t>(alignment_));
     return Error::MemoryAllocationFailed;
   }
 
   auto err = load_into(offset, size, segment_info, aligned_buffer);
   if (err != Error::Ok) {
-    ET_ALIGNED_FREE(aligned_buffer);
+    et_aligned_free(aligned_buffer, alignment_);
     return err;
   }
 
-  // Pass the aligned_buffer pointer as context to FreeSegment.
-  return FreeableBuffer(aligned_buffer, size, FreeSegment, aligned_buffer);
+  // Pass the alignment as context to FreeSegment.
+  return FreeableBuffer(
+      aligned_buffer,
+      size,
+      FreeSegment,
+      // NOLINTNEXTLINE(performance-no-int-to-ptr)
+      reinterpret_cast<void*>(static_cast<uintptr_t>(alignment_)));
 }
 
 Result<size_t> FileDataLoader::size() const {

extension/data_loader/file_data_loader.h

@@ -58,7 +58,7 @@ class FileDataLoader final : public executorch::runtime::DataLoader {
         fd_(rhs.fd_) {
     const_cast<const char*&>(rhs.file_name_) = nullptr;
     const_cast<size_t&>(rhs.file_size_) = 0;
-    const_cast<size_t&>(rhs.alignment_) = 0;
+    const_cast<std::align_val_t&>(rhs.alignment_) = {};
     const_cast<int&>(rhs.fd_) = -1;
   }
 
@@ -86,7 +86,7 @@ class FileDataLoader final : public executorch::runtime::DataLoader {
       const char* file_name)
       : file_name_(file_name),
         file_size_(file_size),
-        alignment_(alignment),
+        alignment_{alignment},
         fd_(fd) {}
 
   // Not safely copyable.
@@ -96,7 +96,7 @@ class FileDataLoader final : public executorch::runtime::DataLoader {
 
   const char* const file_name_; // Owned by the instance.
   const size_t file_size_;
-  const size_t alignment_;
+  const std::align_val_t alignment_;
   const int fd_; // Owned by the instance.
 };
 

extension/data_loader/file_descriptor_data_loader.cpp

@@ -123,17 +123,29 @@ FileDescriptorDataLoader::fromFileDescriptorUri(
 }
 
 namespace {
+
+inline void* et_aligned_alloc(size_t size, std::align_val_t alignment) {
+  return ::operator new(size, alignment);
+}
+
+inline void et_aligned_free(void* ptr, std::align_val_t alignment) {
+  return ::operator delete(ptr, alignment);
+}
+
 /**
  * FreeableBuffer::FreeFn-compatible callback.
  *
- * `context` is the original buffer pointer. It is allocated with
- * ET_ALIGNED_ALLOC, and must be freed with ET_ALIGNED_FREE.
+ * `data` is the original buffer pointer.
+ * `context` is the original alignment.
  *
- * `data` and `size` are unused.
+ * `size` is unused.
  */
 void FreeSegment(void* context, void* data, ET_UNUSED size_t size) {
-  ET_ALIGNED_FREE(context);
+  et_aligned_free(
+      data,
+      static_cast<std::align_val_t>(reinterpret_cast<uintptr_t>(context)));
 }
+
 } // namespace
 
 Result<FreeableBuffer> FileDescriptorDataLoader::load(
@@ -160,24 +172,31 @@ Result<FreeableBuffer> FileDescriptorDataLoader::load(
   }
 
   // Allocate memory for the FreeableBuffer.
-  void* aligned_buffer = ET_ALIGNED_ALLOC(alignment_, size);
+  void* aligned_buffer = et_aligned_alloc(size, alignment_);
   if (aligned_buffer == nullptr) {
     ET_LOG(
         Error,
-        "Reading from %s at offset %zu: ET_ALIGNED_ALLOC(%zd) failed",
+        "Reading from %s at offset %zu: et_aligned_alloc(%zu, %zu) failed",
         file_descriptor_uri_,
         offset,
-        size);
+        size,
+        static_cast<size_t>(alignment_));
     return Error::MemoryAllocationFailed;
   }
 
   auto err = load_into(offset, size, segment_info, aligned_buffer);
   if (err != Error::Ok) {
-    ET_ALIGNED_FREE(aligned_buffer);
+    et_aligned_free(aligned_buffer, alignment_);
     return err;
   }
 
-  return FreeableBuffer(aligned_buffer, size, FreeSegment, aligned_buffer);
+  // Pass the alignment as context to FreeSegment.
+  return FreeableBuffer(
+      aligned_buffer,
+      size,
+      FreeSegment,
+      // NOLINTNEXTLINE(performance-no-int-to-ptr)
+      reinterpret_cast<void*>(static_cast<uintptr_t>(alignment_)));
 }
 
 Result<size_t> FileDescriptorDataLoader::size() const {

extension/data_loader/file_descriptor_data_loader.h

@@ -56,7 +56,7 @@ class FileDescriptorDataLoader final : public executorch::runtime::DataLoader {
         fd_(rhs.fd_) {
     const_cast<const char*&>(rhs.file_descriptor_uri_) = nullptr;
     const_cast<size_t&>(rhs.file_size_) = 0;
-    const_cast<size_t&>(rhs.alignment_) = 0;
+    const_cast<std::align_val_t&>(rhs.alignment_) = {};
     const_cast<int&>(rhs.fd_) = -1;
   }
 
@@ -84,7 +84,7 @@ class FileDescriptorDataLoader final : public executorch::runtime::DataLoader {
       const char* file_descriptor_uri)
       : file_descriptor_uri_(file_descriptor_uri),
         file_size_(file_size),
-        alignment_(alignment),
+        alignment_{alignment},
         fd_(fd) {}
 
   // Not safely copyable.
@@ -94,7 +94,7 @@ class FileDescriptorDataLoader final : public executorch::runtime::DataLoader {
 
   const char* const file_descriptor_uri_; // Owned by the instance.
   const size_t file_size_;
-  const size_t alignment_;
+  const std::align_val_t alignment_;
   const int fd_; // Owned by the instance.
 };
 

runtime/platform/compiler.h

@@ -171,101 +171,6 @@
 using ssize_t = ptrdiff_t;
 #endif
 
-/**
- * Platform-specific aligned memory allocation and deallocation.
- *
- * Usage:
- *   void* ptr = ET_ALIGNED_ALLOC(alignment, size);
- *   // use ptr...
- *   ET_ALIGNED_FREE(ptr);
- *
- * Note: alignment must be a power of 2 and size must be an integral multiple of
- * alignment.
- */
-#if defined(_MSC_VER)
-#include <malloc.h>
-#define ET_ALIGNED_ALLOC(alignment, size) \
-  _aligned_malloc(((size + alignment - 1) & ~(alignment - 1)), (alignment))
-#define ET_ALIGNED_FREE(ptr) _aligned_free(ptr)
-#elif defined(__APPLE__)
-#include <stdlib.h> // For posix_memalign and free
-inline void* et_apple_aligned_alloc(size_t alignment, size_t size) {
-  void* ptr = nullptr;
-  // The address of the allocated memory must be a multiple of sizeof(void*).
-  if (alignment < sizeof(void*)) {
-    alignment = sizeof(void*);
-  }
-  if (posix_memalign(
-          &ptr, alignment, (size + alignment - 1) & ~(alignment - 1)) != 0) {
-    return nullptr;
-  }
-  return ptr;
-}
-#define ET_ALIGNED_ALLOC(alignment, size) \
-  et_apple_aligned_alloc((alignment), (size))
-#define ET_ALIGNED_FREE(ptr) free(ptr)
-#elif __has_builtin(__builtin_aligned_alloc) || defined(_ISOC11_SOURCE)
-// Linux and posix systems that support aligned_alloc and are >= C++17.
-#include <cstdlib>
-#define ET_ALIGNED_ALLOC(alignment, size) \
-  ::aligned_alloc(alignment, (size + alignment - 1) & ~(alignment - 1))
-#define ET_ALIGNED_FREE(ptr) free(ptr)
-#else
-// If the platform doesn't support aligned_alloc, fallback to malloc.
-#include <stdint.h>
-#include <cstdlib>
-inline void* et_aligned_malloc(size_t alignment, size_t size) {
-  // Place to store the offset to the original pointer.
-  size_t offset_size = sizeof(uint16_t);
-
-  // Malloc extra space for offset + alignment.
-  size_t alloc_size = size + offset_size + alignment - 1;
-  void* ptr = std::malloc(alloc_size);
-
-  if (ptr == nullptr) {
-    // Malloc failed.
-    return nullptr;
-  }
-
-  uintptr_t addr = reinterpret_cast<uintptr_t>(ptr);
-  // Align the address past addr + offset_size bytes.
-  // This provides space to store the offset before the aligned pointer.
-  addr = addr + offset_size;
-  uintptr_t aligned_ptr = (addr + alignment - 1) & ~(alignment - 1);
-
-  // Check that alignment didn't overflow the buffer.
-  if (reinterpret_cast<uintptr_t>(aligned_ptr) + size >
-      reinterpret_cast<uintptr_t>(ptr) + alloc_size) {
-    std::free(ptr);
-    return nullptr;
-  }
-
-  // Store the offset to the original pointer.
-  // Used to free the original allocated buffer.
-  *(reinterpret_cast<uint16_t*>(aligned_ptr) - 1) =
-      (uint16_t)(reinterpret_cast<uintptr_t>(aligned_ptr) -
-                 reinterpret_cast<uintptr_t>(ptr));
-
-  return reinterpret_cast<uint16_t*>(aligned_ptr);
-}
-
-inline void et_aligned_free(void* ptr) {
-  if (ptr == nullptr) {
-    return;
-  }
-
-  // Get the original pointer using the offset.
-  uint16_t* original_ptr = reinterpret_cast<uint16_t*>(
-      reinterpret_cast<uintptr_t>(ptr) -
-      *(reinterpret_cast<uint16_t*>(ptr) - 1));
-  std::free(original_ptr);
-}
-
-#define ET_ALIGNED_ALLOC(alignment, size) et_aligned_malloc((alignment), (size))
-#define ET_ALIGNED_FREE(ptr) et_aligned_free(ptr)
-
-#endif
-
 // DEPRECATED: Use the non-underscore-prefixed versions instead.
 // TODO(T199005537): Remove these once all users have stopped using them.
 #define __ET_DEPRECATED ET_DEPRECATED