Use __cpp_aligned_new instead of hand-rolling the implementation (#11293)

Summary:

The current feature detection seems to be incorrect, resulting in fallback to
the manual implementation, which in turn returns untagged pointers under
hwasan.

Use new expressions instead to delegate the work to the C++ compiler.

Differential Revision: D75806635

Author: Tamir Duberstein

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,30 @@ 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,
+      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,30 @@ 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,
+      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