|
| 1 | +// Copyright 2025 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 | +#include "src/operators/fingerprint_cache.h" |
| 7 | + |
| 8 | +#include <assert.h> |
| 9 | +#include <stddef.h> |
| 10 | +#include <stdint.h> |
| 11 | +#include <stdlib.h> |
| 12 | + |
| 13 | +#include "include/experimental.h" |
| 14 | +#include "include/xnnpack.h" |
| 15 | +#include "src/operators/fingerprint_id.h" |
| 16 | +#include "src/xnnpack/allocator.h" |
| 17 | +#include "src/xnnpack/cache.h" |
| 18 | +#include "src/xnnpack/init-once.h" |
| 19 | +#include "src/xnnpack/mutex.h" |
| 20 | + |
| 21 | +#define XNN_FINGERPRINT_MAX_COUNT 256 |
| 22 | +static struct xnn_fingerprint fingerprint_vector[XNN_FINGERPRINT_MAX_COUNT]; |
| 23 | +static uint32_t fingerprint_vector_size = 0; |
| 24 | + |
| 25 | +static struct xnn_mutex mutex; |
| 26 | +XNN_INIT_ONCE_GUARD(mutex); |
| 27 | + |
| 28 | +static void init_mutex_config() { xnn_mutex_init(&mutex); } |
| 29 | + |
| 30 | +const struct xnn_fingerprint* xnn_get_fingerprint( |
| 31 | + const enum xnn_fingerprint_id id) { |
| 32 | + XNN_INIT_ONCE(mutex); |
| 33 | + uint32_t i = 0; |
| 34 | + xnn_mutex_lock(&mutex); |
| 35 | + for (; i < fingerprint_vector_size; ++i) { |
| 36 | + if (fingerprint_vector[i].id == id) { |
| 37 | + break; |
| 38 | + } |
| 39 | + } |
| 40 | + xnn_mutex_unlock(&mutex); |
| 41 | + return i < fingerprint_vector_size ? fingerprint_vector + i : NULL; |
| 42 | +} |
| 43 | + |
| 44 | +void xnn_set_fingerprint(const struct xnn_fingerprint fingerprint) { |
| 45 | + XNN_INIT_ONCE(mutex); |
| 46 | + uint32_t i = 0; |
| 47 | + xnn_mutex_lock(&mutex); |
| 48 | + for (; i < fingerprint_vector_size; ++i) { |
| 49 | + if (fingerprint_vector[i].id == fingerprint.id) { |
| 50 | + fingerprint_vector[i] = fingerprint; |
| 51 | + } |
| 52 | + } |
| 53 | + xnn_mutex_unlock(&mutex); |
| 54 | + if (i < fingerprint_vector_size) { |
| 55 | + return; |
| 56 | + } |
| 57 | + assert(fingerprint_vector_size < XNN_FINGERPRINT_MAX_COUNT); |
| 58 | + fingerprint_vector[fingerprint_vector_size++] = fingerprint; |
| 59 | +} |
| 60 | + |
| 61 | +void xnn_clear_fingerprints() { |
| 62 | + XNN_INIT_ONCE(mutex); |
| 63 | + xnn_mutex_lock(&mutex); |
| 64 | + fingerprint_vector_size = 0; |
| 65 | + xnn_mutex_unlock(&mutex); |
| 66 | +} |
| 67 | + |
| 68 | +// The context for an XNNPack weight cache provider that we pass to operator |
| 69 | +// `create` functions when we want to fingerprint them. |
| 70 | +struct fingerprint_cache_context { |
| 71 | + void* buffer; |
| 72 | + size_t bytes; |
| 73 | + uint32_t hash; |
| 74 | +}; |
| 75 | + |
| 76 | +static size_t fingerprint_cache_look_up( |
| 77 | + void* context, const struct xnn_weights_cache_look_up_key* cache_key) { |
| 78 | + return XNN_CACHE_NOT_FOUND; |
| 79 | +} |
| 80 | + |
| 81 | +static void* fingerprint_cache_reserve_space(void* const context, size_t n) { |
| 82 | + struct fingerprint_cache_context* const ctx = context; |
| 83 | + assert(ctx); |
| 84 | + if (ctx->buffer && ctx->bytes < n) { |
| 85 | + xnn_release_simd_memory(ctx->buffer); |
| 86 | + ctx->buffer = NULL; |
| 87 | + } |
| 88 | + if (ctx->buffer == NULL) { |
| 89 | + ctx->buffer = xnn_allocate_simd_memory(n); |
| 90 | + ctx->bytes = ctx->buffer ? n : 0; |
| 91 | + } |
| 92 | + return ctx->buffer; |
| 93 | +} |
| 94 | + |
| 95 | +static size_t fingerprint_cache_look_up_or_insert( |
| 96 | + void* context, const struct xnn_weights_cache_look_up_key* cache_key, |
| 97 | + void* ptr, size_t size) { |
| 98 | + assert(context); |
| 99 | + struct fingerprint_cache_context* const ctx = context; |
| 100 | + ctx->hash = murmur_hash3(ptr, size, /*seed=*/ctx->hash); |
| 101 | + return 0; |
| 102 | +} |
| 103 | + |
| 104 | +static bool fingerprint_cache_is_finalized(void* context) { return false; } |
| 105 | + |
| 106 | +static void* fingerprint_cache_offset_to_addr(void* context, size_t offset) { |
| 107 | + assert(context); |
| 108 | + struct fingerprint_cache_context* const ctx = context; |
| 109 | + return ctx->buffer; |
| 110 | +} |
| 111 | + |
| 112 | +static enum xnn_status fingerprint_cache_delete_cache(void* context) { |
| 113 | + struct fingerprint_cache_context* const ctx = context; |
| 114 | + if (ctx) { |
| 115 | + if (ctx->buffer) { |
| 116 | + xnn_release_simd_memory(ctx->buffer); |
| 117 | + } |
| 118 | + *ctx = (struct fingerprint_cache_context){0}; |
| 119 | + } |
| 120 | + return xnn_status_success; |
| 121 | +} |
| 122 | + |
| 123 | +struct fingerprint_context create_fingerprint_context( |
| 124 | + const enum xnn_fingerprint_id fingerprint_id) { |
| 125 | + struct fingerprint_context context = { |
| 126 | + .status = xnn_status_uninitialized, |
| 127 | + .fingerprint_id = fingerprint_id, |
| 128 | + .cache = |
| 129 | + (struct xnn_weights_cache_provider){ |
| 130 | + .context = NULL, |
| 131 | + .look_up = fingerprint_cache_look_up, |
| 132 | + .reserve_space = fingerprint_cache_reserve_space, |
| 133 | + .look_up_or_insert = fingerprint_cache_look_up_or_insert, |
| 134 | + .is_finalized = fingerprint_cache_is_finalized, |
| 135 | + .offset_to_addr = fingerprint_cache_offset_to_addr, |
| 136 | + .delete_cache = fingerprint_cache_delete_cache}, |
| 137 | + .op = NULL, |
| 138 | + }; |
| 139 | + if (context.fingerprint_id == xnn_fingerprint_id_unknown) { |
| 140 | + context.status = xnn_status_unsupported_parameter; |
| 141 | + } else if (xnn_get_fingerprint(context.fingerprint_id)) { |
| 142 | + context.status = xnn_status_success; |
| 143 | + } else { |
| 144 | + // Do this after the checks to avoid a memory allocation when unnecessary. |
| 145 | + context.cache.context = xnn_allocate_zero_memory(sizeof(struct fingerprint_cache_context)); |
| 146 | + } |
| 147 | + return context; |
| 148 | +} |
| 149 | + |
| 150 | +static void free_fingerprint_cache_provider( |
| 151 | + struct xnn_weights_cache_provider* const provider) { |
| 152 | + if (provider) { |
| 153 | + provider->delete_cache(provider->context); |
| 154 | + if (provider->context) { |
| 155 | + xnn_release_memory(provider->context); |
| 156 | + } |
| 157 | + } |
| 158 | +} |
| 159 | + |
| 160 | +void finalize_fingerprint_context(struct fingerprint_context* const context) { |
| 161 | + assert(context); |
| 162 | + if (context->status == xnn_status_uninitialized) { |
| 163 | + xnn_set_fingerprint((struct xnn_fingerprint){ |
| 164 | + .id = context->fingerprint_id, |
| 165 | + .value = fingerprint_cache_get_fingerprint(&context->cache)}); |
| 166 | + } |
| 167 | + free_fingerprint_cache_provider(&context->cache); |
| 168 | + if (context->op) { |
| 169 | + xnn_delete_operator(context->op); |
| 170 | + } |
| 171 | +} |
| 172 | + |
| 173 | +uint32_t fingerprint_cache_get_fingerprint( |
| 174 | + const struct xnn_weights_cache_provider* const provider) { |
| 175 | + assert(provider); |
| 176 | + assert(provider->context); |
| 177 | + return ((struct fingerprint_cache_context*)provider->context)->hash; |
| 178 | +} |
0 commit comments