diff options
author | 2018-03-26 11:47:50 -0700 | |
---|---|---|
committer | 2018-03-26 11:50:29 -0700 | |
commit | d2604f8dcb8a63ca063f712c24ce5aa63403b0aa (patch) | |
tree | 6109456e2238ab20f647639aa8f05c0aba5b128d /tensorflow/contrib/lite/kernels/gemm_support.cc | |
parent | 6d46c21e9f300d07e30a2185671f07d34fac3999 (diff) |
Revert to initializing number of threads when SetNumThreads is called. Requiring it
to happen before OpInit() is way too confusing for users.
PiperOrigin-RevId: 190499644
Diffstat (limited to 'tensorflow/contrib/lite/kernels/gemm_support.cc')
-rw-r--r-- | tensorflow/contrib/lite/kernels/gemm_support.cc | 6 |
1 files changed, 6 insertions, 0 deletions
diff --git a/tensorflow/contrib/lite/kernels/gemm_support.cc b/tensorflow/contrib/lite/kernels/gemm_support.cc index 76a5165d14..95f45ea768 100644 --- a/tensorflow/contrib/lite/kernels/gemm_support.cc +++ b/tensorflow/contrib/lite/kernels/gemm_support.cc @@ -61,5 +61,11 @@ gemmlowp::GemmContext* GetFromContext(TfLiteContext* context) { return ptr->gemm_context_; } +void SetNumThreads(TfLiteContext* context, int num_threads) { + IncrementUsageCounter(context); + GetFromContext(context)->set_max_num_threads(num_threads); + DecrementUsageCounter(context); +} + } // namespace gemm_support } // namespace tflite |