From 887a81a25b5468b7141b691464d1592c465b327f Mon Sep 17 00:00:00 2001 From: Haichen Shen Date: Wed, 26 Jun 2019 13:55:19 -0700 Subject: [PATCH 01/12] add more ops --- include/tvm/expr_operator.h | 1 + python/tvm/intrin.py | 16 ++++ python/tvm/relay/backend/interpreter.py | 6 +- python/tvm/relay/frontend/common.py | 11 ++- python/tvm/relay/frontend/mxnet.py | 40 +++++++-- python/tvm/relay/op/_tensor.py | 1 + python/tvm/relay/op/tensor.py | 16 ++++ src/codegen/intrin_rule.cc | 3 + src/codegen/intrin_rule_cuda.cc | 3 + src/relay/op/tensor/unary.cc | 12 +++ tests/python/relay/test_op_level1.py | 2 + topi/include/topi/elemwise.h | 1 + topi/python/topi/math.py | 17 ++++ topi/python/topi/x86/dense.py | 113 ++++++++++++------------ topi/src/topi.cc | 6 +- topi/tests/python/test_topi_basic.py | 1 + topi/tests/python/test_topi_math.py | 2 + 17 files changed, 181 insertions(+), 70 deletions(-) diff --git a/include/tvm/expr_operator.h b/include/tvm/expr_operator.h index b887897546b5..bef9705e4749 100644 --- a/include/tvm/expr_operator.h +++ b/include/tvm/expr_operator.h @@ -512,6 +512,7 @@ TVM_DLL Expr trunc(Expr x); } \ TVM_DECLARE_INTRIN_UNARY(exp); +TVM_DECLARE_INTRIN_UNARY(erf); TVM_DECLARE_INTRIN_UNARY(tanh); TVM_DECLARE_INTRIN_UNARY(sigmoid); TVM_DECLARE_INTRIN_UNARY(sqrt); diff --git a/python/tvm/intrin.py b/python/tvm/intrin.py index e8f164c1c579..d45a2d394e21 100644 --- a/python/tvm/intrin.py +++ b/python/tvm/intrin.py @@ -211,6 +211,22 @@ def exp(x): return call_pure_intrin(x.dtype, "exp", x) +def erf(x): + """Take gauss error function of the input x. + + Parameters + ---------- + x : Expr + Input argument. + + Returns + ------- + y : Expr + The result. + """ + return call_pure_intrin(x.dtype, "erf", x) + + def tanh(x): """Take hyperbolic tanh of input x. diff --git a/python/tvm/relay/backend/interpreter.py b/python/tvm/relay/backend/interpreter.py index 491720de5cda..79738e9c7189 100644 --- a/python/tvm/relay/backend/interpreter.py +++ b/python/tvm/relay/backend/interpreter.py @@ -182,9 +182,9 @@ def _convert_args(self, expr, args, kwargs): if i < num_of_args: if kwargs.get(name): raise Exception( - "duplicate argument supplied in \ - both positional args (at position: {0}), \ - and keyword argument (with name: {1})".format(i, name)) + "duplicate argument supplied in" + + "both positional args (at position: {0})," + + "and keyword argument (with name: {1})".format(i, name)) else: cargs.append(kwargs[name]) diff --git a/python/tvm/relay/frontend/common.py b/python/tvm/relay/frontend/common.py index 4da3a528fc51..81c4ce50154c 100644 --- a/python/tvm/relay/frontend/common.py +++ b/python/tvm/relay/frontend/common.py @@ -124,7 +124,16 @@ def get_int_tuple(self, key, default=RequiredAttr()): """ if key in self.attrs: tshape = self.attrs[key] - return tuple(int(x.strip()) for x in tshape.strip('()[]').split(',') if x) + 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) if isinstance(default, RequiredAttr): raise AttributeError("Required attribute {} not found.".format(key)) return default diff --git a/python/tvm/relay/frontend/mxnet.py b/python/tvm/relay/frontend/mxnet.py index 36c4fb895874..576f979a5b8e 100644 --- a/python/tvm/relay/frontend/mxnet.py +++ b/python/tvm/relay/frontend/mxnet.py @@ -55,10 +55,18 @@ def _mx_fully_connected(inputs, attrs): use_flatten = attrs.get_bool("flatten", True) if has_flatten and use_flatten: inputs[0] = _op.nn.batch_flatten(inputs[0]) + data_shape = _infer_type(inputs[0]).checked_type.shape + weight_shape = _infer_type(inputs[1]).checked_type.shape + if len(data_shape) > 2: + inputs[0] = _op.reverse_reshape(inputs[0], [-1, 0]) res = _op.nn.dense(inputs[0], inputs[1], units=units) if use_bias: assert len(inputs) == 3 res = _op.nn.bias_add(res, inputs[2], axis=-1) + if len(data_shape) > 2: + new_shape = data_shape[:-1] + new_shape.append(weight_shape[0]) + res = _op.reshape(res, new_shape) return res @@ -241,8 +249,8 @@ def _mx_layer_norm(inputs, attrs): def _mx_slice(inputs, attrs): new_attrs = {} - begin = attrs.get_int_tuple('begin', None) - end = attrs.get_int_tuple('end', None) + begin = list(attrs.get_int_tuple('begin', None)) + end = list(attrs.get_int_tuple('end', None)) stride = attrs.get_int_tuple('step', None) if begin is None: raise tvm.error.OpAttributeRequired( @@ -251,11 +259,12 @@ def _mx_slice(inputs, attrs): raise tvm.error.OpAttributeRequired( 'Attribute "end" not found in operator Slice.') if None in begin: - raise tvm.error.OpAttributeInvalid( - 'Value None in attribute "begin" of operator Slice is not valid.') - if None in end: - raise tvm.error.OpAttributeInvalid( - 'Value None in attribute "end" of operator Slice is not valid.') + 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] new_attrs = {'begin': begin, 'end': end} if stride is not None: new_attrs['strides'] = stride @@ -497,7 +506,8 @@ def _mx_arange(inputs, attrs): 'Attribute "repeat" is not supported in operator arange.') new_attrs = {} new_attrs["start"] = _expr.const(attrs.get_float("start", 0.0)) - new_attrs["stop"] = _expr.const(attrs.get_float("stop")) + stop = attrs.get_str("stop", "None") + new_attrs["stop"] = None if stop == "None" else _expr.const(eval(stop)) new_attrs["step"] = _expr.const(attrs.get_float("step", 1.0)) new_attrs["dtype"] = attrs.get_str("dtype", "float32") return _op.arange(**new_attrs) @@ -752,6 +762,19 @@ def _mx_rnn_param_concat(inputs, _): return [inputs] +def _mx_contrib_div_sqrt_dim(inputs, _): + assert len(inputs) == 1 + ndim = len(_infer_type(inputs[0]).checked_type.shape) + dim = _op.take(_op.shape_of(inputs[0]), _expr.const(ndim-1, dtype="int32")) + sqrt_dim = _op.sqrt(dim.astype('float32')) + out = inputs[0] / sqrt_dim + return out + + +def _mx_sequence_mask(inputs, attrs): + return inputs[0] + + def _mx_rnn_layer(inputs, attrs): def _rnn_cell(data, states, i2h_weight, h2h_weight, i2h_bias, h2h_bias, activation): i2h = _op.nn.bias_add(_op.nn.dense(data, i2h_weight), i2h_bias, axis=-1) @@ -910,6 +933,7 @@ def _mx_one_hot(inputs, attrs): _identity_list = [ "log", "exp", + "erf", "sqrt", "floor", "ceil", diff --git a/python/tvm/relay/op/_tensor.py b/python/tvm/relay/op/_tensor.py index 2e342339ba83..c1b95d8cd2d9 100644 --- a/python/tvm/relay/op/_tensor.py +++ b/python/tvm/relay/op/_tensor.py @@ -30,6 +30,7 @@ register_schedule("cos", schedule_broadcast) register_schedule("sin", schedule_broadcast) register_schedule("exp", schedule_broadcast) +register_schedule("erf", schedule_broadcast) register_schedule("sqrt", schedule_broadcast) register_schedule("rsqrt", schedule_broadcast) register_schedule("sigmoid", schedule_broadcast) diff --git a/python/tvm/relay/op/tensor.py b/python/tvm/relay/op/tensor.py index 14a6ba81584c..a2db71e006ad 100644 --- a/python/tvm/relay/op/tensor.py +++ b/python/tvm/relay/op/tensor.py @@ -92,6 +92,22 @@ def exp(data): return _make.exp(data) +def erf(data): + """Compute elementwise error function of data. + + Parameters + ---------- + data : relay.Expr + The input data + + Returns + ------- + result : relay.Expr + The computed result. + """ + return _make.erf(data) + + def sqrt(data): """Compute elementwise sqrt of data. diff --git a/src/codegen/intrin_rule.cc b/src/codegen/intrin_rule.cc index 680a4fc6dd7b..6afe3f350038 100644 --- a/src/codegen/intrin_rule.cc +++ b/src/codegen/intrin_rule.cc @@ -31,6 +31,9 @@ namespace intrin { TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.exp") .set_body(DispatchExtern); +TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.erf") +.set_body(DispatchExtern); + TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.log") .set_body(DispatchExtern); diff --git a/src/codegen/intrin_rule_cuda.cc b/src/codegen/intrin_rule_cuda.cc index 65498dcc5ee2..c507e5b28e0c 100644 --- a/src/codegen/intrin_rule_cuda.cc +++ b/src/codegen/intrin_rule_cuda.cc @@ -92,6 +92,9 @@ TVM_REGISTER_GLOBAL("tvm.intrin.rule.cuda.round") TVM_REGISTER_GLOBAL("tvm.intrin.rule.cuda.exp") .set_body(DispatchExtern); +TVM_REGISTER_GLOBAL("tvm.intrin.rule.cuda.erf") +.set_body(DispatchExtern); + TVM_REGISTER_GLOBAL("tvm.intrin.rule.cuda.log") .set_body(DispatchExtern); diff --git a/src/relay/op/tensor/unary.cc b/src/relay/op/tensor/unary.cc index 826fe693e6e6..cb5a116aee69 100644 --- a/src/relay/op/tensor/unary.cc +++ b/src/relay/op/tensor/unary.cc @@ -85,6 +85,18 @@ RELAY_REGISTER_UNARY_OP("exp") .set_support_level(1) .set_attr("FTVMCompute", RELAY_UNARY_COMPUTE(topi::exp)); + +RELAY_REGISTER_UNARY_OP("erf") +.describe(R"code(Returns the error function value for input array, computed element-wise. + +.. math:: + \erf(x) + +)code" TVM_ADD_FILELINE) +.set_support_level(1) +.set_attr("FTVMCompute", RELAY_UNARY_COMPUTE(topi::erf)); + + RELAY_REGISTER_UNARY_OP("sqrt") .describe(R"code(Returns the sqrt input array, computed element-wise. diff --git a/tests/python/relay/test_op_level1.py b/tests/python/relay/test_op_level1.py index c25393cf4026..dcd4ab7010f3 100644 --- a/tests/python/relay/test_op_level1.py +++ b/tests/python/relay/test_op_level1.py @@ -16,6 +16,7 @@ # under the License. import numpy as np import tvm +import scipy from tvm import relay from tvm.relay import transform from tvm.relay.testing import ctx_list @@ -67,6 +68,7 @@ def check_single_op(opfunc, ref): for opfunc, ref in [(tvm.relay.log, np.log), (tvm.relay.exp, np.exp), + (tvm.relay.erf, scipy.special.erf), (tvm.relay.sqrt, np.sqrt), (tvm.relay.rsqrt, rsqrt), (tvm.relay.sigmoid, sigmoid), diff --git a/topi/include/topi/elemwise.h b/topi/include/topi/elemwise.h index 449766cb8888..b1e27c4444bf 100644 --- a/topi/include/topi/elemwise.h +++ b/topi/include/topi/elemwise.h @@ -46,6 +46,7 @@ using namespace tvm; } TOPI_DECLARE_UNARY_OP(exp); +TOPI_DECLARE_UNARY_OP(erf); TOPI_DECLARE_UNARY_OP(sigmoid); TOPI_DECLARE_UNARY_OP(sqrt); TOPI_DECLARE_UNARY_OP(log); diff --git a/topi/python/topi/math.py b/topi/python/topi/math.py index 5e3f687940d3..6f44b854d952 100644 --- a/topi/python/topi/math.py +++ b/topi/python/topi/math.py @@ -74,6 +74,23 @@ def exp(x): return tvm.compute(x.shape, lambda *i: tvm.exp(x(*i))) +@tvm.tag_scope(tag=tag.ELEMWISE) +def erf(x): + """Take gauss error function of input x. + + Parameters + ---------- + x : tvm.Tensor + Input argument. + + Returns + ------- + y : tvm.Tensor + The result. + """ + return tvm.compute(x.shape, lambda *i: tvm.erf(x(*i))) + + @tvm.tag_scope(tag=tag.ELEMWISE) def tanh(x): """Take hyperbolic tanh of input x. diff --git a/topi/python/topi/x86/dense.py b/topi/python/topi/x86/dense.py index e22ad444b719..986330053e0c 100644 --- a/topi/python/topi/x86/dense.py +++ b/topi/python/topi/x86/dense.py @@ -28,12 +28,19 @@ @autotvm.register_topi_compute(nn.dense, "cpu", "direct") def _declaration_dense(cfg, data, weight, bias=None, out_dtype=None): - batch, _ = get_const_tuple(data.shape) + target = tvm.target.current_target() + if "cblas" in target.libs: + C = cblas.matmul(data, weight, False, True) + if bias is not None: + C = tvm.compute((M, N), lambda i, j: C[i, j] + bias[j].astype(out_dtype), + tag=tag.BROADCAST) + return C + M, _ = get_const_tuple(data.shape) # For small batch sizes, don't pack weight into cache-friendly layout # because of overhead in packing and limited reuse from batch dimension # TODO(icemelon9): use a more systematic way to determine which schedule to use - if batch <= 16: + if M <= 16: return _declaration_dense_nopack(cfg, data, weight, bias, out_dtype) return _declaration_dense_pack(cfg, data, weight, bias, out_dtype) @@ -41,35 +48,31 @@ def _declaration_dense(cfg, data, weight, bias=None, out_dtype=None): # Declare dense compute with packing weight into cache-friendly layout @autotvm.register_topi_compute(nn.dense, "cpu", "direct_pack") def _declaration_dense_pack(cfg, data, weight, bias=None, out_dtype=None): - target = tvm.target.current_target() - if "cblas" in target.libs: - C = cblas.matmul(data, weight, False, True) - else: - if out_dtype is None: - out_dtype = data.dtype - batch, in_dim = get_const_tuple(data.shape) - out_dim, _ = get_const_tuple(weight.shape) - # create tuning space - cfg.define_split("tile_y", batch, num_outputs=3) - cfg.define_split("tile_x", out_dim, num_outputs=3) - cfg.define_split("tile_k", in_dim, num_outputs=2) - if cfg.is_fallback: - _default_dense_pack_config(cfg, batch, out_dim, in_dim) - - packw_bn = cfg["tile_x"].size[-1] - packw_shape = (out_dim // packw_bn, in_dim, packw_bn) - packw = tvm.compute(packw_shape, - lambda z, y, x: weight[z * packw_bn + x, y], name="packed_weight") - - k = tvm.reduce_axis((0, in_dim), name="k") - C = tvm.compute((batch, out_dim), - lambda y, x: tvm.sum( - data[y, k].astype(out_dtype) * - packw[x // packw_bn, k, x % packw_bn].astype(out_dtype), - axis=k), - tag="dense_pack") + if out_dtype is None: + out_dtype = data.dtype + M, K = get_const_tuple(data.shape) # batch, in_dim + N, _ = get_const_tuple(weight.shape) # out_dim + # create tuning space + cfg.define_split("tile_y", M, num_outputs=3) + cfg.define_split("tile_x", N, num_outputs=3) + cfg.define_split("tile_k", K, num_outputs=2) + if cfg.is_fallback: + _default_dense_pack_config(cfg, M, N, K) + + packw_bn = cfg["tile_x"].size[-1] + packw_shape = (N // packw_bn, K, packw_bn) + packw = tvm.compute(packw_shape, + lambda z, y, x: weight[z * packw_bn + x, y], name="packed_weight") + + k = tvm.reduce_axis((0, K), name="k") + C = tvm.compute((M, N), + lambda y, x: tvm.sum( + data[y, k].astype(out_dtype) * + packw[x // packw_bn, k, x % packw_bn].astype(out_dtype), + axis=k), + tag="dense_pack") if bias is not None: - C = tvm.compute((batch, out_dim), lambda i, j: C[i, j] + bias[j].astype(out_dtype), + C = tvm.compute((M, N), lambda i, j: C[i, j] + bias[j].astype(out_dtype), tag=tag.BROADCAST) return C @@ -77,34 +80,30 @@ def _declaration_dense_pack(cfg, data, weight, bias=None, out_dtype=None): # Declare dense compute without packing weight @autotvm.register_topi_compute(nn.dense, "cpu", "direct_nopack") def _declaration_dense_nopack(cfg, data, weight, bias=None, out_dtype=None): - target = tvm.target.current_target() - if "cblas" in target.libs: - C = cblas.matmul(data, weight, False, True) - else: - if out_dtype is None: - out_dtype = data.dtype - batch, in_dim = get_const_tuple(data.shape) - out_dim, _ = get_const_tuple(weight.shape) - # create tuning space - cfg.define_split("tile_x", out_dim, num_outputs=2) - cfg.define_split("tile_y", batch, num_outputs=2) - cfg.define_split("tile_k", in_dim, num_outputs=2) - if cfg.is_fallback: - _default_dense_nopack_config(cfg, batch, out_dim, in_dim) - - vec = cfg["tile_k"].size[-1] - k = tvm.reduce_axis((0, in_dim // vec), "k") - CC = tvm.compute((batch, out_dim, vec), - lambda z, y, x: tvm.sum( - data[z, k * vec + x].astype(out_dtype) * - weight[y, k * vec + x].astype(out_dtype), axis=k)) - - kk = tvm.reduce_axis((0, vec), "kk") - C = tvm.compute((batch, out_dim), - lambda y, x: tvm.sum(CC[y, x, kk], axis=kk), - tag="dense_nopack") + if out_dtype is None: + out_dtype = data.dtype + M, K = get_const_tuple(data.shape) + N, _ = get_const_tuple(weight.shape) + # create tuning space + cfg.define_split("tile_y", M, num_outputs=2) + cfg.define_split("tile_x", N, num_outputs=2) + cfg.define_split("tile_k", K, num_outputs=2) + if cfg.is_fallback: + _default_dense_nopack_config(cfg, M, N, K) + + vec = cfg["tile_k"].size[-1] + k = tvm.reduce_axis((0, K // vec), "k") + CC = tvm.compute((M, N, vec), + lambda z, y, x: tvm.sum( + data[z, k * vec + x].astype(out_dtype) * + weight[y, k * vec + x].astype(out_dtype), axis=k)) + + kk = tvm.reduce_axis((0, vec), "kk") + C = tvm.compute((M, N), + lambda y, x: tvm.sum(CC[y, x, kk], axis=kk), + tag="dense_nopack") if bias is not None: - C = tvm.compute((batch, out_dim), lambda i, j: C[i, j] + bias[j].astype(out_dtype), + C = tvm.compute((M, N), lambda i, j: C[i, j] + bias[j].astype(out_dtype), tag=tag.BROADCAST) return C diff --git a/topi/src/topi.cc b/topi/src/topi.cc index 7e47b62a8af4..5649ad266d69 100644 --- a/topi/src/topi.cc +++ b/topi/src/topi.cc @@ -148,6 +148,11 @@ TVM_REGISTER_GLOBAL("topi.exp") *rv = exp(args[0]); }); +TVM_REGISTER_GLOBAL("topi.erf") +.set_body([](TVMArgs args, TVMRetValue *rv) { + *rv = erf(args[0]); + }); + TVM_REGISTER_GLOBAL("topi.cos") .set_body([](TVMArgs args, TVMRetValue *rv) { *rv = cos(args[0]); @@ -157,7 +162,6 @@ TVM_REGISTER_GLOBAL("topi.sin") .set_body([](TVMArgs args, TVMRetValue *rv) { *rv = sin(args[0]); }); - TVM_REGISTER_GLOBAL("topi.tanh") .set_body([](TVMArgs args, TVMRetValue *rv) { *rv = tanh(args[0]); diff --git a/topi/tests/python/test_topi_basic.py b/topi/tests/python/test_topi_basic.py index 3d3160ed7102..103cfc720715 100644 --- a/topi/tests/python/test_topi_basic.py +++ b/topi/tests/python/test_topi_basic.py @@ -36,6 +36,7 @@ def test_apply(func, name): assert B.op.body[0].name == name test_apply(topi.exp, "exp") + test_apply(topi.exp, "erf") test_apply(topi.tanh, "tanh") test_apply(topi.sigmoid, "sigmoid") test_apply(topi.log, "log") diff --git a/topi/tests/python/test_topi_math.py b/topi/tests/python/test_topi_math.py index a095733757f0..f5162e06d47a 100644 --- a/topi/tests/python/test_topi_math.py +++ b/topi/tests/python/test_topi_math.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. import numpy as np +import scipy import tvm import topi import topi.testing @@ -86,6 +87,7 @@ def check_device(device): test_apply(topi.rsqrt, "rsqrt", lambda x: np.ones_like(x) / np.sqrt(x), 0, 100, skip_name_check=True) test_apply(topi.cos, "cos", np.cos, -2.0*np.pi, 2.0*np.pi) test_apply(topi.sin, "sin", np.sin, -2.0*np.pi, 2.0*np.pi) + test_apply(topi.erf, "erf", scipy.special.erf, -.1, .1, dtype="float32") def test_cast(): From d196cccd43dc9736988963cb5ee683c9aaabc5d2 Mon Sep 17 00:00:00 2001 From: Haichen Shen Date: Tue, 2 Jul 2019 11:59:51 -0700 Subject: [PATCH 02/12] stop vectorization for erf --- src/pass/vectorize_loop.cc | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/pass/vectorize_loop.cc b/src/pass/vectorize_loop.cc index 988aef5195a5..005b74050ed6 100644 --- a/src/pass/vectorize_loop.cc +++ b/src/pass/vectorize_loop.cc @@ -263,11 +263,27 @@ class Vectorizer : public IRMutator { {cond, t, f}, op->call_type, op->func, op->value_index); } } + Expr MutateErf_(const Call* op, const Expr& e) { + Expr arg = this->Mutate(op->args[0]); + if (arg.type().is_vector()) { + need_scalarize_ = true; + return e; + } + if (arg.same_as(op->args[0])) { + return e; + } else { + return Call::make( + op->type, op->name, {arg}, op->call_type, op->func, op->value_index); + } + } // Call Expr Mutate_(const Call* op, const Expr& e) final { if (op->name == intrinsic::tvm_if_then_else) { return MutateIfThenElseExpr_(op, e); } + if (op->name == "erf") { + return MutateErf_(op, e); + } int lane = 0; Array new_args = MutateArray(op->args, &lane); From 33ba2cfa38b8d09795920d542602c03cb9fc0ded Mon Sep 17 00:00:00 2001 From: Haichen Shen Date: Fri, 2 Aug 2019 18:56:36 -0700 Subject: [PATCH 03/12] x --- python/tvm/relay/backend/interpreter.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tvm/relay/backend/interpreter.py b/python/tvm/relay/backend/interpreter.py index 79738e9c7189..687519282419 100644 --- a/python/tvm/relay/backend/interpreter.py +++ b/python/tvm/relay/backend/interpreter.py @@ -182,8 +182,8 @@ def _convert_args(self, expr, args, kwargs): if i < num_of_args: if kwargs.get(name): raise Exception( - "duplicate argument supplied in" + - "both positional args (at position: {0})," + + "duplicate argument supplied in " + + "both positional args (at position: {0}), " + "and keyword argument (with name: {1})".format(i, name)) else: cargs.append(kwargs[name]) From cae6f698ccef375b312063f82efccc62a1f3e61e Mon Sep 17 00:00:00 2001 From: Haichen Shen Date: Fri, 2 Aug 2019 22:02:54 -0700 Subject: [PATCH 04/12] cleanup --- python/tvm/relay/frontend/mxnet.py | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/python/tvm/relay/frontend/mxnet.py b/python/tvm/relay/frontend/mxnet.py index 576f979a5b8e..aa1ffafc37ff 100644 --- a/python/tvm/relay/frontend/mxnet.py +++ b/python/tvm/relay/frontend/mxnet.py @@ -762,19 +762,6 @@ def _mx_rnn_param_concat(inputs, _): return [inputs] -def _mx_contrib_div_sqrt_dim(inputs, _): - assert len(inputs) == 1 - ndim = len(_infer_type(inputs[0]).checked_type.shape) - dim = _op.take(_op.shape_of(inputs[0]), _expr.const(ndim-1, dtype="int32")) - sqrt_dim = _op.sqrt(dim.astype('float32')) - out = inputs[0] / sqrt_dim - return out - - -def _mx_sequence_mask(inputs, attrs): - return inputs[0] - - def _mx_rnn_layer(inputs, attrs): def _rnn_cell(data, states, i2h_weight, h2h_weight, i2h_bias, h2h_bias, activation): i2h = _op.nn.bias_add(_op.nn.dense(data, i2h_weight), i2h_bias, axis=-1) From c282a8be5a7591a2fcd2f9385007791da22d135b Mon Sep 17 00:00:00 2001 From: Haichen Shen Date: Fri, 2 Aug 2019 22:10:52 -0700 Subject: [PATCH 05/12] fix --- python/tvm/relay/backend/interpreter.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/python/tvm/relay/backend/interpreter.py b/python/tvm/relay/backend/interpreter.py index 687519282419..cf7516c760ed 100644 --- a/python/tvm/relay/backend/interpreter.py +++ b/python/tvm/relay/backend/interpreter.py @@ -170,8 +170,8 @@ def _convert_args(self, expr, args, kwargs): return args if kwargs and not isinstance(expr, Function): - raise Exception("can only supply keyword parameters for a \ - relay.Function, found {0}".format(expr)) + raise Exception("can only supply keyword parameters for a " + "relay.Function, found {0}".format(expr)) params = expr.params param_names = [p.name_hint for p in params] @@ -182,16 +182,16 @@ def _convert_args(self, expr, args, kwargs): if i < num_of_args: if kwargs.get(name): raise Exception( - "duplicate argument supplied in " + - "both positional args (at position: {0}), " + + "duplicate argument supplied in " + "both positional args (at position: {0}), " "and keyword argument (with name: {1})".format(i, name)) else: cargs.append(kwargs[name]) if len(cargs) != len(params): raise Exception( - "insufficient arguments, expected" \ - " {0}, provided {1}".format(len(cargs), len(params))) + "insufficient arguments, expected " + "{0}, provided {1}".format(len(cargs), len(params))) return tuple(cargs) From 1a58535f73b5dad35c4b4772b3181f41ecdfa69b Mon Sep 17 00:00:00 2001 From: Haichen Shen Date: Tue, 3 Sep 2019 20:14:57 -0700 Subject: [PATCH 06/12] add whitelist for vectorizable intrin --- include/tvm/ir.h | 6 ++++ python/tvm/relay/frontend/mxnet.py | 2 +- src/lang/ir.cc | 15 ++++++++ src/pass/vectorize_loop.cc | 52 +++++++++++++++------------- topi/tests/python/test_topi_basic.py | 2 +- 5 files changed, 50 insertions(+), 27 deletions(-) diff --git a/include/tvm/ir.h b/include/tvm/ir.h index a9c6c4b10e28..994570e7f9df 100644 --- a/include/tvm/ir.h +++ b/include/tvm/ir.h @@ -556,6 +556,9 @@ class Call : public ExprNode { name == intrin_name); } + /*! \return Whether call node can be vectorized. */ + bool is_vectorizable() const; + static constexpr const char* _type_key = "Call"; TVM_DECLARE_NODE_TYPE_INFO(Call, ExprNode); @@ -571,6 +574,9 @@ class Call : public ExprNode { static constexpr const char* likely = "likely"; static constexpr const char* glsl_texture_store = "glsl_texture_store"; static constexpr const char* prefetch = "prefetch"; + + /*! \brief Vectorizable intrinsic list. */ + static const char* vectorizable_intrinsics[]; }; /*! diff --git a/python/tvm/relay/frontend/mxnet.py b/python/tvm/relay/frontend/mxnet.py index aa1ffafc37ff..e496881e5f22 100644 --- a/python/tvm/relay/frontend/mxnet.py +++ b/python/tvm/relay/frontend/mxnet.py @@ -507,7 +507,7 @@ def _mx_arange(inputs, attrs): new_attrs = {} new_attrs["start"] = _expr.const(attrs.get_float("start", 0.0)) stop = attrs.get_str("stop", "None") - new_attrs["stop"] = None if stop == "None" else _expr.const(eval(stop)) + new_attrs["stop"] = None if stop == "None" else _expr.const(float(stop)) new_attrs["step"] = _expr.const(attrs.get_float("step", 1.0)) new_attrs["dtype"] = attrs.get_str("dtype", "float32") return _op.arange(**new_attrs) diff --git a/src/lang/ir.cc b/src/lang/ir.cc index 628a75f08b85..760aedb593c2 100644 --- a/src/lang/ir.cc +++ b/src/lang/ir.cc @@ -176,6 +176,21 @@ Expr Let::make(Var var, Expr value, Expr body) { return Expr(node); } +const char* Call::vectorizable_intrinsics[] = { + "floor", "ceil", "sign", "trunc", "abs", "round", "exp", "tanh", "sqrt", + "log", "sin", "cos", "likely" +}; + +bool Call::is_vectorizable() const { + size_t cnt = sizeof(Call::vectorizable_intrinsics) / sizeof(char*); + for (size_t i = 0; i < cnt; ++i) { + if (name == Call::vectorizable_intrinsics[i]) { + return true; + } + } + return false; +} + Expr Call::make(DataType type, std::string name, Array args, diff --git a/src/pass/vectorize_loop.cc b/src/pass/vectorize_loop.cc index 005b74050ed6..aed92d692e53 100644 --- a/src/pass/vectorize_loop.cc +++ b/src/pass/vectorize_loop.cc @@ -263,37 +263,39 @@ class Vectorizer : public IRMutator { {cond, t, f}, op->call_type, op->func, op->value_index); } } - Expr MutateErf_(const Call* op, const Expr& e) { - Expr arg = this->Mutate(op->args[0]); - if (arg.type().is_vector()) { - need_scalarize_ = true; - return e; - } - if (arg.same_as(op->args[0])) { - return e; - } else { - return Call::make( - op->type, op->name, {arg}, op->call_type, op->func, op->value_index); - } - } // Call Expr Mutate_(const Call* op, const Expr& e) final { if (op->name == intrinsic::tvm_if_then_else) { return MutateIfThenElseExpr_(op, e); } - if (op->name == "erf") { - return MutateErf_(op, e); - } - int lane = 0; - Array new_args = MutateArray(op->args, &lane); - - // normal code path. - if (op->args.same_as(new_args)) { - return e; + if (!op->is_vectorizable()) { + // Cannot vectorize this op + Array new_args; + for (auto arg : op->args) { + auto new_arg = this->Mutate(arg); + if (new_arg.type().is_vector()) { + need_scalarize_ = true; + return e; + } + new_args.push_back(new_arg); + } + if (op->args.same_as(new_args)) { + return e; + } else { + return Call::make( + op->type, op->name, new_args, op->call_type, op->func, op->value_index); + } } else { - return Call::make( - op->type.with_lanes(lane), op->name, new_args, - op->call_type, op->func, op->value_index); + int lane = 0; + Array new_args = MutateArray(op->args, &lane); + // normal code path. + if (op->args.same_as(new_args)) { + return e; + } else { + return Call::make( + op->type.with_lanes(lane), op->name, new_args, + op->call_type, op->func, op->value_index); + } } } // Load diff --git a/topi/tests/python/test_topi_basic.py b/topi/tests/python/test_topi_basic.py index 103cfc720715..e04979cd5d78 100644 --- a/topi/tests/python/test_topi_basic.py +++ b/topi/tests/python/test_topi_basic.py @@ -36,7 +36,7 @@ def test_apply(func, name): assert B.op.body[0].name == name test_apply(topi.exp, "exp") - test_apply(topi.exp, "erf") + test_apply(topi.erf, "erf") test_apply(topi.tanh, "tanh") test_apply(topi.sigmoid, "sigmoid") test_apply(topi.log, "log") From 0ed579e6867a7585917f102be7fb71b52c4e226c Mon Sep 17 00:00:00 2001 From: Haichen Shen Date: Tue, 3 Sep 2019 20:23:08 -0700 Subject: [PATCH 07/12] add tf converter --- python/tvm/relay/frontend/tensorflow.py | 1 + tests/python/frontend/tensorflow/test_forward.py | 9 +++++++++ 2 files changed, 10 insertions(+) diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index 56cd652d739b..6ef95d4938f7 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -1258,6 +1258,7 @@ def _impl(inputs, attr, params): 'DepthToSpace' : _depth_to_space(), 'Equal' : _broadcast('equal'), 'Elu' : _elu(), + 'Erf' : AttrCvt('erf'), 'Exp' : AttrCvt('exp'), 'ExpandDims' : _expand_dims(), 'Fill' : _fill(), diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index 176ff4f7cdd0..43a18094e678 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -1844,6 +1844,14 @@ def test_forward_zeros_like(): _test_forward_zeros_like((2, 3, 11), "float32") _test_forward_zeros_like((2, 3, 11), "float64") +def test_forward_erf(): + ishape = (1, 3, 10, 10) + inp_array = np.random.uniform(-5, 5, size=ishape).astype(np.float32) + with tf.Graph().as_default(): + in1 = tf.placeholder(shape=inp_array.shape, dtype=inp_array.dtype) + tf.math.erf(in1) + compare_tf_with_tvm(inp_array, 'Placeholder:0', 'Erf:0') + def _test_forward_reverse_v2(in_shape, axis, dtype): np_data = np.random.uniform(-10, 10, size=in_shape).astype(dtype) tf.reset_default_graph() @@ -2244,6 +2252,7 @@ def test_forward_one_hot(): test_forward_log_softmax() test_forward_bias_add() test_forward_zeros_like() + test_forward_erf() # Reductions test_forward_argminmax() From a828ffbbb6a5d0721017d6f92557b02523c4deec Mon Sep 17 00:00:00 2001 From: Haichen Shen Date: Tue, 3 Sep 2019 21:30:22 -0700 Subject: [PATCH 08/12] fix dense --- topi/python/topi/x86/dense.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/topi/python/topi/x86/dense.py b/topi/python/topi/x86/dense.py index 986330053e0c..27cadfa04cce 100644 --- a/topi/python/topi/x86/dense.py +++ b/topi/python/topi/x86/dense.py @@ -32,7 +32,7 @@ def _declaration_dense(cfg, data, weight, bias=None, out_dtype=None): if "cblas" in target.libs: C = cblas.matmul(data, weight, False, True) if bias is not None: - C = tvm.compute((M, N), lambda i, j: C[i, j] + bias[j].astype(out_dtype), + C = tvm.compute(C.shape, lambda i, j: C[i, j] + bias[j].astype(out_dtype), tag=tag.BROADCAST) return C From af79ade9276df9d5df399b5dc385f296a6ba901f Mon Sep 17 00:00:00 2001 From: Haichen Shen Date: Wed, 4 Sep 2019 09:18:56 -0700 Subject: [PATCH 09/12] fix --- src/lang/ir.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lang/ir.cc b/src/lang/ir.cc index 760aedb593c2..dbb64bc786b7 100644 --- a/src/lang/ir.cc +++ b/src/lang/ir.cc @@ -178,7 +178,7 @@ Expr Let::make(Var var, Expr value, Expr body) { const char* Call::vectorizable_intrinsics[] = { "floor", "ceil", "sign", "trunc", "abs", "round", "exp", "tanh", "sqrt", - "log", "sin", "cos", "likely" + "log", "sin", "cos", "popcount", "likely" }; bool Call::is_vectorizable() const { From d2ba40262580d3e1c360d25a5473c7e451273ddb Mon Sep 17 00:00:00 2001 From: Haichen Shen Date: Wed, 4 Sep 2019 10:16:04 -0700 Subject: [PATCH 10/12] add missing intrin --- src/lang/ir.cc | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/lang/ir.cc b/src/lang/ir.cc index dbb64bc786b7..e66d21bba1f7 100644 --- a/src/lang/ir.cc +++ b/src/lang/ir.cc @@ -177,8 +177,9 @@ Expr Let::make(Var var, Expr value, Expr body) { } const char* Call::vectorizable_intrinsics[] = { - "floor", "ceil", "sign", "trunc", "abs", "round", "exp", "tanh", "sqrt", - "log", "sin", "cos", "popcount", "likely" + "floor", "ceil", "sign", "trunc", "fabs", "round", "exp", "tanh", "sqrt", + "log", "sin", "cos", "pow", ir::Call::shift_left, ir::Call::shift_right, + ir::Call::likely, ir::Call::popcount }; bool Call::is_vectorizable() const { From d5aed24ba07b86c0d4eaabb65d374928a314d4f8 Mon Sep 17 00:00:00 2001 From: Haichen Shen Date: Wed, 4 Sep 2019 11:18:17 -0700 Subject: [PATCH 11/12] fix mxnet frontend --- python/tvm/relay/frontend/mxnet.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/tvm/relay/frontend/mxnet.py b/python/tvm/relay/frontend/mxnet.py index e496881e5f22..e5940f950da0 100644 --- a/python/tvm/relay/frontend/mxnet.py +++ b/python/tvm/relay/frontend/mxnet.py @@ -56,7 +56,6 @@ def _mx_fully_connected(inputs, attrs): if has_flatten and use_flatten: inputs[0] = _op.nn.batch_flatten(inputs[0]) data_shape = _infer_type(inputs[0]).checked_type.shape - weight_shape = _infer_type(inputs[1]).checked_type.shape if len(data_shape) > 2: inputs[0] = _op.reverse_reshape(inputs[0], [-1, 0]) res = _op.nn.dense(inputs[0], inputs[1], units=units) @@ -65,7 +64,7 @@ def _mx_fully_connected(inputs, attrs): res = _op.nn.bias_add(res, inputs[2], axis=-1) if len(data_shape) > 2: new_shape = data_shape[:-1] - new_shape.append(weight_shape[0]) + new_shape.append(units) res = _op.reshape(res, new_shape) return res From 8b876fc2eff95de09489325695b5cb95d9efbdba Mon Sep 17 00:00:00 2001 From: Haichen Shen Date: Thu, 5 Sep 2019 02:58:23 +0000 Subject: [PATCH 12/12] fix nvptx --- src/codegen/llvm/intrin_rule_nvptx.cc | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/codegen/llvm/intrin_rule_nvptx.cc b/src/codegen/llvm/intrin_rule_nvptx.cc index b23c0e8f77a4..598b575b5205 100644 --- a/src/codegen/llvm/intrin_rule_nvptx.cc +++ b/src/codegen/llvm/intrin_rule_nvptx.cc @@ -64,6 +64,9 @@ TVM_REGISTER_GLOBAL("tvm.intrin.rule.nvptx.fabs") TVM_REGISTER_GLOBAL("tvm.intrin.rule.nvptx.exp") .set_body(DispatchExternLibDevice); +TVM_REGISTER_GLOBAL("tvm.intrin.rule.nvptx.erf") +.set_body(DispatchExternLibDevice); + TVM_REGISTER_GLOBAL("tvm.intrin.rule.nvptx.fma") .set_body(DispatchExternLibDevice);