This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 3b8aa040e4 [Unity][op] Relax rsqrt and sinh (#14479)
3b8aa040e4 is described below
commit 3b8aa040e46fc849a747eeef213296ede0580736
Author: Yong Wu <[email protected]>
AuthorDate: Tue Apr 4 06:22:46 2023 -0700
[Unity][op] Relax rsqrt and sinh (#14479)
---
python/tvm/relax/frontend/torch/fx_translator.py | 5 +-
python/tvm/relax/op/unary.py | 24 ++++
python/tvm/relax/transform/legalize_ops/unary.py | 2 +
python/tvm/script/ir_builder/relax/ir.py | 2 +
src/relax/op/tensor/unary.cc | 1 +
src/relax/op/tensor/unary.h | 3 +
tests/python/relax/test_frontend_from_fx.py | 5 +-
tests/python/relax/test_op_unary.py | 2 +
.../relax/test_transform_legalize_ops_unary.py | 146 +++++++++++++++++++++
.../relax/test_tvmscript_parser_op_arith_cmp.py | 1 +
10 files changed, 184 insertions(+), 7 deletions(-)
diff --git a/python/tvm/relax/frontend/torch/fx_translator.py
b/python/tvm/relax/frontend/torch/fx_translator.py
index 58c7563419..f1ffbeac4b 100644
--- a/python/tvm/relax/frontend/torch/fx_translator.py
+++ b/python/tvm/relax/frontend/torch/fx_translator.py
@@ -157,10 +157,7 @@ class TorchFXImporter:
arg = self.env[node.args[0]]
if isinstance(arg, (int, float)):
arg = relax.const(arg, "float32")
- sqrt = self.block_builder.emit(relax.op.sqrt(arg))
- return self.block_builder.emit(
- relax.op.divide(relax.const(1, sqrt.struct_info.dtype), sqrt)
- )
+ return self.block_builder.emit(relax.op.rsqrt(arg))
def _round(self, node: fx.node.Node) -> relax.Expr:
if "decimals" in node.kwargs and node.kwargs["decimals"] != 0:
diff --git a/python/tvm/relax/op/unary.py b/python/tvm/relax/op/unary.py
index 866d2a8273..f885c103e7 100644
--- a/python/tvm/relax/op/unary.py
+++ b/python/tvm/relax/op/unary.py
@@ -303,6 +303,30 @@ def round(x: Expr) -> Expr:
return _ffi_api.round(x) # type: ignore
+def rsqrt(x: Expr) -> Expr:
+ """Compute element-wise reciprocal square root of the input data.
+
+ .. math::
+
+ 1/sqrt(x)
+
+ Parameters
+ ----------
+ x : relax.Expr
+ The input data
+
+ Returns
+ -------
+ result : relax.Expr
+ The computed result.
+
+ Note
+ ----
+ The input tensor is required to have float dtype
+ """
+ return _ffi_api.rsqrt(x) # type: ignore
+
+
def sigmoid(x: Expr) -> Expr:
"""Compute element-wise sigmoid of the input data.
diff --git a/python/tvm/relax/transform/legalize_ops/unary.py
b/python/tvm/relax/transform/legalize_ops/unary.py
index ca84cbf0ad..9d19abe16e 100644
--- a/python/tvm/relax/transform/legalize_ops/unary.py
+++ b/python/tvm/relax/transform/legalize_ops/unary.py
@@ -28,8 +28,10 @@ register_legalize("relax.exp",
_call_topi_without_attr(topi.exp, "tir_exp"))
register_legalize("relax.floor", _call_topi_without_attr(topi.floor,
"tir_floor"))
register_legalize("relax.negative", _call_topi_without_attr(topi.negative,
"tir_negative"))
register_legalize("relax.round", _call_topi_without_attr(topi.round,
"tir_round"))
+register_legalize("relax.rsqrt", _call_topi_without_attr(topi.rsqrt,
"tir_rsqrt"))
register_legalize("relax.sigmoid", _call_topi_without_attr(topi.sigmoid,
"tir_sigmoid"))
register_legalize("relax.sign", _call_topi_without_attr(topi.sign, "tir_sign"))
+register_legalize("relax.sinh", _call_topi_without_attr(topi.sinh, "tir_sinh"))
register_legalize("relax.sin", _call_topi_without_attr(topi.sin, "tir_sin"))
register_legalize("relax.sqrt", _call_topi_without_attr(topi.sqrt, "tir_sqrt"))
register_legalize("relax.tanh", _call_topi_without_attr(topi.tanh, "tir_tanh"))
diff --git a/python/tvm/script/ir_builder/relax/ir.py
b/python/tvm/script/ir_builder/relax/ir.py
index c6d36c8587..9857853a80 100644
--- a/python/tvm/script/ir_builder/relax/ir.py
+++ b/python/tvm/script/ir_builder/relax/ir.py
@@ -103,6 +103,7 @@ from tvm.relax.op import (
tensor_to_shape,
shape_to_tensor,
round,
+ rsqrt,
shape_of,
std,
strided_slice,
@@ -624,6 +625,7 @@ __all__ = [
"tensor_to_shape",
"shape_to_tensor",
"round",
+ "rsqrt",
"shape",
"shape_of",
"ShapeExpr",
diff --git a/src/relax/op/tensor/unary.cc b/src/relax/op/tensor/unary.cc
index f1117c1826..5d4a39067f 100644
--- a/src/relax/op/tensor/unary.cc
+++ b/src/relax/op/tensor/unary.cc
@@ -51,6 +51,7 @@ RELAX_REGISTER_UNARY_ARITH_OP_AND_IMPL(floor,
/*require_float_dtype=*/false);
RELAX_REGISTER_UNARY_ARITH_OP_AND_IMPL(log, /*require_float_dtype=*/true);
RELAX_REGISTER_UNARY_ARITH_OP_AND_IMPL(negative,
/*require_float_dtype=*/false);
RELAX_REGISTER_UNARY_ARITH_OP_AND_IMPL(round, /*require_float_dtype=*/false);
+RELAX_REGISTER_UNARY_ARITH_OP_AND_IMPL(rsqrt, /*require_float_dtype=*/true);
RELAX_REGISTER_UNARY_ARITH_OP_AND_IMPL(sigmoid, /*require_float_dtype=*/true);
RELAX_REGISTER_UNARY_ARITH_OP_AND_IMPL(sign, /*require_float_dtype=*/false);
RELAX_REGISTER_UNARY_ARITH_OP_AND_IMPL(sin, /*require_float_dtype=*/true);
diff --git a/src/relax/op/tensor/unary.h b/src/relax/op/tensor/unary.h
index 8f6404c5d9..60dd5148e2 100644
--- a/src/relax/op/tensor/unary.h
+++ b/src/relax/op/tensor/unary.h
@@ -100,6 +100,9 @@ Expr negative(Expr x);
/*! \brief Rounds each element of the input data to nearest integer. */
Expr round(Expr x);
+/*! \brief Compute element-wise reciprocal square root of data. */
+Expr rsqrt(Expr x);
+
/*! \brief Compute element-wise sigmoid of data. */
Expr sigmoid(Expr x);
diff --git a/tests/python/relax/test_frontend_from_fx.py
b/tests/python/relax/test_frontend_from_fx.py
index 9e07ff7b59..2285131e39 100644
--- a/tests/python/relax/test_frontend_from_fx.py
+++ b/tests/python/relax/test_frontend_from_fx.py
@@ -2685,9 +2685,8 @@ def test_rsqrt():
inp_0: R.Tensor((256, 256), dtype="float32")
) -> R.Tensor((256, 256), dtype="float32"):
with R.dataflow():
- lv: R.Tensor((256, 256), dtype="float32") = R.sqrt(inp_0)
- lv1: R.Tensor((256, 256), dtype="float32") =
R.divide(R.const(1, "float32"), lv)
- gv: R.Tensor((256, 256), dtype="float32") = lv1
+ lv: R.Tensor((256, 256), dtype="float32") = R.rsqrt(inp_0)
+ gv: R.Tensor((256, 256), dtype="float32") = lv
R.output(gv)
return gv
diff --git a/tests/python/relax/test_op_unary.py
b/tests/python/relax/test_op_unary.py
index 45336661a1..c8751dbaca 100644
--- a/tests/python/relax/test_op_unary.py
+++ b/tests/python/relax/test_op_unary.py
@@ -44,6 +44,7 @@ def test_op_correctness():
assert relax.op.log(x).op == Op.get("relax.log")
assert relax.op.negative(x).op == Op.get("relax.negative")
assert relax.op.round(x).op == Op.get("relax.round")
+ assert relax.op.rsqrt(x).op == Op.get("relax.rsqrt")
assert relax.op.sigmoid(x).op == Op.get("relax.sigmoid")
assert relax.op.sin(x).op == Op.get("relax.sin")
assert relax.op.sinh(x).op == Op.get("relax.sinh")
@@ -75,6 +76,7 @@ unary_arith_op, require_float_dtype = tvm.testing.parameters(
(relax.op.log, True),
(relax.op.negative, False),
(relax.op.round, False),
+ (relax.op.rsqrt, True),
(relax.op.sigmoid, True),
(relax.op.sign, False),
(relax.op.sin, True),
diff --git a/tests/python/relax/test_transform_legalize_ops_unary.py
b/tests/python/relax/test_transform_legalize_ops_unary.py
index 27103e5f8f..1a5d474c3e 100644
--- a/tests/python/relax/test_transform_legalize_ops_unary.py
+++ b/tests/python/relax/test_transform_legalize_ops_unary.py
@@ -676,6 +676,76 @@ def test_round_symbolic():
tvm.ir.assert_structural_equal(mod, Expected)
+def test_rsqrt():
+ # fmt: off
+ @tvm.script.ir_module
+ class Rsqrt:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3),
"float32"):
+ gv: R.Tensor((2, 3), "float32") = R.rsqrt(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3),
"float32"):
+ gv = R.call_tir(Expected.tir_rsqrt, (x,), R.Tensor((2, 3),
dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_rsqrt(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)),
"float32"), compute: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.rsqrt(rxplaceholder[i0_1, i1_1])
+ # fmt: on
+
+ mod = LegalizeOps()(Rsqrt)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_rsqrt_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Rsqrt:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"),
"float32"):
+ m = T.int64()
+ n = T.int64()
+ gv: R.Tensor((m, n), "float32") = R.rsqrt(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"),
"float32"):
+ m = T.int64()
+ n = T.int64()
+ gv = R.call_tir(Expected.tir_rsqrt, (x,), R.Tensor((m, n),
dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_rsqrt(var_rxplaceholder: T.handle, var_compute: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.int64()
+ n = T.int64()
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n],
dtype="float32")
+ compute = T.match_buffer(var_compute, [m, n], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.rsqrt(rxplaceholder[i0_1, i1_1])
+ # fmt: on
+
+ mod = LegalizeOps()(Rsqrt)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
def test_sigmoid():
# fmt: off
@tvm.script.ir_module
@@ -917,6 +987,82 @@ def test_sin_symbolic():
tvm.ir.assert_structural_equal(mod, Expected)
+def test_sinh():
+ # fmt: off
+ @tvm.script.ir_module
+ class Sinh:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3),
"float32"):
+ gv: R.Tensor((2, 3), "float32") = R.sinh(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3),
dtype="float32"):
+ cls = Expected
+ gv = R.call_tir(cls.tir_sinh, (x,), out_sinfo=R.Tensor((2, 3),
dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_sinh(
+ rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"),
+ compute: T.Buffer((T.int64(2), T.int64(3)), "float32"),
+ ):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("compute"):
+ v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[v_i0, v_i1])
+ T.writes(compute[v_i0, v_i1])
+ compute[v_i0, v_i1] = T.sinh(rxplaceholder[v_i0, v_i1])
+ # fmt: on
+
+ mod = LegalizeOps()(Sinh)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_sinh_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Sinh:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"),
"float32"):
+ m = T.int64()
+ n = T.int64()
+ gv: R.Tensor((m, n), "float32") = R.sinh(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(
+ x: R.Tensor(("m", "n"), dtype="float32")
+ ) -> R.Tensor(("m", "n"), dtype="float32"):
+ m = T.int64()
+ n = T.int64()
+ cls = Expected
+ gv = R.call_tir(cls.tir_sinh, (x,), out_sinfo=R.Tensor((m, n),
dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_sinh(var_rxplaceholder: T.handle, var_compute: T.handle):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ m, n = T.int64(), T.int64()
+ rxplaceholder = T.match_buffer(var_rxplaceholder, (m, n))
+ compute = T.match_buffer(var_compute, (m, n))
+ for i0, i1 in T.grid(m, n):
+ with T.block("compute"):
+ v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[v_i0, v_i1])
+ T.writes(compute[v_i0, v_i1])
+ compute[v_i0, v_i1] = T.sinh(rxplaceholder[v_i0, v_i1])
+ # fmt: on
+
+ mod = LegalizeOps()(Sinh)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
def test_sqrt():
# fmt: off
@tvm.script.ir_module
diff --git a/tests/python/relax/test_tvmscript_parser_op_arith_cmp.py
b/tests/python/relax/test_tvmscript_parser_op_arith_cmp.py
index d43e9a626b..a0c9218b2f 100644
--- a/tests/python/relax/test_tvmscript_parser_op_arith_cmp.py
+++ b/tests/python/relax/test_tvmscript_parser_op_arith_cmp.py
@@ -51,6 +51,7 @@ def _check(
(relax.op.log,),
(relax.op.negative,),
(relax.op.round,),
+ (relax.op.rsqrt,),
(relax.op.sigmoid,),
(relax.op.sign,),
(relax.op.sin,),
