/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
you may not use this file except in compliance with the License.
|
You may obtain a copy of the License at
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
Unless required by applicable law or agreed to in writing, software
|
distributed under the License is distributed on an "AS IS" BASIS,
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
See the License for the specific language governing permissions and
|
limitations under the License.
|
==============================================================================*/
|
#ifndef TENSORFLOW_LITE_SIMPLE_MEMORY_ARENA_H_
|
#define TENSORFLOW_LITE_SIMPLE_MEMORY_ARENA_H_
|
|
#include <list>
|
#include <memory>
|
#include "tensorflow/lite/c/c_api_internal.h"
|
|
namespace tflite {
|
|
// This little structure holds the offset and the size for a dynamic memory
|
// allocation in the memory arena. When the arena is committed and the
|
// underlying buffer is set, the alloc can be resolved into an actual memory
|
// pointer.
|
struct ArenaAlloc {
|
ArenaAlloc() : offset(0), size(0) {}
|
|
size_t offset;
|
size_t size;
|
|
inline bool operator<(const ArenaAlloc& other) const {
|
return offset < other.offset;
|
}
|
};
|
|
// This small class is responsible for allocating, deallocating and reusing
|
// dynamic memory from a common underlying buffer. The arena can be used in
|
// scenarios when the pattern of memory allocations and deallocations is
|
// repetitive, e.g. running NN inference in multiple iterations. Note that
|
// zero-sized allocations are explicitly allowed, and will resolve to null.
|
class SimpleMemoryArena {
|
public:
|
explicit SimpleMemoryArena(size_t arena_alignment)
|
: committed_(false),
|
arena_alignment_(arena_alignment),
|
high_water_mark_(0),
|
underlying_buffer_size_(0),
|
allocs_() {}
|
|
TfLiteStatus Allocate(TfLiteContext* context, size_t alignment, size_t size,
|
ArenaAlloc* new_alloc);
|
|
TfLiteStatus Deallocate(TfLiteContext* context, const ArenaAlloc& alloc);
|
|
inline size_t RequiredBufferSize() {
|
// Add in a small amount of padding to reduce the chance of resize events
|
// for small allocations.
|
size_t padding = arena_alignment_;
|
return arena_alignment_ + high_water_mark_ + padding;
|
}
|
|
TfLiteStatus Commit(TfLiteContext* context);
|
|
TfLiteStatus ResolveAlloc(TfLiteContext* context, const ArenaAlloc& alloc,
|
char** output_ptr);
|
|
TfLiteStatus Clear();
|
|
int64_t BasePointer() const {
|
return reinterpret_cast<int64_t>(underlying_buffer_aligned_ptr_);
|
}
|
|
private:
|
bool committed_;
|
size_t arena_alignment_;
|
size_t high_water_mark_;
|
std::unique_ptr<char[]> underlying_buffer_;
|
size_t underlying_buffer_size_;
|
char* underlying_buffer_aligned_ptr_;
|
// TODO(maciekc): add list iterator to the ArenaAlloc to lookup quickly.
|
std::list<ArenaAlloc> allocs_;
|
};
|
|
} // namespace tflite
|
|
#endif // TENSORFLOW_LITE_SIMPLE_MEMORY_ARENA_H_
|
