add back missed conv2d_output_shape by mistake

python/tvm/contrib/cudnn.py

@@ -172,6 +172,70 @@ def conv2d_w_shape(in_channel,
     """
     return [out_channel, in_channel, filter_h, filter_w]
 
+def conv2d_output_shape(tensor_format,
+                        pad_h,
+                        pad_w,
+                        stride_h,
+                        stride_w,
+                        dilation_h,
+                        dilation_w,
+                        x_shape,
+                        w_shape):
+    """Get output shape of 2D convolution
+
+    Paramters
+    ---------
+    tensor_format: int
+        0: CUDNN_TENSOR_NCHW
+        1: CUDNN_TENSOR_NHWC
+        2: CUDNN_TENSOR_NCHW_VECT_C
+    pad_h: int
+        height pad
+    pad_w: int
+        weight pad
+    stride_h: int
+        height stride
+    stride_w: int
+        width stride
+    dilation_h: int
+        height dilation
+    dilation_w: int
+        width dilation
+    x_shape: list
+        input shape
+    w_shape: list
+        weight shape
+
+    Returns
+    -------
+    oshape: list
+        output shape
+    """
+    assert isinstance(x_shape, list)
+    assert isinstance(w_shape, list)
+    assert len(x_shape) == 4
+    assert len(w_shape) == 4
+    oshape = np.zeros((len(x_shape)), dtype=np.int32)
+    func = _get_global_func("tvm.contrib.cudnn.conv2d.output_shape")
+    func(tensor_format,
+         pad_h,
+         pad_w,
+         stride_h,
+         stride_w,
+         dilation_h,
+         dilation_w,
+         x_shape[0].value,
+         x_shape[1].value,
+         x_shape[2].value,
+         x_shape[3].value,
+         w_shape[0].value,
+         w_shape[1].value,
+         w_shape[2].value,
+         w_shape[3].value,
+         _get_np_int32_array_handle(oshape))
+    return list(oshape)
+
+
 def conv2d_find_algo(tensor_format,
                      pad_h,
                      pad_w,
@@ -321,6 +385,7 @@ def conv2d_forward(x,
             ins[1],
             outs[0]), name="y")
 
+
 def conv3d_w_shape(in_channel,
                    out_channel,
                    filter_d,

tests/python/relay/test_op_level2.py

@@ -867,6 +867,7 @@ def test_bitpack_infer_type():
     test_conv2d_transpose_run()
     test_conv2d_run()
     test_conv2d_winograd()
+    test_conv3d_run()
     test_bitserial_conv2d_infer_type()
     test_batch_flatten()
     test_upsampling()