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"
@@ -2064,6 +2065,58 @@ DefineNodeFunc getDefineNodeFunc(fb_xnnpack::XNodeUnion nodeType) {
20642065}
20652066#undef _DEFINE
20662067
2068+ /*
2069+ Serialized id of an xvalue, or XNN_INVALID_VALUE_ID if it is not a tensor.
2070+ */
2071+ uint32_t getSerializedValueId (ValuePtr value) noexcept {
2072+ const fb_xnnpack::XNNTensorValue* tensor_value = nullptr ;
2073+ if (value->xvalue_union_type () == fb_xnnpack::XValueUnion::XNNTensorValue) {
2074+ tensor_value = value->xvalue_union_as_XNNTensorValue ();
2075+ } else if (
2076+ value->xvalue_union_type () ==
2077+ fb_xnnpack::XValueUnion::XNNQuantizedTensorValue) {
2078+ tensor_value =
2079+ value->xvalue_union_as_XNNQuantizedTensorValue ()->tensor_value ();
2080+ }
2081+ return tensor_value != nullptr ? tensor_value->id_out ()
2082+ : XNN_INVALID_VALUE_ID ;
2083+ }
2084+
2085+ /*
2086+ Serialized ids of the values that XNNPACK copies into its own packed storage
2087+ while the runtime is created. Every other constant value is referenced by raw
2088+ pointer for the lifetime of the runtime.
2089+ */
2090+ std::unordered_set<uint32_t > getPackedValueIds (GraphPtr flatbuffer_graph) {
2091+ std::unordered_set<uint32_t > packed_ids;
2092+ auto insert_weights = [&packed_ids](uint32_t filter_id, uint32_t bias_id) {
2093+ packed_ids.insert (filter_id);
2094+ if (bias_id != XNN_INVALID_VALUE_ID ) {
2095+ packed_ids.insert (bias_id);
2096+ }
2097+ };
2098+
2099+ // Deliberately an if-chain rather than a switch: XNodeUnion has ~50
2100+ // enumerators and -Wswitch-enum requires every one of them to be listed.
2101+ for (auto node : *flatbuffer_graph->xnodes ()) {
2102+ auto type = node->xnode_union_type ();
2103+ if (type == fb_xnnpack::XNodeUnion::XNNFullyConnected) {
2104+ auto n = node->xnode_union_as_XNNFullyConnected ();
2105+ insert_weights (n->filter_id (), n->bias_id ());
2106+ } else if (type == fb_xnnpack::XNodeUnion::XNNConv2d) {
2107+ auto n = node->xnode_union_as_XNNConv2d ();
2108+ insert_weights (n->filter_id (), n->bias_id ());
2109+ } else if (type == fb_xnnpack::XNodeUnion::XNNDepthwiseConv2d) {
2110+ auto n = node->xnode_union_as_XNNDepthwiseConv2d ();
2111+ insert_weights (n->filter_id (), n->bias_id ());
2112+ } else if (type == fb_xnnpack::XNodeUnion::XNNConvTranspose2d) {
2113+ auto n = node->xnode_union_as_XNNConvTranspose2d ();
2114+ insert_weights (n->filter_id (), n->bias_id ());
2115+ }
2116+ }
2117+ return packed_ids;
2118+ }
2119+
20672120/*
20682121Builds the xnnpack runtime object using the buffer pointer. The buffer pointer
20692122must be a valid pointer to the serialized xnnpack object. It also fills the
@@ -2168,15 +2221,19 @@ ET_NODISCARD Error XNNCompiler::compileModel(
21682221 // Invalid ids do not need to be remapped
21692222 remapped_ids.emplace (XNN_INVALID_VALUE_ID , XNN_INVALID_VALUE_ID );
21702223
2171- // If weight cache is not on we hold onto all the unpacked buffers
2172- // and we free them at the end
2224+ // Buffers loaded from the named data map for values whose data XNNPACK packs
2225+ // during runtime creation. They stay alive until the runtime exists and are
2226+ // freed afterwards.
21732227 std::vector<FreeableBuffer> unpacked_buffers;
2228+ const std::unordered_set<uint32_t > packed_value_ids =
2229+ getPackedValueIds (flatbuffer_graph);
21742230
21752231 // External Ids for inputs and outputs
21762232 std::vector<uint32_t > input_ids;
21772233 std::vector<uint32_t > output_ids;
21782234 Error err = Error::Ok;
21792235 for (auto value : *flatbuffer_graph->xvalues ()) {
2236+ size_t prev_buffers = unpacked_buffers.size ();
21802237 err = defineTensor (
21812238 subgraph.get (),
21822239 remapped_ids,
@@ -2195,6 +2252,17 @@ ET_NODISCARD Error XNNCompiler::compileModel(
21952252 if (err != Error::Ok) {
21962253 return err;
21972254 }
2255+
2256+ // Operators that don't pack (PReLU, for example) keep raw pointers into
2257+ // the constant data, so hand their buffers to the executor. A single value
2258+ // can contribute more than one buffer: a quantized weight also loads its
2259+ // scales.
2260+ if (packed_value_ids.count (getSerializedValueId (value)) == 0 ) {
2261+ for (size_t i = prev_buffers; i < unpacked_buffers.size (); i++) {
2262+ executor->unpacked_buffers_ .push_back (std::move (unpacked_buffers[i]));
2263+ }
2264+ unpacked_buffers.resize (prev_buffers);
2265+ }
21982266 }
21992267
22002268 for (auto node : *flatbuffer_graph->xnodes ()) {
@@ -2249,6 +2317,7 @@ ET_NODISCARD Error XNNCompiler::compileModel(
22492317 " Failed to finalize weights cache after creating the xnn runtime" );
22502318 packed_weights_names = std::move (packed_weights_names_result.get ());
22512319 } else {
2320+ // XNNPACK has copied these into its own packed storage.
22522321 for (auto & buffer : unpacked_buffers) {
22532322 buffer.Free ();
22542323 }
0 commit comments