[TOPI, Relay refactor] Move Dilation2d from nn to image namespace (ap…

…ache#5110)

docs/api/python/topi.rst

@@ -57,7 +57,6 @@ List of operators
    topi.nn.relu
    topi.nn.leaky_relu
    topi.nn.dilate
-   topi.nn.dilation2d
    topi.nn.pool
    topi.nn.global_pool
    topi.nn.adaptive_pool
@@ -106,6 +105,7 @@ List of operators
    topi.layout_transform
    topi.image.resize
    topi.image.crop_and_resize
+   topi.image.dilation2d
    topi.argsort
    topi.topk
    topi.sequence_mask
@@ -198,7 +198,6 @@ topi.nn
 .. autofunction:: topi.nn.upsampling
 .. autofunction:: topi.nn.softmax
 .. autofunction:: topi.nn.dense
-.. autofunction:: topi.nn.dilation2d
 .. autofunction:: topi.nn.batch_matmul
 .. autofunction:: topi.nn.log_softmax
 .. autofunction:: topi.nn.conv2d_nchw

docs/langref/relay_op.rst

@@ -70,7 +70,6 @@ This level enables typical convnet models.
    tvm.relay.nn.conv2d
    tvm.relay.nn.conv2d_transpose
    tvm.relay.nn.dense
-   tvm.relay.nn.dilation2d
    tvm.relay.nn.max_pool2d
    tvm.relay.nn.max_pool3d
    tvm.relay.nn.avg_pool2d
@@ -171,6 +170,7 @@ This level enables additional math and transform operators.
 
    tvm.relay.image.resize
    tvm.relay.image.crop_and_resize
+   tvm.relay.image.dilation2d
    tvm.relay.vision.multibox_prior
    tvm.relay.vision.multibox_transform_loc
    tvm.relay.vision.nms
@@ -250,7 +250,6 @@ Level 2 Definitions
 .. autofunction:: tvm.relay.nn.conv2d
 .. autofunction:: tvm.relay.nn.conv2d_transpose
 .. autofunction:: tvm.relay.nn.dense
-.. autofunction:: tvm.relay.nn.dilation2d
 .. autofunction:: tvm.relay.nn.max_pool2d
 .. autofunction:: tvm.relay.nn.max_pool3d
 .. autofunction:: tvm.relay.nn.avg_pool2d
@@ -339,6 +338,7 @@ Level 5 Definitions
 -------------------
 .. autofunction:: tvm.relay.image.resize
 .. autofunction:: tvm.relay.image.crop_and_resize
+.. autofunction:: tvm.relay.image.dilation2d
 .. autofunction:: tvm.relay.vision.multibox_prior
 .. autofunction:: tvm.relay.vision.multibox_transform_loc
 .. autofunction:: tvm.relay.vision.nms

include/tvm/relay/attrs/image.h

@@ -91,6 +91,41 @@ struct CropAndResizeAttrs : public tvm::AttrsNode<CropAndResizeAttrs> {
   }
 };
 
