MIOpen: Honor Max Dim (#222)

* MIOpen: Batchnorm - Allow half/half and half/float, disallow double

* MIOpen: Honor DIM_MAX

* Limit MIOpen batchnorm to same-precision

Author: mwootton

aten/src/ATen/native/Convolution.cpp

@@ -3,6 +3,9 @@
 
 #include "ATen/Config.h"
 
+static const int MIOPEN_DIM_MAX = 4;
+static const bool MIOPEN_ENABLED = getenv("DISABLE_MIOPEN") != NULL;
+
 namespace at { namespace native {
 
 struct ConvParams {
@@ -123,7 +126,8 @@ auto ConvParams::use_miopen(const at::Tensor& input) const -> bool {
   return ((input.type().scalarType() == at::kFloat) || (input.type().scalarType() == at::kHalf))
          && detail::getCUDAHooks().compiledWithMIOpen()
          && input.type().is_cuda()
-         && cudnn_enabled
+         && input.dim() > MIOPEN_DIM_MAX
+         && MIOPEN_ENABLED
          ;
 }
 

aten/src/ATen/native/Normalization.cpp

@@ -7,6 +7,9 @@
 
 #include <vector>
 
+static const int MIOPEN_DIM_MAX = 4;
+static const bool MIOPEN_ENABLED = getenv("DISABLE_MIOPEN") != NULL;
+
 namespace at { namespace native {
 
 namespace {
@@ -67,12 +70,14 @@ Tensor batch_norm(
   }
 
   bool use_miopen = (input.type().is_cuda()
-               && (input.type().scalarType() != at::kHalf
-                 || weight.type().scalarType() == at::kFloat)
+               && input.dim() < MIOPEN_DIM_MAX
+               && input.type().scalarType() != at::kDouble
+               && (input.type().scalarType() == weight.type().scalarType())
                && weight.defined() && bias.defined()
                && ((running_mean.defined() && running_var.defined())
                  || (!running_mean.defined() && !running_var.defined() && training))
                && detail::getCUDAHooks().compiledWithMIOpen()
+               && MIOPEN_ENABLED
                );
 
   if (use_miopen) {

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

@@ -67,9 +67,7 @@ std::tuple<Tensor, Tensor, Tensor> miopen_batch_norm(
     checkAllDefined(c, {running_mean, running_var});
   }
   checkAllSameGPU(c, {input, weight, bias, running_mean, running_var});
-  if (input->type().scalarType() == ScalarType::Half) {
-    checkScalarType(c, weight, ScalarType::Float);
-  } else {
+  if (input->type().scalarType() != ScalarType::Half) {
     checkAllSameType(c, {input, weight});
   }
   checkAllSameType(c, {weight, bias, running_mean, running_var});