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/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/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..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}), \ - 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]) 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) 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..e5940f950da0 100644 --- a/python/tvm/relay/frontend/mxnet.py +++ b/python/tvm/relay/frontend/mxnet.py @@ -55,10 +55,17 @@ 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 + 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(units) + res = _op.reshape(res, new_shape) return res @@ -241,8 +248,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 +258,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 +505,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(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) @@ -910,6 +919,7 @@ def _mx_one_hot(inputs, attrs): _identity_list = [ "log", "exp", + "erf", "sqrt", "floor", "ceil", 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/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/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); diff --git a/src/lang/ir.cc b/src/lang/ir.cc index 628a75f08b85..e66d21bba1f7 100644 --- a/src/lang/ir.cc +++ b/src/lang/ir.cc @@ -176,6 +176,22 @@ Expr Let::make(Var var, Expr value, Expr body) { return Expr(node); } +const char* Call::vectorizable_intrinsics[] = { + "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 { + 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 988aef5195a5..aed92d692e53 100644 --- a/src/pass/vectorize_loop.cc +++ b/src/pass/vectorize_loop.cc @@ -268,16 +268,34 @@ class Vectorizer : public IRMutator { if (op->name == intrinsic::tvm_if_then_else) { return MutateIfThenElseExpr_(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/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/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() 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..27cadfa04cce 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(C.shape, 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..e04979cd5d78 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.erf, "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():