From 634c600cc37182ca5164af0847ebc208d8ef8059 Mon Sep 17 00:00:00 2001
From: Sachin Joglekar <srjoglekar@google.com>
Date: Thu, 9 Jul 2020 09:52:21 -0700
Subject: [PATCH] Shift memory allocation parameters to lite/util. Also make
specifying these parameters explicit in runtime code.
PiperOrigin-RevId: 320416196
Change-Id: I3e321b7ca6a83a968600d64278e991a6280ebd80
---
tensorflow/lite/BUILD | 1 +
tensorflow/lite/arena_planner.h | 6 ++----
tensorflow/lite/core/subgraph.cc | 3 ++-
tensorflow/lite/util.h | 6 ++++++
4 files changed, 11 insertions(+), 5 deletions(-)
diff --git a/tensorflow/lite/BUILD b/tensorflow/lite/BUILD
index c1179f8b1c6..06bfaf348fb 100644
--- a/tensorflow/lite/BUILD
+++ b/tensorflow/lite/BUILD
@@ -106,6 +106,7 @@ cc_library(
":graph_info",
":memory_planner",
":simple_memory_arena",
+ ":util",
"//tensorflow/lite/c:common",
],
)
diff --git a/tensorflow/lite/arena_planner.h b/tensorflow/lite/arena_planner.h
index 548d1c3e4c0..5f7b9901087 100644
--- a/tensorflow/lite/arena_planner.h
+++ b/tensorflow/lite/arena_planner.h
@@ -23,13 +23,11 @@ limitations under the License.
#include "tensorflow/lite/graph_info.h"
#include "tensorflow/lite/memory_planner.h"
#include "tensorflow/lite/simple_memory_arena.h"
+#include "tensorflow/lite/util.h"
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.
@@ -53,7 +51,7 @@ class ArenaPlanner : public MemoryPlanner {
// them until the end of inference.
ArenaPlanner(TfLiteContext* context, std::unique_ptr<GraphInfo> graph_info,
bool preserve_inputs, bool preserve_intermediates,
- int tensor_alignment = kDefaultTensorAlignment);
+ int tensor_alignment);
~ArenaPlanner() override;
ArenaPlanner(const ArenaPlanner&) = delete;
ArenaPlanner& operator=(const ArenaPlanner&) = delete;
diff --git a/tensorflow/lite/core/subgraph.cc b/tensorflow/lite/core/subgraph.cc
index 0f11af51488..5ef9b45514b 100644
--- a/tensorflow/lite/core/subgraph.cc
+++ b/tensorflow/lite/core/subgraph.cc
@@ -859,7 +859,8 @@ TfLiteStatus Subgraph::PrepareOpsAndTensors() {
if (!memory_planner_) {
memory_planner_.reset(new ArenaPlanner(
&context_, std::unique_ptr<GraphInfo>(new InterpreterInfo(this)),
- /*preserve_inputs=*/true, /*preserve_intermediates*/ false));
+ /*preserve_inputs=*/true, /*preserve_intermediates*/ false,
+ kDefaultTensorAlignment));
memory_planner_->PlanAllocations();
}
diff --git a/tensorflow/lite/util.h b/tensorflow/lite/util.h
index 71b69c755fc..cf4576db834 100644
--- a/tensorflow/lite/util.h
+++ b/tensorflow/lite/util.h
@@ -29,6 +29,12 @@ limitations under the License.
namespace tflite {
+// Memory allocation parameter used by ArenaPlanner.
+// Clients (such as delegates) might look at this to ensure interop between
+// TFLite memory & hardware buffers.
+// NOTE: This only holds for tensors allocated on the arena.
+constexpr int kDefaultTensorAlignment = 64;
+
// The prefix of Flex op custom code.
// This will be matched agains the `custom_code` field in `OperatorCode`
// Flatbuffer Table.