Fix std::aligned_alloc feature detection. (pytorch#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.

Replace the feature detection with a check on `__cplusplus`.

While I'm here, use smart pointers to manage memory where it is simple to do
so and remove allocation size rounding where it isn't required.

Differential Revision: D75806635

Author: Tamir Duberstein

extension/data_loader/file_data_loader.cpp

@@ -13,6 +13,7 @@
 #include <cstddef>
 #include <cstring>
 #include <limits>
+#include <memory>
 
 #include <executorch/runtime/platform/compat_unistd.h>
 #include <fcntl.h>
@@ -149,7 +150,9 @@ Result<FreeableBuffer> FileDataLoader::load(
   }
 
   // Allocate memory for the FreeableBuffer.
-  void* aligned_buffer = ET_ALIGNED_ALLOC(alignment_, size);
+  std::unique_ptr<void, void (*)(void*)> aligned_buffer(
+      ET_ALIGNED_ALLOC(alignment_, size),
+      [](void* ptr) { ET_ALIGNED_FREE(ptr); });
   if (aligned_buffer == nullptr) {
     ET_LOG(
         Error,
@@ -161,14 +164,14 @@ Result<FreeableBuffer> FileDataLoader::load(
     return Error::MemoryAllocationFailed;
   }
 
-  auto err = load_into(offset, size, segment_info, aligned_buffer);
+  auto err = load_into(offset, size, segment_info, aligned_buffer.get());
   if (err != Error::Ok) {
-    ET_ALIGNED_FREE(aligned_buffer);
     return err;
   }
 
   // Pass the aligned_buffer pointer as context to FreeSegment.
-  return FreeableBuffer(aligned_buffer, size, FreeSegment, aligned_buffer);
+  void* ptr = aligned_buffer.release();
+  return FreeableBuffer(ptr, size, FreeSegment, ptr);
 }
 
 Result<size_t> FileDataLoader::size() const {

extension/data_loader/file_descriptor_data_loader.cpp

@@ -13,6 +13,7 @@
 #include <cstddef>
 #include <cstring>
 #include <limits>
+#include <memory>
 
 #include <fcntl.h>
 #include <sys/stat.h>
@@ -160,24 +161,27 @@ Result<FreeableBuffer> FileDescriptorDataLoader::load(
   }
 
   // Allocate memory for the FreeableBuffer.
-  void* aligned_buffer = ET_ALIGNED_ALLOC(alignment_, size);
+  std::unique_ptr<void, void (*)(void*)> aligned_buffer(
+      ET_ALIGNED_ALLOC(alignment_, size),
+      [](void* ptr) { ET_ALIGNED_FREE(ptr); });
   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(%zd, %zd) failed",
         file_descriptor_uri_,
         offset,
+        alignment_,
         size);
     return Error::MemoryAllocationFailed;
   }
 
-  auto err = load_into(offset, size, segment_info, aligned_buffer);
+  auto err = load_into(offset, size, segment_info, aligned_buffer.get());
   if (err != Error::Ok) {
-    ET_ALIGNED_FREE(aligned_buffer);
     return err;
   }
 
-  return FreeableBuffer(aligned_buffer, size, FreeSegment, aligned_buffer);
+  void* ptr = aligned_buffer.release();
+  return FreeableBuffer(ptr, size, FreeSegment, ptr);
 }
 
 Result<size_t> FileDescriptorDataLoader::size() const {

runtime/platform/compiler.h

@@ -182,38 +182,36 @@ using ssize_t = ptrdiff_t;
  * Note: alignment must be a power of 2 and size must be an integral multiple of
  * alignment.
  */
+#include <cstdlib>
 #if defined(_MSC_VER)
-#include <malloc.h>
-#define ET_ALIGNED_ALLOC(alignment, size) \
-  _aligned_malloc(((size + alignment - 1) & ~(alignment - 1)), (alignment))
+#define ET_ALIGNED_ALLOC(alignment, size) _aligned_malloc((size), (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) {
+  if (posix_memalign(&ptr, alignment, size) != 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>
+#elif __cplusplus >= 201703L
 #define ET_ALIGNED_ALLOC(alignment, size) \
-  ::aligned_alloc(alignment, (size + alignment - 1) & ~(alignment - 1))
-#define ET_ALIGNED_FREE(ptr) free(ptr)
+  std::aligned_alloc(alignment, (size + alignment - 1) & ~(alignment - 1))
+#define ET_ALIGNED_FREE(ptr) std::free(ptr)
 #else
+#if defined(__has_feature) && __has_feature(hwaddress_sanitizer)
+#error native aligned allocation is not available on this platform, and the fallback is not compatible with HWASAN.
+#endif
+
 // If the platform doesn't support aligned_alloc, fallback to malloc.
-#include <stdint.h>
-#include <cstdlib>
+#include <cstdint>
 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);
@@ -263,7 +261,6 @@ inline void et_aligned_free(void* 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.