+/*! \brief Attributes used in dilation operators */
+struct Dilation2DAttrs : public tvm::AttrsNode<Dilation2DAttrs> {
+  Array<IndexExpr> strides;
+  Array<IndexExpr> padding;
+  Array<IndexExpr> dilations;
+  std::string data_layout;
+  std::string kernel_layout;
+  DataType out_dtype;
+
+  TVM_DECLARE_ATTRS(Dilation2DAttrs, "relay.attrs.Dilation2DAttrs") {
+    TVM_ATTR_FIELD(strides).set_default(Array<IndexExpr>({1, 1}))
+        .describe("Specifies the strides of the sliding window. [stride_height, stride_width].");
+    TVM_ATTR_FIELD(padding).set_default(Array<IndexExpr>({0, 0}))
+        .describe("If padding is non-zero, then the input is implicitly zero-padded"
+                  "Padding support both symmetric and asymmetric as"
+                  "one int : same padding used on all sides"
+                  "two int : bottom, right will use same padding as top, left"
+                  "four int : padding width in the order of (top, left, bottom, right)");
+    TVM_ATTR_FIELD(dilations).set_default(Array<IndexExpr>({1, 1}))
+        .describe("Specifies the dilation rate to use. [dilation_height, dilation_width]");
+    TVM_ATTR_FIELD(data_layout).set_default("NCHW")
+        .describe("Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
+                  "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
+                  "dimensions respectively. Convolution is applied on the 'H' and"
+                  "'W' dimensions.");
+    TVM_ATTR_FIELD(kernel_layout).set_default("IHW")
+        .describe("Dimension ordering of weight. Can be 'IHW', 'HWI', etc."
+                  "'I', 'H', 'W' stands for input_channel, height, and width"
+                  "dimensions respectively.");
+    TVM_ATTR_FIELD(out_dtype)
+        .set_default(NullValue<DataType>())
+        .describe("Output data type, set to explicit type under mixed precision setting");
+  }
+};
+
 }  // namespace relay
 }  // namespace tvm
 #endif  // TVM_RELAY_ATTRS_IMAGE_H_

include/tvm/relay/attrs/nn.h

@@ -155,43 +155,6 @@ struct Conv2DAttrs : public tvm::AttrsNode<Conv2DAttrs> {
   }
 };
 
-
-/*! \brief Attributes used in dilation operators */
-struct Dilation2DAttrs : public tvm::AttrsNode<Dilation2DAttrs> {
-  Array<IndexExpr> strides;
-  Array<IndexExpr> padding;
-  Array<IndexExpr> dilations;
-  std::string data_layout;
-  std::string kernel_layout;
-  DataType out_dtype;
-
-  TVM_DECLARE_ATTRS(Dilation2DAttrs, "relay.attrs.Dilation2DAttrs") {
-    TVM_ATTR_FIELD(strides).set_default(Array<IndexExpr>({1, 1}))
-        .describe("Specifies the strides of the sliding window. [stride_height, stride_width].");
-    TVM_ATTR_FIELD(padding).set_default(Array<IndexExpr>({0, 0}))
-        .describe("If padding is non-zero, then the input is implicitly zero-padded"
-                  "Padding support both symmetric and asymmetric as"
-                  "one int : same padding used on all sides"
-                  "two int : bottom, right will use same padding as top, left"
-                  "four int : padding width in the order of (top, left, bottom, right)");
-    TVM_ATTR_FIELD(dilations).set_default(Array<IndexExpr>({1, 1}))
-        .describe("Specifies the dilation rate to use. [dilation_height, dilation_width]");
-    TVM_ATTR_FIELD(data_layout).set_default("NCHW")
-        .describe("Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
-                  "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
-                  "dimensions respectively. Convolution is applied on the 'H' and"
-                  "'W' dimensions.");
-    TVM_ATTR_FIELD(kernel_layout).set_default("IHW")
-        .describe("Dimension ordering of weight. Can be 'IHW', 'HWI', etc."
-                  "'I', 'H', 'W' stands for input_channel, height, and width"
-                  "dimensions respectively.");
-    TVM_ATTR_FIELD(out_dtype)
-        .set_default(NullValue<DataType>())
-        .describe("Output data type, set to explicit type under mixed precision setting");
-  }
-};
-
-
 /*! \brief Attributes used in winograd weight transformation operators */
 struct Conv2DWinogradWeightTransformAttrs :
     public tvm::AttrsNode<Conv2DWinogradWeightTransformAttrs> {

python/tvm/relay/op/image/_image.py

@@ -20,6 +20,9 @@
 
 import topi
 from .. import op as reg
+from .. import strategy
+from ..op import OpPattern
+
 
 # resize
 @reg.register_compute("image.resize")
@@ -47,3 +50,8 @@ def compute_crop_and_resize(attrs, inputs, out_type):
                                        extrapolation_value, out_dtype)]
 
 reg.register_injective_schedule("image.crop_and_resize")
