micro_allocator.h

00001 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
00002 
00003 Licensed under the Apache License, Version 2.0 (the "License");
00004 you may not use this file except in compliance with the License.
00005 You may obtain a copy of the License at
00006 
00007     http://www.apache.org/licenses/LICENSE-2.0
00008 
00009 Unless required by applicable law or agreed to in writing, software
00010 distributed under the License is distributed on an "AS IS" BASIS,
00011 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
00012 See the License for the specific language governing permissions and
00013 limitations under the License.
00014 ==============================================================================*/
00015 #ifndef TENSORFLOW_LITE_EXPERIMENTAL_MICRO_MICRO_ALLOCATOR_H_
00016 #define TENSORFLOW_LITE_EXPERIMENTAL_MICRO_MICRO_ALLOCATOR_H_
00017 
00018 #include "tensorflow/lite/c/c_api_internal.h"
00019 #include "tensorflow/lite/core/api/error_reporter.h"
00020 #include "tensorflow/lite/core/api/flatbuffer_conversions.h"
00021 #include "tensorflow/lite/experimental/micro/simple_memory_allocator.h"
00022 #include "tensorflow/lite/schema/schema_generated.h"
00023 
00024 namespace tflite {
00025 
00026 typedef struct {
00027   TfLiteNode node;
00028   const TfLiteRegistration* registration;
00029 } NodeAndRegistration;
00030 
00031 // Allocator responsible for allocating memory for all intermediate tensors
00032 // necessary to invoke a model.
00033 class MicroAllocator {
00034  public:
00035   // The lifetime of the model, tensor allocator and error reporter must be at
00036   // least as long as that of the allocator object, since the allocator needs
00037   // them to be accessible during its entire lifetime.
00038   MicroAllocator(TfLiteContext* context, const Model* model,
00039                  uint8_t* tensor_arena, size_t arena_size,
00040                  ErrorReporter* error_reporter);
00041 
00042   // Specify a particular tensor as pre-allocated.  This means that this tensor
00043   // will internally point to the supplied buffer, and no new memory will be
00044   // provided.  The buffer must live at least as long as the allocator, since
00045   // the buffer will be used every time an op is invoked which uses the
00046   // specified tensor.  Most commonly this is useful when a platform-provided
00047   // DMA buffer is used as an input, and it is desirable to avoid unnecessarily
00048   // allocating a new buffer and copying from the DMA buffer. The user must
00049   // ensure the buffer is valid throughout each interpreter run, and is not
00050   // prematurely overwritten.
00051   TfLiteStatus RegisterPreallocatedInput(uint8_t* buffer, size_t input_index);
00052 
00053   // Sets up all of the data structure members for a runtime tensor based on the
00054   // contents of a serialized tensor. This method doesn't allocate any memory,
00055   // all allocations happen subsequently in AllocateTensors.
00056   TfLiteStatus InitializeRuntimeTensor(
00057       const tflite::Tensor& flatbuffer_tensor,
00058       const flatbuffers::Vector<flatbuffers::Offset<Buffer>>* buffers,
00059       ErrorReporter* error_reporter, TfLiteTensor* result,
00060       uint8_t* preallocated_buffer = nullptr);
00061 
00062   // Run through the model and allocate all necessary input, output and
00063   // intermediate tensors except for those already provided via calls to
00064   // registerPreallocatedInput.
00065   // WARNING: doing any allocation after calling is method has the risk of
00066   // corruption tensor data so this method is the last method to be called in
00067   // this class.
00068   TfLiteStatus FinishTensorAllocation();
00069 
00070   // Run through the model to allocate nodes and registrations. We need to keep
00071   // them for the entire life time of the model to allow persistent tensors.
00072   // This method needs to be called before FinishTensorAllocation method.
00073   TfLiteStatus AllocateNodeAndRegistrations(
00074       const OpResolver& op_resolver,
00075       NodeAndRegistration** node_and_registrations);
00076 
00077  private:
00078   const Model* model_;
00079   SimpleMemoryAllocator memory_allocator_;
00080   ErrorReporter* error_reporter_;
00081   TfLiteContext* context_;
00082   uint8_t* arena_;
00083   size_t arena_size_;
00084   // Indicating if the allocator is ready for allocation.
00085   bool active_ = false;
00086 
00087   const SubGraph* subgraph_;
00088   const flatbuffers::Vector<flatbuffers::Offset<Operator>>* operators_;
00089   const flatbuffers::Vector<flatbuffers::Offset<Tensor>>* tensors_;
00090 };
00091 
00092 }  // namespace tflite
00093 #endif  // TENSORFLOW_LITE_EXPERIMENTAL_MICRO_MICRO_ALLOCATOR_H_