Add support for MXNet pad operator, for parameter 'None' in slice operator, 1D Convolution and 1D Deconvolution

include/tvm/expr_operator.h

@@ -521,6 +521,7 @@ TVM_DECLARE_INTRIN_UNARY(log);
 TVM_DECLARE_INTRIN_UNARY(popcount);
 TVM_DECLARE_INTRIN_UNARY(cos);
 TVM_DECLARE_INTRIN_UNARY(sin);
+TVM_DECLARE_INTRIN_UNARY(atan);
 
 // Implementation details after this
 inline bool is_const(const Expr& x) {

include/tvm/relay/attrs/nn.h

@@ -405,13 +405,18 @@ struct UpSamplingAttrs : public tvm::AttrsNode<UpSamplingAttrs> {
 struct PadAttrs : public tvm::AttrsNode<PadAttrs> {
   double pad_value;
   Array<Array<IndexExpr> > pad_width;
+  std::string pad_mode;
 
   TVM_DECLARE_ATTRS(PadAttrs, "relay.attrs.PadAttrs") {
     TVM_ATTR_FIELD(pad_value).set_default(0.0)
-      .describe("Specifies the strides of the convolution.");
+      .describe("The value used for padding when mode is 'constant'.");
     TVM_ATTR_FIELD(pad_width)
       .describe("Number of values padded to the edges of each axis, "
                 "in the format of ((before_1, after_1), ..., (before_N, after_N))");
+    TVM_ATTR_FIELD(pad_mode).set_default("constant")
+      .describe("Padding type to use. \"constant\" pads with constant_value, "
+                "\"edge\" pads using the edge values of the input array, "
+                "\"reflect\" pads by reflecting values with respect to the edges.");
   }
 };
 

python/tvm/intrin.py

@@ -304,6 +304,21 @@ def sin(x):
     """
     return call_pure_intrin(x.dtype, "sin", x)
 
+def atan(x):
+    """Take atan of input x.
+
+    Parameters
+    ----------
+    x : Expr
+        Input argument.
+
+    Returns
+    -------
+    y : Expr
+        The result.
+    """
+    return call_pure_intrin(x.dtype, "atan", x)
+
 def sqrt(x):
     """Take square root of input x.
 

python/tvm/relay/frontend/common.py

@@ -138,16 +138,8 @@ def get_int_tuple(self, key, default=RequiredAttr()):
         """
         if key in self.attrs:
             tshape = self.attrs[key]
-            ret = []
-            for x in tshape.strip('()[]').split(','):
-                x = x.strip()
-                if not x:
-                    continue
-                if x == "None":
-                    ret.append(None)
-                else:
-                    ret.append(int(x))
-            return tuple(ret)
+            return tuple(int(x) if x.strip("- ").isdigit() else None
+                         for x in tshape.strip('()[]').split(',') if x)
         if isinstance(default, RequiredAttr):
             raise AttributeError("Required attribute {} not found.".format(key))
         return default

python/tvm/relay/frontend/mxnet.py

@@ -112,11 +112,55 @@ def _mx_zeros(inputs, attrs):
     return _op.zeros(shape=shape, dtype=dtype)
 
 
+def _mx_conv(inputs, attrs):
+    kernel_size = attrs.get_int_tuple("kernel")
+    if len(kernel_size) == 2:
+        return _mx_conv2d(inputs, attrs)
+    elif len(kernel_size) == 1:
+        return _mx_conv1d(inputs, attrs)
+    else:
+        raise tvm.error.OpAttributeInvalid(
+            '1D or 2D kernels only are supported for operator Convolution')
+
+def _mx_conv1d(inputs, attrs):
+    kernel_size = attrs.get_int_tuple("kernel")
+    if len(kernel_size) != 1:
+        raise tvm.error.OpAttributeInvalid(
+            'Non 1D or 2D kernels are not supported for operator Convolution')
+    data_layout = attrs.get_str("layout", "NCW")
+    # MXNet Conv1D only supports ‘NCW’ layout for now.
+    if data_layout != "NCW":
+        raise tvm.error.OpAttributeInvalid(
+            'Only "NCW" data layout is supported for 1D Convolution')
+    data_layout = "NCHW"
+    channel_axis = 1
+    kernel_layout = "OIHW"
+
+    new_attrs = {}
+    new_attrs["channels"] = attrs.get_int("num_filter")
+    new_attrs["kernel_size"] = (1,) + kernel_size
+    new_attrs["strides"] = (1,) + attrs.get_int_tuple("stride", (1,))
+    new_attrs["padding"] = (0,) + attrs.get_int_tuple("pad", (0,))
+    new_attrs["dilation"] = (1,) +  attrs.get_int_tuple("dilate", (1,))
+    new_attrs["groups"] = attrs.get_int("num_group", 1)
+    new_attrs["data_layout"] = data_layout
+    new_attrs["kernel_layout"] = kernel_layout
+    use_bias = not attrs.get_bool("no_bias", False)
+    data = _op.expand_dims(inputs[0], axis=2)
+    kernel = _op.expand_dims(inputs[1], axis=2)
+    res = _op.nn.conv2d(data, kernel, **new_attrs)
+    if use_bias:
+        assert len(inputs) == 3
+        res = _op.nn.bias_add(res, inputs[2], axis=channel_axis)
+    res = _op.squeeze(res, axis=[2])
+    return res
+
+
 def _mx_conv2d(inputs, attrs):
     kernel_size = attrs.get_int_tuple("kernel")
     if len(kernel_size) != 2:
         raise tvm.error.OpAttributeInvalid(
-            'Non-2D kernels are not supported for operator Conv2D.')
+            'Non 1D or 2D kernels are not supported for operator Convolution')
     data_layout = attrs.get_str("layout", "NCHW")
     channel_axis = _get_channel_axis(data_layout, "conv2d")
 
@@ -142,6 +186,51 @@ def _mx_conv2d(inputs, attrs):
     return res
 
 
+def _mx_conv_transpose(inputs, attrs):
+    kernel_size = attrs.get_int_tuple("kernel")
+    if len(kernel_size) == 2:
+        return _mx_conv2d_transpose(inputs, attrs)
+    elif len(kernel_size) == 1:
+        return _mx_conv1d_transpose(inputs, attrs)
+    else:
+        raise tvm.error.OpAttributeInvalid(
+            '1D or 2D kernels only are supported for operator Convolution')
+
+
+def _mx_conv1d_transpose(inputs, attrs):
+    if "target_shape" in attrs.attrs:
+        raise tvm.error.OpAttributeUnImplemented(
+            'Attribute "target_shape" is not supported for operator Conv2D-transpose.')
+    data_layout = attrs.get_str("layout", "NCW")
+    if data_layout != "NCW":
+        raise tvm.error.OpAttributeInvalid(
+            'Only "NCW" data layout is supported for 1D Convolution')
+    data_layout = "NCHW"
+    channel_axis = 1
+    kernel_layout = "OIHW"
+
+    new_attrs = {}
+    new_attrs["channels"] = attrs.get_int("num_filter")
+    new_attrs["kernel_size"] = (1,) + attrs.get_int_tuple("kernel")
+    new_attrs["strides"] = (1,) + attrs.get_int_tuple("stride", (1,))
+    new_attrs["output_padding"] = (0,) + attrs.get_int_tuple("adj", (0,))
+    new_attrs["padding"] = (0,) + attrs.get_int_tuple("pad", (0,))
+    new_attrs["dilation"] = (1,) +  attrs.get_int_tuple("dilate", (1,))
+    new_attrs["groups"] = attrs.get_int("num_group", 1)
+    new_attrs["data_layout"] = data_layout
+    new_attrs["kernel_layout"] = kernel_layout
+    use_bias = not attrs.get_bool("no_bias", True)
+    data = _op.expand_dims(inputs[0], axis=2)
+    kernel = _op.expand_dims(inputs[1], axis=2)
+    res = _op.nn.conv2d_transpose(data, kernel, **new_attrs)
+
+    if use_bias:
+        assert len(inputs) == 3
+        res = _op.nn.bias_add(res, inputs[2], axis=channel_axis)
+    res = _op.squeeze(res, axis=[2])
+    return res
+
+
 def _mx_conv2d_transpose(inputs, attrs):
     if "target_shape" in attrs.attrs:
         raise tvm.error.OpAttributeUnImplemented(
@@ -257,13 +346,7 @@ def _mx_slice(inputs, attrs):
     if end is None:
         raise tvm.error.OpAttributeRequired(
             'Attribute "end" not found in operator Slice.')
-    if None in begin:
-        data_shape = _infer_type(inputs[0]).checked_type.shape
-        for i, beg in enumerate(begin):
-            if beg is None:
-                assert end[i] is None
-                begin[i] = 0
-                end[i] = data_shape[i]
+    begin = tuple(x if x is not None else 0 for x in begin)
     new_attrs = {'begin': begin, 'end': end}
     if stride is not None:
         new_attrs['strides'] = stride
@@ -373,6 +456,27 @@ def _mx_expand_dims(inputs, attrs):
     axis = attrs.get_int("axis")
     return _op.expand_dims(inputs[0], axis=axis)
 
+def _mx_pad(inputs, attrs):
+    pad_mode = attrs.get_str('mode', None)
+    if pad_mode is None:
+        raise tvm.error.OpAttributeRequired(
+            'Attribute "mode" not found in operator pad.')
+    if pad_mode not in ['constant', 'edge', 'reflect']:
+        raise tvm.error.OpAttributeInvalid(
+            'Value ' + mode + ' in attribute "mode" is not valid')
+    pad_width = attrs.get_int_tuple('pad_width', None)
+    if pad_width is None:
+        raise tvm.error.OpAttributeRequired(
+            'Attribute "pad_width" not found in operator pad.')
+    if None in pad_width:
+        raise tvm.error.OpAttributeInvalid(
+            'Value None in attribute "pad_width" of operator Slice is not valid.')
+    constant_value = attrs.get_float('constant_value', 0.0)
+    padding = tuple(tuple((b, a)) for b, a in zip(pad_width[::2], pad_width[1::2]))
+    return _op.nn.pad(data=inputs[0],
+                      pad_width=padding,
+                      pad_value=constant_value,
+                      pad_mode=pad_mode)
 
 def _mx_leaky_relu(inputs, attrs):
     act_type = attrs.get_str("act_type")
@@ -931,6 +1035,8 @@ def _mx_one_hot(inputs, attrs):
     "ones_like",
     "where",
     "gather_nd",
+    "cos",
+    "sin"
 ]
 
 _convert_map = {
@@ -943,6 +1049,7 @@ def _mx_one_hot(inputs, attrs):
     "broadcast_mod"          : _rename(_op.mod),
     "broadcast_maximum"      : _rename(_op.maximum),
     "broadcast_minimum"      : _rename(_op.minimum),
+    "arctan"                 : _rename(_op.atan),
     "broadcast_equal"        : _mx_compare(_op.equal, _rename),
     "broadcast_not_equal"    : _mx_compare(_op.not_equal, _rename),
     "broadcast_greater"      : _mx_compare(_op.greater, _rename),
@@ -1018,9 +1125,9 @@ def _mx_one_hot(inputs, attrs):
     "_zeros"        : _mx_zeros,
     "FullyConnected": _mx_fully_connected,
     "Activation"    : _mx_activations,
-    "Convolution"   : _mx_conv2d,
+    "Convolution"   : _mx_conv,
     "Convolution_v1": _mx_conv2d,
-    "Deconvolution" : _mx_conv2d_transpose,
+    "Deconvolution" : _mx_conv_transpose,
     "Pooling"       : _mx_pooling,
     "Pooling_v1"    : _mx_pooling,
     "Dropout"       : _mx_dropout,
@@ -1044,6 +1151,8 @@ def _mx_one_hot(inputs, attrs):
     "_full"         : _mx_full,
     "repeat"        : _mx_repeat,
     "tile"          : _mx_tile,
+    "pad"           : _mx_pad,
+    "Pad"           : _mx_pad,
     "take"          : _mx_take,
     "reverse"       : _mx_reverse,
     "squeeze"       : _mx_squeeze,

python/tvm/relay/frontend/tensorflow.py

@@ -1,3 +1,4 @@
+
 # Licensed to the Apache Software Foundation (ASF) under one
 # or more contributor license agreements.  See the NOTICE file
 # distributed with this work for additional information

python/tvm/relay/op/_tensor.py

@@ -29,6 +29,7 @@
 register_schedule("log1p", schedule_broadcast)
 register_schedule("cos", schedule_broadcast)
 register_schedule("sin", schedule_broadcast)
+register_schedule("atan", schedule_broadcast)
 register_schedule("exp", schedule_broadcast)
 register_schedule("erf", schedule_broadcast)
 register_schedule("sqrt", schedule_broadcast)

python/tvm/relay/op/_tensor_grad.py

@@ -60,6 +60,12 @@ def sin_grad(orig, grad):
     x = orig.args[0]
     return [grad * cos(x)]
 
+@register_gradient("atan")
+def atan_grad(orig, grad):
+    """Returns [grad * 1 / (1 + x ^ 2)]"""
+    x = orig.args[0]
+    a = const(2.0)
+    return [grad * ones_like(x) / (ones_like(x) + power(x, a))]
 
 @register_gradient("exp")
 def exp_grad(orig, grad):

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

@@ -673,7 +673,8 @@ def prelu(data, alpha, axis=1):
 
 def pad(data,
         pad_width,
-        pad_value=0.0):
+        pad_value=0.0,
+        pad_mode='constant'):
     r"""Padding
 
     This operator takes in a tensor and pads each axis by the specified
@@ -688,13 +689,16 @@ def pad(data,
         of ((before_1, after_1), ..., (before_N, after_N))
     pad_value: float, optional, default=0.0
         The value used for padding
-
+    pad_mode: 'constant', 'edge', 'reflect'
+        'constant' pads with constant_value pad_value
+        'edge' pads using the edge values of the input array
+        'reflect' pads by reflecting values with respect to the edge
     Returns
     -------
     result : tvm.relay.Expr
         The computed result.
     """
-    return _make.pad(data, pad_width, pad_value)
+    return _make.pad(data, pad_width, pad_value, pad_mode)
 
 
 def mirror_pad(data,

python/tvm/relay/op/tensor.py

@@ -76,6 +76,21 @@ def sin(data):
     """
     return _make.sin(data)
 
+def atan(data):
+    """Compute elementwise atan of data.
+
+    Parameters
+    ----------
+    data : relay.Expr
+        The input data
+
+    Returns
+    -------
+    result : relay.Expr
+        The computed result.
+    """
+    return _make.atan(data)
+
 def exp(data):
     """Compute elementwise exp of data.
 

src/codegen/intrin_rule.cc

@@ -46,6 +46,9 @@ TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.cos")
 TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.sin")
 .set_body(DispatchExtern<FloatSuffix>);
 
+TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.atan")
+.set_body(DispatchExtern<FloatSuffix>);
+
 TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.sqrt")
 .set_body(DispatchExtern<FloatSuffix>);
 

src/codegen/intrin_rule_cuda.cc

@@ -104,6 +104,9 @@ TVM_REGISTER_GLOBAL("tvm.intrin.rule.cuda.cos")
 TVM_REGISTER_GLOBAL("tvm.intrin.rule.cuda.sin")
 .set_body(DispatchExtern<CUDAFastMath>);
 
+TVM_REGISTER_GLOBAL("tvm.intrin.rule.cuda.atan")
+.set_body(DispatchExtern<CUDAMath>);
+
 TVM_REGISTER_GLOBAL("tvm.intrin.rule.cuda.tanh")
 .set_body(DispatchExtern<CUDAMath>);
 

src/codegen/llvm/intrin_rule_nvptx.cc

@@ -88,6 +88,9 @@ TVM_REGISTER_GLOBAL("tvm.intrin.rule.nvptx.cos")
 TVM_REGISTER_GLOBAL("tvm.intrin.rule.nvptx.sin")
 .set_body(DispatchExternLibDevice);
 
+TVM_REGISTER_GLOBAL("tvm.intrin.rule.nvptx.atan")
+.set_body(DispatchExternLibDevice);
+
 }  // namespace llvm
 }  // namespace codegen
 }  // namespace tvm

src/relay/op/nn/pad.cc

@@ -102,14 +102,21 @@ Array<Tensor> PadCompute(const Attrs& attrs,
   }
   const auto* out_ttype = out_type.as<TensorTypeNode>();
   return Array<Tensor>{ topi::pad(inputs[0], pad_before, pad_after,
-                                  tvm::make_const(out_ttype->dtype, param->pad_value)) };
+                                  tvm::make_const(out_ttype->dtype, param->pad_value),
+                                  "T_pad",
+                                  topi::kElementWise,
+                                  param->pad_mode) };
 }
 
 // Handler to create a call to the padding op used by front-end FFI
-Expr MakePad(Expr data, Array<Array<IndexExpr> > pad_width, double pad_value) {
+Expr MakePad(Expr data,
+             Array<Array<IndexExpr> > pad_width,
+             double pad_value,
+             std::string pad_mode) {
   auto attrs = make_node<PadAttrs>();
   attrs->pad_value = pad_value;
   attrs->pad_width = std::move(pad_width);
+  attrs->pad_mode = std::move(pad_mode);
   static const Op& op = Op::Get("nn.pad");
   return CallNode::make(op, {data}, Attrs(attrs), {});
 }