Align TFLite tensors to 16 bytes for EIGEN_DONT_ALIGN

PiperOrigin-RevId: 205015541

1 files changed, 9 insertions, 1 deletions

diff --git a/tensorflow/contrib/lite/arena_planner.h b/tensorflow/contrib/lite/arena_planner.h
index 82c866734f..65a9730b70 100644
--- a/tensorflow/contrib/lite/arena_planner.h
+++ b/tensorflow/contrib/lite/arena_planner.h
@@ -25,6 +25,10 @@ limitations under the License.
 
 namespace tflite {
 
+// Memory allocation tuning
+constexpr const int kDefaultArenaAlignment = 64;
+constexpr const int kDefaultTensorAlignment = 16;
+
 struct AllocationInfo;
 
 // A memory planner that makes all the allocations using arenas.
@@ -47,7 +51,8 @@ class ArenaPlanner : public MemoryPlanner {
   // graph will not share memory with any other tensor, effectively preserving
   // them until the end of inference.
   ArenaPlanner(TfLiteContext* context, std::unique_ptr<GraphInfo> graph_info,
-               bool preserve_inputs, bool preserve_intermediates);
+               bool preserve_inputs, bool preserve_intermediates,
+               int tensor_alignment = kDefaultTensorAlignment);
   ~ArenaPlanner() override;
   ArenaPlanner(const ArenaPlanner&) = delete;
   ArenaPlanner& operator=(const ArenaPlanner&) = delete;
@@ -112,6 +117,9 @@ class ArenaPlanner : public MemoryPlanner {
   // If true, then no overlapping of memory areas is done, meaning intermediates
   // results can be queried after running (modulo running delegates).
   bool preserve_intermediates_;
+
+  // Number of bytes that tensor buffers should be aligned to.
+  int tensor_alignment_;
 };
 
 }  // namespace tflite