Skip to content

Split tensor arena into persistent and non-persistent arenas #164

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
andresovela wants to merge 5 commits into
base: develop
Choose a base branch
from
Open

Split tensor arena into persistent and non-persistent arenas #164

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
33fd2e4
Split tensor arena into persistent and non-persistent arenas
andresovela Jul 17, 2024
7257738
Fix compiler errors
andresovela Jul 18, 2024
2925a8b
Rename misnamed variables
andresovela Jul 18, 2024
63f24c3
Remove unused variable
andresovela Jul 18, 2024
198ef5f
Declar kPersistentTensorArenaSize only if SPLIT_PERSISTENT_TENSOR_ARE…
andresovela Jul 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion tflite_micro_compiler/src/Api.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,19 @@ std::string TFLMC_Compiler::getTensorName(int tensorIndex, int sg) const {
return compiler_->getTensorName(tensorIndex, sg);
}

// Returns tensor arena size
// Returns tensor arena size (persistent + non-persistent)
size_t TFLMC_Compiler::getTensorArenaSize() const {
return compiler_->getTensorArenaSize();
}

// Returns non-persistent arena size
size_t TFLMC_Compiler::getNonPersistentArenaSize() const {
return compiler_->getNonPersistentArenaSize();
}

// Returns persistent arena size
size_t TFLMC_Compiler::getPersistentArenaSize() const {
return compiler_->getPersistentArenaSize();
}

} // namespace tflmc
6 changes: 6 additions & 0 deletions tflite_micro_compiler/src/Api.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,12 @@ class TFLMC_Compiler {
// Returns tensor arena size
size_t getTensorArenaSize() const;

// Returns non-persistent arena size
size_t getNonPersistentArenaSize() const;

// Returns persistent arena size
size_t getPersistentArenaSize() const;

private:
Compiler *compiler_;
};
Expand Down
77 changes: 61 additions & 16 deletions tflite_micro_compiler/src/Compiler.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,7 @@ bool tflmc::Compiler::init(const void *modelData) {
arena_buf_.resize(SUFFICIENT_ARENA_SIZE);

g_arena_size = SUFFICIENT_ARENA_SIZE;
// TODO: This looks unused?
std::vector<uint8_t> arena_buf(g_arena_size);
g_arenaPtr = arena_buf_.data();

Expand Down Expand Up @@ -348,13 +349,15 @@ bool tflmc::Compiler::init(const void *modelData) {
}
}

// This includes:
// Non-persistent arena consists of:
// - Tensors
// - Scratch buffers
nonPersistentArenaSize_ = ramTensorBufferSize;

// Persistent arena consists of:
// - Persistent buffers
// - Variable tensors (which are allocated as persistent buffers separately)
arenaBufferSize_ =
ramTensorBufferSize + totalRuntimeAllocSize + varTensorsSize;
persistentArenaSize_ = totalRuntimeAllocSize + varTensorsSize;

if (debugPrint_) {
interpreter_->allocator_.memory_planner()->PrintMemoryPlan();
Expand Down Expand Up @@ -545,14 +548,29 @@ namespace xcore {
wr << R"(

constexpr int kTensorArenaSize = )"
<< arenaBufferSize_ << R"(;
#ifndef SHARED_TENSOR_ARENA
<< persistentArenaSize_ + nonPersistentArenaSize_ << R"(;

#ifdef EXTERN_TENSOR_ARENA

#ifdef SPLIT_PERSISTENT_TENSOR_ARENA
constexpr int kPersistentTensorArenaSize = )"
<< persistentArenaSize_ << R"(;
extern uint8_t persistent_tensor_arena[];
#endif // SPLIT_PERSISTENT_TENSOR_ARENA

extern uint8_t tensor_arena[];

#else

#ifdef SPLIT_PERSISTENT_TENSOR_ARENA
#error "Split persistent/non-persistent tensor arenas require externally defined tensor arenas (EXTERN_TENSOR_ARENA must be defined) "
#endif

namespace {
uint8_t tensor_arena[kTensorArenaSize] ALIGN(8);
}
#else
extern uint8_t tensor_arena[];
#endif

#endif // EXTERN_TENSOR_ARENA

