| 20 | namespace tflite { |
| 21 | |
| 22 | MicroAllocator::MicroAllocator(TfLiteContext* context, const Model* model, |
| 23 | uint8_t* tensor_arena, size_t arena_size, |
| 24 | ErrorReporter* error_reporter) |
| 25 | : model_(model), |
| 26 | tensor_allocator_(tensor_arena, arena_size), |
| 27 | error_reporter_(error_reporter), |
| 28 | context_(context) { |
| 29 | auto* subgraphs = model->subgraphs(); |
| 30 | if (subgraphs->size() != 1) { |
| 31 | error_reporter->Report("Only 1 subgraph is currently supported.\n"); |
| 32 | return; |
| 33 | } |
| 34 | subgraph_ = (*subgraphs)[0]; |
| 35 | tensors_ = subgraph_->tensors(); |
| 36 | operators_ = subgraph_->operators(); |
| 37 | |
| 38 | context_->tensors_size = tensors_->size(); |
| 39 | context_->tensors = |
| 40 | reinterpret_cast<TfLiteTensor*>(tensor_allocator_.AllocateMemory( |
| 41 | sizeof(TfLiteTensor) * context_->tensors_size, 4)); |
| 42 | |
| 43 | // Null all inputs so we can later perform a null check to avoid re-allocating |
| 44 | // registered pre-allocated inputs. |
| 45 | for (size_t i = 0; i < subgraph_->inputs()->size(); ++i) { |
| 46 | const int tensor_index = subgraph_->inputs()->Get(i); |
| 47 | context_->tensors[tensor_index].data.raw = nullptr; |
| 48 | } |
| 49 | } |
| 50 | |
| 51 | TfLiteStatus MicroAllocator::RegisterPreallocatedInput(uint8_t* buffer, |
| 52 | size_t input_index) { |