1515#include < cinttypes>
1616#include < string>
1717#include < unordered_map>
18+ #include < unordered_set>
1819#include < vector>
1920
2021#pragma clang diagnostic ignored "-Wmissing-prototypes"
@@ -1989,6 +1990,58 @@ DefineNodeFunc getDefineNodeFunc(fb_xnnpack::XNodeUnion nodeType) {
19891990}
19901991#undef _DEFINE
19911992
1993+ /*
1994+ Serialized id of an xvalue, or XNN_INVALID_VALUE_ID if it is not a tensor.
1995+ */
1996+ uint32_t getSerializedValueId (ValuePtr value) noexcept {
1997+ const fb_xnnpack::XNNTensorValue* tensor_value = nullptr ;
1998+ if (value->xvalue_union_type () == fb_xnnpack::XValueUnion::XNNTensorValue) {
1999+ tensor_value = value->xvalue_union_as_XNNTensorValue ();
2000+ } else if (
2001+ value->xvalue_union_type () ==
2002+ fb_xnnpack::XValueUnion::XNNQuantizedTensorValue) {
2003+ tensor_value =
2004+ value->xvalue_union_as_XNNQuantizedTensorValue ()->tensor_value ();
2005+ }
2006+ return tensor_value != nullptr ? tensor_value->id_out ()
2007+ : XNN_INVALID_VALUE_ID ;
2008+ }
2009+
2010+ /*
2011+ Serialized ids of the values that XNNPACK copies into its own packed storage
2012+ while the runtime is created. Every other constant value is referenced by raw
2013+ pointer for the lifetime of the runtime.
2014+ */
2015+ std::unordered_set<uint32_t > getPackedValueIds (GraphPtr flatbuffer_graph) {
2016+ std::unordered_set<uint32_t > packed_ids;
2017+ auto insert_weights = [&packed_ids](uint32_t filter_id, uint32_t bias_id) {
2018+ packed_ids.insert (filter_id);
2019+ if (bias_id != XNN_INVALID_VALUE_ID ) {
2020+ packed_ids.insert (bias_id);
2021+ }
2022+ };
2023+
2024+ // Deliberately an if-chain rather than a switch: XNodeUnion has ~50
2025+ // enumerators and -Wswitch-enum requires every one of them to be listed.
2026+ for (auto node : *flatbuffer_graph->xnodes ()) {
2027+ auto type = node->xnode_union_type ();
2028+ if (type == fb_xnnpack::XNodeUnion::XNNFullyConnected) {
2029+ auto n = node->xnode_union_as_XNNFullyConnected ();
2030+ insert_weights (n->filter_id (), n->bias_id ());
2031+ } else if (type == fb_xnnpack::XNodeUnion::XNNConv2d) {
2032+ auto n = node->xnode_union_as_XNNConv2d ();
2033+ insert_weights (n->filter_id (), n->bias_id ());
2034+ } else if (type == fb_xnnpack::XNodeUnion::XNNDepthwiseConv2d) {
2035+ auto n = node->xnode_union_as_XNNDepthwiseConv2d ();
2036+ insert_weights (n->filter_id (), n->bias_id ());
2037+ } else if (type == fb_xnnpack::XNodeUnion::XNNConvTranspose2d) {
2038+ auto n = node->xnode_union_as_XNNConvTranspose2d ();
2039+ insert_weights (n->filter_id (), n->bias_id ());
2040+ }
2041+ }
2042+ return packed_ids;
2043+ }
2044+
19922045/*
19932046Builds the xnnpack runtime object using the buffer pointer. The buffer pointer
19942047must be a valid pointer to the serialized xnnpack object. It also fills the
@@ -2093,15 +2146,19 @@ ET_NODISCARD Error XNNCompiler::compileModel(
20932146 // Invalid ids do not need to be remapped
20942147 remapped_ids.emplace (XNN_INVALID_VALUE_ID , XNN_INVALID_VALUE_ID );
20952148
2096- // If weight cache is not on we hold onto all the unpacked buffers
2097- // and we free them at the end
2149+ // Buffers loaded from the named data map for values whose data XNNPACK packs
2150+ // during runtime creation. They stay alive until the runtime exists and are
2151+ // freed afterwards.
20982152 std::vector<FreeableBuffer> unpacked_buffers;
2153+ const std::unordered_set<uint32_t > packed_value_ids =
2154+ getPackedValueIds (flatbuffer_graph);
20992155
21002156 // External Ids for inputs and outputs
21012157 std::vector<uint32_t > input_ids;
21022158 std::vector<uint32_t > output_ids;
21032159 Error err = Error::Ok;
21042160 for (auto value : *flatbuffer_graph->xvalues ()) {
2161+ size_t prev_buffers = unpacked_buffers.size ();
21052162 err = defineTensor (
21062163 subgraph.get (),
21072164 remapped_ids,
@@ -2120,6 +2177,17 @@ ET_NODISCARD Error XNNCompiler::compileModel(
21202177 if (err != Error::Ok) {
21212178 return err;
21222179 }
2180+
2181+ // Operators that don't pack (PReLU, for example) keep raw pointers into
2182+ // the constant data, so hand their buffers to the executor. A single value
2183+ // can contribute more than one buffer: a quantized weight also loads its
2184+ // scales.
2185+ if (packed_value_ids.count (getSerializedValueId (value)) == 0 ) {
2186+ for (size_t i = prev_buffers; i < unpacked_buffers.size (); i++) {
2187+ executor->unpacked_buffers_ .push_back (std::move (unpacked_buffers[i]));
2188+ }
2189+ unpacked_buffers.resize (prev_buffers);
2190+ }
21232191 }
21242192
21252193 for (auto node : *flatbuffer_graph->xnodes ()) {
@@ -2174,6 +2242,7 @@ ET_NODISCARD Error XNNCompiler::compileModel(
21742242 " Failed to finalize weights cache after creating the xnn runtime" );
21752243 packed_weights_names = std::move (packed_weights_names_result.get ());
21762244 } else {
2245+ // XNNPACK has copied these into its own packed storage.
21772246 for (auto & buffer : unpacked_buffers) {
21782247 buffer.Free ();
21792248 }
0 commit comments