miopen tune api for batchnorm

Author: dnikolaev-amd

aten/src/ATen/native/miopen/BatchNorm_miopen.cpp

@@ -115,11 +115,18 @@ std::tuple<Tensor, Tensor, Tensor> miopen_batch_norm(
   Constant zero(dataType, 0);
   Tensor save_mean, save_var;
 
+  auto tuning_policy = at::globalContext().benchmarkCuDNN() ? miopenTuningPolicySearch : miopenTuningPolicyNone;
+  miopenTuningPolicy_t previous_policy;
+  MIOPEN_CHECK(miopenGetTuningPolicy(handle, &previous_policy));
+  if (tuning_policy != previous_policy) {
+    MIOPEN_CHECK(miopenSetTuningPolicy(handle, tuning_policy));
+  }
+
   if (training) {
     int64_t num_features = input_t.size(1);
     save_mean = at::empty({ num_features }, weight_t.options());
     save_var = at::empty({ num_features }, weight_t.options());
-    MIOPEN_CHECK(miopenBatchNormalizationForwardTraining(
+    auto status = miopenBatchNormalizationForwardTraining(
       handle, mode, &one, &zero,
       idesc.desc(), input->const_data_ptr(),
       idesc.desc(), output->data_ptr(),
@@ -134,11 +141,15 @@ std::tuple<Tensor, Tensor, Tensor> miopen_batch_norm(
       at::maybe_data_ptr(running_var),
       epsilon,
       save_mean.mutable_data_ptr(),
-      save_var.mutable_data_ptr()));
+      save_var.mutable_data_ptr());
+    if (tuning_policy != previous_policy) {
+      MIOPEN_CHECK(miopenSetTuningPolicy(handle, previous_policy));
+    }
+    MIOPEN_CHECK(status);
   } else {
     save_mean = at::empty({0}, weight_t.options());
     save_var = at::empty({0}, weight_t.options());
-    MIOPEN_CHECK(miopenBatchNormalizationForwardInference(
+    auto status = miopenBatchNormalizationForwardInference(
       handle, mode, &one, &zero,
       idesc.desc(), input->const_data_ptr(),
       idesc.desc(), output->data_ptr(),
@@ -150,7 +161,11 @@ std::tuple<Tensor, Tensor, Tensor> miopen_batch_norm(
       const_cast<void*>(bias->const_data_ptr()),
       running_mean->data_ptr(),
       running_var->data_ptr(),
-      epsilon));
+      epsilon);
+    if (tuning_policy != previous_policy) {
+      MIOPEN_CHECK(miopenSetTuningPolicy(handle, previous_policy));
+    }
+    MIOPEN_CHECK(status);
   }
 
   // save_mean and save_var can be undefined
@@ -223,7 +238,13 @@ std::tuple<Tensor, Tensor, Tensor> miopen_batch_norm_backward(
   Constant one(dataType, 1);
   Constant zero(dataType, 0);
 
-  MIOPEN_CHECK(miopenBatchNormalizationBackward(
+  auto tuning_policy = at::globalContext().benchmarkCuDNN() ? miopenTuningPolicySearch : miopenTuningPolicyNone;
+  miopenTuningPolicy_t previous_policy;
+  MIOPEN_CHECK(miopenGetTuningPolicy(handle, &previous_policy));
+  if (tuning_policy != previous_policy) {
+    MIOPEN_CHECK(miopenSetTuningPolicy(handle, tuning_policy));
+  }
+  auto status = miopenBatchNormalizationBackward(
     handle, mode, &one, &zero, &one, &zero,
     idesc.desc(), input->const_data_ptr(),
     idesc.desc(), grad_output->const_data_ptr(),
@@ -233,7 +254,11 @@ std::tuple<Tensor, Tensor, Tensor> miopen_batch_norm_backward(
     grad_bias_t.data_ptr(),
     epsilon,
     save_mean->const_data_ptr(),
-    save_var->const_data_ptr()));
+    save_var->const_data_ptr());
+  if (tuning_policy != previous_policy) {
+    MIOPEN_CHECK(miopenSetTuningPolicy(handle, previous_policy));
+  }
+  MIOPEN_CHECK(status);
 
   return std::tuple<Tensor,Tensor,Tensor>{grad_input_t, grad_weight_t, grad_bias_t};
 }