namespace {
template <int SZ, class T> struct TfArray {
Expand Down Expand Up @@ -899,6 +917,7 @@ static TfLiteStatus RequestScratchBufferInArena(struct TfLiteContext *context, s

static void *GetScratchBuffer(struct TfLiteContext *context,
int buffer_idx) {
// Scratch buffers are allocated in the non-persistent tensor arena
return tensor_arena + scratch_buffer_offsets[buffer_idx];
}

Expand Down Expand Up @@ -1044,7 +1063,12 @@ TfLiteStatus )"
<< prefix_ << R"(init(void *flash_data) {
// Clear and initialize
scratch_buffer_idx = 0;

#ifdef SPLIT_PERSISTENT_TENSOR_ARENA
persistentBufferPtr = persistent_tensor_arena + kPersistentTensorArenaSize;
#else
persistentBufferPtr = tensor_arena + kTensorArenaSize;
#endif

// Set flash data in xcore context config
xc_config.flash_data = flash_data;
Expand Down Expand Up @@ -1348,14 +1372,35 @@ void tflmc::Compiler::writeHeader(std::ostream &out) {

#include "tensorflow/lite/c/common.h"

#ifdef SHARED_TENSOR_ARENA
#ifndef LARGEST_TENSOR_ARENA_SIZE
#define LARGEST_TENSOR_ARENA_SIZE )" +
std::to_string(arenaBufferSize_) + R"(
#elif LARGEST_TENSOR_ARENA_SIZE < )" +
std::to_string(arenaBufferSize_) + R"(
#define LARGEST_TENSOR_ARENA_SIZE )" +
std::to_string(arenaBufferSize_) + R"(
#ifdef EXTERN_TENSOR_ARENA
#ifdef SPLIT_PERSISTENT_TENSOR_ARENA
#ifndef LARGEST_TENSOR_ARENA_SIZE
#define LARGEST_TENSOR_ARENA_SIZE )" +
std::to_string(nonPersistentArenaSize_) + R"(
#elif LARGEST_TENSOR_ARENA_SIZE < )" +
std::to_string(nonPersistentArenaSize_) + R"(
#define LARGEST_TENSOR_ARENA_SIZE )" +
std::to_string(nonPersistentArenaSize_) + R"(
#endif

#ifndef LARGEST_PERSISTENT_TENSOR_ARENA_SIZE
#define LARGEST_PERSISTENT_TENSOR_ARENA_SIZE )" +
std::to_string(persistentArenaSize_) + R"(
#elif LARGEST_PERSISTENT_TENSOR_ARENA_SIZE < )" +
std::to_string(persistentArenaSize_) + R"(
#define LARGEST_PERSISTENT_TENSOR_ARENA_SIZE )" +
std::to_string(persistentArenaSize_) + R"(
#endif

#else
#ifndef LARGEST_TENSOR_ARENA_SIZE
#define LARGEST_TENSOR_ARENA_SIZE )" +
std::to_string(persistentArenaSize_ + nonPersistentArenaSize_) + R"(
#elif LARGEST_TENSOR_ARENA_SIZE < )" +
std::to_string(persistentArenaSize_ + nonPersistentArenaSize_) + R"(
#define LARGEST_TENSOR_ARENA_SIZE )" +
std::to_string(persistentArenaSize_ + nonPersistentArenaSize_) + R"(
#endif
#endif
#endif

Expand Down
11 changes: 9 additions & 2 deletions tflite_micro_compiler/src/Compiler.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,13 @@ class Compiler {
std::string getTensorName(int tensorIndex, int sg) const;

// Returns tensor arena size
size_t getTensorArenaSize() const { return arenaBufferSize_; }
size_t getTensorArenaSize() const { return persistentArenaSize_ + nonPersistentArenaSize_; }

// Returns persistent arena size
size_t getPersistentArenaSize() const { return persistentArenaSize_; }

// Returns non-persistent arena size
size_t getNonPersistentArenaSize() const { return nonPersistentArenaSize_; }

private:
bool init(const void *modelData);
Expand Down Expand Up @@ -105,7 +111,8 @@ class Compiler {
std::unique_ptr<tflite::MicroInterpreter> interpreter_;
MemMap memMap_;

size_t arenaBufferSize_ = 0;
size_t persistentArenaSize_ = 0;
size_t nonPersistentArenaSize_ = 0;
size_t varTensors_count = 0;
// Vector of vector is for subgraphs
std::vector<std::vector<TensorInfo>> tensors_;
Expand Down