1 | /* Copyright 2020 The TensorFlow Authors. All Rights Reserved. |
---|---|
2 | |
3 | Licensed under the Apache License, Version 2.0 (the "License"); |
4 | you may not use this file except in compliance with the License. |
5 | You may obtain a copy of the License at |
6 | |
7 | http://www.apache.org/licenses/LICENSE-2.0 |
8 | |
9 | Unless required by applicable law or agreed to in writing, software |
10 | distributed under the License is distributed on an "AS IS" BASIS, |
11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | See the License for the specific language governing permissions and |
13 | limitations under the License. |
14 | ==============================================================================*/ |
15 | #include "tensorflow/lite/tflite_with_xnnpack_optional.h" |
16 | |
17 | #include <memory> |
18 | |
19 | #include "tensorflow/lite/c/common.h" |
20 | #include "tensorflow/lite/core/macros.h" |
21 | |
22 | #ifdef TFLITE_BUILD_WITH_XNNPACK_DELEGATE |
23 | #include "tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h" |
24 | #endif |
25 | |
26 | namespace tflite { |
27 | |
28 | using TfLiteDelegatePtr = |
29 | std::unique_ptr<TfLiteDelegate, void (*)(TfLiteDelegate*)>; |
30 | |
31 | #ifdef TFLITE_BUILD_WITH_XNNPACK_DELEGATE |
32 | TfLiteDelegatePtr MaybeCreateXNNPACKDelegate(int num_threads) { |
33 | auto opts = TfLiteXNNPackDelegateOptionsDefault(); |
34 | // Note that we don't want to use the thread pool for num_threads == 1. |
35 | opts.num_threads = num_threads > 1 ? num_threads : 0; |
36 | return TfLiteDelegatePtr(TfLiteXNNPackDelegateCreate(&opts), |
37 | TfLiteXNNPackDelegateDelete); |
38 | } |
39 | #else |
40 | // Using weak symbols to create a delegate allows automatic injection of the |
41 | // delegate simply by adding it as a dependency. See the strong override in |
42 | // lite/tflite_with_xnnpack.cc, |
43 | TFLITE_ATTRIBUTE_WEAK TfLiteDelegatePtr |
44 | AcquireXNNPACKDelegate(int num_threads) { |
45 | return TfLiteDelegatePtr(nullptr, [](TfLiteDelegate*) {}); |
46 | } |
47 | |
48 | TfLiteDelegatePtr MaybeCreateXNNPACKDelegate(int num_threads) { |
49 | return AcquireXNNPACKDelegate(num_threads); |
50 | } |
51 | #endif |
52 | |
53 | } // namespace tflite |
54 |