wehub-resource-sync
132 lines
4.1 KiB
C++
132 lines
4.1 KiB
C++
/* Copyright 2021 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_PLANNER_H_
|
|
#define TENSORFLOW_LITE_SIMPLE_PLANNER_H_
|
|
|
|
#include <cassert>
|
|
#include <cstddef>
|
|
#include <cstdint>
|
|
#include <cstdlib>
|
|
#include <memory>
|
|
#include <vector>
|
|
|
|
#include "tensorflow/lite/core/c/common.h"
|
|
#include "tensorflow/lite/graph_info.h"
|
|
#include "tensorflow/lite/memory_planner.h"
|
|
#include "tensorflow/lite/util.h"
|
|
|
|
namespace tflite {
|
|
|
|
// A structure to keep heap allocation records. This structure is used by
|
|
// SimplePlanner::allocs_.
|
|
struct SimpleAlloc {
|
|
SimpleAlloc() { reset(); }
|
|
|
|
// Size of allocation.
|
|
size_t size;
|
|
// The index of the node that first needs to use this tensor.
|
|
int32_t node;
|
|
// Allocated heap memory address of allocation.
|
|
char* ptr;
|
|
|
|
// Reset member variables.
|
|
inline void reset() {
|
|
size = 0;
|
|
node = 0;
|
|
ptr = nullptr;
|
|
}
|
|
|
|
// Allocate heap memory for a tensor with the given size and first_node
|
|
// information.
|
|
inline bool alloc(size_t new_size, int32_t new_first_node) {
|
|
if (new_size == 0) {
|
|
return false;
|
|
}
|
|
size = new_size;
|
|
node = new_first_node;
|
|
assert(ptr == nullptr);
|
|
ptr = static_cast<char*>(malloc(new_size));
|
|
return true;
|
|
}
|
|
|
|
// Free allocated heap memory and reset member variables.
|
|
inline void free() {
|
|
if (ptr) {
|
|
::free(ptr);
|
|
}
|
|
reset();
|
|
}
|
|
|
|
// Free allocated heap memory but preserve size and node information.
|
|
void release() {
|
|
if (ptr) {
|
|
::free(ptr);
|
|
ptr = nullptr;
|
|
}
|
|
}
|
|
};
|
|
|
|
// A memory planner that makes all the allocations using malloc()/free().
|
|
//
|
|
// This is simple implementation of MemoryPlanner which uses malloc()/free()
|
|
// instead of memory areana. This planner is designed for AddressSanitizer.
|
|
class SimplePlanner : public MemoryPlanner {
|
|
public:
|
|
// Ownership of 'context' is not taken and it must remain util the
|
|
// ArenaPlanner is destroyed. The inputs to the graph will not share
|
|
// memory with any other tensor, effectively preserving them until the end
|
|
// of inference.
|
|
SimplePlanner(TfLiteContext* context, std::unique_ptr<GraphInfo> graph_info);
|
|
~SimplePlanner() override;
|
|
SimplePlanner(const SimplePlanner&) = delete;
|
|
SimplePlanner& operator=(const SimplePlanner&) = delete;
|
|
|
|
TfLiteStatus ResetAllocations() override;
|
|
TfLiteStatus ResetAllocationsAfter(int node) override;
|
|
TfLiteStatus PlanAllocations() override;
|
|
TfLiteStatus ExecuteAllocations(int first_node, int last_node) override;
|
|
TfLiteStatus ReleaseNonPersistentMemory() override;
|
|
TfLiteStatus AcquireNonPersistentMemory() override;
|
|
bool HasNonPersistentMemory() override { return true; };
|
|
void DumpDebugInfo(const std::vector<int>& execution_plan) const override{};
|
|
void GetAllocInfo(size_t* arena_size,
|
|
size_t* arena_persist_size) const override{};
|
|
|
|
private:
|
|
// Free all the all allocations.
|
|
void FreeAllAllocations();
|
|
|
|
// Assign absolute memory location to a tensor.
|
|
TfLiteStatus ResolveTensorAllocation(int tensor_index);
|
|
|
|
TfLiteContext* context_;
|
|
std::unique_ptr<GraphInfo> graph_info_;
|
|
|
|
// Stores allocation data for all tensors.
|
|
std::vector<SimpleAlloc> allocs_;
|
|
|
|
// First node, that uses the tensor. It needs to be allocated before
|
|
// execution of the node's operation.
|
|
std::vector<int32_t> alloc_node_;
|
|
|
|
// Last node, that uses the tensor. It can be deallocated after execution of
|
|
// the node's operation.
|
|
std::vector<int32_t> dealloc_node_;
|
|
};
|
|
|
|
} // namespace tflite
|
|
|
|
#endif // TENSORFLOW_LITE_SIMPLE_PLANNER_H_
|