Add parameter to BuiltinDataAllocator::Allocate for requesting alignment.

Update MallocDataAllocator to use aligned_alloc if that is available (otherwise ignore it, use malloc). In TFLM's MicroBuiltinDataAllocator, update logic to align to the appropriate multiple, potentially saving some memory.

PiperOrigin-RevId: 307907552
Change-Id: Ia67c6f1df50671bfc784028fd2f8e6c008f6e9d6

tensorflow/lite/core/api/flatbuffer_conversions.h

@@ -29,7 +29,7 @@ namespace tflite {
 // Interface class for builtin data allocations.
 class BuiltinDataAllocator {
  public:
-  virtual void* Allocate(size_t size) = 0;
+  virtual void* Allocate(size_t size, size_t alignment_hint) = 0;
   virtual void Deallocate(void* data) = 0;
 
   // Allocate a structure, but make sure it is a POD structure that doesn't
@@ -41,7 +41,7 @@ class BuiltinDataAllocator {
     // TODO(b/154346074): Change this to is_trivially_destructible when all
     // platform targets support that properly.
     static_assert(std::is_pod<T>::value, "Builtin data structure must be POD.");
-    void* allocated_memory = this->Allocate(sizeof(T));
+    void* allocated_memory = this->Allocate(sizeof(T), alignof(T));
     return new (allocated_memory) T;
   }
 

tensorflow/lite/core/api/flatbuffer_conversions_test.cc

@@ -47,7 +47,7 @@ class MockErrorReporter : public ErrorReporter {
 class MockDataAllocator : public BuiltinDataAllocator {
  public:
   MockDataAllocator() : is_allocated_(false) {}
-  void* Allocate(size_t size) override {
+  void* Allocate(size_t size, size_t alignment_hint) override {
     EXPECT_FALSE(is_allocated_);
     const int max_size = kBufferSize;
     EXPECT_LE(size, max_size);

tensorflow/lite/interpreter_builder.cc

@@ -34,6 +34,15 @@ limitations under the License.
 #include "tensorflow/lite/profiling/platform_profiler.h"
 #endif
 
+// aligned_alloc is available (via cstdlib/stdlib.h) with C++17/C11.
+#if __cplusplus >= 201703L || __STDC_VERSION__ >= 201112L
+#if !defined(__ANDROID__) || __ANDROID_API__ >= 28
+#if !defined(__APPLE__)  // Apple does not provide aligned_alloc.
+#define TFLITE_USE_STD_ALIGNED_ALLOC
+#endif
+#endif
+#endif
+
 namespace tflite {
 
 namespace {
@@ -197,7 +206,13 @@ std::vector<int> FlatBufferIntArrayToVector(T* flat_array) {
 // Used to determine how the op data parsing function creates its working space.
 class MallocDataAllocator : public BuiltinDataAllocator {
  public:
-  void* Allocate(size_t size) override { return malloc(size); }
+  void* Allocate(size_t size, size_t alignment_hint) override {
+#ifdef TFLITE_USE_STD_ALIGNED_ALLOC
+    return aligned_alloc(alignment_hint, size);
+#else
+    return malloc(size);
+#endif
+  }
   void Deallocate(void* data) override { free(data); }
 };
 

tensorflow/lite/micro/micro_allocator.cc

@@ -44,29 +44,13 @@ struct AllocationInfo {
 // requirement for SIMD extensions.
 constexpr int kBufferAlignment = 16;
 
-// If building with GNU clib from GCC 4.8.x or lower, `max_align_t` is not a
-// member of `std`. If using a newer version of clib, we import `max_align_t`
-// into the local anonymous namespace to be able to use it like the global
-// `max_align_t` from the older clib.
-#if defined(__GNUC__) && defined(__GNUC_PREREQ)
-#if __GNUC_PREREQ(4, 9)
-using std::max_align_t;
-#endif
-#else
-// We assume other compiler/clib configurations don't have this issue.
-using std::max_align_t;
-#endif
-
 class MicroBuiltinDataAllocator : public BuiltinDataAllocator {
  public:
   explicit MicroBuiltinDataAllocator(SimpleMemoryAllocator* memory_allocator)
       : memory_allocator_(memory_allocator) {}
 
-  void* Allocate(size_t size) override {
-    // Align to an address that is proper for all primitive types, but no more
-    // than the size.
-    return memory_allocator_->AllocateFromTail(
-        size, std::min(size, alignof(max_align_t)));
+  void* Allocate(size_t size, size_t alignment_hint) override {
+    return memory_allocator_->AllocateFromTail(size, alignment_hint);
   }
   void Deallocate(void* data) override {
     // Do not deallocate, builtin data needs to be available for the life time