1 // Copyright 2022 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 <stddef.h> // For size_t. 9 #include <stdint.h> // For uint32_t. 10 #include <xnnpack.h> // For xnn_status. 11 #include <xnnpack/allocator.h> // For xnn_code_buffer. 12 #include <xnnpack/mutex.h> // For xnn_mutex. 13 14 #ifdef __cplusplus 15 extern "C" { 16 #endif 17 18 #define XNN_CACHE_NOT_FOUND SIZE_MAX // Return value when code is not found in the cache. 19 20 // A cache for arbitrary bytes. 21 // The implementation is similar to a hash table with open addressing and linear 22 // probing, but restricted to our use cases. 23 24 // Similar to buckets in a hash table implementation, this is an entry in the 25 // cache. It stores "metadata" about the generated code (size and offset). The 26 // actual bytes are in the cache's buffer. 27 struct xnn_cache_bucket { 28 // A hash for quick comparison. 29 uint32_t hash; 30 // Size of bytes. 31 size_t size; 32 // Offset of bytes, relative to cache's buffer. 33 size_t offset; 34 }; 35 36 enum xnn_cache_type { 37 xnn_cache_type_invalid = 0, 38 xnn_cache_type_code, 39 xnn_cache_type_weights, 40 }; 41 42 struct xnn_cache { 43 enum xnn_cache_type type; 44 // A growing buffer that is used to keep all generated code or repacked weights. 45 union { 46 struct xnn_code_buffer code; 47 struct xnn_weights_buffer weights; 48 }; 49 50 // Entries in the cache. 51 struct xnn_cache_bucket* buckets; 52 // Capacity of the cache, when the load factor (num_entries/num_buckets) grows 53 // beyond a limit, the cache is expanded. 54 size_t num_buckets; 55 size_t num_entries; 56 size_t hits; 57 size_t misses; 58 }; 59 60 // A cache for JIT generated microkernel code. 61 struct xnn_code_cache { 62 struct xnn_cache cache; 63 }; 64 65 enum xnn_status xnn_init_code_cache(struct xnn_code_cache* cache); 66 enum xnn_status xnn_release_code_cache(struct xnn_code_cache* cache); 67 // Looks up `ptr` in the cache, returns offset into cache's buffer if found. 68 // `ptr` should already point into cache->buffer. 69 // If it already exists within the cache, the buffer will be rewound, so we can 70 // reuse the same section of the buffer. 71 size_t xnn_get_or_insert_code_cache(struct xnn_code_cache* cache, void* ptr, size_t size); 72 73 // The state of weights cache finalization. 74 enum xnn_cache_state { 75 // Not finalized. 76 xnn_cache_state_not_finalized, 77 // The underlying memory is trimmed to be as compact as possible. 78 xnn_cache_state_hard_finalized, 79 // The underlying memory has some extra space at the end. 80 xnn_cache_state_soft_finalized, 81 }; 82 83 // A cache for repacked weights. 84 struct xnn_weights_cache { 85 struct xnn_cache cache; 86 // Protects updates of `cache`, it has the same lifetime as `cache`, and so should be initialized/destroyed together 87 // with the `cache`. 88 struct xnn_mutex mutex; 89 // Maximum size of packed weights that have been inserted into the cache. 90 size_t max_weights_size; 91 enum xnn_cache_state finalization_state; 92 }; 93 94 enum xnn_status xnn_init_weights_cache(struct xnn_weights_cache* cache); 95 enum xnn_status xnn_init_weights_cache_with_size(struct xnn_weights_cache* cache, size_t size); 96 // Finalizes the weights cache, so that we cannot insert any more entries into the cache. 97 enum xnn_status xnn_finalize_weights_cache( 98 struct xnn_weights_cache* cache, 99 enum xnn_weights_cache_finalization_kind finalization_kind); 100 enum xnn_status xnn_release_weights_cache(struct xnn_weights_cache* cache); 101 // Ensures that cache has enough space for `n` bytes, locks the mutex to protect future updates. Mutex must be unlocked 102 // using xnn_get_or_insert_weights_cache. 103 void* xnn_reserve_space_in_weights_cache(struct xnn_weights_cache* cache, size_t n); 104 // Looks up packed weights at `ptr` in the cache. If it is found, reuse it. Otherwise, it is added to the cache. Mutex 105 // must already be locked before calling this, it will be unlocked at the end of this function. 106 size_t xnn_get_or_insert_weights_cache(struct xnn_weights_cache* cache, void* ptr, size_t size); 107 bool xnn_weights_cache_is_finalized(struct xnn_weights_cache* cache); 108 109 struct xnn_caches { 110 struct xnn_code_cache *code_cache; 111 struct xnn_weights_cache *weights_cache; 112 }; 113 114 #ifdef __cplusplus 115 } // extern "C" 116 #endif 117