[RELAY] Port winograd ops to relay

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

@@ -1,11 +1,12 @@
 # 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 @@ def get_bool(self, key):
         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 @@ def keys(self):
         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,14 +119,17 @@ 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
     ----------
     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)