Make CUDNN an alias of MIOPEN for HIP ops (pytorch#12278)

Summary:
This is mostly for reusing all the cudnn test cases in our python operator_tests.
Pull Request resolved: pytorch#12278

Differential Revision: D10842592

Pulled By: bddppq

fbshipit-source-id: 4b3ed91fca64ff02060837b3270393bc2f9a9898

caffe2/core/operator.h

@@ -915,7 +915,8 @@ C10_DECLARE_REGISTRY(
   C10_REGISTER_CLASS(HIPOperatorRegistry, name##_ENGINE_##engine, __VA_ARGS__)
 
 #define REGISTER_MIOPEN_OPERATOR(name, ...) \
-  REGISTER_HIP_OPERATOR_WITH_ENGINE(name, MIOPEN, __VA_ARGS__)
+  REGISTER_HIP_OPERATOR_WITH_ENGINE(name, MIOPEN, __VA_ARGS__) \
+  REGISTER_HIP_OPERATOR_WITH_ENGINE(name, CUDNN, __VA_ARGS__) // Make CUDNN an alias of MIOPEN for HIP ops
 
 // StaticLinkingProtector is a helper class that ensures that the Caffe2
 // library is linked correctly with whole archives (in the case of static

caffe2/operators/hip/spatial_batch_norm_op_miopen.cc

@@ -18,6 +18,7 @@
 #include "caffe2/core/hip/context_hip.h"
 #include "caffe2/core/hip/miopen_wrapper.h"
 #include "caffe2/operators/spatial_batch_norm_op.h"
+#include "caffe2/operators/hip/spatial_batch_norm_op_hip_impl.cuh"
 #include "caffe2/utils/math.h"
 
 const double MIOPEN_BN_MIN_EPSILON = 1e-6;
@@ -33,8 +34,6 @@ class MIOpenSpatialBNOp final : public SpatialBNOp<HIPContext> {
         alpha_(OperatorBase::GetSingleArgument<float>("alpha", 1.0)),
         beta_(OperatorBase::GetSingleArgument<float>("beta", 0.0)),
         mode_(miopenBNSpatial) {
-    CAFFE_ENFORCE_EQ(
-        order_, StorageOrder::NCHW, "Only NCHW order is supported right now.");
     MIOPEN_ENFORCE(miopenCreateTensorDescriptor(&data_desc_));
     MIOPEN_ENFORCE(miopenCreateTensorDescriptor(&bn_param_desc_));
     if (epsilon_ <= MIOPEN_BN_MIN_EPSILON) {
@@ -247,6 +246,10 @@ bool MIOpenSpatialBNOp::DoRunWithType() {
   return true;
 }
 bool MIOpenSpatialBNOp::RunOnDevice() {
+  // Fall back to HIP for multi batch spatial BNorm and NHWC layout
+  if (num_batches_ > 1 || order_ == StorageOrder::NHWC) {
+    return SpatialBNOp<HIPContext>::RunOnDevice();
+  }
   if (Input(0).IsType<float>()) {
     return DoRunWithType<float, float>();
   } else {
@@ -327,6 +330,10 @@ bool MIOpenSpatialBNGradientOp::DoRunWithType() {
   return true;
 }
 bool MIOpenSpatialBNGradientOp::RunOnDevice() {
+  // Fall back to HIP for multi batch spatial BNorm and NHWC layout
+  if (num_batches_ > 1 || order_ == StorageOrder::NHWC) {
+    return SpatialBNGradientOp<HIPContext>::RunOnDevice();
+  }
   if (Input(0).IsType<float>()) {
     return DoRunWithType<float, float>();
   } else {

caffe2/python/hip_test_util.py

@@ -0,0 +1,18 @@
+## @package hip_test_util
+# Module caffe2.python.hip_test_util
+"""
+The HIP test utils is a small addition on top of the hypothesis test utils
+under caffe2/python, which allows one to more easily test HIP/ROCm related
+operators.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+from caffe2.proto import caffe2_pb2
+
+def run_in_hip(gc, dc):
+    return (gc.device_type == caffe2_pb2.HIP) or (
+        caffe2_pb2.HIP in {d.device_type for d in dc})

caffe2/python/operator_test/conv_test.py

@@ -11,6 +11,7 @@
 
 from caffe2.proto import caffe2_pb2
 from caffe2.python import brew, core, workspace
+import caffe2.python.hip_test_util as hiputl
 import caffe2.python.hypothesis_test_util as hu
 from caffe2.python.model_helper import ModelHelper
 import caffe2.python.serialized_test.serialized_test_util as serial
@@ -19,7 +20,6 @@
 import unittest
 import os
 
-
 def _cudnn_supports(
         dilation=False,
         nhwc=False,
@@ -213,9 +213,12 @@ def test_convolution_gradients(
         dkernel = dilation * (kernel - 1) + 1
 
         if engine == 'CUDNN':
-            assume(_cudnn_supports(dilation=(dilation > 1),
-                                   nhwc=(order == 'NHWC'),
-                                   backward=True))
+            if hiputl.run_in_hip(gc, dc):
+                assume((order == "NCHW") and not (dilation > 1 and group > 1))
+            else:
+                assume(_cudnn_supports(dilation=(dilation > 1),
+                                       nhwc=(order == 'NHWC'),
+                                       backward=True))
 
         assume(engine != "MKLDNN" or use_bias is True)
 
@@ -373,7 +376,7 @@ def test_3d_convolution_nchw(self, input_channels, output_channels,
            force_algo_fwd=_cudnn_convolution_algo_count("fwd"),
            force_algo_dgrad=_cudnn_convolution_algo_count("dgrad"),
            force_algo_wgrad=_cudnn_convolution_algo_count("wgrad"),
-           **hu.gcs)
+           **hu.gcs_no_hip)     # MIOPEN doesn't support 3D conv yet
     def test_3d_convolution_cudnn_nchw(self, op_type, batch_size, stride, size,
                                        kernel, dilation, pad, use_bias,
                                        force_algo_fwd, force_algo_dgrad,
@@ -461,8 +464,12 @@ def test_convolution_layout(self, op_type, stride, pad, kernel, dilation,
 
         for order in ["NCHW", "NHWC"]:
             engine_list = ['']
-            if _cudnn_supports(dilation=(dilation > 1), nhwc=(order == 'NHWC')):
-                engine_list.append('CUDNN')
+            if hiputl.run_in_hip(gc, dc):
+                if order == 'NCHW':
+                    engine_list.append('MIOPEN')
+            else:
+                if _cudnn_supports(dilation=(dilation > 1), nhwc=(order == 'NHWC')):
+                    engine_list.append('CUDNN')
 
             for engine in engine_list:
                 op = core.CreateOperator(
@@ -649,6 +656,8 @@ def test_use_cudnn_engine_interactions(self):
     def test_1x1_conv(self, op_type, N, G, DX, DY, H, W, use_bias, order,
                       force_algo_fwd, force_algo_dgrad,
                       force_algo_wgrad, gc, dc):
+        if hiputl.run_in_hip(gc, dc):
+            assume(order == "NCHW")
         if order == "NHWC":
             G = 1
 

caffe2/python/operator_test/group_conv_test.py

@@ -8,12 +8,12 @@
 
 from caffe2.proto import caffe2_pb2
 from caffe2.python import core
+import caffe2.python.hip_test_util as hiputl
 import caffe2.python.hypothesis_test_util as hu
 
 import unittest
 import os
 
-
 class TestGroupConvolution(hu.HypothesisTestCase):
 
     @given(stride=st.integers(1, 3),
@@ -36,8 +36,13 @@ def test_group_convolution(
             input_channels_per_group, output_channels_per_group, batch_size,
             order, engine, use_bias, gc, dc):
         assume(size >= kernel)
-        # TODO: Group conv in NHWC not implemented for GPU yet.
-        assume(group == 1 or order == "NCHW" or gc.device_type != caffe2_pb2.CUDA)
+
+        if hiputl.run_in_hip(gc, dc):
+            if order == "NHWC":
+                assume(group == 1 and engine != "CUDNN")
+        else:
+            # TODO: Group conv in NHWC not implemented for GPU yet.
+            assume(group == 1 or order == "NCHW" or gc.device_type != caffe2_pb2.CUDA)
         input_channels = input_channels_per_group * group
         output_channels = output_channels_per_group * group
 

caffe2/python/operator_test/pooling_test.py

@@ -9,9 +9,9 @@
 import unittest
 
 from caffe2.python import core, workspace
+import caffe2.python.hip_test_util as hiputl
 import caffe2.python.hypothesis_test_util as hu
 
-
 class TestPooling(hu.HypothesisTestCase):
     # CUDNN does NOT support different padding values and we skip it
     @given(stride_h=st.integers(1, 3),
@@ -126,6 +126,9 @@ def test_pooling_3d(self, stride, pad, kernel, size, input_channels,
                         batch_size, order, op_type, engine, gc, dc):
         assume(pad < kernel)
         assume(size + pad + pad >= kernel)
+        # Currently MIOpen Pooling only supports 2d pooling
+        if hiputl.run_in_hip(gc, dc):
+            assume(engine != "CUDNN")
         # some case here could be calculated with global pooling, but instead
         # calculated with general implementation, slower but should still
         # be corect.
@@ -159,6 +162,9 @@ def test_pooling_3d(self, stride, pad, kernel, size, input_channels,
            **hu.gcs)
     def test_global_pooling_3d(self, kernel, size, input_channels,
                                batch_size, order, op_type, engine, gc, dc):
+        # Currently MIOpen Pooling only supports 2d pooling
+        if hiputl.run_in_hip(gc, dc):
+            assume(engine != "CUDNN")
         # pad and stride ignored because they will be infered in global_pooling
         op = core.CreateOperator(
             op_type,
@@ -276,6 +282,9 @@ def test_pooling(self, stride, pad, kernel, size,
                      input_channels, batch_size,
                      order, op_type, engine, gc, dc):
         assume(pad < kernel)
+        if hiputl.run_in_hip(gc, dc) and engine == "CUDNN":
+            assume(order == "NCHW" and op_type != "LpPool")
+
         op = core.CreateOperator(
             op_type,
             ["X"],
@@ -306,6 +315,9 @@ def test_global_pooling(self, size, input_channels, batch_size,
                             order, op_type, engine, gc, dc):
         # CuDNN 5 does not support deterministic max pooling.
         assume(workspace.GetCuDNNVersion() >= 6000 or op_type != "MaxPool")
+
+        if hiputl.run_in_hip(gc, dc) and engine == "CUDNN":
+            assume(order == "NCHW" and op_type != "LpPool")
         op = core.CreateOperator(
             op_type,
             ["X"],

caffe2/python/operator_test/spatial_bn_op_test.py

@@ -5,21 +5,17 @@
 
 from caffe2.proto import caffe2_pb2
 from caffe2.python import brew, core, workspace
+import caffe2.python.hip_test_util as hiputl
 import caffe2.python.hypothesis_test_util as hu
 from caffe2.python.model_helper import ModelHelper
 import caffe2.python.serialized_test.serialized_test_util as serial
 
-from hypothesis import given
+from hypothesis import given, assume
 import hypothesis.strategies as st
 import numpy as np
 import unittest
 
 
-def _run_in_hip(gc, dc):
-    return (gc.device_type == caffe2_pb2.HIP) or (
-        caffe2_pb2.HIP in {d.device_type for d in dc})
-
-
 class TestSpatialBN(serial.SerializedTestCase):
 
     @serial.given(size=st.integers(7, 10),
@@ -30,11 +26,13 @@ class TestSpatialBN(serial.SerializedTestCase):
            epsilon=st.floats(min_value=1e-5, max_value=1e-2),
            inplace=st.booleans(),
            engine=st.sampled_from(["", "CUDNN"]),
-           # Currently HIP SpatialBN only supports 2D
-           **hu.gcs_no_hip)
+           **hu.gcs)
     def test_spatialbn_test_mode_3d(
             self, size, input_channels, batch_size, seed, order, epsilon,
             inplace, engine, gc, dc):
+        # Currently MIOPEN SpatialBN only supports 2D
+        if hiputl.run_in_hip(gc, dc):
+            assume(engine != "CUDNN")
         op = core.CreateOperator(
             "SpatialBN",
             ["X", "scale", "bias", "mean", "var"],
@@ -79,11 +77,13 @@ def reference_spatialbn_test(X, scale, bias, mean, var):
            epsilon=st.floats(min_value=1e-5, max_value=1e-2),
            inplace=st.booleans(),
            engine=st.sampled_from(["", "CUDNN"]),
-           # Currently HIP SpatialBN only supports 2D
-           **hu.gcs_no_hip)
+           **hu.gcs)
     def test_spatialbn_test_mode_1d(
             self, size, input_channels, batch_size, seed, order, epsilon,
             inplace, engine, gc, dc):
+        # Currently MIOPEN SpatialBN only supports 2D
+        if hiputl.run_in_hip(gc, dc):
+            assume(engine != "CUDNN")
         op = core.CreateOperator(
             "SpatialBN",
             ["X", "scale", "bias", "mean", "var"],
@@ -129,8 +129,8 @@ def test_spatialbn_test_mode(
             self, size, input_channels, batch_size, seed, order, epsilon,
             inplace, engine, gc, dc):
         # Currently HIP SpatialBN only supports NCHW
-        if _run_in_hip(gc, dc) and (order != 'NCHW'):
-            return
+        if hiputl.run_in_hip(gc, dc):
+            assume(order == "NCHW")
 
         op = core.CreateOperator(
             "SpatialBN",
@@ -179,8 +179,8 @@ def test_spatialbn_train_mode(
             self, size, input_channels, batch_size, seed, order, epsilon,
             momentum, inplace, engine, gc, dc):
         # Currently HIP SpatialBN only supports NCHW
-        if _run_in_hip(gc, dc) and (order != 'NCHW'):
-            return
+        if hiputl.run_in_hip(gc, dc):
+            assume(order == "NCHW")
 
         op = core.CreateOperator(
             "SpatialBN",
@@ -220,8 +220,8 @@ def test_spatialbn_train_mode_gradient_check(
             self, size, input_channels, batch_size, seed, order, epsilon,
             momentum, engine, gc, dc):
         # Currently HIP SpatialBN only supports NCHW
-        if _run_in_hip(gc, dc) and (order != 'NCHW'):
-            return
+        if hiputl.run_in_hip(gc, dc):
+            assume(order == "NCHW")
 
         op = core.CreateOperator(
             "SpatialBN",
@@ -255,11 +255,13 @@ def test_spatialbn_train_mode_gradient_check(
            epsilon=st.floats(min_value=1e-5, max_value=1e-2),
            momentum=st.floats(min_value=0.5, max_value=0.9),
            engine=st.sampled_from(["", "CUDNN"]),
-           # Currently HIP SpatialBN only supports 2D
-           **hu.gcs_no_hip)
+           **hu.gcs)
     def test_spatialbn_train_mode_gradient_check_1d(
             self, size, input_channels, batch_size, seed, order, epsilon,
             momentum, engine, gc, dc):
+        # Currently MIOPEN SpatialBN only supports 2D
+        if hiputl.run_in_hip(gc, dc):
+            assume(engine != "CUDNN")
         op = core.CreateOperator(
             "SpatialBN",
             ["X", "scale", "bias", "mean", "var"],