+
+
+# dilation2d
+reg.register_strategy("image.dilation2d", strategy.dilation2d_strategy)
+reg.register_pattern("image.dilation2d", OpPattern.OUT_ELEMWISE_FUSABLE)

python/tvm/relay/op/image/image.py

@@ -113,3 +113,60 @@ def crop_and_resize(data,
     """
     return _make.crop_and_resize(data, boxes, box_indices, crop_size,
                                  layout, method, extrapolation_value, out_dtype)
+
+
+def dilation2d(data,
+               weight,
+               strides=(1, 1),
+               padding=(0, 0),
+               dilations=(1, 1),
+               data_layout="NCHW",
+               kernel_layout="IHW",
+               out_dtype=""):
+    r"""Morphological Dilation 2D.
+    This operator takes the weight as the dilation kernel and dilates it with
+    data to produce an output. In the default case, where the data_layout is `NCHW`
+    and kernel_layout is `OIHW`, dilation2d takes in a data Tensor with shape
+    `(batch_size, in_channels, height, width)`, and a weight Tensor with shape
+    `(channels, kernel_height, kernel_width)` to produce an output Tensor
+    with the following rule:
+
+    .. math::
+        \mbox{out}[b, c, y, x] = \max_{dy, dx}
+           \mbox{data}[b, c, \mbox{strides}[0] * y  + dy, \mbox{strides}[1] * x + dx] +
+           \mbox{weight}[c, dy, dx]
+
+    Padding and dilation are applied to data and weight respectively before the computation.
+    This operator accepts data layout specification. Semantically, the operator
+    will convert the layout to the canonical layout
+    (`NCHW` for data and `IHW` for weight) and perform the computation.
+
+    weight : tvm.relay.Expr
+        The weight expressions.
+
+    strides : Optional[Tuple[int]]
+        The strides of convolution.
+
+    padding : Optional[Tuple[int]]
+        The padding of convolution on both sides of inputs before convolution.
+
+    dilations : Optional[Tuple[int]]
+        Specifies the dilation rate to be used for dilated convolution.
+
+    data_layout : Optional[str]
+        Layout of the input.
+
+    kernel_layout : Optional[str]
+        Layout of the weight.
+
+    out_dtype : Optional[str]
+        Specifies the output data type.
+
+    Returns
+    -------
+    result : tvm.relay.Expr
+        The computed result.
+    """
+
+    return _make.dilation2d(data, weight, strides, padding, dilations, data_layout,
+                            kernel_layout, out_dtype)

python/tvm/relay/op/nn/_nn.py

@@ -178,9 +178,6 @@ def legalize_conv2d_transpose(attrs, inputs, types):
 reg.register_strategy("nn.conv3d", strategy.conv3d_strategy)
 reg.register_pattern("nn.conv3d", OpPattern.OUT_ELEMWISE_FUSABLE)
 
-# dilation2d
-reg.register_strategy("nn.dilation2d", strategy.dilation2d_strategy)
-reg.register_pattern("nn.dilation2d", OpPattern.OUT_ELEMWISE_FUSABLE)
 
 # conv1d_transpose
 reg.register_strategy("nn.conv1d_transpose", strategy.conv1d_transpose_strategy)

python/tvm/relay/op/nn/nn.py

@@ -2463,60 +2463,3 @@ def adaptive_avg_pool3d(data,
     """
     output_size = [] or output_size
     return _make.adaptive_avg_pool3d(data, output_size, layout)
-
-
-def dilation2d(data,
-               weight,
-               strides=(1, 1),
-               padding=(0, 0),
-               dilations=(1, 1),
-               data_layout="NCHW",
-               kernel_layout="IHW",
-               out_dtype=""):
-    r"""Dilation 2D.
-    This operator takes the weight as the dilation kernel and dilates it with
-    data to produce an output. In the default case, where the data_layout is `NCHW`
-    and kernel_layout is `OIHW`, dilation2d takes in a data Tensor with shape
-    `(batch_size, in_channels, height, width)`, and a weight Tensor with shape
-    `(channels, kernel_height, kernel_width)` to produce an output Tensor
-    with the following rule:
-
-    .. math::
-        \mbox{out}[b, c, y, x] = \max_{dy, dx}
-           \mbox{data}[b, c, \mbox{strides}[0] * y  + dy, \mbox{strides}[1] * x + dx] +
-           \mbox{weight}[c, dy, dx]
-
-    Padding and dilation are applied to data and weight respectively before the computation.
-    This operator accepts data layout specification. Semantically, the operator
-    will convert the layout to the canonical layout
-    (`NCHW` for data and `IHW` for weight) and perform the computation.
-
-    weight : tvm.relay.Expr
-        The weight expressions.
-
-    strides : Optional[Tuple[int]]
-        The strides of convolution.
-
-    padding : Optional[Tuple[int]]
-        The padding of convolution on both sides of inputs before convolution.
-
-    dilations : Optional[Tuple[int]]
-        Specifies the dilation rate to be used for dilated convolution.
-
-    data_layout : Optional[str]
-        Layout of the input.
-
-    kernel_layout : Optional[str]
-        Layout of the weight.
-
-    out_dtype : Optional[str]
-        Specifies the output data type.
-
-    Returns
-    -------
-    result : tvm.relay.Expr
-        The computed result.
-    """
-
-    return _make.dilation2d(data, weight, strides, padding, dilations, data_layout,
-                            kernel_layout, out_dtype)

python/tvm/relay/op/op_attrs.py

@@ -44,11 +44,6 @@ class Conv2DWinogradNNPACKWeightTransformAttrs(Attrs):
     """Attributes for nn.contrib_conv2d_winograd_nnpack_weight_transform"""
 
 
-@tvm._ffi.register_object("relay.attrs.Dilation2DAttrs")
-class Dilation2DAttrs(Attrs):
-    """Attributes for nn.dilation2d"""
-
-
 @tvm._ffi.register_object("relay.attrs.GlobalPool2DAttrs")
 class GlobalPool2DAttrs(Attrs):
     """Attributes for nn.global_pool"""
@@ -124,10 +119,17 @@ class DeformableConv2DAttrs(Attrs):
 class ResizeAttrs(Attrs):
     """Attributes for image.resize"""
 
+
 @tvm._ffi.register_object("relay.attrs.CropAndResizeAttrs")
 class CropAndResizeAttrs(Attrs):
     """Attributes for image.crop_and_resize"""
 
+
+@tvm._ffi.register_object("relay.attrs.Dilation2DAttrs")
+class Dilation2DAttrs(Attrs):
+    """Attributes for image.dilation2d"""
+
+
 @tvm._ffi.register_object("relay.attrs.ArgsortAttrs")
 class ArgsortAttrs(Attrs):
     """Attributes for algorithm.argsort"""

python/tvm/relay/op/strategy/generic.py

@@ -479,13 +479,13 @@ def dilation2d_strategy(attrs, inputs, out_type, target):
     if layout == "NCHW":
         assert kernel_layout == "IHW"
         strategy.add_implementation(
-            wrap_compute_dilation2d(topi.nn.dilation2d_nchw),
+            wrap_compute_dilation2d(topi.image.dilation2d_nchw),
             wrap_topi_schedule(topi.generic.schedule_dilation2d_nchw),
             name="dilation2d_nchw.generic")
     elif layout == "NHWC":
         assert kernel_layout == "HWI"
         strategy.add_implementation(
-            wrap_compute_dilation2d(topi.nn.dilation2d_nhwc),
+            wrap_compute_dilation2d(topi.image.dilation2d_nhwc),
             wrap_topi_schedule(topi.generic.schedule_dilation2d_nhwc),
             name="dilation2d_nhwc.generic")
     else: