miopen tune api for batchnorm

Author: dnikolaev-amd

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

@@ -115,6 +115,13 @@ 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());
@@ -135,6 +142,9 @@ std::tuple<Tensor, Tensor, Tensor> miopen_batch_norm(
       epsilon,
       save_mean.mutable_data_ptr(),
       save_var.mutable_data_ptr()));
+    if (tuning_policy != previous_policy) {
+      MIOPEN_CHECK(miopenSetTuningPolicy(handle, previous_policy));
+    }
   } else {
     save_mean = at::empty({0}, weight_t.options());
     save_var = at::empty({0}, weight_t.options());
@@ -151,6 +161,9 @@ std::tuple<Tensor, Tensor, Tensor> miopen_batch_norm(
       running_mean->data_ptr(),
       running_var->data_ptr(),
       epsilon));
+    if (tuning_policy != previous_policy) {
+      MIOPEN_CHECK(miopenSetTuningPolicy(handle, previous_policy));
+    }
   }
 
   // save_mean and save_var can be undefined
@@ -223,6 +236,12 @@ std::tuple<Tensor, Tensor, Tensor> miopen_batch_norm_backward(
   Constant one(dataType, 1);
   Constant zero(dataType, 0);
 
+  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));
+  }
   MIOPEN_CHECK(miopenBatchNormalizationBackward(
     handle, mode, &one, &zero, &one, &zero,
     idesc.desc(), input->const_data_ptr(),
@@ -234,6 +253,9 @@ std::tuple<Tensor, Tensor, Tensor> miopen_batch_norm_backward(
     epsilon,
     save_mean->const_data_ptr(),
     save_var->const_data_ptr()));
+  if (tuning_policy != previous_policy) {
+    MIOPEN_CHECK(miopenSetTuningPolicy(handle, previous_policy));
+  }
 
   return std::tuple<Tensor,Tensor,Tensor>{grad_input_t, grad_weight_t, grad_bias_t};
 }