1 // Copyright 2020 Google LLC 2 // 3 // This source code is licensed under the BSD-style license found in the 4 // LICENSE file in the root directory of this source tree. 5 6 #pragma once 7 8 #include <stdbool.h> 9 #include <stddef.h> 10 #include <stdint.h> 11 12 #include <xnnpack.h> 13 #include <xnnpack/common.h> 14 #include <xnnpack/allocator.h> 15 16 #ifdef __cplusplus 17 extern "C" { 18 #endif 19 20 struct xnn_value_usage { 21 // The index (to xnn_subgraph_t->nodes) of the first xnn_node that uses this xnn_value. 22 uint32_t first_node; 23 // The index of the last xnn_node that uses this xnn_value. 24 uint32_t last_node; 25 // Note that 'tensor_size' includes the padding of XNN_EXTRA_BYTES. 26 size_t tensor_size; 27 // The memory offset of this xnn_value from the beginning of a memory buffer. 28 size_t alloc_offset; 29 }; 30 31 // Track the memory allocation in a memory arena for a subgraph. 32 struct xnn_value_allocation_tracker { 33 xnn_subgraph_t subgraph; 34 size_t mem_arena_size; 35 // Representing the lifecycle of xnn_values in the 'subgraph', and the array size is 'subgraph->num_values'. 36 struct xnn_value_usage* usage; 37 // The range of value ids (i.e. the index to subgraph->values) whose memory might need to be allocated. 38 size_t min_value_id; 39 size_t max_value_id; 40 }; 41 42 // Initialize the memory allocation tracker for xnn_values. 43 XNN_INTERNAL void xnn_init_value_allocation_tracker(struct xnn_value_allocation_tracker* tracker, 44 const xnn_subgraph_t subgraph); 45 xnn_release_value_allocation_tracker(struct xnn_value_allocation_tracker * tracker)46inline static void xnn_release_value_allocation_tracker(struct xnn_value_allocation_tracker* tracker) { 47 xnn_release_memory(tracker->usage); 48 } 49 50 // Add a to-be-allocated xnn_value (referred by 'value_id') of size 'tensor_size' to the allocation tracker. 51 // Note: this function assumes 'value_id's added in increasing order for simplicity as it's called inside a loop 52 // iterating over 'subgraph->values'. 53 XNN_INTERNAL void xnn_add_value_allocation_tracker(struct xnn_value_allocation_tracker* tracker, 54 uint32_t value_id, size_t tensor_size); 55 56 // Plan the exact the memory allocation for intermediate tensors according to the xnn_value allocation tracker. 57 XNN_INTERNAL void xnn_plan_value_allocation_tracker(struct xnn_value_allocation_tracker* tracker); 58 59 #ifdef __cplusplus 60 } // extern "C" 61 #endif 62