Adapt TFLite to new ruy::Context API.
PiperOrigin-RevId: 307889520
Change-Id: I20da08901e119599dcc8ff2f52204c59ab1646fa
diff --git a/tensorflow/lite/kernels/cpu_backend_context.cc b/tensorflow/lite/kernels/cpu_backend_context.cc
index 5128421..0fa4175 100644
--- a/tensorflow/lite/kernels/cpu_backend_context.cc
+++ b/tensorflow/lite/kernels/cpu_backend_context.cc
@@ -56,7 +56,7 @@
gemmlowp_context_(new gemmlowp::GemmContext) {
SetMaxNumThreads(kDefaultNumThreadpoolThreads);
#ifdef TFLITE_WITH_RUY_GEMV
- ruy_context_->cache_policy = ruy::kCacheLHSOnNarrowMul;
+ ruy_context_->set_cache_policy(ruy::CachePolicy::kCacheLHSOnNarrowMul);
#endif
}
@@ -66,7 +66,7 @@
const int target_num_threads =
max_num_threads > -1 ? max_num_threads : kDefaultNumThreadpoolThreads;
max_num_threads_ = target_num_threads;
- ruy_context_->max_num_threads = target_num_threads;
+ ruy_context_->set_max_num_threads(target_num_threads);
gemmlowp_context_->set_max_num_threads(target_num_threads);
}
diff --git a/tensorflow/lite/kernels/cpu_backend_threadpool.h b/tensorflow/lite/kernels/cpu_backend_threadpool.h
index ff03d37..39eafd5 100644
--- a/tensorflow/lite/kernels/cpu_backend_threadpool.h
+++ b/tensorflow/lite/kernels/cpu_backend_threadpool.h
@@ -37,7 +37,8 @@
void Execute(int tasks_count, TaskType* tasks,
CpuBackendContext* cpu_backend_context) {
TFLITE_DCHECK_LE(tasks_count, cpu_backend_context->max_num_threads());
- cpu_backend_context->ruy_context()->workers_pool.Execute(tasks_count, tasks);
+ cpu_backend_context->ruy_context()->mutable_thread_pool()->Execute(
+ tasks_count, tasks);
}
#else // not TFLITE_WITH_RUY
