[RELAY] Port winograd ops to relay (#2356)

include/tvm/relay/attrs/nn.h

@@ -38,7 +38,7 @@ struct Conv2DAttrs : public tvm::AttrsNode<Conv2DAttrs> {
   IndexExpr channels;
   Array<IndexExpr> kernel_size;
   std::string data_layout;
-  std::string weight_layout;
+  std::string kernel_layout;
   std::string out_layout;
   DataType out_dtype;
 
@@ -68,7 +68,7 @@ struct Conv2DAttrs : public tvm::AttrsNode<Conv2DAttrs> {
                   "'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(weight_layout).set_default("OIHW")
+    TVM_ATTR_FIELD(kernel_layout).set_default("OIHW")
         .describe("Dimension ordering of weight. Can be 'OIHW', 'OIHW16o16i', etc."
                   "'O', 'I', 'H', 'W' stands for num_filter, input_channel, height, and width"
                   "dimensions respectively.");
@@ -84,13 +84,85 @@ struct Conv2DAttrs : public tvm::AttrsNode<Conv2DAttrs> {
   }
 };
 
+
+/*! \brief Attributes used in winograd weight transformation operators */
+struct Conv2DWinogradWeightTransformAttrs :
+    public tvm::AttrsNode<Conv2DWinogradWeightTransformAttrs> {
+  int tile_size;
+
+  TVM_DECLARE_ATTRS(Conv2DWinogradWeightTransformAttrs,
+      "relay.attrs.Conv2DWinogradWeightTransformAttrs") {
+    TVM_ATTR_FIELD(tile_size)
+      .describe("Tile size of winograd. E.g. 2 for F(2x2, 3x3) and 4 for F(4x4, 3x3)");
+  }
+};
+
+/*! \brief Attributes used in convolution operators with winograd algorithm */
+struct Conv2DWinogradAttrs : public tvm::AttrsNode<Conv2DWinogradAttrs> {
+  int tile_size;
+  Array<IndexExpr> strides;
+  Array<IndexExpr> padding;
+  Array<IndexExpr> dilation;
+  int groups;
+  IndexExpr channels;
+  Array<IndexExpr> kernel_size;
+  std::string data_layout;
+  std::string kernel_layout;
+  std::string out_layout;
+  DataType out_dtype;
+
+  TVM_DECLARE_ATTRS(Conv2DWinogradAttrs, "relay.attrs.Conv2DWinogradAttrs") {
+    TVM_ATTR_FIELD(tile_size)
+      .describe("The tile size of winograd. E.g. 2 for F(2x2, 3x3) and 4 for F(4x4, 3x3)");
+    TVM_ATTR_FIELD(strides).set_default(Array<IndexExpr>({1, 1}))
+        .describe("Specifies the strides of the convolution.");
+    TVM_ATTR_FIELD(padding).set_default(Array<IndexExpr>({0, 0}))
+        .describe("If padding is non-zero, then the input is implicitly zero-padded"
+                  "on both sides for padding number of points");
+    TVM_ATTR_FIELD(dilation).set_default(Array<IndexExpr>({1, 1}))
+        .describe("Specifies the dilation rate to use for dilated convolution.");
+    TVM_ATTR_FIELD(groups).set_default(1)
+        .describe("Controls the connections between inputs and outputs."
+                  "At groups=1, all inputs are convolved to all outputs."
+                  "At groups=2, the operation becomes equivalent to having two convolution"
+                  "layers side by side, each seeing half the input channels, and producing"
+                  "half the output channels, and both subsequently concatenated.");
+    TVM_ATTR_FIELD(channels)
+        .describe("The number of output channels in the convolution."
+                  " If it is not set, inferred by shape of the weight.")
+        .set_default(NullValue<IndexExpr>());
+    TVM_ATTR_FIELD(kernel_size)
+        .describe("Specifies the dimensions of the convolution window.")
+        .set_default(NullValue<Array<IndexExpr> >());
+    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("OIHW")
+        .describe("Dimension ordering of weight. Can be 'OIHW', 'OIHW16o16i', etc."
+                  "'O', 'I', 'H', 'W' stands for num_filter, input_channel, height, and width"
+                  "dimensions respectively.");
+    TVM_ATTR_FIELD(out_layout).set_default("")
+        .describe("Dimension ordering of output. Can be 'NCHW', 'NHWC', etc."
+                  "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
+                  "dimensions respectively. Default to be same as input layout.");
+
+    // use 0 bits to indicate none.
+    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 softmax operators */
 struct SoftmaxAttrs : public tvm::AttrsNode<SoftmaxAttrs> {
   int axis;
 
   TVM_DECLARE_ATTRS(SoftmaxAttrs, "relay.attrs.SoftmaxAttrs") {
-      TVM_ATTR_FIELD(axis).set_default(-1)
-          .describe("The axis to sum over when computing softmax.");
+    TVM_ATTR_FIELD(axis).set_default(-1)
+      .describe("The axis to sum over when computing softmax.");
   }
 };
 
@@ -104,7 +176,7 @@ struct Conv2DTransposeAttrs : public tvm::AttrsNode<Conv2DTransposeAttrs> {
   Array<IndexExpr> dilation;
   int groups;
   std::string data_layout;
-  std::string weight_layout;
+  std::string kernel_layout;
   std::string out_layout;
   DataType out_dtype;
 
@@ -136,7 +208,7 @@ struct Conv2DTransposeAttrs : public tvm::AttrsNode<Conv2DTransposeAttrs> {
                 "'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(weight_layout).set_default("OIHW")
+    TVM_ATTR_FIELD(kernel_layout).set_default("OIHW")
       .describe("Dimension ordering of data and weight. Can be 'OIHW', 'OIHW16o16i', etc."
                 "'O', 'I', 'H', 'W' stands for num_filter, input_channel, height, and width"
                 "dimensions respectively.");

nnvm/python/nnvm/to_relay.py

 # pylint: disable=no-else-return, unidiomatic-typecheck, invalid-name, unused-argument
 """Convert an NNVM graph to Relay."""
 import json
+import numpy
+
 from tvm import relay, nd
 from tvm.relay import op, expr, var
 from tvm.relay.frontend.common import StrAttrsDict
 from tvm.relay.frontend.nnvm_common import _rename
-import numpy
 from .symbol import Symbol
 from .compiler import graph_attr
 from .graph import create as graph_create
@@ -42,7 +43,7 @@ def _conv2d(children, attrs, odtype='float32'):
     dilation = attrs.get_int_tuple('dilation', (1, 1))
     groups = attrs.get_int('groups', 1)
     data_layout = attrs.get_str('layout', 'NCHW')
-    weight_layout = attrs.get_str('kernel_layout', 'OIHW')
+    kernel_layout = attrs.get_str('kernel_layout', 'OIHW')
     out_layout = ''
     out_dtype = attrs.get_str('out_dtype', '')
 
@@ -54,7 +55,7 @@ def _conv2d(children, attrs, odtype='float32'):
         dilation=dilation,
         groups=groups,
         data_layout=data_layout,
-        weight_layout=weight_layout,
+        kernel_layout=kernel_layout,
         out_layout=out_layout,
         out_dtype=out_dtype)
 
@@ -77,7 +78,7 @@ def _conv2d_transpose(children, attrs, odtype='float32'):
     dilation = attrs.get_int_tuple('dilation', (1, 1))
     groups = attrs.get_int('groups', 1)
     data_layout = attrs.get_str('layout', 'NCHW')
-    weight_layout = attrs.get_str('kernel_layout', 'OIHW')
+    kernel_layout = attrs.get_str('kernel_layout', 'OIHW')
     out_dtype = attrs.get_str('out_dtype', '')
 
     out_conv2d = op.nn.conv2d_transpose(
@@ -88,7 +89,7 @@ def _conv2d_transpose(children, attrs, odtype='float32'):
         dilation=dilation,
         groups=groups,
         data_layout=data_layout,
-        weight_layout=weight_layout,
+        kernel_layout=kernel_layout,
         out_dtype=out_dtype)
 
     if use_bias:

nnvm/python/nnvm/top/attr_dict.py

@@ -138,7 +138,7 @@ class AttrDict(object):
         else:
             raise ValueError("Wrong bool format for key %s" % key)
 
-    def get_string(self, key):
+    def get_str(self, key):
         """Get string from attr dict
 
         Parameters

nnvm/python/nnvm/top/nn.py

@@ -153,7 +153,25 @@ def schedule_conv2d(attrs, outs, target):
 
 @reg.register_alter_op_layout("conv2d")
 def alter_conv2d_layout(attrs, inputs, tinfos):
-    return topi.nn.conv2d_alter_layout(attrs, inputs, tinfos)
+    """Replace conv2d op with other layouts or algorithms"""
+    import nnvm.symbol as sym
+
+    # map relay op names to nnvm op names
+    sym.contrib_conv2d_winograd_without_weight_transform = \
+            sym.contrib.conv2d_winograd_without_weight_transform
+    sym.contrib_conv2d_winograd_weight_transform = \
+            sym.contrib.conv2d_winograd_weight_transform
+    sym.nn = sym
+
+    # map relay argument names to nnvm argument names
+    raw_reshape = sym.reshape
+    def _reshape(*args, **kwargs):
+        if "newshape" in kwargs:
+            kwargs['shape'] = kwargs.pop('newshape')
+        return raw_reshape(*args, **kwargs)
+    sym.reshape = _reshape
+
+    return topi.nn.conv2d_alter_layout(attrs, inputs, tinfos, sym)
 
 reg.register_pattern("conv2d", OpPattern.OUT_ELEMWISE_FUSABLE)
 
@@ -166,9 +184,9 @@ def compute_contrib_conv2d_NCHWc(attrs, inputs, _):
     dilation = attrs.get_int_tuple("dilation")
     out_channel = attrs.get_int("channels")
     groups = attrs.get_int("groups")
-    layout = attrs.get_string("layout")
-    out_layout = attrs.get_string("out_layout")
-    out_dtype = attrs.get_string("out_dtype")
+    layout = attrs.get_str("layout")
+    out_layout = attrs.get_str("out_layout")
+    out_dtype = attrs.get_str("out_dtype")
     out_dtype = inputs[0].dtype if out_dtype == "same" else out_dtype
     if layout == "NCHW":
         _, in_channel, _, _ = get_const_tuple(inputs[0].shape)
@@ -227,8 +245,8 @@ def compute_contrib_conv2d_winograd_without_weight_transform(attrs, inputs, _):
     strides = attrs.get_int_tuple("strides")
     dilation = attrs.get_int_tuple("dilation")
     groups = attrs.get_int("groups")
-    layout = attrs.get_string("layout")
-    out_dtype = attrs.get_string("out_dtype")
+    layout = attrs.get_str("layout")
+    out_dtype = attrs.get_str("out_dtype")
     tile_size = attrs.get_int("tile_size")
     out_dtype = inputs[0].dtype if out_dtype == "same" else out_dtype
     assert dilation == (1, 1), "Do not support dilate now"
@@ -262,7 +280,7 @@ def compute_conv2d_transpose(attrs, inputs, _):
     strides = attrs.get_int_tuple("strides")
     dilation = attrs.get_int_tuple("dilation")
     groups = attrs.get_int("groups")
-    out_dtype = attrs.get_string("out_dtype")
+    out_dtype = attrs.get_str("out_dtype")
     layout = attrs["layout"]
     out_dtype = inputs[0].dtype if out_dtype == "same" else out_dtype
 

python/tvm/attrs.py

@@ -33,6 +33,45 @@ class Attrs(NodeBase):
         for field in fields:
             yield field.name
 
+    def get_int_tuple(self, key):
+        """Get a python int tuple of a key
+
+        Parameters
+        ----------
+        key: str
+
+        Returns
+        -------
+        value: Tuple of int
+        """
+        return tuple(x.value for x in self.__getattr__(key))
+
+    def get_int(self, key):
+        """Get a python int value of a key
+
+        Parameters
+        ----------
+        key: str
+
+        Returns
+        -------
+        value: int
+        """
+        return self.__getattr__(key)
+
+    def get_str(self, key):
+        """Get a python int value of a key
+
+        Parameters
+        ----------
+        key: str
+
+        Returns
+        -------
+        value: int
+        """
+        return self.__getattr__(key)
+
     def __getitem__(self, item):
         return self.__getattr__(item)
 

python/tvm/relay/build_module.py

@@ -119,7 +119,7 @@ def _bind_params_by_name(func, params):
     return expr.bind(func, bind_dict)
 
 
-def optimize(func, params=None):
+def optimize(func, target, params=None):
     """Perform target invariant optimizations.
 
     Parameters
@@ -127,6 +127,9 @@ def optimize(func, params=None):
     func : tvm.relay.Function
         The input to optimization.
 
+    target: :any:`tvm.target.Target`
+        The optimization target. Some optimization passes are target specific.
+
     params : Optional[Dict[str, tvm.nd.NDArray]]
         Input parameters to the graph that do not change
         during inference time. used for constant folding.
@@ -164,7 +167,11 @@ def optimize(func, params=None):
         func = ir_pass.infer_type(func)
         func = ir_pass.canonicalize_ops(func)
         func = ir_pass.infer_type(func)
-        func = ir_pass.alter_op_layout(func)
+        with target:
+            func = ir_pass.alter_op_layout(func)
+
+    if cfg.pass_enabled("FoldConstant"):
+        func = ir_pass.fold_constant(func)
 
     return func
 
@@ -222,7 +229,7 @@ def build(func,
     cfg = BuildConfig.current
 
     with tophub_context:
-        func = optimize(func, params)
+        func = optimize(func, target, params)
         # Fuse ops before running code gen
         func = ir_pass.infer_type(func)
         func = ir_pass.fuse_ops(func, cfg.opt_level)

python/tvm/relay/frontend/mxnet.py

@@ -72,9 +72,9 @@ def _mx_conv2d(inputs, attrs):
     channel_axis = _get_channel_axis(data_layout, "conv2d")
 
     if "kernel_layout" in attrs.attrs:
-        weight_layout = attrs.get_str("kernel_layout")
+        kernel_layout = attrs.get_str("kernel_layout")
     else:
-        weight_layout = "HWIO" if data_layout == "NHWC" else "OIHW"
+        kernel_layout = "HWIO" if data_layout == "NHWC" else "OIHW"
 
     new_attrs = {}
     new_attrs["channels"] = attrs.get_int("num_filter")
@@ -84,7 +84,7 @@ def _mx_conv2d(inputs, attrs):
     new_attrs["dilation"] = attrs.get_int_tuple("dilate", (1, 1))
     new_attrs["groups"] = attrs.get_int("num_group", 1)
     new_attrs["data_layout"] = data_layout
-    new_attrs["weight_layout"] = weight_layout
+    new_attrs["kernel_layout"] = kernel_layout
     use_bias = not attrs.get_bool("no_bias", False)
     res = _op.nn.conv2d(inputs[0], inputs[1], **new_attrs)
     if use_bias:
@@ -103,9 +103,9 @@ def _mx_conv2d_transpose(inputs, attrs):
     channel_axis = _get_channel_axis(data_layout, "conv2d_transpose")
 
     if "kernel_layout" in attrs.attrs:
-        weight_layout = attrs.get_str("kernel_layout")
+        kernel_layout = attrs.get_str("kernel_layout")
     else:
-        weight_layout = "HWIO" if data_layout == "NHWC" else "OIHW"
+        kernel_layout = "HWIO" if data_layout == "NHWC" else "OIHW"
 
     new_attrs = {}
     new_attrs["channels"] = attrs.get_int("num_filter")
@@ -116,7 +116,7 @@ def _mx_conv2d_transpose(inputs, attrs):
     new_attrs["dilation"] = attrs.get_int_tuple("dilate", (1, 1))
     new_attrs["groups"] = attrs.get_int("num_group", 1)
     new_attrs["data_layout"] = data_layout
-    new_attrs["weight_layout"] = weight_layout
+    new_attrs["kernel_layout"] = kernel_layout
     use_bias = not attrs.get_bool("no_bias", False)
     res = _op.nn.conv2d_transpose(inputs[0], inputs[1], **new_attrs)
 

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

@@ -55,7 +55,7 @@ def compute_conv2d(attrs, inputs, out_type, target):
     dilation = get_const_tuple(attrs.dilation)
     groups = attrs.groups
     layout = attrs.data_layout
-    weight_layout = attrs.weight_layout
+    kernel_layout = attrs.kernel_layout
     out_dtype = attrs.out_dtype
     out_dtype = (inputs[0].dtype if (out_dtype == "same" or out_dtype == "")
                  else out_dtype)
@@ -70,13 +70,13 @@ def compute_conv2d(attrs, inputs, out_type, target):
             inputs[0], inputs[1], strides, padding,
             dilation, layout, out_dtype=out_dtype)
     elif layout == "NCHW" and \
-         weight_layout == "OIHW" and \
+         kernel_layout == "OIHW" and \
          get_const_int(inputs[1].shape[0]) == groups and \
          get_const_int(inputs[1].shape[1]) == 1:
         out = topi.nn.depthwise_conv2d_nchw(
             inputs[0], inputs[1], strides, padding, dilation, out_dtype=out_dtype)
     elif layout == "NHWC" and \
-         weight_layout == "HWOI" and\
+         kernel_layout == "HWOI" and\
          get_const_int(inputs[1].shape[2]) == groups and \
          get_const_int(inputs[1].shape[3]) == 1:
         out = topi.nn.depthwise_conv2d_nhwc(
@@ -91,7 +91,7 @@ def schedule_conv2d(attrs, outs, target):
     """Schedule definition of conv2d"""
     groups = attrs.groups
     layout = attrs.data_layout
-    kernel_layout = attrs.weight_layout
+    kernel_layout = attrs.kernel_layout
     with target:
         if groups == 1 and layout == "NCHW":
             return topi.generic.schedule_conv2d_nchw(outs)
@@ -111,7 +111,8 @@ def schedule_conv2d(attrs, outs, target):
 @reg.register_alter_op_layout("nn.conv2d")
 def alter_op_layout_conv2d(attrs, inputs, tinfos):
     """Alternate the layout of conv2d"""
-    return None
+    from ... import op
+    return topi.nn.conv2d_alter_layout(attrs, inputs, tinfos, op)
 
 reg.register_pattern("nn.conv2d", OpPattern.OUT_ELEMWISE_FUSABLE)
 
@@ -249,7 +250,7 @@ def schedule_l2_normalize(attrs, outs, target):
 
 reg.register_pattern("nn.l2_normalize", OpPattern.OUT_ELEMWISE_FUSABLE)
 
-# Upsampling
+# upsampling
 reg.register_schedule("nn.upsampling", reg.schedule_injective)
 def schedule_upsampling(_, outs, target):
     """Schedule definition of upsampling"""
@@ -257,3 +258,50 @@ def schedule_upsampling(_, outs, target):
         return topi.generic.schedule_injective(outs)
 # pad
 reg.register_schedule("nn.pad", schedule_broadcast)
+
+# winograd related operators
+@reg.register_compute("nn.contrib_conv2d_winograd_without_weight_transform")
+def compute_contrib_conv2d_winograd_without_weight_transform(attrs, inputs, out_dtype, target):
+    """Compute definition of conv2d_winograd_without_weight_transform"""
+    # pylint: disable=assignment-from-no-return
+    padding = attrs.get_int_tuple("padding")
+    strides = attrs.get_int_tuple("strides")
+    dilation = attrs.get_int_tuple("dilation")
+    groups = attrs.get_int("groups")
+    data_layout = attrs.get_str("data_layout")
+    out_dtype = attrs.get_str("out_dtype")
+    tile_size = attrs.get_int("tile_size")
+    out_dtype = inputs[0].dtype if out_dtype == "" else out_dtype
+    assert dilation == (1, 1), "Do not support dilate now"
+    assert groups == 1, "Do not supoort arbitrary group number"
+
+    out = topi.nn.conv2d_winograd_without_weight_transform(
+        inputs[0], inputs[1], strides, padding, dilation, data_layout,
+        out_dtype, tile_size)
+
+    return [out]
+
+@reg.register_schedule("nn.contrib_conv2d_winograd_without_weight_transform")
+def schedule_contrib_conv2d_winograd_without_weight_transform(attrs, outs, target):
+    """Schedule definition of conv2d_winograd_without_weight_transform"""
+    with target:
+        return topi.generic.schedule_conv2d_winograd_without_weight_transform(outs)
+
+reg.register_pattern("nn.contrib_conv2d_winograd_without_weight_transform",
+                     OpPattern.OUT_ELEMWISE_FUSABLE)
+
+
+@reg.register_compute("nn.contrib_conv2d_winograd_weight_transform")
+def compute_contrib_conv2d_winograd_weight_transform(attrs, inputs, out_dtype, target):
+    """Compute definition of contrib_conv2d_winograd_weight_transform"""
+    out = topi.nn.conv2d_winograd_weight_transform(inputs[0], attrs.get_int('tile_size'))
+    return [out]
+
+@reg.register_schedule("nn.contrib_conv2d_winograd_weight_transform")
+def schedule_contrib_conv2d_winograd_weight_transform(attrs, outs, target):
+    """Schedule definition of contrib_conv2d_winograd_weight_transform"""
+    with target:
+        return topi.generic.schedule_conv2d_winograd_weight_transform(outs)
+
+reg.register_pattern("nn.contrib_conv2d_winograd_weight_transform",
+                     OpPattern.OUT_ELEMWISE_FUSABLE)

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

@@ -13,7 +13,7 @@ def conv2d(data,
            channels=None,
            kernel_size=None,
            data_layout="NCHW",
-           weight_layout="OIHW",
+           kernel_layout="OIHW",
            out_layout="",
            out_dtype=""):
     r"""2D convolution.
@@ -23,7 +23,7 @@ def conv2d(data,
 
 
     In the default case, where the data_layout is `NCHW`
-    and weight_layout is `OIHW`, conv2d takes in
+    and kernel_layout is `OIHW`, conv2d takes in
     a data Tensor with shape `(batch_size, in_channels, height, width)`,
     and a weight Tensor with shape `(channels, in_channels, kernel_size[0], kernel_size[1])`
     to produce an output Tensor with the following rule:
@@ -70,7 +70,7 @@ def conv2d(data,
     data_layout : str, optional
         Layout of the input.
 
-    weight_layout : str, optional
+    kernel_layout : str, optional
         Layout of the weight.
 
     out_layout : str, optional
@@ -86,7 +86,7 @@ def conv2d(data,
     """
     return _make.conv2d(data, weight, strides, padding, dilation,
                         groups, channels, kernel_size, data_layout,
-                        weight_layout, out_layout, out_dtype)
+                        kernel_layout, out_layout, out_dtype)
 
 
 def conv2d_transpose(data,
@@ -98,7 +98,7 @@ def conv2d_transpose(data,
                      channels=None,
                      kernel_size=None,
                      data_layout="NCHW",
-                     weight_layout="OIHW",
+                     kernel_layout="OIHW",
                      output_padding=(0, 0),
                      out_dtype=""):
     """Two dimensional trnasposed convolution operator.
@@ -126,7 +126,7 @@ def conv2d_transpose(data,
     data_layout : str, optional
         Layout of the input.
 
-    weight_layout : str, optional
+    kernel_layout : str, optional
         Layout of the weight.
 
     output_padding : Tuple[int], optional
@@ -142,7 +142,7 @@ def conv2d_transpose(data,
     """
     return _make.conv2d_transpose(data, weight, strides, padding, dilation,
                                   groups, channels, kernel_size, data_layout,
-                                  weight_layout, output_padding, out_dtype)
+                                  kernel_layout, output_padding, out_dtype)
 
 
 def softmax(data, axis=-1):
@@ -765,3 +765,96 @@ def batch_norm(data,
                               center,
                               scale)
     return TupleWrapper(result, 3)
+
+
+def contrib_conv2d_winograd_without_weight_transform(data,
+                                                     weight,
+                                                     tile_size,
+                                                     strides=(1, 1),
+                                                     padding=(0, 0),
+                                                     dilation=(1, 1),
+                                                     groups=1,
+                                                     channels=None,
+                                                     kernel_size=None,
+                                                     data_layout="NCHW",
+                                                     kernel_layout="OIHW",
+                                                     out_layout="",
+                                                     out_dtype=""):
+    r"""2D convolution with winograd algorithm.
+
+    The basic parameters are the same as the ones in vanilla conv2d.
+    It assumes the weight is pre-transformed by nn.contrib_conv2d_winograd_weight_transform
+
+    Parameters
+    ----------
+    data : tvm.relay.Expr
+        The input data to the operator.
+
+    weight : tvm.relay.Expr
+        The weight expressions.
+
+    tile_size : int
+        The Tile size of winograd. E.g. 2 for F(2x2, 3x3) and 4 for F(4x4, 3x3)
+
+    strides : tuple of int, optional
+        The strides of convoltution.
+
+    padding : tuple of int, optional
+        The padding of convolution on both sides of inputs before convolution.
+
+    dilation : tuple of int, optional
+        Specifies the dilation rate to be used for dilated convolution.
+
+    groups : int, optional
+        Number of groups for grouped convolution.
+
+    channels : int, optional
+        Number of output channels of this convolution.
+
+    kernel_size : tuple of int, optional
+        The spatial of the convolution kernel.
+
+    data_layout : str, optional
+        Layout of the input.
+
+    kernel_layout : str, optional
+        Layout of the weight.
+
+    out_layout : str, optional
+        Layout of the output, by default, out_layout is the same as data_layout
+
+    out_dtype : str, optional
+        Specifies the output data type for mixed precision conv2d.
+
+    Returns
+    -------
+    result : tvm.relay.Expr
+        The computed result.
+    """
+    return _make.contrib_conv2d_winograd_without_weight_transform(
+        data, weight, tile_size, strides, padding, dilation,
+        groups, channels, kernel_size, data_layout,
+        kernel_layout, out_layout, out_dtype)
+
+
+def contrib_conv2d_winograd_weight_transform(weight,
+                                             tile_size):
+    r"""Weight Transformation part for 2D convolution with winograd algorithm.
+
+    We separate this as a single op to enable pre-compute for inference.
+    Use this together with nn.contrib_conv2d_winograd_without_weight_transform
+
+    Parameters
+    ----------
+    weight : tvm.relay.Expr
+        The weight expressions.
+
+    tile_size : int
+        The Tile size of winograd. E.g. 2 for F(2x2, 3x3) and 4 for F(4x4, 3x3)
+
+    Returns
+    -------
+    result : tvm.relay.Expr
+        The computed result.
+    """
+    return _make.contrib_conv2d_winograd_weight_transform(weight, tile_size)

python/tvm/relay/op/op_attrs.py

@@ -5,10 +5,20 @@ from ..base import register_relay_attr_node
 
 @register_relay_attr_node
 class Conv2DAttrs(Attrs):
-    """Attribute of a Convolution Operator"""
+    """Attribute of nn.conv2d"""
+    pass
+
+@register_relay_attr_node
+class Conv2DWinogradAttrs(Attrs):
+    """Attribute of nn.contrib_conv2d_winograd_without_weight_transform"""
+    pass
+
+@register_relay_attr_node
+class Conv2DWinogradWeightTransformAttrs(Attrs):
+    """Attribute of nn.contrib_conv2d_winograd_weight_transform"""
     pass
 
 @register_relay_attr_node
 class GlobalPool2DAttrs(Attrs):
-    """Attribute of a Global 2D Pooling Operator"""
+    """Attribute of nn.global_pool"""
     pass

python/tvm/relay/testing/__init__.py

@@ -11,4 +11,6 @@ from . import inception_v3
 from . import squeezenet
 from . import vgg
 from . import densenet
+
 from .config import ctx_list
+from .init import create_workload

src/relay/op/layout.h

@@ -60,7 +60,7 @@ class Layout : public NodeRef {
   Layout() : Layout("__undef__") {} // NOLINT(*)
 
   /*! \brief construct from a string */
-  Layout(const char* str) : Layout(std::string(str)) {} // NOLINT(*)
+  Layout(const char* name) : Layout(std::string(name)) {} // NOLINT(*)
 
   /*!
    * \brief construct from a string.
@@ -70,11 +70,64 @@ class Layout : public NodeRef {
    *        indicates the split dimension.
    *        return undefined layout if "__undef__" is passed.
    */
-  Layout(const std::string& layout) { // NOLINT(*)
-    if (layout.length() != 0) {
-      Parse(layout);
-    } else {
-      Parse("__undef__");
+  Layout(const std::string& name) { // NOLINT(*)
+    node_ = make_node<LayoutNode>();
+
+    std::vector<uint32_t> superdim_pos(kUniqueDim, -1);
+    std::vector<uint32_t> subdim_pos(kUniqueDim, -1);
+    std::vector<uint32_t> subdim_size(kUniqueDim, -1);
+    std::vector<char> layout_simplified;
+
+    if (name != "__undef__") {  // parse layout string
+      int32_t factor = 0;
+      uint32_t curr = 0;
+      for (size_t i = 0; i < name.size(); ++i) {
+        const LayoutDim c = name.at(i);
+        if (IsSuperdim(c)) {
+          int pos = c - 'A';
+          CHECK_EQ(factor, 0) << "Invalid layout " << name
+                              << ": invalid factor size " << factor
+                              << " before dimension " << c;
+          CHECK_EQ(superdim_pos[pos], -1) << "Invalid layout " << name
+                                          << ": duplicate dimension " << c;
+          superdim_pos[pos] = curr++;
+          layout_simplified.push_back(c);
+        } else if (IsSubdim(c)) {
+          int pos = c - 'a';
+          CHECK_GT(factor, 0) << "Invalid layout " << name << ": invalid factor size "
+                              << factor << " for dimension " << c;
+          CHECK_EQ(subdim_pos[pos], -1) << "Invalid layout " << name
+                                        << ": duplicate dimension " << c;
+          CHECK_EQ(subdim_size[pos], -1) << "Invalid layout " << name
+                                         << ": duplicate dimension " << c;
+          subdim_pos[pos] = curr++;
+          subdim_size[pos] = factor;
+          layout_simplified.push_back(c);
+          factor = 0;
+        } else if (c >= '0' && c <= '9') {
+          CHECK(factor >= 0) << "Invalid layout " << name << ": _ is adjacent to a number.";
+          factor = factor * 10 + c - '0';
+        } else {
+          LOG(FATAL) << "Invalid layout " << name;
+        }
+      }
+      for (LayoutDim dim : layout_simplified) {
+        CHECK(IsSuperdim(dim) || superdim_pos[dim-'a'] >= 0)
+          << "Invalid layout " << name << ": missing axis "
+          << static_cast<char>(dim - 'a' + 'A');
+      }
+    }
+
+    LayoutNode *node = operator->();
+    node->name = name;
+
+    for (uint32_t i = 0; i < kUniqueDim; ++i) {
+      node->superdim_pos.push_back(superdim_pos[i]);
+      node->subdim_pos.push_back(subdim_pos[i]);
+      node->subdim_size.push_back(subdim_size[i]);
+    }
+    for (LayoutDim dim : layout_simplified) {
+      node->layout_simplified.push_back(dim);
     }
   }
 
@@ -177,7 +230,6 @@ class Layout : public NodeRef {
     const Array<Integer>& layout_simplified = operator->()->layout_simplified;
     if (pos > ndim()) return Layout::Undef();
     if (pos + len > ndim()) len = ndim() - pos;
-    if (len == 0) return Layout::Undef();
     std::ostringstream new_layout;
     for (size_t i = pos; i < pos + len; ++i) {
       if (IsSubdim(layout_simplified[i]->value)) {
@@ -349,69 +401,6 @@ class Layout : public NodeRef {
   }
 
   using ContainerType = LayoutNode;
-
- private:
-  void Parse(const std::string &layout) {
-    node_ = make_node<LayoutNode>();
-
-    std::vector<uint32_t> superdim_pos(kUniqueDim, -1);
-    std::vector<uint32_t> subdim_pos(kUniqueDim, -1);
-    std::vector<uint32_t> subdim_size(kUniqueDim, -1);
-    std::vector<char> layout_simplified;
-
-    if (layout != "__undef__") {  // parse layout string
-      int32_t factor = 0;
-      uint32_t curr = 0;
-      for (size_t i = 0; i < layout.size(); ++i) {
-        const LayoutDim c = layout.at(i);
-        if (IsSuperdim(c)) {
-          int pos = c - 'A';
-          CHECK_EQ(factor, 0) << "Invalid layout " << layout
-                              << ": invalid factor size " << factor
-                              << " before dimension " << c;
-          CHECK_EQ(superdim_pos[pos], -1) << "Invalid layout " << layout
-                                          << ": duplicate dimension " << c;
-          superdim_pos[pos] = curr++;
-          layout_simplified.push_back(c);
-        } else if (IsSubdim(c)) {
-          int pos = c - 'a';
-          CHECK_GT(factor, 0) << "Invalid layout " << layout << ": invalid factor size "
-                              << factor << " for dimension " << c;
-          CHECK_EQ(subdim_pos[pos], -1) << "Invalid layout " << layout
-                                        << ": duplicate dimension " << c;
-          CHECK_EQ(subdim_size[pos], -1) << "Invalid layout " << layout
-                                         << ": duplicate dimension " << c;
-          subdim_pos[pos] = curr++;
-          subdim_size[pos] = factor;
-          layout_simplified.push_back(c);
-          factor = 0;
-        } else if (c >= '0' && c <= '9') {
-          CHECK(factor >= 0) << "Invalid layout " << layout << ": _ is adjacent to a number.";
-          factor = factor * 10 + c - '0';
-        } else {
-          LOG(FATAL) << "Invalid layout " << layout;
-        }
-      }
-      CHECK(!layout_simplified.empty()) << "Invalid layout " << layout;
-      for (LayoutDim dim : layout_simplified) {
-        CHECK(IsSuperdim(dim) || superdim_pos[dim-'a'] >= 0)
-          << "Invalid layout " << layout << ": missing axis "
-          << static_cast<char>(dim - 'a' + 'A');
-      }
-    }
-
-    LayoutNode *node = operator->();
-    node->name = layout;
-
-    for (uint32_t i = 0; i < kUniqueDim; ++i) {
-      node->superdim_pos.push_back(superdim_pos[i]);
-      node->subdim_pos.push_back(subdim_pos[i]);
-      node->subdim_size.push_back(subdim_size[i]);
-    }
-    for (LayoutDim dim : layout_simplified) {
-      node->layout_simplified.push_back(dim);
-    }
-  }
 };
 
 /*!

src/relay/pass/alter_op_layout.cc

@@ -166,7 +166,7 @@ Call CallAlter(const Call& ref_call,
   }
   if (!modified) {
     new_e = CallNode::make(ref_call->op, new_args,
-                           ref_call->attrs, ref_call->type_args);
+                           ref_call->attrs);
   }
 
   const CallNode *new_call = new_e.as<CallNode>();
@@ -184,30 +184,35 @@ Expr AlterOpLayoutRewrite(const Call &ref_call,
   // NOTE: discard the "const" qualifier
   TransformMemorizer memorizer = Downcast<TransformMemorizer>(ctx);
 
-  // fill incomplete state and expand tuple
-  for (auto new_arg : new_args) {
-    auto push_back_one_arg = [&](Expr arg) {
-      // We always expect LayoutAlternatedExpr.
-      // This is used to convert the normal Expr to LayoutAlternatedExpr.
-      if (const LayoutAlternatedExprNode *inp = arg.as<LayoutAlternatedExprNode>()) {
-        inputs.push_back(GetRef<LayoutAlternatedExpr>(inp));
-        normal_new_args.push_back(inp->value);
-      } else {
-        auto inode = make_node<LayoutAlternatedExprNode>();
-        inode->value = arg;
-        inode->memorizer = memorizer;
-        inputs.push_back(LayoutAlternatedExpr(inode));
-        normal_new_args.push_back(arg);
-      }
-    };
+  // fill incomplete state and flatten tuple
+  auto push_back_one_arg = [&inputs, memorizer](Expr arg) {
+    // We always expect LayoutAlternatedExpr.
+    // This is used to convert the normal Expr to LayoutAlternatedExpr.
+    if (const LayoutAlternatedExprNode *inp = arg.as<LayoutAlternatedExprNode>()) {
+      inputs.push_back(GetRef<LayoutAlternatedExpr>(inp));
+      return inp->value;
+    } else {
+      auto inode = make_node<LayoutAlternatedExprNode>();
+      inode->value = arg;
+      inode->memorizer = memorizer;
+      inputs.push_back(LayoutAlternatedExpr(inode));
+      return arg;
+    }
+  };
 
+  for (auto new_arg : new_args) {
+    // NOTE: do not support nested tuple
     if (new_arg->is_type<TupleNode>()) {
       Tuple tuple_new_arg = Downcast<Tuple>(new_arg);
+      std::vector<Expr> fields;
       for (auto x : tuple_new_arg->fields) {
-        push_back_one_arg(x);
+        Expr tmp = push_back_one_arg(x);
+        fields.push_back(tmp);
       }
+      normal_new_args.push_back(TupleNode::make(fields));
     } else {
-      push_back_one_arg(new_arg);
+      Expr tmp = push_back_one_arg(new_arg);
+      normal_new_args.push_back(tmp);
     }
   }
 
@@ -219,7 +224,7 @@ Expr AlterOpLayoutRewrite(const Call &ref_call,
   }
 
   for (auto arg : ref_call->args) {
-    if (arg->is_type<TupleNode>()) {  // expand tuple
+    if (arg->is_type<TupleNode>()) {  // flatten tuple
       Tuple tuple_arg = Downcast<Tuple>(arg);
       for (auto x : tuple_arg->fields) {
         input_shapes.push_back(x->type_as<TensorTypeNode>()->shape);
@@ -263,17 +268,30 @@ Expr AlterOpLayoutRewrite(const Call &ref_call,
 
   // if (new_in != new_in2): insert transform (new_in -> new_in2)
   Array<Expr> transformed_args;
-  for (size_t i = 0; i < inputs.size(); ++i) {
-    transformed_args.push_back(memorizer.Transform(new_call->args[i], new_in[i], new_in2[i]));
+  size_t pt = 0;
+  for (auto arg : new_call->args) {
+    if (arg->is_type<TupleNode>()) {  // unflatten tuple
+      Tuple tuple_arg = Downcast<Tuple>(arg);
+      std::vector<Expr> transformed_tuple_arg;
+      for (auto arg_item : tuple_arg->fields) {
+          transformed_tuple_arg.push_back(
+                  memorizer.Transform(arg_item, new_in[pt], new_in2[pt]));
+          pt++;
+      }
+      transformed_args.push_back(TupleNode::make(transformed_tuple_arg));
+    } else {
+      transformed_args.push_back(
+              memorizer.Transform(arg, new_in[pt], new_in2[pt]));
+      pt++;
+    }
   }
+  CHECK_EQ(pt, inputs.size());
 
   // state[node] = (old_out, new_out)
-  CHECK(ref_call->checked_type_.defined())
-    << "Call infer_type pass before alter_op_layout pass";
-
+  // (handle tuple output)
   if (ref_call->checked_type()->is_type<TupleTypeNode>()) {
     Expr tuple_output = CallNode::make(new_call->op, transformed_args,
-                                       new_call->attrs, new_call->type_args);
+                                       new_call->attrs);
     Array<Expr> fields;
     for (size_t i = 0; i < new_out.size(); ++i) {
       auto rnode = make_node<LayoutAlternatedExprNode>();
@@ -288,7 +306,7 @@ Expr AlterOpLayoutRewrite(const Call &ref_call,
     auto rnode = make_node<LayoutAlternatedExprNode>();
     CHECK_EQ(new_out.size(), 1);
     rnode->value = CallNode::make(new_call->op, transformed_args,
-                                  new_call->attrs, new_call->type_args);
+                                  new_call->attrs);
     rnode->old_layout = old_out[0];
     rnode->new_layout = new_out[0];
     rnode->memorizer = memorizer;
@@ -296,6 +314,9 @@ Expr AlterOpLayoutRewrite(const Call &ref_call,
   }
 }
 
+// Limiations:
+// 1. the altered op should have the same number of arguments as the previous one
+// 2. do not support nested tuple arguments
 TVM_REGISTER_API("relay._ir_pass.AlterOpLayout")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
   TransformMemorizer transformMemorizer(make_node<TransformMemorizerNode>());

src/relay/pass/combine_parallel_conv2d.cc

@@ -91,13 +91,13 @@ class BranchGroupFinder : private ExprVisitor {
     CHECK(attrs_b);
     const auto* tweight_a = a->args[1]->type_as<TensorTypeNode>();
     const auto* tweight_b = b->args[1]->type_as<TensorTypeNode>();
-    const auto shape_a = ConvertLayout(tweight_a->shape, attrs_a->weight_layout, kOIHW);
-    const auto shape_b = ConvertLayout(tweight_b->shape, attrs_b->weight_layout, kOIHW);
+    const auto shape_a = ConvertLayout(tweight_a->shape, attrs_a->kernel_layout, kOIHW);
+    const auto shape_b = ConvertLayout(tweight_b->shape, attrs_b->kernel_layout, kOIHW);
 
     return eq(attrs_a->strides, attrs_b->strides) && eq(attrs_a->padding, attrs_b->padding) &&
            eq(attrs_a->dilation, attrs_b->dilation) && eq(attrs_a->groups, attrs_b->groups) &&
            eq(attrs_a->data_layout, attrs_b->data_layout) &&
-           eq(attrs_a->weight_layout, attrs_b->weight_layout) &&
+           eq(attrs_a->kernel_layout, attrs_b->kernel_layout) &&
            eq(attrs_a->out_dtype, attrs_b->out_dtype) &&
            eq(attrs_a->out_layout, attrs_b->out_layout) && eq(shape_a[2], shape_b[2]) &&
            eq(shape_a[3], shape_b[3]);
@@ -159,7 +159,7 @@ class ParallelConv2DCombiner {
       auto channels = GetConv2DSuperChannelsDim(conv2d);
       num_filters += channels;
     }
-    auto index = branches[0][0]->attrs.as<Conv2DAttrs>()->weight_layout.find('O');
+    auto index = branches[0][0]->attrs.as<Conv2DAttrs>()->kernel_layout.find('O');
     CHECK_NE(index, std::string::npos);
     return std::make_tuple(MakeConcatenate(TupleNode::make(weights), index),
                            MakeConstScalar(Int(32), num_filters));
@@ -182,7 +182,7 @@ class ParallelConv2DCombiner {
     new_attrs->groups = attrs->groups;
     new_attrs->kernel_size = attrs->kernel_size;
     new_attrs->data_layout = attrs->data_layout;
-    new_attrs->weight_layout = attrs->weight_layout;
+    new_attrs->kernel_layout = attrs->kernel_layout;
     new_attrs->out_layout = attrs->out_layout;
     new_attrs->out_dtype = attrs->out_dtype;
     new_attrs->channels = new_channels;

src/relay/pass/fold_scale_axis.cc

@@ -384,7 +384,7 @@ Array<AxesSet> Conv2DForwardPrep(const Call& call, AxesSet out) {
   const auto* param = call->attrs.as<Conv2DAttrs>();
   CHECK(param != nullptr);
   Layout data_layout(param->data_layout);
-  Layout weight_layout(param->weight_layout);
+  Layout kernel_layout(param->kernel_layout);
   int c_big_axis = data_layout.Indexof('C');
   int c_small_axis = data_layout.Indexof('c');
 
@@ -397,8 +397,8 @@ Array<AxesSet> Conv2DForwardPrep(const Call& call, AxesSet out) {
   //
   // only handle depthwise or full conv2d.
   // TODO(tvm-team) handle grouped conv by reshape + bcast
-  bool is_depthwise_conv2d = IsDepthwiseConv2D(call, param, weight_layout);
-  if (weight_layout.Indexof('i') < 0 &&
+  bool is_depthwise_conv2d = IsDepthwiseConv2D(call, param, kernel_layout);
+  if (kernel_layout.Indexof('i') < 0 &&
       c_small_axis < 0 &&
       (param->groups == 1 || is_depthwise_conv2d)) {
     data_axes = {c_big_axis};
@@ -418,19 +418,19 @@ Expr Conv2DForwardRewrite(const Call& ref_call,
   const auto* param = ref_call->attrs.as<Conv2DAttrs>();
   CHECK(param != nullptr);
   Layout data_layout(param->data_layout);
-  Layout weight_layout(param->weight_layout);
+  Layout kernel_layout(param->kernel_layout);
   int c_big_axis = data_layout.Indexof('C');
   CHECK_GE(c_big_axis, 0);
   // For now, we only support simple pattern (no folded weight/data)
   // TODO(tvm-team) support general data layout
-  CHECK_EQ(weight_layout.Indexof('i'), -1);
+  CHECK_EQ(kernel_layout.Indexof('i'), -1);
   CHECK(sdata->axes.size() == 1 &&
         c_big_axis == sdata->axes[0]->value);
-  int big_oc_axis = weight_layout.Indexof('O');
-  int big_ic_axis = weight_layout.Indexof('I');
+  int big_oc_axis = kernel_layout.Indexof('O');
+  int big_ic_axis = kernel_layout.Indexof('I');
 
   // Check it must be depthwise or full conv2d.
-  bool is_depthwise_conv2d = IsDepthwiseConv2D(ref_call, param, weight_layout);
+  bool is_depthwise_conv2d = IsDepthwiseConv2D(ref_call, param, kernel_layout);
   CHECK(param->groups == 1 || is_depthwise_conv2d);
 
   Expr weight = new_args[1];
@@ -438,11 +438,11 @@ Expr Conv2DForwardRewrite(const Call& ref_call,
   // match the ic_axis
   if (is_depthwise_conv2d) {
     Expr scale = ExpandBiasToMatchAxis(
-        sdata->scale, weight_layout.ndim(), {big_oc_axis});
+        sdata->scale, kernel_layout.ndim(), {big_oc_axis});
     weight = Multiply(weight, scale);
   } else {
     Expr scale = ExpandBiasToMatchAxis(
-        sdata->scale, weight_layout.ndim(), {big_ic_axis});
+        sdata->scale, kernel_layout.ndim(), {big_ic_axis});
     weight = Multiply(weight, scale);
   }
   // return transformed conv2d
@@ -799,11 +799,8 @@ RELAY_REGISTER_OP("multiply")
 AxesSet Conv2DBackwardPrep(const Call& call, const Array<AxesSet>& in_axes) {
   const auto* param = call->attrs.as<Conv2DAttrs>();
   CHECK(param != nullptr);
-  Layout out_layout(param->out_layout);
-  if (!out_layout.defined()) {
-    out_layout = Layout(param->data_layout);
-  }
-  Layout weight_layout(param->weight_layout);
+  Layout kernel_layout(param->kernel_layout);
+  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
   int c_big_axis = out_layout.Indexof('C');
   int c_small_axis = out_layout.Indexof('c');
 
@@ -815,9 +812,9 @@ AxesSet Conv2DBackwardPrep(const Call& call, const Array<AxesSet>& in_axes) {
   //
   // only handle depthwise or full conv2d.
   // TODO(tvm-team) handle grouped conv by reshape + bcast
-  bool is_depthwise_conv2d = IsDepthwiseConv2D(call, param, weight_layout);
-  if (weight_layout.Indexof('o') < 0 &&
-      weight_layout.Indexof('i') < 0 &&
+  bool is_depthwise_conv2d = IsDepthwiseConv2D(call, param, kernel_layout);
+  if (kernel_layout.Indexof('o') < 0 &&
+      kernel_layout.Indexof('i') < 0 &&
       c_small_axis < 0 &&
       (param->groups == 1 || is_depthwise_conv2d)) {
     return {c_big_axis};
@@ -836,23 +833,20 @@ Expr Conv2DBackwardTransform(const Call& call,
   }
   const auto* param = call->attrs.as<Conv2DAttrs>();
   CHECK(param != nullptr);
-  Layout out_layout(param->out_layout);
-  if (!out_layout.defined()) {
-    out_layout = Layout(param->data_layout);
-  }
-  Layout weight_layout(param->weight_layout);
+  Layout kernel_layout(param->kernel_layout);
+  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
   int c_big_axis = out_layout.Indexof('C');
   CHECK_GE(c_big_axis, 0);
   // For now, we only support simple pattern (no folded weight/data)
   // TODO(tvm-team) support general data layout
-  CHECK_EQ(weight_layout.Indexof('o'), -1);
-  CHECK_EQ(weight_layout.Indexof('i'), -1);
+  CHECK_EQ(kernel_layout.Indexof('o'), -1);
+  CHECK_EQ(kernel_layout.Indexof('i'), -1);
   CHECK(axes.size() == 1 &&
         c_big_axis == axes[0]->value);
 
-  int big_oc_axis = weight_layout.Indexof('O');
+  int big_oc_axis = kernel_layout.Indexof('O');
   // Check it must be depthwise or full conv2d.
-  bool is_depthwise_conv2d = IsDepthwiseConv2D(call, param, weight_layout);
+  bool is_depthwise_conv2d = IsDepthwiseConv2D(call, param, kernel_layout);
   CHECK(param->groups == 1 || is_depthwise_conv2d);
 
   Expr data = transformer->Transform(
@@ -861,7 +855,7 @@ Expr Conv2DBackwardTransform(const Call& call,
       call->args[1], NullValue<AxesSet>(), NullValue<Expr>());
   // scale on input for deptwise.
   Expr wscale = ExpandBiasToMatchAxis(
-      scale, weight_layout.ndim(), {big_oc_axis});
+      scale, kernel_layout.ndim(), {big_oc_axis});
   weight = Multiply(weight, wscale);
   return CallNode::make(
       call->op, {data, weight}, call->attrs, call->type_args);

src/relay/pass/pattern_util.h

@@ -112,11 +112,11 @@ inline Expr ExpandBiasToMatchAxis(Expr bias,
  */
 inline bool IsDepthwiseConv2D(const Call& call,
                               const Conv2DAttrs* param,
-                              const Layout& weight_layout) {
+                              const Layout& kernel_layout) {
   static const Layout kOIHW("OIHW");
   auto wshape = ConvertLayout(
       call->args[1]->type_as<TensorTypeNode>()->shape,
-      weight_layout, kOIHW);
+      kernel_layout, kOIHW);
   return is_const_int(wshape[0], param->groups) &&
       is_const_int(wshape[1], 1);
 }
@@ -129,7 +129,7 @@ inline bool IsDepthwiseConv2D(const Call& call,
 inline int64_t GetConv2DSuperChannelsDim(const CallNode* call) {
     auto param = call->attrs.as<Conv2DAttrs>();
     auto tweight = call->args[1]->type_as<TensorTypeNode>();
-    auto index = param->weight_layout.find('O');
+    auto index = param->kernel_layout.find('O');
     CHECK_NE(index, std::string::npos);
     auto channels = as_const_int(tweight->shape[index]);
     return *channels;

tests/python/relay/test_op_level2.py

@@ -41,7 +41,7 @@ def test_conv2d_infer_type():
                         padding=(1, 1),
                         channels=16,
                         data_layout="NCHW4n4c",
-                        weight_layout="OIHW4o4i",
+                        kernel_layout="OIHW4o4i",
                         out_dtype="int32")
     yy = relay.ir_pass.infer_type(y)
     assert yy.checked_type ==  relay.TensorType(

tests/python/relay/test_pass_alter_op_layout.py

@@ -91,7 +91,7 @@ def test_alter_layout():
         data, weight = inputs
         new_attrs = dict(attrs)
         new_attrs['data_layout'] = 'NCHW16c'
-        new_attrs['weight_layout'] = 'OIHW16i'
+        new_attrs['kernel_layout'] = 'OIHW16i'
         return relay.nn.conv2d(data, weight, **new_attrs)
 
     def expected():
@@ -105,7 +105,7 @@ def test_alter_layout():
                             channels=64,
                             kernel_size=(3, 3),
                             padding=(1, 1),
-                            weight_layout="OIHW16i",
+                            kernel_layout="OIHW16i",
                             data_layout="NCHW16c")
         b = relay.expand_dims(bias, axis=1, num_newaxis=2)
         b = relay.layout_transform(b, "CHW", "CHW16c")
@@ -269,7 +269,7 @@ def test_alter_layout_broadcast_op():
         y = relay.Function(free_vars(y), y)
         return y
 
-    @register_alter_op_layout("nn.conv2d", level=102)
+    @register_alter_op_layout("nn.conv2d", level=105)
     def alter_conv2d(attrs, inputs, tinfos):
         data, weight = inputs
         new_attrs = dict(attrs)
@@ -305,6 +305,107 @@ def test_alter_layout_broadcast_op():
 
     assert(alpha_equal(a, b))
 
+def test_alter_layout_scalar():
+    """Test alternating the layout of a conv2d.
+    The layout of broadcast operators and the weight should be changed accordingly.
+    """
+    def before():
+        x = relay.var("x", shape=(1, 64, 56, 56))
+        weight = relay.var("weight")
+        y = relay.nn.conv2d(x, weight, channels=64, kernel_size=(3, 3), padding=(1, 1))
+        y = relay.add(y, relay.const(1, "float32"))
+        y = relay.Function(free_vars(y), y)
+        return y
+
+    @register_alter_op_layout("nn.conv2d", level=106)
+    def alter_conv2d(attrs, inputs, tinfos):
+        data, weight = inputs
+        new_attrs = dict(attrs)
+        new_attrs['data_layout'] = 'NCHW16c'
+        return relay.nn.conv2d(data, weight, **new_attrs)
+
+    def expected():
+        x = relay.var("x", shape=(1, 64, 56, 56))
+        w = relay.var("weight")
+
+        y = relay.layout_transform(x, "NCHW", "NCHW16c")
+        y = relay.nn.conv2d(y, w,
+                            channels=64,
+                            kernel_size=(3, 3),
+                            padding=(1, 1),
+                            data_layout="NCHW16c")
+        y = relay.add(y, relay.const(1.0, "float32"))
+
+        y = relay.layout_transform(y, "NCHW16c", "NCHW")
+        y = relay.Function(free_vars(y), y)
+        return y
+
+    a = before()
+    a = infer_type(a)
+    a = canonicalize_ops(a)
+    a = infer_type(a)
+    a = alter_op_layout(a)
+    a = infer_type(a)
+
+    b = expected()
+    b = infer_type(b)
+
+    assert(alpha_equal(a, b))
+
+def test_alter_layout_concatenate():
+    """ """
+    def before():
+        x = relay.var("x", shape=(1, 64, 56, 56))
+        weight1 = relay.var('weight1')
+        weight2 = relay.var('weight2')
+        y = relay.nn.conv2d(x, weight1,
+                            channels=32,
+                            kernel_size=(3, 3),
+                            padding=(1, 1))
+        y1 = relay.nn.conv2d(y, weight2,
+                             channels=32,
+                             kernel_size=(3, 3),
+                             padding=(1, 1))
+        ret = relay.concatenate([y, y1], axis=1)
+        y = relay.Function(free_vars(ret), ret)
+        return y
+
+    @register_alter_op_layout("nn.conv2d", level=107)
+    def alter_conv2d(attrs, inputs, tinfos):
+        data, weight = inputs
+        new_attrs = dict(attrs)
+        new_attrs['data_layout'] = 'NCHW16c'
+        return relay.nn.conv2d(data, weight, **new_attrs)
+
+    def expected():
+        x = relay.var("x", shape=(1, 64, 56, 56))
+        weight1 = relay.var('weight1')
+        weight2 = relay.var('weight2')
+        y = relay.layout_transform(x, "NCHW", "NCHW16c")
+        y = relay.nn.conv2d(y, weight1,
+                            channels=32,
+                            kernel_size=(3, 3),
+                            padding=(1, 1),
+                            data_layout="NCHW16c")
+        y1 = relay.nn.conv2d(y, weight2,
+                             channels=32,
+                             kernel_size=(3, 3),
+                             padding=(1, 1),
+                             data_layout='NCHW16c')
+        ret = relay.concatenate([y, y1], axis=1)
+        ret = relay.layout_transform(ret, "NCHW16c", "NCHW")
+        y = relay.Function(free_vars(ret), ret)
+        return y
+
+    a = before()
+    a = infer_type(a)
+    a = alter_op_layout(a)
+    a = infer_type(a)
+
+    b = expected()
+    b = infer_type(b)
+
+    assert(alpha_equal(a, b))
 
 if __name__ == "__main__":
     test_alter_op()
@@ -313,3 +414,5 @@ if __name__ == "__main__":
     test_alter_layout_dual_path()
     test_alter_layout_resnet()
     test_alter_layout_broadcast_op()
+    test_alter_layout_scalar()
+    test_alter_layout_concatenate()

tests/python/relay/test_pass_fold_scale_axis.py

@@ -67,14 +67,14 @@ def test_fold_fwd_dual_path():
                              channels=channels,
                              kernel_size=(3, 3),
                              data_layout="NHWC",
-                             weight_layout="HWIO",
+                             kernel_layout="HWIO",
                              groups=channels,
                              padding=(1, 1))
         y2 = relay.nn.conv2d(x, conv_weight,
                              channels=channels,
                              kernel_size=(3, 3),
                              data_layout="NHWC",
-                             weight_layout="HWIO",
+                             kernel_layout="HWIO",
                              groups=channels,
                              padding=(1, 1))
         z = relay.add(y1, y2)
@@ -90,7 +90,7 @@ def test_fold_fwd_dual_path():
                              channels=channels,
                              kernel_size=(3, 3),
                              data_layout="NHWC",
-                             weight_layout="HWIO",
+                             kernel_layout="HWIO",
                              groups=channels,
                              padding=(1, 1))
         y2 = relay.nn.conv2d(x,
@@ -98,7 +98,7 @@ def test_fold_fwd_dual_path():
                              channels=channels,
                              kernel_size=(3, 3),
                              data_layout="NHWC",
-                             weight_layout="HWIO",
+                             kernel_layout="HWIO",
                              groups=channels,
                              padding=(1, 1))
         z = relay.add(y1, y2)

topi/python/topi/arm_cpu/conv2d.py

@@ -523,9 +523,25 @@ def schedule_conv2d_winograd_without_weight_transform_(cfg, outs):
 
 ##### REGISTER ALTER OP LAYOUT #####
 @conv2d_alter_layout.register(["arm_cpu"])
-def _alter_conv2d_layout_arm(attrs, inputs, tinfos):
-    """Alter op layout for pre-computing kernel transformation"""
-    import nnvm.symbol as sym
+def _alter_conv2d_layout_arm(attrs, inputs, tinfos, F):
+    """Alter op layout for pre-computing kernel transformation
+
+    Parameters
+    ----------
+    attrs : nnvm.top.AttrDict or tvm.attrs.Attrs
+        Attributes of current convolution
+    inputs : nnvm.symbol or tvm.relay.Expr
+        Grouped input symbols
+    tinfos : list
+        Input shape and dtype
+    F: symbol
+        The context, can be either nnvm.sym or relay.op
+
+    Note
+    ----
+    Unlike other TOPI functions, this function operates on both graph level and operator level,
+    so we have to pass 'F' to make it support our two versions of graph IR, NNVM and Relay.
+    """
     copy_inputs = [s for s in inputs]
 
     new_attrs = {k: attrs[k] for k in attrs.keys()}
@@ -534,9 +550,11 @@ def _alter_conv2d_layout_arm(attrs, inputs, tinfos):
     strides = attrs.get_int_tuple("strides")
     padding = attrs.get_int_tuple("padding")
     groups = attrs.get_int('groups')
-    layout = attrs["layout"]
+    data_layout_key = "data_layout" if "data_layout" in new_attrs else "layout"
+    layout = attrs[data_layout_key]
     out_dtype = attrs["out_dtype"]
-    out_dtype = tinfos[0].dtype if out_dtype == "same" else out_dtype
+    if out_dtype == "" or out_dtype == "same":
+        out_dtype = tinfos[0].dtype
 
     if layout != 'NCHW' or groups != 1:
         return None
@@ -570,7 +588,7 @@ def _alter_conv2d_layout_arm(attrs, inputs, tinfos):
             [new_data, new_kernel, strides, padding, dilation, 'NCHW', out_dtype], conv2d)
         dispatch_ctx.update(target, new_workload, cfg)
 
-        return sym.conv2d(*copy_inputs, **new_attrs)
+        return F.nn.conv2d(*copy_inputs, **new_attrs)
     else:  # pre-compute weight transformation in winograd
         if "-device=arm_cpu" in target.options:
             tile_size = 4
@@ -580,10 +598,10 @@ def _alter_conv2d_layout_arm(attrs, inputs, tinfos):
             tile_size = _pick_tile_size(tinfos[0], tinfos[1])
             VC = cfg['tile_bna'].val
 
-        weight = sym.contrib.conv2d_winograd_weight_transform(copy_inputs[1], tile_size=tile_size)
-        weight = sym.reshape(weight,
-                             shape=(KH + tile_size - 1, KW + tile_size - 1, CO // VC, VC, CI))
-        weight = sym.transpose(weight, axes=[0, 1, 2, 4, 3])
+        weight = F.nn.contrib_conv2d_winograd_weight_transform(copy_inputs[1], tile_size=tile_size)
+        weight = F.reshape(weight,
+                           newshape=(KH + tile_size - 1, KW + tile_size - 1, CO // VC, VC, CI))
+        weight = F.transpose(weight, axes=[0, 1, 2, 4, 3])
 
         copy_inputs[1] = weight
         new_attrs['tile_size'] = tile_size
@@ -594,8 +612,8 @@ def _alter_conv2d_layout_arm(attrs, inputs, tinfos):
                                      kernel.dtype)
         new_workload = autotvm.task.args_to_workload(
             [new_data, new_weight, strides, padding, dilation,
-             new_attrs['layout'], out_dtype, tile_size],
+             new_attrs[data_layout_key], out_dtype, tile_size],
             conv2d_winograd_without_weight_transform)
         dispatch_ctx.update(target, new_workload, cfg)
 
-        return sym.contrib.conv2d_winograd_without_weight_transform(*copy_inputs, **new_attrs)
+        return F.nn.contrib_conv2d_winograd_without_weight_transform(*copy_inputs, **new_attrs)

topi/python/topi/cuda/conv2d_winograd.py

@@ -330,23 +330,40 @@ def schedule_conv2d_winograd_without_weight_transform_cuda(cfg, outs):
 
 ##### REGISTER ALTER OP LAYOUT #####
 @nn.conv2d_alter_layout.register(["cuda", "gpu"])
-def _alter_conv2d_layout(attrs, inputs, tinfos):
-    """Alter op layout for pre-computing kernel transformation"""
+def _alter_conv2d_layout(attrs, inputs, tinfos, F):
+    """Alter op layout for pre-computing kernel transformation
+
+    Parameters
+    ----------
+    attrs : nnvm.top.AttrDict or tvm.attrs.Attrs
+        Attributes of current convolution
+    inputs : nnvm.symbol or tvm.relay.Expr
+        Grouped input symbols
+    tinfos : list
+        Input shape and dtype
+    F: symbol
+        The context, can be either nnvm.sym or relay.op
+
+    Note
+    ----
+    Unlike other TOPI functions, this function operates on both graph level and operator level,
+    so we have to pass 'F' to make it support our two versions of graph IR, NNVM and Relay.
+    """
     if 'cudnn' in tvm.target.current_target().libs or 'miopen' in tvm.target.current_target().libs:
         return None
 
-    import nnvm.symbol as sym
     copy_inputs = [s for s in inputs]
-
     new_attrs = {k: attrs[k] for k in attrs.keys()}
 
     strides = attrs.get_int_tuple("strides")
     padding = attrs.get_int_tuple("padding")
     dilation = attrs.get_int_tuple("dilation")
     groups = attrs.get_int('groups')
-    layout = attrs["layout"]
+    data_layout_key = "data_layout" if "data_layout" in new_attrs else "layout"
+    layout = attrs[data_layout_key]
     out_dtype = attrs["out_dtype"]
-    out_dtype = tinfos[0].dtype if out_dtype == "same" else out_dtype
+    if out_dtype == "" or out_dtype == "same":
+        out_dtype = tinfos[0].dtype
 
     data, kernel = tinfos[0:2]
     N, CI, H, W = get_const_tuple(data.shape)
@@ -371,7 +388,7 @@ def _alter_conv2d_layout(attrs, inputs, tinfos):
         if cfg.template_key == 'int8':
             assert 'cuda' in target.keys
             new_layout = 'NCHW4c'
-            new_attrs['layout'] = new_layout
+            new_attrs[data_layout_key] = new_layout
             new_attrs['out_layout'] = new_layout
             new_attrs['kernel_layout'] = 'OIHW4o4i'
             ic_block_factor = oc_block_factor = 4
@@ -386,7 +403,7 @@ def _alter_conv2d_layout(attrs, inputs, tinfos):
                 conv2d
             )
             dispatch_ctx.update(target, new_workload, cfg)
-            return sym.conv2d(*copy_inputs, **new_attrs)
+            return F.nn.conv2d(*copy_inputs, **new_attrs)
 
         if attrs.get_int_tuple("dilation") != (1, 1):
             warnings.warn("Does not support weight pre-transform for dilated convolution.")
@@ -395,9 +412,9 @@ def _alter_conv2d_layout(attrs, inputs, tinfos):
         # pre-compute weight transformation in winograd
         tile_size = _infer_tile_size(tinfos[0], tinfos[1])
 
-        weight = sym.contrib.conv2d_winograd_weight_transform(copy_inputs[1],
-                                                              tile_size=tile_size)
-        weight = sym.transpose(weight, axes=[0, 1, 3, 2])
+        weight = F.nn.contrib_conv2d_winograd_weight_transform(copy_inputs[1],
+                                                               tile_size=tile_size)
+        weight = F.transpose(weight, axes=[0, 1, 3, 2])
         copy_inputs[1] = weight
         new_attrs['tile_size'] = tile_size
 
@@ -410,7 +427,7 @@ def _alter_conv2d_layout(attrs, inputs, tinfos):
             conv2d_winograd_without_weight_transform
         )
         dispatch_ctx.update(target, new_workload, cfg)
-        return sym.contrib.conv2d_winograd_without_weight_transform(*copy_inputs, **new_attrs)
+        return F.nn.contrib_conv2d_winograd_without_weight_transform(*copy_inputs, **new_attrs)
     elif groups != CI:
         workload = autotvm.task.args_to_workload(
             [tinfos[0], tinfos[1], strides, padding, dilation, groups, out_dtype],
@@ -424,7 +441,7 @@ def _alter_conv2d_layout(attrs, inputs, tinfos):
         if cfg.template_key == 'int8':
             assert 'cuda' in target.keys
             new_layout = 'NCHW4c'
-            new_attrs['layout'] = new_layout
+            new_attrs[data_layout_key] = new_layout
             new_attrs['out_layout'] = new_layout
             new_attrs['kernel_layout'] = 'OIHW4o4i'
             ic_block_factor = oc_block_factor = 4
@@ -440,7 +457,7 @@ def _alter_conv2d_layout(attrs, inputs, tinfos):
                 group_conv2d_nchw
             )
             dispatch_ctx.update(target, new_workload, cfg)
-            return sym.conv2d(*copy_inputs, **new_attrs)
+            return F.nn.conv2d(*copy_inputs, **new_attrs)
 
     # do nothing for depthwise convolution
     return None

topi/python/topi/intel_graphics/conv2d.py

@@ -3,6 +3,7 @@
 
 from __future__ import absolute_import as _abs
 
+import warnings
 import tvm
 
 from .. import generic
@@ -37,8 +38,13 @@ def tile_and_bind3d(s, tensor, z, y, x, z_factor=2, y_factor=None, x_factor=None
     return xi, thread_z, thread_y, thread_x
 
 @conv2d_alter_layout.register(["intel_graphics"])
-def _alter_conv2d_layout(attrs, inputs, tinfos):
+def _alter_conv2d_layout(attrs, inputs, tinfos, F):
     import nnvm.symbol as sym
+    if F != sym:
+        warnings.warn("Only support alter layout for intel graphics in NNVM now. "
+                      "This pass is ignored in relay.")
+        return None
+
     copy_inputs = [s for s in inputs]
 
     data = tinfos[0]

topi/python/topi/mali/conv2d.py

@@ -465,9 +465,9 @@ def schedule_conv2d_winograd_without_weight_transform_(cfg, outs):
 
 ##### REGISTER ALTER OP LAYOUT #####
 @conv2d_alter_layout.register(["mali"])
-def _alter_conv2d_layout(attrs, inputs, tinfos):
+def _alter_conv2d_layout(attrs, inputs, tinfos, F):
     try:
-        return _alter_conv2d_layout_arm(attrs, inputs, tinfos)
+        return _alter_conv2d_layout_arm(attrs, inputs, tinfos, F)
     except KeyError:  # to filter out fallback opencl templates
         return None
 

topi/python/topi/nn/conv2d.py

@@ -57,17 +57,24 @@ def conv2d(input, filter, strides, padding, dilation, layout='NCHW', out_dtype=N
 
 
 @tvm.target.generic_func
-def conv2d_alter_layout(attrs, inputs, tinfos):
+def conv2d_alter_layout(attrs, inputs, tinfos, F):
     """Change Conv2D layout.
 
     Parameters
     ----------
-    attrs : nnvm.top.AttrDict
+    attrs : nnvm.top.AttrDict or tvm.attrs.Attrs
         Attributes of current convolution
-    inputs : nnvm.symbol
+    inputs : nnvm.symbol or tvm.relay.Expr
         Grouped input symbols
     tinfos : list
         Input shape and dtype
+    F: symbol
+        The context, can be either nnvm.sym or relay.op
+
+    Note
+    ----
+    Unlike other TOPI functions, this function operates on both graph level and operator level,
+    so we have to pass 'F' to make it support our two versions of graph IR, NNVM and Relay.
     """
     # not to change by default
     return None

topi/python/topi/x86/conv2d.py

 # pylint: disable=invalid-name,unused-variable,unused-argument,no-member
 """Conv2D schedule on x86"""
+import warnings
+
 import tvm
 from tvm import autotvm
 from tvm.autotvm.task.topi_integration import deserialize_args
@@ -281,8 +283,13 @@ def _topi_nn_conv2d_NCHWc(*args, **kwargs):
 
 
 @conv2d_alter_layout.register("cpu")
-def _alter_conv2d_layout(attrs, inputs, tinfo):
+def _alter_conv2d_layout(attrs, inputs, tinfo, F):
     import nnvm.symbol as sym
+    if F != sym:
+        warnings.warn("Only support alter layout for x86 in NNVM now. "
+                      "This pass is ignored in relay.")
+        return None
+
     copy_inputs = [s for s in inputs]
     new_attrs = {k : attrs[k] for k in attrs.keys()}
     data, kernel = tinfo[0], tinfo[1]