diff options
author | Yu-Cheng Ling <ycling@google.com> | 2018-07-20 16:23:51 -0700 |
---|---|---|
committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-07-20 16:28:54 -0700 |
commit | a4bab4517eddef07236529c0141e85bcae06ad74 (patch) | |
tree | 4ac8bf545068adbe6007ee2dd35273df58cb5cd3 /tensorflow/contrib/lite/arena_planner.h | |
parent | 41781bad97698c29cd74203cef465d2adb2f04e8 (diff) |
Align TFLite tensors to 64 bytes for EIGEN_DONT_ALIGN
PiperOrigin-RevId: 205471025
Diffstat (limited to 'tensorflow/contrib/lite/arena_planner.h')
-rw-r--r-- | tensorflow/contrib/lite/arena_planner.h | 10 |
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..55003cf4e9 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 = 64; + 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 |