1 /* Copyright 2019 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
16 #ifndef TENSORFLOW_LITE_KERNELS_CPU_BACKEND_THREADPOOL_H_
17 #define TENSORFLOW_LITE_KERNELS_CPU_BACKEND_THREADPOOL_H_
18
19 #include "tensorflow/lite/kernels/cpu_backend_context.h"
20 #include "tensorflow/lite/kernels/internal/compatibility.h"
21
22 #ifdef TFLITE_WITH_RUY
23 #include "ruy/context.h" // from @ruy
24 #include "ruy/thread_pool.h" // from @ruy
25 #else
26 #include "public/gemmlowp.h"
27 #endif
28
29 namespace tflite {
30 namespace cpu_backend_threadpool {
31
32 #ifdef TFLITE_WITH_RUY
33
34 using Task = ruy::Task;
35
36 template <typename TaskType>
Execute(int tasks_count,TaskType * tasks,CpuBackendContext * cpu_backend_context)37 void Execute(int tasks_count, TaskType* tasks,
38 CpuBackendContext* cpu_backend_context) {
39 TFLITE_DCHECK_LE(tasks_count, cpu_backend_context->max_num_threads());
40 cpu_backend_context->ruy_context()->mutable_thread_pool()->Execute(
41 tasks_count, tasks);
42 }
43
44 #else // not TFLITE_WITH_RUY
45
46 using Task = gemmlowp::Task;
47
48 template <typename TaskType>
49 void Execute(int tasks_count, TaskType* tasks,
50 CpuBackendContext* cpu_backend_context) {
51 TFLITE_DCHECK_LE(tasks_count, cpu_backend_context->max_num_threads());
52 cpu_backend_context->gemmlowp_context()->workers_pool()->Execute(tasks_count,
53 tasks);
54 }
55
56 #endif
57
58 } // namespace cpu_backend_threadpool
59 } // namespace tflite
60
61 #endif // TENSORFLOW_LITE_KERNELS_CPU_BACKEND_THREADPOOL_H_
62