diff options
author | 2018-07-20 16:23:51 -0700 | |
---|---|---|
committer | 2018-07-20 16:28:54 -0700 | |
commit | a4bab4517eddef07236529c0141e85bcae06ad74 (patch) | |
tree | 4ac8bf545068adbe6007ee2dd35273df58cb5cd3 /tensorflow/contrib/lite/arena_planner.cc | |
parent | 41781bad97698c29cd74203cef465d2adb2f04e8 (diff) |
Align TFLite tensors to 64 bytes for EIGEN_DONT_ALIGN
PiperOrigin-RevId: 205471025
Diffstat (limited to 'tensorflow/contrib/lite/arena_planner.cc')
-rw-r--r-- | tensorflow/contrib/lite/arena_planner.cc | 25 |
1 files changed, 9 insertions, 16 deletions
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<GraphInfo> 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; } |