1/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations 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
29namespace tflite {
30namespace cpu_backend_threadpool {
31
32#ifdef TFLITE_WITH_RUY
33
34using Task = ruy::Task;
35
36template <typename TaskType>
37void 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
46using Task = gemmlowp::Task;
47
48template <typename TaskType>
49void 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