Adding autotune to the Cudnn conv algorithm selection.

For now, use TF_CUDNN_USE_AUTOTUNE=1 to enable this feature. Once it is mature
enough, it will be turned on by default.
Support for the backward steps will be added later.
Change: 121769364

1 files changed, 9 insertions, 0 deletions

diff --git a/tensorflow/stream_executor/stream_executor_pimpl.cc b/tensorflow/stream_executor/stream_executor_pimpl.cc
index acaa0efcb2..fe32039d71 100644
--- a/tensorflow/stream_executor/stream_executor_pimpl.cc
+++ b/tensorflow/stream_executor/stream_executor_pimpl.cc
@@ -286,6 +286,15 @@ bool StreamExecutor::SupportsDnn() const {
   return implementation_->SupportsDnn();
 }
 
+bool StreamExecutor::GetConvolveAlgorithms(
+    std::vector<dnn::AlgorithmType> *out_algorithms) {
+  dnn::DnnSupport *dnn_support = AsDnn();
+  if (!dnn_support) {
+    return false;
+  }
+  return dnn_support->GetConvolveAlgorithms(out_algorithms);
+}
+
 dnn::DnnSupport *StreamExecutor::AsDnn() {
   mutex_lock lock{mu_};
   if (dnn_ != nullptr) {