From a4bab4517eddef07236529c0141e85bcae06ad74 Mon Sep 17 00:00:00 2001 From: Yu-Cheng Ling Date: Fri, 20 Jul 2018 16:23:51 -0700 Subject: Align TFLite tensors to 64 bytes for EIGEN_DONT_ALIGN PiperOrigin-RevId: 205471025 --- tensorflow/contrib/lite/arena_planner.cc | 25 +++++++++---------------- 1 file changed, 9 insertions(+), 16 deletions(-) (limited to 'tensorflow/contrib/lite/arena_planner.cc') diff --git a/tensorflow/contrib/lite/arena_planner.cc b/tensorflow/contrib/lite/arena_planner.cc index 16a0e71624..02442575b3 100644 --- a/tensorflow/contrib/lite/arena_planner.cc +++ b/tensorflow/contrib/lite/arena_planner.cc @@ -17,14 +17,6 @@ limitations under the License. namespace tflite { -namespace { - -// Memory allocation tuning -constexpr const int kDefaultArenaAlignment = 64; -constexpr const int kDefaultTensorAlignment = 4; - -} // namespace - struct AllocationInfo { // The node index requesting this allocation. int node; @@ -36,13 +28,16 @@ struct AllocationInfo { ArenaPlanner::ArenaPlanner(TfLiteContext* context, std::unique_ptr graph_info, - bool preserve_inputs, bool preserve_intermediates) + bool preserve_inputs, bool preserve_intermediates, + int tensor_alignment) : context_(context), graph_info_(std::move(graph_info)), arena_(kDefaultArenaAlignment), persistent_arena_(kDefaultArenaAlignment), preserve_inputs_(preserve_inputs), - preserve_intermediates_(preserve_intermediates) {} + preserve_intermediates_(preserve_intermediates), + tensor_alignment_(tensor_alignment) {} + ArenaPlanner::~ArenaPlanner() {} int64_t ArenaPlanner::BasePointer(TfLiteAllocationType type) { @@ -264,14 +259,12 @@ TfLiteStatus ArenaPlanner::ResolveTensorAllocation(int tensor_index) { TfLiteStatus ArenaPlanner::CalculateTensorAllocation(int tensor_index) { TfLiteTensor& tensor = *graph_info_->tensor(tensor_index); if (tensor.allocation_type == kTfLiteArenaRw) { - TF_LITE_ENSURE_STATUS(arena_.Allocate(context_, kDefaultTensorAlignment, - tensor.bytes, - &allocs_[tensor_index])); + TF_LITE_ENSURE_STATUS(arena_.Allocate( + context_, tensor_alignment_, tensor.bytes, &allocs_[tensor_index])); } if (tensor.allocation_type == kTfLiteArenaRwPersistent) { - TF_LITE_ENSURE_STATUS( - persistent_arena_.Allocate(context_, kDefaultTensorAlignment, - tensor.bytes, &allocs_[tensor_index])); + TF_LITE_ENSURE_STATUS(persistent_arena_.Allocate( + context_, tensor_alignment_, tensor.bytes, &allocs_[tensor_index])); } return kTfLiteOk; } -- cgit v1.2.3