Align TFLite tensors to 16 bytes for EIGEN_DONT_ALIGN

PiperOrigin-RevId: 205015541

1 files changed, 5 insertions, 3 deletions

diff --git a/tensorflow/contrib/lite/arena_planner_test.cc b/tensorflow/contrib/lite/arena_planner_test.cc
index 1adb426d58..7d7c41289c 100644
--- a/tensorflow/contrib/lite/arena_planner_test.cc
+++ b/tensorflow/contrib/lite/arena_planner_test.cc
@@ -24,6 +24,8 @@ limitations under the License.
 namespace tflite {
 namespace {
 
+constexpr const int kTensorAlignment = 4;
+
 // A simple op to be used in tests, as syntactic sugar.
 class TestOp {
  public:
@@ -156,7 +158,7 @@ class ArenaPlannerTest : public ::testing::Test {
     context_.ReportError = ReportError;
     planner_.reset(new ArenaPlanner(
         &context_, std::unique_ptr<GraphInfo>(new TestGraphInfo(graph)),
-        preserve_inputs, /*preserve intermediates*/ false));
+        preserve_inputs, /*preserve intermediates*/ false, kTensorAlignment));
     CHECK(planner_->ResetAllocations() == kTfLiteOk);
     CHECK(planner_->PlanAllocations() == kTfLiteOk);
   }
@@ -178,8 +180,8 @@ class ArenaPlannerTest : public ::testing::Test {
     const TfLiteTensor& tensor = (*graph_->tensors())[tensor_index];
     int64_t offset = GetOffset(tensor_index) + tensor.bytes;
     // We must make sure the offset is aligned to kDefaultArenaAlignment.
-    if (offset % 4 != 0) {
-      offset += 4 - offset % 4;
+    if (offset % kTensorAlignment != 0) {
+      offset += kTensorAlignment - offset % kTensorAlignment;
     }
     return offset;
   };